Hoffman 2021 - ChEES Criterion and Algorithm 1

Summary

The ChEES criterion is $\frac{1}{4}\mathbb{E}[(\Vert {\theta }^{\prime }-\mathbb{E}\theta {\Vert }^2-\Vert \theta -\mathbb{E}\theta {\Vert }^2{)}^2]$: an ESJD of the centred, squared norm. It is shift/rotation invariant, dominated by high-variance directions and focused on second moments. Algorithm 1 ascends its gradient with Adam on $\log T$ using one jittered length per iteration shared by all chains, weights chains by acceptance probability, and tunes step size by dual averaging on the harmonic-mean acceptance probability (target 0.651).

Overview

Section 2.1 argued for minimising first-order autocorrelation of $((\theta -\mathbb{E}\theta {)}^{\top }v{)}^2$ with $v$ the highest-variance direction. We cannot know $v$, so ChEES is a heuristic satisfying three desiderata: (i) invariance to shifts and rotations, (ii) insensitivity to low-variance directions, (iii) focus on second-moment estimation.

Main Content

ChEES (eq. 7)

With $\theta \sim p$, $r\sim \mathcal{N}(0,I)$ at equilibrium and $({\theta }^{\prime },r^{\prime })={\text{leapfrog}}_{ϵ,L}(\theta ,r)$:

$$\text{ChEES}\triangleq \frac{1}{4}\mathbb{E}\left[\right(\Vert {\theta }^{\prime }-\mathbb{E}[\theta ]{\Vert }^2-\Vert \theta -\mathbb{E}[\theta ]{\Vert }^2{)}^2].$$

Centring gives shift invariance; dependence only on squared norms gives rotation invariance.

ChEES decomposition (eqs 8-9)

Take $\mathbb{E}[\theta ]=0$, $\mathbb{E}[{\theta }_i{\theta }_j]=0$ for $i\ne j$, and ${\sigma }_{\cdot }^2={\sum }_d{\sigma }_d^2$. Then

$$\text{ChEES}=\frac{1}{2}\mathbb{E}[(\Vert \theta {\Vert }^2-{\sigma }_{\cdot }^2{)}^2]-\frac{1}{2}\mathbb{E}[(\Vert {\theta }^{\prime }{\Vert }^2-{\sigma }_{\cdot }^2)(\Vert \theta {\Vert }^2-{\sigma }_{\cdot }^2)].$$

The first term is constant in $T$. The second equals ${\sum }_i(\mathbb{E}[{\theta }_i^4]-{\sigma }_i^4){\rho }_i+{\sum }_{i,j\ne i}\mathbb{E}[(({\theta }_i^{\prime }{)}^2-{\sigma }_i^2)({\theta }_j^2-{\sigma }_j^2)]$ where ${\rho }_i$ is the first-order autocorrelation of ${\theta }_i^2$. The cross terms are small if dimensions are approximately independent after rotation. The weights $\mathbb{E}[{\theta }_i^4]-{\sigma }_i^4$ (the variance of ${\theta }_i^2$) are generally proportional to ${\sigma }_i^4$, so the criterion aggressively weights high-variance directions, as desired.

Gradient with respect to $T$ (eq. 10)

$$\frac{d}{dT}{\int }_0^1{\mathbb{E}}_r\left[\frac{1}{4}\right(\Vert {\theta }^{\prime }(uT,\theta ,r)-\mathbb{E}\theta {\Vert }^2-\Vert \theta -\mathbb{E}\theta {\Vert }^2{)}^2]du={\int }_0^{1}u{\mathbb{E}}_r[(\Vert {\theta }^{\prime }(uT)-\mathbb{E}\theta {\Vert }^2-\Vert \theta -\mathbb{E}\theta {\Vert }^2)({\theta }^{\prime }(uT)-\mathbb{E}\theta {)}^{\top }{r}^{\prime }(uT)]du.$$

Sketch. Differentiate ${\theta }^{\prime }(uT)$ along Hamilton’s equations: $\partial {\theta }^{\prime }/\partial t=r^{\prime }$. The integral over $u$ is the Monte Carlo average over the uniform jitter $t=uT$.

Algorithm 1 (ChEES-HMC) in words

Inputs: unnormalised target, initial states ${\theta }_0^{(m)}$ for chains $m=\mathrm{1..}M$, initial step size ${ϵ}_0$, $N$ samples, $N_{\text{warmup}}$, Halton sequence $h_{1:N}$. Initialise $T_0={ϵ}_0$ and moving averages $\bar{T}=\overline{ϵ}=0$.

For $n=\mathrm{1..}N$:

1. Draw momenta $r^{(m)}\sim \mathcal{N}(0,I)$ and pick the jittered length $t_n=h_n{T}_{n-1}$.
2. Propose ${\tilde{\theta }}^{(m)},{\tilde{r}}^{(m)}={\text{leapfrog}}_{{ϵ}_{n-1},\lceil t_n/{ϵ}_{n-1}\rceil }({\theta }_{n-1}^{(m)},r^{(m)})$.
3. Accept with ${\alpha }_n^{(m)}=\min \{1,\exp (\mathcal{L}(\tilde{\theta })-\mathcal{L}({\theta }_{n-1}))\mathcal{N}(\tilde{r};0,I)/\mathcal{N}(r;0,I)\}$.
4. If $n<N_{\text{warmup}}$: compute the harmonic-mean acceptance $\bar{\alpha }=((1/M){\sum }_{m}1/{\alpha }^{(m)}{)}^{-1}$ and update ${ϵ}_n$ by dual averaging toward $\bar{\alpha }=0.651$; estimate means $\hat{\tilde{\theta }}=\frac{1}{M}\sum {\tilde{\theta }}^{(m)}$, $\hat{\theta }=\frac{1}{M}\sum {\theta }_{n-1}^{(m)}$; compute per-chain gradient estimates $${\hat{g}}^{(m)}=t_n(\Vert {\tilde{\theta }}^{(m)}-\hat{\tilde{\theta }}{\Vert }^2-\Vert {\theta }_{n-1}^{(m)}-\hat{\theta }{\Vert }^2)({\tilde{\theta }}^{(m)}-\hat{\tilde{\theta }}{)}^{\top }{\tilde{r}}^{(m)},$$ combine $\hat{g}={\sum }_m{\alpha }^{(m)}{\hat{g}}^{(m)}/{\sum }_m{\alpha }^{(m)}$, update $\log T_n$ with Adam, and set $\overline{ϵ}\leftarrow 0.9\overline{ϵ}+0.1{ϵ}_n$, $\bar{T}\leftarrow 0.9\bar{T}+0.1T_n$.
5. At $n=N_{\text{warmup}}$ freeze ${ϵ}_{n:N}=\overline{ϵ}$, $T_{n:N}=\bar{T}$.

Design choices and their rationale

• Adam on $\log T$ because the scale of $\hat{g}$ is problem dependent: step size $\alpha =0.025$ (not sensitive), ${\beta }_1=0$ (no momentum), ${\beta }_2=0.95$.
• One shared $t$ for all chains so every chain does the same work per iteration (simpler, faster code on SIMD).
• Weight by acceptance probability so obviously invalid states do not drive $\hat{g}$.
• Estimate $\mathbb{E}[\theta ]$ from the initial and final integrator states across chains.
• Harmonic-mean acceptance prevents a “straggler” chain stuck in a region requiring small steps from being ignored by an arithmetic average; if any chain has tiny $\alpha$ the harmonic mean pulls $ϵ$ down for all.
• Round up $t$ to an integer multiple of $ϵ$; adapt $T$ only during warmup to avoid bias from hyperparameters depending on recent states.
• Halton sequence (Owen 2017) instead of i.i.d. uniform jitter for more evenly spread lengths. It is determined a priori so the stationary distribution is unaffected.

Examples

Numpy sketch of ChEES-HMC on a 3-D Gaussian (tested)

Target $\mathcal{N}(0,\text{diag}(1,3,0.3{)}^2)$, 64 chains, 1500 iterations, 1000 warmup. In three seeds the adapted $T$ ended at about 6.0, 6.2 and 6.5, close to the $2.25{\sigma }_{\max }=6.75$ rule from Hoffman 2021 - Gaussian Autocorrelation Analysis and Trajectory Jitter (the match is approximate, not guaranteed; this is an illustration, not the paper’s code).

import numpy as np
def halton(i, b=2):
    f, r = 1.0, 0.0
    while i > 0: f /= b; r += f*(i % b); i //= b
    return r
# per iteration n, with theta (M,D), eps, T, adam state:
#   t = halton(n)*T ; L = ceil(t/eps) ; r ~ N(0,I)
#   th2, r2 = leapfrog(theta, r, eps, L, gradL)
#   a = min(1, exp(logp(th2)-logp(theta) - .5|r2|^2 + .5|r|^2))
#   g_m = t*(|th2-mean(th2)|^2 - |theta-mean(theta)|^2) * sum((th2-mean(th2))*r2, -1)
#   g = sum(a*g_m)/sum(a);  Adam step on log T with lr .025, b1=0, b2=.95
#   eps via dual averaging on harmonic mean of a (target .651)
#   accept per chain with prob a;  freeze eps,T at warmup end (use EMAs)

Examples (PyMC-style)

Motivating example E86 (examples page): ChEES-HMC is available in blackjax. Untested sketch:

import blackjax, jax, jax.numpy as jnp
# logdensity_fn: a JAX log-density over a flat position vector (e.g. from pymc.sampling.jax or hand-written)
warmup = blackjax.chees_adaptation(logdensity_fn, num_chains=128)
key = jax.random.key(0)
(last_states, parameters), _ = warmup.run(key, initial_positions, step_size=0.1, optim=None, num_steps=500)
kernel = blackjax.dynamic_hmc(logdensity_fn, **parameters).step   # then jax.vmap over chains

API names and arguments may differ across blackjax versions; check the installed docs before use.

Connections

Follows from Hoffman 2021 - Gaussian Autocorrelation Analysis and Trajectory Jitter. Results: Hoffman 2021 - Experiments - ESS per Gradient and GPU Throughput. MEADS, the next step, replaces trajectory-length tuning by tuning damping in generalised HMC: Hoffman 2022 - MEADS Heuristics - Step Size, Damping, Drift and Eigenvalue Estimation.

See Also

Hoffman 2021 - ChEES-HMC - Overview, Efficient MCMC