PDE Generalization of In-Context Operator Networks: A Study on 1D Scalar Nonlinear Conservation Laws

TL;DR

The paper addresses generalization in PDE learning by proposing in-context operator learning with ICON-LM, applied to 1D scalar nonlinear conservation laws described by $\partial_t u + \partial_x f(u)=0$. It demonstrates that a single transformer-based ICON can infer and apply forward and reverse operators from data prompts without weight updates, using a data-prompt framework based on cubic-flux conservation laws. It also shows generalization to PDEs with new flux forms by employing change-of-variables and varying-stride prompts, expanding the range of solvable problems. This work advances toward a foundation model for PDE-related tasks under the in-context operator learning framework.

Abstract

Can we build a single large model for a wide range of PDE-related scientific learning tasks? Can this model generalize to new PDEs, even of new forms, without any fine-tuning? In-context operator learning and the corresponding model In-Context Operator Networks (ICON) represent an initial exploration of these questions. The capability of ICON regarding the first question has been demonstrated previously. In this paper, we present a detailed methodology for solving PDE problems with ICON, and show how a single ICON model can make forward and reverse predictions for different equations with different strides, provided with appropriately designed data prompts. We show the positive evidence to the second question, i.e., ICON can generalize well to some PDEs with new forms without any fine-tuning. This is exemplified through a study on 1D scalar nonlinear conservation laws, a family of PDEs with temporal evolution. We also show how to broaden the range of problems that an ICON model can address, by transforming functions and equations to ICON's capability scope. We believe that the progress in this paper is a significant step towards the goal of training a foundation model for PDE-related tasks under the in-context operator learning framework.

Figures (12)

• Figure 1: Operator Learning v.s. In-Context Operator Learning. For classic operator learning, one model is limited to approximate a single operator, with the need of fine-tuning when the operator changes to a "close" new one. For in-context operator learning, a single model can approximate a wide range of operators. It can serve as a "foundation model" that could be directly applied without fine-tuning for different PDE-related tasks in the training distribution, or even beyond due to generalization in the operator space. It can also be fine-tuned to strengthen its expertise in particular operator domains, if necessary.
• Figure 2: Illuatration of training and inference of ICON for PDEs, using conservation laws as examples.
• Figure 3: Illustration of the scheme for recursive predictions.
• Figure 4: Averaged error v.s. the number of examples used for in-context operator learning. (a) Forward error. (b) Reverse error.
• Figure 5: Illuatration of in-context operator learning for $\mathcal{F}_{0.5u^3 + 0.5u^2 + 0.5u,0.1}$, $\mathcal{R}_{0.5u^3 + 0.5u^2 + 0.5u,0.1}$, $\mathcal{F}_{-0.5u^3 -0.5u^2 -0.5u,0.1}$ and $\mathcal{R}_{-0.5u^3 -0.5u^2 -0.5u,0.1}$. For each case, the prompted five condition-QoI examples are shown with dotted color lines. The forward predictions shown with dashed red lines overlap with the ground truth QoIs shown with solid black lines. Since there are no unique ground truth solutions for the reverse operators, we apply the exact forward operators to the predicted QoIs by forward simulation, and show the recovered conditions with dashed blue lines, which overlap with the question conditions shown with solid black lines.