PaperlayerView exact revision on arXiv

Source-linked AI summary

Normalizing Flow Maps

Francesco M. Ruscio, Timothy James Hitge, Fred Zhangzhi Peng, Danyal Rehman, Zander W. Blasingame, Nicholas M. Boffi, Michael M. Bronstein, Alexander Tong, Avishek Joey Bose

arXiv:2610.22711v1cs.LG

TL;DR

Flow maps enable few-step generation but lack tractable likelihoods, limiting their use in scientific applications such as Boltzmann Generators. Normalizing Flow Maps combine invertible few-step transport with exact likelihoods, and achieve best overall performance in conditional image generation while improving over prior continuous flow approaches for Boltzmann Generators.

  • Problem

    Flow maps reduce inference costs but lack tractable likelihoods, limiting their use in scientific applications such as Boltzmann Generators.

    Source: §1 · p. 2

  • Method

    Normalizing Flow Maps learn invertible few-step transport maps with exact likelihoods, trained through pointwise regression or distributional matching.

    Source: §1 · p. 2

  • Results

    NFMs achieve best overall performance in conditional image generation and improve over prior continuous flow approaches for Boltzmann Generators.

    Source: §1 · p. 3, §4.1 · p. 8

  • Takeaways & Limitations

    NFMs support importance sampling in Boltzmann Generators and RL-based post-training of image-generation models.

    Source: §6 · p. 10

    Image post-training relies on imperfect reward proxies that may amplify reward-model biases or favor images misaligned with user intent.

    Source: ETHICS STATEMENT · p. 11

Abstract

from arXiv · show

Scalable approaches to learning neural dynamical systems have placed diffusion, flow matching, and few-step generators in flow maps at the forefront of generative modeling. Such state-of-the-art models have limited access to explicit model densities despite evolving from classical generative modeling approaches that leveraged tractable likelihoods, e.g., Normalizing Flows and Autoregressive models. In this paper, we combine these threads of few-step generation and tractable sample likelihoods and introduce Normalizing Flow Maps (NFM), a class of exactly invertible pushforward maps that generalizes normalizing flows and flow maps. More precisely, NFMs learn the integrator of an underlying dynamical system, but crucially, they provide efficient exact likelihood evaluation by construction. To enable scalable training of NFMs, we derive two complementary axes of consistency inherent to NFMs that yield (self) distillation objectives: (1) pointwise matching, which unlocks $\ell_2$-based flow-map style regression losses, and (2) distributional matching, which matches pushforward distributions to the marginals of a prescribed forward process. On conditional image generation on CIFAR-10 and ImageNet-64, NFMs achieve state-of-the-art sample quality among parameter-matched normalizing flows, while supporting likelihood-based RLHF. For equilibrium sampling in Boltzmann Generators, NFMs outperform prior state-of-the-art normalizing flows on both single-system and transferable peptide benchmarks.

1 INTRODUCTION

Normalizing Flow Maps (NFM) combine few-step transport with exact sample likelihoods through invertible maps along a prescribed stochastic interpolant. The paper introduces pointwise and distributional matching objectives to train these maps and evaluates them on image generation, post-training, and equilibrium sampling.

Motivation

  • Repeated ODE or SDE integration makes generation, steering, and post-training computationally expensive, while exact flow-model likelihoods require costly Jacobian calculations and stochastic trace estimation.Path-space reward-tilting objectives also rely on high-variance value estimates and can lead to unstable post-training and mode collapse.Source: §1 · pp. 1–2
  • Flow maps compress numerical simulation into few network evaluations but lack invertibility and tractable Jacobian determinants, while noise-injected normalizing flows lose cheap sample likelihoods.Together, these limitations motivate a model combining flow-map sampling efficiency with exact likelihoods.Source: §1 · p. 2

Current work

  • NFM learns invertible pushforward maps that support few-step sampling and exact sample likelihoods, generalizing normalizing flows and flow maps.Restricting the maps to prior and data endpoints recovers a normalizing flow; querying arbitrary time pairs enables few-step transport. applications include Boltzmann generators.Source: §1 · pp. 2–3
  • NFM-R uses pointwise matching for ℓ2 regression, while NFM-D matches pushforward distributions to marginals of the known forward process.Both training approaches use complementary matching principles to produce self-distillation objectives.Source: §1 · pp. 2–3
  • On CIFAR-10 and ImageNet-64, NFM improves over MLE-trained normalizing flows in FID and FLD; for Boltzmann Generators, it improves over prior continuous flow approaches.The paper also reports that NFM's exact tractable likelihoods enable direct post-training unlike diffusion models.Source: §1 · p. 3

2 BACKGROUND AND PRELIMINARIES

Dynamical generative models learn maps that push a tractable base distribution toward the data; the background situates this setup in stochastic interpolants, flow maps, and normalizing flows with exact likelihoods.

2.1 STOCHASTIC INTERPOLANTS AND FLOW MAPS

2.2 NORMALIZING FLOWS

  • Normalizing flows parameterize diffeomorphisms with exactly invertible architectures, enabling exact likelihood evaluation through the change-of-variables formula.Source: §2.2 · p. 4
  • Certain normalizing flows are universal density approximators and compose elementary diffeomorphisms for maximum-likelihood training; scalable approaches include Transformer autoregressive flows with permutation layers.The cited universality results are Teshima et al. (2020) and Huang et al. (2018); the Transformer autoregressive flow approach is TARFlow.Source: §2.2 · p. 4

3 METHOD

Normalizing Flow Maps learn invertible transport along a prescribed probability path, combining few-step generation with exact clean-sample likelihoods. The method derives pointwise and distributional consistency objectives for distillation and enables likelihood-based post-training.

3.1 NORMALIZING FLOW MAP

  • NFMs use invertible two-time maps along a probability path for few-step transport and exact density evaluation; restricting the map to (1, 0) recovers a normalizing flow.Unlike noise-augmented MLE, which provides exact likelihoods at the noisy distribution pσ, NFMs amortize transport across noise levels and provide clean-sample likelihoods p0.Source: §3.1 · pp. 4–5

3.2 POINTWISE AND DISTRIBUTIONAL MATCHING

  • Pointwise matching aligns particle trajectories with the probability-flow ODE, while distributional matching aligns NFM pushforwards with the forward process marginals.With the tractable reference p1, the induced intermediate density can be evaluated by change of variables rather than optimized through a diffusion-model weighted ELBO.Source: §3.2 · pp. 5–6

3.3 DISTILLATION LOSSES

  • NFM-R uses intermediate-time regression for an exact-likelihood one-step map and controls endpoint 2-Wasserstein error; NFM-D uses conditional supervision to learn marginal pushforwards without evaluating pt.Conditional NFM-D upper-bounds the marginal objective while sharing its minimizers; minimizing the single-time objective controls the generated endpoint when marginal discrepancy is locally regular. Its two-time extension uses a frozen single-time teacher as a surrogate source density.Source: §3.3 · pp. 6–7

3.4 POST-TRAINING NORMALIZING FLOW MAPS

  • Tractable endpoint densities let NFM post-training optimize distribution-level reward and regularization objectives using score-function gradients and density ratios.These signals support non-differentiable rewards and density-ratio methods such as DPO and GRPO.Source: §3.4 · p. 7

4 EXPERIMENTS

Experiments evaluate NFM on conditional image generation and RL post-training, plus Boltzmann Generator equilibrium sampling. Results show image-quality gains and improved molecular sampling, including transferable and temperature-tuned settings.

4.1 IMAGE GENERATION AND POST-TRAINING

  • NFM-R achieves the best overall image-generation performance, while NFM-D and NFM-R improve FID and FLD over MLE-trained iTARFlow and NFM-D beats SoftFlow on CIFAR-10.Maintaining NFM-R invertibility requires noise annealing during training and FP64 inference; multi-step sampling improves quality but breaks NFM-R invertibility, whereas NFM-D remains invertible across steps.Source: §4.1 · p. 8, Table 1 · p. 8
  • Training iTARFlow on clean data preserves invertibility but substantially worsens image sample quality.Source: §4.1 · p. 8
  • NFM-D post-training yields gains on multiple rewards against best-of-N and sequential Monte Carlo baselines over training.NFM-R uses maximum-entropy regularization and pathwise gradients, while NFM-D uses KL-regularized REINFORCE.Source: §4.1 · pp. 8–9, Figure 3 · p. 8

4.2 BOLTZMANN GENERATORS

  • On Chignolin, NFM-D substantially outperforms SBG and is competitive with autoregressive ArBG, with no other continuous-flow baseline reported as comparable.Chignolin is the most challenging single-peptide benchmark because of its β-hairpin structure.Source: §4.2 · p. 9, Table 2 · p. 9
  • For transferable sampling, NFM-D improves over MLE-trained TARFlow and beats Prose and Robin on five of six metrics, using fewer parameters than Prose.NFM-D uses 118M parameters versus Prose's 288M; the evaluation covers unseen 4- and 8-residue peptides.Source: §4.2 · p. 9, Table 3 · p. 10
  • Temperature tuning further improves transferable performance, with NFM-D achieving the best performance and substantially improving E-W2 on 8-residue peptides.NFM-D tunes the base-distribution temperature, while Robin tunes its attention temperature.Source: Table 4 · p. 9, §4.2 · p. 10

5 RELATED WORK

Related work contrasts invertible generative models with exact likelihoods against high-fidelity stochastic-interpolant and flow-map methods whose likelihood evaluation remains costly.

Normalizing flows

Stochastic interpolants and flow maps

6 CONCLUSION · Appendix · A THEORETICAL RESULTS

Normalizing Flow Maps combine few-step sampling with exact likelihoods, enabling applications in Boltzmann generators and post-training image models. The conclusion highlights a tradeoff between sample quality and numerical invertibility, while the ethics statement notes limits and risks in the evaluated applications.

6 CONCLUSION

  • NFMs demonstrate superior sample quality to existing normalizing flows while avoiding Tweedie denoising, which requires spatial gradients and breaks invertibility guarantees.Source: §6 · p. 10
  • NFMs provide cheap likelihoods used for importance sampling in Boltzmann generators and straightforward RL-based post-training of image generation models.Source: §6 · p. 10
  • NFM-R achieves the highest sample quality but can become numerically non-invertible, impeding RL methods that rely on score-function or KL-regularization.The conclusion points to training schemes that combine pointwise and distributional consistency as important for balancing these benefits.Source: §6 · p. 10

AI USE STATEMENT

  • The authors used generative AI to improve author-written text and mathematical presentation, identify literature, assist implementation, and provide critical feedback, but not to draft sections or generate synthetic datasets.Source: AI USE STATEMENT · p. 11
  • The authors reviewed AI-assisted work, read AI-suggested literature, used separate LLM reviewers for code checks and tests, and take responsibility for the final content and artifacts.Source: AI USE STATEMENT · p. 11

ETHICS STATEMENT

  • Image experiments post-train class-conditional generators with text-to-image reward models, whose imperfect proxies may amplify biases or reward images misaligned with user intent.Source: ETHICS STATEMENT · p. 11
  • Molecular experiments assess equilibrium conformation sampling on existing peptide benchmarks using molecular dynamics data; the studied molecules are considered safe and non-toxic, while other systems need use-specific evaluation.Source: ETHICS STATEMENT · p. 11

A.1 NFM-R PROOFS

The appendix proves endpoint control for pointwise distillation under finite-second-moment and probability-flow assumptions. Its proof couples the true and predicted endpoint maps using the same terminal sample.

A.1 NFM-R PROOFS

  • With diagonal time pairs sampled with probability η, the objective combines flow matching and distillation as LR(θ) = ηLFM + (1 −η)Ldistill.The distillation expectation is conditional on an off-diagonal time pair.Source: §A.1 · p. 21
  • Proposition 3 establishes endpoint control for pointwise distillation under finite second moments and a probability flow whose trajectories follow the marginal velocity and end at the identity.The proof’s conclusion is stated in the supplied text only as a reference to combining equations; the bound itself is not included.Source: §A.1 · pp. 21, 23
  • The proof evaluates the true and predicted endpoint maps on the same terminal sample to form a coupling, using equality of one-time marginals rather than equality of sample paths.The supplied passages identify this shared-sample coupling as the basis for controlling squared Wasserstein distance.Source: §A.1 · p. 22
  • For an unweighted-sum objective, the proof gives a corresponding bound; under uniform source-time sampling, it reduces to (Ctr + bL2)LR.This is the loss-normalization remark’s stated specialization.Source: §A.1 · p. 23

A.2 NFM-D PROOFS … C.1 BASELINES

The appendix establishes properties of NFM-D objectives and describes a tractable two-time extension, while the accompanying sections explain the likelihood tradeoff in denoised flows and the Boltzmann Generator setting and baselines.

A.2 NFM-D PROOFS

  • With finite expected mutual information, conditional and marginal NFM-D objectives share minimizers, and the conditional objective is bounded above by the marginal objective plus a parameter-independent term.The added expected mutual information depends only on the prescribed corruption process, not on θ.Source: §A.2 · pp. 23–24
  • The appendix states that the empirical likelihood-form conditional NFM-D objective is bounded below under the stated time-sampling condition.The bound is finite and independent of θ under the proposition’s assumption.Source: §A.2 · pp. 24–25
  • Single-time NFM-D endpoint control is stated under local Lipschitz regularity at t = 0 and a time-sampling density bounded below near the endpoint.The proposition also invokes a finite expected single-time objective.Source: §A.2 · p. 26
  • For two-time transport, a frozen single-time NFM-D model supplies a tractable surrogate source density when the intermediate marginal is intractable.The single-time objective learns the marginal path from the reference distribution, while the two-time objective extends supervision across intermediate times.Source: §A.2 · pp. 26–27

B EXACT LIKELIHOOD WITHOUT TWEEDIE DENOISING

  • Noise augmentation and Tweedie denoising improve iTARFlow sample quality, but the original flow likelihood does not describe denoised outputs because the transformation’s Jacobian is unaccounted for.Without noise, samples are visibly poor despite exact invertibility; Tweedie denoising is not guaranteed to be invertible.Source: B EXACT LIKELIHOOD WITHOUT TWEEDIE… · p. 28, Figure 5 · p. 28

C EXTENDED BACKGROUND ON BOLTZMANN GENERATORS

  • Molecular equilibrium sampling is difficult because high-dimensional, multimodal energy landscapes slow local methods; Boltzmann Generators learn tractable proposal distributions for sampling and density evaluation.Their samples can support Monte Carlo corrections for mismatch with the target distribution (Noé et al., 2019).Source: C EXTENDED BACKGROUND ON BOLTZMANN… · p. 29
  • Boltzmann Generator proposals can be corrected with self-normalized importance sampling or sequential Monte Carlo bridging and resampling.For importance sampling, the unknown partition function cancels in the expectation ratio.Source: C EXTENDED BACKGROUND ON BOLTZMANN… · p. 29

C.1 BASELINES

C.2 METRICS

NFM evaluation compares generated samples with MD references using energy-based, torus-based, and TICA-based metrics to assess local geometric fidelity, conformational structure, and distributional agreement.

C.2 METRICS

  • Energy 2-Wasserstein Distance (E −W2) compares generated and reference energy marginals, detecting local distortions in bond lengths, angles, and short-range interactions.Source: §C.2 · p. 30
  • Torus 2-Wasserstein Distance (T-W2) compares backbone dihedral configurations using shortest periodic angular displacements, capturing large-scale structure and mode coverage.Source: §C.2 · p. 30
  • TICA 2-Wasserstein Distance (TICA-W2) compares generated and reference distributions along leading TICA components, measuring agreement in slow conformational and metastable structure.Source: §C.2 · p. 30

D IMPLEMENTATION DETAILS … E.1 IMAGE GENERATION

The implementation sections describe invertible transformer flow architectures, objectives and guidance strategies, numerical-invertibility and reinforcement-learning procedures, and image-generation training configurations. Together, they specify how NFM-D and NFM-R are trained and used across image-generation experiments.

D.1 TRANSFORMER AUTOREGRESSIVE FLOWS

  • TARFlows compose affine transformations parameterized by causal transformers, alternating autoregressive direction after the first block to construct exactly invertible networks.Each block can be reversed by autoregressive rollout.Source: §D.1 · p. 31

D.2 NFM ARCHITECTURE

  • NFM-D adds sinusoidal time conditioning to attention and adaptive layer norms, while NFM-R uses multi-token in-context conditioning and general-purpose Transformer improvements.NFM-R removes adaptive layer norms; in-context conditioning supports class, timestep, and classifier-free-guidance conditions and may improve parameter efficiency.Source: §D.2 · p. 31

D.3.1 IMPROVED MEAN FLOW OBJECTIVE; D.4 PSEUDOCODE

  • Improved Mean Flow replaces the unavailable true average velocity with sample-conditional instantaneous velocity and a model estimate at the boundary; pseudocode specifies NFM-R JVP and NFM-D likelihood objectives.The NFM-D procedures include one-time and two-time objectives with Gaussian-prior and log-Jacobian terms.Source: §D.3.1 · p. 32, §D.4 · pp. 34–35

D.3.2 CONDITIONING AND GUIDANCE

  • NFM-D applies TARFlow-style guidance to affine parameters, whereas NFM-R conditions average velocity on guidance strength and interval for flexible few-step inference without added cost.NFM-R trains across varying guidance strengths and interval bounds; its CFG conditioning samples ω, smin, and smax.Source: §D.3.2 · pp. 32–33

D.3.3 NUMERICAL INVERTIBILITY

  • NFM-R’s regression objective can cause numerical non-invertibility in practice; annealing additive Gaussian noise on x0 during training can restore numerical invertibility for FP64 inference.The strategy enabled CIFAR-10 NFM-R fine-tuning with REINFORCE.Source: §D.3.3 · p. 34

D.5 REINFORCEMENT LEARNING FINETUNING

  • When numerically invertible, NFM reinforcement learning uses KL-penalized REINFORCE with a leave-one-out baseline and exact KL; when NFM-R is non-invertible, it uses reward gradients and maximum-entropy regularization.CIFAR-10 experiments target green frogs, while ImageNet-64 uses mixed reward functions and distinct NFM-R and NFM-D training setups.Source: §D.5 · p. 36

E ADDITIONAL RESULTS

E.1 IMAGE GENERATION

  • Experiments evaluate models on CIFAR-10 and ImageNet-64×64; NFM-R uses ReFlow couplings from EDM (Karras et al., 2022) and EDM 2 (Karras et al., 2024), respectively.CIFAR-10 noise-annealed NFM-R starts with σ=0.5, annealed to zero over 300 of 960 epochs; numerically non-invertible NFM-R uses 450 flow-matching epochs then 200 full-iMF epochs, while ImageNet-64 NFM-R uses 200 flow-matching epochs then 25 full-iMF epochs. Two-step CIFAR-10 NFM-D uses equally weighted one-time and composed likelihood objectives; inference follows 1 →0.75 →0, with an unconditional first step and CFG scale 3.2 on the second.Source: §E.1 · pp. 37, 39

E.1.1 REINFORCEMENT LEARNING EXPERIMENTS · E.2 BOLTZMANN GENERATORS · E.2.1 SINGLE PEPTIDES

The experiments show NFM-R improving over sampling baselines after reinforcement-learning fine-tuning, while NFM-D performs strongly on peptide benchmarks and benefits from more effective resampling. The evaluations span image-generation rewards and system-specific peptide sampling, including larger and transferable settings.

E.1.1 REINFORCEMENT LEARNING EXPERIMENTS

  • With a fixed noise seed and a reward favoring green frogs, fine-tuning shifts NFM-D outputs clearly and adjusts NFM-R outputs subtly.Figure 8 shows samples before and after REINFORCE fine-tuning.Source: §E.1.1 · p. 40, Figure 8 · p. 40
  • Across all four rewards and 10 seeds, post-trained NFM-R clearly improves on best-of-N and SMC baselines.The comparison uses the prompt “bighorn sheep standing on a sandy beach beside the sea.”Source: Figure 9 · p. 40, §E.1.1 · p. 40

E.2 BOLTZMANN GENERATORS

  • Boltzmann Generator evaluations cover system-specific AL3, AL4, AL6, and Chignolin, plus transferable sampling on ManyPeptidesMD.The alanine systems have three, four, and six residues, respectively; Chignolin is a 10-residue miniprotein.Source: §E.2 · p. 41

E.2.1 SINGLE PEPTIDES

  • NFM-D outperforms SBG on AL3, AL4, and AL6 and closes the gap to ArBG, though many results—especially T-W2—are statistically indistinguishable.The authors identify a need for more single-system equilibrium datasets, including systems up to 20 residues.Source: Table 9 · p. 41
  • The paper reports energy histograms for AL3, AL4, and AL6 and Ramachandran plots for each system, comparing NFM-D angular distributions with reference MD data.The energy histograms show ground-truth MD data, the proposal, and the resampled distribution.Source: §E.2.1 · p. 41, Figure 10 · p. 42, Figure 11 · p. 42
  • NFM-D and NFM-R have comparable raw AL3 metrics, but NFM-D substantially outperforms NFM-R after resampling.The paper attributes the resampling advantage to NFM-D’s better likelihood estimates.Source: Table 10 · p. 43, §E.2.1 · p. 43
  • For Chignolin, the study compares SBG and NFM-D, ablates architecture and training settings, and uses a smaller model than SBG to avoid overfitting.Table 11 also includes an SBG-iTARFlow reimplementation to control for architecture, optimizer, and training-recipe changes.Source: §E.2.1 · p. 43, Table 11 · p. 43

E.3 TRANSFERABLE SAMPLERS

On transferable peptide sampling, NFM-D generally lowers discrepancies relative to TARFlow, while additional unseen-octapeptide TICA plots demonstrate zero-shot performance. Training choices also improve the learned energy representation, with architecture and loss changes providing gains beyond optimizer and batch-size changes.

Comparison with MLE training

  • Muon training with a larger batch modestly improves most metrics; iTARFlow yields a substantially larger gain, and NFM-D loss further improves learned energy representations.Table 12 reports comparisons on unseen 4- and 8-residue peptides with a budget of 104 energy evaluations at T = 1.Source: §E.3 · p. 43, Table 12 · p. 43

TICA plots for additional unseen octapeptides

  • Additional TICA plots for six unseen octapeptides further demonstrate NFM-D's zero-shot performance.The plots show TICA projections for ground-truth MD and NFM-D.Source: §E.3 · p. 43, Figure 14 · p. 44

Peptide-level performance

  • NFM-D generally achieves lower E-W2, T-W2, and TICA-W2 discrepancies than TARFlow across 4AA and 8AA test peptides, with improvements varying by sequence and metric.The per-sequence plots compare both methods, with lower values indicating better agreement with the reference distribution.Source: §E.3 · p. 44, Figure 15 · p. 44, Figure 16 · p. 45

F ADDITIONAL IMAGENET-64 SAMPLES · F.1 UNCURATED NFM-R B/2 SAMPLES · F.2 UNCURATED NFM-D B/2 SAMPLES

The additional uncurated ImageNet-64 sample passages identify monarch butterflies and electric locomotives for both NFM-R B/2 and NFM-D B/2.

F.1 UNCURATED NFM-R B/2 SAMPLES

  • NFM-R B/2 samples include class 323, monarch butterfly.Source: §F.1 · p. 46
  • NFM-R B/2 samples include class 547, electric locomotive.Source: §F.1 · p. 46

F.2 UNCURATED NFM-D B/2 SAMPLES

  • NFM-D B/2 samples include class 323, monarch butterfly.Source: §F.2 · p. 47
  • NFM-D B/2 samples include class 547, electric locomotive.Source: §F.2 · p. 47
Loading 2610.22711v1…