UniFluids: Unified PDE Operator Learning

• UniFluids is a conditional flow-matching framework that unifies operator learning for heterogeneous PDE datasets by embedding trajectories in a canonical 4D spatiotemporal tensor.
• It employs a diffusion Transformer backbone with x-prediction to generate multiple future frames simultaneously, yielding low nRMSE across 1D, 2D, and 3D scenarios.
• The framework demonstrates strong scalability and cross-scenario generalization, outperforming specialized models on benchmark tasks in fluid and continuum mechanics.

Searching arXiv for the specific UniFluids paper and closely related unified PDE operator-learning work. UniFluids is a conditional flow-matching framework for PDE operator learning that is designed to unify solution-operator prediction across heterogeneous PDE datasets with varying dimensionality and physical variables. It adopts a diffusion Transformer backbone, a unified four-dimensional spatiotemporal representation, and an $x$-prediction variant within flow-matching in order to model 1D, 2D, and 3D time-dependent PDE trajectories in a single architecture (Li et al., 19 Mar 2026). In the broader literature, UniFluids also sits within a longer research trajectory on unified treatments of fluids, solids, and multiphase systems, including hyperbolic continuum formulations, unified SPH discretizations, mixed compressible–incompressible interfacial algorithms, and physics-pretrained operator models such as OmniFluids (Peshkov et al., 2017, Kincl et al., 2022, Denner et al., 2022, Zhang et al., 12 Jun 2025). In the narrow sense of the name, however, UniFluids denotes the 2026 neural-operator framework introduced for unified operator learning with conditional flow-matching (Li et al., 19 Mar 2026).

1. Definition and scope

UniFluids addresses the problem of learning solution operators for multiple PDE families within a single model rather than training separate task-specific surrogates. The framework is presented as a response to the heterogeneity of PDE datasets, particularly variation in spatial dimensionality, temporal rollout structure, and the number of physical channels. Its core claim is that a single conditional model can be trained jointly on diverse PDEs by representing all trajectories in a canonical four-dimensional spatiotemporal tensor and conditioning prediction of future frames on a learned encoding of history frames (Li et al., 19 Mar 2026).

The approach differs from autoregressive PDE foundation models by using flow-matching for parallel sequence generation. The paper characterizes this as the first such approach for unified operator learning, with the practical consequence that multiple future frames are generated in one shot rather than by stepwise rollout (Li et al., 19 Mar 2026). This places UniFluids within the neural-operator and scientific foundation-model literature, but with a design emphasis on cross-dimensional unification rather than specialization to a single PDE family.

A plausible implication is that UniFluids is best understood not as a replacement for classical PDE solvers in general, but as a surrogate-operator framework targeting scalability, shared representation learning, and cross-scenario transfer across heterogeneous benchmarks. That interpretation is consistent with the large-scale multi-dataset evaluation and the explicit comparison to unified-pretraining baselines such as MPP, DPOT, OmniArch, and OmniFluids (Li et al., 19 Mar 2026, Zhang et al., 12 Jun 2025).

2. Unified four-dimensional representation

The technical basis of UniFluids is a canonical 4D spatiotemporal representation. Each solution trajectory

$$u:\;\{1,\dots,T\}\times X\;\longrightarrow\mathbb R^C,\quad X\subset\mathbb R^d$$

is regridded into

$$U\in\mathbb R^{T\times H\times W\times D\times C_{\max}},$$

with

$$(H,W,D)= \begin{cases} (H,1,1) & d=1,\ (H,W,1) & d=2,\ (H,W,D) & d=3. \end{cases}$$

Missing axes are zero-padded, and missing physical variables are zeroed and indicated by a binary mask $m\in\{0,1\}^{C_{\max}}$. A discrete indicator $\mathrm{dim\_type}\in\{1,2,3\}$ records the spatial dimensionality (Li et al., 19 Mar 2026).

This representation is patchified in space and time using non-overlapping 4D patches of size $(P_t,P_h,P_w,P_d)$. Each patch is vectorized and linearly projected to a token $z_i\in\mathbb R^{d_{\rm model}}$, and the patch-origin coordinate

$p_i=(t_i,h_i,w_i,d_i)\in\mathbb N^4$

is used to parameterize a 4D rotary positional encoding in every Transformer block (Li et al., 19 Mar 2026).

The authors identify several consequences of this design. First, 1D, 2D, and 3D data live in the same 4D patch grid and share a single Transformer backbone. Second, no per-dimension special-casing of positional encodings or channels is needed. Third, every original degree of freedom is preserved up to the uniform regridding (Li et al., 19 Mar 2026). This unified representation is also the basis for the reported multi-resolution generalization on 1D Burgers, where a single trained model generalizes to $N_x=256,512,1024,2048$ with negligible error variation (Li et al., 19 Mar 2026).

A common misconception is that such unification implies dimensional equivalence of the underlying physics. The paper does not make that claim. Rather, the unification is representational and architectural: heterogeneous PDE trajectories are embedded into a common tensor-and-token format that permits shared operator learning.

3. Conditional flow-matching formulation

UniFluids models the conditional distribution of future frames $$u:\;\{1,\dots,T\}\times X\;\longrightarrow\mathbb R^C,\quad X\subset\mathbb R^d$$0 given history frames $$u:\;\{1,\dots,T\}\times X\;\longrightarrow\mathbb R^C,\quad X\subset\mathbb R^d$$1. Denoting the history tensor by $$u:\;\{1,\dots,T\}\times X\;\longrightarrow\mathbb R^C,\quad X\subset\mathbb R^d$$2 and the future tensor by $$u:\;\{1,\dots,T\}\times X\;\longrightarrow\mathbb R^C,\quad X\subset\mathbb R^d$$3, the framework chooses a linear path in data space,

$$u:\;\{1,\dots,T\}\times X\;\longrightarrow\mathbb R^C,\quad X\subset\mathbb R^d$$4

with $$u:\;\{1,\dots,T\}\times X\;\longrightarrow\mathbb R^C,\quad X\subset\mathbb R^d$$5. The corresponding target velocity is

$$u:\;\{1,\dots,T\}\times X\;\longrightarrow\mathbb R^C,\quad X\subset\mathbb R^d$$6

A neural network $$u:\;\{1,\dots,T\}\times X\;\longrightarrow\mathbb R^C,\quad X\subset\mathbb R^d$$7 predicts a clean-field estimate $$u:\;\{1,\dots,T\}\times X\;\longrightarrow\mathbb R^C,\quad X\subset\mathbb R^d$$8, where $$u:\;\{1,\dots,T\}\times X\;\longrightarrow\mathbb R^C,\quad X\subset\mathbb R^d$$9 denotes the conditioning from the history encoder, and the model velocity is defined by

$$U\in\mathbb R^{T\times H\times W\times D\times C_{\max}},$$0

Inference then solves the probability-flow ODE

$$U\in\mathbb R^{T\times H\times W\times D\times C_{\max}},$$1

to recover the predicted future field (Li et al., 19 Mar 2026).

Training uses the standard flow-matching velocity-regression objective,

$$U\in\mathbb R^{T\times H\times W\times D\times C_{\max}},$$2

with $$U\in\mathbb R^{T\times H\times W\times D\times C_{\max}},$$3, where the paper notes that $$U\in\mathbb R^{T\times H\times W\times D\times C_{\max}},$$4 is a suitable noise-level sampling distribution such as logit-normal (Li et al., 19 Mar 2026).

The paper emphasizes that UniFluids employs $$U\in\mathbb R^{T\times H\times W\times D\times C_{\max}},$$5-prediction rather than direct $$U\in\mathbb R^{T\times H\times W\times D\times C_{\max}},$$6-prediction or $$U\in\mathbb R^{T\times H\times W\times D\times C_{\max}},$$7-prediction. The rationale is linked to the observation that the effective dimension of the PDE dataset is much lower than its patch dimension, and prior work is cited for the proposition that when data lie on a low-dimensional submanifold, directly predicting clean data is easier than predicting $$U\in\mathbb R^{T\times H\times W\times D\times C_{\max}},$$8 or $$U\in\mathbb R^{T\times H\times W\times D\times C_{\max}},$$9 (Li et al., 19 Mar 2026). In the reported ablations, $$(H,W,D)= \begin{cases} (H,1,1) & d=1,\ (H,W,1) & d=2,\ (H,W,D) & d=3. \end{cases}$$0-prediction is stable and yields the lowest nRMSE across 1D, 2D, and 3D; $$(H,W,D)= \begin{cases} (H,1,1) & d=1,\ (H,W,1) & d=2,\ (H,W,D) & d=3. \end{cases}$$1-prediction is less robust; and $$(H,W,D)= \begin{cases} (H,1,1) & d=1,\ (H,W,1) & d=2,\ (H,W,D) & d=3. \end{cases}$$2-prediction fails under the unified pipeline, with loss going to NaN (Li et al., 19 Mar 2026).

This is a significant architectural distinction from autoregressive operator models. Instead of stepwise temporal propagation, UniFluids frames future-trajectory generation as conditional transport from Gaussian noise to future fields. The claimed advantage is parallel sequence generation, not a change in the underlying operator-learning target.

4. Architecture and training pipeline

UniFluids contains two Transformer modules: a unified condition encoder and a conditional flow-matching operator. The condition encoder takes history patch tokens with 4D-RoPE, processes them through a stack of self-attention and MLP blocks, and outputs both dense conditioning $$(H,W,D)= \begin{cases} (H,1,1) & d=1,\ (H,W,1) & d=2,\ (H,W,D) & d=3. \end{cases}$$3 and a compact global conditioning vector $$(H,W,D)= \begin{cases} (H,1,1) & d=1,\ (H,W,1) & d=2,\ (H,W,D) & d=3. \end{cases}$$4. The compact context is built by temporal aggregation through Fourier-weighted pooling in $$(H,W,D)= \begin{cases} (H,1,1) & d=1,\ (H,W,1) & d=2,\ (H,W,D) & d=3. \end{cases}$$5, followed by spatial aggregation through Perceiver-style cross-attention (Li et al., 19 Mar 2026).

The conditional flow-matching operator takes noisy future tokens, dense context $$(H,W,D)= \begin{cases} (H,1,1) & d=1,\ (H,W,1) & d=2,\ (H,W,D) & d=3. \end{cases}$$6, global context $$(H,W,D)= \begin{cases} (H,1,1) & d=1,\ (H,W,1) & d=2,\ (H,W,D) & d=3. \end{cases}$$7, the time level $$(H,W,D)= \begin{cases} (H,1,1) & d=1,\ (H,W,1) & d=2,\ (H,W,D) & d=3. \end{cases}$$8 embedded via Fourier features, and the dimension indicator $$(H,W,D)= \begin{cases} (H,1,1) & d=1,\ (H,W,1) & d=2,\ (H,W,D) & d=3. \end{cases}$$9. Each block comprises 4D-RoPE self-attention among future tokens, cross-attention from future tokens to dense context, AdaLN modulated by $m\in\{0,1\}^{C_{\max}}$0, and an MLP with residual pathways. A final linear projection and unpatchify operation reconstruct

$m\in\{0,1\}^{C_{\max}}$1

(Li et al., 19 Mar 2026).

The model is instantiated in four scales:

Model	$m\in\{0,1\}^{C_{\max}}$2	Encoder layers $m\in\{0,1\}^{C_{\max}}$3	Flow-matching layers $m\in\{0,1\}^{C_{\max}}$4	Total parameters
S	256	6	12	$m\in\{0,1\}^{C_{\max}}$5M
M	512	12	24	$m\in\{0,1\}^{C_{\max}}$6M
L	768	24	48	$m\in\{0,1\}^{C_{\max}}$7M
XL	1024	32	64	$m\in\{0,1\}^{C_{\max}}$8M

Training uses PDEBench subsets spanning 1D, 2D, and 3D PDEs. The datasets listed are 1D Advection, Burgers, and 1D compressible Navier–Stokes; 2D compressible Navier–Stokes in four regimes, incompressible Navier–Stokes with particle tracer, shallow-water, and reaction–diffusion; and 3D compressible Navier–Stokes with turbulent and random initialization (Li et al., 19 Mar 2026). All trajectories are resampled to unified spatial grids: $m\in\{0,1\}^{C_{\max}}$9 in 1D, $\mathrm{dim\_type}\in\{1,2,3\}$0 in 2D, and $\mathrm{dim\_type}\in\{1,2,3\}$1 in 3D (Li et al., 19 Mar 2026).

Optimization uses AdamW with base learning rate $\mathrm{dim\_type}\in\{1,2,3\}$2, weight decay $\mathrm{dim\_type}\in\{1,2,3\}$3, linear warm-up for 5% of steps, cosine decay to $\mathrm{dim\_type}\in\{1,2,3\}$4, global-norm gradient clipping at 1, and mixed precision. Noise levels are sampled from $\mathrm{dim\_type}\in\{1,2,3\}$5. Condition dropout is 10%, with inference classifier-free guidance scale 2.0. The ODE solver is Euler with 40 steps by default and 10 steps for fast validation (Li et al., 19 Mar 2026).

Figure 1 in the paper is reported to show consistent gains from S to XL. The XL model attains approximately $\mathrm{dim\_type}\in\{1,2,3\}$6 lower validation loss than S after 150 epochs (Li et al., 19 Mar 2026). This suggests that model-scaling behavior is a central empirical claim of the framework, rather than a secondary implementation detail.

5. Quantitative evaluation and empirical properties

The main quantitative metric reported for UniFluids is relative $\mathrm{dim\_type}\in\{1,2,3\}$7 error, denoted nRMSE. UniFluids-XL is compared to task-specific baselines such as PINNs, U-Net, and FNO, as well as unified-pretraining baselines including MPP, DPOT, and OmniArch (Li et al., 19 Mar 2026). The reported nRMSE values for UniFluids-XL are:

Benchmark	UniFluids-XL nRMSE
1D Advection	$\mathrm{dim\_type}\in\{1,2,3\}$8
1D Burgers	$\mathrm{dim\_type}\in\{1,2,3\}$9
1D CFD	$(P_t,P_h,P_w,P_d)$0
2D CFD regime 1	$(P_t,P_h,P_w,P_d)$1
2D CFD regime 2	$(P_t,P_h,P_w,P_d)$2
2D CFD regime 3	$(P_t,P_h,P_w,P_d)$3
2D CFD regime 4	$(P_t,P_h,P_w,P_d)$4
SWE	$(P_t,P_h,P_w,P_d)$5
Incom.2D	$(P_t,P_h,P_w,P_d)$6
3D CFD	$(P_t,P_h,P_w,P_d)$7

The paper also reports the percentage improvement over the best unified-pretraining baseline in the rightmost column of Table 2, though some entries are negative and the 3D CFD row is marked with a dash rather than a value (Li et al., 19 Mar 2026). Because the paper presents these entries directly, they should be read as table-reported benchmark differentials rather than as a uniform dominance claim across all tasks.

Several additional empirical properties are highlighted. Scaling from the 23M-parameter S model to the 512M-parameter XL model yields approximately $(P_t,P_h,P_w,P_d)$8 lower errors on 2D CFD and approximately $(P_t,P_h,P_w,P_d)$9 lower on Incom-2D (Li et al., 19 Mar 2026). In zero-shot generalization on unseen 2D shock, Kelvin–Helmholtz, and Orszag–Tang vortex flows, UniFluids-XL attains nRMSE approximately $z_i\in\mathbb R^{d_{\rm model}}$0–$z_i\in\mathbb R^{d_{\rm model}}$1, outperforming FNO, U-Net, and MPP (Li et al., 19 Mar 2026). In-domain adaptation experiments show that fine-tuning the pretrained S model on CFD subsets further reduces error by 8–22% relative to training from scratch (Li et al., 19 Mar 2026).

The ablation studies are structurally important to the interpretation of the method. They show that only $z_i\in\mathbb R^{d_{\rm model}}$2-prediction is stable across all 1D, 2D, and 3D regimes under the unified pipeline; that the unified 4D representation enables multi-resolution behavior on 1D Burgers; and that pre-training yields lower loss in fewer epochs and better final nRMSE upon fine-tuning (Li et al., 19 Mar 2026). These ablations support the claim that the framework’s performance is not solely due to scale, but also to specific design choices in representation and target parameterization.

6. Relation to unified continuum and surrogate-modeling frameworks

UniFluids belongs to a wider lineage of “unified” formulations in fluid and continuum mechanics, but it operates at a different level of abstraction. Earlier work on unified continuum PDEs sought a single governing system for viscous fluids, elastic solids, and elastoplastic solids. The Godunov–Peshkov–Romenski or SHTC family uses state variables such as density, velocity, distortion field, entropy, and in some variants thermal impulse, with a single energy potential generating fluxes and thermodynamic forces (Peshkov et al., 2017, Jackson et al., 2019, Kincl et al., 2022). In these formulations, the strain-dissipation time $z_i\in\mathbb R^{d_{\rm model}}$3 controls the asymptotic transition between viscous-fluid and elastic-solid behavior, with $z_i\in\mathbb R^{d_{\rm model}}$4 yielding Navier–Stokes-like limits and $z_i\in\mathbb R^{d_{\rm model}}$5 yielding purely hyperbolic elastodynamics (Peshkov et al., 2017, Kincl et al., 2022).

That older usage of “unified” is ontological and constitutive: one PDE system describes multiple material classes. UniFluids, by contrast, is unified at the operator-learning level. It does not propose a new continuum model; it learns mappings between histories and futures over many existing PDE datasets (Li et al., 19 Mar 2026). The distinction matters because the framework’s generality depends on shared representation and conditional modeling rather than on a universal conservation-law derivation.

A related distinction appears in numerical unification. The SHTC-SPH method provides a single meshless solver covering elastic solids, Newtonian fluids, and via $z_i\in\mathbb R^{d_{\rm model}}$6, visco-elasto-plastic materials, while preserving mass, momentum, energy, and thermodynamic compatibility at the discrete level (Kincl et al., 2022). The mixed compressible–incompressible interfacial algorithm of Denner and collaborators similarly unifies finite-volume treatment across Mach regimes within one pressure-based framework (Denner et al., 2022). The two-scale gas–liquid multi-fluid model of 2024 unifies separated and disperse regimes by coupling resolved interface variables and subscale geometric variables with capillarity and inter-scale mass transfer (Loison et al., 2024). These are solver or model unifications; UniFluids is a learned surrogate unification.

The closest neural-operator analogue in the supplied literature is OmniFluids, which integrates physics-only pre-training, coarse-grid operator distillation, and few-shot fine-tuning using factorized Fourier layers, a mixture of operators, and a multi-frame decoder (Zhang et al., 12 Jun 2025). OmniFluids emphasizes exact PDE embedment via spectral-Crank–Nicolson losses and limited-data adaptation, whereas UniFluids emphasizes unified 4D tokenization, conditional flow-matching, and one-shot future generation (Zhang et al., 12 Jun 2025, Li et al., 19 Mar 2026). A plausible implication is that the two frameworks represent complementary strategies within the emerging literature on unified PDE surrogate modeling: one physics-pretrained and operator-distilled, the other generative and conditional-flow-based.

7. Interpretation, limitations, and significance

UniFluids is significant primarily because it reframes unified PDE operator learning as a conditional generative modeling problem over heterogeneous spatiotemporal tensors. Its design synthesizes three ideas: canonical cross-dimensional representation, Transformer-based conditioning with 4D positional structure, and flow-matching with $z_i\in\mathbb R^{d_{\rm model}}$7-prediction (Li et al., 19 Mar 2026). The framework’s reported strengths are strong prediction accuracy, scalability to 512M parameters, and cross-scenario generalization capability (Li et al., 19 Mar 2026).

At the same time, several boundaries of interpretation should be kept clear. The reported results are benchmark results on PDEBench subsets and zero-shot tests on selected unseen 2D flows, not a proof of universal cross-physics transfer. The paper reports strong prediction accuracy and good scalability, but does not claim formal stability, convergence, or physical consistency guarantees of the sort associated with classical discretizations (Li et al., 19 Mar 2026). By comparison, older unified continuum frameworks explicitly derive entropy production, hyperbolicity, or asymptotic limits to Navier–Stokes or elastodynamics (Peshkov et al., 2017, Jackson et al., 2019).

Another potential misconception is that “unified” means identical treatment of all PDE semantics. UniFluids still relies on resampling, channel masking, dimension indicators, and conditioning machinery to encode differences among tasks (Li et al., 19 Mar 2026). The unification is therefore conditional and representation-based, not an erasure of PDE-specific structure. In this respect it resembles foundation-model approaches in other scientific domains more than it resembles a universal constitutive theory.

Within the broader research landscape, UniFluids marks a shift in what “unification” can mean in computational mechanics and scientific machine learning. In earlier work, the term usually referred to shared governing equations or shared numerical infrastructure for fluids, solids, interfaces, or multimaterials (Peshkov et al., 2017, Kincl et al., 2022, Denner et al., 2022, Jackson et al., 2019). In UniFluids, the unification target is the learned operator itself: a single model trained over diverse PDEs, dimensions, and channel structures (Li et al., 19 Mar 2026). This suggests a widening of the concept from continuum theory to representation learning, with operator surrogates becoming a new site of unification in computational science.