---
title: Annealed Importance Sampling Bootstrap
url: https://www.emergentmind.com/topics/annealed-importance-sampling-bootstrap
type: topic
---

# Annealed Importance Sampling Bootstrap

Annealed Importance Sampling Bootstrap (FAB) refers to a family of training algorithms for normalizing flows that embed annealed importance sampling (AIS) within the flow optimization loop, specifically targeting mass-covering objectives to reduce the variance of importance weights. The central innovation is the bootstrapping interaction between the flow and the AIS sampler: AIS leverages the current flow as proposal, providing weighted samples that sharpen low-variance $\alpha$-divergence gradient estimates, which in turn iteratively refine the flow. The methodology has established state-of-the-art performance for learning unnormalized, multimodal, and high-dimensional densities from only black-box target density evaluations, with significant applications in computational physics and high-energy simulation [2111.11510, 2208.01893, 2411.16234].

## 1. Principle of FAB: Annealed Importance Sampling–Flow Bootstrap

FAB augments normalizing flow training with annealed importance sampling. A ladder of intermediate densities $(\pi_0,\ldots,\pi_M)$ interpolates between the flow surrogate $q_\theta(x)$ and a final “optimal” mass-covering target (e.g., $p(x)^2/q_\theta(x)$ for $\alpha=2$ in the $\alpha$-divergence). The intermediates are

\[
\pi_i(x) \propto q_\theta(x)^{1-2\beta_i} p(x)^{2\beta_i}, \quad \beta_0 = 0 < \beta_1 < \dots < \beta_M = 1
\]

Particles are initialized by drawing from $q_\theta$, then propagated through this sequence via Markov transition kernels $K_i$ (often HMC), each stage accompanied by incremental weight update

\[
w_i(x, y) = \frac{\pi_i(y)}{\pi_{i-1}(y)}
\]

The full AIS weight for each sample chain is

\[
w^{(n)} = \prod_{i=1}^M w_i(x^{i-1,(n)}, x^{i,(n)})
\]

Normalized weights then define an empirical (importance-weighted) estimator for any function of interest. The bootstrapping occurs because as $q_\theta$ improves, the efficiency and stability of AIS—and thus the gradient estimation—simultaneously improve [2208.01893, 2111.11510, 2411.16234].

## 2. Mass-Covering $\alpha$-Divergence Objective and Its Surrogate

FAB directly targets the minimization of the $\alpha$-divergence for $\alpha=2$,

\[
D_2(p \,\|\, q_\theta) = \frac{1}{2} \int \frac{p(x)^2}{q_\theta(x)} dx
\]

Minimization of $D_2$ is equivalent to minimizing the variance of the importance weights $w(x) = p(x)/q_\theta(x)$. 

Instead of drawing samples only from $q_\theta$, FAB samples from the AIS kernel targeting $p(x)^2/q_\theta(x)$ and uses these (together with their normalized AIS weights) to form a low-variance, mass-covering surrogate loss:

\[
\mathcal{S}(\theta) = -\sum_{i=1}^N \bar w^{(i)} \log q_\theta(x^{(i)})
\]
where $\bar w^{(i)} = w^{(i)} / \sum_j w^{(j)}$. In practice, gradients are stopped with respect to the AIS weights and samples, and only backpropagated through $q_\theta$ [2208.01893, 2111.11510, 2411.16234].

## 3. Algorithmic Procedure and Practical Implementation

A high-level FAB iteration with optional sample replay buffer is as follows:

```python
# Pseudocode: One FAB Iteration
1. Draw {x⁰ⁿ}ⁿ₌₁…ᴺ ∼ q_θ
2. For n=1...N: Run AIS chain:
     for t=1...M:
         xᵗⁿ ∼ Kᵗ(x^{t-1,n},·)  # e.g., HMC targeting πᵗ
         wᵗⁿ = πᵗ(xᵗⁿ)/π^{t-1}(xᵗⁿ)
     w_AISⁿ = ∏_{t=1}^M wᵗⁿ
   Store (x^{M,n}, w_AISⁿ) in buffer 𝓑
3. Sample L minibatches {(\bar xᵢ, \bar wᵢ)} from 𝓑 ∝ w_AIS
4. For each, update θ ← θ - η ∇_θ [ -∑ᵢ \bar wᵢ log q_θ(\bar xᵢ) ]
```

The replay buffer allows for efficient gradient reuse and substantially reduces the number of required expensive target density evaluations. Each kernel $K^t$ is typically a short HMC sequence, adaptively tuned (target acceptance $\approx$ 0.65), followed by optional resampling based on weight degeneracy.

In applications, small numbers of intermediate distributions ($M=2\text{--}4$) suffice even at moderate-high dimension, and flow architectures typically use coupling/spline layers with moderate capacity (e.g., 14 layers and 400 hidden units per conditioner in 8D) [2411.16234, 2208.01893].

## 4. Empirical Validation and Comparative Performance

FAB demonstrates superior performance to maximum-likelihood flow training (fKLD), reverse-KL flow training (rKLD), and grid-based VEGAS+ sampling. Typical evaluation metrics include:

- Forward KL $KL(p\|q_\theta)$
- Importance sampling efficiency $\epsilon$ (normalized ESS)
- Integral estimates via weighted samples

For an 8D matrix element sampling task [2411.16234]:

| Method          | $\mathbb{E}_p[\log(p/q)] \downarrow$ | $\epsilon$ (%) $\uparrow$ | \# Target evals  |
|-----------------|-------------------|------------------|---------------|
| rKLD            | 7.74 (0.02)       | 56.5             | $1\times10^8$ |
| FAB (w/o buf)   | 7.79 (0.03)       | 84.3             | $1\times10^7$ |
| FAB (w/ buf)    | 7.747 (0.002)     | 90.6             | $1\times10^7$ |

On lower-dimensional multimodal problems, only FAB and ML-trained flows cover all modes, achieving high ESS (e.g., $>70$\% on 25-mode 2D Gaussians, $>90$\% for 30D Boltzmann distributions), where alternative methods fail with near-degenerate ESS and significant mode collapse [2111.11510, 2208.01893, 2411.16234].

## 5. AIS Bootstrap Variants and Extensions

FAB admits flexible choices for the intermediate density ladder, annealing (linear or geometric in $\beta$), and MCMC kernel (HMC, Metropolis). Alternative $\alpha$-divergences ($\alpha>1$) can be substituted, though $\alpha=2$ yields optimal mass covering empirically.

Replay-buffer FAB variants further reduce the number of required target evaluations by reusing AIS-weighted samples through priority sampling and log-density reweighting. There are also suggestions for combining FAB with other flow architectures such as CRAFT and SNF, and adapting the AIS procedure to differentiable variants that further reduce gradient variance [2208.01893, 2411.16234].

## 6. Limitations and Domain-Specific Constraints

FAB requires access to $\nabla_x \log p(x)$ to enable HMC in the annealing steps. This is now feasible in differentiable matrix element modeling frameworks (e.g., MadJAX, ComPWA), but constrains applicability in non-differentiable domains [2411.16234]. HMC and flow capacity hyperparameters must be empirically tuned to balance mixing, computational cost, and acceptance rates. Numerical stability issues may arise if $q_\theta(x)$ is too small in high-density regions of $p(x)$; defensive regularization or early NaN filtering may be required. 

Compute costs are dominated by HMC and target evaluations, making replay buffering essential in high-dimensional or expensive physical simulation tasks.

## 7. Theoretical Significance and Broader Context

FAB demonstrates that integrating AIS into the training loop of normalizing flows with mass-covering divergence minimization results in both enhanced robustness to mode collapse (discovering underrepresented or missing modes) and marked reductions in the variance of importance weights. Theoretical underpinnings guarantee that, under mixing, the approach yields consistent gradient estimates for loss minimization; practical surrogate biases can be made arbitrarily small with increasing minibatch size [2111.11510]. 

Within computational physics, especially for tasks such as particle-physics matrix element simulation, FAB provides a scalable, sampler-independent framework for high-fidelity surrogate learning without reliance on large pre-computed MC datasets. The method’s principled combination of flow-based modeling, importance sampling, and replay reweighting outperforms conventional approaches in both sample efficiency and distributional accuracy [2411.16234, 2208.01893, 2111.11510].

Source: https://www.emergentmind.com/topics/annealed-importance-sampling-bootstrap