Hoffman 2021 - Gaussian Autocorrelation Analysis and Trajectory Jitter

Summary

For HMC on a Gaussian with exact dynamics, first-order autocorrelation is $\cos (t/{\sigma }_d)$, which cannot be small in all directions for any fixed $t$. Uniform jitter $t\sim U(0,T)$ turns cosines into sincs. The second-moment autocorrelation is $\frac{1}{2}+\frac{1}{2}\text{sinc}(2T/{\sigma }_d)$, minimised near $T\approx 2.25{\sigma }_d$; using $T=2.25{\sigma }_{\max }$ bounds the worst-case first- and second-moment autocorrelations at 0.346 and 0.564.

Overview

The paper builds intuition on a $D$-dimensional zero-mean Gaussian with diagonal covariance ${\Sigma }_{dd}={\sigma }_d^2$, assuming $ϵ$ is small enough that leapfrog $\approx$ exact Hamiltonian dynamics. This case matters because “a sampler can hardly be called general-purpose if it fails on multivariate Gaussians” and, via the Bayesian central limit theorem, approximately Gaussian posteriors are common.

Main Content

Phase-space representation (eq. 2)

For dimension $d$ define the magnitude-angle pair by $a_d^2={\sigma }_d^{-2}{\theta }_d^2+r_d^2$ and ${\varphi }_d={\tan }^{-1}({\theta }_d/({\sigma }_d{r}_d))$, so $r_d=a_d\cos {\varphi }_d$ and ${\theta }_d=a_d{\sigma }_d\sin {\varphi }_d$. After trajectory length $t=ϵL$ the state is $r_d^{\prime }=a_d\cos ({\varphi }_d+t/{\sigma }_d)$, ${\theta }_d^{\prime }=a_d{\sigma }_d\sin ({\varphi }_d+t/{\sigma }_d)$. This solves Hamilton’s equations: each coordinate rotates in phase space with angular frequency $1/{\sigma }_d$.

At equilibrium $a_d^2\sim \text{Exponential}(\text{rate }1/2)$ (i.e. $\mathbb{E}[a_d^2]=2$), ${\varphi }_d\sim \text{Uniform}(0,2\pi )$, independent.

First-order autocorrelation of fixed-length HMC (eq. 3)

${\sigma }_d^{-2}\mathbb{E}[{\theta }_d{\theta }_d^{\prime }]=\mathbb{E}[a_d^2\sin {\varphi }_d\sin ({\varphi }_d+t/{\sigma }_d)]=\cos (t/{\sigma }_d)$. Proof sketch. $\sin \varphi \sin (\varphi +\delta )=\frac{1}{2}[\cos \delta -\cos (2\varphi +\delta )]$, $\mathbb{E}[a_d^2]=2$, and $\mathbb{E}[\cos (2{\varphi }_d+\delta )]=0$ for uniform ${\varphi }_d$.

Consequence (resonance). With many dimensions there is almost surely some $d$ with $t/{\sigma }_d\mathrm{mod}2\pi \approx 0$, so there is no $t$ that gives low autocorrelation in every direction (Neal 2011, sec. 3.2: this can interfere with ergodicity unless the trajectory length is randomised).

Jittered first-moment autocorrelation (eq. 4)

If $t\sim \text{Uniform}(0,T)$, the autocorrelation is $\frac{1}{T}{\int }_0^T\cos (t/{\sigma }_d)dt=\frac{\sin (T/{\sigma }_d)}{T/{\sigma }_d}$, which tends to 0 as $T/{\sigma }_d\to \infty$. Its minimum is about $-0.22$ at $T\approx 4.5{\sigma }_d$. If $T=4.5{\sigma }_{\max }$ then ${\max }_d\sin (T/{\sigma }_d)/(T/{\sigma }_d)\le 0.129$.

A chain with $T=4.5{\sigma }_{\max }$ is anti-correlated, good for first moments. But for nonlinear functionals, e.g. posterior variance, one needs the autocorrelation of ${\theta }_d^2$ (mean ${\sigma }^2$, variance $2{\sigma }^4$ at equilibrium).

Second-moment autocorrelation (eqs 5-6)

Fixed $t$: $\frac{1}{2{\sigma }_d^4}\mathbb{E}[(({\theta }_d^{\prime }{)}^2-{\sigma }_d^2)({\theta }_d^2-{\sigma }_d^2)]=\frac{1}{2}(1+\cos (2t/{\sigma }_d))$. Uniform jitter on $(0,T)$: $\frac{1}{2}+\frac{1}{2}\frac{\sin (2T/{\sigma }_d)}{2T/{\sigma }_d}$. This no longer goes to 0 but to $\frac{1}{2}$, and it decays twice as fast, with a minimum at $T\approx 2.25{\sigma }_d$. Taking $T:=2.25{\sigma }_{\max }$ gives worst-case first-moment autocorrelation $\le 0.346$ (eq. 4) and worst-case second-moment autocorrelation $\le 0.564$ (eq. 6).

Numerical check of the constants

The minimum of $\sin x/x$ is $-0.2172$ at $x\approx 4.493$; the minimum of $\frac{1}{2}+\frac{1}{2}\sin (2x)/(2x)$ is $0.391$ at $x\approx 2.247$; $\sin (2.25)/2.25=0.3458$; and the largest positive value of $\sin y/y$ beyond $y=4.5$ is $0.128$ (at $y\approx 7.725$), giving $\frac{1}{2}+\frac{1}{2}(0.128)=0.564$. These reproduce the paper’s 0.129, 0.346 and 0.564 up to rounding.

Figure 2 interpretation

Ten log-spaced scales $\sigma$. Top row: first moment, bottom row: second moment. Left: fixed $t$ (cosines with different periods; the worst-case black line is always near 1 somewhere). Right: jittered $T$ (sincs converging to a constant). Second moments raise the minimum achievable autocorrelation from $-1$ to $0$ (fixed $t$) and from about $-0.22$ to about $0.39$ (jittered).

Design lessons the paper draws.

1. To get low-variance estimators of many expectations at different posterior scales, jitter trajectory lengths and use a maximum length of the order of the largest posterior scale.
2. For small-scale dimensions (${\sigma }_d\ll {\sigma }_{\max }$) the second-moment autocorrelation is at least about $0.5$ whatever you do, so use the shortest, cheapest $T$ that gets there: $T=2.25{\sigma }_{\max }$.
3. Maximising plain ESJD yields trajectories up to twice as long as optimal, because ESJD is fooled by anti-correlated samples that help means but hurt variances. Maximising ESJD of the centred and squared parameters gives the shorter, efficient length. This is the origin of ChEES.

import numpy as np
x = np.linspace(1e-3, 30, 300000)
sinc = np.sin(x)/x                      # jittered first-moment autocorrelation, x = T/sigma
second = 0.5 + 0.5*np.sin(2*x)/(2*x)    # jittered second-moment autocorrelation
print(x[sinc.argmin()], sinc.min())     # ~4.49, -0.217
print(x[second.argmin()], second.min()) # ~2.25, 0.391

Examples

Sampling a $D=100$ Gaussian whose eigenvalues span five orders of magnitude (the paper’s ill-conditioned Gaussian target) exhibits exactly this resonance problem: fixed-$L$ HMC is “disastrous” with too little jitter (Hoffman 2021 - Jitter Ablations, Discussion and Many-Chain Practice).

Connections

Hoffman 2021 - HMC Background and the GPU Problem for NUTS sets up leapfrog; Hoffman 2021 - ChEES Criterion and Algorithm 1 turns the “squared, centred ESJD” idea into an adaptive algorithm. Related vault note: Efficient MCMC.

See Also

Hoffman 2021 - ChEES-HMC - Overview, Hoffman 2022 - GHMC and Neal Persistent Slice Sampling