Annealed EWFM: Efficient CNF Training

• Annealed EWFM is a training paradigm that extends energy-weighted flow matching for generating samples from Boltzmann distributions without needing target samples.
• It employs an iterative temperature annealing strategy to progressively adapt CNFs to complex energy landscapes, significantly reducing energy evaluations.
• The method stabilizes self-normalized importance sampling gradients and achieves high sample efficiency and robustness compared to other energy-based generative approaches.

Annealed Energy-Weighted Flow Matching (aEWFM) is a training paradigm for continuous normalizing flows (CNFs) targeting Boltzmann sampling in high-dimensional, multimodal energy landscapes. By incorporating temperature annealing into the core Energy-Weighted Flow Matching (EWFM) framework, aEWFM enables efficient, sample-free learning of generative models solely from energy evaluations, with substantial improvements in sample efficiency and robustness over previous energy-only methods. The method is applicable across physical, chemical, and combinatorial domains, where sampling from complex unnormalized distributions is fundamental.

1. Formal Definition and Objective

The aEWFM algorithm extends EWFM by optimizing CNFs to produce samples from Boltzmann distributions

$\mu_\text{target}(x) \propto \exp(-E(x)/T)$

without access to samples from $\mu_\text{target}$. Standard conditional flow matching methodologies require direct target samples, but EWFM circumvents this by drawing samples from a proposal $\mu_\text{prop}$ and reweighting via importance sampling:

$w(x) = \frac{\exp(-E(x)/T)}{\mu_\text{prop}(x)}$

The key loss function in EWFM is

$$\mathcal{L}_{\text{EWFM}}(\theta; \mu_{\text{prop}}) = \mathbb{E}_{t, X_t, X_1}\left[ \frac{w(X_1)}{\mathbb{E}_{X'_1 \sim \mu_{\text{prop}}}[w(X'_1)] } \| u_t^\theta(X_t) - u_t(X_t|X_1) \|^2 \right]$$

where $u_t^\theta$ is the CNF velocity field parameterized by $\theta$, and $u_t(\cdot|\cdot)$ is the ideal vector field along the stochastic path ending at $X_1$.

In the annealed extension (aEWFM), a temperature schedule $T_0 > T_1 > \ldots > T_K = T$ is set, interpolating from a high $T_0$ (flattened energy landscape) towards the target temperature $T$. At each annealing stage, the above objective is computed for the temperature $T_k$.

2. Annealing Mechanism and Training Dynamics

The annealing strategy in aEWFM exploits the fact that the Boltzmann distribution at high temperature $T_0$ is smoother and has greater overlap with trivial proposals (e.g., standard Gaussian), minimizing variance in the importance weights. Training commences at $T_0$ and steadily lowers the temperature in geometric or custom schedules. At each temperature level, the current trained CNF $q_\theta$ serves both as a sampler and as the incremental proposal for the next stage. This bootstrapped cooling process enables the CNF to progressively adapt to sharper modes and barriers in the energy landscape.

This annealing not only regularizes the objective but also stabilizes the self-normalized importance sampling (SNIS) gradients encountered during training, improving convergence and robustness against gradient bias that otherwise arises when the proposal and target are poorly matched. In the full iterative annealed EWFM protocol, the model leverages amortized sample buffers—reusing samples for multiple epochs—which allows estimates with up to three orders of magnitude fewer energy calculations than previous state-of-the-art energy-based learning methods.

3. Comparison with Related Energy-Based Methods

aEWFM fundamentally differs from other energy-based generative methodologies:

Method	Target Sampling	Energy Evaluation Scaling	CNF Expressivity
F-AIS-Bootstrap (FAB)	Required	High	Limited (single step)
Iterated Denoising Energy Matching (iDEM)	Not required	High	Limited
EWFM	Not required	Low (via SNIS and reweighting)	High (CNF-based)
aEWFM	Not required	Ultra-low (annealed schedule, sample reuse)	High (CNF-based, annealing schedule)

aEWFM achieves comparable NLL and 2-Wasserstein sample quality to the best energy-only approaches (e.g., FAB, iDEM) on benchmark physical systems (such as 55-particle Lennard-Jones clusters) while requiring up to three orders of magnitude fewer energy evaluations; though, the computational cost for CNF density evaluations (Jacobian/Hutchinson trace) remains a bottleneck in practice. The algorithm consistently outperforms naive iterative methods in both efficiency and ability to sample multi-modal distributions.

4. Explicit Algorithmic Workflow

The aEWFM algorithm proceeds as follows:

1. Initialize Proposal at High Temperature: Set $T_0$ such that $\mu_{T_0}(x) \propto \exp(-E(x)/T_0)$ is nearly flat.
2. Iterative Annealing Loop:
   • For $k = 0$ to $K$ (annealing steps):
     • Sample $x_1 \sim q_\theta^{(k)}(x)$, using the previous model as the proposal.
     • Compute energy-based importance weights for temperature $T_k$.
     • Optimize the EWFM objective $\mathcal{L}_{\text{EWFM}}(\theta; q_\theta^{(k)})$ for $T_k$, using SNIS gradient estimation.
     • Record sample buffer for possible reuse in future stages.
     • Update proposal to current $q_\theta^{(k+1)}$.
3. Final Target Sampling: After convergence at target temperature $T_K$, $q_\theta^{(K)}$ serves as the generative model for $\mu_{\text{target}}$.

For practical implementations, density calculations in CNF can leverage unbiased estimators (e.g., Hutchinson trace) to control computational cost.

5. Demonstrated Applications

aEWFM is especially suited for physical and chemical systems with expensive or inaccessible target samples:

• Lennard-Jones Clusters: Sampling low-NLL high-fidelity configurations in 55-particle clusters with dramatically fewer energy evaluations.
• Protein Folding and Molecular Dynamics: Enables probabilistic exploration in high-dimensional rugged energy landscapes where conventional MCMC gets trapped.
• Boltzmann Sampling for General Statistical Physics: Rapid estimation of equilibrium statistics in models defined only by energy (no explicit target samples).

6. Limitations and Prospective Improvements

The primary limitations of aEWFM stem from:

• The need for proposal density evaluations via CNFs (computational overhead, potential bias from estimator inaccuracies).
• Observable bias in self-normalized gradient estimates for mid-dimensional systems where proposal and target are far apart.
• Increased wall-clock training times due to the annealing schedule, despite energy evaluation efficiency.

Potential improvements include:

• Adoption of mixture-model proposals or adaptive importance sampling strategies to mitigate density evaluation costs.
• Development of stabilized or alternative gradient estimators for self-normalized objectives in regimes with high importance-weight variance.
• Hybrid strategies combining a small set of true target data (if available) for warmer starts in especially challenging landscapes.

7. Theoretical Significance and Implications

By connecting variational flow optimization directly to energy-only objectives and leveraging annealed schedules for stability, aEWFM establishes a general route for learning expressive generative models in regimes previously inaccessible to CNF architectures. The theoretical identification of direct equivalence between EWFM objectives at each annealing level and conditional flow matching—modulo importance reweighting—ensures strong fidelity to the target Boltzmann measure even without sample access. The framework also clarifies relationships between high-temperature and low-temperature training regimes, underscoring the role of annealing in improving SNIS-based gradient estimation and ultimately driving sample quality and diversity.

In summary, aEWFM advances the state of the art in sample-efficient, expressive energy-based modeling for high-dimensional Boltzmann sampling, with explicit annealing-based mitigations for optimization and estimation barriers endemic to such domains (Dern et al., 3 Sep 2025).