def update_fn(gradient, state, params=None): del params lr = step_size_fn(state.count) lr_sqrt = jnp.sqrt(lr) noise_std = jnp.sqrt(2 * (1 - momentum_decay)) preconditioner_state = preconditioner.update_preconditioner( gradient, state.preconditioner_state) noise, new_key = tree_utils.normal_like_tree(gradient, state.rng_key) noise = preconditioner.multiply_by_m_sqrt(noise, preconditioner_state) def update_momentum(m, g, n): return momentum_decay * m + g * lr_sqrt + n * noise_std momentum = jax.tree_map(update_momentum, state.momentum, gradient, noise) updates = preconditioner.multiply_by_m_inv(momentum, preconditioner_state) updates = jax.tree_map(lambda m: m * lr_sqrt, updates) return updates, OptaxSGLDState( count=state.count + 1, rng_key=new_key, momentum=momentum, preconditioner_state=preconditioner_state)
def sample_parms_fn(params, state): mean = params["mean"] std = jax.tree_map(jax.nn.softplus, params["inv_softplus_std"]) noise, new_key = tree_utils.normal_like_tree(mean, state["mfvi_key"]) params_sampled = jax.tree_multimap(lambda m, s, n: m + n * s, mean, std, noise) new_mfvi_state = { "net_state": copy.deepcopy(state["net_state"]), "mfvi_key": new_key } return params_sampled, new_mfvi_state
def adaptive_hmc_update(dataset, params, net_state, log_likelihood, state_grad, key, step_size, n_leapfrog_steps, target_accept_rate, step_size_adaptation_speed, do_mh_correction=True): normal_key, uniform_key, jitter_key = jax.random.split(key, 3) momentum, _ = tree_utils.normal_like_tree(params, normal_key) new_params, net_state, new_momentum, new_grad, new_log_likelihood = ( leapfrog(dataset, params, net_state, momentum, state_grad, step_size, n_leapfrog_steps)) accept_prob = get_accept_prob(log_likelihood, params, momentum, new_log_likelihood, new_params, new_momentum) accepted = jax.random.uniform(uniform_key) < accept_prob step_size = adapt_step_size(step_size, target_accept_rate, accept_prob, step_size_adaptation_speed) if do_mh_correction: params = jax.lax.cond(accepted, _first, _second, (new_params, params)) log_likelihood = jnp.where(accepted, new_log_likelihood, log_likelihood) state_grad = jax.lax.cond(accepted, _first, _second, (new_grad, state_grad)) else: params, log_likelihood, state_grad = (new_params, new_log_likelihood, new_grad) return (params, net_state, log_likelihood, state_grad, step_size, accept_prob, accepted)