Preconditioned Flow Matching

0:00 / 1:30

The idea in 90 seconds

A reversible map P moves the data to a representation where the flow's intermediate distributions are well conditioned; the flow is trained there, and samples are mapped back with P⁻¹. Both a single anisotropic Gaussian and a six-mode mixture are carried through the same four steps, and compared with standard FM at an identical budget. Rendered from real runs of the released code (three seeds, fixed budget).

Abstract

Geometric optimization, not model expressiveness, is the bottleneck

Flow matching learns vector fields by regressing stochastic velocity targets along intermediate distributions $p_t$. We identify a geometric optimization bottleneck in this regression problem: when the covariance $\Sigma_t$ of $p_t$ is ill-conditioned, gradient-based training rapidly fits high-variance directions while making slow progress along low-variance ones.

In an exactly solvable Gaussian setting, we prove that the excess risk is weighted by $\Sigma_t$, and that both gradient descent and stochastic gradient descent inherit condition-number-dependent convergence. We then extend the analysis to Gaussian mixtures, showing that multimodality does not average away this effect; instead, the slowest and worst-conditioned component can control optimization.

Motivated by this analysis, we propose preconditioned flow matching, a precondition-then-match framework that transforms the target distribution into a more isotropic representation, trains the main flow in the transformed space, and maps generated samples back through the inverse transformation. We show theoretically that preconditioning reshapes the intermediate FM path and improves its conditioning. Across controlled Gaussian and Gaussian-mixture experiments, latent MNIST and other high-resolution image datasets up to 512×512, preconditioning improves path-conditioning diagnostics, low-eigenvalue recovery, FID, MMD, precision and recall. Compute-matched baselines and preconditioner-quality ablations further show that the gains are not explained merely by additional preconditioner parameters, but by improved geometry of the downstream flow matching problem.

Three observations

Why a small loss can hide a bad model

Excess risk is weighted by $\Sigma_t$

For linear predictors the population risk splits exactly into an irreducible floor and a quadratic form in the covariance of the intermediate distribution.

$$L_t(A) - L_t^{\star} = \mathrm{Tr}\!\left[(A - A^{\star})\,\Sigma_t\,(A - A^{\star})^{\top}\right]$$

Where $\sigma_i(t)$ is small, a loss below $\varepsilon$ still permits a velocity error up to $\sqrt{\varepsilon/\sigma_i(t)}$.

Convergence scales with $\kappa(\Sigma_t)$

Gradient descent on the same quadratic converges linearly, but the iteration count to reduce the error by ε grows with the condition number of the intermediate covariance.

$$k = O\!\left(\kappa(\Sigma_t)\cdot\log(1/\varepsilon)\right)$$

At a fixed $t$ the conditioning is

$$\kappa(\Sigma_t) = \frac{(1-t)^2 + t^2\lambda_{\max}}{(1-t)^2 + t^2\lambda_{\min}}$$

with $\lambda$ the eigenvalues of the target covariance.

Scalar time reweighting rescales the loss but cannot remove this directional stiffness.

Mixtures do not average it away

With component-wise dynamics, the whole mixture is governed by the slowest eigendirection across all components — one badly conditioned mode can set the time scale for the entire model.

$$\mathrm{rate} \;\propto\; \min_{k,i}\,\sigma_{k,i}(t)$$

Anisotropy gets amplified rather than diluted by multimodality.

Method

Precondition, match, map back

Preconditioning does not change the generative model or its sampling ODE. It changes the representation in which the flow is trained, so that the regression problem it solves is well conditioned.

Step 1

Precondition

Learn a reversible map P (normalizing flow, or a low-capacity flow model) that pushes the target toward an isotropic representation:

$$\tilde{x}_1 = P(x_1), \qquad \mathrm{Cov}(\tilde{x}_1) \approx I$$

Step 2

Match

Train the unchanged FM objective on the transformed path, from the same Gaussian reference. With the sincos schedule used in every experiment here, $c(t)=\cos(\pi t/2)$, $s(t)=\sin(\pi t/2)$:

$$\tilde{x}_t = \cos\!\frac{\pi t}{2}\,x_0 + \sin\!\frac{\pi t}{2}\,\tilde{x}_1, \qquad x_0 \sim \mathcal{N}(0, I)$$

Step 3

Map back

Integrate the learned ODE and invert the preconditioner. Nothing else changes at sampling time:

$$\hat{x}_1 = P^{-1}\!\left(\hat{\tilde{x}}_1\right)$$

Algorithm 1 — Precondition-then-Match

Learn a reversible preconditioner $P_\phi$ on target data, so that $\mathrm{Cov}(P_\phi(x_1)) \approx I$.
Transform the targets: $\tilde{x}_1 = P_\phi(x_1)$.
Train $v_\theta$ with the unchanged FM objective on the preconditioned path $\tilde{x}_t = \cos\!\frac{\pi t}{2}\,x_0 + \sin\!\frac{\pi t}{2}\,\tilde{x}_1$.
Sample $z \sim \mathcal{N}(0,I)$ and solve $d\tilde{x}/dt = v_\theta(\tilde{x}, t)$ to obtain $\hat{\tilde{x}}_1$.
Return $\hat{x}_1 = P_\phi^{-1}(\hat{\tilde{x}}_1)$.

Results

Better conditioning, better samples

Every comparison keeps the downstream generator fixed and changes only the representation in which it is trained.

2D Gaussian

Four panels: standard FM transport, preconditioned FM transport, MMD versus training step, and condition number versus time. — **2D Gaussian transport with and without preconditioning.** (a) Standard FM transports samples from an isotropic source to an elongated target and struggles along the narrow direction. (b) Preconditioned FM transports to a whitened target and maps back with the inverse preconditioner. (c) Preconditioning lowers MMD and avoids early stagnation. (d) It lowers the condition number of $\Sigma_t$ along the path. Mean over 10 seeds; shaded regions are one standard deviation.

Gaussian mixtures

Four-row grid over GMM configurations. Columns: transported samples without preconditioning, transported samples with preconditioning, MMD during training, relative low-eigenvalue recovery error, and per-component path condition number. — **Preconditioning improves Gaussian-mixture flow matching across controlled 2D configurations.** We compare standard and preconditioned flow matching on GMM targets with varying numbers of components, anisotropy levels and transport distances; each row is one configuration. Columns show transported samples *without* preconditioning, transported samples *with* preconditioning, MMD during training, relative low-eigenvalue recovery error, and maximum per-component path condition number. Across all configurations, preconditioning produces better-aligned transported samples, lower final MMD, smaller low-eigenvalue error and substantially reduced path conditioning — the unpreconditioned model typically captures a component's high-variance direction but distorts its low-variance one. Plots show the mean over 10 independent runs; error bars and shaded regions denote one standard deviation.

Latent MNIST

MNIST samples generated by latent flow matching with no preconditioning, normalizing-flow preconditioning, and flow-matching preconditioning. — **MNIST qualitative comparison.** Latent-space flow matching under different preconditioning strategies: no preconditioning (FID 5.97), normalizing-flow preconditioning (FID 2.74), flow-matching preconditioning (FID 3.98).

Table 1 — MNIST preconditioning strategies

Compute-matched baselines enlarge the unpreconditioned generator to match the total parameters of preconditioner + flow.

Preconditioning method	Compute matched	Params	FID ↓	MMD ↓ (×10⁻³)	Precision ↑	Recall ↑
No preconditioning	No	2.32M	5.97	77.10	0.5774	0.2856
No preconditioning	Yes	3.37M	5.42	75.46	0.5869	0.2618
ZCA whitening	No	2.32M	4.20	6.25	0.6022	0.3107
ZCA whitening	Yes	3.37M	4.02	7.02	0.6117	0.2885
Normalizing flow	No	3.37M	2.74	2.95	0.7271	0.3561
Flow matching preconditioner	No	3.37M	3.98	3.02	0.8440	0.5486

Image benchmarks

Table 2 — FID, precision and recall across datasets

Latent space of a pretrained VAE for every dataset except CIFAR-10, which is modelled in pixel space.

Dataset	Method	FID ↓	Precision ↑	Recall ↑
LSUN Churches 256×256	No preconditioner	19.53	0.4540	0.3239
LSUN Churches 256×256	FM preconditioner	14.47	0.5313	0.3255
Oxford Flowers-102 256×256	No preconditioner	23.82	0.5213	0.2761
Oxford Flowers-102 256×256	FM preconditioner	22.64	0.5428	0.2854
AFHQ Cats 512×512	No preconditioner	7.11	0.7764	0.5396
AFHQ Cats 512×512	FM preconditioner	4.86	0.8320	0.5297
CIFAR-10 32×32	No preconditioner	4.15	0.8097	0.5878
CIFAR-10 32×32	FM preconditioner	4.02	0.7856	0.5892
ImageNet-1k 64×64	No preconditioner	8.66	0.7800	0.6339
ImageNet-1k 64×64	FM preconditioner	4.98	0.8172	0.6432

LSUN Churches samples: standard flow matching above, flow-based preconditioning below. — **LSUN Churches, 256×256.** **Top:** standard FM. **Bottom:** FM with a flow preconditioner.

Oxford Flowers-102 samples: standard flow matching above, preconditioned below. — **Oxford Flowers-102, 256×256.** **Top:** standard FM. **Bottom:** FM with a flow preconditioner.

AFHQ Cats samples: standard flow matching above, preconditioned below. — **AFHQ Cats, 512×512.** **Top:** standard FM. **Bottom:** FM with a flow preconditioner.

CIFAR-10 samples: standard flow matching left, preconditioned right. — **CIFAR-10, 32×32.** **Left:** standard FM. **Right:** FM with a flow preconditioner.

ImageNet-1k samples: standard flow matching above, preconditioned below. — **ImageNet-1k, 64×64.** **Left:** standard FM. **Right:** FM with a flow preconditioner.

Conditioning in real representations

Condition number of the intermediate covariance over time for MNIST and ImageNet-1k latent spaces, with and without preconditioning. — **Conditioning of intermediate flow matching distributions.** The condition number $\kappa(\Sigma_t)$ over $t$ for **(a)** MNIST and **(b)** ImageNet-1k latent representations. Without preconditioning the intermediate distributions become increasingly ill-conditioned as $t \to 1$; both NF and FM preconditioning substantially reduce $\kappa(\Sigma_t)$, especially near the data endpoint.

(a) FID versus number of preconditioner parameters for NF and FM preconditioners, against the compute-matched and 2.32M baselines. — **Is it just extra parameters? Preconditioner quality ablation on MNIST.** We evaluate image generation quality as a function of preconditioner capacity and training progress using FID computed on 50k generated samples. **(a)** Increasing the number of preconditioner parameters improves generation quality for both normalizing flow (NF) and flow matching (FM) preconditioners. The blue dashed line denotes the original unpreconditioned FM baseline with 2.32M parameters, and the black dash-dotted curve shows a compute-matched unpreconditioned baseline whose parameter count is increased by approximately the same number of parameters as the preconditioner. Preconditioning improves FID more substantially than simply allocating the same additional parameter budget to the unpreconditioned generator. **(b)** Increasing the number of preconditioner training epochs also improves FID, with the gains eventually beginning to saturate. The dotted horizontal line shows the compute-matched unpreconditioned baseline with 3.37M parameters, and the dashed horizontal line shows the original 2.32M-parameter baseline. Across both ablations, NF preconditioning yields the strongest gains, while FM preconditioning also consistently improves over the unpreconditioned baselines.

(b) FID versus number of preconditioner training epochs for NF and FM preconditioners, against both baselines. — **Is it just extra parameters? Preconditioner quality ablation on MNIST.** We evaluate image generation quality as a function of preconditioner capacity and training progress using FID computed on 50k generated samples. **(a)** Increasing the number of preconditioner parameters improves generation quality for both normalizing flow (NF) and flow matching (FM) preconditioners. The blue dashed line denotes the original unpreconditioned FM baseline with 2.32M parameters, and the black dash-dotted curve shows a compute-matched unpreconditioned baseline whose parameter count is increased by approximately the same number of parameters as the preconditioner. Preconditioning improves FID more substantially than simply allocating the same additional parameter budget to the unpreconditioned generator. **(b)** Increasing the number of preconditioner training epochs also improves FID, with the gains eventually beginning to saturate. The dotted horizontal line shows the compute-matched unpreconditioned baseline with 3.37M parameters, and the dashed horizontal line shows the original 2.32M-parameter baseline. Across both ablations, NF preconditioning yields the strongest gains, while FM preconditioning also consistently improves over the unpreconditioned baselines.

What we do not claim

Preconditioning is not a new backbone, sampler or loss: the downstream FM model and its ODE are unchanged.
It does not mainly accelerate the early drop of the loss; it prevents optimization from plateauing along suppressed directions.
It is complementary to latent modelling, optimal flow matching and schedule design.

Citation

Cite this work

@inproceedings{ahamed2026preconditioned,
  title     = {Preconditioned Flow Matching},
  author    = {Ahamed, Shadab and Gal, Eshed and Siddiqui, Md Shahriar Rahim and
               Ghyselincks, Simon and Eliasof, Moshe and Haber, Eldad},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2026}
}

The paper is also on arXiv: 2603.02337.

License

Code — Apache 2.0

The released implementation, including this project page and the script that renders its video, is distributed under the Apache License, Version 2.0. See LICENSE in the repository.

Paper, figures and video — CC BY 4.0

The paper text, its figures and the concept video are licensed under Creative Commons Attribution 4.0. Reuse is welcome with attribution to the authors.

Third-party assets (datasets, pretrained VAEs, and the sample grids generated from them) remain under the licenses of their original sources.