DISCO: learning to DISCover an evolution Operator for multi-physics-agnostic prediction

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.

Since we cannot upload an updated manuscript for ICML, we will describe the changes we intend to make in response to your review below.

Overall writing:

• 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."
• 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”.
• 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) (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”.
• Macro-formatting: these have been resolved.

Related works. We thank the reviewer for the several relevant references!

• The references to PDE-Net (Long et al., 2021) and Kochkov, Smith et al. (2021) will be added line 99, as well as line 184 left column.
• We will add the following sentence line 425 in the conclusion. “In particular, several papers, such as Lino et al. (2022) and Cao et al. (2023), 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 ...”. To remain consistent and focused on the subject of our paper we will not include PointNet (Qi, Su, et al. 2017) or PointNet++ (Qi et al. 2017), although we recognize them as important references.

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.

Here are our answers to your questions:

• We did try using a FNO (Li et al., 2020) 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.
• 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)), Poseidon (Herde et al., 2024), MPP (McCabe et al. 2024) in multi-step rollouts. Notably, the state-of-the-art model (MPP) does not report any metrics for multi-step predictions.
  While 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.
• 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). Here are two additional examples of rollouts on the same dataset: figure1, figure2. A model overfitting the next step prediction would have a hard time predicting correctly this behavior after $16$ steps.

Here are some additional comments.

• 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.
• We will add the two relevant references you mentioned to our "Related works" section, but note the major differences: Alesiani et al., 2022 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 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.
• The minor typos you pointed out will be corrected in the manuscript.

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.

Here are our answers to your questions

1. 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) 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$.
2. 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) and GEPS (Koupaï et al., 2024) require fine-tuning by design, even when applied to the same PDE. MPP (Mccabe et al., 2024) does not claim zero-shot generalization, and the only zero-shot result shown (Fig. 1) is not compelling.
3. This is a good suggestion. Given that we obtain better estimates compared to baselines such as MPP (Mccabe et al., 2024), our model by design better preserves the conservation laws, which is not surprising (see figure here). Since we use an autoregressive model that does not explicitly enforce conservation laws, these laws will gradually become less respected over time. Incorporating these conservation laws into our training, as done in PINNs (Cai et al., 2021), 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.
   More details on the figure: 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).

Here are comments on other points you raised.

• The suggested reference to MeshfreeFlowNet (Jiang, Esmaeilzadeh et al., 2020), 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).
• “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.