DISCO: learning to DISCover an evolution Operator for multi-physics-agnostic prediction
Authors: Rudy Morel
Jiequn Han
Edouard Oyallon
Venue: ICML 2025 poster (2025)
Keywords: PDEs
prediction
hypernetwork
transformer
neuralPDE
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: For multi-physics-agnostic prediction, we propose DISCO, a model that learns a hypernetwork to DISCover an evolution Operator from a short trajectory, which is then integrated via a neural PDE to predict future states.

Abstract: We address the problem of predicting the next states of a dynamical system governed by *unknown* temporal partial differential equations (PDEs) using only a short trajectory. While standard transformers provide a natural black-box solution to this task, the presence of a well-structured evolution operator in the data  suggests a more tailored and efficient approach. Specifically, when the PDE is fully known, classical numerical solvers can evolve the state accurately with only a few parameters. Building on this observation, we introduce DISCO, a model that uses a large hypernetwork to process a short trajectory and generate the parameters of a much smaller operator network, which then predicts the next states through time integration. Our framework decouples dynamics estimation -- i.e., DISCovering an evolution Operator from a short trajectory -- from state prediction -- i.e., evolving this operator. Experiments show that pretraining our model on diverse physics datasets achieves state-of-the-art performance while requiring significantly fewer epochs. Moreover, it generalizes well to unseen initial conditions and remains competitive when fine-tuned on downstream tasks.

DISCO: learning to DISCover an evolution Operator for multi-physics-agnosie prediction
Rudy Morel
Jiequn Han
Edouard Oyallon 12
Abstract
In this paper, consider dynamical systems described by Partial Differential Equations (PDEs). The conventional data-driven approach typically involves "learning" a fixed dynamic model from a large number of trajectories obtained from various initializations, which is often referred to as Neural Operator learning (Kovachki al., 2023; Boullé & Townsend, 2023). However, in many practical scenar- ios, the exact governing physical laws are unknown and may vary across different trajectories, making such methods unsuitable and requiring more flexibility.
We address the problem of predicting the next states of a dynamical system governed by un- known temporal partial differential equations (PDEs) using only a short trajectory. While standard transformers provide natural black- box solution this task, the presence of well- structured evolution operator in the data suggests a more tailored and efficient approach. Specif- ically, when the PDE is fully known, classi- cal numerical solvers can evolve the state ac- curately with only a few parameters. Build- ing on this observation, we introduce DISCO, model that uses large hypernetwork to pro- cess a short trajectory and generate the parame- ters of a much smaller operator network, which then predicts the next states through time inte- gration. Our framework decouples dynamics es- timation - i.e., DISCovering an evolution Op- erator from short trajectory - from state pre- diction - i.e., evolving this operator. Experi- ments show that pretraining our model on diverse physics datasets achieves oofe--ofthheart perfor- mance while requiring significantly fewer epochs. Moreover, generalizes well to unseen initial con- ditions and remains competitive when fine-tuned on downstream tasks. The code will be made publicly available upon publication https: / //githut hub. com/ RudyMorel/DISCO.
To address this limitation, we study the problem of predict- ing the next states of PDE-governed dynamical system from a few successive past state observations. where the underlying PDE is unknown and can vary across different trajectories. This task called mulli-physics--annxiic predic- fion has recently gained attention with the effort to build foundational models for PDEs (McCabe et al., 2024; Herde et al., 2024). It requires both estimating the underlying dynamic (inverse problem) and integrating it over time (for- ward problem), which are typically treated independently in existing literature (Hsieh et al., 2019; Blanke & Lelarge, prediction requires
2023). Tackling muult--hysscc--uee facing at least three types of variability in the data:
(1) The PDE class for the dynamics (e.g.. Navier-Stokes),
(2) The coefficients of the PDE (e.g., Reynolds number),
(3) The initial conditions of the trajectory.
A natural way handle these three levels of variability is to use models that process context i.e., sequence of preceding states to infer the underlying dynamics. Transformers, with their strong sequential learning capabili- ties (Vaswani et al., 2017), excel at this through In-Context Learning (Brown et al., 2020), allowing them to adapt to different dynamics using prior state information. This adapt- ability has enabled transformers to achieve remarkable per- formance in natural language processing, where context plays crucial role (Dubey et al., 2024). Recent works (Yang et al., 2023; Liu et al., 2023; Yang & Osher, 2024; McCabe et al., 2024; Serrano et al., 2024) have also shown that transformers, when trained on diverse PDEs and initial conditions, can predict future states across different con- texts, establishing them as powerful tools for data-driven
1. Introduction
Modeling dynamical systems from data is highly active area of research with the potential to significantly reduce computational costs (Kidger, 2021) and limit the need for ad-hoc engineering to predict future states (Farlow, 2012).
'Center for Computational Mathematics, Flatiron Institute, New York 'Sorbonne Université, CNRS, ISIR, Paris- France. Corre- spondence to: Rudy Morel < flatiooninsiiute.org.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

DISCO: learning to DISCover evolution Operator for multi-phyysie--gnnstic prediction
operator network
operator network
transformer
hypereework
- encode
3 decode
(a) task-speciitt operator
(b) standard transformer
(c) DISCO
>: generate
weights
neural ODE
Figure 1. (a) A specifc learned operator must be re-trained on unseen Physics. (b) A transformer (e.g., MPP) must learn both an encoder and a decoder, adding reconstruction error. (c) DISCO infers operator parameters without relying on an auto-encoder.
modeling of dynamical systems in physical domains.
the solver respects the underlying pattinuous-time nature of the physics and preserves spatial translation equivariance through the use of a U-Net (Ronneberger et al., 2015). Its parameters are fully populated by the hypernetwork, of- fering greater flexibility, especially estimating spatial derivatives (Bar-Sinai et al., 2019). Furthermore, the oper- ator lives in a low-dimensional space (of dimension 384), far smaller than the size of a data, creating an information bottleneck (Tishby et al., 2000), that forces the model to focus on the most essential aspects of the dynamics, prevent- ing overfitting to the initial conditions (level (3) variability above). DISCO is pretrained on two collections of PDE datasets. On PDEBench (Takamoto et al., 2022), it provides state-of-the-art performances on next frame prediction, with very few epochs. On The Well (Ohana et al., 2024), model compares favorably to both standard transformer ar- chitectures, as well as meta-learning frameworks.
However, applying standard transformers to physical sys- tems is challenging. In order to process a context of po- tentially high-resolution frames, they typically require to reduce spatial resolution through an encoder, which must then be inverted by decoder. Learning the decoder on physics data remains a highly non-trivial task (Olmo al., 2022). As result, transformers often require vast amounts of data to avoid overfitting and can struggle with predicting nut-of-distribution trajectories, leading to instabilities over multi-step rollouts (McCabe et al., 2023; 2024). In con- trast, classical solvers excel when the governing dynamies are known, evolving the system accurately with minimal parameters and without any data training. Such methods often preserve key structural properties of physics, such as continuous- evolution and translation equivariance as described in (Mallat, 1999) - where a spatial translation of the initial condition results in the same translation of the solution in the absence of boundary conditions. In contrast, transformers do not naturally inherit these properties. This lack of structure preservation contributes to transformers* limitations in physical applications, particularly in terms of training efficiency and stability.
Contributions. Our key contributions are as follows: (a) To our knowledge, this is the first work to combine a trans- former hypernetwork with a neural PDE solvers for spa- tiotemporal prediction. (b) We propose a model that tackles multi muult-phyysies-anoste prediction with tailored inductive bias for physical systems, providing improved interpretabil- ity by aligning more closely with classical numerical meth- ods. (c) Our method achieves state-of-the-art performance on next-step prediction across multiple physical systems in the well-studied PDEBench dataset, requiring significantly fewer epochs than baseline models, and (d) achieves better numerical accuracy on multi-step rollouts compared to both standard transformer models and meta-learning frameworks. (e) Our model is the first to be scaled to such a diverse range of PDEs in 1D, 2D, 3D, with spatial resolutions up to 1024, 512 x 512, 64 x 64 x 64 respectively. In contrast, most models tackling multi-physsccs-gnosiie prediction are trained on at most 2D data at 128 128 resolution. Upon publication, we will open-source our implementation.
In this work, we propose a novel model named DISCO that combines the best of both worlds by learning to DISCover an evolution Operator from a context of successive states of a trajectory. Specifically, large transformer hypernet- work processes data context generate parameters for much lower-dimensional PDural-PDE-like solver (Chen et al., 2018), which then predicts the next states through time integration. In this manner, we decouple context dy- namics estimation from state prediction, unlike previous transformer- based methods where both are entangled. Our approach can be viewed as a meta-learming algorithm (Thrun & Pratt, 1998; Finn al., 2017; Koupaï et al., 2024), de- signed to self-adapt to specific contexts, or tasks. However, unlike recent meta-learning methods, which faces only vari- ability (2) and (3) (Koupaï et al., 2024), our framework inherently handles all levels of variability (1), (2), and (3). Moreover, prior meta-learning approaches have not been scaled to high-resolution data, where directly learning an operator remains challenging. The operator network used in
2. Related works
Neural solvers and operators. While classical PDE solvers remain the state-of-the-an art for achieving high preci- sion, neural network- based surrogate solvers (Long et al.,
2

---

DISCO: learning to DISCover evolution Operator for mulli-phyysics-agnnsti prediction
Table 1. The datasets from PDERench (Takamoto et 2022) and The Well (Ohana al., 2024) used in this paper.
PHYSICAL DIMENSION
# OF FIELDS
RESOLUTION (TIME)
RESOLUTION (SPACE)
BOUNDARY
DATASET NAME
CONDITIONS
BURGERS
1D
200
1024
PERIODIC
SHALLOW WATER EQUATION
2D
100
128
128
OPEN
DIFFUSION- DIFFUSION-REACTION REACTION
2D
100
128
128
NEUMANN
INCOMP. NAVIER-STOKES (INS)
2D
1000
512
512
DIRICHLET
CОMP. NAVIER-STOKES (CNS)
2D
21
512
PERIODIC
ACTIVE MATTER
2D
81
256
PERIODIC
BULER AULTI-QQADRANTS
2D
512
PERIODIC OPEN
GRAY-SCOTT REACTIONNDDIFUUSOON
2D
1001
128 X 128
PERIODIC
RAYLEIGH-BÉNAAR
2D
4
200
512 128
PERIODIC DIRICHLET
SHEAR FLOW
2D
200
256 512
PERIODIC
TURBULENCE GRAVITY COOLING
3D
50
64 X x 64
OPEN
MHD
3D
100
64 X 64 x 64
PERIODIC
RAYLEIGH-TAYIE INSTABILITY
3D
120
64 X 64 X 64
PERIODIC X PERIODIC x SLIP
SUPERNOVA EXPLOSION
3D
6
59
64 x 64 64
OPEN
2018; Hsieh et al., 2019; Li et al., 2020; Gelbrecht et al., 2021; Kovachki et al., 2023; Verma et al., 2024; Liu- Schiaffini et al., 2024) have opened up new possibilities for inferring approximate solutions quickly for certain PDEs. However, they require training on samples from the same PDE. Some variants, such as Karniadakis al. (2021); Kochkov al. (2021), incorporate corrective terms to ap- proximate trajectory dynamics, but still lack adaptability to the context, which is key feature of our method. Sym- bolic regression (Lemos et al., 2023) separates trajectory dynamics inference from integration but struggles with high- dimensional data and handle large search spaces. One could bypass estimating dynamics from context; for instance, Herde et al. (2024) train large operator to predict the next frame from the previous one across various physical sys- tems. However, this requires retraining at inference for each unseen PDE to allow the model to determine the dynamics.
classes (level (1) in the introduction). They also require significant retraining at inference, making them impractical for large-scale data at inference- Koupai et al. (2024) tested their approach on 2D data up to 256 256 resolution.
Transformers and in-context models. Without explicitly integrating operator, transformer models trained to pre- dict the next frame from context of previous frames via attention have been applied to diction (Yang et al., 2023; Liu et al., 2023; Yang & Osher, 2024; McCabe et al., 2024; Hao et al., 2024; Cao et al., 2024; Serrano et al., 2024). These architectures operate in a latent (i.e., token) space and, since they directly out- put the frame, require learning decoder. As result, they must tackle not only the already challenging multi- physics-agnostic prediction task but also a reconstruction task, which is highly non-trivial for physics. In contrast, DISCO processes frames directly in the sample-space.
pre-
Meta-learning strategies for dynamical systems. Rather than solving each PDE separately, meta-learning leverages shared weights across tasks, yet existing methods face key limitations. First, they often rely additional infor- mation, such as explicit equation forms (Linial et al., 2021; Sun et al., 2025) or weak supervision from PDE coefficients or domain-specific quantities (Wang et al., 2022; Alesiani et al., 2022), while others assume prior knowledge of PDE symmetries (Mialon al., 2023). Second, their adaptation mechanisms are often restrictive. c.g., affine predictabil- ity (Blanke & Lelarge, 2023) or limiting adaptation only the first layer of an operator (Zhang et al., 2024). In con- trast, DISCO outputs an operator solely from prior frames, without assumptions on the PDE. Certain meta-learning strategies for PDEs (Kirchmeyer et al., 2022; Koupaï et al., 2024) circumvent these limitations but were trained only on variability levels (2) and (3), not across different PDE
3. Methodology
3.1. Preliminaries
Let (1,2r) R denote time and space. We consider PDEs defined in spatial domain Q C g for RR x Ω TM of the form
(x)((x) = 9 g(u((z), (zu(x), Viu((),...)) (1)
B[)( 0, x E 0n,
uo(x)) Unnit (x). Ω
where different corresponds to different physics (differ- ent equations and/or parameters) governing the system, and the next two equations correspond to the boundary and ini- tial conditions. Here, the spatial dimension, typically
3

---

DISCO: learning to DISCover an evolution Operator for multi-phyyiis--guussic prediction
ranging from 1 to 3 in physical systems, and m is the num- ber of physical variables described by the PDE. For nota- tional simplicity, we assume the PDE is time-homogeneou (i.e.. does not depend on t) though our methods can be straightforwardly extended to the time-inhomogeneous case. For PDEs with higher-order time derivatives, specifi- cally of order r, we can redefine the function of interest as (u, 0,u,. ................................ 0. u), which allows us to rewrite the PDE in the form of Eq. 1. Given spatiotemporal trajectory with unknown g. our goal is to predict the state U using the preceding T successive states Ut-T+1,-- -+1,--.. Me, where T is referred to as the context length. In practice, the observed u, is discretized over grid points (assumed fixed and uniformly- distributed in this paper) rather than being a continuous spatial function, yet we still use continuous notation.
significantly larger effective receptive field. Thus, using CNN fo to directly approximate the solution in the form of uu+ (e)) = fa(ue), as proposed in (Wang et al., 2022), would require a much larger network with more parameters, challenging to train.
3.2. Method: DISCovering an Operator from context
Our framework. We propose model called DISCO that combines the best of both worlds from transformer models and classical numerical schemes (see Fig. 1): the parameters 0 E Rd of a small network fe-referred to as the operator network are generated from an ad-hoc model th referred to hypernetwork (Ha et al., 2016), which uses the context to estimate 0, leading to:
úc+1 = ue + J
pt+1
fo
Finite difference and neural PDEs. Finite difference method is a classical approach for numerically solving an explicitly given PDE by discretization on a grid. This con- nection forms a basis for leveraging translation-equivarian networks to approximate PDE solutions. To illustrate this relationship, consider the standard 2D diffusion equation (x) = (z), boundary conditions, where B > 0 is the diffusion co- efficient. Let the state be discretized over a uniform grid with spacing Ar, denoted by Ut,i.j for 0 I i,j I N. Using a standard centered finite difference scheme, the Laplacian A can be approximated as u4 * 0, where 3 con-
(4)
0 = vo (ux (u.TT+,,.., Uz),
where a E TRdd are learnable parameters, while O are param- eters predicted by y, with d1 and d2 being the sizes of the operator network and hypernetwork, respectively. This for- mulation significantly structures the predictor, which now has convolutional structure aligned with Eq. 1: the spatial derivatives are necessarily approximated using the convolu- tional kernels of fo-
1] subject to periodic
Numerical integration. The integral from i to L + 1 in Eq. 4 is discretized using third order adaptive Runge- Kutta method (Bogacki & Shampine, 1996). Gradient backprop- agation is performed using an adjoint sensitivity method which scales linearly in the number of integration steps and has low memory cost (Chen et al., 2018). The solver uses at maximum 32 integration steps, but can use less steps, thus automatically adapting to the effective "speed" of the current trajectory. An alternative would be to fix a budget of 32 frames and let the solver predict longer in the future if possible. It would provide a cheap way of training on long rollouts at training time, which we leave for future work.
0 3/(Az)²(0, 1, 0; 1, 1,,1,,1,, 0, 1, 0] is a small volution filter, and * denotes the convolution operation with periodic padding. Defining fa(u)= == 0, we have uy-1 R ue I fo(u.) ds, which indicates that when the physies is explicitly known, we only need single convolu- tion layer with few parameters to evolve the solution u, accurately.
Extending this concept to machine learning, Bar-Sinai et al. (2019) proposed replacing these fixed convolutional coef- ficients with learnable parameters, resulting in improved accuracy on certain grid sizes. For more complex PDEs involving nonlinear dynamics, additional layers with nonlin- ear activation functions are necessary to capture the under- lying effects (Long et al., 2018; Kochkov et al., 2021). This connection between finite difference methods and CNNs highlights that relatively small CNN architectures, com- pared to transformer-based models, can effectively represent temporal PDEs while preserving spatial translation equiv- ariance. This perspective has been explored in previous work (Ruthotto & Haber, 2020) and recently applied to de- sign architectures for solving PDEs when the governing physics are known (Liu et al., 2024). This connection holds primarily for continuous-timg time or small discrete-time steps. For larger time steps At, the exact solution to the above heat equation still retains a convolutional structure, with a
Operator network architecture fo. The operator net- work fo in DISCO has to be fast (because it will be inte- grated) and oow-dimensional (to enforce a bottleneck). We used a U-Net with 4 downsampling blocks, a MLP in the deepest layers and 4 upsampling blocks with their respective skip connections. The input channel dimension is mapped to 8 before any downsampling, In the deepest layer, it in- creased to 128 = 8-24, and goes through a MLP with hidden dimension of 256. We used GeL.U activation (Hendrycks & Gimpel, 2016) and group normalization (Wu & He, 2018) with groups. For datasets with nor-poiiodic boundary conditions, we use reflection-paddingg for the spatial con- volution and we add a mask representing the boundary as
4

---

DISCO: learning to DISCover an evolution Operator for multi-phyysic--agnosie prediction
validation NRMSE
volidation NRMSE
validttinn NRMSE
validation NRMSE
vadiditiin NRMSE
validation NRMSE
DESCD (oaa)
KIP
akatity
30
20
20
epects
spachh
spuads
spodds
opechh
(a) Burgers
(b) Shallow water
(c) Diff.-reaction
(d) Incomp. NS
(e) Comp. NS
(l) Comp. NS (alone)
Figure 2. Learning curves on PDEBench datasets. DISCO achieves already decent predictions after one epoch, while an auto-encoder
transformer architecture such as Poseidon or MPP requires many epochs before beating a simple identical prediction de+1 = us (black dashed line). On shallow- waterr dataset for example, our model achieves 10 performance after one epoch, while Poseidon and MPP models require more than 150 epochs to reach same accuracy. This an illustration of the better inductive bias implemented our framework, which also avoids having to train a decoder.
additional channel as input to fo. We end up with an opera- tor network with roughly d1 = 200k parameters @ and have an intrinsic dimension of 384 as described below.
paper, we choose the loss function the normalized root mean square error (NRMSE); see App. D for the detailed definition. After the weights a are trained, predicting a specific dynamic from a given trajectory is a downstream task that can be tackled with a simple forward in our model, as no retraining is involved, unlike standard meta-learning models such as (Koupaï et al., 2024).
Hypernetwork architecture v/. Multiple choices of hy- pernetwork could be used, but we decided to exploit trans- formers due to their favorable scaling properties (Lin et al., 2022). In fact, we emphasize that any type of meta learning approaches, like directly optimizing f via gradient descent, could be considered yet it might not be computationally fa- vorable. Our hypernetwork v/ consists of three components: a CNN encoder that reduces spatial resolution, a processor composed of attention layers across time or space, and a final parameter generation block that outputs the parameters 0 of the operator fo- We refer the reader to App. C for details the architecture. The output of the processor is averaged over both time and space, resulting in a single to- ken of dimension 384, which is the intrinsic dimensionality of our operator network. Finally, we feed this token to a MLP with two hidden layers that progressively increases the channel dimension recover the expected parameter shape of f. The output is then normalized so to account for the different layers in the operator network fe and fed to a sigmoid to restrict the range of each parameter 9, which is key in making the training stable. The resulting hypernet- work has roughly d2 R 120m parameters a when trained on PDEBench. We thus have di < d2, meaning that the size of the operator network fe much smaller than the size of the hypernetook vi (see Eq. 4). This is in order limit overfitting and any memorization phenomenon by forcing the transformer hypernetwork to estimate the parameters of the trajectory rather than the next state.
4. Numerical experiments
Multiple physics datasets. In order to showcase the per- formances of DISCO on multi-physscs-agnone tion, we consider two collections of datasets (see Tab. I and App A for details). The PDEBench dataset has been well studied in the literature with SOTA performances on multi- physics-agnosti prediction obtained by MPP (Mc- Cabe et al., 2024). The Well datasets (Ohana et al., 2024) are recently released datasets exhibiting many more exam- ples of Physics evolution operators including 3D operators. Our study is one of the first to test a large model on this rich datasets. Both collections of datasets include simulations of PDEs with varying parameters and varying initial con- ditions. All contexts are kept at their original resolutions, except for the Rayleigh- Taylor instability dataset which orig- inally figures 128 x 128 x 128 and was downsampled to 64 x 64 x 64 for memory purposes. Details on the equations, initial conditions, boundary conditions, and data generation can be found App. A.
predic-
Multiple physics training. When trained on multiple physics, our model is made agnostic to the dataset name and only sees the spatial dimensionality of the data, ID, 2D, 3D, the fields names, e.g., pressure, velocity and the type of peri- odic boundary conditions (see Tab. 1). Most of the weights in the hypernetwork & are thus shared across different con- texts, except for the first field embedding layer (McCabe et al., 2024), which contains field-specife weights, and the very last layer, which is learned per context dimension (ID, 2D, and 3D). The operator network fa hyperparameters, such as number of layers, number of downsamplings, hid- den dimension, are unchanged across all the experiments
End-to-end training. The two equations in Eq. 4 can be learned jointly in an end-to-end manner by solving the optimization problem
min D
Loss(u alt41)
RED
where each lo in the dataset D is a data context, that is a small trajectory, in the form of (ut- T+ ut, utt u In this
5

---

DISCO: learning to DISCover evolution Operator for multi-physsic--agnosic prediction
Table Comparison with SOTA model MPP on PDEBench datasets. Next step prediction performance measured NRMSE (lower is better). Our model trained for 300 epochs achieves SOTA performances set by the MPP model trained on 500 epochs on most datasets. Due to the smaller time resolution of the CNS dataset compared to the other PDEBench datasets (see Tab. 1), our model does not reach state-of-the-art performance. However, when trained on this dataset alone, DISCO does achieve better performance (see Fig. 2 for the training curves): 0.041 vs 0.072 after 50 epochs.
Model
Epochs
Burgers
SWE
DiffRe2D
INS
CNS
MPP (retrained)
300
0.0046
0.00091
0.0077
0.0068
0.037
MPP (retrained)
500
0.0029
0.00075
0.0058
0.0056
0.031
MPP (McCabe et al., 2024)
500
0.00240
0.0106
0.0227
DISCO (ours)
300
0.0027
0.00056
0.0047
0.0033
0.095
conducted in this paper. The operator network size will however vary in size with the context, e.g. if the context is 3D then the convolution layers are 3D layers and thus have more weights. Also, the first layer and last layer of the operator adapts to the fields present in the context. We refer the reader to App. C for further detail on the exact architectures and number of weights.
2015) and FNO (Li et al., 2020)) for finctuning experiment presented later in the paper. See App. C.
Next steps prediction performances. Tab. 2 shows that our model beats state-of-the-ar performances on next step prediction on PDEBench, reported in (McCabe et al., 2024), for most of the datasets. It is remarkable that only a few epochs are required to beat the performances on most of the datasets. As reported in Tab. 5, it only takes 35 epochs to beat the performances of a transformer model trained on 500 epochs. These results are confirmed by observing the validation loss along training (Fig. 2). As we can see, the Poseidon and MPP baselines, can take up 150 epochs to several epochs to beat naive identity prediction (next step is predicted as the current step). This, we believe, is due to the fact that these baselines rely on an auto-encoding of the data, which is in itself a non-trivial task that requires long training. In contrast, DISCO does not employ a decoder, we learn to output the weights of an operator that acts in the data space directly. When trained jointly on the PDEBench datasets, we are do not compare favorably to MPP on the compressible Navier-Stoke dataset. This is due to the fact that this dataset has much coarser time resolution (see Tab. 1) relatively to the other datasets. As a matter of fact, when trained on this dataset alone, DISCO outperforms a standard transformer, see Tab. 2 and Fig. 2f. In order for our model to better adapt to contexts with diverse "speeds" of evolution, one could imagine harmonizing the temporal resolution the beginning of the training and then increasing it slowly for the contexts allowing it, which we leave for future work.
Baselines. In order to cover the different frameworks used in the literature to tackle sellii-hhssis-aagnooe prediction (see Sec. 2), we implemented the following three baselines
the GEPS framework (Koupai et al., 2024) is meta- learning approach extending CODA (Kirchmeyer et al., 2022) which Iearns an operator fa on data from differ- ent contexts (called "environments" in their paper) by training decomposing 0 into shared weights across the contexts and context-specific weights which need to be re-trained at inference. We use this framework our U-Net operator fa. Since we tackle multi-physics- agnostic prediction, the common weights are shared across all contexts, thus incorporating the three types of variability in our task.
the Poseidon model (Herde et al.. 2024) is based on hierarchical multiscale vision transformer, trained pairs of frames (i.e. contexts of a single frame) to predict a next frame in the future. This models also takes a continuous horizon variable n to decide which future frame Uttn it predicts. It is thus an example of large operator network that is trained on multiple physics at the same time and must be finetuned, on unseen data, so to estimate the new Physics.
Tab. 6 assesses the performances of the models on predict- ing the next steps t + n for n 2 1. At training time, all the models were trained on predicting only one step in the future + 1, except Poseidon which was trained on pairs (t, t+ + n) as in (Herde et al., 2024). As expected this model per- forms decently on rollouts. As we can see, DISCO remains competitive on predicting future steps on PDE data. Note that improving the rollout performance of an autoregressive model is an active area of research (McCabe et al., 2023), and one could complement our model with techniques such as noise injection (Hao al., 2024), among others.
MPP (McCabe et al., 2024) is transformer- bassdd model which processes contexts of frames to directly predict the next frame. We did not implement (Ser- rano et al., 2024), which can be seen as an in-context enhanced model of MPP, since we did not focus on in-context capabilities in our paper.
On top of these models we also compared our results to 2 additional baselines (a large U-Net (Ronneberger et al.,
6

---

DISCO: learning to DISCover evolution Operator for multi-phyysics-agnnsti prediction
Table Rollout performance models trained on Well datasets. models are trained for 50 epochs. use NRMSE
(lower is better). DISCO, which relies on an operator and does not require learning a decoder, achieves of art performances on large-scale diverse dataset as the Well (Ohana al., 2024). See Fig. learning curves across training, Figs. 11, 12, 13 for rollout examples.
No RETR AINING
DECODER FREE
MODEL
DATASET
NRMSE
AT INFERENCE.
[+] 1
444
1+8
L 16
ACTIVE MATTER
0.436
0.852
0.878
0.894
BULER MULTI-QUADRANT
0.759
0.873
0,920
0.939
GRAY-SCOTT REACTION-DIFFUSIO
0.158
0.480
0.615
0.718
RAYLEIGH
RAYLEIGH-BÉNARD BÉNARD
>1
>1
GEPS
X
SHEAR FLOW
>1
>1
>1
>
TURBULENCE GRAVITY COOLING
0,509
0.961
>
>1 0.992
MHD
0.977
0.949
0.979
RAYLEIGH-TAYYLOR INSTABILITY
>1
>1
>
>1
SUPERNOVA EXPLOSION
>
>1
>
>1
ACTIVE
0.134
0.442
0.778
S1
EULER MULTI-QUEE BULTI-QUADRANTE
0.070
0.166
0.285
0.446
GRAY-SCOTT REACTION-DIFFUSION
0.0652
0.118
0.197
0.348
SYYLEIGH-BÉNARE
0.115
0.243
0.396
0.636
MPP
X
SHEAR FLOW
0.0522
0.138
0.261
0.501
TURBULENCE GRAVITY COOLING
0.254
0.460
0.670
0.938
MHD
0.435
0.655
0.876
>
HELLEIG--TALLOR INSTABILITY
0.366
0.506
0.720
0.995
SUPERNOVA EXPLOSION
0.673
0.803
0.913
>1
ACTIVE MATTER
0.114
0.367
0.664
0.995
BULER MULTI QUADRANTS
0.0637
0.157
0.248
0.401
GRAY-SCOTT REACTION DIFFUSION
0.00990
0.0373
0.0789
0.215
RAYLEIGH--ÉÉNAR
0.0741
0.208
0.369
0.758
DISCO (OURS)
SHEAR FLOW
0.0249
0.0846
0.157
0.307
TURBULENCE, GRAVITY COOLING
0.165
0.418
0.605
0.822
MHD
0.298
0.684
0.905
>1
RAYLEIGH- TAYLOR INSTABILITY
BAYLIIGH-TYYLOR
0.0789
0.259
0.458
>1
SUPERNOVA EXPLOSION
0.387
0.656
0.778
0.982
While PDEBench (Takamoto et al., 2022) is an early dataset designed for evaluating prediction models in Physics, the Well dataset (Ohana et al., 2024) is a unified collection of more diverse scenarios covering biological systems, fluid dy- namics, including astrophysics simulations. The challenge lies in the diverse operators in the data, making unelear whether model
applied on our operator network. This framework learns shared parameters across the various context during training
and then adapts these parameters through gradient descent at inference time. While this framework seems competitive on contexts originating from the same class of PDE (Koupai et al., 2024), our hypothesis is that on a multi-physics- agnostic prediction large-scale data it is too difficutt in practice to learn shared parameters of relatively "small" operator fo different classes of PDEs. Thus, what we show in this paper is that it is possible to learn shared param- eters, but as the ones of the hypernetoor, not as parameters of the operator itself.
be trained such varied physics. Tab. 3 reports the results of DISCO on the Well, as well as two of our baselines (Poseidon was not implemented for 3D data, nor was for non-square 2D data) for model trained on 50 epochs. As we can see, we compare favorably on all the datasets we trained on, especially on rollouts: incorporating the right inductive biases with an operator and avoiding train auto-encoder enable to beat standard spatiotemporal prediction frameworks for PDEs. To date, the first paper training jointly on such diverse dataset of evolution operators in Physics.
Information bottleneck and operator space. DISCO employs low-dimensional operator fo (see Eq. 4) with an intrinsic dimension of 384, very much smaller than the typical input sizes. We study the impact of this information bottleneck. When trained on multiple datasets, the param- eters 9 returned by our hypernetwork depend on both the initial conditions and the PDE governing the given context.
one of
Let us note that the GEPS adaptation framework does not manage obtain competitive performances, least when
7

---

DISCO: learning to DISCover evolution Operator for multi-phyysic--agnosie prediction
TRUTH
U-Net
FNO
GFPS
Poseidon
MPP
DISCO (ours)
Figure 3. Step t + 16 predictions for different models on the Euler dataset with fixed PDE coefficients. U-net and FNO are trained from scratch while GEPS, Poseidon, MPP and DISCO finetuned from pretrained models on PDEBench (see also Tab. 4).
epoch 0
epoch 10
epoch 290
Table 4. Finetuning on unseen Physics. Performances in NRMSE of pretrained models, GEPS, Poseidon, MPP, when fine-tuned on the unseen Euler equations dataset with single, fixed, set of coeffi- cients (7 = 1.4) after 20 epochs of finetuning. For comparison, we show the performances of two other models: U-Net, FNO, trained from scratch.
valaa oi 1i:0
Cano aotis
Umap axis
Umap ands
Umap axis
Figure 4. Visualization of the parameter space using UMAP compressible Navier-Stokes data. Each point represents a set for given
UNET FNO GEPS
POSEIDON
MPP
DISCO
of parameters 0 predicted by the hypernetwork context. The hypernetwork tends to predict similar 9 values when the context is derived from PDEs with the same parameters n, but different initial conditions. This demonstrates the generalization capability of our approach to varying initial conditions.
0.050
0.067
0.36
0.052
0.032
0.029
conditions not encountered during pretraining, Since there is only one fixed PDE in the data, we also compare against networks such as U-Net (Ronneberger et al., 2015) (not to be confused with the smaller U-Net fo we used), Fourier Neu- ral Operators (FNO) (Li et al., 2020), and autoregressive- diffusion models(Kohl et al., 2024), which are designed to learn a fixed operator from data. Tab. 4 and Fig. 3 show our model achieves the best performance when fine-tuned on this new dataset, demonstrating its ability to generalize to a new PDE and novel initial conditions better than a standard transformer (due to better inductive biases) and other neural operator methods (thanks to pretraining).
We show that our framework is capable of reducing the variability introduced by the initial condition, allowing it to focus primarily on the PDE dynamics. For that, we study the parameter space 0 for certain 2D contexts, through di- mensionality reduction. The size of 0 varies slightly from context to the other. This is because the first and last layer of fe. which is a U-Net, adapt to the number of channels in the data. The remaining parameters have the same size 157 544 across all 2D contexts, and we perform dimensionality re- duction via Umap (McInnes et al., 2018) on them, for 128 contexts from compressible Navier-Stokes dataset, with two different shear viscosity 7) = 0.01 and 1 = 0.1 (see App. A). Fig. 4 visualizes the parameter space across different stages of the training. showing that our model progressively identi- fies two distinct clusters, corresponding to the two physical parameter values. Thus, the hypernetwork clusters together contexts originating from the same PDE coefficient, despite different initial conditions, which is the key for generalizing initial conditions and redueing unnecessary variability of level (3) (see Sec. 1).
Ablation on model size. Our model size depends both on the size of the operator network fo which is integrated to predict future frames, and the size of the hypernetwork t which predict weights 0 from the context. Tab. 7 shows vari- ations in the performances with changes on both networks. The model trained only on the diffusion-reaction dataset from PDEBench.
5. Conclusion
We introduce DISCO, a general and efficient framework for multi-physies-agnostio prediction. It uses transformer- based hypernetwork that adapt to varying contexts to gener- ate the parameters of neural PDE solver, which leverages continuous-time dynamics and spatial translation equivari- ance. Compared to pure transformers, DISCO achieves superior generalization and fine-tuning performance. The operator network in our framework, implemented using con- volutions, is primarily inspired by finite difference schemes on a uniform mesh. However, many challenging problems
Fine-tuning on unseen physics. We further assess the generalization properties of DISCO to an unseen PDE with fine-tune GEPS
unseen initial conditions. Specifically, model, the Poseidon model and the MPP model, both ini- tially trained on PDEBench, on the Euler dataset with a single set of parameters so that this unseen Physics exhibits only one operator in the data. The dataset is governed by compressible
equations and contains initial
8

---

DISCO: learning to DISCover evolution Operator for multi-physsies-agosti prediction
in physics involve non-uniform meshes or arbitrary geome- tries. In such cases, finite volume or finite element schemes can be implemented using graph neural networks instead (Pfaff et al., 2020; Zhou et al., 2022; Brandstetter et al., 2022; Zhou et al., 2025). In particular, several papers, such as (Esmaeilzadch et al., 2020; Lino et al., 2022; Cao et al., 2023), propose U-Net-like graph neural network architec- tures, which are natural candidates for our operator class.
tial equations. Proceedings of the National Academy of Sciences, 116631):15344-55333 2019.
Blanke, M. and Lelarge, M. Interpretable meta-learning of physical systems. arXiv preprint arXi::2112.00477, 2023.
Bogacki, P. and Shampine, L. F. An efficient runge-kutta (4, 5) pair. Computers & Mathematics with Applications, 32 (6):15-28, 1996.
Promising directions for expanding the capabiliiies of DISCO include multi-step rollout during training, that is, predicting future steps (c.g., f 1, + 2, 3) with the same neural solver and backpropagating from a loss av- eraged on all of them. This is made particularly advanta- geous in DISCO as applying the small neural solver re- quires little memory. While our method is validated on time-indepennden dynamics, extending it to time dependennd systems requires integrating temporal inputs into our hy- pernetwork and is left for future work. Additionally, our convolutional operator network is inspired by finite differ- ence schemes; exploring spectral methods could lead to architectures similar to FNOs (Li et al., 2020), leveraging new inductive biases
Boullé, N. and Townsend, A. A mathematical guide to oper- ator learning. arXiv preprint arXiv:2312.14688 2023.
Brandstetter, J., Worrall, D. E., and Welling, M. Message passing neural PDE solvers. In International Conference on Leurning Represeniations, 2022.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. Advances in Neural Information Processing Systems, volume 33, pp. 1877 1901, 2020.
Acknowledgement
We thank our close colleagues from Polymathic AI for their valuable feedhack on our paper, including Michael Mc- Cabe, Ruben Ohana, Lucas Meyer, Miles Cranmer, François Lanusse, Bruno Régaldo-San -aint Blancard, Alberto Bietti, and Shirley Ho. We also thank Patrick Gallinari and Lawrence Saul for insightful diseussions, well as Louis Serrano and Armand Kassai Koupai. We also thank Géraud Krawezik and the Scientific Computing Core at the Flatiron Institute, a division of the Simons Foundation, for providing com- putational resources and support. Part of this work was supported by PEPR IA on grant SHARP ANR-23-PEIA- 0008.
Burns, K. J., Vasil, G. M., Oishi, J. S., Lecoanet, D., and Brown, B. P. Dedalus: A flexible framework for numeri- cal simulations with spectral methods. Physical Review Research, 2(2):023068, 2020.
Cao, Y., Chai, M., Li, M., and Jiang, C. Efficient learning of meshh-aseed physical simulation with bi-stride multi- scale graph neural network. In International conference on machine learning, pp. 3541-3558. PMLR, 2023.
Cao, Y., Liu, Y., Yang, L., Yu, R., Schaeffer, H., and Os- her, S. Vicon: Vision in-context operator networks for multi-physic fluid dynamics prediction. arXiv preprint arXiv:2411. 2024.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle,
References
A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The Llama 3 herd of models. arXiv preprint arXiv:2407.22 2 1783, 2024.
Alesiani, F., Takamoto, M., and Niepert, M. Hyperfno: Improving the generalization behavior of fourier neural operators. In NewrIPS 2022 Workshop on Machine Learn- ing and Physical Sciences, 2022.
Esmaeilzadeh, S., Azizzadenesheli, K., Kashinath, K., Mustafa, M., Tchelepi, H. A., Marcus, P., Prabhat, M., Anandkumar, A., et al. Meshfreeflownet: A physics- constrained deep continuous space-time super-resolution framework. In SC20: International conference for high
Bar-Sinai, Y. Hoyer S., Hickey, J., and Brenner, M. P. Learning data-driven discretizations for partial differen-

---

DISCO: learning to DISCover evolution Operator for multi-phyysics-agnosie prediction
performance computing, networking, storage and analy- sis, pp. 1-15. IEEE, 2020.
Kochkov, D., Smith, A., Alieva, A., Wang, Q., Brenner, M. P., and Hoyer, S. Machine putational fluid dynamies. Proceedings of the National Academy of Sciences, 11(21)):2101784111, 2021.
com-
Farlow, S. J. Partial differential equations for scientists and engineers. Courier Corporation, 2012.
Kohl, G., Chen, L., and Thuerey. N. Benchmarking au- toregressive conditional diffusion models for turbulent flow simulation. In ICML 2024 Al for Science Workshop, 2024.
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. In Interna- tional conference on machine learning. pp. 1126-1135. PMLR, 2017.
Koupai, A. K., Benet, M., Yin, Y., Vittaut, J.-N., and Gallinari, P. Geps: Boosting generalization in parametric pde neural solvers through adaptive conditioning, arXiv preprint arXiv:2410.23889 2024.
Gelbrecht, M., Boers, N., and Kurths, J. Neural partial differential equations for chaotic systems. New Journal of Physics, 23(4):043005, 2021.
Ha, D., Dai, A., and Le, Q. V. Hypernetworks. arXiv preprint arXiv:1609.000 09006, 2016.
Kovachki, N., Li, Z., Liu, B., Arizzadenesheli, K., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Neural operator: Learning maps between function spaces with applications to PDEs. Journal of Machine Learning Re- search, 199,
Hao, Z., Su, C., Liu, S., Berner, J., Ying, C., Su, H., Anand- kumar, A., Song, J., and Zhu, J. Dpot: Auto-regressive denoising operator transformer for large-scale pde pre- training, arXiv preprint arXiv:2403.03542 2024.
Lemos, P., Jeffrey, N., Cranmer, M., Ho, S., and Battaglia, P. Rediscovering orbital mechanics with machine learning. Machine Learning: Science and Technology, 4(4):045002, 2023.
Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.0841: 2016.
Herde, M., Raonić, B., Rohner, T., Käppeli, R., Moli- naro, R., de Bézenac, E., and Mishra, S. Poseidon: Efficient foundation models for pdes. arXiv preprint arXiv:2405.19101, 2024.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions.
arXiv preprint arXiv:2010.08899 2020.
Holl, P., Koltun, V., Um, K., and Thuerey, N. phiflow: A differentiable PDE solving framework for deep learning via physical simulations. In NeurlPS workshop, volume 2, 2020.
Lin, T., Wang. Y., Liu, X., and Qiu, X. A survey of trans- formers. Al open, 3:111-132, 2022.
Linial, O., Ravid, N., Eytan, D., and Shalit, U. Generative ODE modeling with known unknowns. In Proceedings of the Conference on Health, Inference, and Learning, pp. 79-94, 2021.
Hsich, J.-T., Zhao, S., Eismann, S., Mirabella, L., and Er- mon, S. Learning neural pde solvers with convergence guarantees. arXiv preprint arXiv: (006:0990..0000 2019.
Lino, M., Fotiadis, S., Bharath, A. A., and Cantwell, C. D. Multi-scale atation-cquivarian graph neural networks for unsteady eulerian fluid dynamics. Physics of Fluids, 34(8), 2022.
Karniadakis, G. E., Kevrekidis, I. G., Lu, L., Perdikaris, P.. Wang, S., and Yang, L. Physics-informed machine learning. Nature Reviews Physies, 3(6) 440, 2021.
Ketcheson, D. L., Mandli, K., Ahmadia, A. J., Alghamdi, A. De Luna, M. Q., Parsani, M., Knepley, M. G., and Emmett, M. Pyclaw: Accessible, extensible, scalable tools for wave propagation problems. SIAM Journal on Scientific Computing, 32(4):C210 2012.
Liu, J. W., Erichson, N. B., Bhatia, K., Mahoney, M. W., and Re, C. Does in-context operator learning generalize to domain-shifted settings? In The Symbiosis of Deep Learning and Differential Equations III, 2023.
Liu, X.- Zhu, M., Lu, L., Sun, H., and Wang, J.-X. Multi- resolution partial differential equations preserved learning framework for spatiotemporal dynamics. Communica- tions Physics, 7(1):1,, 2024.
Kidger, P. On Neural Differential Equations. PhD thesis, University of Oxford, 2021.
Kirchmeyer, M., Yin, Y., Donå, J., Baskiotis, N., Rako- lomamonjy, A., and Gallinari, P. Generalizing to new physical systems via context-informed dynamics model. In International Conference on Machine Learning, pp. 11283-11301 PMLR, 2022.
Stu-Schiafffni, M., Berner, J., Bonev, B., Kurth, T., Aziz- zadenesheli, K., and Anandkumar, A. Neural operators with localized integral and differential kernels. arXiv preprint arXiv:2402. 16845, 2024.
10

---

DISCO: learning to DISCover evolution Operator for mulii-phyyscc--gnussic prediction
Long, Z., Lu, Y., Ma, X., and Dong, B. Pde-net: Learning pdes from data. In International conference on machine learning, pp. 3208-3216. PMLR, 2018.
Serrano, L.,, Koupaï, A. K., Wang, T. X., Erbacher, P., and Gallinari, P. Zebra: In-context and generative pre- training for solving parametric pdes. arXiv preprint arXiv:2410.03477,0 2024.
Mallat, S. A wavelet tour of signal processing. Academic Press, 1999.
Sun, J., Liu, Y., Zhang, Z., and Schaeffer, H. Towards foundation model for partial differential equations: Mul- tioperator learning and extrapolation. Physical Review E, 1 (3):033334 2025.
McCabe, M., Harrington, P., Subramanian, S., and Brown, J. Towards stability of autoregressive neural operators. arXiv preprint arXiv:2306. 106.9966611...... 2023.
Takamoto, M., Praditia, T., Leiteritz, R., MacKinlay, D., Alesiani, F,, Pfltiger, D., and Niepert, M. Pdebench: An extensive benchmark for scientific machine learning. Advances Neural Information Processing Systems, 35: 1596-1611, 2022.
McCabe, M., Blancard, B. R.-S., Parker, L. H., Ohana, R., Cranmer, M., Bietti, A., Eickenberg, M., Golkar, S., Krawezik, G., Lanusse, F, et al. Multiple physics pretrain- ing for physical surrogate models. Advances in Neural Information Processing Systems, pp. to appear, 2024.
Thrun, and Pratt, L. Learning to learn: Introduction and overview. In Learning learn, pp. 3-17. Springer, 1998.
McInnes, L., Healy, J., and Melville, J. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv:/802 032663426, 2018.
Tishby, N., Pereira, C., and Bialek, W. The informa- tion bottleneck method. arXiv preprint physics/0004057, 2000.
Mialon, G., Garrido, Q., Lawrence, H., Rehman, D., LeCun, Y., and Kiani, B. Self-supervised learning with lie symme- tries for partial differential equations. Advances in Neural Information Processing Systems, 01229933-29004 2023.
Toro, E. F., Spruce, M., and Speares, W. Restoration the contact surface in the hll-riemann solver. Shock waves, 4: 25-34, 1994.
Ohana, R., McCabe, M., Meyer, L., Morel, R., al. The well: large-scale collection of diverse physics simula- tions for machine learning. In Proceedings of the 38th Conference on Neural Information Processing Systems (NeurIPS 2024). NeurlPS, 2024.
Van Leer, B. Towards the ultimate conservative difference scheme. second-order sequel godunov's method, Journal of computational Physics, 32(1):101-136, 1979.
Vaswani, A., Shazeer, N.,, Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I. At- tention is all you need. Advances in Neural Information Processing Systems, 2017.
Olmo, A., Zamzam, A., Glaws, A., and King, R. Physics- driven convolutional autoencoder approach for cfd data compressions. arXiv preprint Xiv:2210000962. 2022.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Verma, Y., Heinonen, M., and Garg, V. Climode: Climate and weather forecasting with physics-informe neural odes. arXiv preprint arXiv:2404. 10024, 2024.
Wang, R., Walters, R., and Yu, R. Meta-learning dynamics forecasting using task inference. Advances in Neural Information Processing Systems, 35:21640-21653, 2022.
Pfaff, T.. Fortunato, M., Sanchez- Bonehez-Gnzzale, A., and Battaglia, P. W. Learning mesh-based simulation with graph networks. arXiv preprint arXiv:2010.03409, 2020.
Wu, Y. and He, K. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp. 3-19, 2018
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmenta- tion. In Medical image computing and computer-asiisted intervention-MICCA. 2015: 18th international confer- ence, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pp. 234-241. Springer, 2015.
Xie, X., Zhou, P., Li, H., Lin, Z., and Yan, S. Adan: Adap- tive nesterov momentum algorithm for faster optimizing deep models. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
Ruthotto, L. and Haber, Deep neural networks motivated by partial differential equations. Journal of Mathematical Imaging and Vision, 62(3):352-366 2020.
Yang, L. and Osher, J. PDE generalization of in-context operator networks: A study on 1d scalar nonlinear con- servation laws. arXiv preprint arXiv:2401.07364 2024.

---

DISCO: learning to DISCover evolution Operator for multi-phyyscs-agnnste prediction
Yang, L., Liu, S., Meng. T., and Osher, S. In-context operator learning with data prompts for differential equa- tion problems. Proceedings of the National Academy of Sciences, 120039):e2310142120, 2023.
Zhang, R., Meng. Q., and Ma, Z.-M. Deciphering and integrating invariants for neural operator learning with various physical mechanisms. National Science Review, 11(4):wad333, 2024.
Zhou, X.-H., Han, J., and Xiao, H. Frame-independent vector-cloud neural network for nonlocal constitutive modeling arbitrary grids. Computer Methods in Ap- plied Mechanics and Engineering, 388:1 14211, 2022.
Zhou, X.H.,, Han, J., Zafar, M. Wolf, E. M., Schrock, C. R., Roy, C. J., and Xiao, H. Neural operator-based super-fidelity: A warm-start approach for accelerating steady-state simulations. Journal of Computational Physics, 529:1 2025.
12

---

DISCO: learning to DISCover evolution Operator for multi-phyyscc--guussic prediction
A. Dataset description
For sake of self-consistency, we provide the description the datasets considered this paper (see Tab. 1) and which originate from two benchmark pupers (Takamoto et al., 2022; Ohana al., 2024), along with their underlying PDEs. The corresponding data is publicly available.
We used 5 datasets from PDEBench, comprising the 2D shallow-water equations, 2D diffusion-reaation, 2D incompressible and compressible Navier-Stokes datasets. These are the datasets used to train MPP (McCabe et al., 2024). In addition, we added the ID Burgers equations dataset in order to showcase our model can be trained on ID and 2D data jointly.
We used 9 datasets from the Well, comprising the 2D active matter, 2D Euler multi-quadrants, 2D Gray-Scott reaction- Bénard, 2D shear flow, 3D turbulence gravity cooling, 3D MHD, 3D Rayleigh-Taylo instability and 3D supernova explosion datasets. We initially excluded some of the datasets present in the Well, such the convective envelope red supergiant dataset, the post neutron star merger dataset, which not stored using cartesian coordinate systems. We provide below the description of the shear flow and Euler qulli--uadrants datasets which were used for specific additional experiments in this paper and refer the reader to (Ohana et al., 2024) for the description of the other datasets.
diffusion, 2D Rayleigh-Bénard,
Examples of trajectories are provided in the rollout Fig. 7, 8, 9 for PDEBench and Fig. 10, 11, 12, 13 for The Well.
A.1. Burgers equations
The Burgers equations model the evolution of ID viscous fluid. They combine nonlinear advection term with linear diffusion term
V
dcu + u, ue = 4
where v is the diffusion coefficient.
The boundary conditions are set to periodic. The dataset is part of PDEBench (Takamoto et al., 2022) and was generated using a temporally and spatially 2nd-order upwind difference scheme for the advection term, and a central difference scheme for the diffusion term. We refer to Fig. for examples of trajectories.
A.2. Shallow water equations
Shallow-water equations, derived from Navier- Stokes equations present suitable framework for modeling free-surface flow problems. They take the form of system of hyperbolic PDEs
0yh + 0gh + 0yhv 0,
1
0, + (, h +
2 gghh?) )) + Syuch - g h00,b
1
0, 0, 59rl")) + O wrt -9r
where h is the water depth, u, v are the velocities in horizontal and vertical direction, is the bathymetry field, and gr describes the gravitational acceleration.
The initial state is circular bump in the center of the domain. The dataset part of PDEBench (Takamoto et al., 2022) and was generated using finite volume solver (Ketcheson al., 2012). We refer Fig. for example trajectory.
A.3. Diffussio reaction equations
2D diffusion-reactioo equation models how a substance spreads and reacts over time, capturing the combined effects of diffusion and chemical or biological reactions in two dimensions
0,u u = D., vx + DuOpyu u + Ra,
0, D2 + D + R..
where Du, Dr are the diffusion coefficients for the activator u and inhibitor u and Ry, R, are the respective reaction functions, which takes the form Ru(u, v) u
n² k and Rs (u, v) = u v where k is a constant. The
13

---

DISCO: learning to DISCover evolution Operator for multi-physscc--annosic prediction
initial states are random Gaussian white noises. The problem involves no-flow Neumann boundary condition, that D 0, D 0, Py PDEBench (Takamoto et al., 2022) and was generated using finite volume method as spatial discretization and fourth-order Runge-Kutte Kutte- method as time discretization. We refer to Fig. 8 for examples of trajectories.
Dr 0 on the edges of the square domain. The dataset is part of
A.4. Incompressible Navier-Stokes
This dataset considers simplification of Navier-Stokes equation that writes
V vv= +(0.++v. Vv) -
where v is the velocity vector field, is the density and u is forcing term and v is constant viscosity.
Initial states and forcing term u are each drawn from isotropic random fields with a certain power-law power-spectrum. The boundary conditions are Dirichlet, imposing the velocity field to be zero at the edges of the square domain. The dataset is part PDEBench (Takamoto et al., 2022) and was generated using a differentiable PDE solver (Holl et al., 2020). We refer Fig. for examples of trajectories.
A.5. Compressible Navier-Stokes
The compressible Davier-Stokes describe the motion of fluid flow
0.p - (pv) = 0
p(0.v + v Vv) --p++ n n (/ -v), 1 (+ + 1 _00") + V ((e + p _ ) v voo') =
where p is the density, is the velocity vector field, p the pressure, p[ 1) is the internal energy, I' 5/3, a' is the viscous stress tensor, and 17, are the shear and bulk viscosity. The boundary conditions are periodic.
The dataset is part of PDEBench (Takamoto et al., 2022) and was generated using 2nd-order HLLC scheme (Toro et al., 1994) with the MUSCL method ( (Van Leer, 1979) for the inviscid part, and the central difference scheme for the viscous part. This dataset contains trajectories with only 5 steps into the future and was used solely for training.
We refer Fig. 9 for an example of trajectory.
A.6. Shear flow
This phenomenon concerns two layers of fluid moving in parallel to each other in opposite directions, which leads various instabilities and turbulence. It is governed by the following incompressible Navier-Stokes equation
du 0t
v Au
Vp = Vu.
where A V V is the spatial Laplacian, with the additional constraints f 0 (pressure gauge). In order to better visualize the shear, we consider passive tracer field governed by the dddvection-diffuiii equation
Ōs
DAs u Vs.
Ot
We also track the vorticity w = V x created by initializing the velocity at different layers of fluid moving in opposite horizontal directions. The fluid equations are parameterized by different viscosity v and tracer diffusivity D.
du dug which measures the local spinning motion of the fluid. The shear is Eu
85
The data was generated using an open-source spectral solver (Burns et al., 2020) with a script that is publicly available. Refer Fig. 12 for an example trajectory.
A.7. Fuler multi-quadrants (compressible)
Euler equations are a simplification of Wavier-Stoker in the absence of viscosity
aU + F(U) + G,) = 0,
14

---

DISCO: learning to DISCover evolution Operator for multi-physsic--agnstt prediction
with
P
pu pu + pwv ule p)
pv pus po +p ule + p)
pu U
F
G=
pv
where p is the density, u, e are the horizontal and vertical velocities, is the pressure and is the energy defined by
P + 200 + c == 7-1
The initial state is piecewise constant signal composed of quadrants, which then evolves in multi-scale shocks. We refer Fig. 10 for example trajectory.
B. Benchmark models hyperparameters
In this work, we compared DISCO with 5 baselines, a U-net (Ronneberger et al., 2015), a Fourier neural operator (Li et al., 2020) implemented using neuralon (Kovachki et al., 2023), GEPS (Koupai et al., 2024) which a meta-learning framework for operator learning, Poseidon (Herde et al., 2024), which is a transformer operator and MPP, which is a transformer based on axial attention (McCabe et al., 2024). We expose here the hyperparameters used in the paper for these models. The number of parameters provided for each model are the ones for the pretraining PDEBench (see Tab. 2) or, if not applicable (U-Net, FNO), are the ones of the models trained on the Fuler multi-quadrants dataset (see Tab. 4).
U-Net. We considered standard U-Net (Ronneberger et al., 2015) with 4 dawn-(and up) sampling blocks, spatial filters of size and initial dimension of 48. The resulting model has 1 17M learnable parameters. Note that this U-Net model is not to be confused with the operator network fo we used in DISCO which is much smaller (around 200k parameters).
FNO. We considered a standard Fourier neural operator (Li et al., 2020) with 4 Fourier blocks, spectral filters of size 16 (number of Fourier modes), and 128 hidden dimensions. The resulting model has 19M learnable parameters.
GEPS. This meta-learing framework (Koupaï et al., 2024) proposes to learn an operator fo by decomposing 0 into shared cavirromment-speiife in their paper). The learnable parameters
parameters across different contexts and context-specific are the shared parameters, the parameters of the linear hypernetwork which process a "code" e for each context, and the codes themselves. A code size is chosen in (Koupai et al., 2024) to be at least the number of parameters the underlying PDE. We chose the codes to be of dimension 12 in our case. This meta-learning strategy is performed on our class of U-Net operator fe which is of size roughly 200k: The total number of trainable parameters in the shared weights and the hypernetwork is around 400k. Given our multi-physsees-anostii task (see Sec. 1) we cannot assume know the specific dataset context is originated from. So that, when trained on collection of datasets, instead of sharing weights per dataset, Aavier-Stokes, one vector for Burgers etc to
which would lead to one vector of shared parameters for compressible comply with the task tackled in this paper and to yield a fair comparison, we chose to share one vector of parameters across all the contexts, thus facing the three levels of variability (1), (2), (3) (see Sec. 1). At inference, all the parameters are frozen, and only the codes are optimized. This optimization is done with a learning rate that tuned on the validation set. We found that the optimal learning rate for optimizing code on single context was 10 for PDEBench and 10 for The Well (we tested three values 10- 10 10-1).
MPP. We considered MPP, which is axial vision transformer (MeCabe et al., 2024) with patch size of 16, 12 attention layers, 12 attention heads per layer, 768 hidden dimensions. The resulting model has 158M learnable parameters. The training
C. Additional details on DISCO architecture
C. OPERATOR NETWORK
The operator network fe in DISCO is a Unet with 4 subsampling layers, starting channel dimension of 8, It first composed of two convolutions with kernel operating on channels. Then followed by the downsampling blocks. the lowest
15

---

DISCO: learning to DISCover evolution Operator for multi-phyysics-anuosic prediction
resolution (the bottleneck), we apply MLP with two layers, going from 128 to 256 channels and going back 128 channels. A downsampling block is composed of first convolution with kernel size 2 and stride 2 followed by GeL.U activation and a convolution of kernel size 3. An upsampling block is first composed of a transpose convolution with kernel of size 2 followed by a GeLU activation and a convolution of kernel size 3. The skip connections between layers of comparable resolution are performed by concatenation. The resulting architecture has roughly 200k parameters.
While most recent U-Net architectures utilize double convolutions as well as ConvNeXt blocks, we could not allow to use them in order to reduce the number of coefficients in the architecture to the minimum so as enforce a bottleneck that is justified in Sec. 3. We also tried adding attention layers in the bottleneck of our U-Net, however, the gain in performance was not worth the additional computational cost.
C.0.2. HYPERNETWORK
The hypernetwork is an axial vision transformer closely following the design of MPP (McCabe et al., 2024) (small version), except that it replaces the decoder by a MLP to generate the parameters of the operator network. It consists of three parts, a CNN encoder with a first field embedding layer, which is a linear transform that adapts to varying number of fields in the input (McCabe et al., 2024). Then, we apply three convolution layers of kernel sizes 4, 2, 2 respectively, with GeLU activation, resulting in an effective patch size of 16. The hidden dimension (token space) is 384. After the encoder, the processor cascades 12 time-space attention blocks, each containing time attention, and axial attentions along each space dimension (McCabe et al., 2024). Each attention block contains 6 heads and uses relative positional encodings. After the attention blocks, we average over time and space to yield a unique token of dimension 384, This token goes through a ML.P with 3 linear layers with intermediate dimensions of 384 and output dimension the dimension of the operator network: roughly 200k. The output, 0 in Eq. 4, is going to populate the different layers of the operator fe. These layers being of different size, their respective weights (0,, 02. ... on Fig. 1) are expected to be of different norm. To enforce this, we apply a signed sigmoid (x) 20(xx) 1 with the following normalization:
0: = siXosi (O.((NAA)
(5)
where 0 is the un-normalized output of the MLP, i is the layer index in the operator network, N, is the norm that the weights should have at initialization (set by PyTorch) for this layer, and À is a factor that was arbitrarily set to 2 for allowing the training to gradually increase the norm of while constraining it to avoid instabilities.
D. Training details
Hy dyperparameters. For joint training multiple datasets, considered the same hyperparameters than McCabe et al. (2024), with batch size of 8, gradient accumulation every 5 batches, epoch size of 2000 batches ("fake epochs"). For single model trainings, we considered batch size of 32 with no gradient accumulation. In certain cases where the optimization was unstable, in particular, when we tried using only intermediate number of steps (see Tab. 7), we used gradient clipping, clipping the total norm of the gradients to a default norm of 1.0. The data is split into train, validation, and test sets with an 80%, 10%, and 10% division, respectively.
Optimization. All trainings were performed using the adaptive Nesterov optimizer (Xie al., 2024) and cosine schedule for the learning rate. Using AdamW optimizer with varying learning rate did not improve overall performance in the cases we tested. The large-scale experiments are run for a fixed number of 300 epochs. We used a weight decay of 0.001 and drop path of 0.1.
Loss. A normalized root mean square error (NRMSE) used for both monitoring the training of the model and assessing the performances in this paper. For two tensors u (target) and u,ạ1 (prediction) with C' channels
C
||ue
0
Loss(u, uu a) = C
(6)
||u||2 +
where the l norm || |12 is averaged along space and e = 10-7 For a batch of data, this loss is simply averaged.
7 is small number added to prevent numerical instabilities.
Software. The model trainings were conducted using python 13.1 and the PyTorch library v2.4.1 (Paszke et al., 2019).
16

---

DISCO: learning to DISCover evolution Operator for multi-phyysics-agnosii prediction
Table Number epochs reach SOTA performance on PDEBench datasets.
Model
# parameters
Burgers
SWE
DiffReID
INS
CNS
CNS (alone)
MPP (retrained)
160m
500
500
500
500
500
50
DISCO (ours)
119m
277
70
55
35
Table Rollout performance of models trained on PDEBench datasets. models are trained for 300 epochs. We use the NRMSE (lower better). DISCO which relies on learning operator explicidy, without having learn decoder, achieves better performance on next step prediction (t + 1). Compared Poseidon (Herde et al., 2024), which is trained on multiple future stop prediction, our model, trained only the step (L + 1) remains competitive. Best performance each dataset in bold, second best is underscored.
No RETRAINING
DECODER FREE
MODEL
TEST DATASET
NRMSE
INFERENCE
++1
+ 4
+ 8
f 16
+ 32
BURGERS
0.508
0.777
0.660
0.629
0.567
SWE
0,0720
0.139
0.238
0.525
0.860
GEPS
DIFFRE2D
0.927
>
>
>1
>1
INS
0.794
0.947
0.986
0.987
0.986
CNS
0.604
0.784
BURGERS
SWE
0.00207
0.00207
0.00237
0.00331
0.0380
POSEIDON
X
X
DIFFRE2D
0.0141
0.0154
0.0227
0.0379
0.0821
INS
0.0250
0.0318
0.0479
0.0869
0.211
CNS
0.135
0.219
BURGERS
0.00556
0.0242
0.0526
0.107
0.191
SWE
0.00145
0.00330
0.00592
0.0112
0.0327
MPP
X
DIFFRE2D
0.0109
0.0283
0.0535
0.109
0.238
INS
0.00936
0.0274
0.0557
0.132
0.379
CNS
0.0349
0.0748
BURGERS
0.00347
0.0130
0.0298
0.0686
0.163
SWE
0.000648
0.00226
0.00384
0.00554
0.00951
DISCO (OURS)
/
DIFFRE2D
0.00448
0.0142
0.0291
0.0501
0.0832
INS
0.00347
0.0139
0.0277
0.0556
0.113
CNS
0.107
0.208
Hardware.
All model trainings were conducted using Distributed Data Parallel across Nvidia H100-80Gb GPUs.
E. Additional experiments
This section contains additional experiments on the models used in this paper.
Translation equivariance standard transformers. As concrete example the issues encountered by a transformer- based architecture, which does not naturally encode translation equivariance, let us consider shifting all the states in a context, as well the following state, by a vector v of increasing norm ||w|]. Fig. 6 shows that the transformer architecture quickly struggles predict the next state, although it competitive when no translation is applied, ||l|| 0.
Additional rollout examples. Figs. 7, 8, 9, show additional rollout trajectories for the PDEBench datasets. Figs. 10, 11, 12, 13 are showing rollout trajectories for The Well datasets.
17

---

DISCO: learning to DISCover evolution Operator for mulii-phyyscc--guussic prediction
Table Inftuence of the model size the accuracy, the diffusion-rcaction dataset from PDEBench. Each table corresponds respectively to variations over a and 0, Gray rows indicate the default values chosen the paper.
hyperparameters
num. weights a
NRMSE
num. layers
num. heads
dim
4
3
192
65m
0.00211
12
6
384
101m
0.00133
12
12
768
280m
0.00203
hyperparameter Cstart
num.
parameters 0
NRMSE
41k
0.00330
8
158k
0.00133
12
353k
0.000971
validation NRMSE
validation NRMSE
validltion NRMSE
DISCO loun)
11
AM
10
epoch
epocl
opeahh
(a) Active matter
(b) Euler
Gray-Scott
validation NRMSE
validation N NRMSK
validation NRMSE
10-1
10
20
50
spocl:
<pocch
epoch
(f) Turbulence gravity cool-
Rayleigh-Bémee Bénnad
(c) Shear flow
ing
validation NRMSE
validation NRMSE
validation NRMSE
Coiem
10"
10
20
30
00
10
spocl
spoth
spocc
(h) Rayleigh-Taylee insta- bility
(g) MHD
(i) Supernova explosion
Figure 5. Learning curves The Well datasets. DISCO achieves better performance next step prediction compared to a MPP model retrained from scratch.
NRMSE f 111
0.3.2
0.224
<h1ft
Figure 6. Performance on the shearflow dataset when context and target are shifted by e E IR? While the transformer performs well 0, it declines than DISCO under shifts due the lack of inductive bias for translation equivariance.
18

---

DISCO: learning to DISCover evolution Operator for multi-phyysics--gnnsti prediction
++1
t + 4
t ++8
+ 16
TRUTH
GEPS
MPP
DISCO
t
t +
1+ 4
t +
1 + 16
TRUTH
GEPS
Poseidon
MPP
DISCO
Figure 7. Examples of rollout trajectory for the Burgers (top) and shallow-water (boltom) datasets from PDEBench.
19

---

DISCO: learning to DISCover evolution Operator for multi-phyysics-aggnsii prediction
+
8
t + 16
GEPS
Poseidin
MPP
DISCO
t+1
++4
++8
+ 16
TRUTH
GEPS
Poseidion
MPP
DISCO
Figure 8. Examples of rollout trajectory for the diffusion-reaace (top) and incompressible Navier-Stokes datasets from PDEBench.
20

---

DISCO: learning to DISCover evolution Operator for muliipphysics-agnont prediction
t +
t ++3 2
++3 3
t+4
TRUTH
GEPS
Poseidon
MPP
DISCO
Figure 9. Example of rollout trajectory for the compressible Navier-Stokes dataset from PDEBench.
21

---

DISCO: learning to DISCover evolution Operator for mulii-phyyics-aguusiic prediction
t
t + 1
t + 4
t + 8
t + 16
TRUTH
GEPS
MPP
DISCO
t t+1 1
++4
++8
t + 16
TRUTH
GEPS
MPP
DISCO
Figure 10. Examples of rollout trajectory for the active matter (top) and euler (bottom) datasets from the Well.
22

---

DISCO: learning to DISCover evolution Operator for multi-physics-agnoste prediction
t+1
t +
++8
+ 16
TRUTH
GEPS
MPP
DISCO
ISCO MPP GEPSTRUTH
t+1
t +
t ++8
++ 16
Figure 11. Examples of rollout trajectory for the Gray-Scot reaction-diffusion (top) and Rayleigh- Bénard (bottom) datasets from the
Well.
23

---

DISCO: learning to DISCover evolution Operator for multi-phyyscs--anustic prediction
t +
t + 4
t + 8
t + 16
TRUTH
GEPS
MPP
DISCO
t + 1
t + 4
t + 8
t + 16
TRUTH
GEPS
MPP
DISCO
Figure 12. Examples of rollout trajectory for the shear flow (top) and MHD (bottom) datasets from the Well.
24

---

DISCO: learning to DISCover evolution Operator for mulli-phyysic--agnssic prediction
t
+
t +
+
+ 16
TRUTH
GEPS
MPP
DISCO
++1 1
++4
++8
+ 16
TRUTH
GEPS
MPP
DISCO
Figure Examples rollout trajectory (2D slices) Rayleigh Taylor instability (top) supernova explosion (bottom) datasets from the Well.
25


Reviews:
[{"review_id": "me5izMbIi8", "reviewer": "Reviewer_Fp3h", "summary": "The paper introduces DISCO, a novel framework for multi-physics-agnostic prediction that combines transformer-based hypernetworks with neural PDE solvers. The key innovation is a two-stage approach where a large transformer hypernetwork processes a context of sequential states to generate parameters for a smaller operator network, which then predicts future states through time integration. This architecture decouples dynamics estimation from state prediction, creating an information bottleneck that helps the model focus on essential dynamics rather than memorizing specific trajectories. The authors demonstrate that DISCO achieves state-of-the-art performance on benchmark datasets (PDEBench and The Well) while requiring significantly fewer training epochs than previous approaches.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive comments on our work, particularly for acknowledging the “particularly creative aspect [of] the architecture's information bottleneck design” and noting that “the claims in the paper are well supported by substantial evidence”. We also appreciate their valuable suggestions and questions.\n\nHere are our answers to your questions \n1. At inference time, on the Well datasets (see Table 1 in our paper), DISCO uses around two times fewer FLOPs than a transformer like MPP ([Mccabe et al., 2024](https://arxiv.org/abs/2310.02994)) to predict the next step $t+1$ (MPP: 4.0 GFLOPs, DISCO: 2.1 GFLOPs on a single context). Additionally, note that when predicting a large time horizon (e.g., $32$ as for this paper), one can use DISCO’s hypernetwork (here, a transformer) to estimate the parameters meta-learned only once (one forward pass of hypernetwork), fix them and apply an integration scheme on the long range-prediction $t+32$. On the contrary, MPP’s transformer, including its encoder and decoder, needs to be applied at each iteration. As a result, MPP requires significantly more FLOPs than DISCO: 128 GFLOPs vs. 37 GFLOPs for a horizon of $t+32$ from time step $t$.\n2. No, without any additional modification, we do not expect DISCO to generalize to new classes of PDEs without fine-tuning (i.e., zero-shot generalization). This is a challenging task, and to our knowledge, no existing method designed for multi-physics agnostic prediction has demonstrated this capability. For comparison, both Poseidon ([Herde et al., 2024](https://proceedings.neurips.cc/paper_files/paper/2024/file/84e1b1ec17bb11c57234e96433022a9a-Paper-Conference.pdf)) and GEPS ([Koupaï et al., 2024](https://proceedings.neurips.cc/paper_files/paper/2024/file/82844e428d9163a9f94830dc03af4f9c-Paper-Conference.pdf)) require fine-tuning by design, even when applied to the same PDE. MPP ([Mccabe et al., 2024](https://arxiv.org/abs/2310.02994)) does not claim zero-shot generalization, and the only zero-shot result shown (Fig. 1) is not compelling.\n3. This is a good suggestion. Given that we obtain better estimates compared to baselines such as MPP ([Mccabe et al., 2024](https://arxiv.org/abs/2310.02994)), our model by design better preserves the conservation laws, which is not surprising (see figure [here](https://postimg.cc/0KvhbXpV)). Since we use an autoregressive model that does not explicitly enforce conservation laws,\nthese laws will gradually become less respected over time.\nIncorporating these conservation laws into our training, as done in PINNs ([Cai et al., 2021](https://arxiv.org/pdf/2105.09506)), would require knowing the specific conservation laws for each context. Such an assumption is incompatible with the multi-physics agnostic prediction task addressed in our paper. \\\nMore details on the [figure](https://postimg.cc/0KvhbXpV): it shows the conservation of mass $\\||\\mathrm{div} u\\||$ and momentum  $\\||\\partial_t u - \\nu\\Delta u + u\\cdot\\nabla u + \\nabla p\\||$ over model rollouts from $t+1$ to $t+64$ averaged on $32$ trajectories from the validation set of the shear flow dataset (incompressible fluid, see Table 1. in the paper). The dashed line represents the conservation law as satisfied in the data. DISCO exhibits smaller deviations from the conservation laws compared to MPP ([Mccabe et al., 2024](https://arxiv.org/abs/2310.02994)).\n\nHere are comments on other points you raised. \n- The suggested reference to MeshfreeFlowNet ([Jiang, Esmaeilzadeh et al., 2020](https://arxiv.org/pdf/2005.01463)), which features a Rayleigh-Bénard convection dataset, very similar to one of the datasets used in our paper, will be added to the “Related Works” section (line 100).\n- “Concrete real world use-cases”: We completely agree with the reviewer on the importance of getting closer to real-world applications. One specific application we are currently exploring is the evolution of an unknown Physical system given limited observational data. In this context, DISCO provides an operator space (see Fig. 3 for a visualization), where all evolution operators encountered during training are mapped. When presented with data from an unseen physical system, one can identify the \"closest\" known operators and leverage them to refine and adapt the model, improving predictions on the new system.", "claims_and_evidence": "The claims in the paper are well supported by substantial evidence from the extensive experimental evaluation. The authors provide detailed empirical results across two comprehensive datasets, with clear performance metrics for both next-step prediction and multi-step rollouts. Particularly compelling is the quantitative demonstration that DISCO achieves state-of-the-art performance with significantly fewer training epochs than transformer-based approaches. The paper also provides convincing evidence for generalization capabilities through visualization of the parameter space and fine-tuning experiments on unseen physics.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are highly appropriate for the multi-physics-agnostic prediction problem. The authors thoughtfully selected benchmark datasets (PDEBench and The Well) that span diverse physical phenomena across different spatial dimensions, providing a comprehensive testbed for their approach. Their evaluation metrics, particularly the normalized root mean square error (NRMSE) for both single-step and multi-step predictions, effectively capture model performance in practical settings. The comparison against multiple strong baselines (including transformers and meta-learning frameworks) strengthens the evaluation framework.", "theoretical_claims": "The paper is primarily empirical in nature and does not present formal mathematical proofs for theoretical claims. The authors establish conceptual connections between their approach and classical numerical methods for solving PDEs, particularly linking their operator network to finite difference schemes, but these are presented as motivational insights rather than rigorous theoretical proofs.", "experimental_designs_or_analyses": "I examined the experimental designs and analyses in the paper and found them to be generally sound. The authors use appropriate datasets (PDEBench and The Well) that represent diverse physical systems, ensuring comprehensive evaluation.", "relation_to_broader_scientific_literature": "The key contributions of DISCO relate to several important research directions in the scientific literature. First, it builds upon recent advances in transformer-based models for physical system modeling (Yang et al., 2023; Liu et al., 2023; McCabe et al., 2024), but addresses limitations in their training efficiency and ability to preserve physical invariances. Second, DISCO connects to neural operator learning approaches (Li et al., 2020; Kovachki et al., 2023) while extending their capabilities to handle unknown and variable dynamics.", "essential_references_not_discussed": "The authors mention neural PDE solvers but don't reference Jiang et al.'s \"MeshfreeFlowNet: A Physics-Constrained Deep Continuous Space-Time Super-Resolution Framework\".", "other_strengths_and_weaknesses": "Pros:\n\nA particularly creative aspect is the architecture's information bottleneck design, which forces the model to learn generalizable physical laws rather than memorizing specific trajectories. This design choice shows originality in addressing overfitting challenges unique to physical system modeling.\n\nCons:\n\nThough the authors mention potential applications, concrete real-world use cases would strengthen the paper's impact", "other_comments_or_suggestions": "I did not identify any significant typos or grammatical errors in the paper", "questions_for_authors": "1. Given that DISCO uses a time integration method for prediction, how does the computational cost of inference compare to direct prediction methods like standard transformers?\n\n2. The paper demonstrates impressive generalization to unseen physics when fine-tuning, but could you clarify whether DISCO can generalize to significantly different classes of PDEs (e.g., from diffusion-dominated to advection-dominated systems) without any fine-tuning?\n\n3. The paper mentions that DISCO achieves better numerical accuracy on multi-step rollouts compared to baselines, but have you analyzed the model's ability to preserve important physical properties like conservation laws or symmetries during long rollouts", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "N9AjSIXLTG", "reviewer": "Reviewer_2f9m", "summary": "This paper proposes a novel method to obtain lightweight surrogate models from physics data. The idea is to use a hypernetwork transformer to learn the parameters of a smaller operator network, which is in charge of performing the time integration. This architecture decouples the learning of the dynamics from the state prediction, which is convenient for generalization to other domains and fine-tuning. The model is tested together with other baselines in two public benchmark datasets: PDEBench and The Well, showing state-of-the-art performance.\n\n## Update after rebuttal\nThe authors addressed all my concerns satisfactorily, and I raised the score to 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback, in particular for acknowledging that “the methods of this paper are clearly written and developed” and that “all the claims of the paper are supported with validation results in benchmark cases”. We also thank the reviewer for their valuable suggestions and for pointing out several typos.\n\nHere are our answers to your questions:\n- We did try using a FNO ([Li et al., 2020](https://arxiv.org/abs/2010.08895)) for our operator network, as well as a U-Net with transformer blocks, but were unable to achieve better results while maintaining the same operator size (~$200$k parameters). To keep the FNO small, we had to reduce the hidden dimension, typically to $32$, and the number of Fourier modes in the spectral convolution layers, typically to $4$ per dimension, which limited the expressiveness of these architectures. A U-Net with transformer blocks faces similar challenges. For a fixed budget of ~$200$k parameters, allocating weights to the attention layers in the skip connections required reducing the hidden dimension, which in turn constrained the model’s expressiveness.\n- We chose not to include rollouts beyond $t+16$ / $t+32$ because neither DISCO nor any baseline models produced satisfactory results at those horizons. For $t+64$, most NRMSE values exceeded $1$, making comparisons difficult. That said, we agree with the reviewer that evaluating longer rollouts is important. Our results demonstrate that DISCO outperforms the main baselines GEPS ([Koupaï et al., 2024)](https://proceedings.neurips.cc/paper_files/paper/2024/file/82844e428d9163a9f94830dc03af4f9c-Paper-Conference.pdf)), Poseidon ([Herde et al., 2024](https://proceedings.neurips.cc/paper_files/paper/2024/file/84e1b1ec17bb11c57234e96433022a9a-Paper-Conference.pdf)), MPP ([McCabe et al. 2024](https://arxiv.org/abs/2310.02994)) in multi-step rollouts. Notably, the state-of-the-art model (MPP) does not report any metrics for multi-step predictions.\\\nWhile tackling very long rollouts is not the goal of our work, we highlight in the conclusion that DISCO can naturally incorporate multi-step rollouts during training. Specifically, one can use DISCO’s hypernetwork (here, a transformer) to estimate the parameters meta-learned only once (one forward pass of hypernetwork), fix them and apply an integration scheme on a longer horizon than just $t+1$. Since the operator network is small, it will have a significantly lower memory usage than applying a large transformer such as MPP at each step in the future.\n- Regarding the last part of your second question, Fig. 11 (Appendix E) in our paper shows that even when trained on predicting the next step $t+1$, our model predicts the position of convection cells (mushroom-like spatial structures) quite decently after $t+8$, $t+16$. This is notable given that the future locations of these structures are highly sensitive to small perturbations in the state at $t$. This task is recognized as particularly challenging in the dataset’s paper (see “Sensitivity to initial conditions” in Appendix D in [Ohana et al., 2024](https://proceedings.neurips.cc/paper_files/paper/2024/file/4f9a5acd91ac76569f2fe291b1f4772b-Paper-Datasets_and_Benchmarks_Track.pdf)). Here are two additional examples of rollouts on the same dataset: [figure1](https://postimg.cc/64d79gYY), [figure2](https://postimg.cc/1V52GCzM). A model overfitting the next step prediction would have a hard time predicting correctly this behavior after $16$ steps. \n\nHere are some additional comments.\n- We appreciate that you have read some of the appendices (specifically, Appendix C as you write in your review), yet you also wrote “There is no supplementary material”. You should know that we have several appendices and that we also provided the code as a .zip file.\n- We will add the two relevant references you mentioned to our \"Related works\" section, but note the major differences: [Alesiani et al., 2022](https://ml4physicalsciences.github.io/2022/files/NeurIPS_ML4PS_2022_89.pdf) assumes the coefficients of the PDE are known and build a hypernetwork that take these coefficients as input, while we only have access to a short context of past states, making the prediction more challenging since our hypernetwork must infer the time evolution. [Lee et al., 2023](https://arxiv.org/abs/2312.15949) indeed uses meta-learning, but on fundamentally different objects and for different purposes. The authors employ a hypernetwork to predict the the parametrization of a future state, in the form of a network which takes the 2d coordinates and returns the value of the state at this location. In particular, such hypernetwork needs to be retrained every time the prediction horizon, the PDE coefficients, or the PDE class are changed. In essence, their hypernetwork predicts a state while on our task, our hypernetwork predicts an actual evolution operator.\n- The minor typos you pointed out will be corrected in the manuscript.", "claims_and_evidence": "All the claims of the paper are supported with validation results in benchmark cases. The method is novel and the results of the presented architecture clearly outperforms concurrent work baselines. I agree with the authors that the key feature of the method relies on the decoder-free structure (lines 290-294), which avoids the challenging (and sometimes, impossible) task of learning an inverse mapping.", "methods_and_evaluation_criteria": "The methods of this paper are clearly written and developed. The model uses a standard U-Net + Axial Vision Transformer architecture which is detailed in Appendix C. The output of the transformer is post-processed such that the U-Net parameters are conveniently normalized. The baseline networks are chosen to be recent state-of-the-art works such as GEPS or MPP. The evaluation criteria follows standard practices in physics-informed machine learning.\n\nThe rollout metrics in Table 3 and 6 are only integrated up to 16/32 timesteps. In my opinion, this is insufficient, as the datasets usually contain several hundreds of snapshots. It is a standard practise in the physics-informed deep learning field to do single-step supervision, but the test is always performed as a rollout from the initial conditions to several times the training time horizon. Single-step prediction error is informative, but useless when trying to evaluate real long-term predictions.", "theoretical_claims": "There are no thoretical proofs in the paper. All the claims are validated experimentally as a pure data-driven procedure.", "experimental_designs_or_analyses": "The experiments are based on two public benchmarks of challenging physics problems. They provide tests of multiple physics training, and finetuning on unseen physics. There is an additional visualization of the parameter space using Umap which shows that similar initial conditions induces similar latent parameters.", "relation_to_broader_scientific_literature": "The literature is correctly discussed, appart from some missing references in regard to hyper-networks for PDEs (see next section).", "essential_references_not_discussed": "It is mentioned in the paper that current literature is limited to known PDE equations or limiting adaptation to only the first layer (lines 138-145). However, one can found several examples in the literature about complete parametrized neural operators in the physics context. For example, a couple involving DeepONets and FNOs:\n\n* [Lee, 2023] HyperDeepONet: learning operator with complex target function space using the limited resources via hypernetwork\n* [Alesiani, 2022] HyperFNO: Improving the Generalization Behavior of Fourier Neural Operators", "other_strengths_and_weaknesses": "I have no more comments.", "other_comments_or_suggestions": "I have found some minor typos:\n* Line 13: \"unkown\" might refer to \"unknown\".\n* Line 233: \"128=8.2^4\" might refer to \"128=8·2^4\"\n* Line 238: Space is missing in \"1D,2D\"\n* Line 270-271: Spaces are missing in \"1D,2D,3D\".\n* Lines 291-292: \"we learn the to output the weights\" might refer to \"we learn to output the weights\".\n* Line 294: \"cmopare\" might refer to \"compare\"\n* Table 3 caption: Spaces missing in \"10,11,12,13\".\n* Lines 620-621: Spaces missing in \"7,8,9\" and \"10,11,12,13\".\n* Line 760: Spaces missing in \"(i),(ii),(iii)\"\n* Line 894: Spaces missing in \"7,8,9\".\n* Figure 8: There is a number \"20\" on the third DISCO result image.", "questions_for_authors": "* Given that the operator network is relatively small, how does it compare to other NO methods, such as FNOs in terms of expresiveness? Have the authors tried to use a more expressive model?\n\n* Table 3 and 6: Have the authors tried to rollout to further timesteps than t+16/t+32? It is very relevant to see how robust is the network and integration scheme over longer rollouts. A network can be overfit to match perfect single-step predictions but might not be robust over integration errors.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "I have no ethical concerns."}}, {"review_id": "w2RjuiWYgt", "reviewer": "Reviewer_uxnz", "summary": "**Summary after rebuttal**\n\nThe authors resolved all of my concerns. I am raising to 4.\n\n**End of Summary after rebuttal**\n\nThe paper introduces DISCO, a novel framework for multi-physics-agnostic prediction of dynamical systems governed by unknown temporal partial differential equations (PDEs). The key contribution is the use of a transformer-based hypernetwork to generate parameters for a smaller operator network, which then predicts the next state through time integration. This approach decouples the estimation of dynamics from state prediction, offering a more efficient and interpretable solution compared to traditional methods.\n\nThe only issues I have is the writing:\n\n- There are some section mis-arrangements: eg, the network structure is introduced in the experimental results section, which should have been in the method. another example is the results section composites of both subsection and paragraph, which is chaotic for reading. The authors can consider only keeping the paragraphs, while moving the not so important sections into appendix.\n\n- Undefined Terms/Symbols: Terms like d1 and d2 (dimensions of hypernet and operator net) are sparsely mentioned and not clearly defined at their first appearance, making it harder for readers to find. The authors should define all symbols at their first appearance and maintain consistency throughout.\n\n- Clarity: The paper uses lots of technical jargon (e.g., \"spatial translation equivariance\") without sufficient explanation, which may hinder readability for a broader audience. I personally do not like this style of writing, and appreciate more clear writing and only keeping core concepts related to the contribution of this paper.\n\n- Macro Formatting: The method name \"DISCO\" sometimes lacks a space before the following content. This could be resolved by adding a tilde (~) in LaTeX when using macros (e.g., \"\\macro~\").\n\nBesides the writing issue, the related work could be strengthened a bit. I will recommend a few citations to add in the below section of this review.\n\nOverall, I give boarderline at this stage. The resolving of above problems will increase my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough reading of our manuscript and for their valuable feedback on the writing and related works, which will help improve the paper.\n\nSince we cannot upload an updated manuscript for ICML, we will describe the changes we intend to make in response to your review below. \n\nOverall writing:\n- We will remove the \"misarrangements\" you mention. The \"operator network architecture\" and \"hypernetwork architecture\" paragraphs will be moved to subsection 3.2 to make this section, which introduces our model, more self-contained. As a result, subsection 4.1, \"Generic architectures\" will be removed. We agree that having only paragraphs in section 4 is beneficial for improving the readability of the paper. With small changes, section 4 will include the following paragraphs: \"Multiple physics training\", which briefly explains how training on multiple physics simultaneously is performed; \"Datasets considered\"; \"Baselines\"; \"Next steps prediction performances\"; \"Information bottleneck and operator space\"; \"Fine-tuning on unseen Physics\"; and \"Model size ablation.\"\n- Undefined Terms/Symbols: The terms $d_1$ and $d_2$ will be defined more clearly and called “operator size” and “hypernetwork size” respectively throughout the paper. In particular, line 165 on the right column will be clarified: “where $\\alpha\\in \\mathbb{R}^{d_2}$ are learnable parameters, while $\\theta\\in\\mathbb{R}^{d_2}$ are parameters predicted by $\\psi$, with $d_1$ and $d_2$ being the sizes of the operator network and hypernetwork respectively”.\n- Clarity: we think the term “translation equivariance”, often used in ML, is a well identified property, satisfied by many PDEs of interest (the right hand-side in Eq. (1) in the manuscript is translation equivariant). This property is reflected in the use of a translation equivariant U-Net as the operator network.  However, to clarify the meaning of this term we added a sentence when first introduced, line 076 “These methods often preserve key structural properties of physics, such as continuous-time evolution and translation equivariance [(Mallat, 1999)](https://www.sciencedirect.com/book/9780123743701/a-wavelet-tour-of-signal-processing) (i.e., a spatial translation of the initial condition results in the same translation of the solution, in the absence of boundary conditions), which transformers do not naturally inherit”. \n- Macro-formatting: these have been resolved.\n\nRelated works. We thank the reviewer for the several relevant references!\n- The references to PDE-Net ([Long et al., 2021](https://proceedings.mlr.press/v80/long18a/long18a.pdf)) and [Kochkov, Smith et al. (2021)](https://www.pnas.org/doi/pdf/10.1073/pnas.2101784118) will be added line 99, as well as line 184 left column.\n- We will add the following sentence line 425 in the conclusion. “In particular, several papers, such as [Lino et al. (2022)](https://pubs.aip.org/aip/pof/article/34/8/087110/2847850) and [Cao et al. (2023)](https://proceedings.mlr.press/v202/cao23a.html), propose U-Net-like graph neural network architectures, which are natural candidates for the class of operator networks in DISCO. There are other promising directions ...”. \nTo remain consistent and focused on the subject of our paper we will not include PointNet ([Qi, Su, et al. 2017](https://openaccess.thecvf.com/content_cvpr_2017/papers/Qi_PointNet_Deep_Learning_CVPR_2017_paper.pdf)) or PointNet++ ([Qi et al. 2017](https://proceedings.neurips.cc/paper_files/paper/2017/file/d8bf84be3800d12f74d8b05e9b89836f-Paper.pdf)), although we recognize them as important references.", "claims_and_evidence": "DISCO achieves state-of-the-art performance on next-step prediction across multiple physical systems in PDEBench.\nEvidence: Table 2 shows DISCO outperforms MPP (McCabe et al., 2024) on most datasets with fewer epochs. For example, DISCO achieves NRMSE of 0.0027 on Burgers, compared to MPP's 0.0029.\n\nDISCO generalizes well to unseen physics and initial conditions.\nEvidence: Fine-tuning experiments on the Euler dataset show DISCO outperforms other models (Table 4). DISCO achieves NRMSE of 0.029, compared to 0.032 for MPP and 0.36 for GEPS.", "methods_and_evaluation_criteria": "DISCO uses a transformer-based hypernetwork to generate parameters for a smaller operator network (U-Net). The operator network is integrated over time to predict the next state. The hypernetwork processes a context of successive states to infer the governing dynamics.\n\n Normalized Root Mean Square Error (NRMSE). \nThe paper evaluates performance on next-step prediction and multi-step rollouts across multiple datasets (PDEBench and The Well).", "theoretical_claims": "The paper does not make strong theoretical claims but focuses on empirical performance. The theoretical justification for the use of hypernetworks and operator networks is grounded in classical numerical methods and finite difference schemes, which makes sense in my mind.", "experimental_designs_or_analyses": "The paper evaluates DISCO on two collections of datasets: PDEBench (5 datasets) and The Well (9 datasets). These datasets cover a wide range of physical systems, including fluid dynamics, reaction-diffusion, and astrophysics.\n\n The paper compares DISCO against several baselines, including MPP (McCabe et al., 2024), Poseidon (Herde et al., 2024), and GEPS (Koupai et al., 2024). DISCO consistently outperforms these baselines in terms of accuracy and training efficiency.", "relation_to_broader_scientific_literature": "The paper situates itself within the broader literature on neural PDE solvers and meta-learning for dynamical systems. It builds on recent work in transformer-based models for PDEs (e.g., MPP, Poseidon) and extends these approaches by introducing a hypernetwork to generate operator parameters.", "essential_references_not_discussed": "The paper could benefit from discussing related work on using ML to learn stencils in PDE or CFD, as well as recent advances in graph neural networks (GNNs) for unstructured grids. Specifically:\n\na) Papers related to using ML to learn stencils in PDE or CFD; these works are highly similar to the current work, but for the meta-learning part:\n- PDE-Net: Learning PDEs from Data\n- Machine learning–accelerated computational fluid dynamics\n- etc, authors can find much more by using citations of above\n\n\nb) The authors mention MeshGraphNet for unstructured grids but could also discuss \"Unet\" for GNNs, similar to their using \"Unet\", instead of stacking CNN here:\n- Multi-scale rotation-equivariant graph neural networks for unsteady Eulerian fluid dynamics\n- Efficient Learning of Mesh-Based Physical Simulation with Bi-Stride Multi-Scale Graph Neural Network\n- Learning Distributions of Complex Fluid Simulations with Diffusion Graph Networks\n- etc, \n\nc) Consider adding PointNet and PointNet++ (\"Unet\" ver of PointNet) as they can be applied to Lagrangian view simulations:\n- PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation\n- PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space\n- etc,", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "NA"}}]

Author rebuttal:

Meta-review: The paper introduces a two stage framework to disentangle dynamics emulation with state prediction through a hypernetwork transformer for the latter and a lightweight neural DE for the former. 

Strengths identified:
* Novel idea and new contribution to the literature in NN based emulation of physical systems
* Comprehensive evaluation on well-known PDE benchmarks with multiple physics and demonstration of good performance to unseen physics

Weaknesses:
* Real-world applications is limited but possibly acceptable for this paper given its contributions are on the methodology side with sufficient evaluation on the benchmarks
* Some concerns on long rollout performance - authors acknowledged this as part of future work
* Writing concerns largely clarified in the rebuttal process.

All reviews concur in the contributions. The authors should carefully revise the writing in the paper as clarified in the rebuttal so that their paper is accessible and understandable to the community.
Decision: Accept (poster)