Boosting Generalization in Parametric PDE Neural Solvers through Adaptive Conditioning

We're thankful for the reviewer's feedback and have addressed the raised concerns below.

The main architecture of the proposed model is similar to LoRA.

The paper primarily focuses on the general adaptation framework rather than the LoRA implementation. The key points are:

• Classical ERM training is inadequate for modeling the variability in dynamical physical processes (framed as parametric PDEs), as demonstrated in section 3.2.
• An alternative inductive principle is necessary; we propose an adaptive conditioning principle: learning from multiple environments to rapidly adapt to a new one (section 3.3). The core issue is developing an effective meta-learning strategy.
• Our implementation of this principle is computationally efficient (fast adaptation to new environments) and scalable w.r.t. the number of training samples, via low-rank and first-order optimization.
• We propose a versatile framework compatible with any NN backbone and demonstrate that physical models can be learned across different environments despite incomplete models and unknown parameters, a novel achievement.
• It improves the state-of-the-art in learning from multiple environments on new proposed datasets. Experiments show our framework works well with MLP, CNN, and FNO backbones.

The motivation of the model design is not well established as well.

The proposed model combines a physics module and an agnostic NN module. The former encodes some prior knowledge available under the form of a PDE that only partly explains the underlying physical phenomenon, while the latter comes as a complement to this partial physical model in order to model the physics not explained by the former. Concerning the order of the combination ($g_a \circ g_p$), this order is natural since the NN component $g_a$ comes as a complement to the partial physics encoded in $g_p$. Alternatives methods exist: [1] propose $g = g_p + g_a$ or [2] proposed a variational formulation. We proposed $g = g_a \circ g_p$, which is more flexible than existing methods, as it does not assume the physical component needed to be completed additive or a probabilistic formulation. Given an initial condition $u_0$, the physical model $g_p$ takes as input $u_0$ and outputs an estimate of the temporal derivative $du/dt$, which is processed by the data driven term $g_a$. The latter then complements the incomplete estimate computed by the physical component.

[1] Yin et al. Augmenting Physical Models with Deep Networks for Complex Dynamics Forecasting, ICLR 2021

[2] Takeishi et al. Physics-Integrated Variational Autoencoders for Robust and Interpretable Generative Modeling, NeurIPS 2021.

The experiment results seem not strong, especially for the Gray-Scott equation

We have added complementary experiments, including a new state-of-the-art transformer baseline (Transolver) in Figure 1 and Table 2, and a new ablation study, Phys-Adaptor (Tables 3 and 4), which uses a shared NN backbone without adaptation. The Transolver's performance is similar to other backbones, and the ablation highlights the importance of the proposed adaptive framework. Regarding Gray-Scott, we respectfully disagree with your conclusions. Our method outperforms all baselines except APHYNITY in the hybrid setting. Figure 15 qualitatively demonstrates the quality of our predicted trajectories for Gray-Scott.

The paper briefly mentions the computational efficiency of the proposed method but lacks a detailed analysis.

Thanks for the suggestion, we have added in the PDF rebuttal page training time and parameters gains for the different multi-environment frameworks, showing the superiority of the GEPS framework.

The sensitivity of the proposed method to various hyper-parameters, e.g. context vector size

Again, this might not be sufficiently indicated in the core paper, but for lack of space, we included the ablation analysis in the appendix. More precisely, in the appendix C.3.1, we studied the impact of the context size on the performance and on the number of parameters required for adaptation. The proposed GEPS reaches higher performance at a lower complexity, in terms of parameters, than the reference baseline CODA. Finally in C.2, we studied how the number of adaptation trajectories impact the adaptation performance.

Why Neural ODE is utilized to generate the full trajectory? Do the experiment settings require continuous-time inference?

Neural ODE here makes reference to a family of time integration methods. This includes several differentiable PDE solvers. Here we made use of a popular RK4 explicit solver available in the torchdiffeq library. Other alternatives are possible with our framework. For the ERM methods, we used the same time-stepping technique method for GEPS than those used in the ERM methods.

We hope we clarified all your concerns and will make sure everything will be made clearer in the final version.

Thank you for taking time to read our answers to your concerns and acknowledging our added results. We would be happy to engage in a further discussion if you allow it.

The motivations served to point out a pitfall in current PDE solvers inductive principle formulation for tackling parametric PDEs. We thus proposed a new inductive principle formulation that could tackle this pointed challenge. As for the specific framework implementation, we aimed to introduce an effective solution applicable to a wide range of neural PDE solvers, supported by two key technical contributions:

• An effective yet computationally efficient meta-learning formulation, as demonstrated in Table 1 of the rebuttal, which compares favorably with existing frameworks.
• A new, flexible approach for learning hybrid models within a meta-learning framework, which, to the best of our knowledge, has not been shown before.

Being backbone agnostic and accepting physics if available allow it to handle a large number of PDE solvers encoutered in the literature.

Regarding the simple baseline « shared backbone and domain-specific weights », as mentioned to reviewer m65H, all our baselines are in fact "shared backbone with domain specific-weights", as shown in Figure 6 of [1] in appendix. We think multi-task and meta-learning frameworks, while different in their problem formulation, solve the same optimization problem, as shown in [2].

As per your requests, we added two simple strategies, where we have a shared backbone but have domain specific-weights for the first-layer or the last layer.

We report the results for Gray-Scott and Burgers for out-domain results (similar conclusions have been observed in-domain):

Concerning the results in Appendix C.2, we actually think this is one particular strength of our method, where we can also tune the low-rank matrices to obtain better accuracy at a reasonable cost if needed. This cost could be significant for large models or for larger number of adaptation samples, justifying why we first proposed to only tune a very small subset of parameters $c^e$.

We hope these results and the already provided experiments answer your concerns. Don’t hesitate to reach out if needed.

[1] Kirchmeyer et al., Generalizing to New Physical Systems via Context-Informed Dynamics Model. ICML 2022

[2] Wang et al., Bridging Multi-Task Learning and Meta-Learning: Towards Efficient Training and Effective Adaptation. ICML 2021

We appreciate the reviewer's thoughtful feedback and recommendations to improve the paper. We have addressed the raised concerns below, including additional experiments as suggested.

Utilizing shared backbone and domain-specific weights for multi-environment or multi-task learning is a widely used design.

Maybe the paper was not clear enough, but the main focus is on the general adaptation framework, not the specific implementation. The main messages are as follows:

• Classical ERM training is inadequate for modeling the variability of dynamical physical processes (parametric PDEs), as demonstrated in section 3.2.
• An alternative inductive principle is needed. We propose an adaptive conditioning principle: learning from multiple environments to adapt rapidly to a new one (section 3.3).

On the technical side, while sharing a backbone and using domain-specific weights is common, current meta-learning methods are costly and do not allow learning PDE solvers on large datasets. Thus, our implementation of this principle:

• is computationally efficient (fast adaptation to new environments) and scalable w.r.t. the number of training samples via a low-rank and first-order optimization.
• is adaptable to any NN backbones. Our experiments show it works well with several backbones (MLP, CNN, FNO).
• can embed physical models learned on different environments, even with incomplete models and unknown parameters, which has never been shown before.

Motivations of the overall hybrid design

The proposed model integrates a physics module $g_p$ encoding prior knowledge from a PDE that only partially explains the physical phenomenon. The NN module $g_a$complements $g_p$ by modeling the unexplained physics.

Regarding the combination order ($g_a \circ g_p$), it is logical that ($g_a$) operates after $g_p$ as it must augment the incomplete knowledge.

Alternative methods exist, such as [1] $g = g_p + g_a$ or [2] a variational formulation. Our approach, $g = g_a \circ g_p$, is more flexible as it doesn't assume an additive or probabilistic formulation. Given an initial condition $u_0$, $g_p$ outputs an estimate of the temporal derivative $du/dt$, which is then refined by $g_a$.

For the implementation, the physical component ($g_p$) is encoded as a differentiable numerical solver. This allows it to be combined with any NN components and the entire architecture to be trained end-to-end. Since the physics component has few free parameters, a low-rank formulation is unnecessary.

[1] Yin et al. Augmenting Physical Models with Deep Networks for Complex Dynamics Forecasting, ICLR 2021

[2] Takeishi et al. Physics-Integrated Variational Autoencoders for Robust and Interpretable Generative Modeling, NeurIPS 2021.

What about using it as an adaptor?

We added a new baseline “Phys-Adaptor” in the rebuttal (pdf companion Table 3 and 4) where only $g_p$ act as an adaptor, thus $g_a$ is shared across the environments. The lower performance of this baseline illustrates again the importance of adapting $g_a$.

SVD-PINN as a new baseline.

PINNs methods are data-free, requiring full knowledge of the physical information, which is not the case here. Our framework is applied on two popular settings:

• Pure data-driven approaches: No prior knowledge of the physics is available, and everything must be learned from data.
• Hybrid formulation: A PDE partially describing the physical system provides some prior knowledge of the physics (section 3.3.3).

These settings correspond to common real-world situations explored in the literature.

Adding an adaptor-based method

Our framework can be used with different backbones as indicated above. [3] has been applied on pre-trained transformers and performs adaptation by adding new “adapter” layers, which increase the number of layers. In our setting, we train a model from scratch to learn to be adapted to new environments efficiently. During adaptation, all baselines keep the same nb. of layers and only tune a small set of parameters, e.g. $c^e$.

[3] Parameter-efficient Multi-task Fine-tuning for Transformers via Shared Hypernetworks, ACL 2021

Experiments with advanced neural operators.

We added the Transolver baseline in our ERM experiments in the rebuttal (Fig 1 and Table 2). The performance is similar to other ERM baselines, verifying that an adaptation framework is necessary regardless of the backbone used.

Adding training times for the different methods

We have added the training times and the number of parameters for each method in the PDF rebuttal page (table 1), showing that GEPS shows a huge time gain compared to gradient-based methods (CAVIA baseline) and that our method is less costly in terms of parameters compared to the reference baselines CoDA and LEADS.

Ensemble baseline with a Transformer-based method.

For the ensemble method, we used the same model (CNN) as in our GEPS framework for a fair comparison. Besides, the above experiments with the Transolver shows that the backbone itself does not make a significant difference.

The "adaptive conditioning" term is unsuitable.

We hope to have clarified this issue. The paper main message is that ERM inductive principle alone is not sufficient for modeling multiple complex physical dynamics and that another inductive principle should be used. Hence the adaptation inductive principle developed in the paper.

Knowledge of the PDE coefficients for ERM experiments

In all our experiments, the PDE coefficients are unknown.

Same knowledge for ERM methods and GEPS?

The assumptions are the same for all the methods (baselines and GEPS). The only difference is in the induction mechanism: ERM methods assume that all data come from the same distribution, while GEPS assumes that data come from different distributions and incorporates an adaptation mechanism to adapt the network to each distribution.

We appreciate your feedback and hope we answered your concerns.

Thank you for your review and for appreciating our new experiments. We would like to engage in a further discussion if you allow it.

We agree that having a « shared backbone and domain specific weights » setting is an important baseline. However, as recent papers suggest [1], meta and multi-task learning, while conceptually different in the problem formulation, share the same optimization problem in reality.

All our baselines, whether meta-learning or multi-task learning, do have a shared backbone and domain specific weights, as you are suggesting. In CoDA [2], the paper provide a good analysis of the differences between all the different multi-environment learning frameworks using a shared backbone with domain specific weights, in Figure 6 in their supplementary material.

However, as per your request, we added two different baselines.

We proposed two different strategies:

• having a shared backbone but having domain specific weights for the first layer
• having a shared backbone but having domain specific weights for the last layer

We report the results for Gray-Scott and Burgers for out-domain results (similar conclusions have been observed in-domain):

We hope this and the already provided new experiments answer your concerns. We are happy to discuss more about it if needed.

[1] Wang et al., Bridging Multi-Task Learning and Meta-Learning: Towards Efficient Training and Effective Adaptation. ICML 2021

[2] Kirchmeyer et al., Generalizing to New Physical Systems via Context-Informed Dynamics Model. ICML 2022

Dear Reviewer,

Please have a look at the last part of the author's rebuttal, the other reviews and indicate if your rating has been updated based on that.

Your Area Chair.

Dear reviewer,

As the discussion period comes to an end, we would like to express our gratitude once again for your time and for highlighting important concerns that we hope have been addressed.

You mentioned the absence of a baseline involving "a sharing backbone with domain-specific weights." We would like to kindly remind that we have addressed this concern in our last response and hope you can take note of it this before the discussion deadline.

If so, we respectfully ask if you could kindly indicate whether your rating has been updated accordingly.

Thank you very much for your time, and we look forward to your response.

We appreciate that our new experiments have addressed your concerns. We would also like to clarify the point you raised:

Retraining the model to fit a new setting is not adaptive.

We understand that the term "adaptive conditioning" might be confusing. In the meta-learning community, "adaptation" refers to the process where a model, trained on a range of environments, is efficiently retrained when exposed to data from a new distribution (or environment) to adjust to that new context.

This is the concept we aimed to convey in our title: rather than performing direct inference on data from new environments (which underperforms compared to our approach), parametric neural PDE solvers should incorporate an adaptation mechanism. This mechanism enables the model to quickly learn from minimal data in these new environments, thereby enhancing its generalization capabilities.

This concept is too broad to accurately describe your method.

We recognize the concern about the breadth of "adaptation." Our paper primarily highlights a limitation in existing methods for solving parametric PDEs and demonstrates that an adaptive mechanism can address this issue, independent of the specific method proposed. While we provide a concrete solution, our main contribution is more general.

However, if you believe the title should be more specific to our method, we can revise it to: Boosting Generalization in Parametric PDE Neural Solvers through Context-Aware Adaptation.

Thank you for the valuable feedback and the constructive discussion, which have helped us strengthen the quality of our paper.

We're thankful for the reviewer's helpful feedback and have addressed the raised concerns below.

Section 3.2 provides the motivation for the proposed method, but I am curious about how this motivation relates to the methods proposed in Sections 3.3.2 and 3.3.3:

The main messages are:

• Classical ERM training is inadequate for modeling the variability in dynamical physical processes (framed as parametric PDEs), as shown in section 3.2.
• An alternative inductive principle is needed. We propose an adaptive conditioning principle: learning from multiple environments to rapidly adapt to a new one (section 3.3). The key focus is developing an effective meta-learning strategy.
• We propose an implementation that is computationally efficient (fast adaptation to new environments) and scalable w.r.t. the number of training samples via low-rank (3.3.1) adaptation and first-order optimization (3.3.2). This framework is adaptable to different NN backbones. Experiments show it works well with MLP, CNN, and FNO.

To further assess the generality and applicability of our framework, we consider two settings:

• Pure data-driven approaches with no prior physics, where everything is learned from data.
• Hybrid formulations with a PDE prior that partially describes the underlying phenomenon (section 3.3.3).

These settings correspond to common real-world situations explored in the literature.

Use of Neural ODE: "Neural ODE" refers to a family of time integration methods. We use a popular RK4 solver, but other solvers could also be used. This is only a relevant technical choice, e.g., for the ERM methods, we used the same time-stepping method proposed in the corresponding papers (FNO, MPPDE, Transolver, CNN) for GEPS.

The improvement in error scale compared to existing meta-learning methods does not seem significant. Figures 14 and 15 also show minimal differences.

We have now provided error bars for in-domain and out-domain experiments (tables 3 and 4 in the pdf companion), highlighting that the results obtained are consistent and significant for most cases. Regarding qualitative results, all the methods are able to reconstruct Burgers trajectory (Fig 14). For Gray-Scott (Fig 15), GEPS clearly outperforms CAVIA and LEADS. Only CoDA matches GEPS, but diverges from the ground-truth for out-range time horizon. For the Kolmogorov dataset (Fig 16), only GEPS is able to predict accurately the trajectory, and starts to diverge from ground truth for out range time horizon.

Advantages or unique aspects of the proposed GEPS method beyond its accuracy.

• We provided training times and number of parameters of GEPS compared to other methods in the PDF rebuttal page (table 1): GEPS is faster to train and requires less parameters than concurrent approaches.
• in appendix C.2, we showed that unlike other methods, GEPS performance increase with the number of adaptation trajectories while other methods performance quickly saturate.
• in appendix C.3.1, we showed that GEPS achieves greater performance and scales better in terms of trainable parameters compared to reference CoDA when increasing the context vector size.
• In C.3.2, we experimentally observed that GEPS can adapt to new environments faster than CoDA.

The overall relative error scale of FNO and MPPDE in Figures 2 and 3 and Table 1 is around 1e-1 which appears to be generally larger than the errors presented in the respective original papers. Is this due to the difficulty of the equations, or is there another reason for the relatively large errors?

The error of FNO and MP-PDE is higher because the datasets used in the respective papers have been generated by a single equation with fixed parameter values while we consider multiple environments (i.e. a distribution of parameters). Whatever the backbone, including FNO, MP-PDE and the new Transolver (see additional experiments below), the ERM induction principle is not adequate for learning from multiple physics environments when only an initial condition is given (Initial value problems).

Note that we have added a SOTA Transformer neural operator (NO) the Transolver, to the experimental comparison (see Fig. 1 and table 2 in the in the pdf companion), with the same conclusions.

How the physical knowledge $g_p$ and the data augmentation $g_a$ in Equation (6) are each encoded?

The physical component of the model is encoded as a differentiable numerical solver. Being differentiable, the latter can be combined with any NN components and the whole architecture trained end to end. As for the combination in equation 6, given an initial condition $u_0$, the physical model $g_p$ takes as input $u_0$ and outputs an estimate of the temporal derivative $du/dt$, which is processed by the data driven term $g_a$. The latter then complements the incomplete estimate computed by the physical component.

What is the significance of the order of $g_p$ and $g_a$ Would it cause problems if $g_a$ were applied before $g_p$

Different approaches have been considered in the literature. [1] considered that the data term and the physical term should be combined in an additive manner: $g = g_a + g_p$. In our case, we think that $g = g_a \circ g_p$ is a more general formulation. As for the order, since one makes the assumption that $g_p$ can be incomplete, it is natural to consider that $g_a$ augments this incomplete knowledge and operates after the physical component.

[1] Yin et al. Augmenting Physical Models with Deep Networks for Complex Dynamics Forecasting, ICLR 2021

We hope it answers your remarks. We will make sure that yours concerns are made clearer in the final version.

Thank you for your thoughtful review and for recognizing the additional experiments we included in response to your concerns.

We would like to emphasize again that, as indicated, our main contribution lies not solely in the technical advancements, but in highlighting the inadequacy of classical ERM for parametric PDEs and proposing an alternative inductive principle. We believe that this change in perspective represents a significant step forward that could impact recent trends in foundation models for PDEs, for example, those aimed at solving parametric PDEs.

In addition to this conceptual advancement, we introduced an effective solution applicable to a wide range of neural PDE solvers, supported by two key technical contributions:

• An effective yet computationally efficient meta-learning formulation, as demonstrated in Table 1 of the rebuttal, which compares favorably with existing frameworks.
• A new, flexible approach for learning hybrid models within a meta-learning framework, which, to the best of our knowledge, has not been shown before.

Respectfully, we would like to note that a score of 5 typically indicates a borderline acceptance, not a weak one.

We will incorporate your feedback into the revised version. Thank you again for your valuable insights, and we are happy to engage in further discussion.

We appreciate the reviewer's thoughtful feedback and have addressed the raised concerns below.

More experiments investigating the physics-aware component of the model would be illuminating:

For the hybrid physics-ML setting, we make two assumptions:

• the physics is only partly known and shall be completed by a NN component for the targeted forecasting task; this is the forward problem. For the pendulum, the physics component ignores the damping term, for Burgers and Kolmogorov, the model ignores both the forcing and the LES terms. The only exception is the Gray-Scott example, for which one assumes full knowledge of the equation, but the physical parameters are unknown and shall be estimated (see below).

• for all datasets, the parameters of the partial physics component are unknown and shall be estimated - this is the inverse problem. The list of parameters to be estimated is provided in appendix D.3 - table 9. For example for the pendulum, 3 parameters $(\omega_0, F, w_f)$ are estimated. Figures 7 to 11, appendix C1, show the error (MAE) for the estimation of the parameters for pendulum and Gray-Scott (this is an average over all the parameters - 3 for the pendulum, 2 for GS).

More ablation studies could be useful to further clarify the importance of each component of the proposed model:

We report 4 different ablations that could justify the importance of our low-rank module.

We have added a new baseline denoted Phys-adaptor, where we only adapt the physical component to new environments while the data-driven component is shared across all environments, without (low-rank) adaptation. The results, indicated on tables 3 and 4 on the pdf file - row Phys-adaptor, show that this model performs worse over all the datasets compared to GEPS, highlighting the importance of adaptation for the NN component.

Other ablations appear in the appendix.

• in section C.2, we evaluated different context-based meta-learning frameworks (CAVIA, CODA, GEPS) during adaptation while increasing the number of samples per environment. We demonstrated that low-rank adaptation performs better than the alternatives for complex PDEs (Kolmogorov). For example, its performance improves with the number of training trajectories, whereas existing methods quickly reach a plateau and do not show further improvement.
• in section C.3.1, we showed that our framework scales better than reference method CoDA in terms of numbers parameters and performance when varying the size of the context.
• Figure 13, appendix C.3.2, shows that our method allows faster convergence during adaptation compared to reference CoDA.

It would also be helpful to include the number of parameters of all methods evaluated.

We reported in the PDF rebuttal page (table 1) the training time and the number of parameters used for each baseline for 1D (Burgers) and 2D (Gray-Scott) PDEs.

Why is it that for Burgers, the hybrid setting of GEPS performs worse than the purely data-driven setting?

We were also surprised by this results, and redo the experiments. We have no clear explanation, but we believe that the high flexibility of our formulation for combining $g_p$ and $g_a$ can make training complex for some PDEs and causes over-fitting, especially considering the low-amount of data considered in our setting.

Two strategies for adapting PDE parameters to a new environment are proposed. Which strategy is used for Tables 2 and 3 in the paper?

For Pendulum and the Burgers, we used the context $c^e$ to learn the physical parameters. For Gray-Scott and the Kolmogorov, we directly learned the physical parameters without using the context vectors. The two strategies provide similar results, we reported in the paper the best results among the two alternatives.

We hope we answered the main concerns which will be made clearer in the final version of the paper.