We introduce ECHO, a transformer-operator framework for generating million-point PDE trajectories. While existing neural operators (NOs) have shown promise for solving partial differential equations, they remain limited in practice due to poor scalability on dense grids, error accumulation during dynamic unrolling, and task-specific design. ECHO addresses these challenges through three key innovations. (i) It employs a hierarchical convolutional encode-decode architecture that achieves a 100 × spatio-temporal compression while preserving fidelity on mesh points. (ii) It incorporates a training and adaptation strategy that enables high-resolution PDE solution generation from sparse input grids. (iii) It adopts a generative modeling paradigm that learns complete trajectory segments, mitigating long-horizon error drift. The training strategy decouples representation learning from downstream task supervision, allowing the model to tackle multiple tasks such as trajectory generation, forward and inverse problems, and interpolation. The generative model further supports both conditional and unconditional generation. We demonstrate state-of-the-art performance on million-point simulations across diverse PDE systems featuring complex geometries, high-frequency dynamics, and long-term horizons.
Figures & tables
Figure 1: Experimental analysis of ECHO. Left: reconstruction error versus compression ratio on Vorticity . Hierarchical, iterative compression improves accuracy, especially at high compression ratios. Middle: long-range rollouts on Gray-Scott . Given the first 3 frames, each model predicts the next 197 frames; frames beyond t=40 lie outside the training horizon. Errors are averaged over test trajectories. Full-trajectory generation accumulates less error than latent autoregressive and physical-space baselines. Right: generative ECHO versus its deterministic variant on forward prediction, interpolation and backward prediction on Rayleigh-Bénard . All panels report the relative L2 error (lower is better).
Figure 2: Architecture of the ECHO framework. ECHO comprises two components: (B) a convolutional auto-encoder and (C) a DiT-based generative process. The auto-encoder uses continuous convolutions to ingest irregular input grids of arbitrary size, maps the dynamics to a regular latent grid, and hierarchically compresses it; the decoder mirrors this hierarchy and applies a final continuous convolution, enabling queries at arbitrary output locations. The DiT module is trained with a flow-matching objective to denoise latent tokens, optionally conditioned on PDE parameters. This design allows ECHO to handle irregular grids and support multiple inference tasks (A).
Figure 3: Three-stage training strategy for million-point trajectory generation. ECHO’s auto-encoder is first trained for low-resolution trajectory reconstruction, then refined on high-resolution single frames. Both stages use reconstruction objectives. The generative processor is then trained separately in the frozen latent space using flow matching.
Rayleigh–Bénard
Gray–Scott
Active Matter
ASM
( ×64 )
( ×32 )
( ×176 )
( ×48 )
Model
Latency (s)
Back.
For.
Back.
For.
Back.
For.
Back.
For.
Deterministic models
FNO
4.42e-2
2.47e3
4.23e-1
–
–
–
–
1.87
1.52
Transolver++
3.52e-1
6.34e-1
3.31e-1
4.43e-1
2.34e-1
7.33e-1
6.91e-1
1.03
9.64e-1
BCAT
1.18e-1
1.91e-1
1.06e-1
2.19e-1
8.82e-2
4.98e-1
4.56e-1
1.95e-1
2.18e-1
Table 1: Comparison of model performance across four dynamical systems. Errors are relative L2 errors computed from a single sample per test trajectory, for backward ( Back. ) and forward ( For. ) prediction. Latency is the time needed to generate a whole trajectory, averaged over the datasets. Lower is better for all metrics. The ratio under each dataset name is the compression ratio of the latent models (ECHO, ENMA); the other baselines operate at full resolution ( ×1 ). Best results in bold , second best underlined ; “–” indicates non-convergence.
MHD
TGC
Model
1
1–8
1–16
1
1–8
1–16
FNO
1.69e-1
–
–
3.15e-2
–
–
AViT
3.40e-1
5.68e-1
7.51e-1
1.26e-1
6.22e-1
1.32
BCAT
3.29e-1
5.43e-1
6.70e-1
1.04e-1
9.89e-1
1.62
ECHO
2.70e-1
3.82e-1
5.50e-1
1.13e-1
5.55e-1
6.67e-1
Table 2: Forward time-stepping on 3D benchmarks. Relative L2 error over the first frame (1), the first 8 frames (1–8) and the first 16 frames (1–16) of the predicted trajectory. Lower is better. Best results in bold , second best underlined ; “–” indicates non-convergence.
Vorticity
Shallow-Water
Eagle
Cylinder Flow
Irr.
Encoder
Rec.
For.
Rec.
For.
Rec.
For.
Rec.
For.
Full resolution ( Xte=100% )
Graph
GINO
9.99e-1
1.00
8.68e-1
1.11
5.58e-1
1.89
7.94e-1
8.65e-1
INR
CORAL
5.13e-1
1.34
2.29e-1
6.89e-1
6.04e-1
1.54
2.94e-1
5.08e-1
Attention
AROMA
5.13e-1
8.42e-1
2.51e-2
4.21e-2
2.81e-1
3.09e-1
2.76e-2
2.21e-1
ENMA
4.38e-1
4.38e-1
7.17e-2
7.35e-2
2.85e-1
3.29e-1
7.50e-2
1.04e-1
Table 3: Comparison of encoder–decoder modules from operator surrogates on irregular scenarios. Rec. and For. denote reconstruction and forward generation; Irr. gives the strategy for handling irregular points. Xte is the spatial sampling ratio of the conditioning frames at inference. Colors indicate encoder scalability: green = no issues; yellow = minor modification; red = model size modification required to fit in GPU memory (see appendix E ). Metrics are relative L2 errors on the full grid of the test set (lower is better). Best results in bold , second best underlined .
Model
GINO
CORAL
AROMA
CALM-PDE
Transolver++
ECHO
For.
OOM
OOM
OOM
OOM
OOM
3.88e-1
Table 4: Super-resolution forecasting on a 1024×1024 Vorticity grid (200% of the training grid, more than 20M points). Relative L2 error of forward generation ( For. ); OOM = out of memory on an 80 GB H100.
Figure 4: Scaling comparison: memory consumption at inference vs. number of points in the input mesh.
Appendix figures & tables42 assets
Supplementary material from the paper’s appendix.
Appendix
Model
Reference
Arbitrary domains
Compression
High-res. encoding
Hierarchical encoding
Generative modeling
Multi-task prediction
Graph
MPPDE
Brandstetter et al. (2022)
✓
✗
✗
✗
✗
✗
RIGNO
Mousavi et al. (2025)
✓
✓
✗
✗
✗
✗
Operator
Transolver++
Luo et al. (2025)
✓
✗
✓
✗
✗
✗
GINO
Li et al. (2023b)
✓
✓
✗
✗
✗
✗
Appendix
Table 5: Comparison of operator learning approaches for spatio-temporal PDE forecasting.
Dataset
# points
Grid
PDE parameters
Size T\times(H\times W{\color[rgb]{0.5,0.5,0.5}\times D})\times C
2D
Vorticity ( Koupaï et al., 2024 )
20M
Dense
Parametric
20×1048576×1
Shallow-Water ( Yin et al., 2022 )
2.6M
Spherical
Single instance
40×32768×2
Gray–Scott ( Ohana et al., 2024 )
1.3M
–
Varying
40×16384×2
Rayleigh–Bénard ( Ohana et al., 2024 )
3.9M
–
Varying
15×65536×4
ASM ( Ohana et al., 2024 )
2.9M
–
Varying
15×65536×3
Appendix
Table 6: Datasets used in the experiments. Sizes are given as T\times(H\times W{\color[rgb]{0.5,0.5,0.5}\times D})\times C , where the depth D only applies to 3D datasets. ASM = Acoustic Scattering Maze; MHD = Magnetohydrodynamics; TGC = Turbulence Gravity Cooling; BC = boundary conditions.
Figure 5: Sample of the Vorticity dataset.
Figure 6: Sample of the Shallow-water dataset.
Figure 7: Sample of the Eagle dataset.
Figure 8: Sample of the Cylinder Flow dataset.
Figure 9: Samples from the Rayleigh-Bénard dataset.
Figure 10: Sample of the Gray-Scott Reaction-Diffusion dataset
Figure 11: Trajectory sample of Active Matter dataset.
Figure 12: Sample of the Acoustic Scattering Maze dataset.
Figure 13: Trajectory sample of MHD dataset sliced along each axis.
Figure 14: Trajectory sample of TGC dataset sliced along each axis.
Figure 15: Encoder–decoder architecture of ECHO. The auto-encoder uses a continuous convolution layer to process irregular input grids with an arbitrary number of points, mapping the dynamics onto a regular latent grid. This latent representation is then compressed by a hierarchical CNN encoder. The decoder mirrors this process with a hierarchical CNN decoder followed by a continuous convolution layer, enabling queries at any desired output location.
Hyperparameter
Value
Patch size
8
Transformer depth
6
Hidden size
512
MLP ratio
2
Number of heads
8
QK normalization
True
Appendix
Table 7: BCAT hyperparameters used across all datasets.
Hyperparameter
Value
Number of layers
8
Hidden size
128
Attention heads
8
MLP ratio
4
Dropout
0
Physics slices
32
Appendix
Table 8: Transolver++ hyperparameters used across all datasets.
Hyperparameter
Active Matter
Rayleigh–Bénard
Gray–Scott
ASM
VAE embedding dim
4
4
4
8
Number of tokens
16
16
64
256
Patch size
2
2
1
2
Spatial Transformer depth
8
8
8
8
Causal Transformer depth
8
8
8
8
Hidden size
512
512
512
512
Appendix
Table 9: ENMA generation hyperparameters.
Hyperparameter
Value
Epochs
200
Batch size
32
Learning rate
10−3 (cosine decay)
Weight decay
10−4
Grad clip norm
1
Betas (β1,β2)
(0.9, 0.95)
Appendix
Table 10: Training hyperparameters used across datasets.
Dataset
Transolver++
FNO
AViT
BCAT
ENMA
ECHO
Rayleigh–Bénard
5.5M
75M
51M
27M
72M
48M
Active Matter
5.5M
75M
53M
28M
72M
48M
Appendix
Table 11: Number of parameters of the models compared in table 1 .
Parameter
Vorticity
Shallow-Water
Eagle
Cylinder Flow
Token dimension
32
32
32
32
Interpolation module
Regularized grid
192×192
64×128×1
64×32
64×16
Hidden dimension
64
64
64
64
Receptive field
0.01
0.01
0.01
0.01
Softmax temperature
1
1
1
1
Appendix
Table 12: Hyper-parameters of the encoder and decoder components of ECHO.
Dataset
GINO
CORAL
AROMA
ENMA
CALM-PDE
ECHO
Vorticity
18M
23M
0.5M
1.9M
0.5M
33M
Shallow-Water
19M
21M
0.5M
5M
0.2M
30M
Appendix
Table 13: Number of parameters of the encoder–decoders compared in table 3 .
Model
Vorticity
Shallow-Water
Eagle
Cylinder Flow
Trajectory
10×(512×512)×1
10×(128×256)×2
200×3000×3
60×2000×3
GINO
10×576×32
40×512×32
100×512×32
60×64×32
CORAL
10×18432
40×16384
200×4096
50×2048
AROMA
10×64×32
40×128×32
200×128×32
60×64×32
CALM-PDE
10×576×32
40×512×32
200×128×32
60×64×32
ECHO
5×24×24×32
10×16×8×32
50×16×8×32
15×16×4×32
Appendix
Table 14: Physical and latent space shapes in the experiments of table 3 , in the format temporal size × spatial size × token dimension. The first row gives the shape of the trajectories. On datasets with many input points, some baselines use a smaller latent shape (in red ) to fit in GPU memory.
Model
Task
Vorticity
Shallow-Water
Eagle
Cylinder Flow
GINO
Rec.
8
16
32
32
TS
4
16
8
16
CORAL
Rec.
1
2
16
32
TS
1
4
16
16
AROMA
Rec.
1
1
4
32
TS
1
1
4
16
Appendix
Table 15: Training batch sizes used across datasets, for reconstruction ( Rec. ) and time-stepping ( TS ).
Hyperparameter
Low resolution
High resolution
Epochs
200
10
Batch size
8
8
Learning rate
10−3 (cosine decay)
10−4 (cosine decay)
Weight decay
10−4
10−4
Grad clip norm
1
1
Betas (β1,β2)
(0.9, 0.95)
(0.9, 0.95)
Appendix
Table 16: Training hyperparameters used across datasets for the low- and high-resolution stages.
Hyperparameter
Generative stage
Epochs
200
Batch size
32
Learning rate
10−3 (cosine decay)
Weight decay
10−4
Grad clip norm
1
Betas (β1,β2)
(0.9, 0.95)
Appendix
Table 17: Training hyperparameters used across datasets for the generative stage.
Setting
Rel. L2
In-distribution
1.17e-1
OOD zero-shot
3.39e-1
OOD few-shot
1.12e-1
Appendix
Table 18: OOD generalization on Vorticity: performance of ECHO under distribution shift, with zero-shot and few-shot adaptation. Relative L2 error (lower is better).
Figure 34
Figure 18: Spectral analysis on a trajectory generation with ECHO on the Vorticity dataset.
Encoder–decoder
#Params
Training throughput (samples/s) ↑
Latency (s) ↓
FPD ↓
GINO
19M
55.7
1.83
1.68e-1
CORAL
21M
4.3
0.75
1.31e-1
AROMA
0.5M
32.8
0.19
1.55e-2
CALM-PDE
2M
60.5
0.24
1.02e-1
ECHO
30M
35.0
0.21
1.12e-3
Appendix
Table 19: Efficiency and generative fidelity on Shallow-Water. Every encoder–decoder is combined with our generative model. Throughput and latency are measured on identical hardware; FPD measures the distributional alignment with ground-truth trajectories. Best results in bold , second best underlined .
Dataset
Task
Model
CRPS ↓
RMSCE ↓
SSR →1
Rel. L2↓
Gray–Scott
Forward
ENMA
8.68e-3
3.16e-1
0.43
3.96e-2
ECHO
8.66e-3
9.34e-2
0.93
3.90e-2
Backward
ECHO
9.69e-3
7.94e-2
1.06
4.41e-2
Rayleigh–Bénard
Forward
ENMA
9.47e-3
2.03e-1
0.53
9.38e-2
ECHO
1.04e-2
8.97e-2
0.86
9.02e-2
Backward
ECHO
1.30e-2
4.77e-2
0.90
1.14e-1
Appendix
Table 20: Probabilistic evaluation with 16 -member ensembles, using the same windows and number of generated frames as table 1 . Rel. L2 is computed on the ensemble mean. Bold marks the best forward result per dataset; ECHO backward rows use identical settings.
Figure 19: Initial value problem on Vorticity ( 512×512 ). Ground truth (GT), ECHO predictions, and per-frame error maps when generating the full trajectory from only the first state and PDE viscosity.
Rayleigh–Bénard
Gray–Scott
Model
Back.
For.
Interp.
Back.
For.
Interp.
Deterministic models
FNO
2.47e3
4.23e-1
–
–
Transolver++
6.34e-1
3.31e-1
4.43e-1
2.34e-1
BCAT
1.91e-1
1.06e-1
2.19e-1
8.82e-2
AViT
4.50e-1
1.01e-1
1.66e-1
7.42e-2
Appendix
Table 21: Backward ( Back. ), forward ( For. ) and temporal interpolation ( Interp. ) results across four dynamical systems. Errors are relative L2 errors (lower is better). Interpolation is evaluated for ECHO only; latencies are reported in table 1 . Best results in bold , second best underlined ; “–” indicates non-convergence.
Vorticity
Shallow-Water
Eagle
Cylinder Flow
Irr.
Encoder
Hier.
Rec.
For.
Rec.
For.
Rec.
For.
Rec.
For.
Intermediate sparsity ( Xte=50% )
Graph
GINO
✗
9.99e-1
1.00
8.69e-1
1.11
5.80e-1
1.88
7.94e-1
8.65e-1
INR
CORAL
✗
5.14e-1
1.34
2.29e-1
6.88e-1
6.11e-1
1.54
3.04e-1
5.07e-1
Attention
AROMA
✗
5.13e-1
8.43e-1
2.69e-2
4.28e-2
3.09e-1
3.34e-1
3.17e-2
2.31e-1
ENMA
✓
4.39e-1
4.39e-1
7.42e-2
7.53e-2
3.02e-1
3.42e-1
7.77e-2
1.09e-1
Appendix
Table 22: Encoder–decoder results from sparse inputs: only a fraction Xte of the input grid is observed at test time and the full grid is reconstructed. The 10% setting is out-of-distribution (training samplings range from 25% to 75% ). Rec. and For. denote reconstruction and forward generation; Irr. gives the strategy for handling irregular points, and Hier. indicates hierarchical (✓) or direct (✗) compression into the latent space. Metrics are relative L2 errors (lower is better). Best results in bold , second best underlined .
Hierarchical level
0
1
2
3
Rel. L2
3.94e-1
3.64e-1
1.09e-1
2.87e-2
Appendix
Table 23: Impact of hierarchical depth in the encoder on reconstruction (relative L2 error). A higher level means more progressive spatial downsampling, with the same final latent size.
w/o refinement
w/ refinement
Rel. L2
2.24e-1
6.88e-2
Appendix
Table 24: Effect of the refinement stage on full-grid Vorticity reconstruction (relative L2 error, lower is better).
Ablation
Configuration
Rel. L2
Solver
Euler
1.62e-1
Midpoint
1.67e-1
RK4
1.68e-1
DOPRI5
1.71e-1
# Steps
2 steps
1.72e-1
5 steps
1.69e-1
Appendix
Table 25: Ablation of the flow-matching generative process on Vorticity. We vary the ODE solver and the number of integration steps used at sampling time. Relative L2 error (lower is better).
Figure 20: Generation using ECHO’s framework on the Vorticity dataset. We generate 1024×1024 spatial points per frame. (Top) ground truth trajectory, (Middle) ECHO’s prediction, (Bottom) Error.
Figure 21: Generation using ECHO’s framework on the Gray-Scott dataset. (Top) ground truth trajectory, (Middle) ECHO’s prediction, (Bottom) Error.
Figure 22: Generation using ECHO’s framework on the Rayleigh-Benard dataset. (Top) ground truth trajectory, (Middle) ECHO’s prediction, (Bottom) Error.
Figure 23: Generation using ECHO’s framework on the Active Matter dataset. (Top) ground truth trajectory, (Middle) ECHO’s prediction, (Bottom) Error.
Figure 24: Generation using ECHO’s framework on the Acoustic Scattering Maze dataset. (Top) ground truth trajectory, (Middle) ECHO’s prediction, (Bottom) Error.
Figure 25: Generation using ECHO’s framework on the Eagle dataset. (Top) ground truth trajectory, (Middle) ECHO’s prediction, (Bottom) Error.
Figure 26: Generation using ECHO’s framework on the Eagle dataset - long trajectory generation. (Top) ground truth trajectory, (Middle) ECHO’s prediction, (Bottom) Error.
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.
Haosen Li, Qi Meng, Jiahao Li +4
Academy of Mathematics and Systems Science, Chinese Academy of Sciences · Microsoft Research Asia · Renmin University of China +1
Partial differential equations (PDEs) are fundamental for modeling complex physical systems, yet classical numerical solvers face prohibitive computational costs in high-dimensional and multi-scale regimes. While Transformer-based neural operators have emerged as powerful data-driven alternatives, they conventionally treat all discretized spatial points as uniform, independent tokens. This monolithic approach ignores the intrinsic scale separation of physical fields, applying computationally prohibitive global attention that redundantly mixes smooth large-scale dynamics with high-frequency fluctuations. Rethinking Transformers through the lens of complex dynamics, we propose DynFormer, a novel dynamics-informed neural operator. Rather than applying a uniform attention mechanism across all scales, DynFormer explicitly assigns specialized network modules to distinct physical scales. It leverages a Spectral Embedding to isolate low-frequency modes, enabling a Kronecker-structured attention mechanism to efficiently capture large-scale global interactions with reduced complexity. Concurrently, we introduce a Local-Global-Mixing transformation. This module utilizes nonlinear multiplicative frequency mixing to implicitly reconstruct the small-scale, fast-varying turbulent cascades that are slaved to the macroscopic state, without incurring the cost of global attention. Integrating these modules into a hybrid evolutionary architecture ensures robust long-term temporal stability. Extensive memory-aligned evaluations across four PDE benchmarks demonstrate that DynFormer achieves up to a 95% reduction in relative error compared to state-of-the-art baselines, while significantly reducing GPU memory consumption. Our results establish that embedding first-principles physical dynamics into Transformer architectures yields a highly scalable, theoretically grounded blueprint for PDE surrogate modeling.
Pengyu Lai, Yixiao Chen, Dewu Yang +3
School of Aeronautics and Astronautics, Shanghai Jiao Tong University, Shanghai, 200240, China
Neural operators have emerged as powerful data-driven solvers for PDEs, offering substantial acceleration over classical numerical methods. However, existing transformer-based operators still face critical challenges when modeling PDEs on complex geometries: directly processing over massive mesh points is computationally expensive, while operating in raw discretization coordinates may obscure the intrinsic geometry where physical interactions are more naturally expressed. To address these limitations, we introduce the Charted Axial Transformer Operator (CATO), a geometry-adaptive and derivative-aware neural operator for PDEs on general geometries. Instead of applying attention directly in the physical coordinate system, CATO learns a continuous latent chart that maps mesh coordinates into a learned chart space, where chart-conditioned axial attention efficiently captures long-range dependencies with reduced computational cost. In addition, CATO introduces a derivative-aware physics loss for steady-state PDEs that jointly supervises solution values, mesh-consistent gradients, and an auxiliary flux-like field, improving physical fidelity and reducing oversmoothing. We further provide a theoretical approximation result showing that, under a favorable chart, charted axial attention can represent low-rank axial solution operators with controlled error, and that small chart perturbations induce bounded approximation degradation. CATO achieves the best performance across all evaluated datasets, yielding an average improvement of approximately 26.76% over the strongest competing baselines while reducing the number of parameters by 81.98%. These results highlight the effectiveness of learning geometry-adaptive charts and derivative-aware physical supervision for accurate and efficient PDE operator learning.