Almanac: HMC sampling with bounded velocity

Technical Summary: ALMANAC: HAMILTONIAN MONTE CARLO SAMPLING WITH BOUNDED VELOCITY

Problem Statement
Hamiltonian Monte Carlo (HMC) is a powerful algorithm for sampling high-dimensional probability distributions, relying on the joint dynamics of potential energy (derived from the target distribution) and kinetic energy (derived from the momentum distribution). Standard implementations typically assume a Gaussian momentum distribution, which implies a quadratic kinetic energy and allows for arbitrarily large velocities. While efficient for Gaussian-like posteriors, this assumption can lead to inefficient exploration in posteriors exhibiting "funnel" geometries or sharp ridges, common in cosmological inference problems such as those involving sky maps and power spectra. In such regions, unbounded velocities can cause the sampler to overshoot steep gradients, resulting in low acceptance rates or the need for computationally expensive small step sizes.

Methodology
The authors investigate alternative momentum distributions based on kinetic energies that naturally bound particle velocities. Specifically, they test two non-Gaussian definitions:

1. Relativistic Kinetic Energy: Modeled after special relativity, where the kinetic energy is hyperbolic in momentum magnitude. This imposes a maximum velocity (analogous to the speed of light, $c$) and allows for tuning via a "rest mass" parameter ($m$).
2. Student's t Kinetic Energy: Defined such that the momentum follows a Student's t-distribution. This results in a sublinear kinetic energy for large momenta, also imposing a finite maximum velocity dependent on the degrees of freedom ($\nu$).

These alternatives are implemented within Almanac, an HMC sampler designed for cosmological posterior distributions of sky maps and power spectra on the sphere. The study utilizes two parameterizations of the posterior:

• Classic: Uses the matrix logarithm of the power spectrum covariance and raw spherical harmonic coefficients.
• Cholesky: Uses the Cholesky decomposition of the covariance matrix and normalized coefficients, designed to mitigate funnel structures.

The authors conduct numerical experiments across three dimensional regimes (low, mid, and high) using Almanac. They evaluate performance using:

• Effective Sample Rate (ESR): The number of near-independent samples per unit of computational time.
• Correlation Length ($\hat{\tau}$): A measure of how many steps are required to obtain an independent sample.
• Fraction of Missing Information (FMI): A diagnostic for convergence and the sampler's ability to explore the full energy submanifolds.

Key Contributions and Results
The paper presents a systematic comparison of Gaussian, Relativistic, and Student's t momentum distributions across varying dimensionalities and parameterizations.

• Parameterization Dominance: The Cholesky parameterization consistently outperformed the Classic parameterization in mid- and high-dimensional settings. It yielded significantly higher ESR (by factors of 5–10) and better FMI values (indicating better convergence). In contrast, the Classic parameterization often produced misleadingly short correlation lengths while failing to converge (FMI < 0.7). However, in a specific low-dimensional, high signal-to-noise regime, the Classic parameterization outperformed Cholesky, aligning with established literature on centered vs. non-centered parameterizations.
• Momentum Distribution Trade-offs: The choice of kinetic energy revealed a trade-off between mixing (robustness) and raw efficiency.
  • Heavy-tailed distributions (e.g., Student's t with low $\nu$, or Relativistic with low mass $m$) improved mixing, evidenced by higher FMI and lower correlation lengths. However, this came at the cost of a lower Effective Sample Rate (slower sampling speed).
  • Moderately heavy-tailed distributions (e.g., Relativistic with tuned mass parameters like $m \approx 1.193$) achieved the best balance, offering competitive ESR while maintaining reasonable mixing properties.
  • Gaussian distributions generally provided the lowest FMI in the mid-dim experiments, indicating poorer convergence in the presence of funnels compared to the bounded-velocity alternatives.
• Scalability: In the high-dimensional regime (millions of parameters), the ESR dropped dramatically (by orders of magnitude) compared to mid-dimensional tests, a direct consequence of increased computational cost per iteration. However, the relative performance trends between parameterizations and momentum distributions remained consistent with the mid-dimensional results.
• Hyperparameter Sensitivity: The performance of the non-Gaussian samplers was found to be highly sensitive to hyperparameters (e.g., rest mass $m$ and speed of light $c$). Small changes in these parameters could shift the sampler from efficient to inefficient, suggesting that fixed defaults are suboptimal.

Significance and Claims
The authors conclude that while modifying the momentum distribution and parameterization can improve convergence and effective sample rates, the gains are generally modest and strongly problem-dependent. The study highlights that:

1. No Universal Solution: There is no single "best" kinetic energy or parameterization; the optimal choice depends on the specific geometry of the posterior (e.g., the presence of funnels) and the dimensionality of the problem.
2. Importance of Tuning: The significant sensitivity to hyperparameters (mass and speed of light) underscores the need for careful tuning. The authors suggest that future work should focus on adaptive and self-tuning strategies for kinetic energy optimization, similar to recent developments in the field, to automate these choices in high-dimensional settings.
3. Robustness of Design: The results reinforce that sampler design is critical. While bounded-velocity distributions offer a mechanism to handle steep gradients and funnel geometries more robustly than standard Gaussian momenta, their utility is maximized only when paired with appropriate parameterizations (like Cholesky) and tuned hyperparameters.

The paper does not claim that these methods solve the scaling challenges of high-dimensional sampling entirely but rather provides evidence that tailored kinetic energy definitions can offer incremental but meaningful improvements in specific cosmological inference contexts.