Hoffman 2022 - MEADS Heuristics - Step Size, Damping, Drift and Eigenvalue Estimation

Summary

MEADS sets, from the neighbouring fold: per-dimension scales ${\hat{\sigma }}_d$; step size $ϵ=\min \{1,0.5/\sqrt{{\lambda }_{\max }(-\bar{H})}\}$ with $-\bar{H}$ estimated by the gradient outer product; damping $\gamma =\max \{1/(tϵ),1/\sqrt{{\lambda }_{\max }(\hat{\Sigma })}\}$ and $\alpha =1-e^{-2ϵ\gamma }$; and slice drift $\delta =\alpha /2$. Largest eigenvalues are estimated by the ratio $\text{tr}({\Sigma }^2)/\text{tr}(\Sigma )$, computable in $O(N^{2}D)$ without an eigendecomposition.

Overview

GHMC (via persistent slice sampling) has four tunable pieces: diagonal preconditioner, $ϵ$, $\alpha$, $\delta$. All are derived below from statistics of one fold, used for the neighbouring fold (see Hoffman 2022 - Ensemble-Chain Adaptation and Parallelism).

Main Content

4.1 Diagonal preconditioning

Estimate marginal standard deviations ${\hat{\sigma }}_d$ from a fold’s states and multiply the neighbour fold’s per-dimension step size by ${\hat{\sigma }}_d$. Subsequent formulas use transformed variables and gradients ${\bar{\theta }}_d={\theta }_d/{\hat{\sigma }}_d$, ${\bar{g}}_d={\nabla }_{{\theta }_d}\log p(\theta ){\hat{\sigma }}_d$.

4.2 Step size

Leapfrog stability on Gaussians

For a quadratic potential $\frac{1}{2}{\theta }^{\top }H\theta$ (Gaussian target with covariance $H^{-1}$) leapfrog stability requires $ϵ\le 2/\sqrt{{\lambda }_{\max }(H)}$ (Leimkuhler and Reich 2004). The accuracy of uncorrected underdamped Langevin MCMC likewise needs $ϵ\propto 1/\sqrt{{\lambda }_{\max }(-\bar{H})}$ (Dalalyan and Riou-Durand 2020). MEADS sets $ϵ:=\frac{1}{2}\frac{1}{\sqrt{{\lambda }_{\max }(-\bar{H})}}$, with $\bar{H}\triangleq \int p(\theta ){\nabla }^2\log p(\theta )d\theta$ the average Hessian (the $\frac{1}{2}$ is a safety margin; Appendix A shows low sensitivity).

Hessian from gradients (eq. 4; proof in Appendix E)

$-\int p(\theta )\frac{{\partial }^2\log p}{\partial \theta \partial {\theta }^{\top }}d\theta =\int p(\theta )(\frac{\partial \log p}{\partial \theta }{)}^{\top }\frac{\partial \log p}{\partial \theta }d\theta$, i.e. the expected negative Hessian equals the expected outer product of the score (the Fisher-information identity applied to the log-density), under mild conditions. Hence $-\hat{H}=\frac{1}{N}{\sum }_n{\bar{g}}_{k,n}{\bar{g}}_{k,n}^{\top }$ from the $N$ transformed gradients of fold $k$; no second derivatives needed.

4.3 Damping factor and slice drift

Momentum memory and damping

Ignoring leapfrog half-steps, $m_{t+i}\approx (1-\alpha {)}^{i/2}{m}_t+{\sum }_{j=1}^iϵ\nabla \log p({\theta }_{t+j})+\sqrt{\alpha }{\xi }_j$ with ${\xi }_j\sim \mathcal{N}(0,I)$ (eq. 5): the influence of $m_t$ decays exponentially in $i$. We want to forget momentum once the chain has had time to traverse the highest-variance direction. On a Gaussian with covariance $\Sigma$ leapfrog needs $O({\lambda }_{\max }(\Sigma )/ϵ)$ steps (in the paper’s notation; dimensionally this is the largest scale) to progress along the least-constrained direction, so choose $-\frac{1}{2}\log (1-\alpha )\propto ϵ/{\lambda }_{\max }$, i.e. for small $\alpha$

$$\gamma :=1/{\lambda }_{\max }(\hat{\Sigma }),\alpha =1-e^{-2ϵ\gamma }(\text{eq. 6}).$$

$\hat{\Sigma }$ is the empirical covariance of the neighbour fold’s (preconditioned) states $\bar{\theta }$. Note. Algorithm 3 line 9 implements $\gamma =1/\sqrt{\text{MAX\_EIG}(\bar{\theta })}$, i.e. the inverse of the largest scale rather than of the largest variance; eq. 6 as typeset omits the square root, so follow Algorithm 3.

Early-iteration damping floor

For stability enforce ${\gamma }_t\ge 1/(t{ϵ}_t)$ early on: if the optimal damping would have the chain forget momentum after more than $t$ steps and fewer than $t$ steps have elapsed, the covariance estimate is not yet trustworthy because chains arrive at the typical set at different rates; those differences must be damped away before a very small damping factor is safe.

Slice drift

Want $u$ stable long enough to travel far without rejecting, but not slower to mix than $\theta ,m$. Set $\delta :=\alpha /2$, giving $u$ a period of $4/\alpha$ steps, after which the decay factor $(1-\alpha {)}^{2/\alpha }\approx e^{-2}\approx 0.14$: the previous period’s momentum has mostly been forgotten.

4.4 Largest eigenvalue estimation

The largest eigenvalue of an unbiased estimate $\hat{\Sigma }$ can be a highly biased estimate of ${\lambda }_{\max }(\Sigma )$ (biased upward for small $N$).

Ratio estimator (eq. 7, Algorithm 2)

$\text{tr}({\Sigma }^2)/\text{tr}(\Sigma )={\sum }_d{\lambda }_d^2/{\sum }_d{\lambda }_d$. If all eigenvalues are 0 or ${\lambda }_{\max }$ this equals ${\lambda }_{\max }$; it approximates ${\lambda }_{\max }$ unless there are many “small-but-not-tiny” eigenvalues. For $X\in {\mathbb{R}}^{N\times D}$ with rows of covariance $\Sigma$: $S=X{X}^{\top }$; ${\lambda }_{\cdot }=\text{tr}(S)/N$; ${\lambda }_{\cdot }^2=\frac{1}{N(N-1)}{\sum }_{n\ne n^{\prime }}{S}_{n,n^{\prime }}^2$; return ${\lambda }_{\cdot }^2/{\lambda }_{\cdot }$. It uses $\mathbb{E}[S_{n,n^{\prime }}^2]=\mathbb{E}[x_n^{\top }{x}_{n^{\prime }}{x}_{n^{\prime }}^{\top }{x}_n]=\text{tr}(\mathbb{E}[X^{\top }X{]}^2)$ for $n\ne n^{\prime }$ (note $\mathbb{E}[X^{\top }X]=N\Sigma$ in the paper’s normalisation) and costs $O(N^{2}D)$, dominated by the $N$ gradient evaluations in practice.

Algorithm 3 (MEADS), in words

Inputs: $K$ folds of $N$ initial states. Initialise $u\sim \text{Uniform}(-1,1)$, $m\sim \mathcal{N}(0,I)$. For $t=\mathrm{1..}T$, for each fold $k\ne t\mathrm{mod}K$: estimate $\hat{\mu },\hat{\sigma }$ from fold $(k-1)\mathrm{mod}K$; standardise states and gradients ($\bar{\theta },\bar{g}$); ${ϵ}_{t,k}=\min \{1,0.5/\sqrt{\text{max\_eig}(\bar{g})}\}$; ${\gamma }_{t,k}=\max \{1/(t{ϵ}_{t,k}),1/\sqrt{\text{max\_eig}(\bar{\theta })}\}$; ${\alpha }_{t,k}=1-e^{-2{ϵ}_{t,k}{\gamma }_{t,k}}$, ${\delta }_{t,k}={\alpha }_{t,k}/2$; run pers_ghmc with step sizes ${ϵ}_{t,k}{\hat{\sigma }}_{t,k}$, damping ${\alpha }_{t,k}$, drift ${\delta }_{t,k}$ on fold $k$‘s chains.

Appendix C: behaviour of the estimator

$N\times D$ matrix $X$ with $X_{n,d}\sim \mathcal{N}(0,{\sigma }_d)$, $D=100$, four spectra each with $\max {\sigma }_d^2=1$: identity; log-spaced $\sigma$ from 0.1 to 1; linear-spaced from 0.1 to 1; bimodal (${\sigma }_d=0$ for $d\le D/2$, $1$ otherwise). Small $N$ biases both the naive top eigenvalue and the ratio upward but the ratio’s bias decays faster. The price is a small asymptotic bias for the log- and linear-spaced spectra: as $N\to \infty$ the ratio converges to $\sum {\sigma }_d^4/\sum {\sigma }_d^2$, a bit below ${\sigma }_{\max }^2$.

Asymptotes and a small-sample check (numpy)

Computed asymptotes: identity 1.000; log-spaced 0.517; linear-spaced 0.606; bimodal 1.000. For the log-spaced spectrum the ratio estimator gave 0.45, 0.53, 0.50, 0.52 at $N=16,64,256,1024$ while the naive top eigenvalue of $X^{\top }X/N$ gave 3.24, 1.81, 1.22, 1.13 against a true ${\lambda }_{\max }=1$.

import numpy as np
def max_eig(X):                 # Algorithm 2; X is (N, D)
    N = X.shape[0]; S = X @ X.T
    lam1 = np.trace(S)/N
    off = S**2; np.fill_diagonal(off, 0.0)
    return (off.sum()/(N*(N-1))) / lam1

Examples (PyMC-style)

Untested sketch for E86 using blackjax’s MEADS adaptation (names may differ by version):

import blackjax, jax
warmup = blackjax.meads_adaptation(logdensity_fn, num_chains=128)
(last_states, parameters), _ = warmup.run(jax.random.key(0), initial_positions, num_steps=2000)
# parameters hold per-chain step_size, momentum_inverse_scale, alpha, delta for the GHMC kernel

Connections

Hoffman 2022 - MEADS - Experiments and Ablations evaluates this recipe; Hoffman 2022 - Appendix - Hogwild Counterexample and Hessian Identity Proof proves eq. 4. Compare with ChEES’s trajectory-length adaptation in Hoffman 2021 - ChEES Criterion and Algorithm 1.

See Also

Hoffman 2022 - MEADS - Overview, Hoffman 2022 - GHMC and Neal Persistent Slice Sampling