LazyHMC: Hamiltonian Monte Carlo Simulation for Lazy, Infinite Dimensional Probabilistic Programs

AI summary

Overview

Research area: Probabilistic programming, non-parametric Bayesian statistics, automatic differentiation, and the semantics of functional programming languages. The paper sits at the intersection of Markov chain Monte Carlo (MCMC) methods, lazy functional programming in Haskell, and measure-theoretic / denotational semantics.

Technical level: Advanced. The paper combines a new smoothness definition on infinite-dimensional spaces, a core lambda calculus with a logical-relations proof, and a Haskell library implementation.

One-sentence scope: The paper shows how to run gradient-based Hamiltonian Monte Carlo (HMC) and No-U-Turn Sampler (NUTS) inference directly on probabilistic programs whose parameter spaces are infinite-dimensional, by exploiting lazy evaluation so that only finitely many dimensions are ever touched.

What This Paper Is About

Hamiltonian Monte Carlo is one of the most successful generic inference methods in probabilistic programming, but it needs gradients and a finite-dimensional parameter space. Meanwhile, programmers can write non-parametric Bayesian models — models with an unknown or unbounded number of parameters, such as clustering with an unknown number of clusters or stochastic processes like random walks — as lazy Haskell programs in the LazyPPL library, where infinite lists stand for random processes. This paper's goal is to bring the efficiency of gradient-based HMC into that lazy, infinite-dimensional setting without giving up either compositionality or the declarative style, by showing that the gradient and the Metropolis-Hastings acceptance ratio collapse to finite objects even though the underlying space is infinite.

Key Contributions

1. A smoothness condition for infinite-dimensional gradients (PACAP). The authors introduce "piecewise analytic under cylindrical analytic partition" (PACAP), a notion of piecewise smoothness for functions on an infinite-dimensional seed space 𝕏 = ℝ^𝒜. It refines earlier finite-dimensional notions (PAP) that are not stable under composition, and it guarantees that locally a function depends on only finitely many coordinates. They prove PACAP holds for a core calculus with real numbers, functions, finite products and countable sums, capturing primitive recursion and corecursion.

2. Infinite-dimensional automatic differentiation. Building on PACAP, the authors adapt a recent idea based on "sparse-map Nagata numbers" to find gradients in this infinite-dimensional setting, showing that although an infinite stream of parameters is maintained lazily, the gradient of the likelihood is non-zero in only a finite subset of dimensions.

3. Lazy HMC and lazy NUTS with a finite Metropolis-Hastings correction. They develop several HMC variants and a lazy No-U-Turn Sampler that operate over the infinite-dimensional parameter space but stay productive because of lazy evaluation. They carefully choose the momentum and structure the Hamiltonian so that the acceptance ratio — ostensibly an infinite product over all dimensions — collapses to a finite product, proposing three different ways of achieving this. Section 5 gives a worked example of piecewise-constant regression showing how the leapfrog integrator and dimension changes interact.

4. An implementation and experimental evaluation in LazyPPL. Lazy HMC and lazy NUTS are implemented as an extension to the LazyPPL library and evaluated on non-parametric models from the literature, including geometric distributions, random walks, and clustering (a Gaussian mixture with a Poisson-distributed number of clusters).

Main Findings

• Gradients have finite support. For programs with a PACAP likelihood, the gradient of the unnormalized density is non-zero in only finitely many dimensions, so leapfrog steps can be performed over just the relevant coordinates while the rest are handled lazily.
• The Metropolis-Hastings ratio becomes a finite product. On unvisited coordinates the leapfrog involution acts as a rotation preserving the standard-normal density, so the corresponding factors cancel, leaving a finite product even though the parameter space is genuinely infinite-dimensional.
• PACAP is compositional. Theorem 3 states that any function f : ⟦(𝒜 → real) → real⟧ definable in the core calculus is PACAP. The proof is by infinitary logical relations, with a "fundamental lemma" proved by induction on typing derivations, using lemmas on support enlargement, intersection, restriction, gluing, and composition.
• Non-smooth points are measure zero. The points where the number of relevant dimensions changes (for example, the geometric distribution's sample value is the number of dimensions used) are shown to form a measure-zero set, so the gradient exists almost everywhere. The authors note this analysis is subtle because "almost surely smooth functions" are not closed under composition (citing Lee et al., 2020), which is why they work with analytic functions and analytic partitions.
• Positive empirical behaviour relative to gradient-free lazy inference. The paper reports that lazy HMC successfully samples from the posterior of the random walk model despite operating over a truly infinite-dimensional space (Figure 3(b)), and that gradient information allows more efficient exploration than the gradient-free lazyLMH method.
• Scope limitation is explicit. The paper states that there is no standard benchmark suite for non-parametric statistics, and that the experiments provide empirical evidence that the methods work as intended. The truncated content provided does not report specific quantitative experimental figures, timings, acceptance rates, or sample counts.
• Functions outside scope. A function such as f(x₁, …, xₙ, …) = Σ 2^(−i)·arctan(xᵢ) is not PACAP because its result depends on every input dimension; the authors note such functions lie outside their scope since the likelihood is then not finitely computable.

Methodology in Plain English

The authors start from the observation that Haskell's lazy evaluation already lets statisticians write down infinite models naturally: an infinite list of independent uniform samples is a stochastic process, and only the elements actually accessed get generated. They work inside the LazyPPL library, which provides a probability monad Prob (whose recursion can be lazy, so infinite random sequences can be built), and a measures monad Meas (whose recursion is not lazy, because observations must terminate to contribute a score). The seed space is fixed as Ω ≅ Ω × Ω with Ω = ℝ^(ℕ*) — effectively infinite rose trees — so randomness can always be split in two, which is what makes the lazy monadic sequencing work.

To do gradient-based inference, they need to make sense of a "gradient" on an infinite product of real coordinates. Their strategy is local: in any finite computation, a program can only inspect finitely many coordinates. They formalise this with the PACAP condition, which says the domain can be carved into countably many "analytic cylinders" (regions defined by finitely many coordinates and finitely many analytic inequalities), each supporting only finitely many coordinates, on which the function agrees with an analytic function of those coordinates. They then prove, using a logical-relations argument over a small typed lambda calculus with countable sums, that every definable function is PACAP — so the property survives being composed in real programs. For programs using unbounded search (like findIndex or find, which the geometric distribution and random walk models use), they verify PACAP directly by an explicit partition instead.

With differentiability established almost everywhere, they build the HMC loop: momentum, leapfrog integration, and a Metropolis correction. The clever part is choosing the momentum distribution and structuring the Hamiltonian so that when a leapfrog step happens only on the finitely many active coordinates, the untouched coordinates undergo a measure-preserving rotation. That cancellation is what makes the infinite product of acceptance terms finite and computable. They also add adaptive trajectory-length selection in a lazy NUTS variant, and implement the whole thing as an extension of LazyPPL, then test it on models where the number of parameters is genuinely unknown.

Why This Matters

Impact on research. The paper removes a long-standing tension in probabilistic programming: previous routes to non-parametric models either sacrificed compositionality (manual truncation to a finite-dimensional approximation), declarativeness (tracking dynamic dimensions explicitly), or efficiency (gradient-free lazy Metropolis-Hastings, as in Dash et al., 2023, and Bowers et al., 2025). By putting gradient-based inference directly into the lazy infinite-dimensional setting, the work connects measure-theoretic semantics and automatic differentiation theory with practical, high-throughput Bayesian inference, and supplies a reusable smoothness notion (PACAP) for reasoning about differentiability over infinite products.

Applications explored in the paper (these are the model classes the authors implement and evaluate, not deployed systems):

• Non-parametric clustering with a Gaussian mixture where the number of clusters is Poisson-distributed and therefore unknown.
• Random walks with observations, including the odometer variant from Mak et al. (2021) where a walker rests on reaching home or exceeding a distance limit.
• Piecewise-constant regression with Poisson-process changepoints, used as the detailed worked example of how leapfrog steps and dimension changes interact.
• IID sequences and geometric distributions, used to expose the subtlety that the sample value itself equals the number of dimensions consumed.

Industry relevance. Gradient-based samplers are the engine behind most practical Bayesian modelling; extending them to unbounded parameter spaces matters for practitioners who need models whose complexity grows with the data — mixture models with unknown component counts, changepoint detection, and process-based time-series models — without hand-tuning truncation levels globally across a model. Because the implementation lives in a Haskell library (LazyPPL) and the smoothness guarantee is derived compositionally for a small core calculus, the results also matter to language designers and compiler writers building differentiable or probabilistic languages.

Future Directions

• Extending the core calculus to unbounded search. Programs such as findIndex and find, used by the geometric and random walk models, fall outside the core calculus and need PACAP verified by hand with an explicit partition. A general, automatic treatment of such search-based programs is an open step.
• Broadening the smoothness condition. Functions that depend on every coordinate, such as the weighted arctangent sum given as a counterexample, are excluded because the likelihood is then not finitely computable. Understanding how much of the non-parametric modelling space lies beyond PACAP is unresolved.
• Benchmarking. The paper explicitly notes there is no standard benchmark suite for non-parametric statistics. Building one would allow the lazy HMC and lazy NUTS variants to be compared systematically against gradient-free lazy sampling and against truncation-based baselines.
• Scaling and comparison of the three Metropolis-Hastings variants. Three different ways of making the infinite-dimensional correction collapse to a finite product are proposed; the paper's own framing invites further study of which is preferable and under what conditions, and of the lazy NUTS extension described in Appendix D.

Target Audience

Researchers and graduate students in probabilistic programming, Bayesian non-parametrics, and programming language semantics; developers of probabilistic languages and compilers who need to implement gradient-based inference over lazily evaluated, unbounded data structures; and mathematically inclined statisticians interested in automatic differentiation and MCMC on infinite-dimensional spaces. Readers need comfort with Haskell-style functional programming, measure theory, and Monte Carlo methodology; the paper is not an introductory treatment.

Authors’ abstract

Hamiltonian Monte Carlo (HMC) is a successful generic inference method in probabilistic programming, but in its ordinary formulation it needs gradients and finite-dimensional parameter spaces. In Haskell, lazy evaluation lets probabilistic programs express stochastic processes and other non-parametric Bayesian models over implicit infinite-dimensional spaces. This paper develops new formulations of gradient-based HMC for this infinite-dimensional setting, via lazy evaluation. For automatic differentiation, we provide an analysis based on a new notion of "piecewise analytic under cylindrical analytic partition" (PACAP), to show that even if a program is infinite-dimensional and defined lazily, the gradient of the likelihood function is finitely supported. For the Monte Carlo method itself, we develop several HMC variants and a No-U-Turn Sampler that operate over the infinite-dimensional parameter space but are still productive because of lazy evaluation. Experiments cover Gaussian mixture clustering, random walks, and piecewise-constant regression with Poisson-process changepoints.