Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows

• Boltzmann Generators use learned proposals and importance sampling, but normalizing flows can have limited expressiveness and overlap, while continuous-flow likelihoods can be costly and discretization-sensitive.Source: §1 · p. 1
• NF2M addresses the gap between expressive stochastic flow maps and Boltzmann Generator correction, which requires likelihoods that implicit models typically lack.Source: §1 · p. 2
• NF2M models clean configurations conditioned on a noisy state with a conditional normalizing flow, yielding a tractable density for few-step generation along a diffusion path.Source: §1 · p. 2
• At each denoising step, NF2M proposes clean states and reweights them against the tractable conditional target, producing an estimator exact as the particle count grows.The estimator is a Rao-Blackwellized mixture that is sampled to denoise the current state.Source: §1 · p. 2
• The paper also introduces conditional normalizing-flow architectures for molecule generation, adapting image and video flow designs.Source: §1 · p. 2
• NF2M shows competitive or superior performance against normalizing-flow and flow-map baselines across molecular systems of varying sizes.Source: §1 · p. 2

• The setup seeks equilibrium samples for an N-atom molecule represented by 3D coordinates, assuming the energy can be evaluated although the Boltzmann normalizer is generally intractable.Source: §2 · p. 2
• Boltzmann Generators propose configurations with a learned model and correct them using self-normalized importance sampling, which is consistent with sufficient proposal support but depends on proposal-target overlap and density evaluation.Source: §2 · pp. 2–3
• Normalizing flows provide sampling and exact density evaluation through invertible transformations, but their architectural constraints can limit distributional expressivity.Source: §2 · p. 3
• Diffusion paths progressively noise clean samples and generate through denoising transitions, with the terminal state approximated by a tractable reference distribution at sufficiently large noise.The background specifies a linear Gaussian noising kernel and a variance-exploding schedule with α_t = 1 and σ_t = t.Source: §2 · p. 3
• A reverse diffusion transition can be formed by marginalizing over possible clean states, while the usual learned-mean Gaussian approximation can be poor at coarse finite steps.Source: §2 · p. 3

3 NORMALIZING FLOW FLOW MAPS (NF2M)

• NF2M corrects each denoising transition along the diffusion path, so individually asymptotically unbiased transitions compose without a final correction.A conditional normalizing flow proposes clean configurations from the current noisy state, and importance sampling corrects them toward the true posterior.Source: §3 · pp. 3–4

3.1 BOLTZMANN SAMPLING ALONG THE DIFFUSION PATH

• The Rao-Blackwellized transition mixes analytically over the endpoint-conditioned Gaussian, attenuating propagated posterior errors by |as,t|.It retains the full conditional law rather than drawing and reweighting one noisy sample per proposal.Source: §3.1 · p. 4

3.2 TRAINING

• A flexible conditional normalizing flow represents multimodal denoising posteriors while retaining the exact likelihood needed for training and inference.The training data are biased samples; their bias is corrected at inference.Source: §3.2 · p. 5

3.3 IMPROVING THE INFERENCE

• NF2M-MIS pools proposals across trajectories, letting each trajectory use CM candidates for the same total of NM flow samples as NF2M-IS.The method requires additional cross-likelihood evaluations for importance weights.Source: §3.3 · p. 6

3.4 THEORETICAL GUARANTEES

• Per-step SNIS bias and MSE decay with particle count M, with their bounds reflecting proposal fit and bridge sensitivity.The bounds assume sub-Gaussian tails for the proposal and its square-weighted tilt.Source: §3.4 · p. 7
• Across K denoising steps, bias accumulates additively and root-MSE is controlled through a triangle inequality.Under the VE schedule, small relative steps near high noise damp posterior-approximation errors; at low noise, MAExL rejuvenation addresses the sharply peaked regime.Source: §3.4 · p. 7

BGs with tractable correction weights

• Sequential Boltzmann Generators anneal a fixed TarFlow proposal, whereas NF2M uses progressive denoising and trains its flow conditionally.Reported ablations find that SMC does not consistently improve over SNIS across systems and metrics.Source: §4 · pp. 7–8
• ArBG uses coordinate-bin probabilities and partial-structure resampling, while NF2M corrects proposals over complete molecular configurations at each noise level.The passage contrasts this with ArBG's binning artifacts and energy evaluations on artificially capped partial structures.Source: §4 · p. 8

BGs with approximate correction

• Several distillation-based BGs use approximate learned densities in importance weights, while NF2M evaluates its conditional proposal densities exactly.The passage states that the approximation error in those learned densities persists regardless of proposal count.Source: §4 · p. 8

Stochastic flow maps and intermediate correction

• Implicit conditional generators in stochastic flow-map methods do not directly provide the proposal densities needed for Boltzmann importance weights; NF2M supplies them with a conditional normalizing flow.Source: §4 · p. 8

5 EXPERIMENTS

• The evaluation reuses the five peptide systems’ MD trajectories, splits, force fields, and protocol from Tan et al. (2025), with energy and dihedral Wasserstein distances as metrics.Table 1 reports these distances; Table 2 reports training and inference times, with systems ordered by increasing size.Source: §5 · p. 8, §5.1 · p. 9, Table 1 · p. 9

5.1 MAIN RESULTS

• NF2M-IS/MIS lead every system except ALA-6 on E-W2, and NF2M has the lowest T-W2 on the largest systems, ALA-6 and chignolin.SNIS has the lowest T-W2 on ALA-2/ALA-3, but its T-W2 rises on larger systems; FALCON has the lowest E-W2 on ALA-6.Source: §5.1 · p. 9, Table 1 · p. 9
• In-loop MAExL improves energy accuracy and stability with little change in torsion distances, while increasing IS from 512 to 1024 steps brings little further benefit.Removing post-hoc MALA changes NF2M’s results relatively little but substantially worsens energy accuracy for several baselines.Source: §5.1 · p. 9
• NF2M-IS and NF2M-MIS perform similarly, with MIS using a smaller candidate budget and reducing inference time most clearly on chignolin.At matched candidate budgets, the methods achieve comparable inner ESS, though both show substantial weight concentration at low noise, especially on chignolin.Source: §5.1 · p. 10
• NF2M inference time is broadly comparable to SNIS, F2D2, and SCALLOP, and much shorter than FALCON and FALCON-A.FALCON inference is roughly 11× slower than SNIS on ALA-2 and 48–129× slower on ALA-3 through ALA-6; on chignolin, its reduced proposal pool is still roughly 60× slower than SNIS or NF2M-IS.Source: §5.1 · pp. 9–10, Table 2 · p. 9

6 CONCLUSION

• NF2M parameterizes a stochastic flow map with a conditional normalizing flow and corrects each denoising transition via tractable SNIS.The conclusion describes consistent samples with theoretical error control and strong empirical performance across peptide systems, without approximate likelihood.Source: §6 · p. 10
• NF2M is currently trained per system, with transferability across peptide sequences and lengths proposed as a next step.The conclusion suggests learning a single conditional flow that generalizes across sequences and lengths.Source: §6 · p. 10
• The conclusion also identifies reverse-KL training as a potential direction for refining the model.Source: §6 · p. 10

• The bridge distribution for Xs conditioned on Xt and X0 is Gaussian, with mean αs,t x0 + bs,t xt and covariance Σs,t.Source: §A.1 · p. 17
• The schedule stops at strictly positive time because the bridge sensitivity is finite for s > 0 and diverges as s approaches zero, where the bridge degenerates to a point mass.Source: §A.1 · p. 17
• The SNIS construction draws candidates from q0|t and forms an approximation of p0|t, then induces an approximation of ps|t.Source: §A.1 · pp. 17–18

• The bridge analysis establishes a contraction property and bounds the oscillation of its operator on a ball, which are combined to obtain the per-step bound.Source: §A.1.1 · pp. 18–20
• The analysis sharpens SNIS bias and MSE bounds by localizing the test function to a set where the proposal places most of its mass.The stated lemma assumes the target measure is absolutely continuous with respect to the proposal and uses their chi-squared divergence.Source: §A.1.1 · p. 18
• Under the stated sub-Gaussian tail assumption, the per-step SNIS bound has a constant branch of order O(ρt(xt)/M), uniformly in γt,s.Source: §A.1.1 · p. 20

A.1.2 ADDING AN MCMC REJUVENATION STEP

• The rejuvenated estimator applies an invariant Markov kernel to the bridge function before self-normalized importance sampling.The kernel leaves the conditional target p0|t(·|xt) invariant, preserving the target expectation.Source: §A.1.2 · pp. 20–22
• The bounds assume exact invariance at a fixed step size; step sizes chosen independently of the chain preserve invariance conditionally, but state-dependent adaptation is not quantified.The production sampler adapts step size online from acceptance history, and that finite-L chain need not be exactly invariant.Source: §A.1.2 · p. 21
• The rejuvenation move improves the bias bound by the mixing factor r_t^L while leaving the MSE of the same order as without rejuvenation.The result applies to the MCMC-integrated estimator; the sampled estimator has an additional conditional variance term.Source: §A.1.2 · p. 23
• MAExL contracts the effective x0-diameter seen by the bridge, while leaving the bridge sensitivity γ_t,s unaltered.The analysis uses invariance, not an exact draw from the conditional target; finite-L dependence on proposal initialization enters through a generally nonzero diameter term.Source: §A.1.2 · p. 24

A.1.3 FROM PER-STEP BOUNDS TO A PATH-ERROR BOUND

• The path-error analysis writes the difference between chained approximate and exact kernels as a sum of per-step kernel errors propagated through the chain.The decomposition is stated for independent random approximating kernels and a given initial law.Source: §A.1.3 · p. 25
• The chained SNIS bounds assume an exact initial law at tK and apply per-step controls along a schedule ending at τ>0.The setup also assumes sub-Gaussian tails at each scheduled state and uses uM = √3 log M.Source: §A.1.3 · p. 26
• With MCMC rejuvenation, the chain uses L-step rejuvenated SNIS approximations at each scheduled transition, with independence covering proposal and MCMC randomness.The corollary assumes the required MAExL constants and tail conditions for the L-step pushforwards.Source: §A.1.3 · p. 26

Define the chained approximation

• The chained bias bound contracts geometrically in L, while the MSE bound remains of the same order as without MCMC rejuvenation.The per-step bias and MSE bounds are applied at each scheduled transition to obtain the path-level result.Source: Define the chained approximation · p. 27
• The chained kernel is random in the current noisy state, with its internal randomness including proposal draws and MCMC randomness.The path argument requires independence across transitions covering both sources of randomness.Source: Define the chained approximation · p. 27
• The guarantees target pτ rather than p0, and starting from νT instead of the exact terminal law adds a total-variation discrepancy independent of M and L.Under the stated VE schedule, the exact terminal marginal is p0 convolved with Gaussian noise.Source: Define the chained approximation · p. 27

• NF2M-MIS pools candidates from selected posterior components so each trajectory uses CM candidates with the same NM flow samples as NF2M-IS.Each component set retains the trajectory’s own posterior, while other slots are selected by posterior proximity.Source: §A.2 · p. 28
• Shared batches are reused across trajectories, but repeated use of one source within a mixture draws a fresh independent batch for each occurrence.This preserves independent candidates within each trajectory’s mixture.Source: §A.2 · p. 29
• The added mixture expressiveness requires cross-likelihood evaluations, but shared flow samples and autoregressive density evaluation keep this overhead comparatively negligible.The additional evaluations are at most NMC(C −1) per level; when C = 1, the mixture reduces to NF2M-IS.Source: §A.2 · pp. 29–30
• At matched candidate budgets, IS and MIS have comparable inner ESS, but both show substantial importance-weight concentration at low noise, especially on chignolin.The authors hypothesize that near-deterministic posteriors make the flow’s affine-coupling scale network difficult to learn; MAExL steps are most accurate in this regime.Source: Figure 4 · p. 31, §A.2 · p. 32
• At matched budgets, NF2M-MIS and NF2M-IS perform on par in downstream distributional distances and inference time.Both methods’ metrics improve monotonically as candidate budgets increase.Source: Figure 5 · p. 31, §A.2 · p. 32

A.3 METROPOLIS-ADJUSTED EXPONENTIAL LANGEVIN

• MAExL integrates the Gaussian-likelihood affine drift exactly, freezes the prior score, and applies a Metropolis–Hastings correction to produce an exact sampler.Its per-chain step sizes are adapted on log ω, with target acceptance 0.574; ω is the proposal-noise scale used to transfer step sizes across noise levels.Source: §A.3 · pp. 32–34
• Under EI-MALA assumptions M1–M2, Corollary 19 provides geometric Wasserstein convergence and pairwise contraction bounds for the MAExL kernel.These bounds apply to the conditional posterior kernel at fixed t and xt.Source: §A.3 · pp. 34–35
• Without MAExL, E-W2 worsens on every system and weighting scheme; on ALA-4 it rises by an order of magnitude, with across-seed standard deviation exceeding the mean.The 512- and 1024-step settings are statistically indistinguishable for NF2M-IS, while T-W2 stays essentially flat across the tested settings.Source: §A.3 · p. 35, Table 4 · p. 36

A.4 CENTRE-OF-MASS ADJUSTMENT FOR THE CONDITIONAL POSTERIOR

• The CoM adjustment extends Gaussian centroid augmentation to the diffusion posterior and uses Cartesian centroid densities for reweighting, rather than SBG’s radial correction.The adjustment derives valid importance weights and separately analyzes when matching the centroid proposal improves asymptotic ESS.Source: §A.4 · p. 36, §A.4.2 · p. 37

A.4.1 SETUP AND NOTATION

• The setup lifts the normalized Boltzmann density on centered configurations with an independent Gaussian centroid of variance 1/N, leaving the centered marginal unchanged.Restricting the Boltzmann density to centered configurations removes the translational degree of freedom.Source: §A.4.1 · pp. 36–37

A.4.3 THE CONDITIONAL GAUSSIAN LIFT

• After conditioning, the lifted posterior factorizes into the centered posterior and a Gaussian centroid posterior, so the centered posterior is unchanged by lifting.Rotational invariance is not required for this factorization.Source: §A.4.3 · pp. 37–38

A.4.4 IMPORTANCE WEIGHTS AND EFFECTIVE LOG-LIKELIHOOD

• The importance weights use the centered posterior and conditional Gaussian centroid factor, yielding consistent self-normalized estimates for integrable centered observables without requiring the flow to factorize.The Gaussian centroid factor must not be included twice when forming the numerator.Source: §A.4.4 · p. 38

A.4.5 WHEN CENTROID MATCHING IMPROVES ESS

• The ESS analysis compares asymptotic fractional ESS for proposals with the same centered component but different centroid proposals, under stated normalization and support conditions.The comparison concerns asymptotic fractional ESS, not the ESS of a particular finite sample.Source: §A.4.5 · p. 38
• Centroid matching gives ρ(q, π) ≤ ρ(qmatch, π), with equality exactly when k = g almost everywhere; the comparison requires sampling from Q gt and does not apply to generic non-factorised flows.For a mismatched factorised proposal, changing only the likelihood denominator is insufficient; the proposal must actually sample from Q gt and evaluate its density.Source: §A.4.5 · pp. 38–39

A.5.1 TARFLOW AND STARFLOW

• TarFlow uses causal-Transformer coupling layers, while StarFlow adds deep–shallow depth allocation and a guarded affine head.The molecular adaptation uses one token per atom and handles SE(3) symmetry externally through centering, a lifted centroid, and random-rotation augmentation; it omits StarFlow’s autoencoder and classifier-free guidance.Source: §A.5.1 · pp. 39–40

A.5.2 MAKING THE FLOW CONDITIONAL ON (xt, t)

• Time enters every block through AdaLN-Zero, while xt enters through cross-attention or a causal prefix that preserves the autoregressive Jacobian.TarFlow conditions every block, whereas StarFlow conditions only its deep block; cross-attention is used by nearly all reported models, with prefix conditioning selected for chignolin.Source: §A.5.2 · p. 40, Figure 7 · p. 42

A.5.3 THE CONDITIONING CACHE

• A preparation step builds conditioning state once per distinct (xt, t), so particles reuse it via row lookup rather than recomputing identical conditioning pathways.The reduction was verified bit-for-bit; with StarFlow prefix conditioning, the prepared state is the same key/value buffer used by autoregressive sampling.Source: §A.5.3 · p. 41

A.5.4 COMPUTATIONAL COMPLEXITY

• The appendix specifies per-call and amortised costs for molecular TarFlow and StarFlow, distinguishing backbone work from cross-attention and prefix-conditioning overhead.Cross-attention adds attention and parameters, whereas prefix conditioning adds a pre-fill and doubled cache; the pre-fill is amortised across particles sharing a conditioning row.Source: §A.5.4 · pp. 41, 43, Table 5 · p. 42, Table 6 · p. 43
• At N = 22–138 atoms and C = 256–512, the NC^2 linear term dominates attention’s N^2C term, while conditioning leaves asymptotic cost unchanged.Both density evaluation and sampling cost O(Ktot(N^2C + NC^2)); sampling remains N sequential steps, and prefix conditioning uses a doubled KV cache, about 5.5 MiB per chignolin conditioning row.Source: §A.5.4 · pp. 42–43

B.1 DATASETS AND METRICS

• Datasets cover alanine dipeptide, trialanine, alanine tetrapeptide, hexa-alanine, and 138-atom chignolin, with some metastable states potentially absent from training.Training uses the first 1 µs, validation the next 0.2 µs, and testing the remainder, except alanine dipeptide follows the split of Klein & Noé (2024).Source: §B.1 · pp. 43–44
• The held-out evaluation uses the same trajectory splits, force fields, and evaluation subsets as Tan et al. (2025) to enable direct comparison with their reported numbers.Each MD simulation uses a 1 fs integration time step.Source: §B.1 · p. 44
• Target energies and gradients are evaluated with a batched PyTorch reimplementation of the OpenMM force-field potential, validated against OpenMM’s double-precision Reference platform.It is 30–400× faster per configuration than OpenMM’s serial loop and returns finite differentiable energies at clashes.Source: §B.1 · p. 44
• Evaluation uses empirical Wasserstein-2 distances between generated and reference samples, subsampling 104 configurations from the held-out trajectory and generated samples when needed.The optimal coupling is computed with the POT library.Source: §B.1 · p. 44
• E-W2 compares force-field energies and is sensitive to fine structural details, while T-W2 compares backbone dihedral angles and captures macrostructure such as metastable-state occupancy.T-W2 is robust to fine-grained noise.Source: §B.1 · p. 44
• The appendix lists unconditional TarFlow and StarFlow architectures, with StarFlow variants differing in the depth of their final two blocks.The first four StarFlow blocks have two layers each; only the final two block depths vary.Source: Table 7 · p. 45, Table 8 · p. 45

NF2M and unconditional-flow training

• Unconditional flows and NF2M use a shared AdamW recipe, while NF2M trains for 1000 epochs and uses an EDM noise schedule.The shared recipe uses bf16 mixed precision, batch size 512, and EMA decay 0.999.Source: §B.2.1 · p. 45

NF2M architecture selection

• NF2M selects among six conditional-flow variants and K = 5 or K = 8 using NF2M-IS validation performance; TarFlow variants fail to train on chignolin.NF2M-MIS selection is restricted to K = 8 and uses NF2M-IS performance.Source: §B.2.1 · pp. 45–46, Table 9 · p. 46

Velocity teacher

• One flow-matching velocity teacher per system supplies the initialization for all four flow-map methods and is also evaluated as a standalone baseline.It is trained on a linear data-to-noise interpolant and downstream methods read its EMA weights.Source: §B.2.1 · p. 46

Flow-map training

• Flow-map training uses ESD pretraining for FALCON variants and LSD pretraining for F2D2 and SCALLOP before their method-specific distillation objectives.Pretraining shares one run per objective between its two student methods.Source: §B.2.1 · pp. 46–47

Training budget

• Flow-map baseline pipelines are time-matched to NF2M despite smaller backbones, yielding 2.7–9.0× as many training-set passes.The flow-map backbones have 4.1M or 6.1M parameters on ALA systems, compared with 18–87M for NF2M flows.Source: §B.2.1 · p. 47

Implementation departures

• The reproduction changes several baseline implementation details, including replacing a fully detached self-distillation correction after it diverged in a synthetic diagnostic.Other reported fixes include divergence-target handling and an unimplemented SCALLOP loss-uncertainty weighting parity item.Source: §B.2.1 · p. 48

Shared initialization

• NF2M shares one first-step denoising approximation across N trajectories by pooling NM candidates, introducing only weak trajectory correlation because the kernel depends weakly on terminal noise.At terminal noise, the first denoising kernel is essentially independent of the terminal state.Source: §B.2.2 · p. 49

Evaluation protocol

• The evaluation uses self-normalized importance sampling and post-hoc MALA, with FALCON variants on chignolin limited to 10^5 proposals by Jacobian cost.Most methods use 10^6 proposals per seed; the chignolin exception uses 10^5 proposals.Source: §B.2.2 · p. 51

Refinement and inference cost

• At matched target-energy budgets, NF2M spends calls on in-loop MAExL and post-hoc MALA, while other methods use post-hoc MALA ladders alone.F2D2 and SCALLOP are 19–22× faster than the velocity teacher, whereas FALCON variants are never faster and reach 6.3× slower on ALA-6.Source: §B.2.2 · p. 52, Table 23 · p. 54

Matched-budget results

• With post-hoc MALA, NF2M leads in energy on every system except ALA-6 and in torsion on ALA-6 and chignolin; without MALA, it leads in energy everywhere and torsion on the three largest systems.NF2M’s own distances change little when its refinement ladder is removed, whereas baseline energy distances can worsen by up to an order of magnitude on larger systems.Source: §B.2.2 · pp. 52–53, Table 24 · p. 55, Figure 8 · p. 56

Comparison with SBG

• On the largest systems NF2M improves on the best SBG-SMC variant on both metrics, while SBG-SMC is strongest in torsion on the smallest peptides.SBG-SMC is less stable across seeds and its annealing ladder raises inference time at comparable training cost.Source: §B.2.2 · pp. 53–54, Table 26 · p. 58