Normalizing Flow Maps and Reward Mixing: How Generative Models Balance Sampling Speed and Exact Likelihood

The core task of generative models can be broken into two problems: how to turn noise into data (the generative mapping), and what reward to use to guide sampling toward the outcomes downstream applications actually want. This Microsoft Research talk revolves around two recent papers, corresponding to these two problems respectively—Normalizing Flow Maps (NFM) and Martini (surrogate reward mixing).

Part 1: Normalizing Flow Maps (NFM)

The Trade-off Dilemma of Existing Generative Model Families

The talk begins with a clear comparative picture. Diffusion models and flow matching models produce high-quality samples, but simulation is expensive and likelihood is hard to obtain; flow maps support few-step sampling with sample quality close to diffusion models, but likewise lack cheap likelihood; normalizing flows are naturally one-step mappings and technically possess likelihood, but if trained directly on clean data, sample quality is not ideal.

The speaker specifically mentioned an often-overlooked detail: training autoregressive flows (TAR flow) "almost necessarily" requires adding a small amount of Gaussian noise to the training distribution, then denoising via the Tweedie identity. On CIFAR-10, training without added noise yields an FID of 40.83, which drops to around 6.55 with a very small amount of added noise. The cost is this: once Tweedie denoising is applied, a many-to-one mapping forms at the endpoints, and exact likelihood on the clean data distribution is lost. This is precisely why such models are difficult to use for scientific sampling and other applications that require likelihood.

The Tweedie identity is a classic result in statistical estimation theory. In the context of diffusion models, it usually refers to: for noised data x_t = α_t x_0 + σ_t ε, there exists an analytic relationship between the optimal denoiser of clean data (i.e., the posterior mean E[x_0|x_t]) and the score function ∇log p_t(x_t). Specifically, the score function can be obtained by multiplying the denoiser by a coefficient related to the noise level. This identity makes "training a denoising network" equivalent to "learning the score of the data distribution," and is the mathematical foundation of the training objectives for diffusion models and flow matching. The phrase "denoising via the Tweedie identity" mentioned in the text refers to substituting the network's denoising estimate back into this identity during sampling or likelihood computation to recover the score. Note that the denoising operation itself is contractive: as noise approaches zero, many noisy samples map to the vicinity of the same clean sample, forming a many-to-one mapping that destroys the invertibility of the mapping, and exact likelihood fails accordingly.

Core Construction of NFM

The definition of NFM is straightforward: along the transport path from noise to data, the mapping between any two time points s and t is strictly invertible. The key change this brings is that the change of variables no longer occurs only between the endpoints 0 and 1, but can occur between any two intermediate times. If we set s=1 and t=0, it degenerates back to an ordinary normalizing flow, so NFM can be viewed as a generalization of normalizing flows that simultaneously combines the capability of few-step operators with exact likelihood.

Regarding parameterization, the speaker presents three schemes:

• Two-time parameterization: directly model the s→t mapping with an invertible network, but it cannot naturally preserve identity (the mapping should be stationary when s=s), requiring an additional penalty term as a constraint.
• Denoiser parameterization: the invertible part is characterized by a two-time delta mapping that can recover the average velocity, thereby obtaining an invertible mapping from any intermediate point to the endpoint; the cost is that an arbitrary s→t mapping is not necessarily invertible, but this buys a richer space of architectural choices.
• Gated two-time parameterization: add s−t gating to each scale and shift transformation of the autoregressive flow, preserving identity while remaining invertible for any s and t.

The change of variables formula is the cornerstone of likelihood computation for normalizing flows and indeed all generative models based on invertible mappings. If z = f(x) is a strictly invertible and differentiable mapping, then the densities of the two distributions satisfy p_X(x) = p_Z(f(x)) · |det ∂f/∂x|. Directly computing the Jacobian determinant in high-dimensional settings is extremely expensive, so practical models typically adopt autoregressive structures or coupling layer designs that make the Jacobian matrix triangular, reducing the determinant to a product of diagonal entries and thereby lowering the O(d³) computation to O(d). The key innovation of NFM in the text is to generalize this formula from "between endpoints" to "between any intermediate times s and t": since time 1 is a known Gaussian distribution, as long as the s→t mapping is invertible, the exact density of the marginal distribution at any time can be written down. This is precisely the mathematical reason why NFM can simultaneously achieve few-step sampling and exact likelihood.

Two Training Perspectives: Particle-Level and Distribution-Level

NFM can be trained at two levels, and at optimum the two agree. From the particle-level perspective, it is just an ordinary flow map, and one can directly apply the regression objectives of mean flow or improved mean flow, with the only difference being that v_st is parameterized by NFM. From the distribution-level perspective (i.e., distribution matching), it exploits the fact that time 1 is a Gaussian distribution, so the marginal distribution at any time t can be written exactly as a Gaussian density plus a log-determinant term. This resembles maximum likelihood training, but along the entire interpolation path and without injecting noise into the data.

The further few-step NFMD adopts a student-teacher structure: a frozen teacher goes from p1 to ps, and the student jumps from s to t and then pulls back to s, using a self-consistency objective to maintain few-step sampling capability along the path.

Experimental Results: From Image Generation to Boltzmann Generators

On CIFAR, particle-level matching NFM-R achieves the best FID, but the speaker candidly acknowledges a practical issue: a mapping that is strictly invertible in theory may be numerically non-invertible in practice due to the training objective, requiring techniques such as noise annealing to zero as assistance. On 64×64 images, among parameter-matched invertible models NFM has the strongest sample quality, and the newly proposed FLD metric also shows that generation diversity is well preserved.

A more convincing application is the Boltzmann generator in the molecular domain. Taking the 10-amino-acid chignolin peptide as an example, NFM-D after reweighting can almost exactly match the target energy distribution, outperforming previous normalizing-flow-based methods and newer autoregressive models. When transferred zero-shot to four-amino-acid and eight-amino-acid systems, NFM-D performs best on five of six metrics. The TICA heatmaps shown by the speaker indicate that NFM-D not only covers all modes of the true molecular dynamics trajectory, but also does not hallucinate spurious modes as autoregressive flows do.

Another result the speaker personally values most: with usable likelihood, one can perform RLHF on continuous image models just as with language models. The team directly applied GRPO, DPO, or maximum-entropy RL, using a large-scale VLM as the reward, guiding "golden retriever" into "golden retriever running on the beach" and "bighorn sheep" into "bighorn sheep by the sea," with obvious distribution shift but a natural process. The reward-versus-training-update curves also outperform baselines such as SMC or best-of-n sampling.

Part 2: Approximating Expensive Oracle Rewards with Surrogate Reward Mixing

Problem Setting: Oracle Rewards Are Too Expensive, So Cheap Surrogates Must Be Used

The post-training mentioned earlier requires a large number of queries to the oracle reward r*, for example in the image setting roughly a thousand-fold order of queries and 256 samples, which is extremely costly. Martini asks: if r* can only be accessed in a limited way through a small feedback set D, can we approximate the tilted distribution q* of r* using a combination function G of multiple cheap surrogate rewards?

The feedback set comes in two forms: independent scalar-value feedback, or ranking preference feedback (corresponding to the Plackett-Luce model, whose binary case is the Bradley-Terry model commonly used in RLHF).

Why Can't We Just Use the Best Single Surrogate?

The speaker uses the example of language model perplexity to illustrate: if one uses only the best single surrogate and keeps increasing the guidance strength, reward hacking occurs and performance actually degrades. Using multiple surrogates simultaneously, since it is impossible to fool all rewards at once, naturally provides protection. This is also why one needs to learn a combination function rather than simply selecting one.

An intuitive counterexample shows this is not a Pareto optimality problem: when both rewards contain symmetric squared terms, the Pareto optimal set can never reach the true optimum located in the interior. Theoretically, in the small-tilt limit, minimizing the mean squared conditional error (MSC) between G∘R and r* is equivalent to minimizing the KL divergence to q* (up to third-order terms), but there always exists an "information floor" V̄ defined by variance, an irrecoverable loss caused by the space that the surrogates cannot cover.

Reward hacking refers to the optimizer finding high-scoring regions of a surrogate objective while deviating from the truly desired behavior. In RLHF, the reward model is only an approximation of human preferences, and reinforcement learning systematically exploits its errors: generating text that scores extremely high according to the reward model but is poorly evaluated by humans. The perplexity example in the text is precisely a quantitative version of this phenomenon—the greater the guidance strength of a single surrogate, the easier it is to enter the surrogate's "blind spots," and true performance first rises and then falls. The protective mechanism of multi-surrogate combination lies in: the error directions of different surrogates are usually not fully correlated, and fooling all surrogates simultaneously is far harder than fooling one, so the combination function acts as an implicit regularization on surrogate errors. This also explains why Martini learns a nonlinear combination function rather than simply selecting or linearly weighting the best-performing surrogate.

Validation: Language, Images, and Protein Docking

In the language modeling experiment, the generation perplexity of Qwen 2.5 1.5B is used as the oracle reward, with BERT, distilled BERT, token entropy, etc. as surrogates, and the feedback set contains only 256 values or rankings. The KL divergence between the learned q_G and q* in Martini is significantly lower than uniform weighting or Pareto-type methods, and perplexity continues to decrease as guidance strength increases.

On the image task, using only 32 prompts, clear improvements are achieved on embedding-space distribution metrics such as Wasserstein and KDE, and both the value-feedback and ranking-feedback settings are effective.

The experiment closest to real-world application is protein-ligand co-folding: with binding affinity as the oracle reward, G is fitted offline on experimental PKD data and then used for structure prediction. Using the Boltz flow map version, the success criterion (RMSD less than 2 Å and PB valid above threshold) rises from 61.8% to 64.5%. This means external affinity data can improve the physical plausibility of generated poses without re-querying the reward.

The Plackett-Luce model is a classic probabilistic model for ranking data, used to characterize the process of "selecting the first place from a set of candidates." If each candidate i has an implicit utility value s_i, the probability of selecting i is exp(s_i) / Σ_j exp(s_j). When there are only two candidates, the model degenerates to the Bradley-Terry model, i.e., P(i beats j) = σ(s_i − s_j), which is exactly the standard form for training preference reward models in RLHF. In the text, ranking preference feedback is mapped to the Plackett-Luce model, meaning that the feedback set D provides not absolute reward values but the relative order of several candidates; Martini needs to fit the combination function G on such data containing only order information. Compared with scalar-value feedback, ranking feedback carries less information but is cheaper to collect and closer to the actual form of human annotation, so this setting is crucial to the practicality of the method.

Open Questions and Boundaries

The Q&A session discussed several key limitations. Complementarity between surrogates is a core assumption: if multiple surrogates provide the same information, it can be identified and removed via correlations on the feedback set; but even if complementary, if the surrogates collectively fail to cover certain regions of the task space, the information floor V̄ cannot be breached—the speaker explicitly stated this "definitely cannot be overcome." Regarding the number of sampling steps for NFM, practice shows that improvements are no longer linear beyond four steps, and few-step NFMD outperforms the single-step version at two to four steps. As for how to use the combination function to post-train exact-likelihood models, the speaker left this as an open question.

One-line Summary

NFM restores exact likelihood via invertible mappings, while Martini approximates expensive oracle rewards with surrogate reward combinations.

Paragraph Summary

The article revolves around two core problems of generative models. The first part introduces Normalizing Flow Maps (NFM), which generalize the invertibility of normalizing flows from endpoints to between any intermediate times, thereby restoring exact likelihood while maintaining few-step sampling capability and avoiding the dilemma that autoregressive flows must inject noise to train. NFM supports both particle-level and distribution-level training perspectives, and achieves method-level results superior to baselines on CIFAR, 64×64 images, Boltzmann generators, and likelihood-based RLHF post-training. The second part introduces Martini, which addresses the problem of excessively high oracle reward query costs by learning a nonlinear combination function of multiple cheap surrogate rewards to approximate the target tilted distribution, avoiding the reward hacking problem of a single surrogate. It is validated on language, image, and protein-ligand co-folding tasks, and points out that the information floor caused by insufficient surrogate coverage is an insurmountable capability ceiling.