In-Context Multi-Operator Learning with DeepOSets

Shao-Ting Chiu; Aditya Nambiar; Ali Syed; Jonathan W. Siegel; Ulisses Braga-Neto

cs.LGDec 18, 2025

In-Context Multi-Operator Learning with DeepOSets

Authors: Shao-Ting Chiu, Aditya Nambiar, Ali Syed, Jonathan W. Siegel, Ulisses Braga-Neto

Organizations: Department of Electrical and Computer Engineering, Texas A&M University · Department of Mathematics, Texas A&M University

Abstract

An important application of neural networks in scientific computing is non-linear operator learning. In this framework, a neural network is trained to fit a non-linear map between infinite dimensional spaces, for example, the solution operator of ordinary and partial differential equations. Recently, inspired by the discovery of in-context learning for large language models, an even more ambitious paradigm has been explored, called multi-operator learning. In this approach, a neural network is trained to learn many different operators at the same time. At inference time, the network is passed a small number of example input and output pairs to disambiguate the desired operator. In this work, we provide a precise mathematical formulation of the multi-operator learning problem, propose a parameter-efficient, attention-free architecture, called DeepOSets, and prove its universality in the in-context multi-operator learning task. We perform experiments demonstrating the efficacy of DeepOSets on several benchmarks, including Burgers' equation, 1D and 2D Darcy flow, and a variety of initial-value and boundary-value differential equations. In a comparison with the transformer-based ICON model on the 1D Darcy problem, under a small training budget (10k steps), DeepOSets attained lower error at every context size with 19% of the parameters, while with an extended budget (100k steps), ICON attains 26-41% lower error, but at the cost of roughly 218 $\times$ more training time, 12 $\times$ higher inference latency, and 5.3 $\times$ more parameters.

Figures & tables

Figure 1: DeepOSets architecture for in-context multi-operator learning.

Panel of Figure 2 — Figure 2: Relative $L_{2}$ error vs in-context examples for 1D Burgers’, 1D Darcy Flow, and 2D Darcy Flow, with red shading representing one standard deviation, all experiments show a decrease in relative error as context increases, however it is much more pronounced in the Darcy examples, perhaps due to the more complicated operator paramatrization.

Problem	Equations	Parameters	Test MSE
Forward ODEConst	$\frac{du(t)}{dt}=a_{1}c(t)+a_{2}$	$a_{1},a_{2}$	2.65e-03 $\pm$ 9.30e-03
Inverse ODEConst			1.15e-02 $\pm$ 2.71e-02
Forward ODE1	$\frac{du(t)}{dt}=a_{1}u(t)c(t)+a_{2}$	$a_{1},a_{2}$	1.65e-03 $\pm$ 8.39e-03
Inverse ODE1			1.44e-02 $\pm$ 3.86e-02
Forward ODE2	$\frac{du(t)}{dt}=a_{1}u(t)+a_{2}c(t)+a_{3}$	$a_{1},a_{2},a_{3}$	1.19e-02 $\pm$ 6.99e-02
Inverse ODE2			1.40e-02 $\pm$ 9.15e-02

Table 1: Summary of tasks. The reported Mean Squared Error (MSE) corresponds to the model trained on all 12 problems simultaneously. All models were tested on individual tasks.

Panel of Figure 3 — Figure 3: Zero-shot generalization to larger context sets. The model was trained with a fixed context size of $N=4$ . During inference, we evaluate performance on context sizes ranging from 3 to 20. (a) Forward operators show a clear trend where providing more examples monotonically decreases the MSE, indicating the model successfully leverages additional data. (b) Inverse operators maintain stability across varying context sizes. Shaded regions denote one standard deviation.

Metric	DeepOSets	ICON
Trainable parameters	894,465	4,744,193
Training time, 10k steps	42.5 s	2.15 h
Training time, 100k steps	144.6 s	8.74 h
Mean time per training step	1.446 ms	314.78 ms
Training throughput	691.7 steps/s	3.18 steps/s
Inference time per sample	0.229 ms	2.750 ms

Table 2: Computational cost of DeepOSets and ICON on the 1D Darcy flow benchmark. Inference time is the mean GPU wall-clock time per forward pass over 500 timed calls (100 per context size) for the models trained for 100k steps.

	10k training steps		100k training steps
$n_{\mathrm{ctx}}$	DeepOSets	ICON	DeepOSets	ICON
1	0.2831	0.7913	0.2096	0.1342
2	0.2085	0.8056	0.1180	0.0826
3	0.1881	0.6299	0.0976	0.0589
4	0.1740	0.6843	0.0821	0.0572
5	0.1459	0.7058	0.0819	0.0603

Table 3: Mean relative

L_{2}

test error of DeepOSets and ICON on the 1D Darcy flow benchmark as a function of the number of in-context examples

n_{\mathrm{ctx}}

, averaged over 100 held-out test samples per condition, under equal-step training schedules of 10k and 100k steps.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Value
Grid points $N$	50 (interior)
Total operators	5,000
Pairs per operator	40
Train / test split (operator level)	90% / 10%
Random seed	42

Appendix

Table 4: Dataset settings for the 1D Darcy-flow benchmark used in the ICON comparison.

Hyperparameter	Value
Model dimension	256
Transformer layers	6
Attention heads	8
FFN widening factor	4
Activation	GELU
Normalization	Pre-norm (LayerNorm)

Appendix

Table 5: ICON hyperparameters used in the comparison.

Metric	DeepOSets	ICON
Trainable parameters	894,465	4,744,193
Training steps completed	10,000	10,000
Training time (s)	42.5	7,746.2
Inference time / sample (ms)	0.670	8.042

Appendix

Table 6: Cost measurements for the 10k-step run of the ICON comparison.

Table 7: Mean relative

L_{2}

test error on the 1D Darcy flow benchmark for DeepOSets and the nearest-neighbor (NN) baseline, on a fixed held-out set of 500 test operators. ICON (trained for 100k steps), is included in (b) as an additional point of reference.

Panel of Figure 4 — Figure 4: Representative testing examples for 1D Burgers, both with 5 in-context examples. Here, the upper graphs with the blue functions represent the in-context query functions, in this case the initial condition $u(x,0)$ . The gray functions are the in-context examples $u(x,0)$ , and the blue one is the final query. The lower graphs represent the 5 in-context solutions that were given to the network, followed by the network prediction in red, and compared with the true FDM solution in black.

Panel of Figure 5 — Figure 5: Representative testing examples for 1D Darcy with 5 and 10 in-context examples for different operators. The graph is formatted in the exact same manner as the one for Burgers’ equation. Here, we see a noticeable reduction in $L_{2}$ error by simply increasing 5 more examples.

Figure 6: 2D Darcy with 5 in-context examples, relative error stated in title, we show the query function, true $u$ , our network prediction, and the difference

Figure 7: 2D Darcy with 10 in-context examples, a slight decrease in relative error, indicating the network still needs additional context

Figure 8: 2D Darcy with 20 in-context examples, a much larger decrease in relative error, showing that the network is benefitting from the additional context

Panel of Figure 9 — Figure 9: Sample results from the 12-task numerical experiment. In each case, the top plot contains the parameter functions while the bottom plot contains the corresponding solution functions. The gray lines indicate in each case the parameter-solution function pairs given in the prompt. The blue line indicates the query parameter function, while the red line indicates the corresponding solution predicted by the trained DeepOSets model. The black line indicates the exact solution corresponding to the query, using the appropriate equation.

Operator Learning

Operator learning studies models that map functions or fields to functions, enabling fast prediction across scientific and dynamical systems.

See all 99 papers

Kernel Methods for Learning Operators with Multiple Inputs and Outputs
Aug 12, 2026Adrien Weihs, Chunyang Liao, Jingmin Sun +1Kernel Methods PDE Operator Learning
Department of Mathematics, University of California Los Angeles, Los Angeles, CA 90095, USA. · Department of Mathematical Sciences, University of Arkansas, Fayetteville, AR, 72701, USA. · Department of Applied Mathematics and Statistics, Johns Hopkins University, Baltimore, MD 21218, USA.
From Theory to Application: A Practical Introduction to Neural Operators in Scientific Computing
Mar 7, 2025Prashant K. JhaNeural Surrogate Modeling Fourier Neural Operator
Department of Mechanical Engineering, South Dakota School of Mines and Technology, Rapid City, SD 57701, USA.
Harness In-Context Operator Learning with Chain of Operators
Jun 10, 2026Minghui Yang, Ling Guo, Liu YangOOD Generalization PDE Operator Learning
Department of Mathematics, Shanghai Normal University, Shanghai, China · Department of Mathematics, National University of Singapore, Singapore

Aug 12, 2026J/K move · Enter open · S save

Kernel Methods for Learning Operators with Multiple Inputs and Outputs

Adrien Weihs, Chunyang Liao, Jingmin Sun, Hayden Schaeffer

Department of Mathematics, University of California Los Angeles, Los Angeles, CA 90095, USA. · Department of Mathematical Sciences, University of Arkansas, Fayetteville, AR, 72701, USA. · Department of Applied Mathematics and Statistics, Johns Hopkins University, Baltimore, MD 21218, USA.

Kernel Methods PDE Operator Learning Operator Learning

Open paper AskPDF arXiv

Learning mappings between infinite-dimensional objects is a central challenge in scientific machine learning. We introduce a general kernel-based encoder-decoder framework for operator learning that separates observation, representation, learning, and reconstruction. We develop this framework for multi-input, multi-output operator learning, where operators map between products of potentially distinct function spaces. Our approximation theory shows that, although the number of inputs and outputs can increase, the convergence rate is governed by the most challenging constituent approximation problem rather than the overall problem dimension. The framework leads to practical kernel methods with closed-form training and inference, combining mathematical tractability with computational efficiency. We further specialize the approach to multiple operator learning by introducing KernelMO, a family of kernel methods with complementary operator-valued and product-space formulations. Across five families of parametric partial differential equations, the proposed methods achieve competitive or state-of-the-art predictive accuracy while reducing training and inference costs relative to neural operator architectures and deep learning based models, offering an efficient and lightweight alternative.

Abstract

Operator Learning

Explore similar work

Kernel Methods for Learning Operators with Multiple Inputs and Outputs

From Theory to Application: A Practical Introduction to Neural Operators in Scientific Computing

Harness In-Context Operator Learning with Chain of Operators