`blackjax.adaptation.chees_adaptation`#

Public API for ChEES-HMC

Module Contents#

Classes#

ChEESAdaptationState

State of the ChEES-HMC adaptation scheme.

Functions#

`base`(→ Tuple[Callable, Callable])	Maximizing the Change in the Estimator of the Expected Square criterion
`chees_adaptation`(→ blackjax.base.AdaptationAlgorithm)	Adapt the step size and trajectory length (number of integration steps / step size)

Attributes#

OPTIMAL_TARGET_ACCEPTANCE_RATE

OPTIMAL_TARGET_ACCEPTANCE_RATE = 0.651[source]#

class ChEESAdaptationState[source]#

State of the ChEES-HMC adaptation scheme.

step_size: Value of the step_size parameter of the HMC algorithm.
log_step_size_moving_average: Running moving average of the log step_size parameter.
trajectory_length: Value of the num_integration_steps * step_size parameter of the HMC algorithm.
log_trajectory_length_moving_average: Running moving average of the log num_integration_steps / step_size parameter.
optim_state: Optax optimizing state for used to maximize the ChEES criterion.
random_generator_arg: Utility array for generating a pseudo or quasi-random sequence of numbers.
step: Current iteration number.

step_size: float[source]#

log_step_size_moving_average: float[source]#

trajectory_length: float[source]#

log_trajectory_length_moving_average: float[source]#

da_state: blackjax.optimizers.dual_averaging.DualAveragingState[source]#

optim_state: optax.OptState[source]#

random_generator_arg: blackjax.types.Array[source]#

step: int[source]#

base(jitter_generator: Callable, next_random_arg_fn: Callable, optim: optax.GradientTransformation, target_acceptance_rate: float, decay_rate: float) → Tuple[Callable, Callable][source]#

Maximizing the Change in the Estimator of the Expected Square criterion (trajectory length) and dual averaging procedure (step size) for the jittered Hamiltonian Monte Carlo kernel [HRS21].

This adaptation algorithm tunes the step size and trajectory length, i.e. number of integration steps / step size, of the jittered HMC algorithm based on statistics collected from a population of many chains. It maximizes the Change in the Estimator of the Expected Square (ChEES) criterion to tune the trajectory length and uses dual averaging targeting an acceptance rate of 0.651 of the harmonic mean of the chain’s acceptance probabilities to tune the step size.

Parameters:

jitter_generator – Optional function that generates a value in [0, 1] used to jitter the trajectory lengths given a PRNGKey, used to propose the number of integration steps. If None, then a quasi-random Halton is used to jitter the trajectory length.
next_random_arg_fn – Function that generates the next random_generator_arg from its previous value.
optim – Optax compatible optimizer, which conforms to the optax.GradientTransformation protocol.
target_acceptance_rate – Average acceptance rate to target with dual averaging.
decay_rate – Float representing how much to favor recent iterations over earlier ones in the optimization of step size and trajectory length.

Returns:

init – Function that initializes the warmup.
update – Function that moves the warmup one step.

chees_adaptation(logdensity_fn: Callable, num_chains: int, *, jitter_generator: Callable | None = None, jitter_amount: float = 1.0, target_acceptance_rate: float = OPTIMAL_TARGET_ACCEPTANCE_RATE, decay_rate: float = 0.5) → blackjax.base.AdaptationAlgorithm[source]#

Adapt the step size and trajectory length (number of integration steps / step size) parameters of the jittered HMC algorthm.

The jittered HMC algorithm depends on the value of a step size, controlling the discretization step of the integrator, and a trajectory length, given by the number of integration steps / step size, jittered by using only a random percentage of this trajectory length.

This adaptation algorithm tunes the trajectory length by heuristically maximizing the Change in the Estimator of the Expected Square (ChEES) criterion over an ensamble of parallel chains. At equilibrium, the algorithm aims at eliminating correlations between target dimensions, making the HMC algorithm efficient.

Jittering requires generating a random sequence of uniform variables in [0, 1]. However, this adds another source of variance to the sampling procedure, which may slow adaptation or lead to suboptimal mixing. To alleviate this, rather than use uniform random noise to jitter the trajectory lengths, we use a quasi-random Halton sequence, which ensures a more even distribution of trajectory lengths.

Examples

An HMC adapted kernel can be learned and used with the following code:

warmup = blackjax.chees_adaptation(logdensity_fn, num_chains)
key_warmup, key_sample = jax.random.split(rng_key)
optim = optax.adam(learning_rate)
(last_states, parameters), _ = warmup.run(
    key_warmup,
    positions, #PyTree where each leaf has shape (num_chains, ...)
    initial_step_size,
    optim,
    num_warmup_steps,
)
kernel = blackjax.dynamic_hmc(logdensity_fn, **parameters).step
new_states, info = jax.vmap(kernel)(key_sample, last_states)