Complexity Analysis of Normalizing Constant Estimation: from Jarzynski Equality to Annealed Importance Sampling and beyond

This paper establishes the first non-asymptotic oracle complexity bounds for annealed importance-based normalizing constant estimation without relying on isoperimetric assumptions and proposes a novel reverse diffusion sampler to overcome the limitations of traditional geometric interpolation in multimodal settings.

The Gist

Imagine you are trying to figure out the total size of a vast, foggy landscape. You can see the hills and valleys (the "energy" of the system), but the fog is so thick that you can't see the whole picture at once. In the world of statistics and machine learning, this "total size" is called the normalizing constant. It's a crucial number needed to make probabilities add up correctly, but calculating it is notoriously difficult, especially when the landscape has many separate peaks (multimodal) or is incredibly high-dimensional.

This paper, presented at ICLR 2026, tackles the question: "How hard is it to calculate this number, and can we do it faster and more reliably?"

Here is a breakdown of their findings using simple analogies.

1. The Problem: The "Foggy Mountain"

Imagine you are a hiker trying to measure the total area of a mountain range.

• The Old Way (Importance Sampling): You pick a spot, look around, and guess the size of the whole range based on that one view. If the mountains are complex (lots of peaks and valleys), your guess is usually terrible because you miss the other peaks entirely. It's like trying to guess the size of a forest by looking at just one tree.
• The "Annealing" Solution: Instead of guessing from one spot, you build a bridge. You start at a simple, flat plain (where you know the size) and slowly transform the landscape into the complex mountain range. You take small steps along this bridge, measuring the changes. This is called Annealing.

2. The Two Main Bridges: JE and AIS

The paper analyzes two popular ways to build this bridge:

• Jarzynski Equality (JE): Think of this as a physics experiment. You pull a rubber band (the system) from a relaxed state to a stretched state very quickly. By measuring the "work" (energy) you put in during many different fast pulls, you can mathematically calculate the difference in energy between the start and end.
• Annealed Importance Sampling (AIS): This is more like a guided tour. You take a group of hikers (samples) and slowly move them from the flat plain to the mountain peaks, stopping at many intermediate campsites. At each stop, you adjust the group's position to match the terrain.

The Paper's Big Discovery:
For a long time, we knew these methods worked well in practice, but we didn't have a precise mathematical rulebook for how long the bridge needs to be to get an accurate answer. The authors created this rulebook. They proved that the difficulty (complexity) of the task depends on something they call the "Action" of the bridge.

• The "Action" Analogy: Imagine the bridge is a path. If the path is smooth and direct, the "Action" is low, and the calculation is easy. If the path is jagged, requires teleporting hikers across huge gaps, or twists violently, the "Action" is high, and the calculation becomes exponentially harder.

3. The Trap of the "Geometric" Bridge

For years, scientists have used a specific type of bridge called Geometric Interpolation. It's popular because it's easy to write down on paper.

• The Paper's Warning: The authors discovered that for complex, multi-peaked landscapes (like a mountain range with two distant peaks), this geometric bridge is actually a trap.
• The "Teleportation" Problem: To get from one peak to another using this specific bridge, the math forces the hikers to "teleport" across the empty space between the peaks. This requires an impossible amount of energy (infinite "Action"). The paper proves mathematically that for certain difficult problems, this method will fail or take an impossibly long time.

4. The New Solution: The "Reverse Diffusion" Elevator

Since the standard bridge is too shaky for complex mountains, the authors propose a new method based on Reverse Diffusion Samplers.

• The Analogy: Imagine the landscape is being slowly covered in fog until it disappears completely into a uniform white mist (a standard Gaussian distribution). This is a "forward" process.
• The Innovation: Instead of building a bridge from the mist to the mountain, the authors suggest running the process in reverse. You start in the uniform mist and slowly "un-cover" the fog, letting the landscape reveal itself naturally.
• Why it works better: This reverse process acts like a guided elevator that gently carries hikers from the mist to the peaks without forcing them to teleport. It naturally handles the "jumps" between peaks that the old method struggled with.

5. The Results: A Race to the Top

The authors tested their new "Reverse Diffusion" method against the old "Geometric" methods (TI and AIS) on two difficult test cases:

1. The Müller Brown Landscape: A classic, tricky mountain range used in physics.
2. The Gaussian Mixture: A landscape with four distinct, separated peaks.

The Outcome:

• Old Methods (TI & AIS): They got stuck. The hikers stayed in the first valley they started in and never found the other peaks. Their estimates of the total size were wildly wrong (biased).
• New Method (Reverse Diffusion): The hikers successfully explored all the peaks. The estimates were accurate, and the "samples" (the hikers' positions) matched the true landscape perfectly.

Summary

This paper provides the first rigorous mathematical proof of how hard it is to calculate these "normalizing constants" without making unrealistic assumptions about the landscape.

1. They showed that the difficulty is determined by the smoothness of the path you take.
2. They proved that the most common path (Geometric Interpolation) is often too jagged and causes "teleportation" failures.
3. They introduced a new, smoother path (Reverse Diffusion) that acts like a gentle elevator, successfully navigating complex, multi-peaked landscapes where old methods fail.

In short: If you need to measure a complex, foggy landscape, don't try to build a shaky bridge across the gaps. Instead, use the new "reverse fog" elevator to reveal the terrain naturally.

Technical Summary

Technical Summary: Complexity Analysis of Normalizing Constant Estimation

Problem Statement

The paper addresses the fundamental problem of estimating the normalizing constant $Z = \int_{\mathbb{R}^d} e^{-V(x)} dx$ (or equivalently, the free energy $F = -\log Z$) for an unnormalized probability density $\pi \propto e^{-V}$. This task is critical in Bayesian statistics (marginal likelihood), statistical mechanics (partition functions), and machine learning (training energy-based models). The problem is particularly challenging in high dimensions or when the target distribution $\pi$ is multimodal, as conventional importance sampling suffers from high variance due to the mismatch between the proposal and target distributions.

While annealing-based methods like the Jarzynski Equality (JE) and Annealed Importance Sampling (AIS) are empirically successful, their theoretical complexity guarantees have remained largely unexplored, particularly in non-asymptotic settings and without assuming strong isoperimetric conditions (e.g., log-concavity) on the target distribution.

Methodology

The authors develop a non-asymptotic analysis framework for estimating $Z$ by leveraging tools from stochastic analysis and optimal transport.

1. Theoretical Framework:

   • Jarzynski Equality (JE): The paper analyzes JE by considering annealed Langevin diffusion (ALD) as a continuous-time process. It establishes a connection between the work functional $W$ and the free energy difference $\Delta F$.
   • Girsanov's Theorem & Optimal Transport: A key technical innovation is the use of Girsanov's theorem to relate the forward path measure (sampling trajectory) and the backward path measure. The analysis avoids explicit dependence on isoperimetric inequalities (like Poincaré or Log-Sobolev) by instead utilizing the action of the curve of probability measures. The action $A$ is defined as the integral of the squared metric derivative in the Wasserstein-2 ($W_2$) distance along the interpolation curve.
   • Annealed Importance Sampling (AIS): The continuous dynamics are discretized to analyze AIS. The authors define a reference path measure and bound the Kullback-Leibler (KL) divergence between the actual sampling path and this reference. This allows for the derivation of oracle complexity bounds for the discrete algorithm.

2. Algorithmic Proposals:

   • Geometric Interpolation Limitations: The paper identifies that the widely used geometric interpolation (linearly interpolating potentials) can lead to an exponentially large action for certain multimodal distributions (e.g., Gaussian mixtures with well-separated modes), causing "mass teleportation" issues where the sampler fails to transition between modes efficiently.
   • Reverse Diffusion Samplers (RDS): To overcome the limitations of geometric interpolation, the authors propose a new class of algorithms based on Reverse Diffusion Samplers. These methods utilize the time-reversal of the Ornstein-Uhlenbeck (OU) process. The OU process naturally transforms the target distribution into a standard Gaussian with a well-behaved action, even for non-log-concave targets. The proposed framework estimates the normalizing constant by simulating this reverse process using score estimators.

Key Contributions

The paper makes the following specific contributions:

1. Novel Complexity Analysis Strategy: It introduces a strategy for analyzing the complexity of normalizing constant estimation applicable to general target distributions that may not satisfy isoperimetric conditions. The analysis relies on the action of the interpolation curve rather than log-concavity.
2. Non-Asymptotic Bounds for JE and AIS:
   • For JE, it proves that running annealed Langevin dynamics for a time $T \propto A/\varepsilon^2$ suffices to estimate $Z$ within $\varepsilon$ relative error with high probability, where $A$ is the action of the curve.
   • For AIS, it establishes the first non-asymptotic oracle complexity bound for normalizing constant estimation without assuming a log-concave target. The bound is derived as $\tilde{O}\left( \frac{d^{4/3}}{\varepsilon^2} \vee \frac{m\beta A^{1/2}}{\varepsilon^2} \vee \frac{d\beta^2 A^2}{\varepsilon^4} \right)$, where $d$ is dimension, $\beta$ is smoothness, $m$ is the second moment, and $A$ is the action.
3. Lower Bound on Geometric Interpolation: The authors prove an exponential lower bound on the action of the geometric interpolation curve for Gaussian mixture targets, providing a quantitative explanation for the inefficiency of standard AIS in multimodal settings.
4. RDS Framework: It proposes a framework for analyzing the oracle complexity of RDS-based algorithms and demonstrates that the OU process offers a curve with significantly better action properties ($O(d\beta + m^2)$) compared to geometric interpolation for multimodal targets.

Results

• Theoretical: The derived complexity bounds for JE and AIS are shown to depend on the action $A$. The analysis confirms that the complexity of estimating the normalizing constant is quantitatively connected to the complexity of sampling, extending previous discrete results to continuous settings without log-concavity.
• Empirical: Experiments on two multimodal distributions in $\mathbb{R}^2$ (a modified Müller Brown distribution and a 4-component Gaussian mixture) demonstrate the superiority of RDS-based methods over TI and AIS.
  • TI and AIS: Both methods produced seriously biased estimates of the normalizing constant and failed to cover all modes of the target distribution (getting stuck in the initial mode).
  • RDS Methods (RDMC, RSDMC, ZODMC, SNDMC): All four RDS-based methods provided accurate estimates of the normalizing constant (relative error close to 1) and generated high-quality samples that covered all modes, as evidenced by low Maximum Mean Discrepancy (MMD) and Wasserstein-2 ($W_2$) distances.

Significance

The paper claims to take a "first step" toward a rigorous non-asymptotic analysis of annealing-based methods for normalizing constant estimation. Its significance lies in:

• Relaxing Assumptions: Moving beyond the restrictive assumption of log-concavity, which has limited the theoretical understanding of these methods in practical, multimodal scenarios.
• Quantifying Inefficiency: Providing a theoretical explanation (via the action lower bound) for why standard geometric annealing fails in multimodal settings, a phenomenon previously observed empirically.
• Bridging Sampling and Estimation: Establishing a quantitative link between the complexity of sampling and the complexity of estimating the normalizing constant in continuous, non-log-concave settings.
• Proposing a Solution: Introducing and analyzing Reverse Diffusion Samplers as a viable alternative that leverages the favorable properties of the OU process to handle multimodality effectively.

The authors note that while their upper bounds are derived, their tightness is unknown, and the practical interpretability of the action metric requires further study. They suggest future work could extend these techniques to other samplers (e.g., Hamiltonian dynamics) and discrete distributions.