ENMA: Tokenwise Autoregression for Continuous Neural PDE Operators
Authors: Armand Kassaï Koupaï
Lise Le Boudec
Louis Serrano
Patrick Gallinari
Venue: NeurIPS 2025 spotlight ()
Keywords: parametric PDE
generative models
autoregressive transformers
operator learning
Primary area: machine_learning_for_sciences
TL;DR: We introduce a generative masked autoregressive neural operator that models parametric PDE dynamics in latent space, enabling efficient trajectory generation, in-context learning, and few-shot generalization from irregular data.

Abstract: Solving time-dependent parametric partial differential equations (PDEs) remains a fundamental challenge for neural solvers, particularly when generalizing across a wide range of physical parameters and dynamics.
When data is uncertain or incomplete—as is often the case—a natural approach is to turn to generative models.
We introduce ENMA, a generative neural operator designed to model spatio-temporal dynamics arising from physical phenomena. ENMA predicts future dynamics in a compressed latent space using a generative masked autoregressive transformer trained with flow matching loss, enabling tokenwise generation. Irregularly sampled spatial observations are encoded into uniform latent representations via attention mechanisms and further compressed through a spatio-temporal convolutional encoder. This allows ENMA to perform in-context learning at inference time by conditioning on either past states of the target trajectory or auxiliary context trajectories with similar dynamics. The result is a robust and adaptable framework that generalizes to new PDE regimes and supports one-shot surrogate modeling of time-dependent parametric PDEs.

ENMA: Tokenwise Autoregression for Generative Neural PDE Operators
Armand Kassaï Koupai! Lise Le Boudec'* Louis Serrano' Patrick Gallinari
Sorbonne Université, CNRS, ISIR, 75005 Paris, France 2 Criteo AI Lab, Paris, France
Abstract
Solving time-dependent parametric partial differential equations (PDEs) remains a fundamental challenge for neural solvers, particularly when generalizing across wide range of physical parameters and dynamics. When data is uncertain incomplete- as is often the case-4 natural approach is to turn to generative models. We introduce ENMA, generative neural operator designed to model spatio-temporal dynamics arising from physical phenomena. ENMA predicts future dynamics in compressed latent space using generative masked autoregressive transformer trained with flow matching loss, enabling tokenwise generation. Irreg- ularly sampled spatial observations via attention mechanisms and further compressed through spatio-temporal convo- lutional encoder. This allows ENMA perform in-context learning inference time by conditioning on either past states of the target trajectory or auxiliary context trajectories with similar dynamics. The result is robust and adaptable framework that generalizes to new PDE regimes and supports one-shot surrogate modeling of time-dependent parametric PDEs. Project page: https://enma- ande.github.io/
encoded into uniform latent representations
Introduction
Neural surrogates for spatio-temporal dynamics and PDE solving have emerged efficient alterna- tives to traditional numerical solvers, driving rapid advances in the field. Early work (de Bezenac al., 2018; Long et al., 2018; Raissi et al., 2019; Wiewel al., 2019) laid the groundwork, followed by wave of models leveraging diverse architectural designs ( (Yin et al., 2022b; Brandstetter et al., 2022; Wu al., 2024a; Ayed al., 2022). Neural operators (NOs) (Chen & Chen, 1995; Li et al.. 2020a; Lu et al., 2021) expanded the paradigm by learning mappings between
function spaces. Recent models have widely adopted this framework, improving scalability and flexibility (Gupta et al., 2021; Hao et al., 2023; Alkin et al., 2024a; Serrano et al., 2024b; Wu et al., 2024a). While earlier approaches focused on fixed PDE instances, current research tackles the more general task of learning parametric PDEs (Kirchmeyer et al., 2022; Nzoyem al., 2025; Koupai et al., 2024), and is now moving toward foundation models capable of handling multi-physics regimes (Subramanian et al., 2023; McCabe et al., 2023; Liu et al., 2025; Cao et al., 2024; Morel & Oyallon, 2025).
Most neural PDE solvers to date have focused on learning deterministic mappings, limiting their ability to capture complex or uncertain physical behaviors. This has spurred interest in stochastic modeling through generative probabilistic methods. A primary motivation arises from chaotic systems like weather forecasting (Price et al., 2025; Couairon et al., 2024) and turbulent flows (Kohl et al., 2024). Another challenge is error accumulation in autoregressive models which hampers long-term predictions and could be mitigated through probabilistic forecasters (Lippe et al., 2023;
'Equal contribution. Correspondence: armand kassaa lat Jisi.. liselebbuude((atjsirrurgg
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Generation process
2
Plosical pook
fames
Flav Matcking MLP
Spass aakingmmtti
Spatial Trassfernser
hames
Continusus WAE detolienzer
Continuous VAE. tokentzer
Spatio-Temporal Causal Transformer
Arngular Appurgra
Zuo
2-3
1Raa 201 1]
Oquargii
Figure 1: ENMA: Continuous CeaatioTemmoral Autoregressive Generation. Given an initial sequence of latent states z a causal transformer predicts the next latent frame Zyya using block- wise attention for temporal autoregression (Sec.3.1.1). This state is concatenated with masked tokens and passed to spatial transformer, which performs masked spatial autoregression by progressively decoding tokens across multiple steps. A lightweight MLP trained with flow matching produces per-token predictions conditioned on spatial transformer predictions (Sec.3.1.2). The completed frame - is appended to the history for rollout. The full latent sequence is then decoded into the physical domain using a continuous VAE decoder (Sec. 3.2).
Kohl et al., 2024). More generally, data-driven surrogates must handle both alcatoric and epistemic uncertainty (Bülte et al., 2025; Wu et al., 2024b), often under conditions of partial observability, sensor noise, coarse space-time resolution. Additionally, they face distribution shifts at test time, where PDE parameters may deviate from training data. Epistemic uncertainty may also stem from model mismatch or limited training data, reducing generalization out-o--distribution regimes (Mouli et al., 2024).
Generative neural surrogates for time-dependent PDEs increasingly draw on advances from computer vision and, more recently, language modeling. These approaches typically fall into two categories. The first includes diffusion models, which operate in continuous spaces- pixcl-level or latent- and model the joint distribution of future states by reversing a noising process. Based on U-Nets or transformer encoders (Ho et al., 2020; Peebles & Xie, 2022), they have been extended to PDE forecasting via iterative denoising (Zhou et al., 2025; Lippe et al., 2023; Kohl et al., 2024; Li et al., 2025). The second family follows autoregressive (AR) paradigm, predicting per-token conditional distributions (Touvron et al., 2023). These models discretize physical fields using vector quantization (Oord al., 2017; Mentzer et al., 2024), and generate sequences either sequentially (Yu al., 2024) or via masked decoding (He et al., 2022; Yu et al., 2023), with recent adaptations for PDE surrogates and in-context learning (Serrano al., 2025). Each approach comes with trade-offs. Diffusion models are robust to distribution drift and support uncertainty quantification via noise injection, but require expensive training and inference (Leng al., 2025). AR models offer efficient, eausality-alinned generation via KV caching and support in-context learning (Brown et al., 2020), but suffer from reduced expressiveness due discretization and often yield uncalibrated uncertainty estimates (Jiang et al., 2020). They are also more sensitive to error accumulation from teacher forcing. Recently, continuous-tokke token AR models have emerged as promising alternative to discrete decoding in vision (Li et al., 2025), motivating their application to scientific modeling, We address fundamental limitations existing generative surrogate models for PDEs-partes partss-aatilularlyy their reliance on discrete tokenization or full-frame diffusion, which hampers scalability, uncertainty modeling, and physical consistency. We introduce ENMA (presented in Figure 1), a continuous autoregressive neural operator for modeling time- deeendepnndett parametric PDEs, where parameters such initial conditions, coefficients, and forcing terms may vary across instances. ENMA operates entirely in continuous latent space and advances both the encoder-decoder pipeline and the generative modeling component crucial to neural PDE solvers. The encoder employs attention mechanisms to process irregular aratio-tempora inputs (i.e., unordered point sets) and maps them onto structured, grid-aligned latent space. A causal spatio- compresses these observations into compact latent tokens spanning multiple states.
convolutional encoder (Yu et al., 2024) then
Generation proceeds in two stages. causal transformer first predicts future latent states autore- gressively. Then, masked spatial transformer decodes each state at the token level using Flow

---

Matching (Lipman et al., 2023, 2024) to model per-token conditional distributions in continuous space providing a more efficient alternative to full-frame diffusion (Lippe al., 2023; Kohl et al., 2024; Li et al., 2025). Finally, the decoder reconstructs the full physical trajectory from the generated latents. ENMA improves state-of-the-art generative surrogates by combining flexible per-token generation with expressive and compact latent modeling. It extends language-inspired AR models (He et al., 2022; Serrano et al., 2025) to continuous tokens, avoiding the limitations of quantization. Compared to diffusion models, it offers lower computational cost by sampling tokens via lightweight MLP instead of full-frame denoising. At inference, ENMA supports uncertainty quantification and in-context adaptation- iniraaiion-conntiionnnng conditioning either past target states or auxiliary trajectories governed by similar dynamics. Our contributions are as follows:
We introduce ENMA, the first neural operator to perform autoregressive generation over continuous latent tokens for physical systems, enabling accurate and scalable modeling of parametric PDEs while avoiding the limitations of discrete quantization.
Using a masked spatial transformer trained with Flow Matching objective to model per-token conditional distributions, ENMA offers principled and efficient alternative to full-frame diffusion models for generation.
ENMA supports probabilistic forecasting via tokenwise sampling and adapts to novel PDE regimes at inference time through temporal or trajectory-basee conditioning-withiee retraining.
To handle irregularly sampled inputs and support multi-state tokenization, ENMA leverages attention-based encoding combined with causal temporal convolutions.
2 Problem Setting
We consider denendetnndent parametric partial differential equations defined over a spatial domain
Ω I2d and time interval [0, T). Each instance is characterized by initial condition u."
(22 Rd))) ) and set of parameters
(b, J c), which include boundary conditions b E L x [0, 27, IR4) ), forcing term f E L (S2 [0, 7', pdd)), and PDE coefficients c. The system of equations
is:
N [w; e, f] (x, 1) 0,
for (xr, E f x (0,7),
Iu: b) (I, t) = 0,
for (xc t) E a2 x [0,7],
(2)
u(x, 0) = u (x),
for x E Ω
where N denotes (potentially nonlinear) differential operator, and B encodes the boundary condi- tions. From an operator learning point of view, the objective is to approximate with a neural network G the temporal evolution operator Gy: u. ++At (u())
For training, we assume access to a dataset of N solution trajectories, [ where each trajectory characterized by an initial condition uy ~ Vuo and environment-specific parameters 7/ Vy, and is observed over a spatial grid X, and on a temporal horizon [0. Since the parameters % are unobserved, we supply the neural network with additional information beyond the current state ut to help resolve this ambiguity. Specifically, we consider two predictive settings:
(i) Temporal conditioning. The model & observes an initial trajectory segment of L states 0:L-1 and must autoregressively forecast future states up until T:
- = G'
u I++1 = (..
(ii) Generalization from context trajectory. In this setting closer the one of classical numerical solvers, the model observes only the initial state u² of the target trajectory, along with a separate context trajectory OLL governed by the same 7, but from different initial condition. The model uses this context to forecast the evolution from uo:
u = "(u" % 2 (.Tee)):
u² = "( 0:1, (ostes)):
In these two settings, the surrogate model must emulate the underlying dynamics from data to unroll the target trajectory. These settings highlight the challenges of forecasting time-dependent systems under partial observability and unobserved parameters.
3

---

3 ENMA
We introduce ENMA,, neural operator tailored for continuous tokenwise autoregressive generation of spatio-temporal dynamics. ENMA follows an mncode-generate-der solution operator G, and can be decomposed in the three corresponding steps G RJ G Dv Pa Ew... The encoder E maps irregularly sampled spatio-temporal inputs u u0=--1 R |x| x observed at X spatial locations over L time steps with c physical channels- into a structured latent representation Z0: IR Lxx The generative model Pa then autoregressively predicts future latent states in tokenwise fashion, and the decoder D. maps the generated latents back into the physical domain. The full process for a one-step prediction can be summarized as:
pipeline to approximate the
U D:L-1 (IL-1ERRR[xLLS encode, 70:L-1 MRMXXx generale, 2L E E Radd decode, ut ERRxixe
(4)
Here, M is the number of spatial latent tokens per state Z and d the latent embedding dimension. The generative model can be repeatedly applied roll out predictions autoregressively over horizon of I' steps. We describe the generative model in Section 3.1, and the encoder-decoder in Section 3 3.2. For simplicity, we describe the model for the temporal conditioning setting, the adaptation for the generalisation from context setting is immediate.
3.1 Tokenwise Autoregressive Generation
The core contribution of ENMA lies in its continuous autoregressive architecture, designed for modeling spatiotemporal dynamics. The generative model proceeds in two stages (5): it first extracts a spatio-temporal representation Z from the encoded trajectory Z using causal transformer, and then predicts the spatial distribution of the next time step via tokenwise decoding mechanism. This distribution is estimated with a fow-matching component, with the help of a spatial transformer and lightweight MLP, inspired by recent autoregressive models such as MAR (Li et al., 2024).
20:L11 E TR MxLxd Causal Transformer
2
E
TR Mxd AR Spatial Generation
24 RRdd (5)
This two-stage design enables ENMA to capture both long-range temporal dependencies through the causal transformer and fine-grained spatial structure with the generative spatial decoder, in scalable, autoregressive manner. This choice is further motivated by the distinction between temporal and spatial dimensions: temporal prediction naturally benefits from the sequential ordering of the trajectory, whereas there is no such ordering for spatial generation. We provide architectural and training details for each stage, respectively, Sections 3.1.1 and 3.1.2, following the temporal conditioning setup introduced in Section 2.
3.1.1 Causal Transformer
The causal transformer designed to extract spatio-temporal representations of trajectory while enabling scalable training. At each time step i, all tokens within the current state Z can attend to one another, as well as all tokens from preceding states Z' for j < i. Formally, given sequence of latent states 2 = (Z" .... z".-1), ), the transformer produces the dynamic context representation for time step t: Z CausalTrransfommer Zass z0: where Zecs learned sequencee (BOS) token. Z can be seen as latent context capturing the dynamics, from the observed time-steps (0, 1). to predict the next step
To support parallel training and efficient inference, this factorization is implemented using a block- wise causal attention mask (see Appendix D). The model's causal structure also enables key-value caching at inference, allowing reuse of past computations and facilitating fast autoregressive rollout,
During training, we apply teacher forcing- conditionnn the model on the full sequence of ground- truth latent states. At inference, the causal Transformer can accept arbitrary sized sequence inputs Z0:L-1 for L E [0. Ziyy is then used context for the spatial transformer described in 1.2.
3.1.2 Masked Autoregressive Generation
To perform tokenwise continuous autoregression, ENMA employs a masked decoding scheme conditioned on the context Ztn- This is implemented with spatial transformer that produces

---

conditioning intermediate representations Z" ( 2M which and temporal context for each token combined with lightweight MLP that models the per- token output distribution p(2 ),, i ,,.., M. We adopt the musked autoregressive (MAR) strategy (Li et al., 2024), which enables cormuutaaionnnnvariant autoregressive generation over the spatial domain. Training and inference pipelines for spatial generation are illustrated in Figure 2.
which capture spatial correlations
Training
Inference
2
Matching
fv
By
2 2M[M]
Bx
2
Discore sconasked DANu
[AG] nelves nolected the
2
spatiol Inllets
Spatial Transformer
aasingrrsm gevrntian
Spatial
Z
Ziss
Trie
[M
forsingg
Zir
[MASK]
Masks
Ground Iruth tokens
Selected takens for generation
Figure 2: Training and inference for the spatial transformer. For training, a random subset of tokens is masked and decoded in single step. At inference, generation proceeds iteratively: starting from fully masked state Z = a subset N₂ of tokens selected each step be generated by the MLP, conditioned on spatial transformer outputs. This process is repeated steps to produce 2
Spatial Transformer. During training, random subset of token indices M C (1,.... M) is selected to be masked in the ground-truth latent frame Zt, with the masking ratio sampled uniformly between 75% and 100%. These tokens are replaced with learned [MASK] embedding to form the partially masked input ZHASK. The spatial transformer processes the concatenated pair (Zorx. Zzisk) and outputs contextual representations Z The tokens ε [MM then used independently to condition a lightweight MLP trained with a Flow Matching objective. For each i E M,, the MLP maps noisy couple-conditionee on -.-- into denoised latent prediction e The loss is computed only over the masked positions M..
into
At inference, tokenwise autoregression proceeds over S autoregressive steps to predict 2 At step s = 0, 2 1.00 is initialized with the mask embedding, i.c. 2 = [MASK] for i E (I...., M). At each step, subset N₂, C {i (1,...,,7 M) 2!! = [MASK] of the masked positions (not previously selected if 0) selected for generation. The spatial transformer processes (Z 20.0) and predicts 7 T.,;; the selected positions be generated 2 are used condition the flow matching MLP, generating the selected tokens $..... These tokens are inserted back into ŻL.S 2. to form 2 Lh,+11
and the process continues until all positions are generated step S, yielding 20 2L.S, Multiple tokens can be generated at each step. The number of decoded tokens per step follows a cosine schedule: [.a.l [ cos? [e ]. Compared to fully parallel decoding, this approach provides fine-gratd tne-grriied control over uncertainty and sample diversity, while enabling efficient generation through vectorized MLP sampling.
Flow Matching for Per-Token Prediction To model per-token conditional distributions in continu- ous latent space, ENMA employs Flow Matching (FM).
Let 2 E Rd denote the ground-truth latent token at position i in the frame zL (zf 2 and let Σ be its contextual embedding produced by the spatial transformer. We drop the superscripts for clarity here. We seek learn the conditional distribution p(z æ) via flow-based transport from a base distribution Po Z) to the target P1 = p(z 2). This is parameterized as an ordinary differential equation (ODE):
dz"
u(z", r),
(6)
dr
5

---

where E [0. 1] is denoising index for the flow matching, and is velocity field implemented by neural network. During training, we sample intermediate points along the probability path using:
x' - e~ <~,(0,
We train a lightweight MIP, conditioned on the context 2, to approximate the velocity field. It takes input (z", rr) and predicts the transport direction. The flow matching training objective is:
Lex
Er.r MMLP("", - +
(8)
which encourages the model to match the velocity that would move z' along the interpolating path toward z. At inference, token generation is performed by solving the ODE:
2 = ODESOLVE (MLP(z", ,,))), with ~/(,,
from r= to r = 1 using the midpoint method. This enables efficient of continuous latent tokens conditioned on their spatial context, without requiring discrete quantization or full-frame denoising. We present the full inference method of ENMA for latent generation in the pscudo-code 1:
Algorithm 1: ENMA Inference: Autoregressive Latent Generation with Cosine Masked Decoding
Input: Encoded latents 20:L-1 from observed inputs u L-L-1
Output: Predicted latents ( 2²}
for f = L to T' do
7 f-
7 0tt-1);;
2' <- [MASK];
for 1 to do
2 SpatialTransformer/ [Zbri 2']):
Tt + CosineSchecule((s.. S);
Sample nt masked indices: N₂ C {i 1 = [MASK]]] );
(0, xd;
2 [WG] ODESOLVE(MLP, Z'[Na),
Z 2%
3.2 Auto-Encoding
Encoding irregular spatio-temporal data into a compact latent representation a central component of ENMA. Inspired by recent advances in vision models (Yu et al., 2023), ENMA combines two key elements. (i) a cross-attention module performs spatial interpolation, mapping the irregular inputs a (X)) defined over an arbitrary spatial grid X' at L time steps, onto regular grid 3, yielding intermediate representations u 0-L-1 (E). (it). a temporally causal CNN processes the intermediate sequence -.11-1 [E), across space and time into tokens 1 Mapping irregular inputs onto regular grid-aligned space, allows ENMA to leverage convolutional inductive biases of the CNN for efficient and coherent compression. Decoding mirrors the encoder structure. The encoder and decoder are optimized jointly with a VAE loss: L = Lrecon + B LKL. To improve robustness to varying input sparsity, randomly subsample the spatial grid at training time: the number of input points varies between 20% and 100% of the full grid X. These two elements are described in more details below.
(i) Interpolation via Cross-Attention For interpolation, we modify the cross-attention module from Serrano al. (2024a) to favor spatial locality, by introducing geometry attention bias. This bias, exploits the geometry of the inputs and induces locality the cross-attes catss-nttention operation, enabling the regular gridded interpolation u 0:L-I1 ( to capture the spatial structure of the physical field.
Formally, we define the atoss-attention from queries located at X' to keys and values located at = as:
[QKT
B) V,
Attention (Q, K, V) = Sootmmx
Big = - -ddit(((;; - Ex), for x E X E E
ViTe

---

where Q = q()),, - V o(E) are query, key and value tokens defined on the respective spaces X' and E. dg denotes the key and query dimension, and m is a scaling factor (either fixed or learned, as in AL.BB (Press et al., 2022)). Intuitively, larger distances between query and key positions induce stronger negative biases, reducing attention weights and promoting locality in the interpolation.
byyer
(ài) Processing via Causal Convolution Following the interpolation mod- ules, ENMA compresses latent trajectories using a causal 3D convolutional encoder, enabling efficient representation across space and time. Unlike prior work (e.g., Wang et al. (2024)), our use of causal convolutions sup- ports variable-length inputs, allowing the model to process both temporally conditioned settings and initial value problems. For 2D spatial domains se- quences, we use 3D convolutions with kernel size (ks, kn, kw) and causal padding in the temporal dimension (Yu et al., 2023). This techniques only applies a k2 - padding in the past of the temporal dimension. This design ensures that outputs at time t depend only imputs up t, enabling inference from a single initial frame (see Figure 3). Spatial and temporal compression are applied independently, with the overall compression rate controlled by the number of stacked blocks each reducing the corresponding dimension by factor of 2. The overall architecture design is detailed in appendix D.
caudal podding
Figure 3: Causal in time convolutional layer.
Decoding Decoding mirrors the encoder structure. Transposed convolutions reverse the spatio-temporal compression, and final cross-attention layer interpolates latent outputs to arbitrary physical grids. Here, the desired output coordinates serve as queries, while the latent tokens provide keys and values, enabling nigh-resolution reconstruction independent of the training grid.
Experiments
We conduct extensive experiments to evaluate ENMA. Section 4.2 assesses the encoder-decodee in terms of reconstruction error, time-stepping accuracy, and compression rate, comparing against standard neural operator baselines. Section 4.3 evaluates ENMA's generative forecasting ability for both temporal conditioning and Initial Value Problem with context trajectory. Dataset, training, and implementation details are provided respectively in Appendices C, E, and E. We also provide additional experiments and ablations in appendix
4.1 Datasets
We evaluate ENMA on five dynamical systems: two 1D and three 2D PDEs. Data is generated in batches where all trajectories share PDE parameters 7, but differ in initial conditions. The PDE description is provided in Appendix C. For each system, we generate 12,000 training and 1,200 test trajectories, using batch size of 10. In ID, use the Combined equation (Brandstetter al., 2022),
where coefficients (a, 3, 7) vary for the terms dar + where both advection speed and initial condition vary. In 2D, consider: the Vorticity equation with varying viscosity y; the Wave equation with varying wave speed c and damping coefficient k; and the Gray-Scott system (Pearson, 1993), where the reaction parameters F and l vary across batches.
au2 88 2h and 8du. and the Advection equation, 2z
4.2 Encoder/ decoder quality evaluation
Setting We evaluate the encoder-decoder quality across baselines that use latent-space repre- sentations key component of neural surrogates. Since prediction accuracy depends on the decoder pair, we first assess its limitations and show ENMA's improvements. We conduct two tests: (i) reconstruction, evaluating auto-encoding fidelity, and (ii) time-stepping. where a small fixed FNO, trained per encoder, performs rollouts in latent space. All models are trained with the same objective as ENMA, using input subsets varying between 20% and 100% of the grid X. At test time, inputs use irregular grids with 20%, 50%, or full (100%) sampling; reconstruction is always evaluated on the full grid. See appendix E for more implementation details and experimental protocol.
7

---

Baselines We compare ENMA's encoder- against representative neural operator architec- tures: two transformer-base models OFormer (Li et al., 2023a) and AROMA (Serrano al.. 2024a); the INR-based CORAL (Serrano et al., 2023); the neural operator GINO (Li et al., 2023b). These baselines differ in compression: OFormer performs no compression (point wise), AROMA, CORAL, and GINO compress spatially, while ENMA compress both spatially and temporally. All baseline are compared with similar token dimension (e.g. d = 4 for datasets and = for 2d datasets), but the compression rate can vary some baselines do not compress in time and/or in space.
Table 1: Reconstruction error- Test results and compression rates. Metries in Relative MSE The compression rate reflects how much the latent representation is reduced compared to the input data. A compression rate of x2 indicates that the latent space contains half as many elements as the input.
Noc
Dataset
+
Advection
Vorticity
Model +
Reconstruction
Time steppigg
Compression rate
Reconstruction
Time- Time-stepping
Compressiio rate
OFormer
1.70-1
1.11440
x0.25
9.99e-1
1.00e40
x0.125
GINO
3.15e-1
8.55c-1
x2
5.63e-1
9.83e-1
x8
T 100%
AROMA
5.41e-3
2.23e-1
X2
1.45e-1
1.13e40
X8
CORAL
1.34e.2
9.64e 1
X2
4.63e-1
9.18e-1
X8
ENMA
1.83e3
1.64e-1
X4
9.20e-2
2.62e-1
X15
OFomme
1.1e+0
9
1.00+00
GINO
3.21e-1
8.64-1
5.69e
9.91e-1 1
50%
AROMA
2.34e-2
2.29e-1
16
1.14++0
CORAL
7.57e-2
9.74e
4.95
9.22el1
ENMA
4.60e-3
1.72e-1
9.90e-2
2.68e-1
OFormer
2.50e-1
1.13e+0
9.99
1,00++0
3.54e-1
9.11e-1
5.90e1
1.04e40
AROMA
1.67-1
3.21e-1
2.29c1
1.14+40
20%
CORAL
4.77e-1
1.06440
6.8%-1
9.37e-1
ENMA
3.05e-2
3.13e-1
1.37e-1
3.11e-1
Results Table 1 presents the reconstruction and time-stepping eπrors for ENMA and baseline methods. These results highlight the robustness of ENMA across different grid sizes and datasets, in both 1D and 2D settings. While most baselines struggle as the grid becomes sparser, ENMA'S encoding decoding strategy consistently provides informative representations of the trajectories, even under limited observations. When the full grid is available (i.e., π = 100%). ENMA outperforms the baselines, reducing the reconstruction error by up to factor of 3x. As the observation rate drops to 50% and 20%, the performance gap widens even further. Notably, at 20%, ENMA maintains strong performance, while other methods degrade significantly. This demonstrates ENMA's ability to generalize and infer latent dynamics under sparse supervision. For the time-stepping task, the higher- quality tokens produced by ENMA significantly improve forecasting accuracy. We also emphasize that ENMA compresses the temporal dimension more effectively than baselines: even with fewer tokens ( x4 compression in 1d and x x15 in 2d), it remains competitive. This improvement is attribbted to ENMA's ability to leverage causal structure during encoding, which leads to more informative and temporally coherent representations. Additional analysis about the ENMA's encoder/decoder architecture provided in appendix F.
4.3 Dynamics forecasting
Setting We evaluate our model the standard dynamics forecasting task, which predicts the future evolution of spatio-temporal systems. We consider two settings: (i) temporal conditioning, where the model observes the first L time steps and predicts up to horizon T; and (ii) initial value problem with context, where only the target's initial state n is given, along with an auxiliary trajectory governed by the same parameters y but from different initialization. The model must infer dynamics from this context and forecast from u". We evaluate both settings under in-distribution (In-D) and -ut-of-distributtio (Out-D) regimes, where the latter involves PDE parameters not seen during training. Each Out-D evaluation uses 120 held-out trajectories. See Appendix C for details.
Baselines We evaluate ENMA against a diverse set of baselines, both deterministic and genera- tive. For temporal conditioning, the deterministic models include the taasformer-basee solvers:
8

---

BCAT (Liu et al., 2025) and AVIT (McCabe et al., 2023) both designed for multi paalti-physiss forecasting. and the Fourier Neural Operator (FNO), a classical reference. For the generative models we consider a continuous autoregressive diffusion transformer AR-DiT (Kohl et al., 2024), and Zebra (Serrano language-model-styil
et al., 2025), language- model-style decoder operating in a quantized latent space. For the initial value problem with context, we compare with the deterministic In-Context ViT that concatenates the context trajectory to the target initial state, and with a [CLS] VIT variant that uses a learned token to summarize the PDE parameters 7,. For the generative baselines, we include Zebra (Serrano et al., 2025).
Table 2: Comparison of model performance for temporal conditioning and initial value problem tasks across 5 dynamical systems. Metrics in Relative MSE. Lower is better.
Dataset
Advection
Combined
Gray-Scott
Wave
Vorticity
Setting
Model
In-D
Our-D
In-D
Our-D
In-D
Out-D
In-D
Our-D
In-D
Out-D
2.47e11
7.95e-1
1.33e11
2.66++1
5.04e-2
1.92e11
6.91e:1
2.64e40
6.07e:2
2.15c-1
BCAT
5.55e-1
9.23e-1
2.68e-1
9.28e-1
3.74e-2
1.57e-1
2.19..1
5.38e-1
5.39p-2
3.00e11
Temporal Conditioning
AVII'
166e-1
5.02e-1
5.67e-2
3.05e11
4.26e-2
168e-1
1.57e1
5.8%e-1
1.76e-1
3.77e-1
2.36e-1
8,56e-1
2.95e-1
1.80+0
3.69e-1
4.99c-1
1.12+00
7.52e+0
1.98c-1
4.80e-1
Zebra
2.04c-1
1.39c+0
L82s-2
2.20c+0
4.21o-2
1.82-1
1.4e-1
3.15e-1
4.43e-2
2.23e-1
ENMA
3,95e-2
5,30e-1
7.86e-3
1.02-1
3.40e-2
1.4c-1
1.45-1
4.89c-1
7.58e-2
3.45e-1
In-Context VIT
1.15e+0
1.20c40
5.79e-1
1.36e40
6.90e-2
1.946-1
1.72o-1
6.246-1
1.53e-1
3.92o-1
[CLS]
1.15e+0
1.36++0
9.60e.2
1.16+40
4.80e
2.19
5.56e
1,02+40
4.Mie-2
2.59e.1
Initial Value Problem
Zehra
3.16e_1
1.47++0
4.78e.2
9.63e-1
4.40e-2
1.22e1
169e.1
3,52e-1
5.90e:2
2.29c-1
ENMA
2.02e-1
8.07e-1
1.56e-1
3.30e-1
4.80e-2
1.34e1
1.5e-1
5.02e-1
8.58e-2
3.20e-1
Table 2 summarizes performance across five PDE benchmarks under both temporal condi- tioning and IVP settings. ENMA achieves state-of-the-aar results on most tasks, outperforming both deterministic solvers (FNO, BCAT, AVIT) and generative baselines (AR-DiT and Zebra). Unlike most methods that operate in physical space, ENMA performs autoregressive generation entirely in continuous latent space (time and space compression), which reduces the computational complexity, but represents a much more challenging setting. Only Zebra shares this property (compression in space only) but relies on quantized tokens, which favor low-frequency reconstruction but can limit expressiveness, ENMA shows strong performance on Advection, Combined, and Gray-Scott, and remains competitive on Wave, despite the added challenge of temporal compression. Performance on Vorticity is slightly lower than the best competitor, likely due to the combined effect of temporal and spatial compression, which makes accurate reconstruction more challenging. We analyze this aspect in appendix F.1.2, where we obsreve that reducing the compression ratio naturally lead to better prediction performance. ENMA's continuous modeling allows finer control over generation and uncertainty estimation. This aspect is explored in appendix F.I.I.
Results
4.4 Evaluation on High-Dimensions Physics Systems
We evaluate ENMA on standard public benchmarks Saayleih-Bénnade and Active Matter, (Ohana et al., 2024)). These datasets feature highly nonlinear spatio-temporal dynamics, multiple interacting physical fields, and dense spatial grids, making them representative of complex, high-dimensional physical systems. We compare against the same competitive deterministic baselines used in section 4.3.
Table 3: Temporal Conditioning setting on complex physical systems (Relative MSE 4). Compression
ratios are reported per dataset.
Rayleigh-Bénard
Active Matter
Model
Temporal Conditioning
Comp.
Time-stepping
Comp. x1
BCAT
1.06e-1
x1
4.56c-1
AVIT
1.01e-1
x1
4.62e-1
x1
ENMA (ours)
9.87c-2
x64
3.33e-1
x176
Results
ENMA attains the lowest error on both systems while operating at high compression (e.g.. x64 on Rayleigh-Bénan and x 176 on Active Matter), demonstrating that accurate time-stepping

---

is achievable even under aggressive latent compaction. Deterministic baselines (BCAT, AVIT) are faster but lack compression and yield higher errors.
4.5 Generative Capabilittess of ENMA
ENMA is generative neural operator capable of producing stochastic and physically consistent trajectories. We highlight two core experiments that demonstrate its generative ability; extended analyses and qualitative visualizations provided in appendix F.1.1.
Uncertainty quantification. ENMA performs un- certainty estimation by sampling multiple trajecto- ries from its continuous latent space through flow matching. Unlike discrete autoregressive baselines such as Zebra, which rely on categorical sampling. ENMA's continuous formulation yields sharper and better-calibrated probabilistic forecasts. As shown in table 4, ENMA achieves the lowest calibration (RM- SCE) and probabilistic (CRPS) errors, confirming its ability produce both reliable and diverse uncertainty estimates.
Table 4: Uncertainty metrics (L is better) on the Combined dataset.
Model AR-DiT
RMSCE /
CRPS
2.68x 10
1.27x10
Zebra
2.19x10-1
9.00
ENMA (ours)
8.68x10-2 22
1.70x10
Data generation. We further assess ENMA'S ability to generate full trajectories without condi- tioning the initial state or PDE parameters. Given only context trajectory, ENMA infers the latent physics and synthesizes coherent apaati-ttemorra fields. table 5, ENMA achieves the lowest Physies Fréchet Distance (FPD) and highest Precision, indicating superior fidelity and physical consistency. Zebra attains slightly higher Recall, reflecting greater diversity but lower sample quality.
Table 5: Generative metrics the Combined dataset. Lower FPD and higher Precision/Recall indicate better quality and diversity.
Model
FPD 4
Precision t
Recall 1
Zebra
1.03 10-1
0.77 0.79
0.86
ENMA (ours)
9.50x10-3
0.78
5 Conclusion
Discussion We presented ENMA, a generative model leveraging continuous latent-space autore- gressive neural operator for modeling dinendependent parametric PDEs. ENMA performs generation directly in compact latent space, objective, enabling efficient forecasting and improves over alternative generative models based either on full-frame diffusion or on discrete quantization. Across 5 dynamical systems, experiments demon- strate that ENMA is strong neural surrogate that competes with neural PDE solvers operating in the physical space, in both temporal conditioning and initial value problem settings. Reconstruction and time-stepping evaluations confirm the effectiveness of its encoder-decoder design, showing improved latent modeling and robustness to unordered point sets compared to existing neural operator baselines.
token level, using masked autoregression and a flow matching
Limitations ENMA's performance is slightly reduced on vorticity, where latent compression can hinder the accurate recovery of fine-scale features. This highlights central trade-off in latent- space surrogates: while compression enables scalability. it may limit expressiveness in certain regimes. Regarding computation efficiency, ENMA is computationally efficient compared to full- frame diffusion next-token models. However, ENMA's model's ss cost increases with respect to the number latent tokens per state. This motivates future work on adaptive generation strategies-for instance, using coarse-to-fine decoding scheme where a token directly synthesizes the frame, and remaining tokens act as refinements (Bachmann et al., 2025).
Broader impact PDE solvers are key to applications in weather, climate, medicine, aerodynamics, and defense. While ENMA not deployed in such settings, offers fast, uncertainty-awar surrogates adaptable to regimes.
10

---

Acknowledgments
We acknowledge the financial support provided by DL4CLIM (ANN-19-CHIAA0001800), DEEPNUM and PEPR Sharp (ANR-23-PEIA-
(ANR-21- 21-CE23-0011--00 PHLUSIM
0008", "ANR"" "FRANCE 2030"). This project was provided with computer and storage resources by GENCI at IDRIS thanks to the grants 1016800R,, 1025-AD01101493600 and 2025- 1 IRI on the supercomputer Jean Zay's A100/H100 partitions.

---

References
Benedikt Alkin, Andreas Fürst, Simon Schmid, Lukas Gruber, Markus Holzleitner, and Johannes Brandstetter. Universal physics transformers. In Neurips, 2024a. URL http://arxiv. org/ aba/ 2402. 12365.
Benedikt Alkin, Andreas Fürst, Simon Schmid, Lukas Gruber, Markus Holzleitner, and Johannes Brandstetter. Universal physics transformers. arXiv preprint arXiv:2402.1224 12365, 2024b.
Ibrahim Ayed, Emmanuel de Bézenac, Arthur Pajot, and Patrick Gallinari. Modelling spatiotemporal dynamics from earth observation data with neural differential equations. Machine Learning, 111:2349-2348 2022. ISSN 15730565. doi: 10. 011777/11994-22--000 10. 1007/ /0000//119994 020 06139-2.
URL https:
Roman Bachmann, Jesse Allardice, David Mizrahi, Enrico Fini, Oğuzhan Fatih Kar, Elmira Amirloo, Alaaeldin El-Nouby, Amir Zamir, and Afshin Dehghan. Flextok: Resampling images into 1d token sequences of flexible length, 2025. URL https: /arraaaxvvvor////.aa.a...../.... herg/abs/2502 13967.
Johannes Brandstetter, Daniel E Worrall, and Max Welling. Message passing neural pde solvers. International Conference on Learning Representations, 2022.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
Christopher Bülte. Philipp Scholl, and Gitta Kutyniok. Probabilistic neural operators for functional uncertainty quantification. TMLR, 2 2025. URL http::/arxiv..rrr/ h2tgaab//500..12000 12902.
Yadi Cao, Yuxuan Liu, Liu Yang, Rose Yu, Hayden Schaeffer, and Stanley Osher. Vicon: Vision in-context operator networks for multi-physics fluid dynamics prediction, 11 2024. http: //arxiv. org/ /bs/ 16063.
Yue Cao, Jiarui Xu, Stephen Lin, Fangyun Wei, and Han Hu. Global context networks, 2020. URL https://arxiv.org www..rgggab........ 13375.
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 11315-11325, 2022.
Tianping Chen and Hong Chen. Universal approximation nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems. Trans, Neur. Netw., (4):911-917, July 1995. ISSN 1045-9227. doi: 10.1 10.1709/7..99225 URL https://doi. org/10. 1109/72.392255.
Wuyang Chen, Jialin Song, Pu Ren, Shashank Subramanian, and Michael W Mahoney. Data-efficient operator learning via unsupervised pretraining and in-context learning. In NeurlPS, 2024. URL
https://github: www.aeleeoom//eeaaaaaaa 1ab- httpd/aaaaaffiiieeeeeeeeaaaeaaae cient_nopt.
Guillaume Couairon, Renu Singh, Anastase Charantonis, Christian Lessig, and Claire Monteleoni. Archesweather & archesweathergen: a deterministic and generative model for efficient ml weather forecasting, arXiv preprint arXiv:2412. 12971, 129,177,,,, dec 2024. URL https://arxiv. org/ 2212. 12971. Published: December 17, 2024.
Emmanuel de Bezenac, Arthur Pajot, and Patrick Gallinari. Deep learning for physical processes: Incorporating prior scientific knowledge. In ICLR, 2018.
Haoge Deng, Ting Pan, Haiwen Diao, Zhengxiong Luo, Yufeng Cui, Huchuan Lu, Shiguang Shan, Yonggang Qi, and Xinlong Wang. Autoregressive video generation without vector quantization. arXiv preprint ar arXiv:2412.1430 2024.
12

---

Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. / 1929, 2020.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. International Conference Learning Representations, 10 2021. URL http: arxiv org/ abs/2010. 11929.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high high-resolution image synthesis. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 12873-12883, 2021.
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet-based operator learning for differential equations. Advances neural information processing systems, 01-24048-24066 2021.
Zhongkai Hao, Zhengyi Wang, Hang Su, Chengyang Ying, Yinpeng Dong, Songming Liu, Ze Cheng, Jian Song, and Jun Zhu. Gnot: A general neural operator transformer for operator learning. In International Conference on Machine Learning, pp. 12556 12569. PMLR, 2023.
Zhongkai Hao, Chang Su, Songming Liu, Julius Berner, Chengyang Ying, Hang Su, Anima Anandku- mar, Jian Song, and Jun Zhu. Dpot: Auto-regressive denoising operator transformer for large-scale pde pre-training. 41th International Conference on Machine Learning (ICML 2024), 2024.
Kaiming Hc, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autocncoders are scalable vision learners. In CVPR, 2022.
Maximilian Herde, Bogdan Raonić, Tobias Rohner, Roger Käppeli, Roberto Molinaro, Emmanuel de Bézenac, and Siddhartha Mishra. Poseidon: Efficient foundation models for pdes, 2024.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge a local nash equílibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffussio probabilistic models. arXiv preprint arxiv:2006.1 11239, 2020.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. Advances in Neural Information Processing Systems, 2022.
Zhengbao Jiang, Frank F. Xu, Jun Araki, and Graham Neubig. How can we know what language models know?, 2020. URL https:///artttt org/ aaaaabb/1991 /911 12543.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers, 2023. URL
https: / //trx//axvvvv//e//////eeeeeaaeva org/ abs/ abss/330 2305. 19466.
Matthieu Kirchmeyer, Yuan Yin, Jérémie Donå, Nicolas Baskiotis, Alain Rakotomamonjy, and Patrick Gallinari. Generalizing to new physical systems via context-informed dynamics model. In International Conference on Machine Learning, pp. 11283-11301. PMLR, 2022.
Georg Kohl, Li-Wei Chen, and Nils Thuerey. Benchmarking autoregressive conditional diffusion models for turbulent flow simulation, 2024. URL https: http://arxvvvvv..raa/as/aaaa orggabs/23099 01745.
Armand Kassaī Koupai, Jorge Mifsut Benet, Yuan Yin, Jean-Noël Vittaut, and Patrick Gallinari. Geps: Boosting generalization in parametric pde neural solvers through adaptive conditioning. Advances in Neural Information Processing Systems, 38, 2024.
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in neural information processing systems, 32, 2019.
13

---

Xingjian Leng, Jaskirat Singh, Yunzhong Hou, Zhenchang Xing, Saining Xie, and Liang Zheng. Repa-e: Unlocking vae for end-to-end tuning with latent diffusion transformers. arXiv preprint arXiv:2504. 10483, 2025.
Tianhong Li, Yonglong Tian, He Li, Mingyang Deng, and Kaiming He. Autoregressive image generation without vector quantization. arXiv preprint urXiv:2406. 1838, 2024.
Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning. Transactions on Machine Learning Research, 2023a. ISSN 2835-8856. URL
Zijie Li, Anthony Zhou, and Amir Barati Farimani. Generative latent neural pde solver using flow matching. arXiv preprint arXiv:2503.2260/ 2025. httphtp:
org/abs//500/ /abbs///..22000.. 22600.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint wrXiv:2010008892 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations, 2020b. URL https / //tp:///aavv///////e///eeaeaiava org/ abs//003, 03485.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. International Conference Learning Representations. 10 2021. URL attp://arxiv org/ abs/2010. 208.0............
Zongyi Li, Nikola Borislavov Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Prakash Otta, Mohammad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Azizzadenesheli, and Anima Anandkumar. Geometry-informed neural operator for large-scale 3d pdes, 2023b. URL https://arxiv. org abt/ 00583.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling, 2023. URL btttps///ariis http://awww...rrabba//aaaaa httgg//b//211.002///
Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky T. Q. Chen, David Lopez-Paz, Heli Ben-Hamu, and Itai Gat. Flow matching guide and code, 2024. URL
org/ abs/ 22222.2222222... 06264.
Phillip Lippe, Bastiaan S. Veeling, Paris Perdikaris, Richard E. Turner, and Johannes Brandstetter. Pde-refiner: Achieving accurate long rollouts with neural pde solvers. In Neurips, pp. 1-36, 2023.
URL http:://aaxiii org/ abs//3/8..222/// 230.0220000532.
Yuxuan Liu, Jingmin Sun, and Hayden Schaeffer. BCAT: A block causal transformer for pde foundation models for fluid dynamics. arXiv preprint arXiv:2501.18972 2025.
Zichao Long, Yiping Lu, Xianzhong Ma, and Bin Dong. PDE-net: Learning PDEs from data. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 3208-3216. PMLR, 10-15 Jul 2018. URL
mlr press/v80/ ht/////01lon//// /woggn/nnnnna html.
Lu Lu. Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. Nat Mach Intell, 3:218 2229, 10 2021. doi: 10. 00.1008/422266-0211 0000000-- URL http: http///axvv.../////v///////aaatl org/ abs/1910. 0w193htth: / //xxx doi. org/ 10. 0205:/8/22256021---222
Michael McCabe, Bruno Régaldo-Sain Saint Blancard, Liam Holden Parker, Ruben Ohana, Miles Cranmer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Geraud Krawezik, Francois Lanusse, et al. Multiple physics pretraining for physical surrogate models. arXiv preprint trXiv:2310.02994 2023.
Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschannen. Finite scalar quantization: VQ-VAE made simple. In The Twelfth International Conference on Learning Representations, 2024. URL https:// //tp://////////eeee//// wwt.:eettettooooeddiiiiiiiee
14

---

Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM, 65(1):99-106, 2021.
Rudy Morel and Edouard Oyallon. Disco: learning to discover an evolution operator for multi- physics-agnostic prediction. In ICML, 2025.
S. Chandra Mouli, Danielle C. Maddix, Shima Alizadeh, Gaurav Gupta, Andrew Stuart, Michael W. Mahoney, and Yuyang Wang. Using uncertainty quantification characterize and improve out-of- domain learning for pdes. In ICML, 2024. URL http://arxiv.
10642.
Thomas Müller, Alex Evans, Christoph Schied, and Alexander Keller. Instant neural graphics primitives with multiresolution hash encoding. ACM Transactions on Graphics, 41(4):1-15, jul 2022. doi: 10.11445/5292233333122 URL attps://doo http:::://oo..or/10...o org/ 115.222221...222222222 3530127.
Roussel Desmond Nzoyem, David A.W. Barton, and Tom Deakin. Neural context flows for meta- learning of dynamical systems. In The Thirteenth International Conference Learning Represen- fations, 2025. URL https: / //ppnr/eeeeeeeeeeeee//ereeeeeeee wtterriiewneeeefooooeeee id
Ruben Ohana, Michael McCabe, Lucas Meyer, Rudy Morel, Fruzsina Agocs, Miguel Beneitez, Marsha Berger, Blakesly Burkhart, Stuart Dalziel, Drummond Fielding, et al. The well: large-scale collection of diverse physics simulations for machine learning. Advances in Neural Information Processing Systems, 37:44989 2024.
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. arXiv preprint arXiv:
2017.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf. Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019. URL http::/ /wwperaapppeeaeeaaa/eee neurips cc/paprr/ pet-pocttcchannn an- an-imperative style- high- wwwppprrrrrroeeeee deep learning---.. library. pdf.
John E. Pearson. Complex patterns in a simple system. Science, 1993. doi: 10.1126/science.261 5118.189.
William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprini 222:02221 2022.
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, and Peter W. Battaglia. Learning mesh- based simulation with graph networks. International Conference on Learning Representations, 10 2021. URL httttt/pp///iiii
httg//bb//0000./ 03409.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022. URL
Ilan Price, Alvaro Sanchez-Gronalezz Ferran Alet, Tom R. Andersson, Andrew EI-Kadi, Do- minic Masters, Timo Ewalds, Jacklynn Stout, Shakir Mohamed, Peter Battaglia, Remi Lam, and Matthew Willson. Probabilistic weather forecasting with machine learning, Nature, 637:84-90, dec 2025. doi: 10. 18835/441556-024-88222222 URL https: //www/ nature. com/ chom/arriclsss s41586- 01-55666------8 02252-9....
M. Raissi, Perdikaris, and G. E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, 2019. ISSN 10902716. doi: j0.1016/ 10.055.
15

---

Robin Rombach, Andreas Blattmann, Dominik Lorenz. Patrick Esser, and Björ Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Louis Serrano, Lise Le Boudec, Armand Kassaï Koupai, Thomas X Wang, Yuan Yin, Jean-Noël Vittaut, and Patrick Gallinari. Operator learning with neural fields: Tackling pdes on general geometries. Neural Information Processing Systems, 2023.
Louis Serrano, Thomas X Wang, Etienne Le Naour, Jean-Noël Vittaut, and Patrick Gallinari. Aroma: Preserving spatial structure for latent pde modeling with local neural fields. arXiv preprint arXiv:2406.02176, 2024a.
Louis Serrano, Thomas X Wang, Etienne Le Naour, Jean-Noël Vittaut, and Patrick Gallinari. Aroma: Preserving spatial structure for latent pde modeling with local neural fields. In NeurIPS, 2024b. 2obs/2406.02 02176.
URL http:://aattii wwt./tsg//////aoooeo/....t..eer/
Louis Serrano, Armand Kassai Koupaï, Thomas X Wang, Pierre ERBACHER, and Patrick Gallinari. Zebra: In-context generative pretraining for solving parametric PDEs. In Forty-second Interna- tional Conference Machine Learning. 2025. URL https://openreete
22kNOkkokU.
Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Im- plicit neural representations with periodic activation functions. Advances in Neural Information Processing Systems, 33:7462-7473, 2020.
Shashank Subramanian, Peter Harrington, Kurt Keutzer, Wahid Bhimji, Dmitriy Morozov, Michael W Mahoney, and Amir Gholami. Towards foundation models for scientific machine learning: Char- acterizing scaling and transfer behavior. Advances in Neural Information Processing Systems 37 (NeurIPS 2023), 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. CoRR, 2023.
Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators. In The Eleventh International Conference on Learning Representations, 2023. URL
httts::/openrrviiee wwwpenrwwvieeeeeeeeeeeeeeeeeeeee wwt.:ee//oooooooeeeee wwwee//oormmiiii het::/i=mmmiiiim imniiMMMPPPP....
Sifan Wang, Jacob H Seidman, Shyam Sankaran, Hanwen Wang, and George Pappas Paris. Cvit: Continuous vision transformer for op-erator learning. arXiv preprint arXiv:2405. 13998, 3, 2024.
Steffen Wiewel, Moritz Becher, and Nils Thuerey, Latent space physics: Towards learning the temporal evolution of fluid flow. Computer Graphics Forum, 38(2):71-82) 2019.
Haixu Wu, Huakun Luo, Haowen Wang, Jianmin Wang, and Mingsheng Long. Transolver: A fast transformer solver for pdes on general geometries. In International Conference on Machine Learning, 2024a.
Tailin Wu, Willie Neiswanger, Hongtao Zheng, Stefano Ermon, and Jure Leskovec. Uncertainty quantification for forward and inverse problems of pdes via latent global evolution. In AAAI, volume 38, pp. 320-328, 2024b. ISBN 1577358872. doi: 10.1609/aaii 1..09/aaiivv3i1..7777
Jiawei Yang, Katie Luo, Jiefeng Li, Congyue Deng, Leonidas Guibas, Dilip Krishnan, Kilian Q Weinberger, Yonglong Tian, and Yue Wang. Denoising vision transformers, 2024. URL https: //arxiv. org/ /201. 02957.
Liu Yang, Siting Liu, Tingwei Meng, and Stanley J Osher. In-context operator learning with data prompts for differential equation problems. Proceedings of the National Academy of Sciences, 120 (39):22311140)) 20222200, 2023.
Yuan Yin, Ibrahim Ayed, Emmanuel de Bézenac, Nicolas Baskiotis, and Patrick Gallinari. Leads: Learning dynamical systems that generalize across environments. Neural Information Processing Systems, 2022a.
16

---

Yuan Yin, Matthieu Kirchmeyer, Jean- Yves Franceschi, Alain Rakotomamonjy, and Patrick Galli- nari Continuous pde dynamics forecasting with implicit neural representations. International Conference on Learning Representations, 9 2022b. http: //axxiv.. org/abs/2205 14855.
Lijun Yu, Yong Cheng, Kihyuk Sohn, José Lezama, Han Zhang, Huiwen Chang, Alexander G Hauptmann, Ming- Yang, Yuan Hao, Irfan Essa, et al. Magvit: Masked generative video transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10459-10469, 2023.
Lijun Yu, José Lezama, Nitesh B. Gundavarapu, Luca Versari, Kihyuk Sohn, David Minnen, Yong Cheng, Vighnesh Birodkar, Agrim Gupta, Xiuye Gu, Alexander G. Hauptmann, Boqing Gong, Ming-Hsuan Yang, Irfan Essa, David A. Ross, and Lu Jiang. Language model beats diffusion tokenizer is key to visual generation, 2024. URL htw.::tttt///ettm /ttp:v//aaiivveeee org/ abs/ abs/2310....//////////////////// 05737.
Anthony Zhou, Zijie Li, Michael Schneier, John R Buchanan and Amir Barati Farimani. Tex12pde: Latent diffusion models for accessible physics simulation, 2025. URL https: / aarxiv org/ abs/2410. 01153.
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We clearly state our claims in the introduction and illustrate them with an extensive experimental analysis.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss some limitations in the conclusion, in dedicated paragraph. More- over additional discussion on computational efficiency wil be provided in the appendices.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sellsspeciiiationn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
18

---

Answer:
[NA]
Justification:
The paper do not contain theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide as much as possible details regarding dataset generation, archi- tectures and training strategy in the appendices. Code will be released upon acceptance for reproductibility.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
19

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: We will release the code upon acceptance for reproducibility. In the mean time, we precisely detail all training and inference details the appendices.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr
onoooooooooooo icy for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nips. wctpublic//aeccccccccccccccc ggwii/gg///////eeee// /Ce///eeembbiiic cttsSSubbissioo.. ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: training details, byperparameters, architecture, optimizers... are listed in the appendices.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Training all models and baselines on dedicated datasets can be computationally demanding. Thus we did not report errors bars of our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
20

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Computational resources details are presented the appendices.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: /ttppp::///errp...... ctpp.cc//puue
Answer: [Yes]
Justification: Authors have reviewed the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss broader impact in the conclusion.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
21

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper do not make use of scraped or personal data.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Origin of existing assets used are detailed in the code (mainly baselines).
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi wwwode.hho....c. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
22

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
Code and instructions for use will be released upon acceptance.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
23

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: Appart from for writing or formatting purposes, we did not make use of LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy (https: : : http///urrr.cccccc httprrpssecc//ooeeeennneee////// for what should or should not be described.
24

---

Appendix Table of Contents
A. Notation
26
B. Related Works
277
Operator Learning
.27
Generative Models
27
Parametric PDEs
28
C. Datasets Details
29
Combined Equation
29
Advection Equation
31
Wave Equation
33
Gray-Scott Equation
34
Vorticity Equation
36
D. Architecture Details
33
ENMA Encoder-
37
* Interpolation via Cross-Attention
37
* Compression via Causal Convolutions
39
Tokenwise Autoregressive Generation Patchification of Latent Codes
41
Causal Transformer
Spatial Transformer
41
E. Implementation Details
Dynamics forecasting: Implementation
*
Architecture Configuration
Baseline Details
.3
Training Details ................................
Implementation Implemeeaaiioo..................
* Architecture Details
46
Baseline Details
Training Details ................................
Dynamics forecasting on Complex systems: Implementation
F. Additional Experiments
49
ENMA Process Task
Generative Ability of ENMA
49
* Forecasting with Less Compression
52
* Inference Speed Comparison
53
* Ablation Studies
Experiments on the Encoder/Decoder
55
Additional Experiments
datasets
55
* OOD Encoding with 10% Input Grid
56
* Super-Resolution
56
* Ablation studies
57
Gearetry-Aware Attention Analysis
Visualization in token space
60
G. Visualization
63
Combined Equation
63
Advection Equation
64
Gray-Scott Equation
65
Wave Equation
66
Vorticity Equation
68
25

---

A
Notation
We summarize the main notations used throughout the paper:
Symbol
Deseription
æ
Spatial coordinate
t
Temporal coordinate
(ar, 1)
Solution of the PDE at (x,
T
Final time
Observation ratio (proportion of available inputs) Input spatial grid
x
X
Test-time input spatial grid
Y
PDE parameters
M
Differential operator
Boundary condition operator
True temporal evolution operator
Neural approximation of the evolution operator
Z
True latent state from VAE tokenizer
M
Number of spatial tokens in Z Spatial token, with Z (ZI,..., ZM) Number observed time steps (history)
20:-1
Latent sequence of past states Predicted latent state from the Causal Transformer
ZoYN
2
Contextualized latent tokens from Spatial Transformer
Spatial token Ž Predicted latent state
2
2
Spatial token in 2 ( EM)
M
Masked token indices during training Number of autoregressive steps
S
s
Autoregressive step index (s E [0, S0,S-1) 1])
N.
Masked token indices to predict at step Flow matching step index Velocity field for flow matching Gaussian noise
Regular interpolation grid Interpolation grid point (& E E)
g
Loss function
Q
Query tokens for attention
K
Key tokens for attention
V
Value tokens for attention
m
Scaling factor in cross-attention bias
26

---

B Related Works
B.1
Operator Learning
Operator learning has emerged as a powerful framework for modeling mappings between infinite- dimensional function spaces. Foundational works such DeepONet and the Fourier Neural Operator (FNO) established neural operators (NOs) as effective tools for learning these mappings (Li et al., 2020a; Lu et al., 2021). Subsequent research has sought to improve both expressiveness and efficiency, exploring factorized representations for reduced complexity, wavelet- based techniques for multi-scale modeling (Gupta et al., 2021), and latent-space formulations to support general geometries (Tran et al., 2023; Li et al., 2023b). Implicit neural representations, as in CORAL, have also been proposed to accommodate variable spatial discretizations at inference time (Serrano al., 2023).
recent direction in operator learning leverages transformer-based architectures, inspired by their success in vision and language tasks. OFormer introduced transformer model for embedding input-output function pairs, demonstrating the potential of transformers for operator learning (Li et al., 2023a). This has since led to more advanced architectures such as GNOT and Transolver, which improve input function encoding and generalization to irregular domains (Hao et al., 2023; Wu al., 2024a). Aroma and UPT further adopt perceiver-style designs to learn compact and adaptable latent neural operators (Serrano et al., 2024a; Alkin et al., 2024b).
One close work is CVTT, which learns compressed spatio-temporal representations using transformer encoders and decoders (Wang et al., 2024). In contrast, our approach first applies an attention- based encoder to map irregularly sampled fields to a uniform latent representation, followed by causal spatio-temporal convolutional encoder that accommodates a flexible number of input states, enabling both regular and irregular conditioning.
B.2 Generative Models
While most operator learning methods deterministic, generative modeling introduces critical capabilities for modeling physical systems-most notably the ability to represent uncertainty and capture one-to-many mappings, which especially relevant in chaotic or partially observed regimes. Two primary generative paradigms have emerged in this context: diffusion models and autoregressive transformers.
Diffusion Transformers Diffusion models synthesize data by learning to reverse a progressive noising process, typically through sequence of denoising steps (Ho et al., 2020). In computer vision, Latent Diffusion Transformers (DiT') have demonstrated strong performance by applying this principle in the latent space of VAE (Peebles & Xie, 2022). More recently, diffusion-based techniques have been adapted to scientific modeling. For example, Kohl et al. (2024) proposed an autoregressive diffusion framework tailored for PDEs, particularly in turbulent settings where capturing stochasticity is essential. Similarly, Lippe et al. (2023) enhanced the modeling of high- frequency chaotic dynamics via noise variance modulation during denoising. Zhou al. (2025) extended DiT's to the physical domain, generating PDE data from textual prompts.
Despite their expressiveness, diffusion models often require many sampling steps and lack efficient in-context conditioning. As an alternative, Flow Matching has recently gained attention. Instead of discrete denoising steps, it learns continuous-time velocity field that transforms one distribution into another via an ODE. This leads to significantly faster sampling with far fewer steps, making well-suited for scientific applications (Lipman et al., 2023, 2024).
Autoregressive Transformers Autoregressive models, originally designed for language modeling. have been successfully extended to image and video domains by treating spatial and temporal data as sequences. These models often couple a VQ-VAE with a causal or bidirectional transformer to model discrete token sequences (Oord et al., 2017; Esser et al., 2021; Chang et al., 2022). In video generation, frameworks such as Magvir and Magvit2 Yu al., 2023, 2024) encode spatiotemporal information via 3D CNNs and generate frames autoregressively over quantized latent tokens. In the context of PDE modeling, Zebra adapts this paradigm by combining a spatial VQ-VAL with a causal transformer for in-context prediction (Seraano et al., 2025). However, the use of discrete codebooks limits expressiveness and may hinder the ability to represent lie-ggrineed physical phenomena, which are inherently continuous.
27

---

To address this, recent advances in vision have introduced masked autoregressive transformers that operate directly on continuous latent tokens (Li al., 2024). These Masked Autoregressive (MAR) models predict subsets of tokens using diffusion-style loss, allowing multi-token generation without relying on vector quantization. Notably, block-wise causal masking enables efficient inference while preserving temporal consistency (Deng et al., 2024). ENMA builds on this approach by adapting MAR for neural operator learning: our model learns to autoregress over continuous latent fields, supporting in-context generalization while maintaining high fidelity and computational efficiency.
B.3 Parametric PDEs
Generalizing to unseen PDE parameters is a central challenge in neural operator learning. Approaches span classical data-driven training, gradient-basee adaptation, and emerging in-context learning strategies.
Classical ML Paradigm A standard approach trains models on trajectories sampled from dis- tribution over PDE parameters, aiming to generalize to unseen configurations. This often involves stacking past states input channels (Li et al., 2021) or along a temporal axis, analogous to video modeling (Ho al., 2022; McCabe et al., 2023). However, such models typically degrade under dis- tribution shifts, where small parameter changes lead to significantly different dynamics. Fine-tuning has been proposed to address this (Subramanian et al., 2023), but often requires substantial data per new PDE instance (Herde et al., 2024; Hao et al., 2024).
Gradient-Based Adaptation To improve generalization, several methods train across multiple PDE environments. LEADS(Yin et al., 2022a) employs shared model with euvironment--pecfffe modules, updated during inference. Similarly, Kirchmeyer et al. (2022) introduce a hypernetwork conditioned on learnable context vectors cx, enabling adaptation to new environments. Building on this, NCF(Nzoyem et al., 2025) uses a Taylor expansion to dynamically adjust context vectors, offering both uncertainty estimation and inter-environment adaptation.
In-Context Learning for PDEs
Inspired by LLMs, recent work explores in-context learning (ICL) for PDEs. Yang al. (2023) propose transformer that processes context-query pairs, but scalability is limited to ID or sparse 2D settings. Cao et al. (2024) extend this to PDEs via patches for a vision transformer. They introduce special conditionning mechanism, where input-output function given as inputs to the transformer, allowing to handle flexible time discretizations at inference. Chen et al. (2024) propose unsupervised pretraining of neural operators, followed by ICL-style inference via trajectory retrieval and -iithouu generative modeling. Serrano et al. (2025) address this by introducing causal transformer over discrete latent tokens, paired with an effective in-context pretraining strategy, allowing adaptation to new PDE regimes via in-context learning and probabilistic predictions.
28

---

C Datasets details
To evaluate the performance of ENMA, we generate several synthetic datasets based on diverse parametric PDEs in ID and 2D. Each dataset consists of 12,000 trajectories for training. For evaluation, we genérate two test sets: 1,200 trajectories for in-distribution (In-D) and 120 for out-of- distribution (Oul-D evaluation. In the In-D case, parameters 7 differ from the training set but sampled from the same distribution; in Out-D, they lie outside this range. Parameter ranges for the in-D and Out-D are presented in Table and further detailed in each dataset section.
Data generation proceeds as follows: for each sampled set of PDE parameters y (as described in section 2), we simulate batch of 10 trajectories using numerical solver from 10 different initial conditions. Trajectories are computed over a time horizon E [0, T on a spatial grid X. This pipeline yields datasets with diverse and complex dynamics.
We detail below the PDEs and parameter ranges used. ID, we consider the Combined equation (appendix C.1) and the Advection equation (appendix C.2); in 2D, we study the Wave equation (appendix C.3), the Gray-Scott system, and the Vorticity equation (appendix C.5).
Table 6: In-distribution (In-D) and out-of-distribttion (Out-D) parameter ranges for each dataset.
Dataset
Parameter
In-D
Out-D
4((0.3,0.5)
4([0.3,0.5]
3
4([000005, 0.5))
2{[000005, 0.5))
Combined
4([0.01,1) 1))
4([0.01,1) 1))
U([0.5, 1])
Advection
5]]
u([77, -5] U [5, 7J)
L((1100, 500))
4((500, 550))
Wave
4400, 50])
(550, 60))
F
u(00.023, 0.045))
4([00045, 0.0467])
Gray-Scott
u([000590, 0.0640))
u([000570, 0.0590J)
Vorticity
u(1103,100
(((10-9,10-9)
C.1 Combined Equation
The Combined equation (Brandstetter et al., 2022) unifies several canonical PDEs-such the heat and Korteweg-de Vries equations-by varying a set of coefficients (ax, 8, 7), enabling the modeling of diverse dynamical behaviors. is defined as:
@u(x, 1) ① Ōt dx (au(z,t) t)² )u(x,() 1) Ox
+ 0u((, = 0, (z, E Ω (0.T),
Ox
u(I, =u (x), ZEΩ,
with parameters sampled uniformly as a 4((003 0.5)). ,, B 1000~~u10000 u(0000055 0.5)), and y 4(00011, 1)) for both training and in in-domain test sets.
Simulations over the domain li x [0, = [0, 2ml x [0, 1]. The initial condition is defined as follows:
N w (x()) = Σ i=1
2mlIx
+ dz).
A sin
L
with domain length 2x, amplitudes A; u(-0..5 0.5)). phases Φ; U(, 2n)), and frequen- cies 1 ~ (1(11,2 2, ()). Each trajectory is simulated for 100 time steps over spatial grid of 256 points. To train our models, we subsample each trajectory to retain only 20 time steps over spatial grid of 128 points.
29

---

Out-Domain
For out-of dis-tiibiitiittionn evaluation, we consider more complex scenario by adding fourth-order spatial derivative. The resulting PDE is:
0u(z, t) Ø
(au(x, 1)2 δt dx
@u(z,,t) t) 7 Or² +
)u((z,t) t)
pu((x,t)
- 0,
dr.
Ox.
with parameters sampled as a ~ 00333 0.5]). 3 ~ (1000005 0.5)). 7 ~ UQ00.11, 1J), and 8 ~ U((55, 1]). Aside from this modification, simulations follow the same configuration the in- distribution setup. Visualizations of the combined equation for in and out-domain trajectories are showed in Figure 4 and 5 respectively.
Combined Equation Trajectory
Combined Equation Trajectory
Naandien: a p-0.5..5...4.
Parambers p=0 p-2..
66-
52-
04-
0.
Thhe
-.4
2.4
-011
0.2
-24
64
0.
0.6
0.6
3.0
06
02
0.4
3.6
0.8
(a)
(b)
Figure 4: Samples from the Combined Dataset.
Combined OOD Equation Trajectory 0. 0.556-
0.8
Spazion
2.4
-0.
0.0
0.2
0.6
0.5
1.0
Figure 5: OOD sample of the Combined equation.
30

---

C.2 Advection Equation
The advection equation models the transport of quantity constant speed. For out-domain sets, the advection speed is sampled from
For our experiments, it is solved over the domain (I2 x [0, ((0, 1] [0, and is defined as:
Öu( t)
du(x,t)
(xx, t) E 2 (0,T),
Ot u(x, 0) u (x)
dæ
x E 92
where the advection speed is sampled from (1([5,,) 5)). For out-of distribbtinnn (Out-D) evaluation, we increase the difficulty by sampling the advection speed a from a uniform distribution 2/[[7, - -5 [5, 7]), explicitly excluding the training range [-5,5) 5). For both in-domain and out-domain sets, generate trajectories defined by three types of initial conditions:
Sine sum:
N Σ
2
2 (x)
A sin
+ Φ1
Cosine sum:
12T120 + Φz
u (x)
A, cos
Sum of sine and cosine:
Σ A sin 2T
1
" (x)
+ thi
008
+ Φ1)
L
where 1, 0([-0.50.55), di ((10, 2T)), and 4 ~ (((1( (1, 3)). Each trajectory is
simulated for 100 time steps over spatial grid of 1024 points. To train our models, we subsample each trajectory to retain only 20 time steps over a spatial grid 128 points. Visualizations of the advection equation for in and out-domain trajectories are shown in Figures 6 and 7.
Advection Equation Trajectory
Advection Equation Trajectory
3.3
Paramitter 0
1.0
3.0
M
0.00
04
0.6
0.4
Time
0.0
0.4
-030
60-
20
-0.25
-00
80
02
0.4
06
0.4
6.
(a)
(b)
Figure 6: In-domain samples from the Advection Dataset.
31

---

Advection Equation Trajectory Parameten. a= -579
0.6
20
S0UABOTL
-01
-03
Figure 7: Out-domain sample from the Advection equation.
32

---

C.3 Wave Equation
We consider 2D damped wave equation defined over the domain [0, [0. 1]2 [0, 0.005.
given by:
(w t) e² Awww t) Ow(z,t) = 0. dt2
(x, E x (0, T)
Ot
w(x, 0) w (x))
92
where ((1000 ) is the wave speed and k ~ 21(10, 50) the damping coefficient for the In-D datasets. For Out-D, we sample (((500, 550)) and k u[50, 60). We focus on learning scalar field w. The initial condition is defined sum of Gaussians:
N
(x 22L)² + (y 3:L)2
(x,
Σ exp
20
with z4. yr ~ u( 1]), o; ~ u([.0025, 0.1]). = 1, and N ~ u(f2, 3, 4}). Simulations are performed on a 64 64 spatial grid with 30 time steps. Two in in-domain trajectories can be visualized in Figure 8 and in Figure 9 for out-domain trajectories.
Wave Equation Trajectory
Parameters: 42.42, k 285.86
1=7
t=14 14
t= 21
28
(a) Wave Equation Trajectory
Parameters: 15.66, 168.69 t-
1-0
21
t-
0.5
00
(b)
Figure 8: Samples from the Wave Dataset.
Wave Equation Trajectory
Parameters: = 52.50, k 525.00
1=7
1:
Salitton
0.0
-0.5
Figure 9: OOD sample of the Wave equation.
33

---

C.4
Gray-Scott Equation
The PDE describes a by the following 2D equations:
system generating complex spatiotemporal patterns, governed
du dt du
Du Au ui² + F(1 - u),
D,Av un (F k)u,
where u, represent the concentrations of two chemical species over 2D spatial domain S. with periodic boundary conditions. The diffusion coefficients are fixed across all trajectories: D, = 0.102 and D, 0.204. Each PDE instance 7 defined by variations the reaction parameters. For In-D datasets, we sample F ~ 24( 0.045)) and k ~ L((000559, 0.0644)). For Out-D evaluation, we sample F ((.0045, 0.0467] and ~ 24(00.570, 0,00590)). The spatial domain is discretized 32 x 32 grid with spatial resolution As 2. Trajectories presented in Appendix C.4 for in-domain and in Figure 11 for out-domain.
Gray-Scott Equation Trajectory
Parameters: F= 0.04, 0.06
1=0
5
I- 10
1 15
1- 19
0.6
0.4
0.2
0.0
1.0
0.8
0.6
0.4
0.2
(a)
Gray- Equation Trajectory
Parameters: F- 0.04, 0.06
1
5=5
1-
T 25
0.8
0.6
0.4
0.2
0.0
80
1.0
0.8
0.6
D.4
0
(b)
Figure 10: Samples from the Gray-Scott Dataset.
34

---

Gray-Scott OOD Equation Trajectory
Parameters: F= 0
1-0
5
t- 10
15
t- 19
0.8
0.6
0.4
o
0.2
0.0
1.0
0.8
0.6
oo
0.
0
0.2
Figure 11: OOD sample of the Gray- Scott equation.
35

---

C.5
Vorticity Equation
We consider 2D turbulence model and focus on the evolution of the vorticity field w, which captures the local rotation of the fluid and is defined w V x u, where u is the velocity field. The governing equation
dw
Ot + (u- Vjw- v vマ = 0,
(9)
where denotes the kinematic viscosity, defined
1/Re. For In-D datasets, sample u([10- 10 2)), while for Out-D, we consider a more challenging turbulent regime with u([10 10 *[). The initial conditions generated from the energy spectrum:
4 V ((k)
h ku
1
h: Ko
exp
(10)
ko
where kn denotes the characteristic wavenumber. Vorticity is linked to energy by the following equation:
B(k)
w(k)
(11)
In-distribution and out-distribution trajectories be visualized in Figure 12 and fig. 13 respectively.
Vorticity Equation Trajectory
Parameter: 4.30
0.0
Vorticity Equation Trajectory
0.5
-.5
(b)
Figure 12: Samples from the Vorticity Dataset.
Vorticity Equation Trajectory
Parameter: v 1.6e
10
Figure 13: OOD sample of the Vorticity equation.
36

---

D Architecture details
D.1 ENMA encoder-decoder
We describe the architecture of the encoder-decoder modulc in detail below. The overall encoding and decoding pipeline is illustrated in Figure 14. Starting from an input physical field u u0 defined on a (possibly irregular) spatial grid Ai.. the encoder first applies an attention-based interpolation module (appendix D.1.1) to project the field onto a regular intermediate grid E, producing u.:L-1(3) This interpolation operates purely in space and is applied independently at each timestep.
Next, (=: passed through causal convolutional network that encodes the data in both space and time (appendix D. 1.1.2), yielding latent tokens 20=2,-1 of length La < L when temporal compression is used. Temporal compression optional, for notation simplicity, will assume do not compress time dimension. These latent tokens form the input to the autoregressive dynamics
model.
The decoder mirrors the encoder's structure. It first upsamples the latent tokens back to the intermedi- grid E (appendix D.1.2), then uses cross-attention module reconstruct the physical field any desired output grid Nous- This yields the predicted trajectory 00:t-1 (Koue). The final stage employs the same cross- attention mechanism the encoder's 's initial interpolation module (appendix D.1.1).
(Xi))
(Rad
Interpoiation via Crose-Attertion
Interpoiation via Crose-Attention
Xa
Reguiar Repessentation (4---1))
Causal compression layers
upsampling layers
Figure 14: General architecture of ENMA's encoder/decoder.
D.1.1 Interpolation via cross-attention
Figure 15 shows the detailed architecture of the interpolation module used in the encoder. The decoder mirrors this structure by directly reversing its operations.
The interpolation module takes as input a spatio-temporal field 0:L-1 E R Xx Lxe defined on a potentially irregular spatial grid Xr,- Positional encodings first applied the input coordinates as in Mildenhall et al. (2021), 1), and the physical field is projected into a higher-dimensional representation using linear layer, yielding embeddings Pe and h, which serve the keys and values for the cross-attention module. Following Serrano et al. (2024a), the queries are learned latent embeddings that act as intermediate representation.
To promote structured and localized interactions, we also initialize a learned coordinate grid for these latent queries. These coordinates are used to compute attention biases, encouraging stronger attention between spatially adjacent tokens. The cross-attention module outputs an intermediate field x which further refined using Physies Attention (Wu et al., 2024a) to reduce computational overhead.
37

---

Irregular input field X c
Irregular input coords
Linear -
Positionnal encoding
Values X h
Keys | x P-
K
Regular intermediate grid | d
Cross-Attention Module with geometry aware bias
[QK*
Regular intermediate field | Greg
Attention(Q, K, V) Softmax
Big = m dist(a; Es), for 2; E 2, C E
Intermediate field
E x h
Residual connexion
Linear
X depth
h
+
Physics self attention
x h
Linear
Interpolated intermediate field [
Figure 15: Detailed architecture of the Interpolation module.
38

---

Geometry-awars encoding bias that promotes attention locality. Appendix D.1.1 illustrates how this penalization applied. Formally, we define the cross-attention from queries located at X to keys and values located at = as:
As presented in section 3.2, we introduce a geometry aware attention
[QKT + ( vd
Attention (Q, K, V) = Softmax
)V, V.
Br,, = - - distiit( - Ey), for 2z E X E =
where Q ((),, K = k(3,, and = v(3) are the query, key, and value embeddings defined over X' and E, respectively; dg is the key/query dimensionality; and m scaling factor, either fixed or learned as in ALiBi (Press et al., 2022). Intuitively. the bias B imposes stronger penalties for distant pairs (zi, E,), reducing their attention weights and encouraging the model to focus on local interactions during interpolation.
different inputs keys/values
different positions keys/values
K k(2)
(2) = δ
Attention
QK"
Attention penalization
high
Figure 16: Geometry-aware Alibi bias.
D.1.2 Compression via causal convolutions
For the architecture of the CNN, we adopt a design similar to that of Yu et al. (2023). our case, it takes as input the output of the interpolation module, u3: and lifts it to higher- dimensional representation tensor is then compressed through stack of three building blocks: residual, and compress, time. Finally, a concluding residual block projects the representation back to the target latent dimension. The up-sumpp up-sampling module strictly mirrors this compression pipeline. This type of architecture has been shown to effective for ppatio-temporal compression (Serrano et al., 2025; Yu et al., 2023). The three types of layers used in the compression module detailed below, and the full architecture is shown in fig. 17.
size hemmp using a causal convolution layer (see fig. 3). The lifted
The residual block processes the input while preserving its original shape. It consists of a causal convolution with kernel size k. followed by a linear layer and a Global Context layer adapted from Cao et al. (2020). If the output dimensionality differs from the input's, ss additional convolution is used to project the spatial channels to the desired size.
residual blocks:
The compress_spen space block reduces spatial resolution by factor of 2d. i.c., each spatial dimension is downsampled by a factor of 2 using a convolutional layer with stride 2. The kernel size & and padding p set accordingly, with k//2. To ensure that only spatial dimensions are compressed, inputs are reshaped so that this operation does not affect the temporal axis.
compress saprress_ppcceeeee blocks:
compress_tar time blocks: The compress, time block performs temporal compression similarly to the spatial case, but operates along the time dimension. To preserve causality, padding is applied
39

---

only to the past, with size = h 1, so that frame at time attends only to frames at times < t. A convolution with stride is used to reduce the temporal resolution by factor of 2.
Interpolated intermediate field
L E x
Residual
Group norm
à h, &
Activation
temporal padding
Group norm
Activation
Global Context h
causal conv hoomp: Rin
Compress space
Reshape 4. E - A,
Compression layers
h. k//2 pading
Residual
Compress space
Compress time
Reshape 5,
Compress time
h, h,
Residual
k 1 Causal padding
convid k 8
Reshhape h,
Compressed tokens
Z 701-1
Figure 17: Detailed architecture of the Compression module.
40

---

D.2
Tokenwise Autoregressive Generation
We provide additional details about our generative process for tokenwise autoregression of physical fields.
D.2.1
Patchification of Latent Codes
Given sequence of physical states ux 0:7-11 the VAE encoder produces a sequence of latent represen- tations z RMxxxx where M the number of spatial tokens. For 2D physical systems, the latent representation has spatial dimensions M,, xM2 with M M1 x M2 and M1 = M2 in our setting. The number tokens per frame plays a critical role in capturing fine-gu ane-grained spatial details. However, it also directly impacts the computational cost of tokenwise autoregressive generation, which scales with the number of tokens in each latent state.
To mitigate this cost while preserving spatial structure, we adopt a patchify strategy, commonly used in vision transformers. Specifically, we apply spatial down-sampling the latent representation by dividing it into non-overlapping patches, reducing the token count per frame while retaining local coherence (see Figure 18). This significantly improves inference efficiency without sacrificing modeling fidelity.
panchyfy with patch sizes PiuPs
mamaaggy
umb
dp
d
R M
Figure 18: Illustration of latent patchification. The spatial latent map ZM,xM is partitioned into coarser patches to reduce the number of autoregressed tokens.
D.2.2 Causal Transformer
The causal Transformer introduced in section 3.1.1 captures the dynamics necessary to predict the next latent state - from the past sequence Zu:L-1, We implement next-state prediction strategy using block-wise causal attention mask, where tokens within each frame can attend to another, and tokens in frame can attend to all tokens from earlier frames j < é. The attention mask used is defined as:
AMXM
Oа 1axM
0Mx M 0мжм
1
Matto
T rows
1MXM
1MXM
1MxM
T columns
This structure ensures that temporal causality is respected while allowing rich intra-frame interactions. It predicts latent code Z Caussa/rrassoommer/ ZBas, 20:L-1) that captures the dynamics given the input sequence history.
For positional information, we apply standard sine-cosine embeddings along spatial dimensions. In the temporal direction, we omit explicit positional encodings, relying instead on the causal structure to encode temporal ordering implicitly (Kazemnejad et al., 2023).
D.3 Spatial Transformer
At inference time, given target number of steps S, the spatial Transformer generates the final latent frame 2 in an autoregressive manner over S steps. Each step selectively predicts subset of masked tokens, conditioned on the dynamic context Z CausalTransformest Zans. 20:L-1). The token prediction schedule follows a cosine scheduler, progressively revealing tokens in smooth, annealed fashion. Figure 19 illustrates how tokens generated at inference for S steps.
41

---

STEP1
STEP
Matching
By
Doley
Matching
fx
2
fx
2 (AN
2-
[&]
2
Wi = nadeas
seketel
Spetial Transformer
Z
[MASK]
Z
MASK]
STEP 3
STEP
Flow Matching
suicy
Flow Matching
MLP
2y
2' [N
2
Z'
Sodess
W. notoss
2
Spatial Transformer
Spatial Transformer
[MASK]
Z
[MASK]
Masked takens
Predicted takens for generation
Selected tokens for generation
Figure 19: Illustration of tokenwise autoregressive generation inference with spatial Transformer over S = steps. At each step, a subset of tokens is selected for generation based on a cosine schedule.
42

---

E
Implementation details
The code has been written in Pytorch (Paszke al., 2019). All experiments were conducted on a A100. We estimate the total compute budget-including development and evaluation-to be approximately 1000 GPU-days.
E.1
Dynamics forecasting: implementation Details
This section outlines the implementation details related to the autoregressive time-stepping and generative experiments with ENMA, along with the corresponding baseline configurations.
E.1.1 Architecture Configuration
To enable fair comparison with existing baselines, we disable the interpolation component of ENMA's encoder-decoder when conducting autoregression fixed grid, consistent with the setup used in Alkin al. (2024a), The hyperparameter settings for ENMA's generation architecture across all datasets are summarized in table 7.
Table Model parametrs configuration for datasets for the ENMA's generation process.
Hyperparameters
Combined
Advection
GS
Wave
Vorticity
Vae embedding dimension
4
8
Number of tokens
16
16
64
256
256
Patch size
1
4
4
Spatial Transformer depth
6
6
Causal Transformer depth
6
6
hidden size
512
512
512
512
512
mlp ratio
2
2
2
heads
8
8
dropout
0
0
QK normalization
True
True
True
True
True
Normalization Type
RMS
RMS
RMS
RMS
RMS
activation
Swiglu
Swiglu
Swiglu
Swiglu
Swiglu
Layer Norm Rank
24
24
24
Positional embedding
Sinus
Sinus
Sinus
Sinus
Sinus
MLP depth
3
3
3
MLP width
512
512
512
512
512
number of steps S FM steps
16
16
16
10
10
10
10
E.1.2
Baseline details
We detail here the architecture of the baselines used to evaluate ENMA's dynamics forecasting experiments. We compareed method FNO (Li et al., 2020a), BCAT (Liu et al., 2025), AVIT (McCabe et al., 2023), AR-DiT (Kohl et al., 2024), Zebra (Serrano et al., 2025), an In-Context ViT and [CLS] ViT done in (Serrano et al., 2025).
FNO For FNO, we followed the authors guidelines and concatenanted the temporal historic directly with the channels. We considered 10 modes for both ID and 2D datascts and used width of 128. We stacked Fourier layers.
BCAT BCAT is deterministic block-wise causal transformer approach for learning spatio-temporal dynamics. has been proposed for tackling multi-physics problems, and we adapted it for parametric PDEs. BCAT performs autoregression in the physical space. As vision transformers (Dosovitskiy
43

---

et al., 2021), it relies patchs to reduce the number of tokens. We report the model hyper- parameters
details for the Transformer in table 8.
Table 8: Model paper-paraeetess configuration for all datasets for the BCAT process.
Hyperparameters
Combined
Advection
GS
Wave
Vorticity
Patch size
8
8
8
Transformer depth
6
6
hidden size
512
512
512
512
512
mlp ratio
2
2
2
2
num heads
8
8
8
dropout
0
0
0
QK normalization
True
True
True
True
True
Normalization Type
RMS
RMS
RMS
RMS
RMS
activation
Swiglu
Swiglu
Swiglu
Swiglu
Swiglu
Positional embedding
Sinus
Sinus
Sinus
Sinus
Sinus
We evaluate against the Axial Vision Transformer (AVIT) introduced in (Müller et al., 2022). which performs attention separately along spatial and temporal dimensions. Their proposed MPP model extends AViT to multi-physics settings by incorporating strategies for handling multi-channel inputs and system-specift normalization. our parametric setting, such mechanisms are unnecessary. We adopt the same hyperparameter configuration as in table 8, which we found to yield the best results.
AVIT
AR-DIT Our autoregressive diffusion transformer follows the setup proposed in Kohl et al. (2024), where known frames are concatenated with noise and denoised progressively to generate the next physical state. As in (Rombach al., 2022), we employ AdaLayerNorm condition the model during the diffusion process. The corresponding hyperparameters are provided in table 9.
Table 9: Model hyper-paramenee parameters configuration for all datasets for the AR-DIT process.
Hyperparameters
Combined
Advection
GS
Wave
Vorticity
Patch size
8
8
8
8
8
Transformer depth
6
hidden size
512
512
512
512
512
mlp ratio
4
4
4
4
num heads
8
8
8
8
8
dropout
0
0
0
0
Normalization Type
AdaLayer Sinus
AdaLayer Sinus
AdaLayer Sinus
AdaLayer Sinus
AdaLayer
Positional embedding
Sinus
Diffusion steps
100
100
100
100
100
Zebra Zebra is token-based autoregressive model that combines VQ- VAE for discretizing spatial fields with a causal transformer for modeling temporal dynamics. At each time step, autoregressively predicts discrete latent tokens corresponding to the next frame, allowing for uncertainty quantification through stochastic sampling in latent space. The full set of hyperparameters for both the VQ-VAE and the transformer components is provided in table 10.
In-Context ViT
implement an in-context Vision Transformer (Dosovitskiy et al., 2020) to address the initial value problem using context trajectory. The model is trained to forecast the next frame in target sequence, conditioned on both the observed context and preceding frames. For architectural configuration, we adopt the same hyperparameters reported in table 8.
44

---

Table 10: Hyperparameters for Zebra's VQVAE and Transformer components.
Hyperparameters
Advection Combined
GS
Wave
Vorticity
VQ-VAE
start hort_hidden_siee sizee
128
128
128
max hiwdhhdddiiiii size
256
256
1024
1024
1024
num
2
3
2
codebook
256
256
2048
2048
2048
node_dim num_codebooks
64
64
16
16
16
1
shared_codebook codebook
True
True
True
True
True
likens_per_frane
32
32
256
128
256
start learning rate
3e-4
3e-4
3e-4
3e-4
3e-4
weight_ decay
le-4
le-4
le-4
le-4
le-4
scheduler
Cosine
Cosine
Cosine
Cosine
Cosine
num_epochs
1000
1000
300
300
300
Transformer
max_context__ot
2048
2048
8192
8192
8192
batch_size
2
2
2
num nun_ggradint_acccuulattt
4
4
4
hidden_sizz
256
256
384
512
384
mlp_ratio
4.0
4.0
4.0
4.0
4.0
depth
8
8
8
heads vocabulary size
8
8
8
264
264
2056
2056
2056
start learning rate
lc-4
le-4
lc-4
lc-4
le-4
weight_deccs scheduler
le-4
le-4
le-4
le-4
le-4
Cosine
Cosine
Cosine
Cosine
Cosine
num_epochs
100
100
30
30
30
[CLS] VIT This variant of the Vision Transformer is adapted to meta-learning setting, where dedicated [CLS] token captures the environment-spect -speciffttt-sttiiieeeeeeeeeeeeee variations across different PDE settings. During inference, the [CL.S] token is updated via 100 gradient steps to adapt to new environments. We use the same model configuration as in table
E.1.3 Training details
Models for ID datasets trained for approximately 8 hours, while training 2D datasets took around 20 hours. Unless otherwise specified, all datasets followed the same training objective.
Optimizer and Learning Rate Schedule We use the AdamW optimizer with 81 0.9 and 32 = 0.95 for all experiments. The learning rate follows cosine decay schedule, starting from an initial value of 10-3 and annealing to 10-5 over the course of training. To stabilize the early training phase, we apply a linear warmup over the first 500 optimization steps.
Baselines All baselines were trained following the same protocol as ENMA. Each model was trained from scratch for 1000 epochs, and the best-performing checkpoint was selected based on training performance. A cosine learning rate scheduler was used across all methods, which consistently improved prediction quality. All models were trained with a learning rate of le except for AR-DiT, which required lower rate of 11 due to unstable training dynamics.
45

---

Table 11: Training syper-parameters configuration for datasets for the ENMA's generation process.
Hyperparameters
Combined
Advection
GS
Wave
Vorticity
Epochs
1000
1000
1000
1000
1000
batch size
128
128
128
64
64
learning rate
le-3
le-3
le-3
1e-3
1e-3
weight Decay
le-4
le-4
le-4
le-4
le-4
grad clip norm
betas
[0.9, 0.95]
[0.9, 0.95]
[0.9, 0.95]
[0.9, 0.95]
[0.9, 0.95]
E.2 Encoder-Decoder implementation protocol
E.2.1
Architecture details
We present in table the hyper parameters for the architecture of ENMA'S auto encoder.
Table 12: Architecture details of the auto encoder as presented in table 1. appendices D.1.1 and D.1.2 for the notation.
refer the reader to
Module
Block
Parameter
Advection
Vorticity
token dim
intermediate grid size
positional encoding
Positional encoding
positional encoding num freq positional encoding max freq
intermediate grid dim Grey
hidden dim h attenntttttnion heads
Interpolation module
Cross attention block
4
4
cross-attention dim heads
alibi scaling n alibi heads
2,
1,2,3,4
Geometry-aware bias
4
depth
2
2
atlention heads
Physies self attention
attention dim heads
physics-attention attention slice
residual compress, space residual
residual
compress, _pace residual
compress_epee
compress_ting
Compression layers
residual compress_ space residual compress. tite res desidual 16
residual
Compression layers
Compression module
i hemmp compression kernel size
16
7
7
causal input layer causal output layer
kernel size kin kernel size kout
7
7
7
7
E.2.2
Baseline details
We detail here the architecture the baselines used to evaluate ENMA's auto encoding. Hyper- parameters are chosen from the original papers and/or closest available implementation. We recall that Oformer acts pointwise the physical space, and thus does not make any spatial nor temporal compression. CORAL, AROMA and GINO act at a frame level i.e. they compress space but not time. Finally, our ENMA's auto encoder architecture performs both a spatial and a temporal compression. We provide comparison in table 13.
46

---

Table 13: Token sizes used evaluate auto encoders as presented in table We show sizes using the following formatting: temporal size spatial size x token dimension.
Model
Advection
Vorticity
Trajectory sizes
50 x 128 1
30 x (64 x 64) x
Oformer
50 128 4
30 x (64 64) x 8
GINO
50 x 16 x
30 x (16 x 16) x
AROMA
50 16 x
30 x 64 x 8
CORAL
50 x 64
30 x 512
ENMA
26 16 x 4
16 x (8 8)
Oformer Oformer models are trained point wise i.e. no spatial neither temporal compression are processed. We used 4-depth layers with Galerkin attention type and 128 latent channels.
GINO GINO architecture is combination of 2 well-known neural operator architecture (GNO (Li et al., 2020b) and FNO (Li et al., 2020a)). The GNO uses the following architecture parameters: a GNO radius of 0.033 with linear transform. The latent grid size is 16 for 1d datasets and 16 x 16 for 2d. As a consequence, we adapt the FNO modes to for datasets and 8 x for 2d's.
AROMA AROMA makes use of a latent token of size 16 in 1d and 64 in 2d. Other architecture parameters are kept similar iee. 4 cross and latent heads with dimension 32. The hidden dimension of AROMA'S architecture is set to 128. The INR has 3 layers of width 64.
CORAL Finally, CORAL is meta-learning baseline that represents solution using an INR (SIREN (Sitzmann et al., 2020)), conditioned from a hyper-network. The hyper-network takes as input latent code, optimized with 3 gradient descent steps. The inner learning rate is set to 0.01. The hyper network has layer with 128 neurons. SIREN models have layers with 256 neurons.
As comparison between baselines, we show on table 13 the latent token sizes that are created with the hyper parameters detailed above.
E,2.3 Training details
Reconstruction performances For training the auto-encoder (as well as the baseline models, unless stated otherwise), we followed a unified training procedure. The loss function combines relative mean squared error term with a KL divergence term: L = Crecon + 3 LKL, where 3 = 0.0001.
To enhance robustness to varying levels of input sparsity, randomly subsample the spatial input grid during training. The number of input points ranges from 20% to 100% of the full grid X. The output grid used for loss computation remains fixed to the full grid. Grid subsampling is performed independently for each sample in a batch and refreshed at every iteration.
Models are optimized using the AdamW optimizer with an initial learning rate of 0.001, sched- uled via a cosine decay down to 1 over the course of training. Training batch sizes for each encoder-decoder architecture are listed in table 14.
ENMA: Trained with learning rate of 0.001 and batch size of 64 in ID. The autoencoder is trained for 1, 000 epochs in ID and 150 epochs in 2D.
AROMA: Training procedure follows ENMA. 2D datasets are trained for 1, 000 epochs.
To stabilize training. the KL loss term is removed. The model is trained standard autoencoder minimizing relative MSE, with an initial learning rate of 1e-6.
CORAL:
Oformer: Follows the same training procedure as ENMA. 2D datasets are trained for 1, 000 epochs.
Training setup is similar ENMA. 2D datasets trained for 100 epochs due to the inner loop over batches in the forward pass.
GINO:
47

---

Table 14: Batch size used for training autoencoders as presented in table
Model
Advection
Vorticity
Oformer
64
4
GINO
64
32
AROMA
128
8
CORAL
128
8
ENMA
64
8
Time-stepping The time-stepping task consists in training small FNO to unroll the dynamics in the latent space created by the different models. This allows us to assess the quality of the extracted features for dynamic modeling. Training proceeds as follows: Starting from trajectory, the entire trajectory. By concatenating this to the PDE parameters, we unroll the dynamic in the latent space, to build full sequence of tokens. We compute the loss (Relative MSE) in the latent space directly and at each step in the auto regressive process. All models are trained using initial learning rate of 0.0001 with and Adam W optimizer. The learning is also scheduled using a cosine scheduler. 250 epochs are performed for all baselines in 1d and in GINO and CORAL performs 50 training epochs, while Oformer and CORAL are trained for 150 epochs. This trainings can differs depending on the computational cost of the encoder.
encode then take the first tokens in the temporal dimension as input for the FNO.
The FNO used for processing is a simple 1d/2d FNO, with 3 layers width 64. FNO modes are set to and the activation function is GELU fonction.
E.3 Dynamics forecasting on complex systems: implementation details
This section outlines the implementation details related to the auto-regressive steppsg experi- ments with ENMA on the two complex physical systems: Kayleigh-Bénan Bénard and Active Matter (Ohana al., 2024).
On these two datasets, we disabled the interpolation component ENMA's encoder-decoder and used standard auto-encoder boost reconstruction performance, instead of using variational formulation. Autoregression is then performed on a fixed grid with a sequence length of size 3. 12 and 17 time-steps are predicted respectively for Rayleigh- Benigh and Active Matter. We report the hyper-parameteres used for the experiments in table 15.
Table 15: Model hyperparameter configuration for ENMA the Active Matter and Rayleigh-Bénan Rayleigh-Bénar
Hyperparameters
Active Matter
VAE embedding dimension
32
32
Number of tokens
256
256
Patch size
2
Spatial Transformer depth
6
6
Causal Transformer depth
6
Hidden size
512
512
MLP ratio
4
Num heads
8
8
QK normalization
True
True
Normalization type
RMS
RMS
Positional embedding
Sinus
Sinus
MLP depth
3
MLP width
512
512
Number of steps
16
16
FM steps
10
10
48

---

F
Additional experiments
We conduct a wide range of experiments to evaluate the capabilities of our model ENMA:
Generative capabilities (appendix F.1.I): evaluation uncertainty quantification, sample diversity, solver fidelity, and latent distribution alignment.
Compression study (appendix F.1.2): Dynamics forecasting evaluation with reduced latent compression on 2D datasets.
Inference speed comparison (appendix F.1.3): benchmarking ENMA against generative baselines.
Ablation studies (appendix F. F.1.4): investigating the impact of history length, autoregressive steps, flow matching iterations and the dynamics modeling objective.
OOD reconstruction (appendix F.2.2): testing under high sparsity with only 10% of the spatial input grid.
Super-resolution (appendix F2.3):: evaluation on finer output resolutions beyond the training grid.
Encoder-decodee variants (appendix F.2.4): ablations on positional bias and time-stepping architectures.
Geometry-Awand Attention Analysis (appendix F.2.5): qualitative inspection of the geometry- aware attention module
Token space visualization (appendix F.2.6): qualitative visualization of the encoded tokens
F1
ENMA process task
F1111
Generative ability of ENMA
We illustrate here additional benefits from the generative capabilities of ENMA through two example tasks: uncertainty quantification and new trajectory generation. For all uncertainty experiments, we focus on the Combined equation, as generating multiple samples can be costly for 2D data.
Mean prediction and confidence interval
Uncertainty quantification ENMA naturally enables uncertainty quantification by generat- ing multiple stochastic samples from the learned conditional distribution. This is achieved by sampling multiple candidates in the latent space via different noise realizations during the spatial decoding process. Specifically, for each autore- gressive step, ENMA injects Gaussian noise into the flow matching sampler, producing diverse plausible predictions for the spatial tokens. Re-
Mean, Time
Gt. Time: 19
Mean, Time:
Mcan 3x5td. Time:
peating this process yields ensemble of tra- jectories. In contrast, discrete AR models rely on sampling from categorical distributions (e.g.. top-k or nucleus sampling), which often leads
to overconfident outputs, limited diversity, and poorly calibrated uncertainty estimates.
An illustration is shown in fig. 20, where the red curve denotes the ground truth, the blue curve is the predicted mean at the final time step, and the shaded region indicates the empirical confidence interval defined by +:3 standard deviations.
Figure 20: Uncertainty quantification using ENMA. Multiple trajectories are sampled, and the final time step is used to compute the point- wise mean (blue), standard deviation (shaded), and ground truth (red).
To evaluate uncertainty quality, we report two
standard metrics:
49

---

CRPS (Continuous Ranked Probability Score) measures the accuracy and sharpness of probabilistic forecasts by comparing the predicted distribution the true outcome. Lower CRPS indicates better probabilistic calibration.
RMSCE (Root Mean Squared Calibration Error) quantifies calibration by comparing pre- dicted confidence intervals with empirical coverage. low RMSCE suggests well-calibrated uncertainty estimates.
CRPS and RMSCE Results ENMA sig- nificantly outperforms both AR-DiT and Ze- bra in terms of uncertainty calibration and probabilistic accuracy. as shown in table 16. These results indicate that ENMA provides both sharper and more reliable uncertainty estimates. The stochastic sampling in la- tent space enabled by flow matching allows ENMA model diverse plausible trajectories while maintaining strong calibration. In con- trast, discrete token-hased methods like Zebra and AR-DiT tend to produce overconfident or underdispersed forecasts. This demonstrates the advantage of per continuous latent-space
Method
Metric
Combined
RMSCE
2.68e-1
AR-DiT
CRPS
1.27e-2
RMSCE
2.19e-1
Zebra
CRPS
9.00e-3
RMSCE
8.68e-2
ENMA (ours)
CRPS
1.70e-3
Table 16: Comparison of uncertainty metrics (4 is better) for Combined.
modeling for uncertainty-aware forecasting.
Temporal Uncertainty Evaluation. To further assess the quality of uncertainty calibration time, we report the evolution of RMSCE and CRPS across autoregressive prediction steps in fig. 21. Ideally, both metrics should remain low and stable well-calibrated uncertainty estimates and accurate probabilistic forecasts throughout the trajectory. For CRPS, all models exhibit increasing scores over time due the compounding errors typical of neural PDE solvers. However, ENMA consistently achieves substantially lower CRPS values than both Zebra and AR-DiT, indicating sharper and more accurate trajectory forecasts. In the case of RMSCE, ENMA displays distinct trend: its calibration error decreases over time. This behavior reflects the model's evolving confidence-initee
time, indicating that the model maintains
predictions, informed by ground truth history, tend to be overly confident, resulting in narrower (and sometimes miscalibrated) confidence intervals. As the model progresses through the trajectory and relies more on its own predictions, becomes more conservative, yielding better-calibrated uncertainty estimates. ENMA maintains the lowest RMSCE every step, underscoring its robustness providing reliable confidence intervals throughout the rollout. Together, these results highlight ENMA's superiority in delivering both accurate and well-calibrated probabilistic forecasts over time.
INSCE exclut on per Oime stop
CAPS evolution per time stap
0.00
0.10%
0.79
0.39
0.03
0.53
0.01
0.00
(a) RMSCE evolution over time. Lower values indi- cate better uncertainty calibration.
(b) CRPS evolution over time. Lower values reflect sharper and more accurate forecasts.
Figure 21: Evolution of uncertainty metrics over time for the Combined dataset. (a) RMSCE and (b)
CRPS demonstrate that ENMA outperforms baselines in both calibration and predictive sharpness.
Data generation As second demonstration of ENMA's generative capabilities, we evaluate its ability to produce plausible and diverse trajectories without conditioning the initial state u².
50

---

Specifically, we consider the setting of generation conditioned on context example. Unlike classical neural approaches that typically require an initial condition and generate future states given y, this setup assumes no access to either u or 7. Instead, the model is provided with a context trajectory from which it must infer the underlying parameter "y and generate a coherent trajectory, including a plausible initial condition. This setting highlights ENMA's capacity to function as generative solver for parametric PDEs.
To benchmark generative performance, we compare ENMA against baseline models using standard metrics from the generative modeling literature. We introduce the Frécher Physics Distance (FPD), physics-adapted variant of Fréchet Inception Distance (FID) Heusel et al. (2017), along with Precision and Recall Kynkäänniemi et al. (2019). These metrics are computed in a compressed feature space extracted by lightweight CNN encoder trained to regress the underlying PDE parameters treated as instance classes.
The CNN processes the full trajectory and encodes it into a 64-dimensional feature vector, enabling semantically meaningful comparisons while mitigating the curse of dimensionality. In this space, FPD estimates the 2- 2-Wasserstein distance between real and generated distributions, while Precision and Recall respectively measure sample fidelity and diversity. Following standard protocol, use the training set as the reference distribution. Results for the Combined dataset are reported in table 17.
Table 17: Comparison of generative (FPD, Precision and Recall) metrics on the Combined dataset.
Model
FPD 4 Precision f Recall f
Zebra
1.03e-1
7.70c-1 7.92e-1
8.61e-1
ENMA
9.50e-3
7.80e-1
Results ENMA achieves the lowest FPD, significantly outperforming Zebra- indicating that ENMA'S generated trajectories are statistically closer to the real data distribution in the learned feature space. This reflects ENMA'S ability to generate high-fidelity samples that are well-aligned with physical ground truth. In terms of Precision, ENMA also outperforms Zebra, demonstrating superior sample quality. However, Zebra obtains higher Recall, suggesting that it covers a broader range of modes from the data distribution. This highlights a trade-off: ENMA prioritizes fidelity and realism, whereas Zebra exhibits slightly more diversity, potentially the cost of precision. Overall, physically plausible trajectories,
these results illustrate ENMA's strength in producing while maintaining reasonable diversity in its generative predictions.
Latent distribution alignment To further analyze the generative behavior of ENMA and Zebra, we visualize PCA projections of feature representations extracted by CNN trained on ground truth data. As shown in fig. 22, ENMA's samples (orange) closely align with the real data (blue), indicating that the generated trajectories remain well-covered within the training distribution. In contrast, Zebra exhibits large number of outliers that fall outside the support of the real data distribution. This mismatch suggests poorer calibration and lower fidelity to the training dynamics, which aligns with the lower FPD and precision scores reported in table 17. These findings confirm ENMA's ability to generate physically plausible and distributionally consistent trajectories.
Fidelity with respect to the numerical solver we take advantage of the known parametric structure of the PDEs. For each sample, we condition ENMA on a context trajectory and let it generate a full trajectory, including the initial state. Since the true PDE parameters y used to generate the context are known, we can pair ENMA's generated initial condition with the ground-truth y and run the numerical solver used during dataset creation. This produces a reference trajectory governed by the true physical dynamics. By comparing ENMA's generated trajectory with the solver-based rollout, we can assess how well ENMA captures the underlying PDE. A close match indicates that ENMA has learned to infer physically consistent initial conditions and dynamics from context.
To evaluate the fidelity of ENMA'S generations,
Results
As shown in fig. 23, the trajectories produced by ENMA closely match those from the solver, both qualitatively and quantitatively. We also report relative L2 error of 0.081 between the
51

---

PCA Features
10
65
Componeeet
Real Deta
Real D33
Data Con
PCA Compprentt1
ENMA
(b) Zebra
Figure 22: PCA projections of CNN features from generated (orange) and real (blue) trajectories at the final timestep.
trajectories generated by ENMA and by the solver, demonstrating that ENMA's generated states not only coherent but also physically meaningful.
Tapectry
Tapcory
Tayeday 2313
0.5
44
05
6.6
0.2
0.
SoMsio
5.4
3.2
-022
5.0
-04
-1.
Schvol Pajedcry #270
Schver Pajester #450
Sclvor The eopry 4313
3.0
5.8
-44
-32
0.
0.7
5.0
1.0
0.0
52
0.6
0.6
0.6
1.0
0.0
0.2
0.6
0.5
0.8
1.0
5.0
02
0.6
0.6
0.6
1.0
Figure 23: Comparison between trajectories generated by ENMA (top row) and the corresponding rollouts obtained from the PDE solver using the generated initial condition (bottom row). ENMA'S predictions yield physically consistent rollouts.
F.1.2 Dynamics forecasting with less compression
As discussed in section 4.3, ENMA exhibits lower performance on the Vorticity dataset compared to baseline models. We hypothesize that this is due to the aggressive latent compression, which limits the VAE's ability to capture fine-scale, high-frequency quality at decoding time. To investigate this, we evaluate ENMA under reduced compression settings. While the main experiments used a spatial compression factor 4 for Vorticity, we report additional results on 2D datasets in the temporal conditioning setting, comparing against AVIT and Zebra with relaxed compression factor.
cruutturs-uulimateee
constraining generation
Results Table shows results on Gray-Scott, Wave, and Vorticity under relaxed compression to match baseline settings. This setup addresses ENMA's lower performance on Vorticity observed in section 4.3, attributed to overly aggressive compression. ENMA outperforms AVIT and Zebra on Gray-Scott and Wave across both In-D and Out-D Dut-D settings. On Vorticity, ENMA remains competitive,
52

---

Table 18: Comparison of model performance on Gray-Scott, Wave, and Vorticity under the temporal conditioning setting. Metrics in Relative MSE (J).
Model
Gray-Scott
Wave
Vorticity
In-D
Out-D
In-D
Out-D
In-D
Out-D
AVIT
4.26e-2
1.68e-1
1.57e-1
5.88e-1
1.76e-1
3.77e-1
Zebra
4.21e-2
1.82e-1
1.40e-1
3.15e-1
4.43e-2
2.23e-1
ENMA
2.23e-2
1.52e-1
7.00e-2
2.83e-1
5.21e-2
2.82e-1
closely matching Zebra. These results highlight ENMA's strong predictive performance under comparable constraints.
E.1.3 Inference speed against generative approaches
One of the primary limitations of generative models compared to deterministic surrogates lies in their inference speed. Diffusion-basee methods, for instance, require multiple iterative passes through large model-typically a Transformer-to generate single sample, resulting in high computational cost. Similarly, autoregressive (AR) approaches like (Serrano et al., 2025) generate tokens sequentially, which becomes especially expensive when modeling high-dimensional spatio- temporal data.
ENMA addresses these inefficiencies by combining the strengths of both paradigms. It adopts an AR framework that supports multi-token generation and leverages key-value caching for faster inference, while also benefiting from the expressiveness of generative methods through its use of Flow Matching. Crucially, rather than relying on large model each step, ENMA performs flow matching through lightweight MLP, substantially reducing the computational overhead compared to traditional diffusion models.
Results We compare the inference time sam- ple for AR-DiT, Zebra, and ENMA in table 19, ENMA achieves the fastest inference across both the Combined and Vorticity datasets, with an aver- age of 2s and 4.5s per sample, respectively. This represents 2x speedup over AR-DiT and 3x speedup over Zebra on Combined. Zebra partic- ularly exhibits slower inference on the Vorticity dataset where it reaches 31s per sample. These results highlight the efficiency of ENMA's continu- ous latent autoregressive generation, which scales favorably compared to existing generative baselines.
Table 19: Inference speed comparison on the Combined and Vorticity datasets. Lower is bet- ter.
Model
Combined
Vorticity
AR-DiT
4s
6.7s
Zebra
6s
31s
ENMA
2s
4.5s
E.1.4
Ablation Studies
Relative 12 vs. Input History Size
To better understand the behavior and flexibility of our generative framework, we conduct ablations ex- amining ENMA's performance across three key di- mensions: (i) sensitivity to the length of the input history, (ii) the number of autoregressive spatial steps S, and (iii) the number of Flow Matching (FM) steps during generation.
0.02
0.000
.000
0.006
Impact of Input History Length In many real- world scenarios, the number of available historical observations varies, making it important for model adapt seamlessly to different history lengths. We evaluate ENMA's predictive performance tion of the number of past frames provided as context and evaluate the predictions quality on the last 10 steps of the trajectory. As shown in fig, 24, the relative L2 error consistently decreases as the input
685
Stoo
Figure 24: Relative L2 error vs. input history size. ENMA benefits from longer histories.
func-
53

---

history length increases. This trend confirms that ENMA effectively leverages additional temporal context to refine its predictions. Notably, the model already achieves competitive performance with as few as 2-3 input frames and continues to improve as more information is made available. This highlights the model's capacity for conditional generalization across varying input regimes.
Impact of the Number of Autoregressive Steps The number of autoregressive steps S di- rectly impacts inference efficiency: fewer steps accelerate generation, but may reduce accuracy. In the case of the Combined equation, each spa- tial state of 128 points is compressed into 16 la- tent tokens. We evaluate ENMA's performance as we vary S from 1 to 16 in fig. 25, where = corresponds to generating all tokens at once, and S = 16 to generating one token per step.
Relative L2 vs. Number of Autoregressive steps
EMPA
Reatuee
Number of autoregessive
We observe that performance is poor when S S 4. but significantly improves at S = 6, after which it quickly plateaus. This behavior reflects the masking ratio used during training: the model is trained to reconstruct frames with 75%-100% of tokens masked. Hence, inference scenarios where majority of tokens are already visible (i.e., large S' values) fall outside the model's training regime, which may limit gains from further increasing S. These results suggest that modest number of autoregressive steps-around 6 b-balaces speed and accuracy. Expanding the training schedule to include lower masking ratios could further improve performance for large S, but would increase training time due to the larger space of masking patterns.
Figure 25: Relative 1.2 error as function of the number of autoregressive steps S. Performance improves markedly around S 6 and plateaus thereafter.
Impact of the Number of Flow Matching Steps The number of flow matching (FM) steps governs the granularity of sampling from the learned conditional distribution each au- toregressive step. While more FM steps can, in principle, yield smoother and more accu- rate trajectories, they also increase inference cost-even if the MLP used for token decod- ing remains lightweight. In fig. 26, we evaluate ENMA's performance as a function of FM steps. We observe sharp drop in relative L2 error be- tween and 5 steps, after which the performance platcaus. Beyond 10 steps, improvements are marginal or even slightly inconsistent. This indi- cates that ENMA captures most of the necessary detail with very few sampling steps.
Relative Number of FM
0.0302
Relotiv
0.0093
0.0000
0.0005
000000
20
Figure 26: Relative L2 error versus number of flow matching steps per token. Performance quickly stabilizes after 5 steps.
From practical standpoint, this suggests that using as few as 5 FM steps can offer an optimal trade-off between speed and accuracy, making ENMA efficient at inference without compromising quality.
Latent Dynamics Modeling We assess the impact of the training objective on the quality of latent trajectory generation by comparing Flow Matching (ENMA) against a Diffusion-bassd latent model and a fully Deterministic variant. As shown in table 20, flow matching yields the lowest relative MSE, outperforming both alternatives. While diffusion-basse generative capacity, we observe convergence instability and degraded accuracy in practice. The deterministic approach performs better than diffusion but remains inferior to flow matching, likely due to its inability to capture stochastic variability in the latent dynamics. Overall, these results highlight the robustness and efficiency of the flow-matching paradigm for continuous stochastic trajectory modeling in ENMA.
training theoretically offers strong
54

---

Table 20: Ablation the training objective on the Vorticity dataset. Metrics are reported in Relative MSE on the test set.
Training Objective
Relative MSE
Flow Matching (ENMA) Diffusion Deterministic
0.0644
0.7579
0.0879
F.2 Experiments on the encoder-decoder
F.2.1 Additional Experiments on new datasets
We provide additional experiment on the encode-decoder of ENMA on a new dataset: Cylinder Flow (Pfaff et al., 2021). This dataset is inherently irregular and contains multiple output channels. The results are shown in table 21 and illustrate the superior performance of ENMA on this additional dataset.
Setting In this example, used tokens with spatial size of 8 x and feature dimension of for all baselines except OFormer, which does not perform spatial compression. ENMA uses a latent grid of size 32 32 with feature dimension of 16, and then applies two 'compress s_s_spe'''"""" and one 'compress_tine simme' layers to reach the desired latent size. All other configurations follow those described in appendix All models trained autoencoders by minimizing the RMSE loss. The FNO architecture follows the same configuration as described in appendix E.2.3. Models were trained for 1,000 epochs on the reconstruction task and 250 epochs for the ume-stepping experiment.
Table 21: Reconstruction error Test results and compression rates. Metrics in Relative MSE. The compression rate reflects how much the latent representation reduced compared the input data. A compression rate of x2 indicates that the latent space is half the size of the input in terms of raw elements.
Dataset
Cylinder Flow
1 Xue
Model 4
Reconstruction
Time-stepping
Compression rale
OFormer
1.71e-1
7.96e-1
x0.38
GINO
8.04e-1
1.56
x10
100%
AROMA
9.38e-2
6.50e-1
x10
CORAL
3.39c-1
5.35e-1
x10
ENMA
8.78e-2
1.71e-1
x20
OFormer
1.75e-1
7.96e-1
GINO
8.04e-1
1.55
50%
AROMA
1.02e-1
6.63e-1
CORAL
3.40c-1
5.40e-1
ENMA
9.06e-2
1.75e-1
OForme
1.92e-1
8.08c-1
GINO
8.03e-1
1.54
20%
AROMA
1.26e-1
7.22e-1
CORAL
3,43e-1
5.50e-1
ENMA
9.98e-2
1.82e-1
Results These additional results on widely known dataset further validate the conclusions drawn in section 4.2. ENMA outperforms the baselines on the reconstruction task while achieving twice the compression. Moreover, tokens capture informative features that enable strong performance on the time-stepping task.
55

---

F.2.2 OOD encoding on 10% of the input grid
We evaluate the performance of the auto-encoder in an duttof-distributiinn (OOD) reconstruction setting by reducing the input grid size to 10% of the original spatial grid. This requires models to reconstruct the full field from very sparse ohservations- only 12 points on the Advection dataset and 409 points the Vorticity dataset. Notably, such extreme sparsity was encountered during training, where input sampling ratios ranged from श 20% to 100%.
Table 22: Reconstruction error using π 10% of the initial grid Test results. Metrics in Relative MSE.
Dataset +
Advection
Vorticity
Cylinder Flow Reconstruction
the
Model +
Reconstruction
Reconstruction
OFormer
5.82e-1
1.00++0
2.19e-1
GINO
2.63e-1
6.46e-1
8.05e-1
AROMA
4.69c-1
3,28e-1
1.51e-1
10%
CORAL
1.13e+0
1.38c+0
3.58c-1 I
ENMA
2.45e-1
2.78e-1
1.15e-1
Results This reconstruction experiment demonstrates that ENMA maintains superior reconstruction quality even when provided with extremely sparse input fields. Among the basclines, AROMA and GINO also exhibit robust performance under such challenging conditions. contrast, OFormer and CORAL fail to reconstruct the physical field, highlighting their limited generalization in low-data regimes.
F.2.3 Super-resolution
In this section, we assess the models' ability generalize to finer spatial grids. For the Advection dataset, we evaluate reconstructions starting from input subsamplings of 20%, 50%, 100%, and query the model on denser grids with म. 200%, 400%. 800% of the original resolution. the Vorticity dataset, we use the same input subsamplings but perform super-resolution only at Tar 200% due to data availability constraints.
For comparison, we also report reconstruction results at गRr = 100%, as in table 1, and present the full results in table 23. The table structured follows: rows correspond to the input grid ratio n, matching the setup of table 1, while columns indicate the relative query grid size Zy. For instance, on the Advection dataset with original grid size of 128, the second row (л 50%) corresponds 64 input points. and the third column (T = 400%) corresponds to querying 512 points. Other rows and columns follow the same interpretation.
Results As shown table 23, all models exhibit performance degradation under the super-resolution setting. Despite this, ENMA consistently outperforms other encoder-decoder architectures resolutions. CORAL demonstrates stable performance as output resolution increases, while OFormer struggles significantly when queried on unseen grids for both datasets. ENMA and AROMA show similar trends as the super-resolution difficulty increases: however, AROMA's performance degrades more rapidly with lower input grid densities. Overall, ENMA's auto-encoder proves be the most robust, both to variations in input sparsity and to changes in query resolution.
all
56

---

Table 23: Reconstruction error on super resolution task. Test results. Metrics in Relative MSE.
Dataset
Advection
Vorticity
4 Aie
100%
200%
400%
800%
100%
200%
Fr Model +
OFormer
1.70e-1
5.20e+0
5.19e+0
5.19e+0
9.99e-1
1.00e+0
GINO
5.74e-2 7.77e-2
8.83e-2
9.43e-2
5.63e-1
5.76e-1
AROMA
5.41c-3 3.78c-2
5.62e-2
6.54e-2
1.45e-1
1.71e-1
π 100%
CORAL
1.34c-2
4.00c-2
5.76c-2
6.66e-2
4.50c-1
4.55e-1
ENMA
1.83e-3
3.71e-2
5.56e-2
6.49e-2
9.20e-2
1.36e-1
OFormer
1.79e-1
5.20c+0
5.19c+0
5.18e+0
9.99e-1
1.00c+0
GINO
6.64e-2
8.38e-2
9.37e-2
9,95e-2
5.69e-1
5,82e-1
AROMA
2.34e-2
4.44e-2
6.09e-2
6.94e-2
1.64e-1
1.89e-1
50%
CORAL
7.57e-2 8.80e-2
9.96e-2
1.06e-1
4.93e-1
4.98e-
ENMA
4.60e-3
3.74e-2
5.59e-2
6.51e-2
9,90e-2
1.41e-1
OFormer
2.50e-1
5,18c+0
5.17e+0
5.16e+0
9.99e-1
1,00c+0
GINO
9.13c-2
1.04e-1
1.12c-1
1.17e-1
5.90e-1
6.01e-1
AROMA
1.67e-1
1.72e-1
1.78e-1
1.81e-1
2.29e-1
2.45e-1
. 20%
CORAL
4.77e-1 4.79e-1
4.82e-1
4.84e-1
7.59e-1
7.62e-1
ENMA
3.05e-2
4.94e-2
6.49e-2
7.31e-2
1.37e-1
1.69e-1
F.2.4 Ablation studies
We conduct two additional ablation studies to evaluate specific architectural components of the incoerr-decode framework: (i) the effect of using a geometry-awant attention bias (table 24), (ii) the impact of the additional temporal causal compression, and (iii) the impact of the time-stepping process used in the decoder (table 26).
Geometry-avee awaare attention bias. Table 24 reports the reconstruction error when ENMA trained with and without the geometry-an encourages spatial locality in attention by penalizing interactions between distant query key pairs. The results show consistent improvements across all input sparsity levels, with relative gains ranging from 25% to 40%. We provide additional analysis and some visualizations appendix F.2.5.
attention bias deseribed in appendix D.I. This positional bias
Table 24: Reconstruction error with and without geometry-aware attention bias. Metrics are Relative MSE. Relative improvement is reported in parentheses.
Input Ratio
Model Variant
Reconstruction
w/o bias
3.09e-3
100%
Full
w/ bias
1.83e-3 (-40%)
w/o bias
6.84e-3
50%
Half
w/ bias
4.60e-3 (-33%)
w/o bias
4.13e-2
20%
Sparse
w/ bias
3.05e-2 (-26%)
Causal Non-causal encoder We compare causal convolution module in the encoder with non causal encoder, to assess its impact on temporal feature representation. Table 25 presents an ablation of the causal component in ENMA. Compared the causal autoencoder, the causal version adds temporal compression, reducing the number of tokens by half. Despite this, it matches outperforms the non causal cersion-especio in low-data settings indicating that the causal layer helps in providing more informative tokens.
57

---

Table 25: Reconstruction error: Ablation of the Causal component on the Advection dataset.
Metrics in Relative MSE on the test set (please refer to table for the sampling operation).
π
Input Ratio
Model Variant
Reconstruction
Causal
5.16e-3
100%
Full
Non causal
4.55e-3
Causal
6.98e-3
50%
Half
Non causal
6.02e-3aa
Causal
1.66e-2
20%
Sparse
Non causal
2.01e-2
Time-stepping architecture. We further examine the impact of the time-stepping module by replacing the FNO used in ENMA with 4-layer U-Net. Both models use identical lifting and projection layers to match the latent token dimensions. The U- Net adopts convolutional layers with kernel size stride 1, and padding 1.
Table 26: Reconstruction error for different time- stepping processes (FNO U-Net) on the Advection dataset. Metrics are Relative MSE.
Model
FNO
U-Net
OFormer
1.1lc+0
1.54c+0
GINO
7.89e-1 7.52e-1
100%
AROMA
2.23e-1 6.43e-1
Table 26 summarizes the reconstruction performance on the Advection dataset across three input sparsity levels. We ob- serve that ENMA maintains superior accu- racy under both architectures. While the U-Net variant exhibits slightly degraded performance compared to the FNO, re- mains competitive and stable. A dash (-) indicates that the model diverged and pro- duced NaNs during training.
CORAL
9.64e-
ENMA
1.64e-1
3.51e-1
OFormer
1.11e+0
1.04e+0
GINO
7.87e-1
7.50e-1
50%
AROMA
2.29c-1
6.49e-1
CORAL
9.74e-
ENMA
1.72c-1
3.58e-1
OFormer
1.13e+0
1.03e+0
GINO
7.96e-1
7.53e-1
20%
AROMA
3.21e-1
7.32e-1
These results confirm the robustness of ENMA across different architectural choices for the steppsg component and further emphasize the benefits of the geometry-aware attention mechanism.
CORAL
1.06e+0
ENMA
3.13e-1
4.15e-1
58

---

F.2.5 Georetry-Aware Attention Analysis
We visualize the attention mechanism in ENMA'S encoder-decoder using the Advection dataset to better understand the effect of the geometry-awand attention bias.
Attention Bias Figure 28 displays the geometry-aware bias introduced in Section D.I. Each plot corresponds one attention head, with rows representing different input sizes (128, 64, and 25 points from top bottom) and columns corresponding different heads. Bias values near zero (white) indicate minimal distance be- tween query and key coordinates, promoting attention, while larger distances produce strongly negative biases (dark blue), which sup- press it. Because the input grid is irregular, attention patterns in the second and third rows are not strictly diagonal. Additionally, since the intermediate coordinates not constrained to maintain a perfectly regular grid, which further contributes these deviations. Notably, the later attention heads appear to penalize distant tokens more strongly-an effect explained by the learned wead-specific scaling factors, following the approach of Press et al. (2022).
Learned Inner Grid Coords
1.0
0.0
learned (see fig. 27), the model is
-0.2
Figure 27: Learned coordi- nates of the intermediate regu- lar grid.
Posiiional Attention Blas Head
Mertoo Ora
Posttional Attention Blas
Apprr Bra
Positional per
Figure 28: Geometry-aware attention bias B for each attention head, shown for 128 (top), 64 (middle), and 25 (bottom) input points.
Attention Scores Figures 29 and 30 present the actual attention scores produced using the geometry- aware bias. These visualizations confirm that attention is concentrated nearby points, enforcing spatially local inductive bias during interpolation. We provide visualization for the advection dataset fig. 29 and vorticity fig. 30.
These results demonstrate that the attention mechanism learns to prioritize spatially proximate inputs, under varying levels of spatial sparsity.
59

---

Attention Head
Hi22
Attontion score per Head
Atlention Head
Figure 29: Attention scores per head using the geometry-aware bias for 128 (top), 64 (middle), and 25 (bottom) input points.
Attention visualization in the physical space We visualize the final attention weights in physical space in fig. 30. For a selected timestep of a given trajectory, we choose a set of spatial tokens (left panel) and display their corresponding attention scores across the physical domain (middle panel). The resulting patterns clearly demonstrate that the cross-attention mechanism, guided by the geometry- aware bias, preserves locality by assigning high attention weights nearby ciations-ffeectvelle focusing on spatially relevant information.
Selected takens
Local attention scores
Physical field
0.5
Attention Score
0.1
Figure 30: Attention scores the vorticity dataset, spatially selected tokens (left) given timestep (f 15). The resulting attention score on the output grid keeps the local behavior (middle). Attention averaged across heads for visualization. The left figure is the physical field considered.
F.2.6 Visualization in the Token Space
Figure 31 shows the latent token representations on the Advection dataset for varying input sparsity (128, 64, and 25 points from top to bottom). The final column displays the corresponding ground-truth physical trajectory.
60

---

Taken dimension 214
Taken dimension 214
Takon dimension 3/4
Traan dimension 414
Plysica Tajectory
0.6
0.4
Taken on dimersion 1/4
Taken dimonsion 2/4
Talen dimonslon 3/4
Toxinn on dimcnsion 414
Physkcal Taaectory
0.
Tajectory
0.4
-02
Figure 31: Token representation when using 128 (top), 64 (middle), and 25 (bottom) input points.
The tokens effectively capture the dynamics in a compressed form, demonstrating the encoder's strong inductive bias. Notably, the first latent dimension remains stable across input resolutions, encoding coarse global structure. In contrast, the remaining dimensions vary with input density, indicating progressive refinement of local details sign of partial disentanglement across token dimensions. This behavior supports the model's ability to encode structure hierarchically and adaptively under compression.
Intermediate Fields in the Encoder-Decoder Figure 32 visualizes the full trajectory of latent transformation within ENMA's encoder-decoder. Each row corresponds to a processing stage, averaged across channels.
Starting from sparse, irregular input field (top row, ग 20%), the first cross-attention layer projects the observations onto a learned latent grid (row 2). effectively interpolating the missing spatial points. The causal CNN then encodes the trajectory into a compact token space (row 3), capturing the trajectory's 's temporal evolution in a lower-dimensional representation. These tokens are subsequently upsampled (row 4) and decoded onto the full spatial grid (row 5). The final reconstruction is smooth and well-aligned with the underlying dynamics, highlighting the effectiveness of ENMA's encoder decoder architecture. Minor artifacts observed in intermediate stages likgelylikey caused by learned positional embeddings (Yang et al., 2024) are effectively corrected during decoding.
61

---

Input Field
0.5
20
40
20
40
60
80
100
120
Mean interpolated Field
0.0
40
-0.5
0
20
40
60
80
100
120
Encoded Token
0.5
20
0.0
0
2
10
12
14
Mean decoded interpolated Field
20
40
60
80
100
120
Output Field
0.5
20-
0.0
-055
40
60
100
120
Figure 32: Intermediate fields in the encoder/decoder blocks: interpolation (rows 1-2), compression
(rows 2-3), upsampling (rows 3-4), and reconstruction (rows 4-5).
62

---

G Visualization
G.1
Combined Equation
Figure 33 and Figure 34 provide qualitative results on the Combined equation, showcasing ENMA's accuracy on in-distribution imputs and its generalization capability out-of-distribution scenarios.
ENMA Ground Truth for 1D PDEs
Sample 1: a 0.8434, & 0.131, y 0.53, 6 0.0
Sample 2: a 0.7475, 0.3855, y 0.59, 6 0.0
Ground Truth
0.0
-05
-0.
-.2
ENMA
0.0
-0.
RMSE Error
0.005
0.006
0.004
0.004
00000
0.002
0.002
0.001
0,000
0.000
Figure 33: Qualitative comparison between ENMA prediction and ground truth for in-distribution examples from the Combined.
ENMA Ground Truth for 1D PDEs
Sample 1: 0.3788, & = 0.498, y = 0.38, 6 0.5 Sample 2: a 0.3101, B = 0.4045, y 0.4, 5 0.8889
Ground Truth
02
-02
-03
X
0.2
ENMA
0.0
RMSE Error
0.05
0.04
0.00
0.00
Figure 34: Qualitative comparison between ENMA prediction and ground truth for an OOD example from the Combined,
63

---

G.2
Advection Equation
Figure 35 and Figure 36 provide qualitative results the Advection equation. These plots highlight ENMA's accurate forecasting on in-distribution samples and its ability to generalize to out-of- distribution scenarios.
ENMA vs Ground Truth for 1D PDEs
Sample -1.001
Sample -0.409
Ground Truth
0.50
0.00
0.25
-0
-0.30
0.75
0.25
ENMA
-0.25
0.25
-0.75
.0.50
RMSE Error
0.015
00.00
0.010
0.50
0.00
Figure 35: Qualitative comparison between ENMA prediction and ground truth for in-distribution examples from the Advection dataset.
ENMA vs Ground Truth for 1D PDEs
Sample 6.7555
Sample -5.348
Ground Truth
0.50
0.0
0.00
-0.2
-B4
-0.50
0.4
ENMA
-03
0.25
-0.25
-0.30
RMSE Eiror
0.06
02
X
Figure 36: Qualitative comparison between ENMA prediction and ground truth for an OOD example from the Advection dataset.
64

---

G.3 Gray-Scott Equation
Figure 37, Figure 38, and Figure 39 present qualitative comparisons between ENMA and ground truth on the Gray-Scott equation. ENMA accurately reconstructs complex spatiotemporal patterns on in-distribution samples and demonstrates good qualitative performance out-of-distribution parameters.
ENMA vs Ground Truth Parameters: F 0.0332, k 0.0607
t-
t-19 19
Ground Truth
0.75
0.5
0.00
ENMA Prediction
0.0
MSE
0.
0.0
Figure 37: Qualitative comparison between ENMA prediction and ground truth for an in-distribution
sample from the Gray-Scott dataset (F 0.0323, k 0.0606).
ENMA Ground Truth Parameters: F= 0.0332, k 0.0604
Ground Truth
8
8
O
0.50
0.25
ENMA Prediction
1.0
P
8
0
MSE Error
0.15
0.10
0.05
0.00
Figure 38: Another in-distribution Gray-Scott example showing the agreement between ENMA
prediction and ground truth (F' 0.0316, & 0.0597).
65

---

ENMA vs Ground Truth Parameters: F 0.0467, 0.058
t-0 O
t-5
t- 10
t- 15
t- 19
Ground Truth
0.75
0.50
0.00
ENMA Prediction
0.0
MSE
Figure 39: Out-of-distribution (OOD) generalization for the Gray-Scott equation. ENMA prediction
remains consistent despite being evaluated at unseen parameters (F 0.0467, & = 0.058).
G.4 Wave Equation
Figure 40, Figure 41, and Figure 42 present qualitative comparisons for the Wave equation. ENMA accurately captures wavefront propagation in in-distribution scenarios and generalizes well to out-of- distribution conditions.
ENMA vs Ground Truth Parameters: c = 140.404, 23.2323
t
t=7
14
= 21
=
Ground Truth
1.5
1.0
0.5
0.0
ENMA Prodiction
1.5
1.0
0.5
0.0
MSE Error
0.1
Figure 40: Qualitative comparison between ENMA prediction and ground truth for an in-distribution
sample from Wave dataset (c 140.404, 23.2323).
66

---

ENMA Ground Truth Parameters: 144.4444, 21.2121
Ground Truth
0.5
0.0
ENMA Prediction
0.5
MSE Error
0.02
Figure 41: Second in-distribution sample from the Wave dataset (c 144.4444, k = 21.2121),
showing ENMA's prediction, ground truth, and the corresponding MSE error.
ENMA Ground Truth Parameters: 500.0, 52.5
t- 21
Ground Truth
0.5
ENMA Prediction
1.0
0.5
0.15
MSE
0.05
Figure 42: Out-of-distribution example from the Wave dataset (c 500.0, k = 52.5), demonstrating
ENMA's generalization capabilities beyond the training regime.
67

---

G.5 Vorticity Equation
Figure 43, Figure 44, and Figure 45 show qualitative comparisons for the Vorticity equation. ENMA reliably captures fluid dynamics behavior on in-distribution samples and maintains accuracy in out-of-distribution regimes.
ENMA Ground Truth Parameters: 0.0019
t=0
t-7
14
21
Ground Truth
0
-10
ENMA Prediction
-10
MSE Error
Figure 43: Qualitative comparison between ENMA prediction and ground truth for an in-distribution
sample from the Vorticity dataset (v 0.0019).
ENMA vs Ground Truth Parameters: == 0.0048
t=
t-7
t=
28
Ground Truth
10
ENMA Prediction
10
1.0
MSE
0.5
0.0
Figure 44: Second in-distribution sample from the Vorticity dataset (v 0.0048), illustrating ENMA's ability to reproduce complex vortex structures.
68

---

ENMA Ground Truth Parameters: 0.0007
14
Ground Truth
ENMA Prediction
10
MSE
Figure 45: Out-of-distribatis example from the Vorticity dataset (u 0.0007), highlighting ENMA's robustness in extrapolating vortex dynamics.
69


Reviews:
[{"review_id": "WXbphkUJPa", "reviewer": "Reviewer_HhGJ", "summary": "A method for dynamics propagation in latent space is introduced. It separates the temporal forecast, achieved by a temporal transformer, from the spatial forecast, achieved by a masked autoregressive (MAR) model, by concatenating the two operations and making the MAR input conditional on the output of the temporal transformer. After the MAR, a flow-matching model aligns the predicted spatial latent with the ground truth. The flow-matching model consists in an ODE \\dot x = v(x) which uses an MLP as velocity field.\n\nThese ingredients are brought together with an encoder that is capable of creating latent vectors on irregularly spaced time grids.\n\nExtensive experiments on simulated data and comparison to many other methods show that the proposed method performs favorably compared to existing literature.", "questions": "Would it be possible to have further experiments on public benchmarks?\n\nWould it be possible to see the benefits and necessity of the individual components of the method by ablating them?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "The paper is easy to read and the diagrams are self-explanatory. The method is presented clearly and set in context.\nThe proposed method is a step forward towards accurate dynamics propagation on potentially lower-resolution latents without drifting off the trajectory.\n\nExperiments are rigorous and extensive. However, it would be very useful to have numbers on some public benchmarks. Are there sub-datasets in PDEbench, PDEarena, or The Well, which would lend themselves to benchmark the proposed method and compare to others?", "was_revised": true, "final_justification": "This work is solid and I am happy to maintain my score of 5 which I think best reflects the character of this work. If asked, I will advocate for acceptance.", "rebuttal": "We sincerely thank the reviewer for the constructive and encouraging feedback. We appreciate the reviewer’s positive feedback on the clarity of the method and the accompanying diagrams. Below, we address the main points raised.\n\n# New experiment: 2 new Public Benchmarks (W1 & Q1)\n\nWe fully agree that evaluating on public datasets enhances the reproducibility and comparability of our results. While our primary goal was to assess ENMA’s capacity to generalize across *parametric time-dependent PDEs*—where PDE parameters may vary significantly from one trajectory to another—most existing public benchmarks exhibit only limited parametric variability. Thus, they do not fully reflect the challenges addressed in our setting. For each physical system, we generated 1200 distinct PDE instances, each with a different set of parameters, and 10 trajectories for each instance with a distinct initial condition, leading to 12000 samples. In comparison, datasets from public benchmarks offer a different setting: each PDE system in these datasets only includes a few number of parametric instances (about 10 vs 1200 for our setting), offering a limited diversity for each PDE. Thus current public benchmark, focus on solving a variety of PDEs, but do not allow to fully evaluate the parametric setting which is our target in the paper.\n\nNonetheless, ENMA is designed to be general and flexible, and it can readily be applied to a wide range of time-dependent PDE problems. To this end, we have included results on two public datasets that capture complementary aspects of ENMA’s capabilities:\n\n- **CylinderFlow** [1]: A benchmark with natively irregular spatial grids, used to assess ENMA’s encoding/decoding capabilities for irregular meshes in comparison to existing neural operator methods (corresponding to Table 1 in the paper). We report reconstruction performance, time-stepping (used with a FNO for all benchmarks for fair comparison) and the compression ratio used. This experiment clearly highlights the capabilities of ENMA to manage irregular grids.  The time-stepping task is much better with ENMA's tokens, likely due to the temporal compression proposed ($\\times 22$ compression rate for ENMA, whereas baselines have $\\times 11$ compression rate).\n\n    Table R1: Results of the reconstruction task on the cylinder flow dataset. Metrics in Relative MSE on the test set. We highlight in **bold** the best performing model.\n    |Sampling|Baseline|Reconstruction|Time-stepping|Compression\n    |-|-|:-:|:-:|:-:| \n    || GINO| 6.85e-1 | 1.59e0|$\\times 11$\n    |100\\%| AROMA| 9.59e-3 | 6.59e-1 | $\\times 11$\n    || CORAL | 5.06e-2 | 7.84e-1 | $\\times 11$\n    || ENMA | **7.78e-3** | **1.44e-1** | **$\\times 22$**\n    |||||\n    || GINO| 6.85e-1 | 1.59e0 |\n    |50\\%| AROMA| 9.86e-3 | 6.78e-1 |\n    || CORAL | 4.96e-2 | 7.83e-1|\n    || ENMA | **7.80e-3** | **1.44e-1**|\n    |||||\n    || GINO| 6.85e-1 | 1.59e0 |\n    |20\\%| AROMA| 1.06e-2 | 6.82e-1 |\n    || CORAL | 5.14e-2 | 7.83e-1 |\n    || ENMA | **7.87e-3** | **1.44e-1** |\n\n- **Active Matter (The Well)** [2]: A complex physical system involving collective behavior, used to evaluate ENMA’s performance against autoregressive transformer baselines for temporal generation (corresponding to Table 2 of the manuscript). In Table R2, we observe that ENMA is close to AVIT (also known as MPP) while operating in a very compressed space. Relaxing this compression ratio would yield to better performance, as we observed in Appendix F.1.2, but we believe that operating latent space for heavy data such as Active Matter is essential.\n\n   Table R2: Results of the temporal generation task on the Active Matter dataset. Metrics in Relative MSE on the test set.\n    |Baseline|Temporal conditioning|Compression ratio|\n    |-|:-:|:-:|\n    |AVIT|0.383| $\\times 1$|\n    |ENMA|0.412|$\\times 10$|\n\n# Ablation Studies: Benefits of ENMA's components (Q2)\n\nWe appreciate the suggestion to further clarify the contribution of ENMA’s architectural components. Below we summarize both existing and new ablation studies. We will add an ablation section in the main text of the final manuscript with references to the different experiments :\n\n### Existing ablations\n- **Compression ratio** (Sec. F.1.2): We show that reducing the compression improves performance on Vorticity, making ENMA competitive. On Wave and Gray-Scott, ENMA was on par with existing baselines at a higher compression rate and now outperforms them. \n- **Sampling schedule** (Sec. F.1.4): We evaluate ENMA’s predictive accuracy as a function of the number of autoregressive and denoising steps, showing that accurate forecasts can be obtained with few steps.\n- **History length** (Sec. F.1.4): We observe consistent performance gains with longer input histories, in line with ENMA’s autoregressive formulation.\n- **Geometry-aware encoding** (Sec. F.2.3): We compare reconstruction performance with and without the geometry-aware component. Across various input sparsity settings (full, half, sparse), the geometry-aware encoder consistently improves reconstruction accuracy.\n\n### New ablations\nWe have evaluated the importance of: (i)  causal encoder, (ii) generative vs deterministic component, (iii) scheduler.\n \n- **Causal CNN**: We compare a causal convolution module in the encoder with a non causal encoder, to assess its impact on temporal feature representation.\n    Table R3 presents an ablation of the causal component in ENMA. Compared to the non causal autoencoder, the causal version adds temporal compression, reducing the number of tokens by half. Despite this, it matches or outperforms the non causal version—especially in low-data settings— indicating that the causal layer helps in providing more informative tokens. \n\n    Table R3: Ablation of the Causal component on the Advection dataset. Metrics in Relative MSE on the test set (please refer to table 1 for the sampling operation). \n    |Sampling|Relative MSE||\n    |-|:-:|:-:| \n    ||Causal|Non causal|\n    |100\\% | 5.16e-3 | 4.55e-3 |\n    ||||\n    |50\\%  | 6.98e-3 | 6.02e-3 |\n    ||||\n    |20\\%  | 1.66e-2 | 2.01e-2 |\n\n- **Latent dynamics modeling**: We compare the flow-matching-based generation to both a diffusion-based latent model and a fully deterministic alternative, to quantify the benefit of continuous stochastic trajectory modeling. In table R4, we observe that a Flow matching framework works better than a deterministic approach, while offering the advantages of generative modeling. We also tested a diffusion training objective, but it didn't converge correctly in our experiments leading to lower performance. Overall flow-matching is easier to train. \n\n    Table R4: Ablation of the generation mode on the Vorticity dataset. Metrics in Relative MSE on the test set. \n    |Training Objective |Relative MSE|\n    |-|:-:| \n    |Flow Matching (ENMA) |0.0644|\n    |Diffusion|0.7579|\n    |Deterministic |0.0879|\n\n- **Scheduler**: We evaluated different scheduling strategies to select the number of masks to generate at a given step. In table R5, we observe that the choice of the scheduler is not of particular importance; all strategies lead to very similar performance.\n\n    Table R5: Ablation of the Scheduler on the Vorticity dataset. Metrics in Relative MSE on the test set. \n    |Scheduler|Relative MSE|\n    |-|:-:| \n    |Linear |0.0619|\n    |Cosine  |**0.0612**|\n    |Exponential  |0.0615|\n    |Reverse exponential |0.0618|\n\nThese new experiments will be included in the final version to further validate ENMA’s design choices.\n\nWe hope this addresses the reviewer’s suggestions and helps clarify the motivation and robustness of the proposed method. We thank again the reviewer for their thoughtful and detailed assessment.\n\n### References\n[1] Battaglia et al, Learning mesh-based simulation with graph networks, International Conference on Learning Representations, 2021. \n\n[2] Ohana et al., The Well: a Large-Scale Collection of Diverse Physics Simulations for Machine Learning, NeurIPS 2024 Track Datasets and Benchmarks Poster, 2024."}, {"review_id": "Ekhs4V4bNl", "reviewer": "Reviewer_ydbq", "summary": "The authors propose a new generative neural operator that is able to both operate on a token level and in a continuous space. The framework is made of the following components: 1) an encoder which projects irregularly sampled points into a fixed size continuous embedding using geometry aware attention, 2) a generative model which uses a causal transformer to predict future latent states temporally, then a spatial transformer with flow matching to performtokenwise spatial generation, and 3)a decoder to map the generated tokens back to the data space. ENMA shows clear empirical improvements over existing baselines.", "questions": "- **justification needed**: can the authors provide additional justification for choosing flow-matching over other generative frameworks? an analysis of the pros and cons of a few alternatives would suffice.\n- **Sharing code**: please share the source code to replicate your results anonymously. \n- **Missing some ablations**: an interesting ablation would be to vary the compression rate and see if less compression would indeed help complex system like vorticity. \n- **More complex setups**: it would be interesting to see how ENMA fares in modelling higher dimensional systems \n- **Small changes**:\n    - can you explain in a few words what relative mse is in the experimental section and why you use it?\n    - can you elaborate more on the FNO architecture used in the main paper?\n- **please clarify**: I understand that the compression rate refers to the size of the latent space over the size of the data space. I don't understand why this is included as a comparison criteria in Table 1. Isn't it a parameter the authors can control in other baselines as well?\n-Nitpicks\n    - Title of 3.1: tokenwise should be capitalized\n    - Table 1 caption: random bolding, was a bit misealding. Either bold all the metrics computed or none.", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n- The paper is well-written, the figures convey clear ideas and help support the ideas of the text\n- The model is tested in various setups (reconstruction for the encoder-decoder, temporal conditioning and ivp with context) with clear improvements in performance.\n- Extensive ablations and additional details are provided in the appendix.\n\n## Weaknesses\nOverall, I am willing to bump up my score if the authors 1) improve the justification of some components in their model, and 2) provide additional empirical evidence testing their model in more complex PDE setups. More detailed in Questions.", "was_revised": true, "final_justification": "The authors have addressed all my concerns.", "rebuttal": "We thank the reviewer for the thoughtful and constructive feedback. Below, we address each of the points raised.\n# Ablations: Components of ENMA (W1, Q1 and Q3)\nWe provide justifications and ablations concerning ENMA's components.\nWe first highlight that several ablations already included in the appendix directly address some of the reviewer’s suggestions (they will be better referenced in the main text):\n### Existing Ablations\n- **Compression ratio**:(Appendix F.1.2): We provide ablations across multiple 2D datasets (Wave, Gray-Scott, Vorticity). Reducing compression substantially improves ENMA’s performance. On Wave and Gray-Scott, ENMA surpasses all baselines, and on Vorticity, it becomes competitive.\n- **Geometry-aware encoding** (Appendix F.2.3): We compare reconstruction performance with and without the geometry-aware component. Across various input sparsity settings (full, half, sparse), the geometry-aware encoder consistently improves reconstruction accuracy.\n### New Ablations\nWe have evaluated the importance of: (i)  causal encoder, (ii) flow-matching vs deterministic component, (iii) scheduler. \n- **Causal CNN**: We ablate the use of the causal convolution module in the encoder to assess its impact on temporal feature representation.\n    Table R1 presents an ablation of the causal component in ENMA. Compared to the non causal autoencoder, the causal version uses causal layers and adds temporal compression, reducing the number of tokens by half. Despite this, it matches or outperforms the non causal version—especially in low-data settings— indicating that the causal layer helps in providing more informative tokens.\n\n    Table R1: Ablation of the Causal component on the Advection dataset. Metrics in Relative MSE on the test set (please refer to table 1 for the sampling strategy). \n    |Sampling|Relative MSE||\n    |-|:-:|:-:| \n    ||Causal|Non causal|\n    |100\\% | 5.16e-3 | 4.55e-3 |\n    ||||\n    |50\\%  | 6.98e-3 | 6.02e-3 |\n    |||||\n    |20\\%  | 1.66e-2 | 2.01e-2 |\n\n- **Latent dynamics modeling (Q1)**: We chose flow matching due to its well-established advantages: more stable training than diffusion models and fewer sampling steps during inference. Appendix F.1.4 shows that ENMA achieves strong results with as few as 5 steps. To further support this choice, we have added a new experiment comparing ENMA with a diffusion-based model and a fully deterministic model. Our results confirm that flow matching leads to more stable training and better accuracy, while maintaining the uncertainty estimation benefits of generative modeling. Diffusion didn't converge correctly in our experiments leading to lower performance. \n\n    Table R2: Ablation of the generation mode on the Vorticity dataset.\n    |Training Objective |Relative MSE|\n    |-|:-:| \n    |Flow Matching (ENMA) |0.0644|\n    |Diffusion|0.7579|\n    |Deterministic |0.0879|\n\n- **Scheduler**:  We evaluated different scheduling strategies for the number of masks to generate at a given step. In table R3, we observe that all strategies lead to similar performance.\n\n    Table R3: Ablation of the Scheduler on the Vorticity dataset.\n    |Scheduler|Relative MSE|\n    |-|:-:| \n    |Linear |0.0619|\n    |Cosine  |**0.0612**|\n    |Exponential  |0.0615|\n    |Reverse exponential |0.0618|\n\nThese new experiments will be included in the final version to further validate ENMA’s design choices.\n# New benchmarks on complex setups (W2 and Q4)\nTo evaluate ENMA on more challenging PDE scenarios, we have added results on two publicly available benchmarks:\n- **CylinderFlow** [1]: A time-dependent PDE on a natively irregular spatial grid. We report encoder-decoder performance and compare to neural operator baselines designed for irregular grids. Table R4 shows the results of the auto encoding and time-stepping tasks corresponding to table 1 in the manuscript. This experiment clearly highlights the capabilities of ENMA to manage irregular grids.  The time-stepping task is much better with ENMA's tokens, likely due to our temporal compression ($\\times 22$ compression rate for ENMA, when baselines have $\\times 11$ compression rate).\n    \n    Table R4: Results of the reconstruction task on the cylinder flow dataset. Metrics in Relative MSE on the test set. In bold, the best performing model.\n    |Sampling|Baseline|Reconstruction|Time-stepping|Compression\n    |-|-|:-:|:-:|:-:| \n    || GINO| 6.85e-1 | 1.59e0|$\\times 11$\n    |100\\%| AROMA| 9.59e-1 | 6.59e-1 | $\\times 11$\n    || CORAL | 5.06e-2 | 7.84e-1 | $\\times 11$\n    || ENMA | **7.78e-3** | **1.44e-1** | **$\\times 22$**\n    |||||\n    || GINO| 6.85e-1 | 1.59e0 |\n    |50\\%| AROMA| 9.86e-3 | 6.78e-1 |\n    || CORAL | 4.96e-2 | 7.83e-1|\n    || ENMA | **7.80e-3** | **1.44e-1**|\n    |||||\n    || GINO| 6.85e-1 | 1.59e0 |\n    |20\\%| AROMA| 1.06e-2 | 6.82e-1 |\n    || CORAL | 5.14e-2 | 7.83e-1 |\n    || ENMA | **7.87e-3** | **1.44e-1** |\n\n- **Active Matter** (The Well [2]) : A complex PDE defined on a $256 \\times 256$ grid with 11 input fields. We evaluate ENMA’s temporal generation ability against autoregressive transformer baselines (Table 2 of the manuscript, temporal conditioning setup). In Table R5, we observe that ENMA is close to AVIT (also known as MPP) while operating in a very compressed space. Relaxing this compression ratio would yield to better performance, as observed in F.1.2, but we believe that operating in such latent space for high dimensional data such as Active Matter is essential.\n\n    Table R5: Results of the temporal generation task on the Active Matter dataset. Metrics in Relative MSE on the test set.\n    |Baseline|Temporal conditioning|Compression ratio|\n    |-|:-:|:-:|\n    |AVIT|0.383| $\\times 1$|\n    |ENMA|0.412|$\\times 10$|\n\n# Sharing code (Q2)\nAccording to NeurIPS policy one cannot add links in the rebuttal, hence we cannot provide the code now. However, we will publicly release both the codebase and datasets upon notification of acceptance.\n# Small changes (Q5): \n- **Relative MSE**: This is a standard metric in PDE learning. It normalizes error by the energy of the ground truth, which is particularly important for comparing PDEs with different magnitude scales. It is defined as: \n\n$$\\text{Relative MSE}=\\frac{1}{n} \\sum_{i=1}^n \\frac{(y_i - \\hat{y}_i)^2}{||y||_2^2}$$\n\nwith $ \\lVert y \\rVert_2^2 = \\sum_{j=1}^{n} y_j^2 $\n- **FNO architecture**: In the main text, table 1 evaluates the influence of ENMA’s encoder-decoder choice  in the reconstruction task. We compare several encoder-decoder components baselines using FNO as the time stepper. This allows us to isolate the role of the encoder-decoder model. Note we also performed tests with a U-Net in Appendix F.2.3, with similar conclusions.  All the details on the architectures are provided in Appendix E.2.3.\nIn table 2 of the manuscript, we benchmarked FNO against other time steppers, architecture details are in Appendix E.1.2.\n\n# Clarifications (Q6): \n- **Compression rate in table 1**: While most frameworks can control the compression rate, they proceed differently.\n OFormer operates directly in the physical space (pointwise, no compression). GINO, AROMA, and CORAL compress only in space. In contrast, ENMA compresses both space and time.  We tuned latent channel dimensions for each method to achieve optimal trade-offs between compression and accuracy. This is indicated by the different compression ratios shown in table 1.  \nFor a better comparison, we provide (Table R6, R7) new experimental results, where the spatial compression is the same for all methods and ENMA performs an additional compression in the time dimension, hence a double compression ratio compared to other methods. This will be discussed in the final version.\n\n    Table R6: Additional results on the Advection dataset with unified spatial compression rate. Metrics in Relative MSE. Models ending with -old indicates results proposed in table 1 of the manuscript, while models ending with -new are the updated results. \n    |Sampling|Baseline|Advection||| \n    |-|-|:-:|:-:|:-:| \n    |Sampling|Baseline|Reconstruction|Time-stepping|Compression|\n    || GINO-new| 3.15e-1 | 8.55e-1 | $\\times 2$| \n    |100\\%| GINO-old| 5.74e-2 | 7.89e-1 | $\\times 0.5$|\n    |||||\n    || GINO-new| 3.21e-1 | 8.64e-1 | | \n    |50\\%| GINO-old| 6.64e-2e-2 | 7.87-1 | | \n    |||||\n    || GINO-new| 3.54e-1 | 9.11e-1 | |\n    |20\\%| GINO-old| 9.13e-2 | 7.96e-1 | |\n    \n    Table R7: Additional results on the Vorticity dataset with unified spatial compression rate. Metrics in Relative MSE. Models ending with -old indicates results proposed in table 1 of the manuscript, while models ending with -new are the updated results. \n    |Sampling|Baseline|Vorticity||| \n    |-|-|:-:|:-:|:-:|\n    |Sampling|Baseline|Reconstruction|Time-stepping|Compression|\n    |100\\%| CORAL-new | 4.63e-1 | 9.18e-1 | $\\times 8$ |\n    || CORAL-old | 4.50e-1 | 9.85e-1 | $\\times 2$|\n    |||||\n    |50\\%| CORAL-new | 4.95e-1| 9.22e-1 |\n    || CORAL-old | 4.93e-1 | 9.85e-1 |\n    |||||\n    |20\\%| CORAL-new | 6.89e-1 | 9.37e-1 |\n    || CORAL-old | 7.59e-1 | 9.87e-1 |\n\n- **Formatting**: \nWe thank the reviewer for catching the issues, such as the capitalization of “Tokenwise” in Section 3.1. We will correct these in the camera-ready version.\n- **Bold formatting in table 1**: \nWe bolded the best score and underlined the second-best. We agree this should be clarified in the caption to avoid confusion.\n\nWe hope this response resolves the reviewer’s concerns. We appreciate the detailed and constructive\nfeedback and would be happy to further engage in discussion\n\n### References\n[1] Battaglia et al, Learning mesh-based simulation with graph networks, International Conference on Learning Representations, 2021. \n\n[2] Ohana et al., The Well: a Large-Scale Collection of Diverse Physics Simulations for Machine Learning, NeurIPS 2024 Track Datasets and Benchmarks Poster, 2024."}, {"review_id": "HIdpbmzJ8Z", "reviewer": "Reviewer_8eQS", "summary": "1) Introduces ENMA, a continuous‐token autoregressive neural operator that models time‐dependent parametric PDEs in a compact latent space via masked spatial and causal temporal Transformers.\n\n2) Encode: irregularly sampled spatio‐temporal inputs are interpolated onto a regular grid and compressed with causal 3D convolutions into latent tokens.\n\n3) Generate: a causal Transformer predicts future latent frames, then a masked spatial Transformer + flow‐matching MLP decodes per‐token next‐state distributions without discrete quantization.\n\n4) Decode: A continuous VAE decoder reconstructs full physical fields from generated latent sequences, supporting in‐context learning and probabilistic forecasting.", "questions": "1) Can the authors provide wall‑clock benchmarks for end‑to‑end inference and memory usage as a function of token count and mask schedule?\n\n2) How does performance change without flow matching (e.g., using a direct MLP regression) or with different masking schedules in the spatial decoder?\n\n3) Have you tested ENMA on irregular meshes or 3D PDEs to validate its claimed domain‐agnostic adaptability?\n\n4) How well do the flow‐matching–derived uncertainties correlate with actual forecasting errors, especially under out‐of‐distribution parameter shifts?\n\n5) Can authors compare the proposed method with DiffusionPDE as a baseline?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\n1) Continuous latent autoregression avoids quantization artifacts and enables fine‐grained uncertainty estimation via flow‐matching.\n\n2) Irregular‐grid compatibility: attention‐based encoder handles unordered inputs and varying sampling rates gracefully.\n\n3) In‐context adaptation: conditions on either past target states or auxiliary trajectories to generalize across parameter regimes without retraining.\n\n4) Scalable compression: combines spatial and temporal token compression while maintaining high reconstruction and forecasting fidelity.\n\n**Weaknesses**\n\n1) Compression trade‑off: High latent compression can degrade accuracy on complex dynamics (e.g., vorticity), highlighting limited expressiveness.\n\n2) Computational overhead: Two-stage Transformers plus flow‑matching sampling incur nontrivial inference cost and memory use.\n\n3) Benchmark scope: evaluation limited to five PDE families; lacks tests on stiff, multiphysics, or 3D systems from the PDEBench dataset.\n\n4) Ablation gap: missing analysis isolating the impacts of masked decoding schedule, flow‐matching vs. simpler autoregression, and causal CNN design.", "was_revised": true, "final_justification": "My main concerns have been addressed, and I hope the authors incorporate all the new results and open-source the codebase as suggested. I am increasing my score from 4 to 5.", "rebuttal": "We thank the reviewer for the detailed and thoughtful comments, which have helped us improve the quality and clarity of the paper. Below, we address each point raised.\n# Compression Trade-off (W1)\nWe acknowledge that latent compression often leads to reduced performance, especially for complex PDEs with high-frequency components.  However, for high dimensional datasets, operating directly in the physical space is usually prohibitive with transformers, hence compressing is often required. We believe that our work highlight a direction for improving this trade-off for solving PDEs.\nIn order to evaluate the influence of compression, in Appendix F.1.2, we ablate the compression ratio of ENMA across 2D datasets. We observe that reducing compression indeed significantly boosts ENMA’s performance, while still reducing the computational complexity compared to operating in the physical space, making it competitive on Vorticity with the best baselines and outperforming them on Wave and Gray-Scott.\n# Computational Cost (W2 & Q1)\nRegarding computational overhead, we agree that ENMA’s generative process incurs higher inference cost than deterministic models. This is inherent to generative methods leveraging iterative denoising or sampling. However, as shown in Appendix F.1.3, this overhead is limited and ENMA achieves strong performance with only a few autoregressive and denoising steps. We report in appendix section F.1.3 wall-clock inference times against other generative baselines (Zebra and AR-DiT), showing that ENMA is faster due to its lightweight design and flow-matching objective. This will be better referenced in the final version. As requested, we report below the memory usage as a function of token count for one sample. The experiment demonstrates a limited increase in memory requirements.\n\nTable R1: Memory consumption (Go) and inference time (s) wrt token count (on Vorticity).  \n| Token Count | Memory (MB) | Time (s) |\n|-|:-:|:-:|\n| 64 | 3896 | 0.5 |\n| 512 | 3896 | 1.5 |\n| 1024 | 3897 | 5.7 |\n| 2048 | 3897 | 10.6 |\n| 4096 | 3898 | 20.9 |\n| 8192 | 3900 | 40.9 |\n| 16384 | 3904 | 83.7 |\n    \n# New experiments on public benchmarks (W3 & Q3):\n\n\nWe fully agree that evaluating on public datasets enhances the reproducibility and comparability of our results. While our primary goal was to assess ENMA’s capacity to generalize across parametric time-dependent PDEs—where PDE parameters may vary significantly from one trajectory to another—most existing public benchmarks exhibit only limited parametric variability. Thus, they do not fully reflect the challenges addressed in our setting. For each physical system, we generated 1200 distinct PDE instances, each with a different set of parameters, and10 trajectories for each instance with a distinct initial condition, leading to 12000 samples. In comparison, datasets from public benchmarks offer a different setting: each PDE system in these datasets only includes a few number of parametric instances (about 10 vs 1200 for our setting), offering a limited diversity for each PDE. Thus current public benchmark, focus on solving a variety of PDEs, but do not allow to fully evaluate the parametric setting which is our target in the paper.\n\nNonetheless, ENMA is designed to be general and flexible, and it can readily be applied to a wide range of time-dependent PDE problems. To this end, we have included results on two public datasets that capture complementary aspects of ENMA’s capabilities:\n- **CylinderFlow**: A time-dependent PDE on an irregular grid. We evaluate ENMA’s encoder-decoder performance and compare to neural operator baselines designed for irregular spatial inputs. Table R2 shows the results of the auto encoding and time-stepping tasks as presented in table 1 in the manuscript. This experiment clearly highlights the capabilities of ENMA to manage irregular grids.  The time-stepping task is much better with ENMA's tokens, likely due to the temporal compression proposed ($\\times 22$ compression rate for ENMA, whereas baselines have $\\times 11$ compression rate).\n\n    Table R2: Results of the reconstruction task on the cylinder flow dataset. Metrics in Relative MSE on the test set. We highlight in bold the best performing model.\n    |Sampling|Baseline|Reconstruction|Time-stepping|Compression\n    |-|-|:-:|:-:|:-:| \n    || GINO| 6.85e-1 | 1.59e0|$\\times 11$\n    |100\\%| AROMA| 9.59e-1 | 6.59e-1 | $\\times 11$\n    || CORAL | 5.06e-2 | 7.84e-1 | $\\times 11$\n    || ENMA | **7.78e-3** | **1.44e-1** | **$\\times 22$**\n    |||||\n    || GINO| 6.85e-1 | 1.59e0 |\n    |50\\%| AROMA| 9.86e-3 | 6.78e-1 |\n    || CORAL | 4.96e-2 | 7.83e-1|\n    || ENMA | **7.80e-3** | **1.44e-1**|\n    |||||\n    || GINO| 6.85e-1 | 1.59e0 |\n    |20\\%| AROMA| 1.06e-2 | 6.82e-1 |\n    || CORAL | 5.14e-2 | 7.83e-1 |\n    || ENMA | **7.87e-3** | **1.44e-1** |\n\n- **Active Matter**: A complex PDE on a 256×256 regular grid with 11 input fields. We test ENMA’s generative capabilities under temporal conditioning and compare to autoregressive transformer baselines (Table 2). In Table R3, we observe that ENMA is close to AVIT (also known as MPP) while operating in a very compressed space. Relaxing this compression ratio would yields to better performance, as we observed for other baselines, but we believe that operating in such latent space for high dimensional data such as Active Matter is essential.\n\n    Table R3: Results of the temporal generation task on the Active Matter dataset. Metrics in Relative MSE on the test set. \n    |Baseline|Temporal conditioning|Compression ratio|\n    |-|:-:|:-:|\n    |AVIT|0.383| $\\times 1$|\n    |ENMA|0.412|$\\times 10$|\n\n# Ablation studies on ENMA's components (W4, Q2)\nWe appreciate the suggestions for further ablations. Note that Appendix F already includes several ablations—e.g., geometry-aware attention (F.2.3), compression ratio (F.1.2), sampling schedule and history length (F.1.4)—we include below new ablations, following the reviewer's suggestion:\n\n- **Masking schedules**: We evaluate decoding under different masking strategies (e.g., linear, exponential). In R4, we observe all strategies lead to similar performance.\n\n    Table R4: Ablation on the masking strategy on the test set (on Vorticity). Metrics in relative MSE. \n    |Scheduler|Relative MSE|\n    |-|:-:| \n    |Linear |0.0619|\n    |Cosine  |**0.0612**|\n    |Exponential  |0.0615|\n    |Reverse exponential |0.0618|\n\n- **Generative modeling choice**: We compare flow-matching with a denoising-based objective and a fully deterministic approach (table R5).  Flow matching works better than a deterministic approach, while offering the advantages of generative modeling (sampling etc).  Our experiments with diffusion training didn't converge correctly, leading to higher loss. Overall flow-matching is easier to train. \n\n    Table R5: Ablation of the generation method (on Vorticity). Metrics in Relative MSE on the test set.\n    |Training Objective |Relative MSE|\n    |-|:-:| \n    |Flow Matching (ENMA) |0.0644|\n    |Diffusion|0.7579|\n    |Deterministic |0.0879|\n\n- **Causal CNN**: Causal vs non causal convolutions in the encoder.\n Table R6 presents an ablation of the causal component in ENMA. Compared to the non causal autoencoder, the causal version uses causal layers and adds temporal compression, reducing the number of tokens by half. Despite this, it matches or outperforms the non causal version—especially in low-data settings— indicating that the causal layer helps in providing more informative tokens.\n\n    Table R6: Ablation of the Causal component (on Advection). Metrics in Relative MSE on the test set. \n    |Sampling|Relative MSE||\n    |-|:-:|:-:| \n    ||Causal|Non causal|\n    |100\\% | 5.16e-3 | 4.55e-3 |\n    ||||\n    |50\\%  | 6.98e-3 | 6.02e-3 |\n    ||||\n    |20\\%  | 1.66e-2 | 2.01e-2 |\n\n# Uncertainty Calibration under Distribution Shifts (Q4)\n\nIn Appendix F.1.1 (Figures 21–22), we evaluate UQ in the in-distribution setting by plotting CRPS and RMSCE across forecast steps. Following your suggestion, we have added the relative L2 error at each timestep in Table R7, highlighting its correlation with the UQ metrics. As expected in autoregressive models, L2 increases over time due to error accumulation. CRPS follows a similar trend, indicating that probabilistic accuracy degrades with the prediction horizon, while RMSCE decreases, suggesting the model is initially overconfident and becomes better calibrated as it increasingly relies on its own predictions—despite the rising forecast error.\n\nTable R7: Uncertainty calibration wrt time steps (on Combined). \n| Metrics | t=0 | t=2 | t=4 | t=6 | t=8 | t=10 | t=14 | t=16 | t=18 |\n|-|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|\n| CRPS |0.0006       |0.0008       |0.001       |0.0013       |0.0015       |0.0018       |0.0020       | 0.0022       |0.0023       |\n| RMSCE |0.141       |0.117       |0.106       |0.096       |0.087       |0.076      |0.0697       |0.064       |0.060       |\n| Relative L2 |0.016       |0.022       |0.025       |0.040       |0.052       |0.050       |0.058       |0.071       |0.072       |\n\nWe will also investigate the behavior in OoD settings in additional experiments. We expect ENMA to yield larger confidence intervals under distribution shift.\n   \n# DiffusionPDE baseline (Q5)\nThanks for mentioning DiffusionPDE—we will quote this work in the paper. DiffusionPDE however belongs to a different family of generative methods: it generates trajectories in one shot and not through an auto-regressive process which is the focus of the paper. It also leverages a PDE loss requiring the knowledge of the PDE equations, when we consider pure data-driven approaches. It is not designed for solving time-dependent parametric PDEs but one PDE instance at a time. Hence direct comparison is not possible.\n\n We thank the reviewer again for their valuable feedback and are happy to engage further if needed."}, {"review_id": "4852Wx4WpU", "reviewer": "Reviewer_aeqi", "summary": "This paper presents ENMA, a generative neural operator for time-dependent parametric PDEs that performs continuous autoregressive generation in latent space. Unlike existing approaches that rely on discrete tokenization or full-frame diffusion, ENMA uses a two-stage generation process: a causal transformer predicts future latent states, followed by masked spatial transformers that generate tokens using flow matching. The method can perform in-context learning by conditioning on historical states or auxiliary trajectories, allowing adaptation to new PDE regimes without retraining.", "questions": "1. The encoder uses cross-attention to interpolate irregularly sampled inputs onto a structured grid, but it’s unclear why this is necessary. Why wouldn’t simple masking of missing spatial points suffice? From my understanding, masking does not inherently result in information loss. The cross-attention mechanism adds architectural complexity—could the authors clarify what is gained from this choice?\n\n2. Equation (4) suggests that the model outputs predictions on the locations as input. Yet the paper claims the model can evaluate on arbitrary spatial grids. Please revise or clarify.\n\n3. The sentence “A compression rate of ×2 indicates that the latent space is half the size of the input in terms of raw elements” is somewhat vague. Does this refer to the number of spatial tokens, total elements across space and time, or some other metric? A more precise definition would help interpret the reported results.\n\n4. The paper emphasizes that ENMA operates with continuous latent tokens, but it’s not clear what this means in contrast to standard dense embeddings. How is this different from using a ViT encoder, where the inputs are mapped to latent embeddings without any vector quantization? If such embeddings also qualify as continuous tokens, what is the specific advantage or novelty of ENMA’s formulation in terms of continuous latent space/tokens?\n\n5. Appendix F contains valuable architectural analysis and ablation results, yet the main paper only refers to it in passing (e.g., “Additional analysis... in Appendix F”). I recommend the authors briefly summarize the key insights and conclusions from Appendix F in the main text to give readers a better sense of what is covered and why it matters.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strength:**\n\n- The proposed approach offers a promising alternative for PDE operator learning, particularly in settings requiring uncertainty quantification and in-context adaptability.\n\n- The experimental validation is relative comprehensive, including evaluations on both standard and challenging scenarios, such as out-of-distribution generalization and sparse input regimes.\n\n**Weakness:**\n\n- The paper feels rushed in several places. For example, in Figure 1, the text \"Continuous VAE tokenizer\" is not centered in its box, and section titles are inconsistently formatted (e.g., Section 3.1 vs. Section 4.2). \n\n-  The method section lacks clear structure, with architectural design, training procedure, and inference details all mixed together. This makes it difficult to understand the overall pipeline. I strongly suggest reorganizing this section into distinct parts to improve clarity.\n\n- While the proposed framework is sensible, it largely reads as a combination of existing techniques—causal transformers, masked spatial decoding, and flow matching. There’s limited discussion or ablation to isolate what’s truly novel or impactful in this particular integration. As a result, the contribution feels incremental rather than innovative.\n\n- One major concern is the choice of benchmarks. The model is not evaluated on standard public datasets such as PDEBench or PDEArena, which include multiscale and stiff systems and present challenges to neural PDE surrogates. Without comparison on these benchmarks, it's hard to assess ENMA's performance relative to strong baselines including DPOT and MPP.   I believe evaluating on at least one of these would substantially strengthen the empirical results.", "was_revised": true, "final_justification": "Through the rebuttal process, my main concerns have been addressed, and I will increase my score.", "rebuttal": "We thank the reviewer for the detailed and constructive feedback. We address the raised concerns point-by-point below.\n## **Typographical and Formatting Issues (W1)**\n We appreciate the reviewer highlighting these issues. All typographical and formatting inconsistencies (e.g., Figure 1 alignment, inconsistent section headers) will be corrected in the revised version.\n## **Clarity of the Method Section (W2)**\nThanks for the suggestion. In the final version, we will restructure this section to clearly distinguish between (i) model architecture, (ii) training procedure, and (iii) inference pipeline. We believe this will improve the readability and understanding of ENMA’s generative framework.\n## **Contribution and novelty (W3)** \nWhile we acknowledge that the components of ENMA—causal transformers, masked decoding, and flow matching—are based on existing ideas, their integration within ENMA offers an entirely new paradigm for PDE generation. Specifically:\n- ENMA performs token-wise autoregressive generation entirely in a continuous latent space. In contrast to prior work relying on discrete tokenization (e.g., Zebra) or full-frame diffusion (e.g., AR-DiT), ENMA combines the fine-grained generation of LLM-style models with the smooth nature of continuous representations.\n- We demonstrate that in-context generation is possible in this setting—an advance over previous generative operator learning frameworks, which typically rely on vector quantization.\n- On the encoding side, we propose a geometry-aware cross-attention encoder to project irregular spatial inputs onto a structured latent grid. This enables us to apply causal 3D CNNs and spatial patchification efficiently (Appendix D.2.1). This strategy is novel and we believe is an interesting direction; a concurrent work [1] proposed the same idea to patchify physical fields, which we will cite in the revised version.\n- We introduce a novel metric, the Fréchet Physics Distance (FPD), inspired by vision benchmarks, to evaluate generative modeling performance in physical systems (Appendix F.1.1).\n\nAs for the ablations, we detail below existing (Appendix F) and some additional experiments, following the reviewer's suggestion. We agree that this should be better referenced in the manuscript, and we will update the main text accordingly. \n### **Existing Ablations** \n- **Compression ratio** (Appendix F.1.2): We show that reducing the compression improves performance on Vorticity, making ENMA competitive while maintaining strong results on Wave and Gray-Scott.\n- **Geometry-aware encoding** (Appendix F.2.3): We compare reconstruction performance with and without the geometry-aware component. Across various input sparsity settings (full, half, sparse), the geometry-aware encoder consistently improves reconstruction accuracy.\n### **New Ablations**: \nWe have evaluated the importance of: (i)  causal encoder, (ii) flow-matching vs deterministic component, (iii) scheduler.\n\n- **Causal CNN**: We ablate the use of the causal convolution module in the encoder to assess its impact on temporal feature representation. Table R1 presents an ablation of the causal component in ENMA. Compared to the non causal autoencoder, the causal version uses causal layers and adds temporal compression, reducing the number of tokens by half. Despite this, it matches or outperforms the non causal version—especially in low-data settings— indicating that the causal layer helps in providing more informative tokens.\n\n    Table R1: Ablation of the Causal component (Advection). Metrics in Relative MSE\n    |Sampling|Relative MSE||\n    |-|:-:|:-:| \n    ||Causal|Non causal|\n    |100\\% | 5.16e-3 | 4.55e-3 |\n    ||||\n    |50\\%  | 6.98e-3 | 6.02e-3 |\n    |||\n    |20\\%  | 1.66e-2 | 2.01e-2 |\n- **Latent dynamics modeling**: We compare the flow-matching-based generation to both a diffusion-based latent model and a fully deterministic alternative. In table R2, we observe that a Flow matching framework works better than a deterministic approach, while offering the advantages of generative modeling. Diffusion didn't converge correctly in our experiments leading to lower performance.\n\n    Table R2: Ablation of the generation mode (Vorticity).\n    |Training Objective |Relative MSE|\n    |-|:-:| \n    |Flow Matching (ENMA) |0.0644|\n    |Diffusion|0.7579|\n    |Deterministic |0.0879|\n- **Scheduler**: We evaluate the importance of the scheduler strategy for choosing the number of tokens to generate. In table R3, we observe that the choice of the scheduler is not of particular importance; all strategies lead to very similar performance.\n\n    Table R3: Ablation of the Scheduler (Vorticity).\n    |Scheduler|Relative MSE|\n    |-|:-:| \n    |Linear |0.0619|\n    |Cosine  |0.0612|\n    |Exponential  |0.0615|\n    |Reverse exponential |0.0618|\n\nThese new experiments will be included in the final version to further validate ENMA’s design choices.\n## **Choice of Benchmarks (W4)**\nOur primary goal was to assess ENMA’s capacity to generalize across parametric PDEs—where PDE parameters vary  from one trajectory to another while most existing public benchmarks exhibit only limited parametric variability: for each physical system, we consider 1200 distinct PDE instances, each with a different set of parameters, when available benchmarks  only include a few (about 10) parametric instances. Thus current public benchmark, do not allow to fully evaluate our target parametric setting.  Nevertheless, following your suggestion, we  include new results on two public benchmarks:\n- **CylinderFlow** [2]: A well-known time-dependent PDE on irregular spatial grids, used to evaluate ENMA’s encoder-decoder stage against baselines (as Table 1 in the main text).\n    Table R4 shows the results of the auto encoding and time-stepping tasks. This experiment clearly highlights the capabilities of ENMA to manage irregular grids, both for reconstruction and time-stepping. ENMA also is more efficient in terms of compression due to the temporal compression proposed ($\\times 22$ compression rate for ENMA, whereas baselines have $\\times 11$ compression rate).\n\n    Table R4: Results of the reconstruction task on the cylinder flow dataset. \n    |Sampling|Baseline|Reconstruction|Time-stepping|Compression\n    |-|-|:-:|:-:|:-:| \n    || GINO| 6.85e-1 | 1.59e0|$\\times 11$\n    |100\\%| AROMA| 9.59e-1 | 6.59e-1 | $\\times 11$\n    || CORAL | 5.06e-2 | 7.84e-1 | $\\times 11$\n    || ENMA | **7.78e-3** | **1.44e-1** | **$\\times 22$**\n    |||||\n    || GINO| 6.85e-1 | 1.59e0 |\n    |50\\%| AROMA| 9.86e-3 | 6.78e-1 |\n    || CORAL | 4.96e-2 | 7.83e-1|\n    || ENMA | **7.80e-3** | **1.44e-1**|\n    |||||\n    || GINO| 6.85e-1 | 1.59e0 |\n    |20\\%| AROMA| 1.06e-2 | 6.82e-1 |\n    || CORAL | 5.14e-2 | 7.83e-1 |\n    || ENMA | **7.87e-3** | **1.44e-1** |\n\n- **Active Matter** (The Well [3]): A complex system on a $256 \\times 256$ regular grid with 11 input channels. We compare ENMA’s generative capacity under temporal conditioning to autoregressive baselines (Table 2 of the manuscript). In Table R5, we observe that ENMA is close to AVIT while operating in a very compressed space. Relaxing this compression ratio would yield to better performance, as we observed in Appendix F.1.2.\n\n    Table R5: Results of the temporal generation task on the Active Matter dataset.\n    |Baseline|Temporal conditioning|Compression ratio|\n    |-|:-:|:-:|\n    |AVIT|0.383| $\\times 1$|\n    |ENMA|0.412|$\\times 10$|\n\nRegarding baselines: Our AViT baseline uses the same architecture as MPP, minus the\nmodules designed for multi-field modeling, which are not required in our setting.\n\n## Questions\n1. **Cross-attention vs masking**: ENMA has been designed to handle inputs defined on point clouds without requiring any prior mesh structure. Cross-attention with geometry-aware bias is a mechanism that allows interpolation from arbitrary, unordered spatial coordinates onto a latent regular grid. This is important in many settings where observations are irregular and may change along time. This cannot be performed through masking which assumes inputs lie on a predefined grid with entries are potentially missing.\n2. **Equation (4)**: Thank you for pointing this out. We used this notation for simplification, but we agree that this is misleading and we will revise the equation and text to clearly reflect ENMA’s ability to evaluate on arbitrary spatial grids.\n3. **Compression ratio**: Compression rate refers to the number of points in the token space compared to the original spatio-temporal resolution. We will clarify this in the revision. \n4. **Continuous tokens**: By \"continuous latent tokens\", we refer to standard dense embeddings—used commonly in neural operator models. The novelty here lies in the use of dense embeddings in a tokenwise generative autoregressive model. All the tokenwise generative models so far leverage token quantization and make use of a classifier as for the generative component like in LLMs. Here we extend this setting to the generation of continuous distributions. ViTs also use continuous tokens but cannot be used in a generative mode, contrarily to our model.\n5. **Appendix F summarization**: We agree that Appendix F contains valuable content and will summarize its key findings in the main paper before the conclusion, and improve referencing throughout the text. \n\nWe hope this clarifies the motivation, novelty, and empirical contributions of our work. We thank the reviewer again for the constructive feedback and would be happy to engage in further discussion.\n\n## References \n\n[1] Wen et al., Geometry-Aware Operator Transformer as an Efficient and Accurate Neural Surrogate for PDEs on Arbitrary Domains, arXiv:2505.18781 (2025).\n\n[2] Fortunato et al. Learning mesh-based simulation with graph networks, International Conference on Learning Representations, 2021. \n\n[3] Ohana et al. The Well: a Large-Scale Collection of Diverse Physics Simulations for Machine Learning, NeurIPS 2024 Track Datasets and Benchmarks Poster, 2024."}]

Author rebuttal:

Meta-review: This paper presents ENMA, a novel generative neural operator that leverages tokenwise autoregression to model complex spatio-temporal dynamics arising from time-dependent parametric PDEs, integrating masked spatial and temporal transformers with continuous flow-matching in latent space, enabling both efficient compression and robust out-of-sample generalization. 
The reviewers agree on the strong technical quality of the work, highlighting its clear motivation, thorough experimental evaluation (including new results on public benchmarks), and extensive ablation analysis. Some details and comparisons (e.g., to diffusion-based approaches) could be discussed in even greater depth.
All reviewers ultimately support acceptance, citing methodological novelty, evaluation quality, and impactful results.
Decision: Accept (spotlight)