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

Abbas Mammadov; So Takao; Bohan Chen; Ricardo Baptista; Morteza Mardani; Yee Whye Teh; Julius Berner

cs.CVMar 7, 2026

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

Authors: Abbas Mammadov, So Takao, Bohan Chen, Ricardo Baptista, Morteza Mardani, Yee Whye Teh, Julius Berner

Organizations: Department of Statistics, University of Oxford · California Institute of Technology · University of Toronto · NVIDIA

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

Figures & tables

Figure 1: One-step conditional generation with Variational Flow Maps (VFM). Given an observation $y$ , VFM learns a noise adapter network $q_{\phi}(z|y)$ , which approximates the noise space posterior $p(z|y)$ via amortized variational inference. Conditional noise samples $z\sim q_{\phi}(z|y)$ are then mapped to data space in a single step via a learned flow map $x=f_{\theta}(z)$ , producing conditional samples that approximate $p(x|y)$ . In VFM, the networks $q_{\phi}$ and $f_{\theta}$ are trained jointly by extending the variational autoencoder framework to learn the correspondence between the triple $(x,y,z)$ . By jointly training, $f_{\theta}$ learns to compensate for the simple Gaussian assumption on $q_{\phi}$ .

Panel of Figure 2 — Figure 2: Prior 2D samples and posterior densities in data space (top row) and noise space (bottom row). We observe the $x$ -component (black dashed lines) with $\sigma=0.1$ . The unconditional samples are color-coded by checkerboard cell; light grey for off-manifold samples. VFM successfully captures the bimodal nature of the posterior, while the baselines struggle to do so.

Figure 3: Qualitative comparison on ImageNet 256 $\times$ 256 box inpainting. Top row: ground truth, measurement, and reconstructions from guidance-based baselines. Bottom row: conditional samples produced by VFM, showing diversity in the inpainted region.

Task	Method	NFE	PSNR ( $\uparrow$ )	SSIM ( $\uparrow$ )	LPIPS ( $\downarrow$ )	FID ( $\downarrow$ )	MMD ( $\downarrow$ )	CRPS DINO ( $\downarrow$ )	CRPS Inc ( $\downarrow$ )	Time (s) ( $\downarrow$ )
Inpaint (box)	Latent DPS	250 $\times$ 2	22.80	0.704	0.349	62.89	0.132	0.511	0.389	7.223
	Latent DAPS	250 $\times$ 2	23.98	0.707	0.348	–	–	0.468	0.365	43.93
	PSLD	250 $\times$ 2	22.61	0.699	0.346	67.22	0.153	0.536	0.435	10.07
	MPGD	250 $\times$ 2	22.76	0.705	0.350	62.35	0.132	0.510	0.388	7.487
	FlowChef	250 $\times$ 2	22.80	0.704	0.349	63.20	0.133	0.512	0.389	7.612
	FlowDPS	250 $\times$ 2	23.21	0.706	0.364	75.62	0.166	0.606	0.482	14.47

Table 1: Quantitative comparison on ImageNet for box inpainting and Gaussian deblurring. Best results are in bold , second best are underlined .

\uparrow

: higher is better,

\downarrow

: lower is better. For VFM, we display results for single samples and average over 10 samples, displayed as {sample} / {average} . VFM achieves the best results on LPIPS, FID, MMD, CRPS, with a significantly reduced wall-clock time.

Figure 4: Unconditional generation on ImageNet $256\times 256$ . Left: unconditional samples from VFM-B/2. Right: unconditional FID comparison versus mean-flow baselines. VFM retains competitive performance despite it being trained for posterior sampling.

Figure 5: One-step reward-aligned generation using VFM fine-tuning. Starting from a pre-trained ImageNet flow map, VFM efficiently adapts the latent noise space and flow trajectories to sample from a reward-tilted distribution, achieving strong visual alignment with a target reward $R(x,c)$ in a single forward pass while preserving image quality.

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 8 — Figure 8: Ablation of VFM with respect to key modeling choices in the loss. Observation in black dots and $\sigma=0.1$ . For VFM ( 8(c) , 8(d) , 8(e) ), we used $\tau=100.0$ , $\alpha=1.0$ and $K=4$ . We observe that 8(a) : frozen- $\theta$ fails to capture the bimodal nature of the posterior; 8(b) : unconstrained- $\theta$ produces many off-manifold samples; 8(c) : removing EMA from the term $\mathcal{L}_{\text{obs}}$ in VFM also produces many off-manifold samples when $\tau$ is large; 8(d) : removing the KL term $\mathcal{L}_{\text{KL}}$ in the VFM loss leads to unstable optimization and results in poor approximations of both the prior and posterior.

Panel of Figure 9 — Figure 9: Ablation of VFM with respect to the $\tau$ parameter. For each plot, we set $\alpha=1.0$ and $K=4$ . We fix $y=0.5$ and $\sigma=0.1$ . We observe that for $\tau\lesssim\sigma$ , the quality of prior/posterior approximations are poor, yielding many off-manifold samples. This is likely due to the difficulty of optimization as we tighten the correspondence between $x$ and $z$ . For $\tau\geq 1$ , we observe significant improvements in results and surprising robustness with respect to large values of $\tau$ .

Panel of Figure 10 — Figure 10: Ablation of VFM with respect to the $\alpha$ parameter. For each plot, we set $\tau=100.0$ and $K=4$ . We observe that the warping of the latent space becomes stronger as $\alpha\rightarrow 1$ , making it easier to sample from the bimodal posterior using the simple Gaussian variational posterior in latent space.

Task	Method	NFE	PSNR ( $\uparrow$ )	SSIM ( $\uparrow$ )	LPIPS ( $\downarrow$ )	FID ( $\downarrow$ )	MMD ( $\downarrow$ )	CRPS DINO ( $\downarrow$ )	CRPS Inc ( $\downarrow$ )	Time (s) ( $\downarrow$ )
Inpaint (random)	Latent DPS	250 $\times$ 2	26.01	0.721	0.337	55.81	0.113	0.472	0.363	7.2164
	Latent DAPS	250 $\times$ 2	25.09	0.671	0.384	–	–	0.474	0.356	44.347
	PSLD	250 $\times$ 2	25.63	0.713	0.338	56.13	0.123	0.462	0.386	10.286
	MPGD	250 $\times$ 2	26.03	0.720	0.339	55.82	0.112	0.470	0.363	7.3512
	FlowChef	250 $\times$ 2	26.01	0.720	0.338	55.73	0.111	0.471	0.364	7.3885
	FlowDPS	250 $\times$ 2	25.80	0.729	0.344	62.62	0.139	0.557	0.453	14.054

Appendix

Table 2: Quantitative comparison on ImageNet for various inverse problems. Best results are in bold , second best are underlined .

\uparrow

: higher is better,

\downarrow

: lower is better.

Figure 11: Quantitative evaluation of generated samples over 10,000 training iterations for varying values of

\lambda

. We report HPSv2 (left), PickScore (middle), and ImageReward (right). Higher scores indicate better alignment with human preferences.

Figure 12: Qualitative comparison on Random Inpainting. We compare one-step VFM samples against seven baselines. VFM recovers fine details and texture consistent with the unmasked regions, while maintaining high perceptual quality.

Figure 13: Qualitative comparison on Box Inpainting. Comparison of VFM against baselines for large occlusions. VFM generates plausible semantic content to fill the missing regions in a single step.

Figure 14: Qualitative comparison on Super-Resolution ( $\times 4$ ). VFM effectively upsamples the low-resolution inputs, leading to sharp edges and realistic textures compared to the often over-smoothed baseline results.

Figure 15: Qualitative comparison on Gaussian Deblurring. VFM successfully restores sharpness from heavily blurred observations ( $\sigma=3.0$ ), and it also avoids the artifacts common in guidance-based methods.

Figure 16: Qualitative comparison on Motion Deblurring. Comparison of deblurring performance on motion-blurred inputs. VFM resolves the motion streaks into coherent structures.

Figure 17: Visualizing the Learned Noise Space. We visualize the outputs of the noise adapter $q_{\phi}(z|y)$ . From left to right: ground truth, measurement, the predicted latent mean $\mu$ , standard deviation $\sigma$ , and three independent latent samples drawn from the distribution. The visible structure in the “noise” confirms that the adapter optimizes the latent initialization to align with the conditional data manifold. From top to bottom, rows correspond to: random inpainting, box inpainting, super-resolution, gaussian deblurring, and motion deblurring.

Figure 18: Posterior Uncertainty Quantification. We display the pixel-wise mean and standard deviation computed from 10 conditional samples generated by VFM. The standard deviation maps (right column) accurately capture the uncertainty inherent in the inverse problem, which highlights ambiguous regions where the model generates diverse solutions.

Figure 19: Unconditional Samples. Curated unconditional samples generated by the VFM.

Figure 20: Posterior Diversity (Sample Set 1). Evaluation of sample diversity on gaussian deblurring. While baselines often collapse to a single mode or fail to produce valid results, VFM generates eight distinct, plausible, and measurement-consistent posterior samples.

Figure 21: Posterior Diversity (Sample Set 2). Additional examples of diverse posterior sampling on box inpainting task. The high variance among the VFM samples reflects the multimodal nature of the posterior distribution for these ill-posed tasks.

Figure 22: Uncurated Samples from Reward Fine-Tuning ( $\lambda=1$ ). Additional uncurated samples generated by the fine-tuned flow map. VFM consistently samples high-quality images from the target reward-tilted distribution, resulting in enhanced aesthetic and perceptual quality in a single forward pass.

Figure 23: Evolution of One-Step Generations during Reward Fine-Tuning. We visualize the progress of generated samples across varying fine-tuning iterations using fixed latent seeds, all produced in a single neural function evaluation (1 NFE). These seeds are drawn from a pure standard Gaussian distribution. Although we tilt the original flow map to accommodate adapter-conditioned noises, the original noise space remains valid. As a result, this short fine-tuning process successfully enhances the one-step generative capabilities of the base DMF model even without the use of an adapter.

Conditional Generative Modeling

Conditional generative modeling learns to produce data distributions steered by labels, inputs, text, or physical and structural constraints.

See all 393 papers

Local Flow Matching Generative Models
Oct 3, 2024Chen Xu, Xiuyuan Cheng, Yao XieFlow Matching Generative Modeling
H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology. · Department of Mathematics, Duke University
Noise-Robust Conditional Flow Matching: Generating Clean Samples from Noisy Datasets
Jul 28, 2026Adrian Urbański, Gabriel della Maggiora, Artur YakimovichFlow Matching Image Generation
Center for Advanced Systems Understanding (CASUS), Görlitz, Germany · Helmholtz-Zentrum Dresden-Rossendorf e. V. (HZDR), Dresden, Germany · Institute of Computer Science, University of Wrocław, Wrocław, Poland +2
SubFlow: Sub-mode Conditioned Flow Matching for Diverse One-Step Generation
Apr 14, 2026Yexiong Lin, Jia Shi, Shanshan Ye +3Flow Matching Flow-Based Generative Modeling
Sydney AI Centre, The University of Sydney, Sydney, NSW, Australia · School of Artificial Intelligence, Xidian University, Xi’an, Shaanxi, China · Australian Artificial Intelligence Institute, University of Technology Sydney, Sydney, NSW, Australia +1

Oct 3, 2024J/K move · Enter open · S save

Local Flow Matching Generative Models

Chen Xu, Xiuyuan Cheng, Yao Xie

H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology. · Department of Mathematics, Duke University

Flow Matching Generative Modeling Continuous Normalizing Flows Flow-Based Generative Modeling

Open paper AskPDF arXiv

Flow Matching (FM) is a simulation-free method for learning a continuous, invertible flow that interpolates between two distributions, and in particular generates data from noise. Inspired by the variational nature of the diffusion process as a gradient flow, we introduce a stepwise FM model, Local Flow Matching (LFM), which sequentially learns a sequence of FM submodels, each matching a diffusion process up to the time-step size in the data-to-noise direction. In each step, the two distributions to be interpolated by the sub-flow model are closer than those in the full-flow matching model, which interpolates data to noise distributions, enabling smaller models with more efficient training. This variational perspective also allows us to prove a theoretical generation guarantee for the proposed flow model in terms of the

χ^2

-divergence between the generated and true data distributions, leveraging the contraction property of the diffusion process. In practice, the stepwise structure of LFM is naturally amenable to model distillation, and various distillation techniques can be applied to accelerate generation. We empirically demonstrate that LFM achieves competitive generative performance compared to FM on unconditional generation of tabular and image datasets, and on conditional generation of robotic manipulation policies.

Abstract

Conditional Generative Modeling

Explore similar work

Local Flow Matching Generative Models

Noise-Robust Conditional Flow Matching: Generating Clean Samples from Noisy Datasets

SubFlow: Sub-mode Conditioned Flow Matching for Diverse One-Step Generation