UniFluids: Unified Neural Operator Learning with Conditional Flow-matching

Haosen Li; Qi Meng; Jiahao Li; Rui Zhang; Ruihua Song; Liang Ma; Zhi-Ming Ma

cs.LGMar 19, 2026

UniFluids: Unified Neural Operator Learning with Conditional Flow-matching

Authors: Haosen Li, Qi Meng, Jiahao Li, Rui Zhang, Ruihua Song, Liang Ma, Zhi-Ming Ma

Organizations: Academy of Mathematics and Systems Science, Chinese Academy of Sciences · Microsoft Research Asia · Renmin University of China · Institute of Zoology, Chinese Academy of Sciences

Abstract

Partial differential equation (PDE) simulation holds extensive significance in scientific research. Currently, the integration of deep neural networks to learn solution operators of PDEs has introduced great potential. In this paper, we present UniFluids, a conditional flow-matching framework that harnesses the scalability of diffusion Transformer to unify learning of solution operators across diverse PDEs with varying dimensionality and physical variables. Unlike the autoregressive PDE foundation models, UniFluids adopts flow-matching to achieve parallel sequence generation, making it the first such approach for unified operator learning. Specifically, the introduction of a unified four-dimensional spatiotemporal representation for the heterogeneous PDE datasets enables joint training and conditional encoding. Furthermore, we find the effective dimension of the PDE dataset is much lower than its patch dimension. We thus employ $x$ -prediction in the flow-matching operator learning, which is verified to significantly improve prediction accuracy. We conduct a large-scale evaluation of UniFluids on several PDE datasets covering spatial dimensions 1D, 2D and 3D. Experimental results show that UniFluids achieves strong prediction accuracy and demonstrates good scalability and cross-scenario generalization capability. The code will be released later.

Figures & tables

Figure 1: Overall architecture. (a) Unified 4D representation: heterogeneous 1D/2D/3D PDE trajectories with different variable sets are aligned to a canonical 4D grid $(t,h,w,d)$ with zero-padding on degenerate axes and a channel mask for missing variables, then space–time patchified into tokens paired with explicit 4D coordinates, enabling 4D RoPE. (b) Unified condition encoder: a Transformer encodes the observed history into $\mathbf{c}_{\text{tok}}$ (dense condition) and a aggregated $\mathbf{c}_{g}$ (compact condition) via temporal/spatial aggregation. (c) Conditional flow-matching operator: starting from noisy future-window tokens, the operator stacks 4D-RoPE self-attention and cross-attention to $\mathbf{c}_{\text{tok}}$ , with AdaLN modulated by $(t,\mathbf{c}_{g},\texttt{dim\_type})$ , to predict the clean solution.

Dim- $d$	$V$	Target	PR $\downarrow$	EV90 $\downarrow$	EV90/ $d$ $\downarrow$
1D	48	$x$	1.14	1	0.0208
1D	48	$\epsilon$	47.60	43	0.8958
1D	48	$v$	1.75	19	0.3958
2D	512	$x$	1.29	2	0.0039
2D	512	$\epsilon$	471.99	428	0.8359
2D	512	$v$	18.05	276	0.5391

Table 1: Effective-dimension diagnostics for PDE patch vectors.

V

denotes the size of the patch vector.

Figure 2: Eigen-spectra of covariance for PDE patch vectors (log-scale), comparing targets $x$ , $\epsilon$ , and $v$ across 1D/2D/3D. Patchified PDE states concentrate variance in a small number of components (fast spectral decay), whereas $\epsilon$ (and partially $v$ ) spans a much larger effective dimension, making $\epsilon$ / $v$ prediction increasingly ill-conditioned as patch dimension grows.

Baselines – Task-specific models
Methods	1D			2D CFD			2D		3D
Methods	CFD	Adv.	Bur.	$M=1$	$M=0.1$	All	SWE	Incom.	CFD (TURB.)
PINNs	–	0.8130	0.9450	–	–	–	0.0170	–	–
U-Net	2.6700	0.7760	0.3201	–	–	–	0.0830	1.1200	0.7989
FNO	1.4100	0.0091	0.0174	–	–	–	0.0044	0.2574	0.3052
Baselines – Unified pre-training

Table 2: Comparison of the relative

L^{2}

-error on the PDEBench subsets used in this work (lower is better). “–” indicates that the result is not reported by the corresponding baseline under the same setting. For the 2D CFD datasets, the parameters

(\eta,\zeta)=(0.1,0.1)

. The improvement is calculated by the best results of UniFluids and best results of the unified pre-training baselines.

Figure 3: Scaling behavior of validation loss. Validation loss (log scale) across training epochs for UniFluids models of increasing scale (S/M/L/XL). Larger models consistently achieve lower validation loss, and the log-scale axis highlights the widening performance gap as training progresses.

Panel of Figure 4 — Figure 4: Validation loss versus epoch on CFD fine-tuning (S model), comparing training from scratch (w/o pretrain) and fine-tuning from the foundation checkpoint (w/ pretrain) across 1D/2D/3D.

Method	1D			2D			3D
Method	Adv.	Bur.	CFD	CFD (All)	SWE .	Incom.	CFD (turb.)
x-pred	0.0052	0.0253	0.0267	0.0561	0.0099	0.0553	0.6749
v-pred	0.0706	0.0243	0.0635	0.8197	0.0334	0.5258	0.5545

Table 3: Prediction-target ablation (S model). nRMSE

\downarrow

on representative PDEBench subsets; “2D CFD (All)” averages the four CFD regimes in our suite.

\epsilon

-pred is omitted because training becomes numerically unstable (loss blows up to NaN/Inf under our unified pipeline), so no valid nRMSE can be obtained.

Method	1D CFD	2D CFD	3D CFD
w/o pretrain	0.0269	0.0063	0.5954
w/ pretrain	0.0245	0.0049	0.5456
Reduction (%)	8.9%	22.2%	8.4%

Table 4: In-domain CFD adaptation (S model). nRMSE

\downarrow

. The 2D result is reported on the subset

(M,\eta,\zeta)=(1,10^{-1},10^{-1})

Methods	2D-Shock	2D-KH	2D-TOV
FNO	0.7484	1.0891	0.5946
U-Net	1.6667	0.1677	0.4217
MPP-L	0.3243	1.3261	0.3025
UniFluids-XL	0.2160	0.3113	0.2159

Table 5: Zero-shot performance on unseen cases (without fine-tuning). nRMSE

\downarrow

(lower is better).

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Subset	Dim	physical variables	$N$	$(T,N_{s})$	Target $N_{s}$
1D_Advection	1D	$(V_{x})$	80000	$(201,\;1024)$	$1024$
1D_Burgers	1D	$(V_{x})$	120000	$(201,\;1024)$	$1024$
1D_CFD (CNS)	1D	$(V_{x},\rho,p)$	50000	$(101,\;1024)$	$1024$
2D_CFD (CNS)	2D	$(V_{x},V_{y},\rho,p)$	40000	$(21,\;512^{2})$	$128^{2}$
diff_react_2d	2D	$(u,v)$	1000	$(101,\;128^{2})$	$128^{2}$
swe_pdb	2D	$(h)$	1000	$(101,\;128^{2})$	$128^{2}$

Appendix

Table 6: PDEBench subsets used in this work (variables and resolutions).

N

T

N_{s}

denote #trajectories, #timesteps, and spatial grid resolution in the PDEBench release; Target

N_{s}

is our unified resolution.

Model	$d$	$L_{e}$	$h_{e}$	$L_{j}$	$h_{j}$	Params (M)	Encoder/Flow-Matching Operator (M)
S	256	4	4	7	4	22.77	6.11 / 16.66
M	384	6	6	10	6	60.49	15.51 / 44.98
L	512	8	8	14	8	135.31	32.29 / 103.01
XL	1024	8	8	14	8	512.28	119.64 / 392.65

Appendix

Table 7: Model scaling configurations used in this work.

Target resolution	nRMSE $\downarrow$
128	0.01077
256	0.01063
512	0.01082
768	0.01099
1024	0.01033

Appendix

Table 8: Multi-resolution evaluation on 1D Burgers (nRMSE

\downarrow

, 110 test samples).

Model	1 step	5 steps	40 steps
UniFluids-S	0.03617 $\pm$ 0.00023	0.16240 $\pm$ 0.00060	1.26670 $\pm$ 0.00355
UniFluids-M	0.05475 $\pm$ 0.00594	0.23163 $\pm$ 0.01024	1.80871 $\pm$ 0.01983
UniFluids-XL	0.06911 $\pm$ 0.00062	0.31227 $\pm$ 0.00126	2.45982 $\pm$ 0.03195

Appendix

Table 9: Inference time (seconds) for one-shot generation of 10 frames on 1D_Advection . Mean

\pm

std over 10 runs.

Panel of Figure 5 — Figure 5: Prediction results on 1D PDEBench subsets. From left to right: 1D Advection, 1D Burgers, and 1D CFD (compressible Navier–Stokes). Each panel shows rollouts for time steps $T{+}1$ to $T{+}6$ ; the top row is ground truth and the bottom row is our prediction.

Panel of Figure 6 — Figure 6: Prediction results on 2D PDEBench subsets. Top row: ground truth; bottom row: our prediction; shown for time steps $T{+}1$ to $T{+}6$ . We visualize (top-left) 2D CFD (compressible Navier–Stokes), (top-right) reaction–diffusion, (bottom-left) shallow-water equation, and (bottom-right) 2D incompressible Navier–Stokes (with particle tracer).

Figure 7: Prediction results on the 3D CFD subset (compressible Navier–Stokes). We display rollouts for time steps $T{+}1$ to $T{+}6$ ; the top row shows ground truth and the bottom row shows our prediction.

PDE Operator Learning

Partial differential equation (PDE) operator learning builds models that map inputs such as coefficients and initial conditions to solution fields.

See all 228 papers

Learning Physical Operators using Neural Operators
Feb 26, 2026Vignesh Gopakumar, Ander Gray, Dan Giles +5PDE Surrogate Modeling Physics-Informed Neural Operators
UCL Centre for Artificial Intelligence · UK Atomic Energy Authority · LIX, CNRS, ´Ecole Polytechnique +2
PDEFlow: Autonomous Agentic PDE Pipelines for Neural Operator Learning and Solver-Free Inference
Jul 6, 2026Akshat Jani, Prathamesh Gadekar, Sakhinana Sagar Srinivas +1PDE Solving PDE Surrogate Modeling
Tata Research Development and Design Center, Pune, India 411057
Compositional Neural Operators for Multi-Dimensional Fluid Dynamics
May 12, 2026Hamda Hmida, Hsiu-Wen Chang, Youssef MesriPDE Operator Learning Physics-Informed Neural Operators
Mines Paris - PSL University, Centre for Material Forming (CEMEF) · Mines Paris - PSL University, Centre for Robotics (CAOR)

Feb 26, 2026J/K move · Enter open · S save

Learning Physical Operators using Neural Operators

Vignesh Gopakumar, Ander Gray, Dan Giles, Lorenzo Zanisi+4

UCL Centre for Artificial Intelligence · UK Atomic Energy Authority · LIX, CNRS, ´Ecole Polytechnique+2

PDE Surrogate Modeling Physics-Informed Neural Operators Computational Fluid Dynamics Modular Neural Networks

Open paper AskPDF arXiv

Neural operators have emerged as promising surrogate models for solving partial differential equations (PDEs), but struggle to generalise beyond training distributions and are often constrained to a fixed temporal discretisation. This work introduces a physics-informed training framework that addresses these limitations by decomposing PDEs using operator splitting methods, training separate neural operators to learn individual non-linear physical operators while approximating linear operators with fixed finite-difference convolutions. This modular mixture-of-experts architecture enables generalisation to novel physical regimes by explicitly encoding the underlying operator structure. We formulate the modelling task as a neural ordinary differential equation (ODE) where these learned operators constitute the right-hand side, enabling continuous-in-time predictions through standard ODE solvers and implicitly enforcing PDE constraints. Demonstrated on incompressible and compressible Navier--Stokes equations, our approach achieves better convergence and superior performance when generalising to unseen physics. The method remains parameter-efficient, enabling temporal extrapolation beyond training horizons, and provides interpretable components whose behaviour can be verified against known physics.

Abstract

PDE Operator Learning

Explore similar work

Learning Physical Operators using Neural Operators

PDEFlow: Autonomous Agentic PDE Pipelines for Neural Operator Learning and Solver-Free Inference

Compositional Neural Operators for Multi-Dimensional Fluid Dynamics