Generative Models | Bohan Chen's Personal Webpage

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

Tue, 07 Jul 2026 00:00:00 +0000

Variational Flow Maps (VFM) recasts conditional generation as a problem of learning the right initial noise distribution for a pretrained or jointly trained one-step flow map. An observation-dependent adapter transforms simple noise before the flow map sends it to data space, enforcing the measurement while retaining the learned data prior.

A variational training objective aligns the adapted noise with the conditional target and supports both amortized inverse-problem solving and reward-based alignment. Because the conditioning mechanism acts before generation, VFM avoids tracing and guiding the long sampling trajectories required by diffusion and iterative flow models.

Across image inverse problems and ImageNet conditional-generation tasks, VFM produces diverse, well-calibrated samples in one or a few network evaluations. The results combine competitive fidelity with orders-of-magnitude faster sampling than iterative diffusion and flow baselines.

Flow Matching for Data Assimilation Accepted in SIAM/ASA JUQ

Mon, 15 Jun 2026 00:00:00 +0000

Table of Contents

I am happy to share that our paper, “Flow Matching for Efficient and Scalable Data Assimilation,” has been accepted for publication in the SIAM/ASA Journal on Uncertainty Quantification.

Paper: ·
Code:
Publication page and BibTeX:
Journal:

Why Flow Matching for Data Assimilation?

Data assimilation combines a dynamical model with noisy observations to estimate a system’s evolving state. Classical ensemble filters are efficient but can struggle with nonlinear or non-Gaussian distributions. More expressive generative approaches can address those settings, but methods such as the ensemble score filter require many sampling steps and can become expensive in large systems.

Our goal was to retain the expressiveness of a generative method while making the sampling process faster, training-free, and easier to scale.

The Ensemble Flow Filter

EnFF transports an ensemble through a guided flow. The F2P variant directly connects one filtering distribution to the next predictive distribution instead of repeatedly returning to a Gaussian reference.

We introduce the ensemble flow filter (EnFF), a flow-matching framework for sequential Bayesian data assimilation. It does not require an offline training stage. Instead, it constructs a Monte Carlo estimate of the marginal flow field from the current ensemble and incorporates each observation through localized guidance.

A central ingredient is the filtering-to-predictive (F2P) flow. Standard generative filters repeatedly transport samples from a Gaussian reference distribution. F2P instead starts from the previous filtering ensemble and transports it toward the next predictive distribution. This uses the structure already present in Bayesian filtering, producing a more direct path and more stable behavior when the number of sampling steps is small.

The framework also connects new generative filters with classical methods. Under appropriate flow choices and assumptions, EnFF recovers the bootstrap particle filter and the ensemble Kalman filter.

Accuracy, Efficiency, and Scale

We evaluate EnFF on Lorenz-63, Lorenz-96, the one-dimensional Kuramoto–Sivashinsky equation, and the two-dimensional Navier–Stokes equations, using both identity and nonlinear observation operators. The experiments range from compact chaotic systems to a Navier–Stokes pressure field on a 256×256 grid.

On the 256×256 Navier–Stokes benchmark, EnFF-F2P closely reconstructs the pressure field with only ten sampling steps, while EnSF shows a much larger error.

The results show that EnFF can match or improve the ensemble score filter while using far fewer sampling steps, with the advantage becoming more pronounced as the sampling budget decreases. EnFF-F2P is also less sensitive to its flow hyperparameters across different step counts.

Accuracy across increasing Kuramoto–Sivashinsky and Navier–Stokes dimensions under identity and nonlinear observations. EnFF variants remain competitive as the state dimension grows.

Across the high-dimensional benchmarks, EnFF offers a favorable balance of filtering accuracy, distributional quality, runtime, and stability. These results suggest that flow design—not only the choice of generative model—is a useful lever for building practical data-assimilation algorithms.

Collaboration and Reproducibility

This work was a rewarding collaboration with Taos Transue, So Takao, and Bao Wang. Taos Transue and I contributed equally. The accompanying contains the implementation and experiment code.

Figures on this page were extracted from the latest version of the paper.

Flow Matching for Efficient and Scalable Data Assimilation

Mon, 15 Jun 2026 00:00:00 +0000

The ensemble flow filter (EnFF) is a training-free data-assimilation framework that uses flow matching to transform a forecast ensemble into samples from the filtering distribution. Its Monte Carlo flow-field estimator and localized observation guidance avoid model training while retaining the flexibility of generative flow design.

The paper introduces a filtering-to-predictive (F2P) flow that uses the previous filtering distribution, rather than a standard Gaussian, as its reference. This path is better aligned with sequential Bayesian filtering and improves efficiency and robustness when only a small number of sampling steps is available. The analysis also shows how EnFF recovers the bootstrap particle filter and ensemble Kalman filter under appropriate choices and assumptions.

Experiments span Lorenz-63, Lorenz-96, the one-dimensional Kuramoto–Sivashinsky system, and two-dimensional Navier–Stokes equations, including state dimensions up to a 256×256 grid. Across these benchmarks, EnFF provides a strong accuracy–cost trade-off and scales to nonlinear, high-dimensional data-assimilation problems.

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

Tue, 10 Mar 2026 00:00:00 +0000

Table of Contents

Update: Our paper, “Variational Flow Maps: Make Some Noise for One-Step Conditional Generation,” has been accepted to the main track of ICML 2026 as a poster.

Paper: ·
Code and checkpoints:
Publication page and BibTeX:
Original announcement:

The Problem: Fast Generation Is Hard to Condition

Flow maps can turn noise into a high-quality image in one forward pass. That makes them dramatically faster than iterative diffusion models, but it also removes the trajectory normally used to inject measurements, constraints, or rewards.

For an inverse problem, we observe a degraded image—perhaps blurred, masked, or downsampled—and want samples that are both consistent with that observation and plausible under the learned image prior. Diffusion methods can repeatedly guide an evolving sample toward the observation. A one-step flow map has no intermediate states to steer: once its initial noise is fixed, its output is fixed too.

VFM addresses this “guidance gap” by changing the question. Instead of asking how to guide the sampling path, we ask:

What is the right noise distribution to start from?

The Key Idea: Learn the Right Noise

Given an observation, a noise adapter produces conditional noise samples; the jointly trained flow map transports them to conditional data samples in one step.

Given an observation or condition, VFM uses a lightweight noise adapter to predict a distribution in latent noise space. A sample from this distribution is then decoded by the flow map in a single forward pass.

The adapter and flow map are trained jointly with a variational objective. This joint training is essential: a simple Gaussian adapter may be too limited when paired with a frozen generator, but the flow map can reshape the noise-to-data coupling so that the same simple adapter represents a complex, multimodal posterior in data space. In a linear-Gaussian setting, our analysis shows that joint training recovers the exact posterior mean, whereas separate training is generically biased.

We also connect the data-reconstruction term in the variational objective to the mean-flow loss. That structural constraint keeps generated samples near the learned data manifold and avoids the off-manifold failure mode seen when the adapter and generator are optimized without a flow-map constraint.

One Model, Multiple Inverse Problems

The adapter can be conditioned on the inverse-problem class, allowing a single model to amortize inference across denoising, random and box inpainting, super-resolution, Gaussian deblurring, and motion deblurring.

ImageNet 256×256 box inpainting. The lower row shows distinct VFM posterior samples that remain consistent with the visible image.

On ImageNet 256×256 inverse problems, VFM combines distributional quality with a large speed advantage:

For box inpainting, one-step VFM reaches an FID of 33.34, compared with 62.35–75.62 for the iterative guidance baselines reported in the paper.
VFM uses 1 neural function evaluation, while those baselines use 250×2 evaluations with classifier-free guidance.
A single VFM sample takes about 0.025–0.027 seconds in our experiments, versus roughly 7–47 seconds for the iterative methods.
VFM performs strongly on distributional and calibration metrics including FID, MMD, CRPS, and LPIPS.

There is an important evaluation lesson here. For posterior sampling, diversity and calibration matter alongside pixel-wise fidelity. PSNR and SSIM often reward posterior-mean-like, smoother reconstructions, so they do not by themselves capture whether a method produces the right distribution of plausible solutions. Averaging several VFM draws improves these pixel-wise metrics, while the individual draws expose meaningful posterior uncertainty.

Beyond Inverse Problems: One-Step Reward Alignment

The same idea extends naturally from observations to general differentiable rewards. VFM learns a context-dependent noise adapter while fine-tuning the flow map toward a reward-tilted distribution. In our experiments, strong reward alignment emerges in less than half an epoch of fine-tuning, and the resulting model still generates with a single forward pass. The uncurated reward-aligned samples shown in this post’s header were generated by that fine-tuned flow map.

Code and Reproducibility

The includes the training and sampling code, released checkpoints, scripts for the ImageNet inverse problems, and a small checkerboard example that illustrates how joint training changes the latent coupling. The conditional sampler supports one-step generation as well as two-to-four-step refinement.

This work was a wonderful collaboration among Abbas Mammadov, So Takao, Bohan Chen, Ricardo Baptista, Morteza Mardani, Yee Whye Teh, and Julius Berner, spanning the University of Oxford, Caltech, PhysicsX, the University of Toronto, and NVIDIA. Abbas Mammadov and So Takao contributed equally.

Figures on this page were extracted from the paper, which is available under a CC BY 4.0 license.