PMosFM: Preconditioned Manifold Matching for One-Step Physics-Constrained Generation

Zhangyong Liang; Haibin Ling

cs.LGSep 30, 2026

PMosFM: Preconditioned Manifold Matching for One-Step Physics-Constrained Generation

Authors: Zhangyong Liang, Haibin Ling

Organizations: Department of Artificial Intelligence Westlake University

Abstract

Physics-constrained generative models aim to generate physical fields that match a target distribution and satisfy prescribed constraints. However, enforcing these constraints often increases sampling costs through iterative corrections or training costs through residual optimization and trajectory unrolling. To address this issue, we introduce \textbf{P}reconditioned \textbf{M}anifold \textbf{o}ne-\textbf{s}tep \textbf{F}low \textbf{M}atching (\textbf{PMosFM}), a preconditioned manifold matching framework for one-step physics-constrained generation. By encoding constraints in a manifold decoder, PMosFM learns transport in intrinsic coordinates without separate residual losses or terminal residual unrolling. A geometric preconditioner rescales coordinates using the decoder-induced metric, while a regularized covariance transform approximately whitens the interpolation-state inputs. A finite-interval objective couples velocity supervision with consistency between decoded endpoints in physical space. We show that exact parameterization removes residual-induced Gauss--Newton curvature, that geometric and covariance effects separate in a local conditioning bound, and that physical flow-map error bounds endpoint distributional error. Controlled ablations examine conditioning, and experiments evaluate optimizer-update time and memory footprint. At inference, PMosFM uses one neural transport evaluation followed by physical decoding. Experiments across benchmarks show lower training and sampling time than the multi-step baselines at comparable physical and distributional fidelity. Code and datasets will be released publicly.

Figures & tables

Figure 1: Overview of PMosFM and representative physics-constrained methods. (a) PCFM projects a pre-correction endpoint

x_{1}^{-}

onto the feasible set

\mathcal{Z}_{c}

using

\Pi_{\mathcal{Z}_{c}}

. (b) PBFM combines flow and residual gradients

g_{\mathrm{FM}}

and

g_{R}

into the ConFIG direction

g_{\mathrm{cf}}

. (c) PMosFM applies precondition–manifold–matching, using

T_{\theta}^{0,1}

for one-step transport and

\Psi_{c}

for physical decoding.

Method	Physics at training	Hard constraints	Gradient-free inference	Complex constraints	No manual physics– fidelity balancing	One step sampling	Residual-factorized physics training	Spectral/Jacobian preconditioning
FFM [ 20 ]	$\times$	$\checkmark$	$\checkmark$	$\times$	$\times$	$\times$	$\times$	$\times$
FM-OT [ 24 ]	$\times$	$\times$	$\checkmark$	$\times$	$\times$	$\times$	$\times$	$\times$
CoCoGen [ 19 ]	$\times$	$\checkmark$	$\checkmark$	$\checkmark$	$\times$	$\times$	$\times$	$\times$
DiffusionPDE [ 18 ]	$\times$	$\times$	$\checkmark$	$\checkmark$	$\times$	$\times$	$\times$	$\times$
PIDM [ 4 ]	$\checkmark$	$\times$	$\checkmark$	$\checkmark$	$\times$	$\times$	$\times$	$\times$
D-Flow [ 5 ]	$\times$	$\times$	$\checkmark$	$\checkmark$	$\times$	$\times$	$\times$	$\times$

Table 1: Comparison between PMosFM and other physics-constrained generative models.

Figure 2: Manifold coordinates and endpoint matching. Schematic comparison of PCFM, PBFM, and PMosFM.

Figure 3: Preconditioned residual geometry.

y_{i}=\mu_{c}+C_{c}^{-1}r_{i}

; dots denote

s

-derivatives, with

\dot{r}_{s}=w

and

\dot{y}_{s}=C_{c}^{-1}w

Figure 4: Physical and sampling quality comparison. Lower divergence and MSE are better.

Benchmark		PBFM					PMosFM † 0 unrolling
Benchmark		1 step	2 steps	3 steps	4 steps	Sum	PMosFM † 0 unrolling
Darcy	[s]	0.067	0.098	0.128	0.159	0.452	0.019
Darcy	[GB]	12.6	23.3	33.7	44.2	114	0.128
Kolmogorov	[s]	0.194	0.302	0.410	0.518	1.42	0.021
Kolmogorov	[GB]	4.52	6.91	9.56	12.2	33.2	0.143
Dynamic Stall	[s]	0.081	0.118	0.155	0.190	0.544	0.056

Table 2: Optimizer-update time (s) and peak CUDA memory (GB) across unrolling configurations.

Figure 5: Preconditioning of feasible transport. (a) No preconditioning map. (b) With preconditioning map. (c) MMD 2 convergence. (d) Low-variance mode error. (e) Interpolation conditioning.

Figure 6: Training-residual convergence across physics-constrained benchmarks.

Figure 7: Generated fields and conservation residuals across methods. The top row compares generated solutions, while the bottom row shows the conservation-residual trajectories.

Benchmark	Metric	FM-OT	CoCoGen	PIDM	DiffusionPDE	D-Flow	ECI	PCFM	PBFM	PMosFM
Dynamic Stall	RE $\times 10^{6}$ $\downarrow$	11.0	376	$2.96\times 10^{3}$	12.2	11.3	40.7	0.143	0.339	0.0421
	WD $\times 10^{4}$ $\downarrow$	2.71	132	179	2.51	3.48	11.1	4.01	1.81	0.842
	JS $\times 10^{2}$ $\downarrow$	0.983	1.69	4.14	1.03	1.01	0.0570	7.46	0.680	0.0368
	NFE $\downarrow$	20	100	100	20	20	200	20	20	1
	IT [s] $\downarrow$	0.0598	0.184	0.0866	0.172	0.139	0.432	3.91	0.0605	0.00338
Darcy Flow	RE $\downarrow$	4.16	1.32	0.0220	3.39	2.29	3.05	$4.19\times 10^{3}$	0.838	0.120

Table 3: Sampling quality and inference efficiency across benchmarks. RE: physical residual error; WD: Wasserstein distance; JS: Jensen–Shannon divergence; NFE: learned network evaluations; IT: inference time.

Figure 8: Structural recovery under sparse observations. PMosFM retains large-scale structure under masking across field, spectral, and velocity-gradient diagnostics.

Benchmark	PMosFM-	PMosFM-	PMosFM-	PMosFM-	PMosFM-	PMosFM-	PMosFM-
Benchmark	MeanFlow	iMF	W-Flow	Drifting	SoFlow	Shortcut	native
Burgers	0.0758	0.0765	0.0518	0.0713	0.0714	0.0849	0.0374
Dynamic Stall	0.1403	0.1315	0.0594	0.0442	0.1352	0.0455	0.0036
Darcy Flow	0.0277	0.0341	0.0163	0.0233	0.0317	0.0158	0.0043
KF Reconstruction	0.0144	0.0145	0.0574	0.0476	0.0249	0.0104	0.0037
KF Generation	2.0645	2.1278	1.6020	2.1451	2.0887	1.9672	1.5922

Table 4: One-step-core ablation within PMosFM. Energy distance (ED;

\downarrow

) is reported; lower is better. KF denotes Kolmogorov flow. Bold denotes the lowest available value in each benchmark row.

Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Name	Description
$(x,\mathcal{X}_{h},c,h)$	state space, condition, mesh	$x\in\mathcal{X}_{h}\subseteq\mathbb{R}^{d}$ is a discretized physical state; $c$ is the condition; $h$ is the spatial discretization.
$R_{h}(x,c)\in\mathbb{R}^{n_{R}}$	discrete residual	Residual that evaluates the hard physical constraint; $n_{R}$ is the number of residual components.
$(\mathcal{Z}_{c},\mathcal{M}_{c}^{\chi})$	residual-zero set	$\mathcal{Z}_{c}=\{x:R_{h}(x,c)=0\}$ ; $\mathcal{M}_{c}^{\chi}=\chi_{c}(\mathcal{Y}_{c})\subseteq\mathcal{Z}_{c}$ is the component represented by the model.
$(y,\mathcal{Y}_{c},m)$	intrinsic coordinates	$y\in\mathcal{Y}_{c}\subset\mathbb{R}^{m}$ represents the independent degrees of freedom; $m=d-q$ under the regular-rank assumption.
$(\chi_{c},\eta_{c})$	constraint-satisfying parameterization and encoder	$R_{h}\circ\chi_{c}\equiv 0$ ; $\eta_{c}$ encodes feasible data or provides a measurable selection.
$(J_{\chi},J_{R})$	coordinate and residual Jacobians	$J_{\chi}=D_{y}\chi_{c}$ , $J_{R}=D_{x}R_{h}$ satisfy $J_{R}J_{\chi}=0$ ; regular-rank conditions identify $\mathop{\mathrm{range}}\nolimits J_{\chi}$ with the tangent space.

Appendix

Table 5: Summary of the notation adopted throughout the paper and appendix.

Table 14

Constraint type	Representative form	Linearity
Affine IC / BC / observations	$\begin{gathered}R_{h}(x,c)=Ax-b=0,\\ x=x_{p}+Ny,\;AN=0\end{gathered}$	Linear
Global conservation	$\begin{gathered}\ell^{\top}x-C=0,\\ \text{conservation-preserving or mean-free coordinates}\end{gathered}$	Linear
Differential compatibility	$\begin{gathered}D_{x}u+D_{y}v=0,\\ (u,v)=(D_{y}\psi,-D_{x}\psi)\end{gathered}$	Linear
Algebraic constitutive relation	$\begin{gathered}g(x,c)=0,\\ \text{eliminate dependent variables}\end{gathered}$	Nonlinear
Implicit PDE constraint	$\begin{gathered}F_{h}(z^{\star}(y,c),y,c)=0,\\ z=z^{\star}(y,c)\end{gathered}$	Potentially nonlinear
General nonlinear coupled constraint	$\begin{gathered}R_{h}(x,c)=0,\\ \text{local parameterization or approximate decoder}\end{gathered}$	Nonlinear

Appendix

Table 6: Summary of constraint types in the experimental suite, categorized by mathematical form and scope.

Benchmark	Metric	FM-OT	CoCoGen	PIDM	DiffusionPDE	D-Flow	ECI	PCFM	PBFM	PMosFM
Dynamic Stall	Train s/iter $\downarrow$	0.0183	0.0090	0.0094	0.0511	0.0470	0.0169	0.0183	0.5440	0.0563
Dynamic Stall	Peak GB $\downarrow$	0.4027	0.9467	0.9465	2.1603	0.9475	0.4027	0.4027	34.1000	0.2580
Darcy Flow	Train s/iter $\downarrow$	0.0289	0.0073	0.0065	0.0517	0.0170	0.0276	0.0289	0.4520	0.0185
Darcy Flow	Peak GB $\downarrow$	0.5621	0.2412	0.2412	2.1569	0.2418	0.5621	0.5621	114.0000	0.1280
Burgers	Train s/iter $\downarrow$	0.0246	0.0249	0.0245	0.0513	0.0240	0.0091	0.0091	0.0126	0.0057
Burgers	Peak GB $\downarrow$	2.2491	2.2465	2.2465	2.1565	2.2465	0.9487	0.9487	0.7245	0.5676

Appendix

Table 7: Training time (s per optimizer update) and peak allocated GPU memory (GB) across the benchmarks and baselines in Table 3 . Lower is better; bold marks PMosFM.

Benchmark	Task	Physical check	Distributional check
Darcy flow	Field generation	PDE/boundary residual	WD/JS; field moments
Dynamic stall	Conditional field	Wall-shear residual	MSE; shock statistics
Burgers	Nonlinear field	Conservation/boundary residual	WD/JS; shock location
Kolmogorov flow generation	Turbulence field	Incompressibility/flow statistics	MMD/WD; spectra; diversity
Kolmogorov flow reconstruction	Conditional inverse	Observation/flow consistency	MSE; CRPS/coverage
Turbulence forecast	Forecasting	Temporal/flow diagnostics	Forecast MSE; spectra; long-horizon stats

Appendix

Table 8: Benchmark summary. Physical checks assess conservation; statistical metrics assess coverage.

Dataset	Metric	FFM	Diffusion PDE	PDM	D-Flow	ECI	PCFM	PBFM	PMosFM
Heat Equation	MMSE / $10^{-2}$	4.56	4.49	0.45	1.97	0.697	0.241	0.043	0.016
	SMSE / $10^{-2}$	3.51	3.93	0.02	1.14	0.973	0.937	0.242	0.003
	CE (IC) / $10^{-2}$	579	599	0	102	0	0	24.2	0
	CE (CL) / $10^{-2}$	2.11	2.06	0	64.8	0	0	57.5	0
	FPD	1.77	1.70	0.17	2.70	1.34	1.22	–	0.0806
Navier–Stokes	MMSE / $10^{-2}$	16.5	17.4	12.21	–	5.23	4.59	–	4.41

Appendix

Table 9: Zero-shot constrained-PDE benchmark. Heat and Navier–Stokes: linear IC/CL; Reaction–Diffusion and Burgers: nonlinear CL with IC/BC. Lower is better; best baseline in bold.

Figure 9: Fixed-IC Burgers conservation diagnostics. Rows compare solution means and standard deviations, followed by mass and boundary residuals for the displayed methods.

Figure 10: Fixed-IC Burgers solution statistics. Mean solution profiles across time and pointwise standard deviations for the reference and displayed methods.

Coordinates	Active $\kappa_{+}$	Steps to $90\%$ $\downarrow$
Raw covariance coordinates	$10^{6}$	$2{,}558{,}427$
Exact state-whitened coordinates	$1$	$1$

Appendix

Table 11: Population tangent-regression conditioning with 64 eigenmodes and the same stable step rule. Lower is better.

Figure 11: Qualitative Dynamic Stall comparison. One held-out condition is shown. Columns show pressure, velocity gradients, temperature, density, and wall shear. Rows compare the reference and model predictions; PMosFM retains the shock-transition location and cross-channel structure.

Figure 12: Dynamic Stall ensemble moments. Conditional means and standard deviations of pressure, tangential velocity gradients, temperature, density, and wall shear for one held-out condition.

Figure 13: Focused

Q

–

R

comparison for turbulence forecasting. Selected physics-constrained baselines are compared with the held-out ground truth.

Method	Field MSE $\downarrow$	Divergence RMS $\downarrow$	$Q$ – $R$ JS $\downarrow$	$Q$ – $R$ TV $\downarrow$	Hellinger $\downarrow$
FM-OT	$1.097$	$1.000\times 10^{-1}$	$9.165\times 10^{-3}$	$7.499\times 10^{-2}$	$1.048\times 10^{-1}$
advNO	$2.205$	$1.248$	$1.317\times 10^{-1}$	$4.171\times 10^{-1}$	$3.798\times 10^{-1}$
FluidFlow	$1.102$	$1.004\times 10^{-1}$	$9.422\times 10^{-3}$	$7.609\times 10^{-2}$	$1.062\times 10^{-1}$
CoCoGen	$2.006\times 10^{2}$	$4.693$	$6.803\times 10^{-2}$	$2.857\times 10^{-1}$	$2.727\times 10^{-1}$
DiffusionPDE	$1.785$	$6.874\times 10^{-1}$	$1.815\times 10^{-1}$	$4.653\times 10^{-1}$	$4.617\times 10^{-1}$
PIDM	$1.695\times 10^{2}$	$4.248$	$6.351\times 10^{-2}$	$2.721\times 10^{-1}$	$2.638\times 10^{-1}$

Appendix

Table 12: Held-out turbulence evaluation for physics-constrained generation uses a common evaluator with fixed

Q

–

R

bins. All methods follow the main-text objective; lower is better.

Figure 14: Full

Q

–

R

comparison for turbulence forecasting. All evaluated methods are compared with the held-out ground truth.

Missing	Method	Missing-region	Physical	CRPS $\downarrow$	Coverage	Interval
Missing	Method	RMSE $\downarrow$	residual $\downarrow$	CRPS $\downarrow$	error $\downarrow$	width
50%	PBFM	0.063159	1.102126	0.024882	0.075232	0.084228
	ECI	0.045158	1.910683	0.019637	0.035143	0.126545
	PCFM	0.067615	1.520453	0.030269	0.099867	0.112164
	PCFT	0.069091	1.081061	0.027396	0.095131	0.082924
	PMosFM	0.029653	0.620392	0.008742	0.010382	0.036809

Appendix

Table 13: TFR sparse-reconstruction metrics for held-out frames 168–186 (32 samples per condition). Lower is better except interval width; bold marks the minimum at each missing rate.

Figure 15: Random-mask reconstruction. One held-out turbulence frame is tested with shared exact-count velocity masks (0–99% missing). Columns show the input, ground truth, PCFM, ECI, PBFM, and PMosFM; lower panels show energy spectra and standardized $\partial_{z}w$ PDFs.

Benchmark	Configuration	Loss Formulation	Field MSE $\downarrow$	Residual / Div. RMS $\downarrow$	Physical Sliced $W_{2}$ $\downarrow$
Kolmogorov Flow ( $128\times 128$ )	PE only ( $\gamma=0.0$ )	Direct Endpoint $\mathcal{L}_{\mathrm{PE}}$	$0.4241\pm 0.0182$	$1629.9\pm 52.4$	$0.3852\pm 0.0124$
	PE-heavy ( $\gamma=0.25$ )	Two-Time Joint	$0.1853\pm 0.0118$	$850.4\pm 38.1$	$0.2821\pm 0.0092$
	PMosFM Balanced ( $\gamma=0.50$ )	Two-Time Joint ( $\mathcal{L}_{\mathrm{FM}}+\mathcal{L}_{\mathrm{PE}}$ )	$\mathbf{0.0531\pm 0.0019}$	$\mathbf{3.45\times 10^{-9}}$	$\mathbf{0.2522\pm 0.0014}$
	FM-heavy ( $\gamma=0.75$ )	Two-Time Joint	$0.2184\pm 0.0145$	$1453.1\pm 46.2$	$0.2789\pm 0.0101$
	FM only ( $\gamma=1.0$ )	Velocity Flow Matching $\mathcal{L}_{\mathrm{FM}}$	$0.4714\pm 0.0210$	$7258.3\pm 112.5$	$0.3621\pm 0.0153$
Darcy Flow ( $64\times 64$ )	PE only ( $\gamma=0.0$ )	Direct Endpoint $\mathcal{L}_{\mathrm{PE}}$	$0.0382\pm 0.0024$	$870.8\pm 42.1$	$0.0284\pm 0.0018$

Appendix

Table 14: Velocity–endpoint loss balance on Kolmogorov and Darcy flow.

	Training terms			Held-out diagnostic
Variant	Velocity	Endpoint	$\lambda_{\nabla}$	Physical MSE $\downarrow$	Divergence RMS $\downarrow$	$Q$ – $R$ JS $\downarrow$	$Q$ – $R$ TV $\downarrow$	Hellinger $\downarrow$
No gradient	yes	yes	–	$5.631\times 10^{-3}$	$1.507\times 10^{-7}$	$6.257\times 10^{-3}$	$5.634\times 10^{-2}$	$8.698\times 10^{-2}$
Full objective	yes	yes	$0.2$	$5.524\times 10^{-3}$	$1.508\times 10^{-7}$	$6.202\times 10^{-3}$	$5.565\times 10^{-2}$	$8.664\times 10^{-2}$

Appendix

Table 15: One-step objective ablation for one-step turbulence forecasting. Lower is better.

Transform	Condition $\kappa(H)$ $\downarrow$	90% Reduction $\downarrow$	Kolmogorov Div Err $\downarrow$	Iteration ratio
Identity (Vanilla)	$1.00\times 10^{6}$	$2{,}558{,}427$	$625.17$	$1.0\times$ (baseline)
Standardization (Diagonal)	$4.82\times 10^{3}$	$1{,}240$	$485.30$	$2{,}063\times$
Preconditioning (PMosFM)	$\mathbf{1.00}$	$\mathbf{1\text{ step}}$	$\mathbf{234.01}$	$\mathbf{>10^{6}\times}$

Appendix

Table 16: Geometric preconditioning effects on conditioning and Kolmogorov-flow convergence.

Iterations $K$	Relative Pressure Error $\downarrow$	Independent Residual RE $\downarrow$	Latency / sample (ms) $\downarrow$
16	$0.0124\pm 0.0012$	$0.0385\pm 0.0035$	0.62 ms
32	$0.0031\pm 0.0004$	$0.0084\pm 0.0009$	0.98 ms
64	$0.0008\pm 0.0001$	$0.0019\pm 0.0002$	1.65 ms
128	$0.0002\pm 0.00003$	$0.00048\pm 0.00005$	2.84 ms
256	$\mathbf{0.00004\pm 0.00001}$	$\mathbf{0.00012\pm 0.00002}$	5.12 ms

Appendix

Table 17: Darcy decoder accuracy and latency across iteration counts.

$N$	$\kappa(A_{h}^{\top}A_{h})$	$\kappa(A_{h}^{\top}W_{h}A_{h})$
32	$4.33\times 10^{4}$	$1.0001$
64	$6.90\times 10^{5}$	$1.0001$
128	$1.10\times 10^{7}$	$1.0001$
256	$1.76\times 10^{8}$	$1.0001$

Appendix

Table 18: Normal conditioning.

Figure 16: Support and measure on a residual manifold. Exact parameterization enforces feasibility, and the co-area correction recovers the target law.

Figure 17: Normal and tangent conditioning. Left: preconditioning controls residual-normal conditioning under grid refinement. Right: whitening accelerates tangent-regression convergence.

Explore similar work

Constraint-Aware Flow Matching: Decision Aligned End-to-End Training for Constrained Sampling
May 12, 2026Jacob K. Christopher, James E. Warner, Ferdinando FiorettoGenerative Models Contrastive Flow Matching
University of Virginia · NASA Langley Research Center
Local Flow Matching Generative Models
Oct 3, 2024Chen Xu, Xiuyuan Cheng, Yao XieGenerative Flow Networks Conditional Flow Matching
H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology. · Department of Mathematics, Duke University
SNAP-FM: Sparse Nonlinear Accelerated Projection for Physics-Constrained Generative Modeling
Jun 30, 2026Alaina Kolli, Theodoros Xenakis, Utkarsh Utkarsh +4Inequality Constraints Solver Iterations
Massachusetts Institute of Technology, Cambridge, MA, USA · Norwegian University of Science and Technology, Trondheim, Norway

May 12, 2026, cs.LGJ/K move · Enter open · S save

Constraint-Aware Flow Matching: Decision Aligned End-to-End Training for Constrained Sampling

Jacob K. Christopher, James E. Warner, Ferdinando Fioretto

University of Virginia · NASA Langley Research Center

Generative Models Contrastive Flow Matching Training-Inference Mismatch Decisions

Open paper AskPDF arXiv

Deep generative models provide state-of-the-art performance across a wide array of applications, with recent studies showing increasing applicability for science and engineering. Despite a growing corpus of literature focused on the integration of physics-based constraints into the generation process, existing approaches fail to enforce strict constraint satisfaction while maintaining sample quality. In particular, training-free constrained sampling methods, while providing per-sample feasibility guarantees, introduce a fundamental mismatch between the training objective and the constrained sampling procedure, often leading to performance degradation. Identifying this training-sampling misalignment as a central limitation of current constrained generative modeling approaches, this paper proposes Constraint-Aware Flow Matching, a novel end-to-end framework that explicitly incorporates constraint projections into the training objective. By aligning the model's learned dynamics with the constrained sampling process, the proposed method mitigates distributional shift induced by projection-based corrections, enabling high-quality constrained generation. The proposed approach is evaluated on three challenging real-world benchmarks, illustrating the generality and efficacy of the method.

Abstract

Explore similar work

Constraint-Aware Flow Matching: Decision Aligned End-to-End Training for Constrained Sampling

Local Flow Matching Generative Models

SNAP-FM: Sparse Nonlinear Accelerated Projection for Physics-Constrained Generative Modeling