Preconditioned Flow Matching

Shadab Ahamed; Eshed Gal; Md Shahriar Rahim Siddiqui; Simon Ghyselincks; Moshe Eliasof; Eldad Haber

Figure 4: Gaussian model illustrating conditioning and optimization effects in flow matching: (a) Growth of the condition number

\kappa(\Sigma_{t}

) along the linear interpolation

\Sigma_{t}=(1-t)^{2}I+t^{2}H

showing how intermediate distributions

p_{t}

gradually inherit the ill-conditioning of

H

; (b) evolution of the eigenvalues

\sigma_{i}(t)

of

\Sigma_{t}

for

\kappa(H)=1000

, highlighting the increasing spectral disparity that creates slow and fast directions; (c) Covariance ellipses of

p_{t}=\mathcal{N}(0,\Sigma_{t})

at selected times, visualizing the transition from an isotropic distribution to a highly anisotropic one; (d) Predicted number of gradient-descent iterations

k(t)

required to reach a fixed error tolerance, demonstrating the sharp slowdown at large

t

due to ill-conditioning; (e) Effect of ideal preconditioning (whitening), which removes conditioning growth and keeps

\kappa(\Sigma_{t})\approx 1

for all

t

; (f) Per-mode gradient-descent error decay at

t=0.9

, showing that large-variance modes converge rapidly while small-variance modes decay extremely slowly; (g) Optimal score field

s^{*}(x,t)=-\Sigma_{t}^{-1}x

, exhibiting strong anisotropy and stiffness along low-variance directions; (h) Optimal flow matching velocity field

v^{\star}(x,t)=A^{\star}(t)x

, which is smoother than the score but still inherits anisotropy from

\Sigma_{t}

.

Figure 5: Flow matching without preconditioning. (a) Ground-truth Swiss roll samples. (b) Generated samples obtained by transporting Gaussian source

z

through a learned flow.

Direction	No preconditioner	NF preconditioner	FM preconditioner
$z\rightarrow x_{1}$	$0.11$	$0.06$	$0.07$
$x_{1}\rightarrow z$	$0.81$	$0.31$	$0.34$

Appendix

Table 3: Sliced-Wasserstein distances for the Swiss roll dataset. The top row compares generated samples with ground truth; the bottom row measures how Gaussian the pushed data is. Lower is better.

Figure 6: Flow matching with preconditioning. Top row: Normalizing flow preconditioning. (a) The data

x_{1}

(Swiss roll) is pushed forward by the normalizing flow preconditioner to

\tilde{x}

. (b) Random Gaussian samples are mapped by the learned downstream flow to the preconditioned data space. (c) For visualization, data in the preconditioned space can also be mapped backward through the learned flow toward Gaussian samples. (d) The points in the preconditioned space are mapped back to the original space using

{\cal P}_{\theta}^{-1}

. Bottom row: Low-capacity flow preconditioning. (e) The data

x_{1}

is pushed forward by a low-capacity flow preconditioner to

\tilde{x}

. (f) Random Gaussian samples are mapped by the learned downstream flow to the preconditioned data space. (g) For visualization, data in the preconditioned space can also be mapped backward through the learned flow toward Gaussian samples. (h) Generated samples in the preconditioned space are mapped back using the inverse low-capacity flow.

Figure 7: Preconditioning improves Gaussian-mixture flow matching across controlled 2D configurations. We compare standard flow matching and preconditioned flow matching on GMM targets with varying numbers of components, anisotropy levels, and transport distances. Each row shows one GMM configuration, while 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. These diagnostics support the theoretical prediction that preconditioning improves flow matching optimization by reducing anisotropy along the interpolation path, especially for the hardest mixture components. Plots show the mean over 10 independent runs with different random seeds; error bars / shaded regions denote one standard deviation.

Figure 8: OFM and preconditioning across six Gaussian-mixture configurations. Final MMD, trajectory straightness, normalized binned ISTE proxy, and MMD at the reported NFE

=8

setting. Endpoint MMD is evaluated after mapping preconditioned samples back to the original target space. Error bars show variability across 10 runs. The corresponding source and target distributions, together with samples transported by each method, are shown in Figure 9 .

Figure 9: Generated samples for FM, OFM, preconditioned FM, and their combination. Rows show Gaussian-mixture configurations; columns show methods. Blue points denote source samples, green points target samples, red points transported samples, and black triangles the transformed target for preconditioned methods. The displayed MMD values are in units of

10^{-3}

.

Figure 10: Conditioning of intermediate flow matching distributions with and without preconditioning. We plot the condition number

\kappa(\Sigma_{t})

of the intermediate covariance matrices

\Sigma_{t}

over time

t

for (a) MNIST and (b) ImageNet-1k latent representations. Without preconditioning, the intermediate distributions become increasingly ill-conditioned as

t\to 1

, indicating that the downstream flow matching regression problem inherits strong anisotropy from the target representation. Both normalizing flow (NF) and flow matching (FM) preconditioning substantially reduce

\kappa(\Sigma_{t})

, especially near the data endpoint, supporting the claim that preconditioning improves the geometry of the FM path rather than merely adding model capacity.

Figure 11: Directional learning dynamics and generation quality on latent MNIST. The left and right columns show standard flow matching and NF-preconditioned flow matching, respectively. Both models use the sine–cosine interpolation

x_{t}=\sin(\pi t/2)x_{1}+\cos(\pi t/2)x_{0}

, where

x_{0}\sim\mathcal{N}(0,I)

and

x_{1}

denotes the target in the latent representation used by each model. Top: Covariance-unweighted directional operator errors along the eight highest-variance (solid) and eight lowest-variance (dashed) eigendirections at

t\in\{0.50,0.75,0.90\}

. Middle: The directional imbalance ratio

R(t,e)=\operatorname{err}_{\mathrm{bottom}\text{-}8}(t,e)/\operatorname{err}_{\mathrm{top}\text{-}8}(t,e)

; the dotted line denotes

R(t,e)=1

. Bottom: FID measured throughout the corresponding MNIST training trajectories. The directional diagnostics are computed within the latent representation used by each model, whereas FID is evaluated in the common decoded image space after applying the inverse preconditioner where necessary.

Figure 12: 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.

Method	Additional training	Additional inference operation	Main-flow NFE	Inverse-prec. NFE	Total NFE	FID $\downarrow$
FM	–	–	60	0	60	5.97
NF preconditioning (Ours)	Train NF preconditioner	One direct NF inverse pass	60	N/A	$60^{\dagger}$	2.74
FM preconditioning (Ours)	Train FM preconditioner	Integrate inverse preconditioner	60	60	120	3.98
RF	–	–	60	0	60	6.86
RF + Reflow	Pair generation and RF retraining	–	60	0	60	5.15

Appendix

Table 4: Comparison with Rectified flow and Reflow on MNIST. NFE counts ODE vector-field evaluations. † The NF inverse is a direct network evaluation and is therefore not included in the ODE NFE.

Dataset	Training space	Baseline $\max_{t}\kappa(\Sigma_{t})\downarrow$	Precond. $\max_{t}\kappa(\widetilde{\Sigma}_{t})\downarrow$	Cond. reduction $\uparrow$	FID: no precond. $\downarrow$	FID: FM precond. $\downarrow$	Relative FID gain $\uparrow$	Diagnostic
MNIST	Latent	52.83	2.02	$26.2\times$	5.97	3.98	33.3%	Preconditioning beneficial
LSUN Churches	Latent	102.63	5.54	$18.53\times$	19.53	14.47	25.9%	Preconditioning beneficial
Oxford Flowers-102	Latent	83.47	47.92	$1.74\times$	23.82	22.64	5.0%	Limited benefit
AFHQ Cats	Latent	64.75	2.36	$27.44\times$	7.11	4.86	31.6%	Preconditioning beneficial
CIFAR-10	Pixel	755.92	453.68	$1.67\times$	4.15	4.02	3.1%	Limited benefit
ImageNet-1k	Latent	84.37	1.90	$44.4\times$	8.66	4.98	42.5%	Preconditioning beneficial

Appendix

Table 5: Relationship between path-conditioning reduction and generation quality across all image datasets. Conditioning reduction is the ratio of the baseline maximum path condition number to its value after preconditioning.

Distribution	$\gamma$	Peak path $\max_{t}\kappa(\Sigma_{t})\downarrow$	$\\|J_{\mathcal{P}_{\gamma}^{-1}}\\|_{2}\downarrow$	$\kappa(J_{\mathcal{P}_{\gamma}^{-1}})\downarrow$	Final MMD $(\times 10^{-3})\downarrow$
Bimodal	$0.00$	$849.75\pm 11.14$	$1.00$	$1.00$	$114.69\pm 4.26$
	$0.25$	$160.02\pm 2.05$	$2.18$	$2.31$	$202.96\pm 10.04$
	$\mathbf{0.50}$	$30.14\pm 0.37$	$4.76$	$5.34$	$55.74\pm 17.32$
	$0.75$	$5.70\pm 0.06$	$10.39$	$12.35$	$66.88\pm 17.59$
	$1.00$	$1.21\pm 0.02$	$22.66$	$28.53$	$324.64\pm 5.02$
Anisotropic ring	$0.00$	$7.95\pm 0.17$	$1.00$	$1.00$	$20.46\pm 2.90$

Appendix

Table 6: Trade-off between path conditioning, inverse sensitivity, and generation quality under partial affine preconditioning. Path conditioning and MMD are reported as mean

\pm

standard deviation over three seeds. Since the preconditioner is affine, its inverse-Jacobian quantities are constant over space. Bold values of

\gamma

identify the lowest final MMD for each target.

Points evaluated	$\\|J_{\mathcal{P}^{-1}}\\|_{2}$ median / p95 / max	$\kappa(J_{\mathcal{P}^{-1}})$ median / p95	Random amplification median / p95	Worst-direction amplification median / p95
Real preconditioned samples	$2.41\,/\,2.99\,/\,3.33$	$275\,/\,388$	$0.60\,/\,0.88$	$2.41\,/\,2.99$
Generated endpoints	$2.32\,/\,3.01\,/\,4.65$	$273\,/\,379$	$0.66\,/\,0.95$	$2.32\,/\,3.01$

Appendix

Table 7: Local sensitivity of the learned nonlinear MNIST inverse. Amplification is measured as

\|\mathcal{P}^{-1}(z+\delta z)-\mathcal{P}^{-1}(z)\|_{2}/\|\delta z\|_{2}

. The worst direction is the leading right singular vector of the inverse Jacobian; p95 denotes the

95

th percentile.

Deblurring
Method	Params/FLOPs	SSIM $\bm{\uparrow}$	PSNR (posterior mean) $\bm{\uparrow}$	Posterior std.	Misfit $(10^{-3})$ $\bm{\downarrow}$
Baseline FM	$1.07\mathrm{M}/1.0\times$	0.803	22.4	0.135	5.8
Parameter-matched	$1.50\mathrm{M}/1.0\times$	0.756	22.6	0.156	5.8
FLOPs-matched	$1.07\mathrm{M}/1.27\times$	0.817	22.5	0.128	5.8
Preconditioned (Ours)	$(1.07+0.44)\mathrm{M}/1.27\times$	0.834	21.7	0.128	6.2
$4\times$ super-resolution

Appendix

Table 8: Conditional image restoration on MNIST. “Params/FLOPs” reports the parameter count and training FLOPs relative to the baseline. For our method, the two parameter counts correspond to the main FM model and preconditioner. PSNR is computed from the posterior mean, while posterior standard deviation summarizes sample diversity. Higher SSIM and PSNR are better; lower measurement misfit is better.

Figure 13: MNIST qualitative comparison. Samples generated by latent-space flow matching under different preconditioning strategies. Left: The no-preconditioning baseline produces recognizable digits but shows occasional shape inconsistencies and reduced sharpness (FID = 5.97). Middle: NF preconditioning improves sample fidelity and yields cleaner digit shapes (FID = 2.74). Right: FM preconditioning also improves over the baseline, producing more consistent and diverse digits (FID = 3.98). Overall, the qualitative trends are consistent with the quantitative FID improvements.

Figure 14: LSUN Churches qualitative comparison. Samples generated by standard flow matching ( top ) and by flow-based preconditioning using an additional flow ( bottom ). Preconditioning produces a sharper global structure and more stable spatial layouts.

Figure 15: Oxford Flowers-102 qualitative comparison. Samples generated by standard flow matching ( top ) and by flow-based preconditioning using an additional flow ( bottom ). Preconditioning yields more coherent structure and improved visual consistency.

Figure 16: AFHQ Cats qualitative comparison. Samples generated by standard flow matching ( top ) and by flow-based preconditioning using an additional flow ( bottom ).

Figure 17: CIFAR-10 qualitative results. Samples generated by standard flow matching ( left ) and by flow-based preconditioning using an additional flow ( right ).

Figure 18: ImageNet-1k qualitative comparison. Samples generated by standard flow matching (top) and by flow-based preconditioning using an additional flow (bottom) . Preconditioning improves visual consistency and sample fidelity while using the same downstream generation setting.

Figure 19: ImageNet-1k qualitative comparison. Samples generated by standard flow matching (top) and by flow-based preconditioning using an additional flow (bottom) . Preconditioning improves visual consistency and sample fidelity while using the same downstream generation setting.

Setting	No Preconditioning	NF Preconditioning	FM Preconditioning
Training space	Latent	Latent	Latent
Latent dim	64	64	64
Interpolation	$s(t)=\sin(\pi t/2)$	$s(t)=\sin(\pi t/2)$	$s(t)=\sin(\pi t/2)$
Interpolation	$c(t)=\cos(\pi t/2)$	$c(t)=\cos(\pi t/2)$	$c(t)=\cos(\pi t/2)$
FM generator model
Network	CNN	CNN	CNN

Appendix

Table 9: Experimental settings and hyperparameters for MNIST latent flow matching experiments.

Setting	LSUN Churches	Oxford Flowers-102	AFHQ Cats	CIFAR-10	ImageNet-1k
Image resolution	$256{\times}256$	$256{\times}256$	$512{\times}512$	$32{\times}32$	$64{\times}64$
Training space	Latent	Latent	Latent	Pixel	Latent
VAE / latent encoder	Stable Diffusion VAE	Stable Diffusion VAE	Stable Diffusion VAE	–	Stable Diffusion VAE
Latent dim	$4\times 32\times 32$	$4\times 32\times 32$	$4\times 64\times 64$	–	$4\times 8\times 8$
Interpolation	$s(t)=\sin(\pi t/2)$	$s(t)=\sin(\pi t/2)$	$s(t)=\sin(\pi t/2)$	$s(t)=t$	$s(t)=t$
Interpolation	$c(t)=\cos(\pi t/2)$	$c(t)=\cos(\pi t/2)$	$c(t)=\cos(\pi t/2)$	$c(t)=(1-t)$	$c(t)=(1-t)$

Appendix

Table 10: Experimental settings and hyperparameters for LSUN Churches, Oxford Flowers-102, AFHQ Cats, CIFAR-10, and ImageNet-1k.

Preconditioning method	Compute matched?	# Params	FID $\downarrow$	MMD $\downarrow$ ( $\times 10^{-3}$ )	Precision $\uparrow$	Recall $\uparrow$
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 (Ours)	No	2.32M	4.20	6.25	0.6022	0.3107
ZCA whitening (Ours)	Yes	3.37M	4.02	7.02	0.6117	0.2885
NF preconditioner (Ours)	No	3.37M	2.74	2.95	0.7271	0.3561
FM preconditioner (Ours)	No	3.37M	3.98	3.02	0.8440	0.5486

Dataset	Method	FID $\downarrow$	Precision $\uparrow$	Recall $\uparrow$
LSUN Churches $(256\times 256)$	No preconditioner	19.53	0.4540	0.3239
LSUN Churches $(256\times 256)$	FM preconditioner (Ours)	14.47	0.5313	0.3255
Oxford Flowers-102 $(256\times 256)$	No preconditioner	23.82	0.5213	0.2761
Oxford Flowers-102 $(256\times 256)$	FM preconditioner (Ours)	22.64	0.5428	0.2854
AFHQ Cats $(512\times 512)$	No preconditioner	7.11	0.7764	0.5396
AFHQ Cats $(512\times 512)$	FM preconditioner (Ours)	4.86	0.8320	0.5297

Preconditioned Flow Matching

Abstract

Figures & tables

Appendix

Flow Matching

Posterior Augmented Flow Matching

Same Flow, Different Paths: Variance Reduction in Flow Matching

Flow Matching is Adaptive to Manifold Structures

Abstract

Flow Matching

Explore similar work

Posterior Augmented Flow Matching

Same Flow, Different Paths: Variance Reduction in Flow Matching

Flow Matching is Adaptive to Manifold Structures