Focusing on Difficult Directions for Learning HMC Trajectory Lengths
Extracted source text ยท arXiv 2110.11576
\documentclass[twoside]{article}
\usepackage[arxiv]{aistats2022}
\usepackage[round]{natbib}
\renewcommand{\bibname}{References}
\renewcommand{\bibsection}{\subsubsection*{\bibname}}
\bibliographystyle{apalike}
\usepackage[title]{appendix}
\usepackage{hyperref}
\usepackage{booktabs}
\usepackage{url}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{array}
\usepackage{graphicx}
\usepackage[nameinlink]{cleveref}
\usepackage{algorithm}
\usepackage{algorithmicx}
\usepackage{algpseudocode}
\usepackage{listings}
\usepackage{xcolor}
\usepackage{courier}
\usepackage{makecell}
\usepackage{capt-of}
\newcommand{\bbR}{\mathbb{R}}
\newcommand{\bbE}{\mathbb{E}}
\newcommand{\tpi}{\tilde{\pi}}
\newcommand{\calC}{\mathcal{C}}
\newcommand{\lnorm}{\left\lVert}
\newcommand{\rnorm}{\right\rVert}
\newcommand{\argmin}{\mathrm{argmin}}
\newcommand{\unif}{\mathrm{Unif}}
\newcommand{\argmax}{\mathrm{argmax}}
\crefname{appsection}{Appendix}{Appendices}
\algrenewcommand\algorithmicindent{0.5em}
\algnewcommand{\LineComment}[1]{\State \(\triangleright\) #1}
\algrenewcommand\algorithmicrequire{\textbf{Input:}}
\algrenewcommand\algorithmicensure{\textbf{Output:}}
\lstset{
belowcaptionskip=1\baselineskip,
breaklines=true,
xleftmargin=\parindent,
language=Python,
showstringspaces=false,
numbers=left,
basicstyle=\footnotesize\ttfamily\color{blue},
keywordstyle=\bfseries\color{green!40!black},
commentstyle=\itshape\color{purple!40!black},
identifierstyle=\color{black},
numberstyle=\tiny\color{gray},
stringstyle=\color{orange},
}
\begin{document}
\twocolumn[
\aistatstitle{Focusing on Difficult Directions for Learning HMC Trajectory Lengths}
\aistatsauthor{ Pavel Sountsov \And Matt D. Hoffman }
\aistatsaddress{ Google Research } ]
\begin{abstract}
Hamiltonian Monte Carlo (HMC) is a premier Markov Chain Monte Carlo (MCMC) algorithm for continuous target distributions. Its full potential can only be unleashed when its problem-dependent hyperparameters are tuned well. The adaptation of one such hyperparameter, trajectory length ($\tau$), has been closely examined by many research programs with the No-U-Turn Sampler (NUTS) coming out as the preferred method in 2011. A decade later, the evolving hardware profile has lead to the proliferation of personal and cloud based SIMD hardware in the form of Graphics and Tensor Processing Units (GPUs, TPUs) which are hostile to certain algorithmic details of NUTS. This has opened up a hole in the MCMC toolkit for an algorithm that can learn $\tau$ while maintaining good hardware utilization. In this work we build on recent advances along this direction and introduce SNAPER-HMC, a SIMD-accelerator-friendly adaptive-MCMC scheme for learning $\tau$. The algorithm maximizes an upper bound on per-gradient effective sample size along an estimated principal component. We empirically show that SNAPER-HMC is stable when combined with mass-matrix adaptation, and is tolerant of certain pathological target distribution covariance spectra while providing excellent long and short run sampling efficiency. We provide a complete implementation for continuous multi-chain adaptive HMC combining trajectory learning with standard step-size and mass-matrix adaptation in one turnkey inference package.
\end{abstract}
\section{Introduction}
Hamiltonian Monte Carlo (HMC) \citep{duane1987hybrid, neal2011mcmc} is an efficient Markov Chain Monte Carlo (MCMC) algorithm for generating approximate samples from a continuous probability distribution with density $\tpi(x)$ supported on $x \in \bbR^D$. For HMC, it is only necessary to be able to evaluate $\tpi(x)$ up to a constant factor and to be able to compute its gradient. Such densities are commonly found as posteriors over parameters in the context of Bayesian inference, where the gradients can be obtained via automatic differentiation if the probabilistic model is implemented in a suitable programming language.
HMC generates proposals for a related distribution $P(x) P(m)$ where $x$ and $m \in \bbR^D$ are identified as a fictitious ``position'' and ``momentum'' of a particle. $P(m)$ is typically chosen to be a zero-centered multivariate Gaussian with the log-density $\log \tpi(m) = -\frac{1}{2} m^T M^{-1} m$ where $M$ is the the covariance matrix, identified as the ``mass'' matrix, which is commonly chosen to be diagonal.
HMC generates a proposal in three steps. Given the current state $x_0$, a $m_0$ is sampled from $P(m)$. Then, Hamilton's equations: $\frac{\partial x}{\partial t} = \frac{\partial H}{\partial m}$, $\frac{\partial m}{\partial t} = -\frac{\partial H}{\partial x}$, where $H(x, m) \triangleq -\log \tpi(x) - \log \tpi(m)$ is the Hamiltonian, are solved for the interval $t = [0, \tau]$ ($\tau$ is commonly known as the trajectory length). Typically an analytic solution is not available, so the dynamics are discretized and a suitable sympletic integrator, commonly the leapfrog integrator, is used with step size parameter $\epsilon$. The number of leapfrog steps $L$ is related to $\tau$ as $L = \lceil \tau / \epsilon \rceil$. The proposed state $(x_\tau, -m_\tau)$ (the momentum is negated to ensure a reversible proposal) is accepted or rejected using the Metropolis-Hastings step \citep{metropolis1953equation, hastings1970mh} that targets the joint $P(x) P(m)$ distribution.
A typical use case of MCMC is to estimate expectations of a test function $f(x)$, i.e. $\theta_f \triangleq \bbE[f(x)]$. The MCMC estimator of $\theta_f$ can be constructed as $\hat{\theta}_f \triangleq \frac{1}{N} \sum_{t=0}^N f(x^{(t)})$ where $x^{(t)}$ is the $t$'th iteration of the MCMC chain. When the MCMC central limit theorem applies, the variance of this estimator can be computed as $Var[\hat{\theta}_f] = \frac{Var[\theta_f]}{ESS}$ where $ESS$ is the Effective Sample Size. A convenient way of expressing ESS is in terms of the autocorrelation function $\rho_f(l) \triangleq \frac{\bbE[((f(x^{(t)}) - \theta_f)(f(x^{(t + l)}) - \theta_f]}{Var[f(x)]}$:
\begin{align}
\label{eq:ess}
ESS_f \triangleq \frac{N}{1 + 2 \sum_{l=1}^{\infty} \rho_f(l)}.
\end{align}
$ESS/grad$ is a typical measure of efficiency, where $grad$ is the number of gradient evaluations needed to compute the estimator. It is typical to assume that the gradient evaluations dominate the wallclock runtime of algorithm. For HMC using the standard leapfrog integrator the number of gradient evaluations is $L$.
The efficiency of the HMC sampler strongly depends on the settings of its hyperparameters $\tau$, $\epsilon$ and $M$. $M$ and $\epsilon$ are typically set either continuously via Adaptive MCMC \citep{andrieu2008tutorial} or windowed adaptation popularized by \citet{carpenter2017stan}. In both cases, $\epsilon$ is adjusted such that the Metropolis-Hastings acceptance probability is between 0.6 and 0.9 \citep{betancourt2014optimizing}. $M$ is typically set to the diagonal sample covariance. How to set $\tau$ well is less obvious. \citet{hoffman2011no} proposed a variant of HMC called the No-U-Turn Sampler (NUTS) that dynamically chooses $\tau$ without breaking detailed balance. This allows it to adapt to local geometry of the distribution, although the requirement for detailed balance requires the algorithm to discard on average half of considered states when generating a proposal. This combination of adaptive methods is the standard way of using HMC, but it is not well adapted to the modern modern SIMD-based accelerators. A core issue is that those accelerators prefer to evaluate multiple chains in parallel, and when using NUTS, different chains may choose different numbers of leapfrog steps. SIMD implementations of NUTS exist \citep{radul2019automatically, phan2019iterative, lao2019unrolled} but all suffer from significant wall-clock and utilization overhead compared to a non-dynamic HMC \citep{radul2019automatically,hoffman2021adaptive}.
Alternate ways of setting $\tau$ are based on the Adaptive MCMC framework, e.g. Adaptive HMC \citep{wang2013adaptive}, eHMC \citep{wu2018faster} and ChEES-HMC \citep{hoffman2021adaptive}. They function by defining a criterion $\calC(\tau)$ as proxy for sampler efficiency and adjusting $\tau$ to maximize it. eHMC and Adaptive HMC optimize Expected Square Jump Distance (ESJD) \citep{pasarica2010adaptively} and $\mathrm{ESJD} / L^{0.5}$ respectively. ESJD is defined as $\bbE \left[ \lnorm x^{(0)} - x^{(t + 1)}(\tau)\rnorm_2^2 \right]$. ChEES-HMC maximizes the Change in the Estimator of the Expected Square (ChEES) which can be expressed by first transforming $x$ with:
\begin{align}
\label{eq:f_chees}
f_{ChEES}(x) = \frac{1}{2}\lnorm z \rnorm_2^2,
\end{align}
where $z \triangleq x - \bbE[x]$, before computing the generalized ESJD:
\begin{align}
\label{eq:esjd}
ESJD_f = \bbE \left[ \lnorm f(x^{(t)}) - f(x^{(t + 1)}(\tau)) \rnorm_2^2 \right].
\end{align}
\citet{hoffman2021adaptive} show that ChEES is superior to the standard ESJD for tuning HMC to maximize $ESS_f$ for a general $f$.
\begin{figure*}[!tb]
\centering
\includegraphics[scale=0.55]{figures/motivation.pdf}
\caption{
Shortcomings of the ChEES criterion. The non-decay of the
criterion with increasing number of leapfrog steps causes
gradient-based adaptation to sometimes overshoot the global optimum
(left). The ChEES-found number of leapfrog steps across 100 runs
are displayed as jittered points, with jitter amount proportional
to local point density. For target distributions where the
covariance has a large number of medium-sized eigenvalues (with
magnitude $\sigma_r$) and few large eigenvalues ($\sigma_0$),
ChEES-HMC finds trajectory lengths smaller than optimal (right),
causing a drop in ESS (middle). The jagged nature of the right plot arises from optimum jumping between peaks (similar to ones in the left panel) as they shift.
}
\label{fig:motivation}
\end{figure*}
ChEES has a number of shortcomings as a criterion. First, it does not incorporate the cost of generating a proposal. This is most evident when using a jittered trajectory length, which is a recommended practice to minimize resonances \citep{neal2011mcmc}. Left panel of \Cref{fig:motivation} plots ChEES versus $\bar{L} \triangleq \bbE[L]$ and the $\bar{L}$'s found by ChEES-HMC (100 separate runs). While most trials find the first peak of the criterion (corresponding to the global optimum for this problem), stochastic optimization can easily overshoot and find overlong $\bar{L}$s. This can be mitigated by the use of a more gentle learning rate at the cost of adaptation speed. This problem also exists, in principle, with regular ESJD that is used by NUTS, but there the algorithm is constructed to stop at the first peak it finds, which is not an easy thing to do with stochastic optimization. This issue is exacerbated if the $P(x) P(m)$ changes over the course of adaptation when $M$ is co-adapted.
The second shortcoming is that ChEES can be fooled by certain eigenspectra of the covariance matrix of $P(x)$ to choose a trajectory length that is too short. Consider a zero-centered multivariate Gaussian $P(x)$ with a diagonal covariance matrix set to $\{\sigma_0^2, \sigma_r^2 ... \sigma_r^2\}$, with the trailing $\sigma_r^2$ repeated 300 times and $\sigma_0 \ge \sigma_r$. For this problem, the optimal mean trajectory length when computing $\min_i ESS_{z_i^2}$ is $\tau^* \approx 1.12 \sigma_0$ (dashed red line, \Cref{fig:motivation}, right).
As we vary $\frac{\sigma_r}{\sigma_0}$ we compute the mean trajectory length $\hat{\tau}_{ChEES}$ that maximizes ChEES (\Cref{fig:motivation}, right panel) and $\min_i ESS_{z_i^2}$ evaluated at $\hat{\tau}_{ChEES}$ (\Cref{fig:motivation}, middle panel). While ChEES's optimum identifies the optimal trajectory length when the $\sigma_r$ is either very small (meaning they barely affect the value of ChEES) or very large (meaning that $P(x)$ is close to being isotropic), it struggles in the intermediate regime. There the repeated small eigenvalues overwhelm the contribution to ChEES from $\sigma_0$, causing ChEES to prefer shorter-than-optimal trajectories. This is caused by ChEES computing an average of per-dimension autocorrelation sums, each weighted by $\sigma_d^4$. While it heavily weights directions with greatest variance, the weighting can be overcome as illustrated here.
In this work we propose an alternative criterion, SNAPER, for learning $\tau$ using the Adaptive-MCMC framework that fixes the aforementioned issues with ChEES, culminating in the following contributions:
\begin{enumerate}
\item We show how SNAPER combined with the adaptation recipe produces a sampler that is efficient in terms of $ESS/grad$ as well as total number of gradient evaluations (including warmup period).
\item We show how to incorporate SNAPER into a unified, windowless, continuous adaptation recipe to learn the main hyperparameters of HMC.
\end{enumerate}
\section{SNAPER Criterion}
\label{sec:snaper}
A natural way to incorporate a notion of cost into ChEES is to normalize it by some increasing function of $L$. \cite{wang2013adaptive} optimize $ESJD / L^{0.5}$ with a heuristic motivation. Since $L$ is a discrete quantity it is not directly usable when using gradient-based adaptation, but it is straightforward to replace it with $\tau$ which, assuming $\epsilon$ is independent of $\tau$ (a reasonable assumption given the symplectic nature of the leapfrog integrator), is approximately proportional to $L$. We thus define the ChEES-Rate (ChEESR) criterion:
\begin{align*}
ChEESR \triangleq \frac{ESJD_{\lnorm z \rnorm_2^2}}{\tau}
\end{align*}
What does ChEESR optimize? It is easy to show that $ESJD_f \propto 1 - \rho_f(1)$. This identity has been used to motivate ESJD, as $ESS_f$ is related to $\rho_f(1)$ via a decreasing nonlinear function (\ref{eq:ess}). It's not obvious that this connections allows using $ESJD_f / \tau$ as a proxy for $ESS_f / grad$ however. To better motivate this, we recall that, given a square integrable $f(x)$ and a reversible MCMC proposal, it is possible to bound $ESS$ from above (see \Cref{sec:ess_max}) by:
\begin{align}
\label{eq:ess_max}
ESS \le \frac{1 - \rho_f(1)}{1 + \rho_f(1)} \triangleq ESS_{max}.
\end{align}
This bound is the same as one gets if one computes ESS by assuming a specific functional form for $\rho_f(l) = \rho_f(0)^l$ and substituting it into \Cref{eq:ess}. If we assume positive $\rho_f(1)$ it is clear that $ESJD$ is in fact an upper bound on $ESS$ and therefore ChEESR is an upper bound on $ESS_{\lnorm z \rnorm_2^2} / \tau$.
When doing Bayesian inference it is common to desire to maximize $\min_i ESS_{z_i^2}$, as individual elements of the problem state typically have semantic meaning assigned to them by the modeler. Starting from ChEES there are three steps of approximation between it and $\min_i ESS_{z_i^2}$: $ChEES \ge_1 ESS_{max} \ge_2 ESS_{\lnorm z \rnorm_2^2} / \tau \approx_3 \min_i ESS_{z_i^2}$. We cannot control the first inequality except to hope that $\rho_{\lnorm z \rnorm_2^2}(1)$ is small. The second inequality can be measured empirically with minimal memory overhead because $\lnorm z \rnorm_2^2$ is a scalar quantity. The tightness of this bound can be recruited as a post-hoc certificate that $\tau$ optimization likely found a good optimum, although a practitioner would not have much recourse to fix the optimization if the bound is loose outside of adjusting the parameterization of the model. The third approximation occasionally has pathologies as shown by the example in \Cref{fig:motivation}, at least when using $f(x) = \lnorm z \rnorm_2^2$. We propose adjusting $f(x)$ to resolve that failure case.
Identifying $\argmin_i ESS_{z_i^2}$ is difficult, but a reasonable heuristic is to focus on optimizing ESS along the direction of greatest variance. This has been part of the motivation behind ESJD and ChEES where they weigh dimensions by $\sigma_d^2$ and $\sigma_d^4$ respectively. It is in fact straightforward to determine this direction directly by identifying it with the first principal component of the covariance matrix and then learning it via Principal Component Analysis. A convenient algorithm in this field is Oja's method \citep{oja1982simplified} which is memoryless and is easy to integrate into the rest of the adaptive MCMC. In practice, we use a minibatch version of the algorithm as we have access to the state of all the chains at once (\Cref{alg:oja}). In summary, we propose to maximize the Squared Norm Along Principal component ESJD Rate (SNAPER):
\begin{align}
\label{eq:snaper}
SNAPER &\triangleq \frac{ESJD_{(z^Tp)^2}}{\tau} \\
f_{SNAPER}(x) &= (z^Tp)^2
\end{align}
where $p$ is the learned first principal component.
\begin{algorithm}
\caption{One iteration of Oja's Algorithm with minibatch support}\label{alg:oja}
\begin{algorithmic}[1]
\Require $x^{(t, b)}$, a minibatch of D-dimensional states indexed by batch index $b \in [0, B]$. $p^{(t)}$ current value of the first principal component. $\eta^{(t)}_p$, current learning rate.
\Ensure $p^{(t + 1)}$, updated value of the principal component.
\State $\tilde{p} \leftarrow {0}_D$
\For{$b=1$ to $B$}
\State $\tilde{p} \leftarrow \tilde{p} + x^{(t, b)} x^{(t, b) T} p^{(t)}$
\EndFor
\State $p' \leftarrow \frac{\tilde{p}}{\lnorm \tilde{p} \rnorm}$
\State $p^{(t + 1)} \leftarrow \frac{p^{(t)} + \eta^{(t)}_p p'}{\lnorm p^{(t)} + \eta^{(t)}_p p' \rnorm}$
\end{algorithmic}
\end{algorithm}
\section{Adaptive SGA-HMC}
We use automatic differentiation to compute the pathwise gradients of a criterion. To avoid differentiating the HMC proposal, we manually derive and annotate this gradient. For the criteria we wish to test, we only require the gradients flowing through the proposed\footnote{Technically \Cref{eq:esjd} uses the accepted MCMC state, but computing the $\tau$ gradients flowing through the Metropolis-Hastings step is trivial.} state $x'^{(t + 1)}$. Specifically, the vector-Jacobian product of the gradients flowing through $x'^{(t + 1)}$ are:
\begin{align}
\nabla^T J(x'^{(t + 1)}(\tau)) = \epsilon \sum_{d=1}^D \nabla_d \frac{\partial H(x, m)}{\partial m_d}.
\end{align}
For the purpose of adapting $\tau$ we ignore the fact that for a given $L$, $\tau$ can only take discrete values due to the discretized dynamics used to implement HMC. Now, given $\tau \sim \unif(0, e^{\phi_\tau})$ and a differentiable criterion are thus able compute $\bbE[\frac{\calC(\tau)}{\phi_\tau}]$. We use ADAM \citep{kingma2015adam} to adapt $\phi_\tau$. When learning $\epsilon$ we first parameterize it as $\epsilon = e^{\phi_\epsilon}$ and then use a synthetic gradient to push the harmonic mean of per-chain acceptance probability $\alpha^{(b)}$ towards a target acceptance rate $\alpha^* = 0.8$:
\begin{align}
\widehat{\frac{\partial\mathcal{A}}{\partial\alpha}} = \alpha^* - \left(\frac{1}{B} \sum_b \alpha_b^{-1}\right)^{-1}.
\end{align}
A harmonic mean is used to make sure that all chains that are evaluated in parallel and share HMC hyperparameters can make progress \citep{hoffman2021adaptive}. $\phi_\epsilon$ is also adapted using ADAM. When HMC is parameterized by $\tau$ rather than $L$ it is important to initialize $\epsilon$ to a reasonable value so as not to waste a lot of compute in the early iterations. Standard protocol is to do a search over $\epsilon$ until a reasonable acceptance probability is obtained \citep{hoffman2011no}. To make more efficient use of gradient evaluations, we instead force HMC to only take a single leapfrog step for the first 100 iterations. This procedure recruits ADAM to perform the search for us, while also carefully navigating the potentially pathological geometry if the algorithm is initialized in a suboptimal location.
The mean $\mu$ and variance $\sigma^2$ of the state are learned using a variant of Welford's algorithm \citep{welford1962note} which uses an exponentially weighted average with a decaying learning rate to smoothly discount early, non-representative iterates. To aid trajectory adaptation, instead of using raw $\sigma^2$ as the value for the diagonal of $M$, we normalize it by the largest component. This has the effect of minimizing the effect of learning $M$ on the optimal value for $\tau$ (which ideally depends primarily on the longest length scale of the distribution). ChEES, ChEESR and SNAPER require (see \Cref{eq:snaper,eq:f_chees}) $x$ to be centered. This center is computed by applying the same algorithm as is used to learn $\mu$ except the update is derived either from the current state $x^{(t)}$ or acceptance-probability weighted proposed state $\frac{1}{\sum_b \alpha_b} \sum_b \alpha_b x'^{(t + 1, b)}$ as appropriate. The pseudocode for the combined algorithm, Adaptive Stochastic Gradient Ascent HMC (SGA-HMC) is listed in \Cref{alg:sga_hmc}. While adaptation is proceeding, we maintain iterate averages of all the hyperparameters for use in post-warmup sampling. \Cref{fig:trace} illustrates the action of SGA-HMC with the SNAPER criterion on one of the example problems.
As a baseline, we also define an Adaptive NUTS algorithm, which follows \Cref{alg:sga_hmc} except that SGA-HMC is replaced by NUTS. Critically, for the first 100 steps we still only take one leapfrog step as with Adaptive SGA-HMC, to make sure NUTS starts out with a reasonable $\epsilon$.
\begin{algorithm}
\caption{Adaptive Stochastic Gradient Ascent HMC}\label{alg:sga_hmc}
\begin{algorithmic}[1]
\Require $x^{(t, b)}$, an a minibatch of D-dimensional states indexed by batch index $b \in [1, B]$. $\phi_\tau^{(t)}, \phi^{(t)}_\epsilon$, parameters describing $\tau$ and $\epsilon$ respectively. $\mu^{(t)}, \sigma^{2(t)}$, running first and second moments. $t$ iteration index.
\Ensure Next state $x^{(t+1)}$. Updated versions of parameters and statistics.
\LineComment{Compute current HMC hyperparameters.}
\State $\epsilon^{(t)} \leftarrow e^{\phi^{(t)}_\epsilon}$
\State $\tau^{(t)} \sim \unif(0, e^{\phi_\tau})$
\State $M^{(t)}_{dd} \leftarrow \sigma^{2(t)}_d / \max_d \sigma^{2(t)}_d$
\LineComment{Generate the next state using HMC.}
\State $x^{(t + 1)}, x'^{(t)}, \alpha^{(t)} \leftarrow \textsc{HMC}(x^{(t)}, \epsilon^{(t)}, \tau^{(t)}, M^{(t)}$)
\LineComment{Clip $\tau^{(t)}$ during the first 100 iterations.}
\If{$t < 100$}
\State $\tau^{(t)} \leftarrow \epsilon^{(t)}$
\EndIf
\LineComment{Compute the criterion value.}
\State $\calC^{(t)} \leftarrow \calC(x^{(t)}, x'^{(t)}, \alpha^{(t)}, \mu^{(t)})$
\LineComment{Update trajectory distribution parameters}.
\State $\phi_\tau^{(t + 1)} \leftarrow \textsc{ADAM}(\phi_\tau^{(t)}, -\frac{\partial{\calC^{(t)}}}{\partial \phi_\tau^{(t)}}, \eta_\tau, \beta_1, \beta_2)$
\LineComment{Update step size parameters.}
\State $\phi_\epsilon^{(t + 1)} \leftarrow \textsc{ADAM}(\phi_\epsilon^{(t)}, \alpha^* - (\frac{1}{B} \sum_b \alpha_b^{-1})^{-1}, \eta_\epsilon)$
\LineComment{Update mean and variance estimates.}
\State $\eta_\mu \leftarrow \frac{1}{\lceil t / \kappa \rceil + 1}$
\State $\mu^{(t + 1)} \leftarrow (1 - \eta_\mu)\mu^{(t)} + \frac{\eta_\mu}{B} \sum_b x^{(t + 1, b)}$
\State $\sigma^{2(t + 1)} \leftarrow (1 - \eta_\mu)\sigma^{2{(t)}} + \frac{\eta_\mu}{B} \sum_b (x^{(t + 1, b)} - \mu^{(t)})^2$
\end{algorithmic}
\end{algorithm}
\begin{figure*}[!tb]
\centering
\includegraphics[scale=0.55]{figures/trace.pdf}
\caption{
Evolution of the learned hyperparameters from a SNAPER-HMC run on the German Credit Sparse Regression problem as a function of MCMC iteration number. During the initial phase (first vertical line), the number of leapfrog steps is set to 1. Then, the adaptation proceeds for a chosen number of steps (second vertical line), after which the hyperparameters are fixed. Sampling can be stopped when $\hat{R} < 1.01$ (third vertical line). The last panel is the sin of the angle between the learned principal component direction and its projection onto the top 15 eigenvectors of the unconstrained state covariance matrix whose scales are greater than 0.85 of the largest scale.
}
\label{fig:trace}
\end{figure*}
\section{Experiments}
We validate SNAPER and ChEESR performance on six problems, three applied to real world data and three applied to synthetic data sampled from the respective model's prior. All the models are taken from the Inference Gym \citep{sountsov2020inference}, which we refer the reader to for model definitions. We summarize the models used in \Cref{tab:models}. For each model, we run SGA-HMC with one of the criteria for a prescribed number of adaptation steps, we then turn off adaptation, switch HMC to use the iterate-averaged hyperparameters and collect statistics like ESS/grad and potential scale reduction \citep{gelman1992,vehtari2019rank}. Some of the models have constraints on their parameters, which we handle by reparameterizing them and adjusting the log-density by the logarithm of the determinant of the Jacobian. SGA-HMC is ran and all of the statistics are computed in the unconstrained space. For each run, we initialize the chains at the prior mean. This severely underdispersed initialization aids algorithms which rely on cross-chain statistics, although it leaves us vulnerable to not detecting all the modes of the posterior. We consider SNAPER, ChEESR and ChEES criteria. For ChEES we considered two settings of the trajectory length learning rate $\eta_\tau = [0.025, 0.05]$, the latter being identified in the plots and tables as ``ChEES fast''. For all other conditions $\eta_\tau$ was set to $0.05$. See \Cref{sec:hyperparameters} for more hyperparameter settings.
\begin{table*}[!tb]
\centering
\begin{tabular}{ lll }
\toprule
Model & $D$ & Inference Gym Snippet \\
\midrule
Logistic Regression & 25 & \lstinline[]$GermanCreditNumericLogisticRegression()$ \\
Sparse Logistic Regression & 51 & \makecell[l]{\lstinline[]$GermanCreditNumericSparseLogisticRegression($ \\ \hspace{0.25in} \lstinline[]$positive_constraint_fn="softplus")$} \\
Brownian Bridge & 32 & \makecell[l]{\lstinline[]$BrownianMotionUnknownScalesMissingMiddleObservations($ \\ \hspace{0.25in} \lstinline[]$use_markov_chain=True)$} \\
Item Response Theory & 501 & \lstinline[]$SyntheticItemResponseTheory()$ \\
Hierarchical Linear Model & 97 & \lstinline[]$RadonContextualEffectsIndiana()$ \\
Stochastic Volatility & 2519 & \lstinline[]$VectorizedStochasticVolatilitySP500()$ \\
\bottomrule
\end{tabular}
\caption{
Bayesian posteriors used for the experiments. The snippets refer to the Python class invocation to construct the model (the classes are all found in the \lstinline{inference_gym.targets} module).
}
\label{tab:models}
\end{table*}
\subsection{Long Run Efficiency}
For the first experiment, we examined the long run sampling efficiency of all the algorithms. This is the setting where we can discount the cost of the warmup iterations and are solely interested in collecting a large number of samples from long chains. First, we performed sweeps of SGA-HMC with fixed $L$ and $M$ and measured $ESS/grad$, $ESS_{max}/grad$ and $ESJD/grad$ for three test functions, $f_{SNAPER}$, $f_{ChEES}$ and $z^2_{argmin ESS}$ (lines in \Cref{fig:asymptotics}). For Item Response Theory and Sparse Logistic Regression $ESS_{max}$ is close to actual $ESS$ when using $f_{SNAPER}$ (left column). Let $\hat{\tau}_{\calC} \triangleq \argmax_\tau \calC(\tau)$. For these models, $\hat{\tau}_{SNAPER}$ is close to $\hat{\tau}_{ ESS/grad}$. The bound is not as tight for Sparse Logistic Regression when using $f_{ChEES}$, middle column, or Brownian Bridge, which causes the performance to be less predictable. We reiterate that the bound of $ESS_{max}/grad \ge ESS/grad$ can be computed empirically with minimal memory overhead, because both $f_{ChEES}$ and $f_{SNAPER}$ are scalar functions. Another phenomenon is how $f_{SNAPER}$ is a superior test function to $f_{ChEES}$ when the posterior covariance spectrum resembles the pathological one from \Cref{fig:motivation}, as it does for Item Response Theory. We see that $\hat{\tau}_{SNAPER}$ matches $\hat{\tau}_{\min_i ESS_{z_i^2} / grad}$ closely, while $\hat{\tau}_{ChEESR}$ is too short.
We then ran 100 separate trials of Adaptive SGA-HMC targeting each criterion, each trial being distinguished by the PRNG seed used for the random sampling events. We warm up for 5000 steps, followed by 1000 steps of sample collection. We see that both ChEESR and SNAPER recover the optima of their respective criteria reliably. ChEES and ChEES-fast have a lot more variability which is shown by both the large spread of learned trajectories for Brownian Bridge (note the log-scale on the x-axis) and the occasional overshoot for Item Response Theory. In general, ChEES-derived metrics either find trajectories that are too long or too short, depending on the problem, causing suboptimal $\min_d ESS_{z_d^2}$. SNAPER results in consistently good performance across all problems (especially when $ESS_{max}$ bound is tight). Additional results can be seen in \Cref{tab:asymptotics_tabl} and \Cref{sec:extra_plots}.
\begin{figure*}[!tb]
\centering
\includegraphics[scale=0.55]{figures/asymptotic.pdf}
\caption{
Long run behavior of different criteria. The rows are three
representative target distributions. The columns are three test
functions: squared projected norm, squared norm and squared
component with the lowest ESS/grad. The lines are computed via
sweeps over the number of mean leapfrogs per MCMC iteration (higher is better). The
points are jittered final mean numbers of leapfrog steps across 100
runs, with jitter amount proportional to local point density.
SNAPER tends to find a good trajectory length as measured by all
ESS/grad metrics.
}
\label{fig:asymptotics}
\end{figure*}
\begin{table*}[!tb]
\centering
\begin{tabular}{ lllllll }
\toprule
& Grid & ChEES & ChEES fast & SNAPER & ChEESR & NUTS \\
\midrule
Brownian Bridge & 2.09e-3 & 1.23e-4 & 1.20e-4 & \textbf{2.56e-3} & 2.28e-3 & 3.84e-4 \\
Hierarchical Linear Model & 4.28e-3 & \textbf{4.35e-3} & 4.01e-3 & 3.79e-3 & 4.03e-3 & 1.28e-3 \\
Item Response Theory & 1.01e-2 & 3.76e-3 & 3.73e-3 & \textbf{9.77e-3} & 2.80e-3 & 4.37e-3 \\
Logistic Regression & 1.20e-1 & 8.35e-2 & 8.45e-2 & 1.05e-1 & \textbf{1.28e-1} & 8.94e-3 \\
Sparse Logistic Regression & 5.23e-3 & 1.08e-3 & 1.14e-3 & 5.15e-3 & \textbf{5.33e-3} & 1.44e-3 \\
Stochastic Volatility & 1.36e-2 & 1.46e-2 & 1.46e-2 & \textbf{1.49e-2} & 1.32e-2 & 5.14e-3 \\
\bottomrule
\end{tabular}
\caption{
Asymptotic $\min_i ESS_{z_i^2} /grad$ per experiment configuration (including grid search), higher is better. The numbers
are 10'th percentile across 100 runs. Bold counts are the best
configurations for each problem (not counting the grid search).
}
\label{tab:asymptotics_tabl}
\end{table*}
\subsection{Short Run Efficiency}
Many questions in Bayesian inference only require less than 1000 independent samples, sometimes as low as 100. If we run 100 chains in parallel, then in principle we could stop and record just the final iteration as soon as we had a certificate that the sampler has converged. $\hat{R}$ is a popular convergence diagnostic that approaches 1 from above as the chain converges to its stationary distribution; $\hat{R} < 1.01$ is considered an acceptable threshold for convergence \citep[e.g.]{vats2018revisting}. Thus, we envision the following protocol: 1) Run Adaptive SGA-HMC for a prescribed number of warmup/adaptation steps. 2) Turn off adaptation and collect samples until $\hat{R} < 1.01$ for all test functions of choice (in this experiment, we use $f(x) = x_d$ for all $d \in [1, D]$). A good algorithm will reach this stopping condition with few gradient evaluations. There is a non-zero optimal number of warmup steps for this procedure: too few, and the sampling stage will proceed too slowly due to a poorly adapted sampler; too many, and the adaptation time is wasted. In \Cref{fig:early_stop} we run Adaptive SGA-HMC for all four settings, and stop adaptation at 500, 750 and 1000 steps. We run 100 separate trials for each experiment configuration. We report both the total number of gradient evaluations (left column) and the number of gradient evaluations just in the sampling phase (right column). We observe that ChEES-fast sometimes will overshoot the optimum (Hierarchical Linear Model, 1000 warmup steps). This deficiency can be ameliorated with a slower learning rate, but at a cost reduced performance in some settings (Hierarchical Linear Model, 500 warmup steps). In general, for ChEES and ChEES-fast, stopping early can recover some of the competitiveness with other critera because the $\phi_\tau$ has not yet converged to its overlong critical point. ChEESR and SNAPER behave similarly to each other except in the case of Item Response Theory where the pathological covariance spectrum causes ChEESR to learn an inefficient sampler. Additional results are available in \Cref{tab:early_stop_table} and the \mbox{\Cref{sec:extra_plots}}. When ChEESR can learn a good sample, it tends to behave well in these experiments. SNAPER sometimes struggles with very short warmup windows, perhaps due to the additional noise caused by needing to learn the value of the first principal component.
\begin{figure*}[!tb]
\centering
\includegraphics[scale=0.55]{figures/early_stop.pdf}
\caption{
Gradient evaluation counts as a function of number of warmup steps.
The three rows are three representative target distributions. The left
column counts the amount of grads for both warmup and sampling
periods, while the right counts sampling grads only. Sampling stops
when the post-warmup chains reach an $\hat{R} < 1.01$. Each setting
of the warmup steps represents a separate run from the other
settings.
}
\label{fig:early_stop}
\end{figure*}
\begin{table*}[!tb]
\centering
\begin{tabular}{ llllll }
\toprule
& ChEES & ChEES fast & SNAPER & ChEESR & NUTS \\
\midrule
Brownian Bridge & 2.7e+5/2.9e+5 & 4.0e+5/7.1e+5 & 8.7e+4/8.3e+4 & \textbf{7.8e+4}/7.9e+4 & 2.8e+5/3.8e+5 \\
Hier. Linear Model & 1.2e+5/7.1e+4 & 5.8e+4/7.3e+4 & \textbf{4.7e+4}/5.8e+4 & 5.9e+4/5.7e+4 & 2.1e+5/2.1e+5 \\
Item Response & 7.4e+4/7.7e+4 & 7.6e+4/7.5e+4 & \textbf{3.3e+4}/4.0e+4 & 9.7e+4/1.0e+5 & 6.8e+4/7.6e+4 \\
Logistic Reg. & 2.9e+3/3.9e+3 & 2.9e+3/3.9e+3 & \textbf{2.8e+3}/3.7e+3 & 2.9e+3/3.7e+3 & 3.6e+4/4.9e+4 \\
Sparse Logistic Reg. & 1.3e+5/3.5e+5 & 3.0e+5/4.1e+5 & \textbf{4.0e+4}/4.9e+4 & 4.3e+4/5.6e+4 & 2.6e+5/3.0e+5 \\
Stochastic Volatility & 2.1e+4/2.2e+4 & 2.0e+4/2.2e+4 & 1.8e+4/1.8e+4 & 1.7e+4/\textbf{1.7e+4} & 5.5e+4/6.8e+4 \\
\bottomrule
\end{tabular}
\caption{
Total count of gradient evaluations, lower is better. The two
counts are the 90'th percentile across 100 runs for 500 and 750 warmup
steps respectively. Bold counts are the best configurations for each
problem.
}
\label{tab:early_stop_table}
\end{table*}
\section{Discussion}
In this work we have introduced a new criterion for gradient-based HMC trajectory length adaptation -- SNAPER. We have illustrated that SNAPER has excellent long and short chain efficiency. In most cases, with sufficient warmup, SNAPER delivers consistently optimal trajectory lengths. SNAPER derives these quantities by having an accurate cost model of HMC sampling efficiency and by using a good heuristic for identifying which directions to focus on when optimizing $\tau$. An epiphenomenon of this robustness is that the criterion tends to have a single, well defined peak (compare \Cref{fig:motivation}, left column, and \Cref{fig:asymptotics}, left column) allowing the use of high learning rates.
When does SNAPER-HMC fail? SNAPER-HMC relies on the principal component identifying the most ``difficult'' direction, but it's easy to come up with a counterexample where a low variance direction has enough nonlinear twists to require a larger $\tau$ to explore well than simpler, high variance dimensions might suggest. If these difficult dimensions can be identified manually, then a practitioner could hand-craft a $p$ (\Cref{eq:snaper}) that focuses on those dimensions as a remedy to this problem. A final limitation is that since SNAPER is relies on the quality of the $ESS_{max}$ bound, when that bound is loose SNAPER cannot be guaranteed to recover a good $\tau$.
SGA-HMC is motivated by the use of SIMD-hardware to run multiple chains in parallel. Sometimes this is not acceptable due to memory requirements. SGA-HMC can still function when few chains are available, due to cross-iteration averaging courtesy of ADAM and Welford's algorithm, but its convergence speed suffers due to increasing estimator variance. This issue is unlikely to affect long run performance, making the procedure still beneficial compared to NUTS.
SNAPER-HMC is particularly well suited to the Adaptive MCMC framework where $M$ is also learned; empirically, we observed that learning a good $M$ is difficult enough that SNAPER to introduce minimal overhead on top that. If $M$ is fixed, learning a $\tau$ via SNAPER-HMC may be less attractive. This is likely a rare use case, which means that SNAPER-HMC is likely a good baseline MCMC method to use with SIMD hardware.
\acknowledgments{
We thank Rif A. Saurous for helpful comments on the manuscript and Dave Moore, Colin Carroll and Sharad Vikram for proofreading the implementation.
}
\clearpage
\bibliography{main}
\newpage
\begin{appendices}
\onecolumn
\crefalias{section}{appsection}
\section{Additional Plots}
\label{sec:extra_plots}
We include two additional plots for the short and long run performance of the considered algorithms, \Cref{fig:asymptotics_all} and \Cref{fig:early_stop_all}.
\begin{minipage}{\textwidth}
\centering
\includegraphics[scale=0.55]{figures/asymptotic_all.pdf}
\captionof{figure}{Long run behavior of different criteria. See \Cref{fig:asymptotics} for details.}
\label{fig:asymptotics_all}
\end{minipage}
\begin{figure*}[!hb]
\centering
\includegraphics[scale=0.55]{figures/early_stop_all.pdf}
\caption{
Gradient evaluation counts as a function of number of warmup steps. See \Cref{fig:early_stop} for details.
}
\label{fig:early_stop_all}
\end{figure*}
\clearpage
\section{$ESS_{max}$ derivation}
\label{sec:ess_max}
We derive the bound $ESS_{f,max} \ge ESS_f$ following \cite{geyer1992practical} and \cite{sokal1997monte}. Given $f(x)$ that satisfies $\int f(x) \pi(dx) < \infty$ and $x^{(t)}$, $t \ge 0$ are forms an ergodic and reversible Markov chain. We can use the spectral theorem to write:
\begin{align}
\rho_f(l) = \int_{-1}^{1} dE_\lambda \lambda^{|l|},
\end{align}
where $dE_\lambda$ is a probability measure supported on $[-1, 1]$. We can then write:
\begin{align}
\sum_{l=-\infty}^{\infty} \rho_f(l) &= \int_{-1}^{1} dE_\lambda \, \frac{1 + \lambda}{1 - \lambda} \\
&\ge \frac{1 + \int_{-1}^{1} dE_\lambda \, \lambda }{1 - \int_{-1}^{1} dE_\lambda \, \lambda } \\
\label{eq:bound}
&\ge \frac{1 + \rho_f(1) }{1 - \rho_f(1) }.
\end{align}
Since $\rho(l) = \rho(-l)$ we have:
\begin{align}
ESS_f &\triangleq \frac{N}{1 + 2 \sum_{l=1}^{\infty} \rho_f(l)} \\
&= \frac{N}{\sum_{l=-\infty}^{\infty} \rho_f(l)} \\
&\le N \frac{1 - \rho_f(1)}{1 + \rho_f(1)} \triangleq ESS_{f,max}.
\end{align}
\section{Hyperparameter settings}
\label{sec:hyperparameters}
The hyperparameters were minimally tuned, and shared across all problems.
\begin{minipage}{\textwidth}
\centering
\begin{tabular}{ lll }
\toprule
Hyperparameter & Meaning & Value \\
\midrule
$B$ & Number of chains ran in parallel & 64 \\
$\kappa$ & Factor controlling moment adaptation & 8 \\
$\eta_\tau$ & Learning rate for $\phi_\tau$ & 0.05, 0.025 (ChEES) \\
$\eta_p^{(t)}$ & Learning rate for $p$ & $\frac{8}{t}$ \\
$\eta_\epsilon$ & Learning rate for $\phi_\epsilon$ & 0.05 \\
$\beta_1$ & ADAM hyperparameter for $\phi_\tau$ & 0 \\
$\beta_2$ & ADAM hyperparameter for $\phi_\tau$ & 0.5 \\
\bottomrule
\end{tabular}
\end{minipage}
\end{appendices}
\end{document}
\begin{thebibliography}{}
\bibitem[Andrieu and Thoms, 2008]{andrieu2008tutorial}
Andrieu, C. and Thoms, J. (2008).
\newblock A tutorial on adaptive mcmc.
\newblock {\em Statistics and computing}, 18(4):343--373.
\bibitem[Betancourt et~al., 2014]{betancourt2014optimizing}
Betancourt, M.~J., Byrne, S., and Girolami, M. (2014).
\newblock Optimizing the integrator step size for hamiltonian monte carlo.
\bibitem[Carpenter et~al., 2017]{carpenter2017stan}
Carpenter, B., Gelman, A., Hoffman, M.~D., Lee, D., Goodrich, B., Betancourt,
M., Brubaker, M., Guo, J., Li, P., and Riddell, A. (2017).
\newblock Stan: A probabilistic programming language.
\newblock {\em Journal of statistical software}, 76(1).
\bibitem[Duane et~al., 1987]{duane1987hybrid}
Duane, S., Kennedy, A.~D., Pendleton, B.~J., and Roweth, D. (1987).
\newblock Hybrid monte carlo.
\newblock {\em Physics letters B}, 195(2):216--222.
\bibitem[Gelman and Rubin, 1992]{gelman1992}
Gelman, A. and Rubin, D.~B. (1992).
\newblock Inference from iterative simulation using multiple sequences.
\newblock {\em Statistical Science}.
\bibitem[Geyer, 1992]{geyer1992practical}
Geyer, C.~J. (1992).
\newblock Practical markov chain monte carlo.
\newblock {\em Statistical Science}, 7(4):473--511.
\bibitem[Hastings, 1970]{hastings1970mh}
Hastings, W.~K. (1970).
\newblock Monte carlo sampling methods using markov chains and their
applications.
\newblock {\em Biometrika}, 57(1):97--109.
\bibitem[Hoffman et~al., 2021]{hoffman2021adaptive}
Hoffman, M., Radul, A., and Sountsov, P. (2021).
\newblock An {Adaptive-MCMC} scheme for setting trajectory lengths in
hamiltonian monte carlo.
\newblock 130:3907--3915.
\bibitem[Hoffman and Gelman, 2011]{hoffman2011no}
Hoffman, M.~D. and Gelman, A. (2011).
\newblock The no-u-turn sampler: Adaptively setting path lengths in hamiltonian
monte carlo.
\newblock {\em The Journal of Machine Learning Research}.
\bibitem[Kingma and Ba, 2015]{kingma2015adam}
Kingma, D. and Ba, J. (2015).
\newblock {Adam}: A method for stochastic optimization.
\newblock In {\em International Conference on Learning Representations}.
\bibitem[Lao and Dillon, 2019]{lao2019unrolled}
Lao, J. and Dillon, J.~V. (2019).
\newblock Unrolled implementation of no-u-turn sampler.
\newblock
\url{https://github.com/tensorflow/probability/blob/main/discussion/technical_note_on_unrolled_nuts.md}.
\bibitem[Metropolis et~al., 1953]{metropolis1953equation}
Metropolis, N., Rosenbluth, A.~W., Rosenbluth, M.~N., Teller, A.~H., and
Teller, E. (1953).
\newblock Equation of state calculations by fast computing machines.
\newblock {\em The journal of chemical physics}, 21(6):1087--1092.
\bibitem[Neal, 2011]{neal2011mcmc}
Neal, R.~M. (2011).
\newblock {MCMC} using {Hamiltonian} dynamics.
\newblock In {\em Handbook of Markov Chain Monte Carlo}. CRC Press New York,
NY.
\bibitem[Oja, 1982]{oja1982simplified}
Oja, E. (1982).
\newblock A simplified neuron model as a principal component analyzer.
\newblock {\em Journal of Mathematical Biology}, 15:267--273.
\bibitem[Pasarica and Gelman, 2010]{pasarica2010adaptively}
Pasarica, C. and Gelman, A. (2010).
\newblock Adaptively scaling the metropolis algorithm using expected squared
jumped distance.
\newblock {\em Statistica Sinica}, pages 343--364.
\bibitem[Phan and Pradhan, 2019]{phan2019iterative}
Phan, D. and Pradhan, N. (2019).
\newblock Iterative {NUTS}.
\newblock \url{https://github.com/pyro-ppl/numpyro/wiki/Iterative-NUTS}.
\bibitem[Radul et~al., 2019]{radul2019automatically}
Radul, A., Patton, B., Maclaurin, D., Hoffman, M.~D., and Saurous, R.~A.
(2019).
\newblock Automatically batching (control-intensive) programs for modern
accelerators.
\bibitem[Sokal, 1997]{sokal1997monte}
Sokal, A.~D. (1997).
\newblock {\em Monte Carlo Methods in Statistical Mechanics: Foundations and
New Algorithms}.
\bibitem[Sountsov et~al., 2020]{sountsov2020inference}
Sountsov, P., Radul, A., and contributors (2020).
\newblock Inference gym.
\newblock \url{https://pypi.org/project/inference_gym}.
\bibitem[Vats and Knudson, 2018]{vats2018revisting}
Vats, D. and Knudson, C. (2018).
\newblock Revisiting the {Gelman-Rubin} diagnostic.
\bibitem[Vehtari et~al., 2019]{vehtari2019rank}
Vehtari, A., Gelman, A., Simpson, D., Carpenter, B., and B{\"{u}}rkner, P.-C.
(2019).
\newblock {Rank-normalization, folding, and localization: An improved R for
assessing convergence of MCMC}.
\bibitem[Wang et~al., 2013]{wang2013adaptive}
Wang, Z., Mohamed, S., and de~Freitas, N. (2013).
\newblock Adaptive hamiltonian and riemann manifold monte carlo samplers.
\bibitem[Welford, 1962]{welford1962note}
Welford, B.~P. (1962).
\newblock Note on a method for calculating corrected sums of squares and
products.
\newblock {\em Technometrics}, 4(3):419--420.
\bibitem[Wu et~al., 2018]{wu2018faster}
Wu, C., Stoehr, J., and Robert, C.~P. (2018).
\newblock Faster hamiltonian monte carlo by learning leapfrog scale.
\end{thebibliography}