Towards a Foundation Model for Partial Differential Equations Across Physics Domains

AI summary

Overview

Research area: Physics-informed machine learning / neural operators for partial differential equations (PDEs), with a focus on foundation-model style pretraining across multiple physics domains.

Technical level: Advanced.

Scope: The paper introduces PDE-FM, a modular foundation model that combines spatial–spectral tokenization, physics-aware conditioning, a Mamba state-space backbone, and a Fourier operator decoder, pretrained on twelve datasets from The Well benchmark and evaluated on twelve 2D and 3D physical simulation tasks.

What This Paper Is About

Most neural operator models for PDEs are built for one dataset or one narrow family of equations, and their accuracy degrades as soon as boundary conditions, scales, or governing dynamics change. The authors ask whether a single model can be pretrained across many different physical systems at once and then transfer to new physics regimes without architectural changes. PDE-FM is their answer: a hybrid architecture that learns "operators as distributions over physics" rather than as isolated mappings.

Key Contributions

1. A modular hybrid architecture. PDE-FM combines (i) dual spatial patch and low-frequency spectral tokenization, (ii) FiLM-based physics-aware conditioning using metadata such as boundary conditions, constitutive parameters, and time grids, (iii) a Mamba state-space backbone for long-range dependencies in linear time, and (iv) a shallow 2D FNO decoder.

2. A multi-dataset pretraining recipe with dataset-specific adapters. Per-dataset 1×1 input/output adapters normalize heterogeneous channel counts into a shared latent channel budget, and mini-batches are drawn with a temperature-scaled, difficulty-aware sampling probability that combines dataset size (exponent τ) with an EMA loss term (exponent α ∈ [0,1]) to reduce overfitting to large or easy datasets and to mitigate negative transfer.

3. A dual spatial–spectral training objective. The loss combines VRMSE (variance-reduced RMSE) with a spectral L2 term over a truncated frequency set, with weights that increase with frequency magnitude, plus optional conservation-loss and PDE-residual terms when invariants are available. The spectral weight λ and, if used, the conservation weights α_j and β are cosine-annealed from higher spectral weight to warm-start global structure.

4. A systematic ablation and benchmark. The authors ablate backbone (Transformer vs. Mamba), decoder (FNO vs. Conv), normalization (Layer vs. None), and three conditioning mechanisms (FiLM, Spectral Tokenizer, Cross-Attention), then benchmark the best configuration against FNO, TFNO, U-net, CNextU-net, and the PhysiX foundation model.

Main Findings

• Best ablation configuration: Mamba + FiLM + FNO + (SpecTok, X-Attn) + LayerNorm achieved the lowest mean VRMSE of 0.2581 across all tasks in the short sweep (EPOCHS = 8, STEPS = 600, BATCH = 8, LR = 10⁻⁴). The Transformer + FNO variant reached 0.2779, and the weakest listed configuration reached 0.3350. FNO decoders consistently outperformed convolutional alternatives, and layer normalization improved convergence, especially for the Mamba backbone.

• State-of-the-art on six datasets: With extended training (30 epochs, 1000 steps per epoch), PDE-FM achieved the lowest VRMSE on rayleigh_benard, shear_flow, turbulence_gravity_cooling, supernova_explosion_64, gray_scott_reaction_diffusion, and post_neutron_star_merger. It ranked second in one dataset (turbulent_radiative_layer_2D, at 0.2321 versus CNextU-net's 0.1956). The introduction states a different count, saying PDE-FM "ranks second in five others."

• Lowest average error across datasets: PDE-FM recorded a mean VRMSE of 0.165, compared with CNextU-net at 0.304, FNO at 0.441, and TFNO at 0.469. The authors note that PhysiX does not report results on 3D PDE domains, so it is not included in this average.

• Relative improvement: The abstract and results section report a 46% reduction in mean VRMSE relative to prior operator-learning baselines; the introduction describes this as "over 40%."

• Strongest in turbulent and advective regimes: On shear_flow PDE-FM scored 0.0345 versus PhysiX's 0.0700 and CNextU-net's 0.8080; on rayleigh_benard it scored 0.0415 versus PhysiX's 0.1470. The authors attribute this to high-frequency retention from spectral tokenization and temporal stability from the Mamba recurrent backbone.

• Astrophysical and relativistic results: On post_neutron_star_merger, PDE-FM reached 0.2995 versus TFNO's 0.3793, which the paper describes as a 19% VRMSE reduction, with linear scaling in memory and time.

• Radiative and multiphase flows: On turbulence_gravity_cooling, PDE-FM set a new best score of 0.0796 VRMSE.

• Limitations identified: The viscoelastic_instability task remains the primary limitation. PDE-FM's VRMSE there is now 0.52, described as halving compared with earlier iterations, but still behind CNextU-net's 0.25 and PhysiX's 0.2370. On linear acoustic problems, PDE-FM scored 0.0487 on acoustic_scattering_maze (versus CNextU-net's 0.0153) and 0.0414 on helmholtz_staircase (versus FNO's 0.00046), which the authors attribute to convolutional priors being preferable for high-frequency precision in stationary or linear cases. On convective_envelope_rsg, PDE-FM's 0.0896 is higher than FNO's 0.0269, TFNO's 0.0283, U-net's 0.0555, and CNextU-net's 0.0799.

• Baseline context: PhysiX has 4.5B parameters and a token-based autoregressive design; it achieved the best overall score on active_matter (0.0904, versus PDE-FM's 0.1974) and competitive performance on elastic systems.

• Generalization pattern: The parity plot shows nearly all points below the y = x diagonal except the viscoelastic and acoustic cases, which the authors present as evidence of broad cross-physics generalization. They also note that VRMSE = 1 corresponds to a trivial mean-field predictor.

Methodology in Plain English

The input to the model is a field of shape (channels × height × width). PDE-FM does four things with it.

First, it builds two views of the input: spatial patches processed by a patch convolution, and a global spectral token produced by taking a truncated 2D FFT that keeps only the lowest m × (m/2 + 1) frequencies per channel, with real and imaginary parts stacked. Physics metadata—boundary conditions, constitutive parameters, time grids—is encoded as a vector and used to modulate the spatial tokens through FiLM (a learned scale and shift), with a learned context token prepended when metadata is present. FFTs run in FP32 for numerical stability, missing metadata defaults to zero vectors, and metadata uses dataset-level standardization.

Second, spatial tokens pass through ConvNeXt-style residual blocks while the spectral token passes through an MLP, and the two exchange information through shallow bidirectional cross-attention. The single spectral token gates global context into the spatial tokens without quadratic cost.

Third, the context, spatial, and spectral tokens are concatenated and processed by L Mamba state-space layers with residual connections. Mamba gives sub-quadratic O(N_p d) compute and memory rather than the O(N_p²) of attention, which is what allows large grids and long contexts. Layer normalization before the backbone and gradient clipping stabilize training.

Fourth, the spatial slice is reshaped back to a latent grid, upsampled, and decoded with a shallow 2D FNO that keeps only low-frequency modes, biasing output toward spectral smoothness.

Training minimizes VRMSE plus a weighted spectral L2 loss over truncated frequencies, with optional conservation and PDE-residual penalties. For multi-dataset pretraining, per-dataset 1×1 adapters map each dataset's channels into a shared latent channel budget; batches are sampled with a probability that mixes temperature scaling by dataset size (τ) with difficulty weighting from each dataset's EMA loss.

The pretraining corpus is twelve nonlinear 2D and 3D datasets from The Well, a 15 TB curated collection of 16 spatiotemporal simulation datasets. All datasets use a unified HDF5 format storing arrays of shape (n_traj, n_steps, H, W, [D]) in fp32 and an 80/10/10 train/validation/test split. Sampling used τ = 0.5. Inputs were interpolated to standardized grids ranging from 128² to 512² for 2D systems and from 64³ to 192 × 128 × 66 for 3D systems. Training used AdamW with an initial learning rate of 5 × 10⁻⁴, cosine-annealing decay, gradient clipping at 1.0, mixed precision, and distributed data-parallel training across multiple GPUs. The evaluation metric throughout is VRMSE, the benchmark's primary metric, which normalizes errors by spatial variance so quantities with different physical scales are comparable.

The twelve datasets in Table 1 are: active_matter (Cartesian 2D, 256 × 256, 81 steps, 360 trajectories), turbulent_radiative_layer_2D (Cartesian 2D, 128 × 384, 101 steps, 90 trajectories), viscoelastic_instability (Cartesian 2D, 512 × 512, variable steps, 260 trajectories), shear_flow (Cartesian 2D, 128 × 256, 200 steps, 1,120 trajectories), gray_scott_reaction_diffusion (Cartesian 2D, 128 × 128, 1,001 steps, 1,200 trajectories), rayleigh_benard (Cartesian 2D, 512 × 128, 200 steps, 1,750 trajectories), post_neutron_star_merger (log-spherical 3D, 192 × 128 × 66, 181 steps, 8 trajectories), supernova_explosion_64 (Cartesian 3D, 64³, 59 steps, 1,000 trajectories), turbulence_gravity_cooling (Cartesian 3D, 64³, 50 steps, 2,700 trajectories), convective_envelope_rsg (spherical 3D, 256 × 128 × 256, 100 steps, 29 trajectories), helmholtz_staircase (Cartesian 2D, 1,024 × 256, 50 steps, 512 trajectories), and acoustic_scattering_maze (Cartesian 2D, 256 × 256, 100 steps, 8,000 trajectories).

For SOTA comparison, no ensembling, test-time augmentation, or extra data were used, and official splits were followed for all datasets. Baseline results for FNO, TFNO, U-net, and CNextU-net were extracted from Ohana et al. (2024) and PhysiX results from Nguyen et al. (2025).

Why This Matters

Impact on research. The paper argues that pretraining across heterogeneous physical processes can yield transferable representations of dynamics, rather than requiring one bespoke surrogate per PDE family. It also shows a way to get properties that are hard to combine: the scalability and generalization of large sequence models alongside the inductive structure of domain-specific solvers. The claim that datasets sharing invariant structures—such as incompressibility or conservation of vorticity—mutually reinforce one another during pretraining suggests a mechanism for emergent transfer across unseen domains. The authors also position PDE-FM as bridging neural operators and foundation models by learning "operators as distributions over physics."

Real-world applications (as suggested by the physics domains covered):

• Geophysical and climate-relevant flow modeling, via the turbulent and convective regimes studied (Rayleigh–Bénard convection, shear flow).
• Astrophysical simulation, including core-collapse supernovae and post-neutron-star-merger relativistic magnetohydrodynamics, where fast surrogates could substitute for expensive volumetric solvers.
• Multiphase and radiative cooling processes, relevant to combustion, energy systems, and thermal management.
• Acoustic scattering in layered media, relevant to sensing, imaging, and wave-propagation engineering.

Industry relevance. The authors are affiliated with IBM Research Brazil (São Paulo and Rio de Janeiro) and IBM Research Zurich. The pretrain-once, adapt-everywhere framing—where a model is transferred to new physical regimes without architectural or data-specific modifications—is directly relevant to industrial simulation workflows that currently require per-case model development. The use of 1×1 adapters and a shared core also matters practically for organizations that want one deployed model across many simulation products, and the linear-time Mamba backbone matters for memory and compute budgets at large grid sizes.

Future Directions

1. Conservation-based and energy-preserving loss regularization to improve stability across long rollouts, building on the optional conservation and PDE-residual terms already sketched in the loss formulation.

2. Adaptive spectral decoders and hybrid neural operators that dynamically allocate resolution across spatial scales, motivated by the finding that stationary and linear problems still favor convolutional priors for high-frequency precision.

3. Curriculum or multi-domain pretraining strategies that better balance data diversity and physical consistency across 2D and 3D regimes.

4. Explicit latent memory or physics-informed temporal embeddings for elasticity-dominated systems, since the viscoelastic instability benchmark suggests long-term stress–strain coupling requires mechanisms beyond the current architecture.

5. Scaling to the full breadth of The Well, including magnetohydrodynamic, elastic, and radiative datasets, toward truly universal representations for physics-informed machine learning.

Target Audience

Researchers and practitioners in scientific machine learning, neural operators, and physics-informed modeling who are interested in foundation-model approaches to simulation. It is also relevant to computational physicists and engineers working in fluid dynamics, astrophysics, radiative transfer, elasticity, and acoustics who want fast surrogate solvers, and to industry teams at organizations running large-scale multi-physics simulation pipelines. Some familiarity with neural operators, Fourier methods, and PDE terminology is assumed.

Authors’ abstract

We present PDE-FM, a modular foundation model for physics-informed machine learning that unifies spatial, spectral, and temporal reasoning across heterogeneous partial differential equation (PDE) systems. PDE-FM combines spatial-spectral tokenization, physics-aware conditioning, and a Mamba-based state-space backbone with an operator-theoretic decoder, enabling scalable and data-efficient modeling of complex physical dynamics. In contrast to task-specific neural operators, PDE-FM is pretrained once on diverse PDE datasets and can be transferred to new physical regimes without architectural or data-specific modifications. Evaluated on twelve 2D and 3D datasets from The Well benchmark - spanning hydrodynamic, radiative, elastic, and astrophysical phenomena - PDE-FM achieves state-of-the-art accuracy in six domains, reducing mean VRMSE by 46% relative to prior operator-learning baselines. The model demonstrates robust cross-physics generalization, excelling in turbulent and radiative systems while maintaining strong performance in linear and steady-state regimes. These results suggest that large-scale pretraining across diverse physical processes can yield transferable representations of dynamics, marking a step toward unified, foundation-level surrogates for multi-physics simulation and scientific discovery.