ATLAS Sampler: Adaptive MCMC for Complex Posteriors

• ATLAS Sampler is an adaptive MCMC method that dynamically adjusts leapfrog step sizes within each orbit to control energy error.
• It uses a dyadic scheduling strategy to select the largest feasible micro-step size locally, ensuring robust exploration of multi-scale targets.
• Empirical results demonstrate competitive effective sample sizes and enhanced robustness over standard NUTS in challenging posterior geometries.

The Within-Orbit Adaptive Leapfrog No-U-Turn Sampler (WALNUTS) is a Markov chain Monte Carlo (MCMC) algorithm designed to enhance Hamiltonian Monte Carlo (HMC) by providing adaptive, localized control over the leapfrog integrator's step size within each simulated orbit. The method generalizes the No-U-Turn Sampler (NUTS), introducing a dynamic mechanism that selects the largest feasible leapfrog step size at fixed intervals along a trajectory, constrained by a user-defined threshold on energy error. This adaptation improves sampling efficiency and robustness in posterior distributions with pronounced multi-scale geometry and local stiffness, while maintaining reversibility, detailed balance, and ergodicity (Bou-Rabee et al., 23 Jun 2025).

1. Background: HMC, NUTS, and the Need for Local Adaptivity

Hamiltonian Monte Carlo (HMC) simulates Hamiltonian dynamics in an augmented phase space $(q,p)$, where $q\in\mathbb{R}^d$ (the variable of interest) and $p\in\mathbb{R}^d$ (auxiliary momentum). The joint density is parameterized by $U(q) = -\log\mu(q)$ (potential energy) and $K(p) = \frac12 p^\top M^{-1}p$ (kinetic energy), yielding a total Hamiltonian $H(q,p) = U(q) + K(p)$. These dynamics are approximated via the leapfrog integrator with step size $\epsilon$, which is efficient and volume-preserving but introduces small energy errors, corrected by Metropolis–Hastings acceptance.

The NUTS algorithm innovates on HMC by automatically tuning the integration time (orbit length) via a geometric "no-U-turn" condition, but still employs a fixed step size $\epsilon$ throughout the trajectory. Complex posteriors with varying curvature or strongly multi-scale structure necessitate local step size adaptation: regions of higher stiffness require smaller steps to control integration error, while smooth regions permit larger, more efficient steps. Existing methods cannot locally adapt $\epsilon$ within an orbit while maintaining the reversibility and detailed balance that guarantee proper stationary distribution sampling. WALNUTS addresses this gap by controlling error accumulation at the level of individual macro-steps within each trajectory (Bou-Rabee et al., 23 Jun 2025).

2. Adaptive Energy-Controlled Leapfrog Scheduling

WALNUTS divides each orbit into fixed-length macro-steps governed by a nominal step size $\epsilon_0$. Within each macro-step, the leapfrog micro-step size is chosen adaptively on a dyadic schedule: candidate micro-step sizes are $q\in\mathbb{R}^d$0 for $q\in\mathbb{R}^d$1. For each candidate partition, the algorithm computes the maximal and minimal Hamiltonian values $q\in\mathbb{R}^d$2 over the $q\in\mathbb{R}^d$3 sub-steps, where $q\in\mathbb{R}^d$4. The total energy spread, $q\in\mathbb{R}^d$5, is compared against the user-specified threshold $q\in\mathbb{R}^d$6.

The largest candidate $q\in\mathbb{R}^d$7 with $q\in\mathbb{R}^d$8 is selected for the current macro-step, ensuring that local integration error remains controlled. This procedure is repeated at each macro-step, allowing the WALNUTS orbit to adapt micro-step size to local stiffness while maximizing progress in regions with favorable geometry (Bou-Rabee et al., 23 Jun 2025).

3. WALNUTS Orbit Construction, State Selection, and Algorithmic Details

WALNUTS constructs each Hamiltonian trajectory on a macro-grid $q\in\mathbb{R}^d$9, adaptively partitioning each macro-interval into $p\in\mathbb{R}^d$0 micro-steps of length $p\in\mathbb{R}^d$1 as determined by the dyadic scheduling criterion. The extension direction for each segment is selected via a sequence of Bernoulli draws, producing a binary doubling process for trajectory expansion (as in NUTS).

In each extension, the "sub-U-turn" condition is checked to avoid reversals before the main U-turn criterion is violated, terminating the expansion when detected. Like NUTS, WALNUTS employs a progressive, biased state selection: when a new segment is added to the orbit, a candidate is drawn via categorical sampling proportional to the segment's weights, and accepted with probability $p\in\mathbb{R}^d$2. This induces a bias toward points further from the starting position, favoring global exploration in the presence of skewed or funnel-like geometries.

Final state selection occurs after all doublings and U-turn checks, with the output state indexed by the last accepted progressive sample. WALNUTS requires no additional Metropolis–Hastings correction step beyond the acceptance in progressive sampling; the overall transition map is reversible and measure-preserving on the extended path space (Bou-Rabee et al., 23 Jun 2025).

4. Theoretical Properties: Invariance, Reversibility, and Complexity

WALNUTS can be analyzed via the auxiliary variable + involution formalism. The orbit construction, parameter selection ($p\in\mathbb{R}^d$3 sequences), and state selection can be embedded in a mapping $p\in\mathbb{R}^d$4 that is an involutive, measure-preserving transformation of the extended state space. The transition kernel defined by WALNUTS leaves the target distribution $p\in\mathbb{R}^d$5 invariant and is reversible with respect to $p\in\mathbb{R}^d$6, as shown by direct construction and appeal to the auxiliary-involution theorem (see Theorem 1 of (Bou-Rabee et al., 23 Jun 2025)). No separate Metropolis–Hastings step is necessary; the algorithm inherently preserves detailed balance.

Ergodicity and irreducibility of the chain follow from the standard conditions on HMC (regularity, positive-definite kinetic energy), with convergence guarantees inherited from the underlying Markov structure.

Computational complexity per iteration is dominated by gradient evaluations from the leapfrog integrator. If an orbit traverses $p\in\mathbb{R}^d$7 macro steps with average $p\in\mathbb{R}^d$8 micro steps each, the total number of gradient evaluations is $p\in\mathbb{R}^d$9. The energy-control criterion ensures that $U(q) = -\log\mu(q)$0 is minimal in well-conditioned regions and increases only when necessitated by local stiffness. The dyadic search for admissible $U(q) = -\log\mu(q)$1 incurs only $U(q) = -\log\mu(q)$2 overhead per macro-step (Bou-Rabee et al., 23 Jun 2025).

5. Empirical Results and Comparative Performance

Benchmark evaluations of WALNUTS conduct comparisons to standard NUTS on multiscale target distributions:

• On high-dimensional Gaussian targets $U(q) = -\log\mu(q)$3, WALNUTS–D (deterministic two-point schedule) matches or slightly exceeds NUTS in effective sample size (ESS) per 1000 gradients for $U(q) = -\log\mu(q)$4. WALNUTS–R2P (random-2-point scheduling) trails NUTS by less than 10% in ESS but exhibits greater robustness to non-equilibrium initializations.
• On Neal’s funnel distribution, NUTS requires significant tuning to avoid divergences and typically fails to explore the narrow "neck" region, as evidenced by near-zero ESS in that domain. WALNUTS, with $U(q) = -\log\mu(q)$5 and $U(q) = -\log\mu(q)$6, correctly recovers the marginal distributions across the entire space with comparable computational cost.
• On the Stock–Watson stochastic volatility model ($U(q) = -\log\mu(q)$7 time steps), WALNUTS (with $U(q) = -\log\mu(q)$8, $U(q) = -\log\mu(q)$9) eliminates pathologies (no divergent transitions), flexibly utilizes large steps in smooth regimes and small steps in stiff sections, and produces comparable or improved ESS per gradient relative to NUTS, which suffers from large energy errors ($K(p) = \frac12 p^\top M^{-1}p$0) in 5–10% of orbits and requires very small steps to function at all (Bou-Rabee et al., 23 Jun 2025).

These results demonstrate that within-orbit adaptive step size selection in WALNUTS yields improved or comparable sampling efficiency and markedly superior robustness, particularly in regions where traditional NUTS integration breaks down.

6. Significance, Limitations, and Future Directions

WALNUTS advances the HMC framework by enabling fine-grained, reversible, within-trajectory step size adaptation based on rigorous energy-error control. This approach effectively addresses longstanding challenges in sampling from distributions with extreme local scale variation or funnel-like geometries, where existing global or pre-orbit step size adaptation is insufficient.

The theoretical structure of WALNUTS ensures that all desirable properties of NUTS—detailed balance, ergodicity, unbiasedness—are preserved, while the empirical findings confirm that localized adaptation improves practical robustness on challenging targets. The method is parameterized by the energy-error tolerance $K(p) = \frac12 p^\top M^{-1}p$1 and the macro-step size $K(p) = \frac12 p^\top M^{-1}p$2, requiring only minor modification of existing NUTS implementations to deploy.

Potential avenues for further research include refinement of the dyadic scheduling policy, integration with Riemannian manifold metrics, and generalization to more complex kinetic energy forms. A plausible implication is that WALNUTS’ trajectory-based adaptation could benefit other Hamiltonian simulation tasks where local control of discretization error is critical (Bou-Rabee et al., 23 Jun 2025).