Efficient Generative Transformer Operators For Million-Point PDEs

Armand Kassaï Koupaï; Lise Le Boudec; Patrick Gallinari

cs.LGDec 4, 2025

Efficient Generative Transformer Operators For Million-Point PDEs

Authors: Armand Kassaï Koupaï, Lise Le Boudec, Patrick Gallinari

Organizations: Sorbonne Université, CNRS, ISIR, 75005 Paris, France · Criteo AI Lab, Paris, France

Abstract

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 $\times$ 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 $L^{2}$ 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.

Deterministic models
		Rayleigh–Bénard		Gray–Scott		Active Matter		ASM
		( $\times 64$ )		( $\times 32$ )		( $\times 176$ )		( $\times 48$ )
Model	Latency (s)	Back.	For.	Back.	For.	Back.	For.	Back.	For.
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

L^{2}

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 (

\times 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

L^{2}

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.

Full resolution ( $\mathcal{X}_{te}=100\%$ )
		Vorticity		Shallow-Water		Eagle		Cylinder Flow
Irr.	Encoder	Rec.	For.	Rec.	For.	Rec.	For.	Rec.	For.
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
Attention	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.

\mathcal{X}_{te}

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

L^{2}

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\times 1024

Vorticity grid (200% of the training grid, more than 20M points). Relative

L^{2}

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

Graph
Model	Reference	Arbitrary domains	Compression	High-res. encoding	Hierarchical encoding	Generative modeling	Multi-task prediction
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\times\bm{1048576}\times 1$
Shallow-Water ( Yin et al., 2022 )	2.6M	Spherical	Single instance	$40\times 32768\times 2$
Gray–Scott ( Ohana et al., 2024 )	1.3M	–	Varying	$40\times 16384\times 2$
Rayleigh–Bénard ( Ohana et al., 2024 )	3.9M	–	Varying	$15\times 65536\times\bm{4}$
ASM ( Ohana et al., 2024 )	2.9M	–	Varying	$15\times 65536\times\bm{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 8: Sample of the Cylinder Flow dataset.

Panel of Figure 9 — 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.

Panel of Figure 13 — Figure 13: Trajectory sample of MHD dataset sliced along each axis.

Panel of Figure 14 — 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 $(\beta_{1},\beta_{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 .

Interpolation module
Parameter	Vorticity	Shallow-Water	Eagle	Cylinder Flow
Token dimension	32	32	32	32
Regularized grid	$192\times 192$	$64\times 128\times 1$	$64\times 32$	$64\times 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\times(512\times 512)\times 1$	$10\times(128\times 256)\times 2$	$200\times 3000\times 3$	$60\times 2000\times 3$
GINO	$10\times 576\times 32$	$40\times 512\times 32$	$100\times 512\times 32$	$60\times 64\times 32$
CORAL	$10\times 18432$	$40\times 16384$	$200\times 4096$	$50\times 2048$
AROMA	$10\times 64\times 32$	$40\times 128\times 32$	$200\times 128\times 32$	$60\times 64\times 32$
CALM-PDE	$10\times 576\times 32$	$40\times 512\times 32$	$200\times 128\times 32$	$60\times 64\times 32$
ECHO	$5\times 24\times 24\times 32$	$10\times 16\times 8\times 32$	$50\times 16\times 8\times 32$	$15\times 16\times 4\times 32$

Appendix

Table 14: Physical and latent space shapes in the experiments of table 3 , in the format temporal size

\times

spatial size

\times

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
GINO	TS	4	16	8	16
CORAL	Rec.	1	2	16	32
CORAL	TS	1	4	16	16
AROMA	Rec.	1	1	4	32
AROMA	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 $(\beta_{1},\beta_{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 $(\beta_{1},\beta_{2})$	(0.9, 0.95)

Appendix

Table 17: Training hyperparameters used across datasets for the generative stage.

Setting	Rel. $L^{2}$
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

L^{2}

error (lower is better).

Panel of Figure 18 — Figure 18: Spectral analysis on a trajectory generation with ECHO on the Vorticity dataset.

Encoder–decoder	#Params	Training throughput (samples/s) $\uparrow$	Latency (s) $\downarrow$	FPD $\downarrow$
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 $\downarrow$	RMSCE $\downarrow$	SSR $\rightarrow 1$	Rel. $L^{2}$ $\downarrow$
Gray–Scott	Forward	ENMA	8.68e-3	3.16e-1	0.43	3.96e-2
	Forward	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
	Forward	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.

L^{2}

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\times 512$ ). Ground truth (GT), ECHO predictions, and per-frame error maps when generating the full trajectory from only the first state and PDE viscosity.

Deterministic models
	Rayleigh–Bénard			Gray–Scott
Model	Back.	For.	Interp.	Back.	For.	Interp.
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

L^{2}

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.

Intermediate sparsity ( $\mathcal{X}_{te}=50\%$ )
			Vorticity		Shallow-Water		Eagle		Cylinder Flow
Irr.	Encoder	Hier.	Rec.	For.	Rec.	For.	Rec.	For.	Rec.	For.
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
Attention	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

\mathcal{X}_{te}

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\%

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

L^{2}

errors (lower is better). Best results in bold , second best underlined .

Hierarchical level	0	1	2	3
Rel. $L^{2}$	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

L^{2}

error). A higher level means more progressive spatial downsampling, with the same final latent size.

	w/o refinement	w/ refinement
Rel. $L^{2}$	2.24e-1	6.88e-2

Appendix

Table 24: Effect of the refinement stage on full-grid Vorticity reconstruction (relative

L^{2}

error, lower is better).

Ablation	Configuration	Rel. $L^{2}$
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

L^{2}

error (lower is better).

Figure 20: Generation using ECHO’s framework on the Vorticity dataset. We generate $1024\times 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.

PDE Surrogate Modeling

PDE surrogate modeling builds fast approximations of partial differential equation solutions for simulation, prediction, and design.

See all 325 papers

UniFluids: Unified Neural Operator Learning with Conditional Flow-matching
Mar 19, 2026Haosen Li, Qi Meng, Jiahao Li +4PDE Operator Learning PDE Surrogate Modeling
Academy of Mathematics and Systems Science, Chinese Academy of Sciences · Microsoft Research Asia · Renmin University of China +1
From Complex Dynamics to DynFormer: Rethinking Transformers for PDEs
Mar 3, 2026Pengyu Lai, Yixiao Chen, Dewu Yang +3PDE Surrogate Modeling Neural Operators
School of Aeronautics and Astronautics, Shanghai Jiao Tong University, Shanghai, 200240, China
CATO: Charted Attention for Neural PDE Operators
May 9, 2026Chun-Wun Cheng, Sifan Wang, Carola-Bibiane Schönlieb +1PDE Operator Learning Attention Mechanisms
DAMTP University of Cambridge · Institute for Foundations of Data Science Yale University · Yau Mathematical Sciences Center Tsinghua University

Mar 19, 2026J/K move · Enter open · S save

UniFluids: Unified Neural Operator Learning with Conditional Flow-matching

Haosen Li, Qi Meng, Jiahao Li, Rui Zhang+3

Academy of Mathematics and Systems Science, Chinese Academy of Sciences · Microsoft Research Asia · Renmin University of China+1

PDE Operator Learning PDE Surrogate Modeling Conditional Flow Matching

Open paper AskPDF arXiv

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.

Abstract

PDE Surrogate Modeling

Explore similar work

UniFluids: Unified Neural Operator Learning with Conditional Flow-matching

From Complex Dynamics to DynFormer: Rethinking Transformers for PDEs

CATO: Charted Attention for Neural PDE Operators