Technical Summary: Physics-Informed Neural PDE Solvers via Spatio-Temporal MeanFlow

1. Problem Statement

Solving Partial Differential Equations (PDEs) efficiently and accurately is fundamental to modeling complex physical phenomena. While deep learning paradigms like Physics-Informed Neural Networks (PINNs) and Neural Operators (e.g., FNO, DeepONet) have advanced the field, they face specific limitations regarding the continuous integral nature of physical systems:

• PINNs rely on pointwise residuals, effectively bypassing the integral perspective of PDE evolution.
• Neural Operators (like FNO) often rely on pre-discretized temporal grids, treating integration intervals as fixed hyperparameters and functioning as discrete, step-by-step integrators.
• Generative Models (e.g., Diffusion models) typically operate on a noise-to-solution paradigm, requiring multi-step sampling or ODE evaluations, which can be computationally expensive.

The core challenge is to develop a solver that directly learns the finite-interval evolution of physical states in a continuous manner, allowing for flexible integration lengths without the error accumulation of autoregressive steps or the rigidity of fixed discretization.

2. Methodology: Spatio-Temporal MeanFlow

Inspired by MeanFlow (Geng et al., 2025), a continuous-time integrator for generative ODEs that learns the average velocity over a finite interval, the authors propose Spatio-Temporal MeanFlow. This framework transforms multi-step numerical integration into an efficient, single-step prediction with a freely controllable integration length.

Core Formulation

The method replaces the generative velocity field $v$ in the MeanFlow identity with the physical PDE operator $f$. For a PDE system defined by $\gamma \frac{\partial u}{\partial t} = f[u, \nabla, x, t]$, the state evolution from $(\xi, \tau)$ to $(x, t)$ is modeled via a straight spatio-temporal path $p$. The Spatio-Temporal MeanFlow function $m$ is defined as the integral average of state evolution along this path:
$$m[u(\xi, \tau), p, a] \triangleq \frac{\int_p du(p)}{\int_p \|dp\|_2}$$
According to the Fundamental Theorem of Line Integrals, this satisfies:
$$l(\xi, \tau, x, t) \cdot m = u(x, t) - u(\xi, \tau)$$
where $l$ is the total spatio-temporal Euclidean distance.

Physics-Informed Constraints

By differentiating the fundamental identity with respect to initial coordinates $(\xi, \tau)$, the authors derive physics-informed constraints that serve as the training objective. These constraints couple time evolution with spatial consistency:

1. Temporal Constraint: Relates the mean flow to the PDE operator and time derivatives.
2. Spatial Constraint: Relates the mean flow to spatial gradients ($\nabla u, \Delta u$) and spatial derivatives of the mean flow.

Space-Time Decoupling Strategy

To address the pathological curvature in the optimization landscape caused by scale mismatches between temporal and spatial dimensions (leading to ill-conditioned Hessians), the framework employs a decoupling strategy. The global loss is decomposed into two independent first-order components:

• Temporal MeanFlow Loss ($L_{T-MF}$): Constrains the network to satisfy temporal evolution along the time-line.
• Spatial MeanFlow Loss ($L_{S-MF}$): Enforces consistency between predicted spatial derivatives and numerical spatial gradients.

Theoretical analysis (Theorem 3.2) proves that the global spatio-temporal consistency is bounded by these decoupled losses, and the decoupling significantly improves the condition number of the Hessian matrix, enabling faster and more robust convergence.

Unified Framework

The formulation naturally accommodates both time-dependent ($\gamma \neq 0$) and stationary ($\gamma = 0$) PDEs. For static systems, the time dimension acts as a dummy variable, and the model specializes to a purely spatial MeanFlow loss, providing a unified framework for diverse PDE types.

3. Key Contributions

1. MeanFlow-Inspired PDE Solver: A novel framework that transforms conventional multi-step numerical integration into finite-interval predictions. It enables flexible integration steps at inference time, allowing the model to predict solutions over arbitrary time intervals in a single step or via a few composed steps.
2. Spatio-Temporal MeanFlow Loss: A physics-informed constraint that extends the original MeanFlow concept from the temporal domain to the spatio-temporal domain. This significantly enhances solving accuracy and generalization, particularly under out-of-distribution (OOD) initial conditions and varying spatial resolutions.
3. Comprehensive Validation: Extensive evaluations on canonical synthetic benchmarks (Burgers, Navier-Stokes, Darcy, Poisson) and a real-world dataset (HYCOM ocean data), demonstrating superior accuracy-efficiency trade-offs compared to representative baselines.

4. Experimental Results

The authors evaluated their method against three categories of baselines: physics-informed coordinate-based solvers (PINN, I-PINN), neural operators (FNO, DeepONet, PINO, CFO), and generative solvers (DiffusionPDE, Rectified Flow, PBFM).

• Accuracy: On the Burgers and Navier-Stokes benchmarks, the proposed method achieved the lowest relative $L_2$ error (e.g., $0.26 \times 10^{-2}$ on Burgers vs. $0.80 \times 10^{-2}$ for the next best, CFO). It also outperformed baselines on steady-state PDEs (Darcy, Poisson).
• Efficiency: The method achieved sub-millisecond inference (0.239 ms/sample on Burgers) with extremely low computational cost (0.0118 GFLOPs/sample), significantly outperforming generative models that require multi-step sampling.
• Generalization:
  • Resolution: Trained on a resolution of $s=128$, the model achieved stable zero-shot generalization to higher resolutions ($s=256, 512$) without retraining, a capability attributed to the spatial MeanFlow loss.
  • OOD Initial Conditions: The model demonstrated robust performance on initial conditions with shifted parameters (e.g., different Gaussian random field parameters), outperforming baselines that rely heavily on data fitting.
• Real-World & Inverse Problems: The framework was successfully applied to the HYCOM ocean dataset (handling noise and assimilation artifacts) and the inverse Darcy flow problem (recovering permeability fields), showcasing versatility beyond synthetic forward prediction.

5. Significance and Claims

The paper claims that Spatio-Temporal MeanFlow offers a unified, efficient, and generalizable approach to PDE solving. By transcending purely data-driven learning and leveraging the continuous integral nature of physical systems, the framework:

• Optimizes the Accuracy-Efficiency Trade-off: It achieves the optimal Pareto frontier, delivering the lowest error with highly competitive inference speeds across all baselines.
• Enables Flexible Integration: Unlike fixed-step solvers, it allows users to control the integration length at inference time, balancing speed and accuracy dynamically.
• Improves Robustness: The spatio-temporal constraints provide strong regularization, enabling the model to generalize to higher resolutions and OOD conditions where traditional methods often fail.

Limitations and Future Work: The authors note that while the mathematical formulation supports continuous spatial evolution, the current empirical implementation focuses on regular grids. Future work aims to activate continuous spatial modeling for irregular domains and extend the framework to datasets lacking analytical PDEs. They also recommend using the framework for iterative design and optimization while relying on mathematically guaranteed numerical methods for final safety certifications in critical domains.