UniFluids: Unified Neural Operator Learning with Conditional Flow-matching

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 (10)

• 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.
• 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.
• 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.
• 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.
• 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).