R: Runs one step of Hamiltonian Monte Carlo.

mcmc_hamiltonian_monte_carlo {tfprobability}

R Documentation

Runs one step of Hamiltonian Monte Carlo.

Description

Hamiltonian Monte Carlo (HMC) is a Markov chain Monte Carlo (MCMC) algorithm that takes a series of gradient-informed steps to produce a Metropolis proposal. This class implements one random HMC step from a given current_state. Mathematical details and derivations can be found in Neal (2011).

Usage

mcmc_hamiltonian_monte_carlo(
  target_log_prob_fn,
  step_size,
  num_leapfrog_steps,
  state_gradients_are_stopped = FALSE,
  step_size_update_fn = NULL,
  seed = NULL,
  store_parameters_in_results = FALSE,
  name = NULL
)

Arguments

`target_log_prob_fn`	Function which takes an argument like `current_state` (if it's a list `current_state` will be unpacked) and returns its (possibly unnormalized) log-density under the target distribution.
`step_size`	`Tensor` or `list` of `Tensor`s representing the step size for the leapfrog integrator. Must broadcast with the shape of `current_state`. Larger step sizes lead to faster progress, but too-large step sizes make rejection exponentially more likely. When possible, it's often helpful to match per-variable step sizes to the standard deviations of the target distribution in each variable.
`num_leapfrog_steps`	Integer number of steps to run the leapfrog integrator for. Total progress per HMC step is roughly proportional to `step_size * num_leapfrog_steps`.
`state_gradients_are_stopped`	`logical` indicating that the proposed new state be run through `tf$stop_gradient`. This is particularly useful when combining optimization over samples from the HMC chain. Default value: `FALSE` (i.e., do not apply `stop_gradient`).
`step_size_update_fn`	Function taking current `step_size` (typically a `tf$Variable`) and `kernel_results` (typically `collections$namedtuple`) and returns updated step_size (`Tensor`s). Default value: `NULL` (i.e., do not update `step_size` automatically).
`seed`	integer to seed the random number generator.
`store_parameters_in_results`	If `TRUE`, then `step_size` and `num_leapfrog_steps` are written to and read from eponymous fields in the kernel results objects returned from `one_step` and `bootstrap_results`. This allows wrapper kernels to adjust those parameters on the fly. This is incompatible with `step_size_update_fn`, which must be set to `NULL`.
`name`	string prefixed to Ops created by this function. Default value: `NULL` (i.e., 'hmc_kernel').

Details

The one_step function can update multiple chains in parallel. It assumes that all leftmost dimensions of current_state index independent chain states (and are therefore updated independently). The output of target_log_prob_fn(current_state) should sum log-probabilities across all event dimensions. Slices along the rightmost dimensions may have different target distributions; for example, ⁠current_state[0, :]⁠ could have a different target distribution from ⁠current_state[1, :]⁠. These semantics are governed by target_log_prob_fn(current_state). (The number of independent chains is tf$size(target_log_prob_fn(current_state)).)

Value

a Monte Carlo sampling kernel