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 387 papers

Oct 3, 2024

Local Flow Matching Generative Models

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.

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

Ask PDF

Jul 28, 2026

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

Generative models learn the statistical properties of their training data, so high-quality generation depends on clean and representative datasets. In scientific imaging, acquisition often yields noisy measurements, while collecting clean references can be costly, impractical or even unattainable. Training directly on these measurements results in a model that reproduces the corrupted data. This can be circumvented by learning the clean population distribution directly from the noisy data. Conditional flow matching (CFM) combines a simple regression objective with stable training, efficient sampling, and strong image-generation performance, making it a natural framework for this setting. We introduce Noise-Robust Conditional Flow Matching (NR-CFM), an unconditional generator that learns from one corrupted observation per image. NR-CFM provides a closed-form clean endpoint correction for additive white Gaussian noise and learns a data-driven correction for general Gaussian corruptions with more complex covariance structure. Across the evaluated corruption settings, NR-CFM outperforms NR-GAN in most cases and remains competitive with Ambient Diffusion in the high-noise regime. We further evaluate NR-CFM on scientific data at signal-to-noise ratios as low as

0.001

, where it generates plausible particle images from severely corrupted measurements.

Adrian Urbański, Gabriel della Maggiora, Artur Yakimovich

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

Flow Matching Image Generation Generative Modeling Flow-Based Generative Modeling

Ask PDF

Apr 14, 2026

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

Flow matching has emerged as a powerful generative framework, with recent few-step methods achieving remarkable inference acceleration. However, we identify a critical yet overlooked limitation: these models suffer from severe diversity degradation, concentrating samples on dominant modes while neglecting rare but valid variations of the target distribution. We trace this degradation to averaging distortion: when trained with MSE objectives, class-conditional flows learn a frequency-weighted mean over intra-class sub-modes, causing the model to over-represent high-density modes while systematically neglecting low-density ones. To address this, we propose SubFlow, Sub-mode Conditioned Flow Matching, which eliminates averaging distortion by decomposing each class into fine-grained sub-modes via semantic clustering and conditioning the flow on sub-mode indices. Each conditioned sub-distribution is approximately unimodal, so the learned flow accurately targets individual modes with no averaging distortion, restoring full mode coverage in a single inference step. Crucially, SubFlow is entirely plug-and-play: it integrates seamlessly into existing one-step models such as MeanFlow and Shortcut Models without any architectural modifications. Extensive experiments on ImageNet-256 demonstrate that SubFlow yields substantial gains in generation diversity (Recall) while maintaining competitive image quality (FID), confirming its broad applicability across different one-step generation frameworks. Project page: https://yexionglin.github.io/subflow.

Yexiong Lin, Jia Shi, Shanshan Ye +3

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

Flow Matching Flow-Based Generative Modeling One-Step Generative Modeling Conditional Generative Modeling Mode Collapse

Ask PDF

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