In-Context Operator Learning with Data Prompts for Differential Equation Problems

1 Introduction

• Earlier approaches include methods for specific equation solutions and methods for approximating solution operators, but both are described as limited when problems change.The cited examples include PINNs and related PDE methods, as well as PDE-Net, DeepONet, and FNO.Source: §1 · pp. 1–2
• A neural network limited to one solution operator may need retraining when equation terms or conditions change.Even a minor equation change can shift the operator when that change is not represented among its inputs.Source: §1 · p. 2
• The authors argue that shared commonalities across solution operators let one neural network avoid retraining and use fewer data when learning a new operator.Their goal extends beyond approximating one operator to learning an operator from data and applying it to new problems.Source: §1 · p. 2
• ICON learns an operator from demonstrations and predicts the quantity of interest for a new condition during inference, without weight updates.The demonstrations are condition–quantity-of-interest pairs; the new input is the question condition.Source: §1 · pp. 3–4

2 Problem Setup

• An operator maps one input function or a tuple of input functions to an output function, called the QoI in this paper.The operator inputs are called conditions.Source: §2 · p. 4
• Because analytical representations may be unavailable, conditions and QoIs are represented by key-value pairs of discrete inputs and corresponding function outputs.The number of pairs and their key arrangement can vary across functions.Source: §2 · pp. 4–5
• At inference, demonstrations share an unknown operator and are used with a question condition to predict its QoI; that operator may be outside the training distribution.The training setup likewise groups condition–QoI pairs sharing an operator, while leaving the operator itself unknown.Source: §2 · p. 5

3 Methodology

• The methodology frames the task as learning an operator from multiple condition–QoI demonstrations and applying it to a question condition.The method section separates prompt/query construction, architecture, data preparation and training, and inference.Source: §3 · p. 5, §3.1 · p. 5

3.1 Prompt and Query

• Prompts concatenate demonstrations and a question condition, while query vectors specify evaluation keys for the question QoI.The prompt matrix distinguishes condition and QoI pairs with index vectors; key rows can encode function terms and coordinates.Source: §3.1 · pp. 5–6

3.2 Neural Network Architecture

• A self-attention transformer encodes the prompt into an operator-and-question embedding, which a cross-attention decoder combines with queries to predict QoI values.A final linear layer adjusts the decoder output to match the question QoI’s dimensionality.Source: §3.2 · pp. 6–7

3.3 Data Preparation and Training

• Training batches mix prompts and labels from problems with different operators, and the network parameters are updated using mean squared error.The preparation procedure samples problem types, parameters, examples, and a question to build training prompts and queries.Source: §3.3 · pp. 8–9

3.4 Inference: Few-shot Learning without Weight Update

• At inference, one forward pass uses demonstrations to learn an operator and predict the question QoI without updating network weights.The input combines examples and a question condition with query keys specifying where the QoI should be evaluated.Source: §3.1 · p. 5, §3.4 · p. 9

4.1 Problems

• Training covers 19 problem types, with 1,000 parameter sets and 100 condition–QoI pairs per operator; prompts use 1–5 demos and 41–50 pairs.That yields 19,000 operators in total.Source: §4 · p. 9, §4.1 · p. 9

4.2 In-Distribution Operators

• Across 19 problems, average relative error stays below 6% with one demo, and most errors are around 2% with five demos.Testing used 500 cases per problem across 100 operators, with five cases per operator.Source: §4.2 · pp. 11–12, Figure 5 · p. 13

4.3 Super-Resolution and Sub-Resolution

• For MFC problem 17, error falls as key–value pairs increase from 10 to 500 and converges below 1%, even with one demo; training used 41–50 pairs.The evaluation varies the number of sampled key–value pairs in each condition/QoI while keeping the number of demos fixed.Source: §4.3 · pp. 12–14, Figure 6 · p. 13

4.4 Out-of-Distribution Operators

• Across four tested ODE and reaction-diffusion problems, predictions remain accurate when operator parameters extend beyond the training region.The tests include forward and inverse problems, with five demos per prompt.Source: Figure 7 · p. 14, §4.4 · pp. 14–15

4.5 Generalization to Equations of New Forms

• With training limited to ODE 2, correct demos perform no better than applying the ODE 2 operator directly or using ODE 2 demos on the new ODE.The new ODE reduces to ODE 2 when b = 0.Source: §4.5 · p. 16
• For the unseen ODE form, errors decrease as related ODE datasets are added, and training on ODEs 1, 2, and 3 performs significantly better than training on ODE 2 alone.The authors describe this as potential to learn operators for unseen equation forms as related training-data size and diversity increase.Source: §4.5 · p. 16
• Training with the full dataset used elsewhere in the paper does not improve performance on the new ODE, whose form is not closely related to the added damped-oscillator, PDE, and MFC data.The authors attribute the lack of improvement to that mismatch in relatedness.Source: §4.5 · p. 16

• Learning need cover only a distribution of question conditions, rather than every possible condition.Source: §5 · p. 16
• For ODEs, shared relations among u’s time derivative, u, and c may let the network identify the ODE from a few demos.Source: §5 · pp. 16–17
• The operator family studied is small and simple; larger families may require more demos, a larger network, and more computation.Source: §5 · p. 17

• ICON acts as an operator learner at inference, learning from given data and applying the operator to new conditions without weight updates.Source: §6 · p. 17
• A single network learned operators from a small number of prompted demos across forward and inverse ODE, PDE, and mean-field control problems without retraining or fine-tuning.Source: §6 · p. 17
• With training key–value counts limited to a narrow range, ICON generalized to a broader test range, with errors decreasing and converging as counts increased.Source: §6 · p. 17
• Experiments provide preliminary evidence of potential to learn and apply operators for previously unseen equation forms.Source: §6 · p. 17
• The experiments are small in scale; future work proposes scaling network size, problem types, key and value dimensions, condition and QoI lengths, and demo capacity.Source: §6 · p. 17

6.1 Details on Differential Equation problems

• Table 3 summarizes operator, condition, and QoI parameters for the differential-equation problems.Source: Table 3 · p. 18, §6.1 · p. 18
• ODE datasets use 50 equidistant key–value pairs per u and c; prompts, queries, and ground truth use the first n−1 c pairs and first n u pairs, with n sampled from 41 to 50.Source: §6.1 · pp. 18–19
• Damped-oscillator and PDE datasets prepare 100 equidistant pairs per condition/QoI, then sample n pairs for prompts, queries, and ground truth, with n from 41 to 50.Source: §6.1 · p. 19
• MFC data use 100 equidistant pairs for g and a 51 × 100 grid for density ρ; n pairs per condition/QoI are sampled for prompts, queries, and ground truth, with n from 41 to 50.Source: §6.1 · p. 19