📊 statistics

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

This paper introduces LazyHMC, a novel framework that extends Hamiltonian Monte Carlo to infinite-dimensional probabilistic programs in Haskell by leveraging deferred evaluation and a new "PACAP" analysis to ensure gradients are finitely supported, thereby enabling efficient gradient-based inference for non-parametric Bayesian models.

Original authors: Maria-Nicoleta Crăciun, C. -H. Luke Ong, Tom Schrijvers, Sam Staton

Published 2026-08-11

📖 8 min read🧠 Deep dive

CC BY 4.0

Original authors: Maria-Nicoleta Crăciun, C. -H. Luke Ong, Tom Schrijvers, Sam Staton

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a giant jigsaw puzzle, but the box says there are infinite pieces. In the world of computer science and statistics, this is a common problem called "non-parametric" modeling. Instead of guessing a fixed number of variables (like "there are exactly 5 clusters of data"), the computer tries to figure out how many pieces it actually needs, potentially forever. To solve these puzzles, scientists use a clever trick called Hamiltonian Monte Carlo (HMC). Think of HMC as a hiker trying to find the deepest valley in a foggy mountain range. Instead of taking tiny, random steps (which is slow and inefficient), the hiker throws a ball with momentum. The ball rolls down the slopes, gathering speed, allowing the hiker to leap across the foggy hills and find the best spot much faster. However, there's a catch: standard hikers need a map with a fixed number of coordinates. If the map is infinite, the hiker gets stuck because they can't calculate the slope of an infinite mountain.

This is where a new paper comes in, offering a solution for computers that are "deferred" in the best possible way. The researchers, working with a programming language called Haskell, realized that while the mountain might be infinite, the hiker only ever needs to look at a few specific rocks to know which way to roll. They developed a new method called DeferredHMC. It's like having a hiker who only looks at the ground directly under their feet and the few steps ahead, ignoring the rest of the infinite mountain until they actually need to step there. This allows the computer to solve complex, infinite puzzles that were previously too difficult for standard methods, using the speed of momentum without getting overwhelmed by infinity.

The Problem: Infinite Mountains and Stuck Hikers

In the world of probabilistic programming, statisticians write code to describe how data is generated. Sometimes, they don't know how many "parameters" (the knobs and dials of the model) they need. Maybe they are clustering data points and don't know if there are 3 groups or 300. Or maybe they are modeling a random walk that goes on forever. In these cases, the "parameter space" is infinite.

Standard HMC is a superstar for finite problems. It uses gradients (slopes) to guide the search. Imagine the likelihood of a model as a landscape: high hills are bad fits, deep valleys are good fits. HMC uses the slope of the hill to roll a ball toward the valley. But to calculate the slope, you need to know the height of the hill at every point. If the hill has infinite dimensions, calculating the slope seems impossible.

Previously, people tried to fix this in three ways, but all had flaws:

Truncation: They just cut off the mountain after a certain height. This is like saying, "We'll only look at the first 100 pieces of the puzzle." It works, but it's messy because you have to guess where to cut, and if you cut too early, you miss the solution.
Dynamic Dimensions: They build the mountain piece by piece as they go. This is better, but it makes the code complicated and less "declarative" (less about what the problem is and more about how to build it).
Deferred Methods without Gradients: They used "deferred" evaluation (only computing what's needed) but gave up on the momentum. This meant the hiker had to take tiny, random steps again, which is very slow.

The paper asks: Can we have the best of both worlds? Can we use the powerful, fast momentum of HMC on an infinite mountain, but only look at the parts of the mountain we actually need?

The Solution: The Deferred Hiker

The authors, Maria-Nicoleta Crăciun and her team, say yes. They developed DeferredHMC, a method that combines the power of gradients with the "deferred" nature of Haskell programming.

Here is the magic trick: Even though the mountain is infinite, the hiker (the computer program) only ever touches a finite number of rocks to figure out the slope.

The "PACAP" Insight: The team proved mathematically that for a wide class of programs, the "slope" (gradient) is only non-zero at a finite number of places. Even if the model has infinite potential dimensions, the specific data you are looking at only "activates" a few of them. It's like a giant library with infinite books, but for a specific question, you only need to read three pages. The rest of the library is irrelevant for that moment.
The Deferred Evaluation: In Haskell, values are only computed when they are needed. The computer doesn't generate the whole infinite list of random numbers; it just generates the first one, then the second one, only when the math demands it.
The Cancellation Trick: This is the most clever part. In HMC, you have to calculate an "acceptance ratio" to decide if a new position is good. Normally, this involves multiplying numbers for every single dimension. If there are infinite dimensions, this product is infinite. But the authors showed that for the dimensions the hiker didn't visit, the math works out perfectly so that those terms cancel each other out (they multiply to 1). So, the infinite product collapses into a tiny, finite product involving only the visited dimensions.

How It Works in Practice

The paper introduces a few different versions of this deferred hiker, but they all follow the same logic:

Step 1: The hiker starts at a position and gets a random "kick" (momentum).
Step 2: The hiker rolls down the slope. As they roll, they only compute the parts of the infinite mountain they actually step on.
Step 3: If the hiker steps on a new rock that wasn't there before (activating a new dimension), the system deferredly generates the necessary data for that spot.
Step 4: The hiker decides whether to stay or go back. Because the unvisited parts of the mountain cancel out in the math, the hiker doesn't need to know about them.

The authors also created a Deferred NUTS (No-U-Turn Sampler). Standard HMC requires you to guess how far to roll (the number of steps). If you roll too little, you don't explore enough; too much, and you waste time. NUTS is a smart algorithm that automatically stops when it starts to double back on itself. The authors made a deferred version of this too, so the computer can decide how far to roll without ever needing to see the whole infinite mountain.

The Experiments: Does It Actually Work?

To test their idea, the team ran simulations on several tricky problems:

Geometric Distribution: A model where the number of steps is unbounded. The deferred HMC methods were at least 10 times faster than the previous best method (NP-HMC) and produced accurate results.
Random Walks: A model where a walker moves forever. The deferred methods successfully found the starting position of the walker, even though the walk is infinite. The deferred NUTS method was slower in raw time but found much better solutions (higher "effective sample size") than the others.
Clustering (Gaussian Mixture): A model where the number of groups is unknown. The deferred methods correctly identified the number of clusters and the position of the centers, performing as well as or better than existing methods.
Polynomial Regression: A model where the complexity (degree) of the curve is unknown. The deferred methods found the correct degree (5) and fit the data well, whereas older methods got stuck on the wrong complexity.

What This Means

The paper doesn't claim to have solved every problem in statistics. It doesn't say this is the final word on all infinite models. However, it demonstrates that gradient-based HMC can work directly on infinite-dimensional spaces without needing to cut them off or manually track dimensions.

The key takeaway is that deferred evaluation is a superpower. By waiting until the very last moment to compute a value, the computer avoids the impossible task of handling infinity. Instead, it handles a manageable, finite slice of the problem, and the math ensures the rest of the infinite world doesn't matter.

The authors suggest that while this works well for the models they tested, real-world systems might still need human guidance to tune the settings (like how big the steps should be). But for now, they have opened the door to a new way of doing Bayesian inference: one where the computer can explore the infinite without getting lost.

Technical Summary: LazyHMC

Problem Statement

Probabilistic programming offers a declarative framework for Bayesian inference, allowing statisticians to specify complex models via code. However, generic inference for non-parametric models—those defined over infinite or unbounded parameter spaces (e.g., stochastic processes, Dirichlet processes)—remains a significant challenge.

Existing approaches to non-parametric inference suffer from specific trade-offs:

Truncation: Approximating infinite dimensions with finite ones sacrifices compositionality, requiring manual adjustment of bounds across the model.
Dynamic Dimensions: Tracking dimensions as they are needed preserves accuracy but sacrifices declarativeness by explicitly managing dimension counts.
Gradient-Free On-Demand Methods: Approaches like on-demand Metropolis-Hastings maintain compositionality and declarativeness but lack the efficiency of gradient-based methods, often resulting in poor mixing and slow convergence.

The central problem addressed by this paper is the absence of a method that combines the efficiency of gradient-based Hamiltonian Monte Carlo (HMC) with the compositional and declarative nature of on-demand, infinite-dimensional probabilistic programming. Standard HMC relies on finite-dimensional gradients and cannot be directly applied to spaces where the dimensionality is unbounded or unknown.

Methodology

The paper proposes LazyHMC, a framework that extends HMC to infinite-dimensional spaces by leveraging on-demand evaluation and a new theoretical analysis of smoothness. The methodology rests on three pillars:

1. Theoretical Foundation: PACAP and Infinite-Dimensional Gradients

To apply HMC, one must compute gradients of the log-likelihood function over an infinite-dimensional seed space (e.g., $R^{\mathbb{N}^*}$ ). The authors introduce a smoothness condition called Piecewise Analytic under Cylindrical Analytic Partition (PACAP).

Definition: A function is PACAP if its domain can be partitioned into countably many "analytic cylinders" (sets defined by constraints on a finite subset of coordinates) where the function is analytic.
Key Insight: For programs definable in a core calculus (covering primitive recursion and corecursion), the likelihood function is PACAP. Crucially, this implies that for any given state, the gradient is finitely supported; it is non-zero only on a finite subset of dimensions that are actually "visited" (evaluated) during the computation of the likelihood.
Implementation: The authors utilize Nagata numbers (a sparse-map representation of dual numbers) for automatic differentiation. This allows the system to automatically discover the finite set of dimensions with non-zero gradients without iterating through all infinite dimensions, a feat impossible with standard forward-mode AD in infinite spaces.

2. On-Demand Hamiltonian Monte Carlo Framework

The paper adapts the Involutive MCMC (iMCMC) framework to infinite-dimensional "rose trees" (infinite maps from addresses to real numbers).

The Challenge: The Metropolis-Hastings acceptance ratio formally involves an infinite product over all dimensions.
The Solution: By carefully structuring the proposal mechanism (the involution), the authors ensure that for all unvisited dimensions (those with zero gradient), the proposal acts as a measure-preserving transformation (specifically a rotation or identity) that leaves the prior density unchanged. Consequently, the terms in the acceptance ratio corresponding to unvisited dimensions cancel out exactly.
Result: The infinite acceptance ratio collapses to a finite product over only the visited dimensions, making the computation feasible.

3. Algorithmic Variants

The paper presents three specific instantiations of LazyHMC, differing in how they handle the involution and visited sets:

LazyHMC1: Uses a rotation-based position update. It requires tracking only the set of dimensions with non-zero gradients ( $u$ ), as the rotation preserves the density on unvisited sites. This is the most purely functional variant.
LazyHMC2: Uses a different Hamiltonian splitting (joint rotation of position and momentum). It shares the cancellation properties of LazyHMC1 but may offer better performance when the prior geometry interacts poorly with LazyHMC1's step size.
LazyHMC3: Uses the standard leapfrog integrator (translation). Since translations do not preserve the standard normal density, this variant requires tracking the full set of visited sites ( $v$ ) to apply a corrective measure-preserving map to unvisited sites, ensuring the acceptance ratio remains valid.
LazyNUTS: Adapts the No-U-Turn Sampler (NUTS) to the on-demand setting, allowing the trajectory length to be determined adaptively without manual tuning, while still evaluating only finitely many sites per iteration.

Key Contributions

PACAP Analysis: A formal proof that a core calculus of probabilistic programs (including recursion and corecursion) yields likelihood functions that are PACAP, guaranteeing finite support for gradients in infinite-dimensional spaces.
On-Demand Automatic Differentiation: An adaptation of Nagata numbers to infinite-dimensional spaces, enabling the automatic discovery of relevant dimensions for gradient computation without explicit dimension tracking.
On-Demand HMC Variants: The development of three HMC algorithms (LazyHMC1, 2, 3) and an On-Demand NUTS sampler that operate over infinite-dimensional parameter spaces while maintaining the correctness guarantees of MCMC (stationarity) via the iMCMC framework.
Correctness Framework: A rigorous measure-theoretic framework (Framework A) demonstrating how the infinite-dimensional acceptance ratio collapses to a finite product under specific involution conditions.

Experimental Results

The authors implemented LazyHMC and LazyNUTS as an extension to the LazyPPL library in Haskell. They evaluated the methods on non-parametric models where standard HMC is inapplicable or requires truncation:

Geometric Distribution: A model with an unbounded number of random choices. LazyHMC variants achieved low Total Variation Distance (TVD) from the ground truth (approx. 0.01–0.03) and were significantly faster (approx. 10x) than the non-truncated NP-HMC baseline.
Random Walk: A model involving an infinite stochastic process. LazyHMC variants demonstrated higher effective sample sizes (ESS) per second compared to NP-HMC and showed comparable posterior densities.
Gaussian Mixture Model (GMM): A clustering model with an unknown, Poisson-distributed number of clusters. The methods successfully inferred the number of components and cluster means, with predictive densities (LPPD) comparable to NP-HMC.
Polynomial Regression: A model with an unbounded polynomial degree. LazyHMC and LazyNUTS successfully concentrated posterior mass on the correct degree (5), whereas the gradient-free on-demand LMH method struggled, often getting stuck on higher degrees.

In general, the on-demand HMC variants were competitive with or faster than NP-HMC (which avoids truncation but lacks on-demand evaluation optimizations) and significantly outperformed gradient-free on-demand Metropolis-Hastings in terms of mixing efficiency.

Significance and Claims

The paper claims to demonstrate that gradient-based HMC can be applied directly to infinite-dimensional parameter spaces without resorting to truncation or dynamic dimension tracking. The significance lies in preserving the compositional and declarative benefits of probabilistic programming while gaining the efficiency of gradient-based inference.

The authors are modest in their claims, noting that:

There is no standard benchmark suite for non-parametric statistics, so the experiments serve as empirical evidence of correctness and feasibility rather than a definitive performance leaderboard.
The current implementation relies on manual tuning of hyperparameters (step size, trajectory length), and future work should incorporate automatic adaptation (e.g., dual averaging).
While the approach is purely declarative, real-world complex models may still require modeller interaction (e.g., guiding the inference engine), similar to other advanced probabilistic programming systems.

Ultimately, the paper establishes a theoretical and practical pathway for performing efficient Bayesian inference on models that naturally live in infinite-dimensional spaces, bridging the gap between the expressiveness of on-demand probabilistic programming and the power of Hamiltonian Monte Carlo.