PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2603.07276v1cs.CVcs.LGstat.ML

TL;DR

Iterative generative models can condition on external observations, but flow maps lack a trajectory for incorporating such constraints. Variational Flow Maps learn observation-dependent initial noise and jointly train a noise adapter with the flow map, enabling well-calibrated conditional sampling in one or few steps. The framework also supports reward alignment and amortization across inverse-problem settings, while relying on a simple Gaussian adapter assumption.

  • Problem

    Flow maps efficiently generate samples but lack an explicit sampling trajectory for incorporating external constraints in conditional generation and inverse problems.

  • Method

    VFMs jointly train an observation-dependent noise adapter and flow map using a principled variational objective that aligns noise, data, and observations.

  • Results

    VFMs produce well-calibrated conditional samples in one or few steps and extend to reward-tilted sampling through fast flow-map fine-tuning.

  • Takeaways & Limitations

    Conditioning can be performed by learning appropriate initial noise rather than steering an iterative sampling trajectory.

  • Takeaways & Limitations

    The framework currently assumes a Gaussian adapter, motivating more expressive noise models for richer non-Gaussian conditional structures.

Abstract

from arXiv · show

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

1. Introduction

Flow maps offer efficient generation but lack a trajectory for incorporating observations during conditional sampling. Variational Flow Maps address this guidance gap by learning observation-dependent noise and jointly training the adapter with the flow map.

  • Motivation: Diffusion and flow-based methods often require dozens to hundreds of sequential function evaluations, creating high computational cost for real-time applications.This sampling-cost limitation motivates methods that generate samples in one or few steps.
  • The guidance gap: Flow maps lack intermediate states or a trajectory to steer, limiting their use for conditional generation and inverse problems.Once noise z is selected, the deterministic mapping to x leaves no opportunity to incorporate measurement information during generation.
  • Variational Flow Maps: VFMs learn an observation-dependent noise sampler so conditional generation can use one or few flow-map steps.The method reformulates conditioning as finding noise z whose mapped data x is consistent with the observation.
  • Joint training: Jointly training the noise adapter and flow map aligns x, y, and z while allowing the flow map to compensate for limited adapter expressivity.The framework extends the variational autoencoder perspective to learn the correspondence among all three variables.
  • Extensions: VFMs also extend to reward alignment by fine-tuning pretrained flow maps to sample reward-tilted distributions in a single step.The paper presents this as a fast and scalable extension of the framework.

2. Background

Flow-based models transport noise to data through learned ODE dynamics, while flow maps learn the transport solution operator for arbitrary-step generation. Bayesian inverse problems and amortized variational inference provide the conditioning and posterior-approximation foundations for VFMs.

  • Flow-based generative models: Flow-based generative models transport samples from a Gaussian prior to the data distribution through an ODE.The learned dynamics connect p1(z) = N(0, I) to pdata(x).
  • Flow maps: Flow maps learn the ODE solution operator directly, enabling one-step generation or composed multi-step refinement after training.A single evaluation fθ(x1, 0, 1) produces a one-step sample, while intermediate evaluations support multi-step generation.
  • Inverse problems: Inverse problems seek signals consistent with noisy observations and a prior distribution, often expressed using a known forward operator and noise level.The observation model uses a forward operator A and σ > 0, while Bayesian inference targets the corresponding posterior.
  • Inverse problems: Guidance-based posterior sampling incorporates likelihood gradients at each denoising step but requires iterative refinement and cannot directly apply to one-step flow maps.This contrast motivates conditioning through noise-space inference rather than trajectory guidance.
  • Variational inference: Amortized inference uses a neural network to predict a variational distribution from the conditioning variable instead of optimizing separately for each instance.With diagonal Gaussian variational families, the network predicts distribution parameters such as μϕ(x) and σϕ(x).

3. Variational Flow Maps (VFMs)

Variational Flow Maps reformulate conditional generation in noise space and jointly train a noise adapter with the flow map. The resulting objective couples observations, data, and noise while supporting single- or few-step posterior sampling.

  • Noise-space reformulation: VFMs approximate inverse-problem posteriors by learning an observation-dependent noise distribution and mapping sampled noise through the flow map.The method replaces trajectory guidance with learning which initial noise values produce data consistent with the observation.
  • Noise-space reformulation: A variational posterior qϕ(z|y) makes inference tractable because the noise prior is simple, commonly N(0, I).The formulation uses a Gaussian decoder likelihood and a tractable KL term for a conjugate Gaussian noise posterior.
  • Joint training: Jointly training fθ and qϕ reshapes the noise-to-data correspondence, compensating for limited variational-posterior expressivity and adding structural flow constraints.The joint strategy aligns x, y, and z through a modified objective rather than holding the pretrained flow map fixed.
  • Joint training: Joint training recovers the true posterior mean exactly in the stated Gaussian setting, whereas separate training almost surely fails to match it.The result concerns the estimator E_z∼qϕ(z|y)[fθ(z)] versus E_p(x|y)[x].
  • Objective and sampling: The objective combines mean-flow, observation, and KL terms, with the mean-flow term linked to reconstruction error and structural flow-map learning.The loss is Lθ,ϕ := 1/(2τ^2) LMF(θ; ϕ) + 1/(2σ^2) Lobs(θ, ϕ) + LKL(ϕ).
  • Objective and sampling: High-quality conditional samples require only a small number of steps, substantially fewer than full generative ODE or SDE integration.Multi-step sampling starts from adapter-generated noise and iteratively applies the mean-flow model over a short partition.

4. Experiments

Experiments show that joint VFM training captures complex conditional posteriors while preserving data-manifold structure, achieving strong ImageNet inverse-problem quality with substantially lower inference cost. VFM also retains competitive unconditional generation and supports fast reward alignment.

  • 2D Example: VFM captures both modes of the toy posterior while preserving the checkerboard support, unlike frozen-θ and unconstrained-θ baselines.Removing LKL makes training unstable because the inverse problem lacks prior regularization.
  • 2D Example: For τ ≥1, CRPS and posterior MMD improve substantially, while SACC and prior MMD approach the pretrained flow's values.α < 1 is more stable and gives better prior fit, whereas α = 1 gives the best posterior fit.
  • Image Inverse Problems: Guidance methods score higher on PSNR and SSIM for single draws, but averaged VFM samples narrow the gap and can exceed baselines on Gaussian deblurring.VFM is competitive without averaging on LPIPS, a perceptual feature-space metric.
  • Image Inverse Problems: VFM requires one sampling step and achieves around two orders of magnitude lower wall-clock time than iterative guidance methods.The reported example compares approximately 0.03 seconds for VFM with nearly one minute for DAPS.
  • Unconditional Generation: VFM remains competitive in unconditional ImageNet generation despite being trained for posterior sampling.The comparison uses 50,000 unconditional samples and reports results against similarly sized baselines.
  • Reward Alignment: VFM fine-tuning adapts a pretrained flow map for arbitrary differentiable rewards, achieving strong alignment in under 0.5 epochs and one neural function evaluation.The resulting model samples from the reward-tilted distribution in a single step.

5. Related Works

Related work addresses variational inference, noise-space posterior inference, consistency training, and one-step conditional sampling through different model and coupling choices. VFM instead uses a simpler adapter with a trainable flow map for conditional posterior sampling.

  • Variational Inference: Prior variational methods use diffusion-based priors with flexible normalizing-flow posteriors, making high-resolution scaling difficult.These methods approximate posteriors in data space or amortize inference over observations.
  • Noise-Space Inference: Noise-space posterior inference with a frozen generator uses neural-SDE adapters, whereas VFM unfreezes the flow map and keeps the adapter simpler.The approaches therefore differ in where model flexibility is introduced.
  • Consistency Methods: Variational Consistency Training learns data-dependent noise couplings for stable unconditional consistency training, while VFM targets amortized conditional posterior sampling.The two methods are conceptually related but pursue different goals.
  • Consistency Methods: Noise Consistency Training performs one-step conditional sampling in a diffusion process over (z, y)-space, rather than directly generalizing to flow maps.Its construction is tied to consistency models.

6. Conclusion

VFM enables efficient posterior sampling and reward fine-tuning in one or a few steps by jointly training a flow map with an amortized noise adapter.

  • Conclusion: VFM infers observation-, class-label-, or prompt-dependent initial noise while jointly adapting the flow map and noise adapter.The stated future directions include more expressive non-Gaussian noise models and extensions to other modalities.

Impact Statement

Reducing inference cost could support fast posterior sampling in scientific and engineering workflows and reduce inference energy costs. The paper frames VFM as a step toward low-cost conditional sampling without sacrificing performance.

  • Impact: Lower-cost conditional generation could benefit scientific and engineering workflows that rely on fast posterior samples.The passage also identifies reduced inference energy cost as a potential benefit.
  • Impact: The VFM objective is derived by matching two representations of p(x, y, z) with KL divergence, yielding bounds involving data, observation, and KL losses.The derivation uses the nonnegative joint-distribution entropy term to obtain the stated upper bound.

A.2. Proof of Proposition 3.1

The linear-Gaussian analysis shows why joint training of the generative map and variational posterior is necessary under diagonal covariance constraints: joint training can recover the exact posterior mean, whereas separate training generally cannot.

  • Setup: The analysis studies separate and joint optimization of a linear generative map and Gaussian noise adapter under diagonal covariance constraints.The generative map is affine, while the adapter is parameterized by observation-dependent Gaussian mean and covariance.
  • Optimal solutions: The optimal generative parameters form a rotational family Kθ = UΛQ and bθ = m for orthogonal Q.This rotational freedom creates multiple equivalent latent coordinate systems.
  • Mean recovery: Joint training recovers the exact Bayesian posterior mean, while separate training misses it almost surely for observations drawn from the model.For joint solutions, the composed mean equals m + K(y − Am), the exact posterior mean.
  • Separate training: Separate training almost surely fails under general-position assumptions because its fixed latent coordinates do not diagonalize the required precision structure.The resulting constrained optimum has a strictly larger loss than the unconstrained optimum.
  • Joint training: Joint training preserves the global optimum while adapting Q so the optimal variational covariance satisfies the diagonal constraint.The learnable rotation can diagonalize the relevant matrix, allowing constrained and unconstrained optima to coincide.
  • Conclusion: The proof formalizes that diagonal constraints create a mean-recovery gap for separate training but not for appropriately jointly trained solutions.This result is established under general-position assumptions on the observation operator and data covariance.

B.1. 2D Checkerboard Data

The checkerboard experiment evaluates conditional generation from noisy partial observations using a synthetic multimodal prior and posterior-sample quality metrics.

  • Data: The synthetic prior is a 2D checkerboard supported on alternating squares in [−2, 2]2.Samples are generated from a uniform grid construction and 20,000 retained samples are used for training.
  • Observation model: The inverse problem observes only the first component with Gaussian observation noise σ = 0.1.The observation is generated through a linear forward operator.
  • Metrics: NLPD evaluates predictive consistency with observations, while CRPS evaluates uncertainty around the ground-truth target.CRPS combines sample accuracy with a diversity term.
  • Metrics: MMD measures distance between true and approximate distributions, and SACC measures the proportion of samples lying on checkerboard support.Both prior and posterior support accuracy are evaluated.
  • Sampling: VFM conditional samples are obtained by drawing z from qϕ(z|y) and mapping each draw through fθ.The same flow map is also used to generate approximate prior samples from standard Gaussian noise.

B.1.4. ABLATION PLOTS

The ablations show that VFM’s conditional sample quality depends on the correspondence-relaxation parameter and loss design, while ImageNet comparisons require extensive tuning of iterative baselines.

  • Loss ablation: Frozen-θ fails to capture posterior bimodality, whereas unconstrained-θ produces many off-manifold samples.VFM captures both modes while preserving the checkerboard pattern.
  • Loss ablation: Removing EMA from Lobs produces many off-manifold samples when τ is large, while removing LKL causes unstable optimization and poor results.These variants isolate the roles of EMA and KL regularization in VFM.
  • τ ablation: For τ ≲ σ, prior and posterior approximations are poor and produce many off-manifold samples.The authors attribute this behavior to optimization difficulty when correspondence between x and z is tightened.
  • τ ablation: For τ ≥ 1, results improve substantially, especially for CRPS and posterior MMD, while support and prior-MMD metrics approach frozen-θ values.The improvement is reported for both single-step and four-step samplers.
  • α ablation: Increasing α toward 1 strengthens latent-space warping and makes bimodal posterior sampling easier for the Gaussian variational posterior.This ablation fixes τ = 100 and K = 4.
  • ImageNet setup: ImageNet comparisons use a 130M-parameter SiT-B/2 flow backbone and a 10M-parameter U-Net-style noise adapter operating in SD-VAE latent space.The adapter is conditioned on inverse-problem class using FiLM.
  • Baseline comparison: Iterative guidance baselines require exhaustive per-task hyperparameter sweeps and are computationally burdensome relative to VFM’s one-step sampling.Unless otherwise stated, baselines use 250 ODE steps.

B.2.3. METRICS AND EVALUATION

The evaluation uses both pixel-space and distributional metrics to assess reconstruction fidelity and posterior quality. VFM is reported to outperform baselines on distributional measures while averaging multiple samples can improve pixel-space scores.

  • Evaluation covers PSNR, SSIM, LPIPS, FID, MMD, and CRPS across pixel-space and semantic or distributional fidelity categories.Distributional metrics are prioritized because pixel-wise scores can favor conditional-mean estimates.
  • Inference-time optimization methods can achieve high PSNR and SSIM by producing smooth conditional-mean estimates that lose texture and realistic detail.
  • The evaluation uses 1,024 ImageNet validation reconstructions for most metrics and 10 reconstructions for each of 128 samples for CRPS.Latent DAPS uses 128 rather than 1,024 samples because of slower generation.
  • VFM significantly outperforms baselines on FID, MMD, and CRPS, indicating better perceptual quality and posterior approximation.
  • Averaging multiple one-step VFM samples produces competitive or better PSNR and SSIM values while smoothing reconstructions.

B.2.4. INVERSE PROBLEMS AND EVALUATION SETUP

VFM is evaluated on diverse linear inverse problems using a shared pretrained SiT-B/2 backbone for fair sampling-method comparisons. The benchmark emphasizes structurally challenging tasks and specifies standardized degradation operators.

  • The jointly trained VFM handles denoising, random inpainting, box inpainting, super-resolution, Gaussian deblurring, and motion deblurring.
  • Quantitative evaluation focuses on inpainting, super-resolution, and deblurring, omitting pure denoising.
  • All baselines use the same pretrained SiT-B/2 backbone as VFM, isolating differences to iterative guidance versus one-step VFM sampling.The setup is intended to make VFM results a reference for SiT-B/2 benchmarking.
  • Table 2 compares ImageNet inverse-problem results, marking best and second-best values and indicating whether higher or lower scores are preferable.
  • Random inpainting masks 30–70% of each image, while box inpainting uses random rectangular masks with independently sampled dimensions from 32 to 128.

B.3. General Reward Alignment with VFM

VFM extends conditional flow-map sampling to reward alignment by treating reward maximization as an inverse problem. Its variational fine-tuning objective improves reward scores, while one-step samples retain stronger alignment than multi-step samples.

  • Reward alignment targets a reward-tilted distribution defined from a base data distribution and a differentiable context-conditioned reward model.
  • The reward objective replaces the inverse-problem observation likelihood with λR(fθ(z), c), treating context as the observation.
  • The variational fine-tuning objective combines expected reward, a KL regularizer, and a data loss, and is derived from the ELBO.At its global optimum, the formulation recovers the true reward-tilted distribution according to the paper.
  • One-step samples achieve the highest reward scores, whereas multi-step samples regress toward the unconditional ImageNet distribution.A short K-step rollout during training is proposed as future work to propagate reward signals into intermediate fields.
  • A default reward strength of λ = 1 achieves strong performance without extensive tuning, and VFM consistently boosts HPSv2, PickScore, and ImageReward.Scores are evaluated every 500 steps and averaged over 64 random generations.

B.4. Additional Results

Additional experiments show that VFM produces sharp single-pass inverse-problem reconstructions, structured conditional noise, uncertainty localized to ambiguous regions, diverse posterior samples, and reward-aligned unconditional generations.

  • Qualitative Comparisons: Across five inverse problems, VFM produces sharp, coherent, measurement-consistent samples in one forward pass, while baselines may require hundreds of function evaluations.
  • Uncertainty Quantification: VFM’s posterior standard-deviation maps localize uncertainty in ambiguous regions such as occlusions and details lost to blur.The maps are computed from 10 posterior samples.
  • Structured Noise: The noise adapter learns structured latent means, standard deviations, and samples that align the latent space with the conditional data manifold.
  • Diversity and Mode Coverage: VFM generates diverse, high-quality samples consistent with measurements, with greater diversity emerging as inverse problems become more ill-posed.The results contrast with baselines that frequently fail or collapse to a single solution.
  • Unconditional Generation: The trained flow map produces curated unconditional samples and reward-aligned images that adapt rapidly and stably across fine-tuning iterations.The reward-aligned generations are produced in a single forward pass using fixed latent seeds.
Loading 2603.07276v1…