Technical Summary: FUND – Density Flow for Sampling Unnormalised Distributions

Problem Statement

Sampling from unnormalized probability distributions, specifically Boltzmann distributions of the form $p(x) \propto e^{-E(x)}$, is fundamental to modeling complex physical systems in statistical physics, quantum mechanics, and computational chemistry. Traditional methods like Markov Chain Monte Carlo (MCMC) and Molecular Dynamics (MD) suffer from critical slowing down, high autocorrelation, and poor mode mixing, particularly when navigating high-dimensional energy landscapes with metastable states separated by high barriers.

While deep generative models, such as Normalizing Flows (NFs), offer efficient sampling, they face significant limitations in this domain:

1. Data Dependency: Standard NF training (e.g., via Maximum Likelihood Estimation/Forward KL) requires large datasets of true samples from the target distribution, which are often unavailable or expensive to generate via MCMC/MD.
2. Mode Collapse: Training without samples using the Reverse KL divergence ($KL(q_\theta \| p)$) biases the model toward low-energy regions, leading to mode collapse where the model fails to cover the full support of the target distribution.
3. Existing Alternatives: Recent methods like Flow Annealed Importance Sampling Bootstrap (FAB) and Iterated Denoising Energy Matching (iDEM) attempt to bypass true sample requirements but introduce substantial computational overhead (e.g., repeated Annealed Importance Sampling or solving reverse Stochastic Differential Equations) or fail to provide exact likelihoods without additional approximation stages.

Methodology: FUND (Density Flow for Unnormalised Distributions)

The authors propose FUND, a framework that learns a tractable sampler $q(x)$ for unnormalized distributions without requiring any true samples from the target distribution. The core innovation lies in shifting the learning paradigm from sample-trajectory matching (as in Flow Matching) to distribution-trajectory matching.

Core Framework

1. Annealed Distribution Trajectory: Instead of modeling the path of individual samples, FUND models a continuous trajectory of probability densities $\{p_t(x)\}_{t \in (0,1]}$ that transforms a simple unimodal prior $p_0(x)$ into the complex target $p_1(x) = p(x)$. The intermediate distributions are defined as:
   $$p_t(x) \propto p(x)^t$$
   As $t$ increases from 0 to 1, the distribution gradually sharpens, and modes separate. At $t \to 0$, the distribution is smooth and unimodal; at $t=1$, it recovers the full complexity of the target.

2. Iterative Learning with Importance Weighting: The model learns a sequence of intermediate distributions $q_t$ corresponding to $p_t$.

   • Initialization: Training begins at a small $t_1$ where $p_{t_1}$ is unimodal. Since no samples from $p_{t_1}$ exist, the model is trained using a Reverse KL (RKL) objective ($L_{RKL}$), which relies only on energy evaluations.
   • Progression: For subsequent steps $t > t_1$, the model generates samples from the previously learned $q_{t'}$ (where $t' < t$). These samples are re-weighted using importance weights $w(x) = p_t(x) / q_{t'}(x)$ to supervise the learning of $q_t$.
   • Loss Function: The training objective combines a supervised term (Forward KL or Score Matching) with the unsupervised RKL term:
     $$L_{final}(\theta) = \lambda_1 L(\theta) + \lambda_2 L_{RKL}(\theta)$$
     • $L(\theta)$: Encourages mode coverage (exploration) using weighted samples.
     • $L_{RKL}(\theta)$: Encourages concentration on dominant modes (exploitation) and corrects density discrepancies.
   • Hyperparameter Annealing: The weights $\lambda_1$ and $\lambda_2$ are annealed. Initially, $\lambda_1 \approx 0$ (relying on RKL). As training progresses, $\lambda_1$ increases to leverage the re-weighted samples for mode coverage, while $\lambda_2$ increases later to refine mode separation.

3. Architecture and Implementation:

   • Single-Block Architecture: Unlike incremental approaches that add coupling blocks at each step, FUND uses a single Normalizing Flow model refined across all time steps. This improves memory efficiency and prevents catastrophic forgetting.
   • Replay Buffer: A persistent buffer stores samples and log-likelihoods from previous intermediate distributions to compute importance weights for subsequent steps, ensuring stable learning and mode preservation.
   • Discretization: The continuous time $t$ is discretized into $N$ steps. The number of steps and the initial $t_1$ are tuned based on the complexity and multimodality of the target.

Key Contributions

1. Sample-Free Learning: FUND enables the training of Normalizing Flows for unnormalized distributions without access to true target samples, eliminating the need for expensive MCMC/MD pre-sampling.
2. Distribution-Trajectory Matching: The paper introduces a novel formulation that models the evolution of the density itself rather than individual sample trajectories, allowing for direct and efficient density evaluation during inference.
3. Balanced Objective: By combining Forward KL/Score Matching (for coverage) with Reverse KL (for concentration) in a sequential annealing schedule, FUND mitigates the trade-off between mode collapse and mode covering, ensuring comprehensive coverage of complex, high-dimensional landscapes.
4. Scalability: The method demonstrates scalability to high-dimensional problems (e.g., 144-dimensional scalar $\phi^4$ theory) where traditional baselines degrade significantly.

Experimental Results

The authors evaluate FUND on three benchmark families:

1. MOG-40: A 2D mixture of 40 Gaussians with widely separated modes.
2. Many-Well (MW-8, MW-16): Synthetic distributions with exponentially growing numbers of modes ($2^{d/2}$).
3. Scalar $\phi^4$ Theory: A lattice field theory model on $8\times8$, $10\times10$, and $12\times12$ grids (up to 144 dimensions).

Performance Metrics:

• Negative Log-Likelihood (NLL): Measures fit to the target. FUND variants consistently achieve the lowest NLL across all benchmarks, outperforming FAB, iDEM, and standard RKL/FKL baselines.
• Reverse NLL (RNLL) & Effective Sample Size (ESS): These metrics assess mode coverage. FUND achieves high ESS and low RNLL, indicating it captures multiple modes effectively. In contrast, RKL baselines show low NLL but poor ESS (mode collapse), while FAB and iDEM show degraded performance in high dimensions.
• Computational Cost:
  • FUND(FKL): Offers a favorable cost-quality tradeoff, with training times comparable to FAB but significantly faster inference (single forward pass vs. AIS or MCMC).
  • FUND(SCORE): Achieves the highest sampling accuracy and density estimation but incurs higher training costs due to score-matching computations.
  • Baselines: MCMC methods have high inference costs; iDEM has high inference costs due to reverse SDE simulation; FAB has high training costs due to repeated AIS.

Significance and Claims

The paper claims that FUND establishes a new paradigm for sample-free learning of complex physical distributions. By progressively shaping the target distribution through intermediate annealed states, the method overcomes the "mode collapse" inherent in energy-based training and the "data dependency" of likelihood-based training.

The authors emphasize that FUND provides a robust, scalable alternative to MCMC for high-dimensional Boltzmann distributions, particularly in lattice field theory and statistical physics where critical slowing down and topological trapping are persistent challenges. The framework enables amortized sampling (fast parallel generation via a single forward pass) while maintaining high fidelity to the target distribution's multimodal structure, offering a practical pathway for applications in Bayesian inference and gauge-theoretic simulations where traditional sampling methods struggle.