Variational Flow Maps: Make Some Noise for One-Step Conditional Generation

Abstract: Flow maps enable high-quality image generation in a single forward pass. However, unlike iterative diffusion models, their lack of an explicit sampling trajectory impedes incorporating external constraints for conditional generation and solving inverse problems. We put forth Variational Flow Maps, a framework for conditional sampling that shifts the perspective of conditioning from "guiding a sampling path", to that of "learning the proper initial noise". Specifically, given an observation, we seek to learn a noise adapter model that outputs a noise distribution, so that after mapping to the data space via flow map, the samples respect the observation and data prior. To this end, we develop a principled variational objective that jointly trains the noise adapter and the flow map, improving noise-data alignment, such that sampling from complex data posterior is achieved with a simple adapter. Experiments on various inverse problems show that VFMs produce well-calibrated conditional samples in a single (or few) steps. For ImageNet, VFM attains competitive fidelity while accelerating the sampling by orders of magnitude compared to alternative iterative diffusion/flow models. Code is available at https://github.com/abbasmammadov/VFM

• The paper demonstrates that joint training of a flow map and a noise adapter enables efficient one-step conditional generation with improved multimodal support.
• It uses a variational framework to align noise with data posteriors, achieving superior perceptual metrics and reduced inference cost on complex tasks.
• Empirical results on tasks like ImageNet inpainting show that VFM delivers fast, diverse, and calibrated sample generation compared to iterative methods.

Variational Flow Maps: Efficient One-Step Conditional Generation with Learned Noise Alignment

Introduction and Motivation

Recent advances in continuous-time generative modeling—including diffusion [ho2020denoising, song_generative_2020, karras2022elucidating] and flow models [lipman2022flow, boffi_flow_2024, geng2025meanflowsonestepgenerative]—have converged to frameworks that transport tractable priors (typically standard normal distributions) to match complex data distributions. Standard algorithms require iterative inference, incurring prohibitive computational cost for applications demanding real-time or interactive conditional generation. Flow maps ameliorate inference cost by directly parameterizing the ODE solution, supporting one-step or few-step sampling. However, their unconditional design restricts applicability to tasks where sample generation must be constrained by external observations or conditions, such as image inverse problems.

The paper "Variational Flow Maps: Make Some Noise for One-Step Conditional Generation" (2603.07276) introduces a paradigm shift for conditional flow-based generation: rather than guiding the evolution path (as in legacy diffusion guidance), VFM reframes conditioning as noise selection—learning an amortized mapping from an observation $y$ (or condition) to a noise distribution $q_\phi(z|y)$. The flow map $f_\theta$ then decodes $z$ to the data space in a single forward pass. This approach closes the "guidance gap" for flow models and yields efficient, well-calibrated conditional samples.

Methodology

Observation-Dependent Noise Adaptation

Given an observation $y$, VFM learns a noise adapter network $q_\phi(z|y)$ to approximate the posterior over latent noises induced by the Bayesian inverse problem:

$$p(z|y) \propto \exp\left(-\frac{\|y - A(f_\theta(z))\|^2}{2\sigma^2}\right) p(z)$$

where $p(z)$ is the simple prior (e.g., $\mathcal{N}(0, I)$) and $A(\cdot)$ is a known measurement operator. Unlike guidance-based diffusion inference—which iteratively infuses likelihood gradients over the generation trajectory—VFM samples a latent $z \sim q_\phi(z|y)$ and decodes $x = f_\theta(z)$ in a single step. The structural insight is that, when $q_\phi$ and $f_\theta$ are jointly optimized, the restricted variational posterior (e.g., Gaussian) can be compensated by warping the flow map such that conditional data posteriors become easy to represent in the noise space.

Figure 1: One-step conditional generation with VFM: the observation-dependent noise adapter $q_\phi(z|y)$ is jointly trained with the flow map $f_\theta$, enabling posterior-aligned single-step generation.

Joint Variational Objective

The core innovation is a joint training objective, extending the VAE framework to triples $(x, y, z)$, matching factorizations

$q_\phi(z|y)p(y|x)p(x) \approx p_\theta(x, y|z)p(z)$

with $p_\theta(x, y|z)$ a relaxed Gaussian decoder. The loss encompasses:

• Data fit: $\mathbb{E}[\|x - f_\theta(z)\|^2]$ aligns decoded samples with true data.
• Observation fit: $\mathbb{E}[\|y - A(f_\theta(z))\|^2]$ promotes consistency with observations.
• KL regularization: $\mathrm{KL}(q_\phi(z|y)\|p(z))$ penalizes posterior mismatch.

A key result demonstrates that, under a linear-Gaussian scenario, only joint optimization recovers the posterior mean for any observation, whereas fixing $f_\theta$ renders Gaussian adapters insufficient and biased.

Structural Constraints and Flow Map Training

To regularize $f_\theta$ and ensure semigroup consistency, the data fit term is upper-bounded by the "mean flow loss" from recent literature [geng2025meanflowsonestepgenerative]. This higher-order constraint, evaluated over interpolated states, is crucial for learning expressive, reversible flow maps compatible with both multi-step and one-step generation.

Empirical Evaluations

Illustration: 2D Checkerboard Bayesian Inference

A didactic two-dimensional inverse problem with a multimodal "checkerboard" prior convincingly shows the weaknesses of fixed decoder + adapter (single-mode collapse) and unconstrained fitting (off-support artifacts). VFM, via joint training, captures both modes and preserves data manifold structure.

Figure 2: Prior and posterior densities in both data and noise spaces for frozen-$\theta$, unconstrained-$\theta$, and VFM. Only VFM exhibits correct multimodal support.

Large-Scale Posterior Sampling on ImageNet

On ImageNet 256×256, VFM is evaluated on a comprehensive suite of linear inverse problems (random/box inpainting, super-resolution, Gaussian and motion deblurring), amortizing over multiple measurement operators. Competing baselines include state-of-the-art iterative guidance methods (DPS, DAPS, PSLD, MPGD, FlowDPS) adapted to comparable backbone architectures. Results highlight:

• VFM dominates in distributional/perceptual metrics (FID, MMD, LPIPS, CRPS), often by a large margin. For instance, FID for box inpainting improves from 63–76 (baselines) to 33.3 (VFM).
• Baselines achieve higher PSNR/SSIM, but are heavily mean-seeking and visually less compelling (over-smoothing).
• Inference with VFM is two orders of magnitude faster (single step vs. 500+ steps for diffusion).

  Figure 3: Qualitative comparison on box inpainting: VFM produces diverse, plausible samples in the masked region, while baselines converge to mean-like or blurred solutions.

Posterior Diversity and Uncertainty Quantification

VFM naturally generates diverse, measurement-consistent conditional samples due to its Bayesian formulation. Diversity is directly visualized for highly ill-posed tasks. Furthermore, pixelwise uncertainty maps (computed across multiple VFM samples) expose interpretable ambiguity in reconstructions.

Figure 4: Visualization of the learned noise space: sampling $z \sim q_\phi(z|y)$ manifests structured latent codes tailored to the conditional task.

Figure 5: Posterior uncertainty quantification: mean and standard deviation of VFM samples capture true posterior variance tied to unobserved regions.

Reward Alignment via Amortized Noise Injection

VFM is extended to general reward alignment: starting from a pre-trained flow map, fine-tuning adapts both the noise adapter and flow map to maximize a scalar reward (e.g., human preference or text-image alignment), forming a reward-tilted distribution. Experimental results demonstrate competitive or superior reward scores (HPSv2, PickScore, ImageReward) with rapid convergence and preserved generation quality, achievable in a single function evaluation (1 NFE).

Figure 6: One-step reward-aligned generation: VFM fine-tuning steers generation toward semantic rewards without compromising realism.

Implications and Theoretical Advances

The main theoretical implication is that data posterior inference via structural flow models is tractable in high-dimensions provided that (1) the coupling between noise and data is learned via joint optimization, and (2) flow maps are regularized to maintain unconditional generative capacity. VFM leverages this to amortize complex inference, obviating expensive iterative guidance and enabling scalable, flexible Bayesian generation.

Practically, this implies that conditional generative modeling (including data-imaging inverse problems and reward-aligned generation such as RLHF) can be performed at negligible sampling cost, which is germane for computationally constrained deployments (e.g., scientific imaging, real-time systems, and rapid prototyping in design applications).

Connections to Literature and Future Directions

VFM conceptually subsumes prior attempts at noise-space posterior inference (e.g., [venkatraman2025outsourced]) by lifting all expressivity bottlenecks at the generator, rather than the adapter. It also aligns with recent variational and consistency approaches, but establishes that joint flow map/adaptor alignment is essential for accurate posterior recovery even with Gaussian adapters. Several extensions are immediate: more expressive (e.g., normalizing flow or transformer-based) adapters, multi-modal or multi-domain generation, application to non-linear and structured inverse problems, and generalization to video (learning temporal-coherent noise adapters).

Conclusion

Variational Flow Maps (2603.07276) provide a rigorous, efficient, and scalable solution to the problem of one-step conditional generation in flow-based models. The central mechanism—learning a noise adapter via joint variational training—unlocks the efficiency of flow maps for practical inverse problems and reward-conditioned generation, with strong empirical performance and theoretical underpinnings. This approach sets a new standard for coupling Bayesian inference and generative modeling, promising substantial impact in both applied and foundational directions.