Thompson Sampling in Function Spaces via Neural Operators
Authors: Rafael Oliveira
Xuesong Wang
Kian Ming A. Chai
Edwin V. Bonilla
Venue: NeurIPS 2025 poster ()
Keywords: neural operators
Thompson sampling
bandits
Bayesian optimization
Primary area: general_machine_learning
TL;DR: An extension of Thompson sampling to optimization problems over function spaces using neural operators as surrogates

Abstract: We propose an extension of Thompson sampling to optimization problems over function spaces where the objective is a known functional of an unknown operator's output. We assume that queries to the operator (such as running a high-fidelity simulator or physical experiment) are costly, while functional evaluations on the operator's output are inexpensive. Our algorithm employs a sample-then-optimize approach using neural operator surrogates. This strategy avoids explicit uncertainty quantification by treating trained neural operators as approximate samples from a Gaussian process (GP) posterior. We derive regret bounds and theoretical results connecting neural operators with GPs in infinite-dimensional settings. Experiments benchmark our method against other Bayesian optimization baselines on functional optimization tasks involving partial differential equations of physical systems, demonstrating better sample efficiency and significant performance gains.

Thompson Sampling in Function Spaces via Neural Operators
Rafael Oliveira" CSIRO's Data61
Xuesong Wang CSIRO's Data61 Sydney, Australia
Kian Ming A. Chai DSO National Laboratories Singapore
Edwin V. Bonilla CSIRO's Data61 Sydney, Australia
Sydney, Australia
Abstract
We propose an extension of Thompson sampling to optimization problems over function spaces where the objective is known functional of an unknown operator's s output. We assume that queries to the operator (such as running a figh-fidelitty simulator physical experiment) are costly, while functional evaluations the operator's ' output are inexpensive. Our algorithm employs approach using neural operator surrogates. This strategy avoids explicit uncertainty quantification by treating trained neural operators as approximate samples from a Gaussian process (GP) posterior. We derive regret bounds and theoretical results connecting neural operators with GPs in infiniteddimnnse benchmark method against other Bayesian optimization baselines functional optimization tasks involving partial differential equations of physical systems, demonstrating better sample efficiency and significant performance gains.
comple-then-opimiiee
settings. Experiments
Introduction
Neural operators have established themselves as versatile models capable of learning complex, nonlinear mappings between function spaces [1]. They have demonstrated success across diverse fields, including climate science [2]. materials engineering [3], and computational fluid dynamics [4]. Although their applications in supervised learning and physical system emulation are well- studied, their potential for online learning and optimization within infiniteddimensional function
spaces remains relatively untapped.
In many scientific contexts, learning operators that map between function spaces naturally arises, such the task of approximating solution operators for partial differential equation (PDE) [1]. However, adaptive methods that efficiently query these operators to optimize functional objectives of their outputs (particularly in active learning setting) still underdeveloped. For example, when designing porous structures, one is often interested in optimizing how liquids flow through the structure using, e.g., Darcy flow PDEs [5], and, in the sciences, inverse problems can be solved by optimization to infer initial conditions or parameters of a physical process from observations [6, 7].
To address this gap, propose framework that integrates neural operator surrogates with Thompson sampling-bast based acquisition strategies [8] to actively optimize objectives of the form:
E argmax argmax (a). aEA
where G. : :A - an unknown operator between function spaces A and U4 and f u + I is known functional. We follow the steps of Bayesian optimization frameworks for composite functions [9, 10], which leverage knowledge of the composite structure to speed-up optimization, extending these frameworks to functional domains. Applying the theoretical results for the infinite-wide limit
"Corresponiing author: rafael. dos.anttodeelliviiraaaaaaaa cww.:eerrrr..aaaaaa.a...o.....
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

of neural networks [11, 12], we show that trained neural operator approximates posterior sample from a vector-valued Gaussian process [13-15] in a sample-then-opee the-thttoooiiiee approach [16]. Therefore, we are able to implement an approximate form of Thompson sampling without the need for expensive uncertainty quantification frameworks for neural operators, such as deep ensembles [17] or mixture density networks [18], and derive theoretical regret bounds on its performance. Experiments evaluate our approach on problems with classic PDE benchmarks against Bayesian optimization baselines.
2 Related work
Bayesian optimization with f'unctionals and operators. Bayesian optimization (BO) has been a successful approach for optimization problems involving expensive-to-evaluat black-box functions [19]. Prior work on BO in function spaces includes Bayesian Functional Optimization (BFO) [20], which uses Gaussian processes to model objectives defined over functions, focusing on scalar functionals without explicitly learning operators. Follow-up work extended the framework to include prior information about the structure of the admissible input functions [21]. Astudillo and Frazier [9] introduced the framework of composite Bayesian optimization, which was later applied by Guilhoto and Perdikaris [10] to optimization problems involving mappings from finite-dimensional inputs to function-valuee outputs. Their objective was to optimize known functional of these vunction-valuedd outputs. Our approach differs by directly working in function spaces, involving function-to-function operators. Despite the availability of GP models for function-to-function mappings [22], we are unaware of BO or GP-based bandit algorithms incorporating such models. Lastly, in the bandits literature, TeanThhah and Yu [23] introduced the problem of functional bandits. Despite the terminology, they deal with the problem of optimizing known functional of the arms rewards distribution, similar to the setting of distributionally robust BO [24], and therefore not directly comparable our case.
Thompson sampling with neural networks. Neural Thompson Sampling (NTS) [25] employs neural networks trained via random initialization and gradient descent to approximate posterior distributions for bandit problems with scalar inputs and outputs, inspiring our use of randomized neural training for operator posterior sampling. The Sample-Then-Opinn Opttimiee Batch NTS (STO- BNTS) variant [16] refines this by defining acquisition functions on functionals of posterior samples, facilitating composite objective optimization. STO-BNTS extends this to batch settings using Neural Tangent Kernel (NTK) and Gaussian process surrogates, relevant for future batched active learning with neural operators. These approaches rely on the NTK theory [11], which shows that infinitely wide neural networks trained via gradient descent behave Gaussian processes. To the best of our knowledge, this approach has not yet been extended to the case of neural network models with function-valued inputs, such as neural operators.
Active learning for neural operators. Pickering et al. [17] applied deep operator networks (Deep- ONets) [26] to the problem of Bayesian experimental design [27]. In that framework, the goal is to select informative inputs (or designs) to reduce uncertainty about an unknown operator. To quantify uncertainty, Pickering al. [17] used an ensemble of DeepONets and quantified uncer- tainty in their predictions based on the variance of the ensemble outputs. Li et al. [18] introduced multi-resolution active learning with Gaussian mixture models derived from Fourier neural operators (28). With probabilistic outputs, mutual information can be directly quantified for active learning and Bayesian experimental design approaches. Lastly, Musekamp al. [29] proposed a benchmark for neural operator active learning and evaluated ensemble- based models with variance-based uncertainty quantification on tasks involving forecasting. In contrast to our focus in this paper, active learning approaches are purely focused on uncertainty reduction, neglecting other optimization objectives.
3 Preliminaries
Problem formulation. Let AA and UI denote two function spaces, and let G. : A + 1 be an unknown target operator² between them. Consider an objective functional f : 2 - I, which is
Here, we use the term unknown loosely, in the sense that it is not fully implementable within the com- putational resources or paradigms accessible to us. For example, the target operator can be simulator in a high- perforrrffm computing facility which we have limited access
2

---

assumed known and cheap to evaluate. Given compact search space S C A, we aim to solve:?
a" E argmax f(G.(a)), aES
while G. is only accessible via expensive oracle queries: for chosen a, observe function- valued output = HG.(a) + E, where H : -> y represents an observation operator, typically the discretization grid, with 3/ being a (usually inite-dimensionas is observation noise, assumed independent and identically distributed (i.i.d.) across queries. The algorithm allowed query the oracle with any function E S for up budget of N queries. For this paper, we focus on problems with finite search space S < co, though the framework is general.
Hilbert space, and &~ (O E)
Neural operators. neural operator is a specialized neural network architecture modeling operators G::A - U between function spaces A and 2 [1]. Assume A C C(X, Rd), and U C C(2, Rd.) where C(S,S') denotes the space of continuous functions between sets S and S'. Given an input function E A. neural operator Gø performs sequence of transformations
u1 UL-1 + uy through L layers of neural networks, where u, : Nr + IR is a continuous function for each layer I E {l..... L}. and X2 == Z is the domain of the output functions and dy du. In one of its general formulations, for a given layer 1 E (1...., L.), the result the transform (or update) at any x X4+1 can be described as:
un (x)) = =a))
14+1(1) = on (1x, Ri(E,,', (T((()))) m((x)) me(x)) dui(')) + W ([)(())) b((s)) (2)
Gola)(z) == 11) (2),
where 11; X1+1 + 2i is a fixed mapping, aj : + denotes an activation function applied elementwise, R. : Ke+1 x Ne x R4 x RR + Rd++ixd defines a (possibly nonlinear or positive- semidefinite) kernel integral operator with respect measure in on X1, W, E 1R4++1x, is weight matrix, and b, X++1 + Rdi+1 is bias function. We denote by 0 the collection of all learnable parameters of the neural operator: the weights matrices Wi, the parameters of the bias functions by and the matrix-valuee kernels R,, for all layers E {1..... L}. Variations to the formulation above correspond to various neural operator architectures based on low-rank kernel approximations, graph structures, Fourier transforms, etc. [1].
Vector-valued Gaussian processes. Vector-valued Gaussian processes extend scalar GPs [13] the case of vector-valued functions [14]. Let A be an arbitrary domain, and let be Hilbert space representing codomain. We consider the case where both the domain A and codomain U might be anfinite-dimnsiin
vector spaces, which leads GPs whose realizations are operators G. A + 2 [15]. To simplify our exposition, we assume that 24 is a separable Hilbert space, though the theoretical framework general enough be extended to arbitrary Banach spaces [30]. vector-valued Gaussian process G. ~ GP(G, K) on A is fully specified by a mean operator G: A + 24 and a cositie-semmdeffniie operator- operator-valued covariance function K : A x A + L(U). where C(4)) denotes the space bounded linear operators on U. Formally, given any a, a' A and any u, u' e 24 it follows that:
[a)] = Cov (ov((G((a) u), (G. = (u, K (a,a')u))
where (-, > denotes the inner product and Cov(,.)) stands for the covariance between scalar variables.
Assume we are given set of observations D. ((a,,)) 1=))/=1 CA x U, where y/ = G. (a.) &, and & ~ (O, Σ) corresponds to Gaussian noise with covariance operator E L(U). The posterior mean and covariance can then be defined by the following recursive relations:
G.(a) = ((- K2(a, = = K,-1(a,a') - az) (KK a.) + Σ)- ( (K-1(ar,
+ Σ) - G
(5)
(6)
'W use "E argmax" acknowledging that the problem may have multiple global optima, forming set of global optimizers. Whenever assume unique minimizer, will use the equality symbol "=", instead.
3

---

Algorithm 1: GP-TS
Algorithm 2: NOTS (ours)
Input:
Search space S, initial data Do
Input: Search space S, initial data Do for = 1...., do
for / E (1...., ,T) do
Sample 9t (P( 12)1)
0: argming (x)), 0r,0 (0 Ex)
Select 2, E argmax EEx Dc(a)
a; E argmax aes f(Go,
Query y/: = fxx) +
yr G. G. + &t
Update D, = D U M1)
D. = Dr-1 U fa Vr}
for any a, a' E A, and E N, which are an extension of the same recursions from the scalar-valued case [31, App. F] to the case of vector-valued processes. Such definition arises from sequentially conditioning the GP posterior Ko == K. This recursion leads to the same matrix-bart based definitions of the usual GP posterior equations [13], but in case it avoids complications with the resulting higher-order tensors that arise when kernels are operator-valued.
each observation, starting from the prior with Go G and
Thompson sampling. Thompson sampling (TS) is a relatively simple randomized strategy for sequential decision making under uncertainty, which has found many successes the Bayesian opti- mization and aulti-armed bandits literature [8, 25, 32, 33]. When applied to optimization problems, the core idea TS is to query an objective function f at points ž₁ sampled from the probability distri- bution of the optimum location To do so, the objective function is modeled as sample from a Bayesian probabilistic model, which is typically a linear model [8] or GP [33], and then TS samples realizations Se of the objective from the model's posterior p( A point which maximizes sampled function 9i then corresponds to a sample from the posterior distribution over the optimum p(=*|D,,-)) The procedure is summarized in Algorithm for the case of a GP. Under mild assumptions, 'S is known to produce sequence of candidates x, such that (x+) asymptotically converges to fx=) [33, 34].
argmax, f(x) given the observations D:-1 := {zy, V= --=
4 Neural operator Thompson sampling
We propose Thompson sampling algorithm for the optimization of functionals of unknown operators in the setting of Eq. 1. Instead of relying extensions of traditional probabilistic methods to operator modeling, our method applies flexible and scalable neural operators as surrogates Gt, training them to approximate posterior samples over the true operator G. conditioned on data. The method is designed to efficiently explore the search space while balancing the
trade-off.
4.1 Approximate posterior sampling
Given data D, = {(4.. (=) we train neural operator Gg with parameters O, that minimize:
()) = 1y HGa(a;))l A|101|
where ||-|| represents the norm in the underlying vector space, and À >>> 0 is regularization factor which relates to the noise process & [35]. We minimize & via gradient descent starting from O..0 (0, Eo). where Eg is diagonal matrix following Kaiming He [36] or LeCun initialization [37], which scale each layer's 's weights initialization variance according to the width of the previous layer. By an extension of standard results on the infinite-width limit of neural networks the neural operator setting, we can show that the trained neural operator approximates posterior sample from vector-valued GP when, e.g., we train only the last linear layer (see App. C.4), which in turn guarantees regret bounds (Sec. 5). The prior over G. is implicitly defined as the vector- valued Gaussian process given by the conjugate kernel [38, 39] associated with the neural operator architecture and the weights initialization distribution. Lastly, we note that, in practice, observations are discretized over finite grid or other fints-aimensional representation [1], so that the observation space is y C m and the difference norms in Eq. 7 reduce to Euclidean distances.

---

4.2 Thompson sampling algorithm
In Algorithm 2, we present the Neural Operator Thompson Sampling (NOTS) algorithm for the optimization of sequentially over T iterations similar standard GP-TS (Algorithm 1). To sample realization from the neural operator posterior, each iteration begins with the random initialization of the parameters of a neural operator that serves as a surrogate model for the true unknown operator. At each iteration, the neural operator model is trained according to Section 4.1, minimizing regularized least- sast-squares loss based on the currently available data, yielding an approximate sample G, == Ge. from the true operator posterior p(G,\Dr-1). The next step involves selecting the input for querying the oracle by maximizing the value of the objective functional f over the neural operator's 'ss predictions G.(a). Finally, the algorithm runs the potentially expensive step of querying the true operator G. with the selected input function a,, which may involve a complex simulation or physical experiment, and updates the dataset with the new (noisy) observation yt- This process repeats for up to T iterations, producing a sequence of function- valued queries ae that approximates the true optimum a' (1).
functionals of black-box operators. The algorithm operates
Computational cost. Each iteration of NOTS incurs a linear computational cost of (t)) due to the retraining of the neural operator model, which can be further reduced by use of minibatch stochastic gradient descent. The reinitialization with randomized weights followed by retraining is what ensures that we have a new approximate posterior sample for TS conditioned on the available data at every iteration. Compared to a more traditional GP-based approach, which applied to our setting would incur (($) cost per step due to the inversion of a covariance matrix of l data points, we achieve a much more computationally efficient and scalable algorithm, despite the cost of retraining the model.
5 Theoretical results
In this section, we establish the theoretical foundation of our proposed method. We show how a randomly initialized neural operator approximates a GP in the infinite-wind wiith limit through the use of the conjugate kernel, also known as the NNGP kernel [38-42]. under certain assumptions. This allows us to extend existing results for GP Thompson Sampling (GP-TS) [33] to our setting.
5.1 Neural operator abstraction
A neural operator models nonlinear operators G A - 21 between possibly
function spaces .A and Ut Current results in NTK [11] and GP limits for neural networks [12] do not immediately apply to this setting, as they rely can leverage an abstraction for neural operator architectures which sees their layers as maps over fiitee-iimensional inputs [43], which result from truncations to make the modeling problem tractable.
minee-dimeniionl domains. However,
Considering a neural operator with single hidden layer, let M E N represent the layer's width, AR A + C(Z, Rdd) denote (fixed) continuous operator, and bo Z - Rede denote (fixed) continuous function. For simplicity, we will assume scalar-valued output functions with d, = In general, with single hidden layer, the model described in Eq. 2 can be rewritten as:
(=) w.a (WR (RRAR(a((s) W. (.))(((==)) + Wxbo(=)), ZE
(8)
where 0 == -eecw Wo: WR. W.,, Ws) MaMatdntdd =: W represents the model's flattened parameters. The finite weight matrix WR representing the kernel convolution integral arises as a result of truncations required in the practical implementation of neural operators (e.g., a finite number of Fourier modes or quadrature points). With this formulation, neural operator architectures [43]. In Appendix B, we discuss how Fourier neural operators [28] fit under this formulation, though the latter is general enough to incorporate other cases. We also highlight that neural operators possess universal approximation properties [44], given sufficient data and computational resources, despite the inherent low-rank approximations in their architecture.
recover most popular
5.2 infinite-width limit of neural operators
With the construction Eq. 8, we can simply see the result of neural operator layer when evaluated at a fixed 2 E Z cquivalently as a M-width feedforward neural network:
(z) he he(v)(a)) ()a()
(9)
5

---

where the input is given by va == (A)( alIo(2)), bo(2)] E V. and v RAR t. db.
Conjugate kernel. We can now derive fifhnitewwidth limits. The conjugate kernel describes the distribution of the untrained neural network he V - R under Gaussian weights initialization, whose infinite-width limit yields Gaussian process [38, 40]. Formally, the conjugate kernel is defined as:
kn (v, v') ;= lim E [ (v)he. (v'),: v,v V.
(10)
M M
Further background the conjugate kernel and the NTK can be found in Appendix A.
Since the composition of the map A 3 (a, 2) H (a) E V with a kernel on V yields a kernel on A x Z [45, Lem. 4.3), the conjugate kernel of Ge is determined by:
(a z, d', 2') == k, (Q), ver( a, a' E A, 2, E.
(11)
where ky is the conjugate kernel of the neural network hg. Such a kernel defines a covariance function for GP over the space of operators mapping A to U. Assume U (v) is closed subspace of the space of functions which are square integrable with respect to a a-finite Borel measure on 2, and let L(4)) denote the space of linear operators on U. The following then defines positive-semidefnite operator-valued kernel KG A x A L(1))
(Ka(a,a')u)(e = - ko(a, 2, d', d')z')((='))))))e)))))))))))" du(z'),
(12)
for any u EU, a, a' E and = E 2. Hence, we can state the following result, whose proof can be found in Appendix C.2.
Proposition 1. Let Go : -u be a neural operator with a single hidden layer, where U S (v) is closed, and is finite Borel measure on Z. Assume w. (N oğ for og > 0 such that og x 1/M, while the remaining paramelers have their entries sampled from a fixed normal distribution. Then, as M + co, on every compact subset of A, the neural operator converges in distribution to a zero-mean vector-valued Gaussian process with operator-valued covariance functiinc given by:
lim Eo @ Go(a')) = Ka(a,a'), a,a' E A,, M M
where Ka : A L(U) is defined in Eq. 12, and & denotes the outer product.
5.3 Bayesian cumulative regret bounds
Bayesian regret. We analyze the performance of sequential decision-making algorithm via its
Bayesian cumulative regret. An algorithm's 's instant regret for querying a E A at iteration 2 1 is:
rt := f(G. (a (a")) - (a.)
(13)
where a" is defined in Eq. 1. The Bayesian cumulative regret after T iterations is then defined as:
RT E
Σ
rz
(14)
where the expectation is over all sources of randomness affecting the decision-making process, i.e., the prior for G. and the observation noise. If the algorithm achieves sub-linear cumulative regret, its simple regret asymptotically vanishes, as limy E [min/e(1....) re] < limy-xo 4 Rr, leading the algorithm's 's queries a, to eventually approach the true optimum a"
Regularity assumptions. For our analysis, we assume 4 S C²(v) is closed subspace of the Hilbert space (1) of square-integrable r-measurable functions, for given finite Borel measure on compact domain Z. We will assume the search space S C A is finite. The true operator G. - l4 will be assumed to be a sample from a vector-valued Gaussian process G. ~ GP(O, K), where the operator-valued kernel K A x A + L(U) is given by the neural operator's 's infinite-width limit in Proposition 1. Observations y = HG, (a) + are assumed to be corrupted by i.i.d. zero-men mean Gaussian noise, & ~ (O E), where Σ is positive definite on y C R.
We adapt ofte-offttheaar regret bounds for GP-TS [33] to an exact version of NOTS. To do so, we first observe that, for a linear functional f E L(U, R). the composition with Gaussian random operator G. GP(G, K) yields ccllar-vauued GP, i.e., f G. GP o f Kf),, where ff (KKf: aa) + f(K((. f)).). We can then extend GP-TS regret bounds to the case of operators.

---

Proposition 2. Let f : + IR be bounded linear functional such that f = foH, where f: + +R linear, and G. GP(O, )). Consider sequential algorithm selecting a, argmax (G)(())
and observing Yr HG.(a.) (a:) + Et. where G, (r,,, and & ~ A((O,I), for (T.....T]. Then, this algorithm's 's expected cumulative regret is such that:
Rr E O((T),
(15)
where Ō() suppresses logarithmic of the O+) asymptotic rate.
This result shows that NOTS can achieve sublinear cumulative regret in the infinite-width limit with an exact GP posterior sample. The result connects existing GP-TS guarantees to NOTS, and it differs from existing guarantees for other neural network based Thompson sampling algorithms [16, 25], which explored the scalar case and a frequentist setting (i.e., the objective function being a fixed element of the reproducing kernel Hilbert space defined by the network's neural tangent kernel). In the Bayesian setting, there is also no need for time-dependent regularization parameter [16], allowing for simpler implementation. Yet we note that Proposition 2 concerns the exact GP case. However, Proposition 1 ensures that a single-hidden-lee the infinite-width limit, and we show in the appendix that training the last layer via gradient descent approximates a posterior sample, as in previous results for conventional neural networks [12, App. DJ. Appendix C presents proofs and further discussions on limitations and extensions, and a validation experiment can be found in Appendix E.
randomly initialized neural operator follows GP in
Experiments
We evaluate our NOTS algorithm on two popular PDE benchmark problems: Darcy flow and shallow water model. Our results are compared against a series of representative Bayesian optimization and neural Thompson sampling baselines. More details about our implementations and further experiment details can be found in Appendix D. Code for our experiments will be made available online.* 4
6.1 Algorithms
We compare NOTS against a series of GP-based and neural network BO algorithms modeling directly the mapping from function-valued inputs A (discretized over regular grid) to the scalar-valued functional evaluations f(G+(a)), besides trivial random search (RS) baseline. NOTS is implemented with standard and spherical FNOs [46]. following default library settings for these PDEs [47]. We first implemented BO with a 3-layer infinite-width ReLU Bayesian neural network (BNN) model, represented as a GP with the corresponding conjugate kernel. According Li et al. [48], these models can achieve optimal performance in high- sigh-dimensional settings when compared to other BNN methods. Two versions of this framework are in our experiments, one with log- log-expected improvement, given its well established competitive performance [49], simply denoted as "BO" in our plots, and one with Thompson sampling (GP-TS) [34]. As our experiments are over finite domains, sampling from a scalar GP boils down to sampling from a multivariate normal distribution. Next, we evaluated a version of Bayesian functional optimization (BFO) by encoding input functions in a reproducing kernel Hilbert space (RKHS) via their minimum- norm interpolant and using squared-exte exponnntcap kernel over functions which takes advantage of the RKHS structure as in the original BFO [20]. Lastly, evaluated sample-then-opting neural Thompson sampling (STO-NTS), training a 2-layer 256-width fully connected neural network with a regularized least-squares loss [16].
6.2 PDE benchmarks
Darcy flow. Darcy flow models fluid pressure in a porous medium [28], with applications in con- taminant control, leakage reduction, and filtration design. In our setting, the input a E C((O, 1)2 R,) is the medium's 's permeability on a Dirichlet boundary, and the operator G. maps a to the pressure field C((O, 1)2. R). To train Gg, solver at 16 x 16 resolution. Two materials are considered, leading to a binary grid for a and a continuum of pressure values for each u grid cell. More details are in Li et al. [28] and Appendix D.
generate 1,000 onput- pairs via finite-difference
*Code repository: bttps::/giite com/ http:/csroooummmm/o
7

---

Cumulative regret trials
Averaged cumulative trials
Averoped Cumulative Regret
RS
Cumilativs Regree
250
BD
NOTS
10a
200
BFO
RS
150
GPTS
6x10
STONTS
NOTS
4 #x10
50
4x10-1
STONTS
GPTS
150 terations
200
300
teeraionn
best candidate in- put function
best candidate out- put function
worst candidate in- function
worst output function
candidate
10
10
15
15
Figure 1: Darcy flow rate optimization. Overlay of cumulative regret (top left) and its average (top right) metrics across trials for the negative total flow rates case in the Darcy flow problem. The shaded areas correspond to standard deviation across 10 trials. The corresponding input-output functions that achieved the best and worst flow rates are presented (bottom). White regions a(x) = means fully open permeability and black regions a(x) 0 represents impermeable pore material. The output function suggests pressure field where brighter color indicates higher pressure.
Shallow water modeling. Shallow water models capture the time evolution of fluid mass and discharge on rotating sphere [46]. The input E C(S² {t 0}. R represents the initial geopotential depth and two velocity components, while the output u E C(S x {t r) R gives the state at time t т. We train GB on 200 random initial conditions on x equiangular grid, using a 1,200 timestep to simulate up to hours.
6.3 Optimization functionals
We introduce several optimization functionals that are meaning in the context of the benchmark problems. As we aim to solve maximization problem, physical quantities to be minimized are defined with negative sign. The first three functionals were applied the Darcy flow problem and the last one to shallow water modeling. Note that both cases, have the same domain for the PDE solutions and input functions a, i.c., Z X.
and clarify their physical
Negative total flow rates [50] [(u a) = - Sax of the domain and the outward pointing unit normal vector of the boundary. This functional integrates the volumetric flux -a(x) (o(x))v((r) Vu(x) along the boundary, which corresponds to the total flow rate the fluid. Such an objective can be optimized for leakage reduction and contaminant control.
n))d.. Here 0.Y' is the boundary
Negative total pressure [51] [())
- fx |u(x))dr. This objective computes the total fluid pressure over the domain in the Darcy flow system.
Negative total potential energy f(u, - fxx a(x) Vu(x)) dr + fx ((z)u(() dr. This func- tional quantifies the system's 's total potential energy, balancing the energy dissipated by fluid friction (the first term) against the potential energy supplied by the uniform fluid source (the second term, where s is assumed). The minimizer a", therefore, consists of the most hydrodynamically efficient design for the given flow constraints.
Inverse problem (() = - Ilu uy 12. "y represents the ground truth solution. This objective is specific to shallow water modeling, as we aim to find the initial condition that generates u, at time which is also simplification of the assimilation objective in weather forecasting [52, 53].

---

Averaged curuueetive regret across trials
Averaged cumulative regret across trials
Avarages Cunulativo Regret
Aweraged Curmuaativn Regiet
4 10
10-1-
10
10
NOTS
10--
BFO
IFO
STONTS
STONTS
50
100
150
200
250
300
100
150
200
250
300
Iterations
Iterations
Pressure
(b) Potential energy
Figure Darcy flow pressure (a) and potential energy (b) optimization problems averaged cumulative regret. The shaded areas correspond to one standard deviation across 10 trials.
Cumulative regret across trials
Averaged cumulative regret across trials
1500
RS
Averaged Cumulative Ragrat
Cumulative Regree
B0
NOTS
BFO
GPTS
109
80
500
GPTS
STONTS
200
250
50
100
150
Figure 3: Shallow water inverse problem. Overlay of cumulative regret (left) and its average (right) metries across trials for the inverse problem the shallow water data. The shaded areas correspond to one standard deviation across 10 trials.
6.4 Results
Our results are presented in Figure to 3, comparing the cumulative regret NOTS against the baselines different settings of PDE problems and functional objectives. Results are summarized in Hr. of each method across the different problems.
Table 1 with the final average regret, i.e., T
In Figure I, we present results for the flow optimization problem in the Darcy flow PDE benchmark. The results clearly show that GP-based BO methods struggle in this high-dimensional setting, while NOTS (ours) able to consistently find optimal solutions. As described Section 6.2, input functions E A for Darcy flow are binary masks representing two materials of different permeability which are discretized over 2D grid of 16-by-16 sampling locations. Hence, when applied to standard GP-based BO methods, the inputs correspond to 256-dimensional vectors, which can be quite -ige-dimensional for standard GPs. The optimization results of the input and output functions also show the effectiveness of our approach. In the case of the "best candidate" which achieves the lowest total flow rate, the input function shows large contiguous impermeable regions ideal
that block fluid outflow and thus generate high interior pressure which can be treated design for leakage control. In contrast, the "worst candidate" exhibits the highest total flow rates. It has smooth, boundary-connecter permeable zones allowing fluid to escape effortlessly. Lastly, figures 2(a) and 2(b) show the results on optimizing pressure and potential energy on Darcy flow. On these functionals, BO and GP-TS can achieve better performance, recalling their use of the infinite-width BNN kernel, which has shown good performance on high-dimensional problems [48]. Yet, we can see significant performance improvements from NOTS with respect to all baselines.
Figure shows our results for the inverse problem on the shallow water PDE benchmark. This setting involves higher dimensional discretized inputs (61 4414emiiia when flattened), leading to an extremely challenging problem GP approaches. inputs kernel is too computationally intensive for BFO, leading it to crash before 250 iterations are completed. We believe that STO-NTS's low performance is due to architectural limitations, as it uses simple fully connected network, which leads to need for higher amounts of data (i.c.,
particular, the evaluation of the functional

---

Table Results summarr: Final average regret of each method and its standard deviation.
Method
Darcy flow rates
Darcy flow energy
Darcy flow pressure
Shallow water
RS
0.872 + 0.022
0.309 + 0.005
0.077 + 0.001
4.632 + 0.876
BO
0.703 + 0.045
0.251 + 0.024
0.047 + 0.001
1.639 + 0.532
BFO
0.788 + 0.066
0.208 + 0.014
0.078 = 0.006
3.076 = 0,886
GP-TS
0.674 + 0.050
0.189 0.093
0.038 = 0.004
1.942 = 0.502
STO-NTS
0.068 + 0.002
0.282 0.011
0.068 + 0.002
2.329 = 0,800
NOTS
0.012 + 0.001
0.125 + 0.042
0.012 + 0.001
0.134 + 0.043
more iterations). NOTS, however, is able learn the underlying physics of the problem aid its predictions, leading to more efficient exploration and higher performance.
7 Conclusion
We have developed Neural operator Thompson sampling (NOTS) for optimization problems in func- tion spaces and shown that it provides significant performance gains in encoding the compositional structure of problems involving black-box operators, such as complex physics simulators or real physical processes. NOTS also comes equipped with theoretical guarantees, connecting the existing literature on Thompson sampling to this novel setting involving neural operators.
Discussion. We have shown empirically that using neural operators as surrogates for Thompson sampling can be effective without the need for expensive uncertainty quantification schemes by relying on theoretical results for infinitely wide deep neural networks and their connection with Gaussian processes. Neural operators have allowed for effective representation learning which scales to very high-dimensional settings, where traditional bandits and Bayesian optimization algorithms would struggle. Although GPs typically perform well on Bayesian modeling tasks with low volumes of data, the functional optimization problems we considered have high-dimensional data as both inputs and outputs, rendering the application of traditional multi-output GP models challenging. The basic computational complexity of inference with a vector-valued GP model scales cubically with both the number of data points and the number of output coordinates [14]. For the shallow water PDE, for example, both inputs and outputs lie in a 6144-dimensional space. With 300 iterations, a multi-output GP would have invert kernel matrix over more than million data points towards the last iterations. Hence, without specialized kernels and computationally efficient approximations, traditional GP approach would be unsuitable due to the very large number of outputs. In contrast, neural operators are specially designed to deal with function-valued input and output data, typically spatial domains, with linearly scaling computational complexity. Therefore, NOTS can better scale to accommodate longer runs or extensions to batched evaluations than a GP approach, even though we limited experiments 300 iterations to allow for comparisons against GP baselines.
Limitations and future work. We note that our current results are focused on the case of finite search spaces and well specified models, which provide a first step towards more general use cases. An extension to continuous domain could, for example, parameterize the set of input functions and optimize such parametric representation or tractable nonparametric extensions [20, 21], which might be application specific. Our theoretical analysis only considered the case of neural operator with a single hidden layer, despite the multi-layer setting in our experiments. These and other limitations are further discussed in Appendix F. As future work, plan to investigate the generalization of our results to more general settings, such as continuous domains and batched evaluations. Lastly, we note that NOTS also offers framework for task-to-task amortization and few-shot learning, operator learning data can be reused across tasks with different objective functionals.
Acknowledgments and Disclosure of Funding
This research was carried out solely using CSIRO's s resources. Chai contributed while on sabbatical leave visiting the Machine Learning and Data Science Unit at Okinawa Institute Technology, and the Department of Statisties in the University of Oxford. This project was supported by resources and expertise provided by CSIRO IMT Scientific Computing.
of Science and
10

---

References
[1] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya. Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. Journal of Machine Learning Research, 24(89), 2023.
[2]
Thorsten Kurth, Shashank Subramanian, Peter Harrington, Jaideep Pathak, Morteza Mardani, David Hall, Andrea Miele, Karthik Kashinath, and Anima Anandkumar. FourCastNet: Acceler- ating global high-resolution weather forecasting using adaptive Fourier neural operators. In Proceedings of the Platform for Advanced Scientific Computing Conference, PASC '23, New York, NY, USA, 2023. Association for Computing Machinery.
[3] Vivek Oommen, Khemraj Shukla, Saaketh Desai, Rémi Dingreville, and George Em Karniadakis. Rethinking materials simulations: Blending direct numerical simulations with neural operators. npj Computational Materials, 10(1):145, 2024.
[4] Zongyi Li, Nikola Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Otta, Moham- mad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Azizzadenesheli, and Anima Anandkumar. Geometry-informer neural operator for large-scale 3d PDEs. Advances in Neural Information Processing Systems, (5355836-5854, 2023.
[5] Niclas Wiker, Anders Klarbring, and Thomas Borrvall. Topology optimization of regions of Darcy and Stokes flow. International journal for numerical methods in engineering, 69(7): 1374-1404, 2007.
[6] V. V. Penenko and E. A. Tsvetova. Inverse problems for the study of climatic and ecological processes under anthropogenic influences. IOP Conference Series: Earth and Environmental Science, 386( 1):01203,, nov 2019.
[7] Dan MacKinlay, Dan Pagendam, Petra M Kuhnert, Tao Cui, David Robertson, and Sreekanth Janardhanan. Model inversion for spatio-temporal processes using the Fourier neural operator. In Fourth Workshop Machine Learning and the Physical Sciences (NeurIPS 2021), 2021.
[8] Daniel Russo and Benjamin Van Roy. Learning to optimize via posterior sampling. Mathematics 2014.
of Operations Research,
[9] Raul Astudillo and Peter I. Frazier. Bayesian optimization of composite functions. In 36th International Conference Machine Learning (ICML), 2019.
[10] Leonardo Ferreira Guilhoto and Paris Perdikaris. Composite Bayesian optimization in function spaces using NEON - Neural Epistemic Operator Networks. Scientific Reports, 14, 2024.
[11] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems, volume 31, Montreal, Canada, 2018.
[12] Jaehoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In 33rd Conference Neural Information Processing Systems (NeurIPS), Vancouver, Canada, 2019.
[13] Carl E. Rasmussen and Christopher K. Williams. Gaussian Processes for Machine Learning. The MIT Press, Cambridge, MA, 2006.
[14] Mauricio A Alvarez, Lorenzo Rosasco, and Neil D Lawrence. Kernels for vector-valued functions: a review. Foundations and Trends in Machine Learning, 4(3), 2012.
[15] Palle E T Jorgensen and James Tian. valuud Gaussian processes and their covariance kernels. Infinite Dimensional Analysis, Quantum Probability and Related Topics, 27(02), 2024.
[16] Zhongxiang Dai, Yao Shu, Bryan Kian Hsiang Low, and Patrick Jaillet. Sample-then-optmmize batch neural Thompson sampling. In Proceedings of the 36th International Conference on Neu- ral Information Processing Systems, NIPS 22, Red Hook, NY, USA, 2022. Curran Associates Inc.
11

---

[17] Ethan Pickering, Stephen Guth, George Em Karniadakis, and Themistoklis P. Sapsis. Discover- ing and forecasting extreme events via active learning in neural operators. Nature Computational Science, 2(12))823-833 2022.
[18] Shibo Li, Xin Yu, Wei Xing, Robert Kirby, Akil Narayan, and Shandian Zhe. Multi-resolution active learning of Fourier neural operators. In Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li, editors, Proceedings of the 27th International Conference on Artificial Intelligence and Statistics, volume 238 of Proceedings of Machine Learning Research, pages 1440-2443 PMLR, 2024.
[19] Bobak Shahriari, Kevin Swersky, Ziyu Wang, Ryan Adams, and Nando De Freitas. Taking the human out of the loop: A review of Bayesian optimization. Proceedings of the IEEE, 104 (1):148-175, 2016.
[20] Ngo Anh Vien and Marc Toussaint. Bayesian functional optimization. In AAAI Conference Artificial Intelligence, pages 4171-4178, New Orleans, LA, USA, 2018.
[21] Pratibha Vellanki, Santu Rana, Sunil Gupta, David de Celis Leal, Alessandra Sutti, Murray Height, and Svetha Venkatesh. Bayesian functional optimisation with shape prior. Proceedings of the AAAI Conference on Artificial Intelligence, 33(01))1617-162, 2019.
[22] Carlos Mora, Amin Yousefpour, Shirin Hosseinmardi. Houman Owhadi, and Ramin Bostanabad. Operator learning with Gaussian processes. Computer Methods in Applied Mechanics and Engineering. 434:1 17581, 2025.
[23] Long Tran- Thanh and Jia Yuan Yu. Functional bandits. arXiv e-prints, art. 1405 2014.
[24] Johannes Kirschner, Iliaa Bogunovic, Stefanie Jegelka, and Andreas Krause. Distributionally robust Bayesian optimization. In Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics (AISTATS), Palermo, Italy, 2020. PMLR: Volume 108.
[25] Weitong Zhang, Dongruo Zhou, Lihong Li, and Quanquan Gu. Neural Thompson sampling. In International Conference on Learning Representations, 2021.
[26] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):118-23) 2021.
[27] Tom Rainforth, Adam Foster, Desi R Ivanova, and Freddie Bickford Smith. Modern Bayesian experimental design. Statistical Science, 39(1):100-114, 2024.
[28] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhat- tacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations. OpenReview, 2021.
(29) Daniel Musekamp, Marimuthu Kalimuthu, David Holzmüller, Makoto Takamoto, and Mathias Niepert. Active learning for neural PDE solvers. In International Conference on Learning Representations (ICLR), Singapore, 2025. OpenReview.
[30] Houman Owhadi and Clint Scovel. Gaussian Measures, Cylinder Measures, and Fields on B. pages 347-359. Cambridge Monographs on Applied and Computational Mathematics. Cambridge University Press, 2019.
[31] Sayak Ray Chowdhury and Aditya Gopalan. On kernelized multi-armed bandits. In Proceedings of the 34th International Conference on Machine Learning (ICML), Sydney, Australia, 2017.
[32] Kirthevasan Kandasamy, Akshay Krishnamurthy, Jeff Schneider, and Barnabas Poczos. Asyn- chronous Parallel Bayesian Optimisation via Thompson Sampling. In Proceedings the 1 International Conference on Artificial Intelligence and Statistics (AISTATS), Lanzarote, Spain, 2018.
12

---

[33] Shion Takeno, Yu Inatsu, Masayuki Karasuyama. and Ichiro Takeuchi. Posterior sampling- based Bayesian optimization with tighter Bayesian regret bounds. In Proceedings of the 41 International Conference on Machine Learning (ICMI. 2024), volume 235, Vienna, Austria, 2024. PMLR.
[34] Daniel Russo and Benjamin Van Roy. An information-thheorticc analysis of Thompson sampling. Journal of Machine Leurning Research (JMLR), 17:1-30, 2016.
[35] Sergio Calvo-Ordoñez, Jonathan Plenk, Richard Bergna, Alvaro Cartea, José Miguel Hernández- Lobato, Konstantina Palla, and Kamil Ciosek. Observation noise and initialization in wide neural networks. In 7th Symposium on Advances in Approximate Bayesian Inference Workshop Truck, 2025.
[36] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), pages 1026-1034, 2015.
[37] Yann LeCun, Leon Bottou, Genevieve B Oir, and Klaus Robert Müller. Efficient BackProp. pages 9-50. Springer Berlin Heidelberg, Berlin, Heidelberg, 1998.
[38] Jachoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-dickstein. Deep neural networks as Gaussian processes. In International Conference on Learning Representations (ICLR), 2018.
[39] Zhengmian Hu and Heng Huang. On the random conjugate kemnl and neural tangent kernel. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 4359.4368. PMLR, 18-24 Jul 2021.
[40] Radford M Neal. Priors for Infinite Networks, chapter 2. pages 29-53. Springer New York, New York, NY, 1996.
[41] Amit Daniely. SGD learns the conjugate kernel class of the network. In I Guyon, U Von Luxburg. S Bengio, H Wallach, R Fergus, S Vishwanathan, and R Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017
[42] Zhou Fan and Zhichao Wang. Spectra of the conjugate kernel and neural tangent kernel for linear-width neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, Curran Associates, Inc., 2020.
[43] Mike Nguyen and Nicole Mücke. Optimal convergence rates for neural operators. arXiv e-prints, art. arXiv:2412.17518, 2024.
[44] Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds for Fourier neural operators. Journal of Machine Learning Research, 22(290):
1-76, 2021.
[45] Ingo Steinwart and Andreas Christmann. Support Vector Machines, chapter 4, pages 110-163. Springer New York, New York, NY, 2008.
[46] Boris Bonev, Thorsten Kurth, Christian Hundt, Jaideep Pathak, Maximilian Baust, Karthik Kashinath, and Anima Anandkumar. Spherical Fourier neural operators: Learning stable dynamics on the sphere. In International conference on machine learning (ICML), pages 2806-2823, PMLR, 2023.
[47] Jean Kossaifi, Nikola Kovachki, Zongyi Li. David Pitt, Miguel Sch-Siiia,,,i, Valentin Duruis- seaux, Robert Joseph George, Boris Bonev, Kamyar Azizzadeneshelli, Julius Berner, and Anima Anandkumar. A Library for Learning Neural Operators. arXiv e-prints, art. arXiv:2412.1000 10354, December 2024. doi: 10.48550/arxiv.24122 10354.
[48] Yucen Lily Li, Tim G. J. Rudner, and Andrew Gordon Wilson. A study of Bayesian neural network surrogates for Bayesian optimization. In 2024 International Conference on Learning Representations (ICLR), Vienna, Austria, 2024. OpenReview.
13

---

[49] Sebastian Ament, Samuel Daulton, David Eriksson, Maximilian Balandat, and Eytan Bak- shy. Unexpected improvements to expected improvement for Bayesian optimization. In 37th Conference on Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA, 2023.
[50] Victor J Katz. The history of Stokes' theorem. Mathematics Magazine, 52(3):146-1555 1979.
[51] SeongHee Jeong and Sanghyun Lee. Optimal control for Darcy's 's equation in a heterogeneous porous media. Applied Numerical Mathematics, 207:303-333 322, 2025.
[52] Florence Rabier, Jean-Noel Thépaut, and Philippe Courtier. Extended assimilation and forecast experiments with a Nour-dimensional variational assimilation system. Quarterly Journal of the Royal Meteorological Society, 124(550);1811-1887) 1998.
[53] Yi Xiao, Lei Bai, Wei Xue, Kang Chen, Tao Han, and Wanli Ouyang. Fengwu-4DVar: Coupling the data-driven weather forecasting model with 4D variational assimilation. arXiv preprint arXiv:2312.12455 12455, 2023.
[54] Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard E. Turner, and Zoubin Ghahramani. Gaussian Process Behaviour in Wide Deep Neural Networks. In International Conference on Learning Representations, Vancouver, Canada, 2018. OpenReview.net.
[55] Chaoyue Liu, Libin Zhu, and Mikhail Belkin, On the linearity of large non-linear models: When and why the tangent kernel is constant. In Advances in Neural Information Processing Systems, volume 33, 2020.
[56] Boris Hanin. Random neural networks in the infinite width limit as Gaussian processes. The Annals of Applied Probability, 33(6A):4798 4819, 2023.
[57] Thomas M. Cover and Joy A. Thomas. Elements of Information Theory. John Wiley & Sons, 2005. ISBN 9780471241959, doi: 10. 10.10022047744882.
[58] Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesen- sky. Bin Bao, Peter Bell, David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison, Zachary DeVito, Elias Ellison, Will Feng, Jiong Gong, Michael Gschwind, Brian Hirsh, Sherlock Huang, Kshiteej Kalambarkar, Laurent Kirsch, Michael Lazos, Mario Lezcano, Yanbo Liang, Jason Liang, Yinghai Lu, CK Luk, Bert Maher, Yunjie Pan, Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk, Michael Suo, Phil Tillet, Eikan Wang. Xiaodong Wang, William Wen, Shunting Zhang, Xu Zhao, Keren Zhou, Richard Zou, Ajit Mathews, Gregory Chanan, Peng Wu, and Soumith Chintala. PyTorch Faster machine learning through dynamic Python bytecode transformation and graph compilation. In 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume (ASPLOS '24). ACM, April 2024. doi: 10.1145/3620665.36400
[59] Niranjan Srinivas, Andreas Krause, Sham M. Kakade, and Matthias Seeger. Gaussian Process Optimization in the Bandit Setting: No Regret and Experimental Design. Proceedings of the 27th International Conference on Machine Learning (ICML 2010), pages 1015-1022, 2010.
[60] Johnathan M Bardsley, Antti Solonen, Heikki Haario, and Marko Laine. Randomize-therr optimize: A method for sampling from posterior distributions in nonlinear inverse problems. SIAM Journal on Scientific Computing, 36(4):A1895 A1910, 2014.
[61] Zhiping Mao and Xuhui Meng. Physies-informed neural networks with residual/gradient- based adaptive sampling methods for solving partial differential equations with sharp solutions. Applied Mathematics and Mechanics, 44(7):1069-108 2023.
[62] Gilles Pisier. Probabilistic methods in the geometry of Banach spaces. Probability and analysis, pages 167-241, 1986.
[63] Daniel Augusto de Souza, Yuchen Zhu, Harry Jake Cunningham, Yuri Saporito, Diego Mesquita, and Marc Peter Deisenroth. Infinite Neural Operators: Gaussian processes on functions. arXiv e-prints, art. arXiv:2510.16675 October 2025. doi: 10.48550/arxiw.2510..66 16675.
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Demonstrated by theoretical and experimental results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how be expected to generalize other settings.
much the results
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Discussion the appendix and the conclusion
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
15

---

Justification: In the appendix (supplementary material), the reader can find the proofs and full assumptions.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Details in the appendix.
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
16

---

Answer: [Yes]
Justification:
Code will be released at htttp:://githbl http://gtthuu...ocs....oiia wommcssiro-fuiii
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips..
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /CodeSubmic / //nips. wttpccccpppccccceeececccccccccct /tte:///ggiiieeeeeeeuuubiiiiii ssionnollll onPolicy))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: In the appendix (supplement)
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Standard deviations reported with the plots
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
17

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Details in the appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: / /ttp///eerpp/.ee/ee cc httpllcc/thhissuiieeeee.
Answer:
[Yes]
Justification:
We have read NeurlPS Code of Ethics and carried our research accordingly.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Discussed in Appendix G.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
18

---

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: NA.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: PDE benchmarks acknowledged in the main paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
19

---

Answer: [Yes]
Justification:
Code will be released at htttp:://gtthen ctt. cotmm//srrofuumml/.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: NA.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: NA.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
20

---

Justification: NA.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
: Please refer to our LLM policy (https://neur)) htw.arrpps.cccoooeeerrrreeee/e/e for what should or should not be described.
21

---

Appendix
We now present detailed theoretical background, proofs, experiment settings, and additional results that complement the main paper. Appendix A reviews essential background on the infinite-width limit of neural networks [12] and how they relate Gaussian processes [13]. We discuss the distinction and applicability of the two main kernel-based frameworks suitable for this type of analysis, namely, the neural tangent kernel (NTK) by Jacot et al. [11] and the conjugate kemnll, also known as the neural network Gaussian process (NNGP) kernel [38, 40], which was the main tool for our derivations. Appendix formulates Fourier neural operators [28] under the mathematical abstraction that allowed us to derive the operator-valuee valued kernel for neural operators. The proofs of the main theoretical results then appear Appendix C, including the construction and properties of the operator-valued valued kernel and the correspondence between trained neural operators and their GP limits. Appendix D) describes the PDE benchmarks considered, namely Darcy flow and shallow water equations, alongside the respective objective functionals for optimization tasks. Experiment details, hyperparameter settings, and baseline implementation details are provided Section D.4. Appendix E presents results on an experiment with a single-hidden-laye neural operator validating our theoretical results. Lastly, we discuss limitations and potential broader impact in sections F and G, respectively.
A Additional background
In this section, we discuss the main differences between the neural tangent kernel [11] and the conjugate kernel, also known as the neural network Gaussian process (NNGP) kernel [12]. Both kernels are used to approximate the behavior of neural networks, but they differ in how they use Gaussian processes to describe the network's 's behavior.
A.1 Conjugate kernel (NNGP)
The conjugate kernel has long been studied in the neural networks literature, describing the cor- respondence neural networks with randomized parameters and their limiting distribution as the network width approaches infinity [38-41, 41, 54]. Neal [40] first showed the correspondence between an infinitely wide single-hidden hidden-layer network and a Gaussian process by applying the central limit theorem. More recent works [38, 41, 54] later showed that the same reasoning can be extended to neural networks with multiple hidden layers. The NNGP kernel is particularly useful for Bayesian inference as it allows us to define GP priors for neural networks and analyze how they change when conditioned on data, providing us with closed-form expressions for an exact GP posterior in the infinite-width limit [38].
Define an L-layer neural network h(,,0): :X' + with (22; 0) == (z 0) via the recursion:
ho (z; 0 == x
(16)
bi), LE {I. where x E X' represents an arbitrary input on a denotes layer's weights matrix, Mi is the width of the Ith layer, b, M bias vector, or R - IR denotes the layer's ss activation function, which is applied elementwise on vector-valued inputs, and = veo( (W,, Wi,betH bi) collects all the network parameters into a vector. Assume
by (x; 0) an
domain X, W. E R M,M
[Wila, (1 Mi..) and [b:]; ~ NN 1), for i {L...., Mi. E l E (1...., L). and let M min min(M,,-. (M,,... Mr)- The NNGP kernel then corresponds to the infinite-width limit of the network outputs covariance function [38] as:
(411) and
kmna (z. x')
lim E[M(x: 0)h(2': 0)), x, X
(17)
where the expectation is taken under the parameters distribution. By an application of the central limit theorem, it can be shown [38, 40] that the neural network converges in distribution to Gaussian process with the kernel defined above, ie.::
ho S h ~ GP kxacP),
(18)
where 4 denotes convergence in distribution as M - oo. In other words, the randomly initialized network follows GP prior in the infinite-widte limit. Moreover, it follows that, when conditioned on
22

---

data Dr := f.; /4 )- assuming = h(z;) + and 6 ~ NO, o²). Bayesian neural network distributed according GP posterior in the -ndth limit as:
h|DN GP(uN, kx)
(19)
(x(x) == B[h(z) Dx] kx(x)) (KN + o;
(20)
(x(a, x') == Cor(h())),, DN] = x(x,x)) - kx(x))) (KN + "k))),
(21)
for any 2,, E X, where KN = [k(ss; 2;)] E RNNN, kx(x)) = [ RNN
YN = b and we set k: :- Kamer to avoid notation clutter. Hence, the NNGP kernel allows us compute exact GP posteriors for neural network models. However, we emphasize that the
conjugate kernel should not be confused with the neural tangent kernel [11], which corresponds to
the infinite-width limit of [zv 0)- Voh(z'; 0)], instead.
A.2 Neural tangent kernel (NTK)
The NTK approximates the behavior of neural network during training via gradient descent by considering the gradients of the network with respect to its parameters [11]. Consider an L-layer feedforward neural network he X - R as defined in Eq. 16. In its original formulation, Jacot et al. [11] applied scaling factor of VMr to the output of each layer to ensure asymptotic convergence in the limit M oo of the network trained via gradient descent. However, later works showed that standard network parameterizations (without explicit output scaling) also converge to the same limit as long as a LeCun or Kaiming/He type of initialization scheme applied the parameters with appropriate scaling of the learning rates [12, 55], which ensure bounded variance in the infinite- -with limit. The NTK describes the limit:
bNTI( x²) lim ([)oho( Voho(z')),
(22)
M-+00
for any v, x' E X, where the expectation is taken under the parameters initialization distribution. Under mild assumptions, the trained network's output distribution converges to Gaussian process described by the NTK [11, 38). Although originally derived for the unregularized case, applying L2 regularization to the parameters norm yields a GP posterior with a term that can account for observation noise [35]. Namely, consider the following loss function:
N
2 (0) = Σ(i - ha:))) A|10 - 00112.
(23)
(=1
where 6g denotes the initial parameters. As the network width grows larger, the NTK tells us that the network behaves like linear model [1 11,, 55] as:
Mx;; 0) R h(:; 0o) + Voh( 0) 0--0, (0 0b), x EE
(24)
The approximation becomes exact in the infinite width limit within any bounded neighborhood B²x(0)) :- (0 0 < Rt of arbitrary radius 0 < < oo around Oo, the second-order error term vanishes [55]. The latter also means that Val(-: 0) converges to fixed feature map ő X + Ho. where Ho is the Hilbert space spanned by the limiting gradient vectors. With this observation, our loss function can be rewritten as:
(0) (u-h(zz
Voh(z 0) (0 00)) + All0
0ol12.
(25)
The minimizer of the approximate loss can be derived in closed form. Applying the NTK then yields the infinite-widtt model:
hx(x) = h(x) + k (z) (K NTK + -yy - hv),
(26)
where h ~ CP (0, kmTaF) denotes the network at its random initialization, as defined above, XTK (x) =
[Kwie x:) >) E R. K : [kwxx(2; (Ir,, xr,)] ) E NNNN and hy == (M(,,, )|N E R.V. Now applying the GP limit to the randomly initialized network h [12, we have that:
hN (x MTK (z) (x) (K + XXI
kNN)
(27)
(28)
( x²) = -k(x,0)) + k (x:) (K + AI) -I) KK + - (x')
(29)
- -KNTK (z) (K + XXI) -kw(')) - ( + AI) -
NIX (z').
23

---

where we again set k := kauer to avoid clutter. However, note that such GP model does not generally correspond Bayesian posterior. An exception where only the last linear layer is trained, while the rest are kept fixed at their random initialization; in which case case, the GP described by the NTK and the exact GP posterior according to the NNGP kernel match in the unregularized setting [12).
A.3 Application Thompson sampling
For our purpose, is important to have a Bayesian posterior in order to apply Gaussian process Thompson sampling (GP-TS) [33] for the regret bounds in Proposition 2. Therefore, are con- strained by existing theories connecting neural networks to Gaussian processes to assume training only the last layer of neural networks of infinite width, which gives a Bayesian posterior of the NNGP after training. In addition, we had to consider the case of a single hidden layer neural operator, as the usual recursive step applied derive the infinite-width limit would require intermediate (iinitteddimmsionaa) function space in our case, making the extension to the multi-layer case not trivial due the usual dimee-dimensionll assumptions [55]. Nonetheless, the NOTS algorithm suggested by our theory has demonstrated competitive performance in our experiments even in more relaxed settings with a multi-layer model. Future theoretical developments Bayesian analysis of neural networks may eventually permit the convergence analysis of the more relaxed settings in our experiments. In any case, we present an experiment with a wide Engle-hidden-layer model with training only on the last layer in Appendix E.
B Fourier neural operators under the abstract representation
Recalling the definition the main paper. we consider single hidden layer neural operator. Let M E N represent the layer's 's width, Ar A - (*)) denote (fixed) continuous operator, and bo Z + ds denote (fixed) continuous function. For simplicity, we assume scalar outputs with du We consider models of the form:
Go(a)(z) w.a (Wn (R))R(()=)) (.a(Hl(2)) + Webu(z)), 2E
(30)
where 0 = :(wo, WR. W.. Wo) E RM, x xdd xIR Mxd. x R Mxd. =: W represents parameters.
Fourier neural operators. As example, show how the formulation above applies to the Fourier neural operator (FNO) architecture [28]. For simplicity, assume that X' is the d-dimensional periodic torus, i.e., X' = [0, 2 ) and 2 X. Then any square-integrable function X Cd. can be expressed as a Fourier series:
a (x)
Σ a(s)e #EZ-
Va E X,
(31)
where := : /-1 E C denotes the imaginary unit, and â(s) are coefficients given by the function's 's Fourier transform F; -2 Cd.) + (2d, Cda) as:
/ a(z)e ((,a) dz, s s Zd.
a(s) = (s) (2T) x
(32)
For translation-invariint kernel R(xxx') R(x x'), applying the convolution theorem, the integral operator can be expressed as:
R R(5z)a(z) Rxa a
x
(F(R Σ R(s)
(33)
#E24
In practice, function observations are only available discrete set of points and the Fourier series is truncated at a maximum frequency Smax E Zd, which allows one to efficiently compute it via the fast Fourier transform (FFT), Considering these facts, FNOs approximate the integral as [28]:
N R(s.))((s))e x EEZ,
f R(a, (')(2
(34)
24

---

where the N values of range from 0 to Smax in all d coordinates. Finally, defining Ar as:
AR : C(2, Cd,) + C( ) (Fa)(s))ettun))
(35)
a (Fa)(sN)e(ton))
and letting WR = [R(s1)...... [((s)) we recover Eq. 30 for FNOs in the complex-valued case.
For real-valued functions, to ensure that the result is again real-valued, symmetry condition imposed R, so that its values for negative frequencies are the conjugate transpose of the corresponding values for positive frequencies. However, we can still represent it via a single matrix of weights, which is simply conjugate transposed for the negative frequencies. Lastly, note that complex numbers can be represented as tuples of real numbers.
C Theoretical Analysis
In this section, we provide the proofs of the theoretical results presented in the main paper.
C.1 Auxiliary results
Definition (Multi-Layer cally-Connected Neural Network). multi-layer fully- fully-connected neural network with L hidden layers, input dimension dg, output dimension d1+1, and hidden layer widths di,..., d,, is defined recursively as follows. For input x E X, the pre-activations and activations at layer 1 1, 1....,,.. L + 1 are:
(x) w x + (0)
(36)
(x)=
bbl-), 1= 1= L,
(37)
v vll (x) a(vlly)
(38)
where w(!) E Rd weight matrices, b Rd:41 are bias vectors, R - R. is coordinate-wii non-linearity, and the network output is f(x) = (x). The weights are initialized as with)
1/2
CW de
W! where W p with mean variance and finite higher (N(0, given fixed constants ow and cs 2 0.
moments, and biases by"
Lemma 1 (Infinite-width limit [56]). Consider a feedforward fully connected neural network as
E + R that is absolutely continuous with polynomially bounded derivative. Fix the input dimension do, the output dimension d1+1, the number of layers L, and a compact set X Rddo As hidden layer widths di....,dy oo, the random field + f(x) converges weakly in C(X,R Bd211) to centered Gaussian process with covariance K (L++)) XxXX Bddet xde+1 defined recursively by:
in Definition with non- con-linearity
K (
c.I ow B(v.v') la(v) @ a(v)l.
(39)
K (2,x) K (I,x²) (x(,)))))
K (x,))) for 2, with the initial condition for = 1
where (v,v') N (0,
determined by the first-layer weights and biases.
Lemma 2 (Thm. 3.1 in Takeno et al. [33]). Let f GP(O, where l X X + R
positive-defmite kernel finite X. Then the Bayesian cumulative regret of GP-TS is such that:
Ri E O((TTT)
where YT denotes the maximum information gain after T iterations with the GP model.
C.2 width neural operator kernel
Assumption The activation function a IR - absolutely continuous with derivative bounded almost everywhere.
25

---

Lemma 3 (Continuity of limiting GP). Let Ge A + C(Z) be neural operator with single hidden layer, as defined as in Eq. 30. Assume W, NO I for d > such that oe & W and let the remaining parumeters have their entries sumpled from fixed normal distribution. Then, as M + oc, the neural operator converges in distribution to a zero-mean Gaussian process with continuous realizations G A' > C(2) on every compact subset A' C.A.
Proof. As shown in Section 5.2, when evaluated at fixed point 2 E 2, neural operator with a single hidden layer can be seen as:
hho()())=h(((((... 2)): a A,
(40)
where 2b) := (=(a) is fixed map 1/> A 2 - V. with V = BedRdd+ and ho is conventional feedforward neural network, as defined in Definition 1. By Assumption 1 and Lemma 1, it follows that, M + xc, he converges in distribution to Gaussian process h GP(O, kx) with continuous sample paths, i.e. P/ E C(V') = 1 on every compact V C V. The continuity of +b A x Z + V then implies that g -== h x is zero-mean GP whose sample paths lie almost surely in C(A' x 2), for a compact AA' C A. as Z is already assumed compact. Therefore, for each E A, we have IP [g(a..) -) E C(Z) = 1, so that G(aa) == g(a,)) -) defines an almost surely continuous operator G A' + C(Z) on compact A'' c A. The verification that is a vector-valued GP trivially follows.
Proposition 1. Let Go A U be neural operator with single hidden layer, where U _ (v) is closed, and v is a finite Borel measure on 2. Assume w. ~ (N, IgT), for og > such that o x 1/M, while the remaining parameters have their entries sampled from a fixed normal distribution. Then, M - co, every compact subset of A, the neural operator converges in distribution to a zero-mean veltor Gaussian process with operator- covariance funciiion given by:
lim Eo @ Go(a')) Ko(a,a'), a, E A, M
where KG A A L(U) is defined in Eq. 12, and & denotes the outer product.
Proof of Proposition 1. We start by noting that any continuous function u E C(Z) is automatically included in (v),,, since ||ul|2 (ull(=((>) = Sz (z) dz(z) S || < co. Hence, any operator mapping into C(2) also maps into L (1) by inclusion.
Applying Lemma 3, it follows that Go $ G. where G is a zero-mean GP, as M + co. Now, given any E U, a, A and = Z, we have that:
(B(G(a) @ (()))lu)(z) = ((G()) w]]
gla,,) / gld,, =')((z') du(=)) )) (z)
V gla, = (') z')u(=') du(z')
E
(41)
f Elg(a, 2) 22 (2)( de(zz)
z
[ kola, z, d', 2')((z))))))))))))()))) du(='),
where we applied the linearity of expectations and the correspondence between g A x 2 - R and the limiting operator G + U. As the choice of elements was arbitrary, it follows that the above defines an operator-valued kernel KG. Linearity follows from the expectations. Given any E A, as a positive-semidefinitt operator, the operator norm of Ko (a, a) is bounded by its trace, such that:
|| (a, a)ll < Tr(Ko a)) = BlIlG()) E [ g² 2()))) < x2 1
(42)
and the last expectation is finite, since g is almost surely continuous. Hence, Ko(a, E L(U).
26

---

C.3 Regret bound
Proposition 2. Let f + IR be bounded linear functional such that f fo where f: +R linear, and G. GP(O, K). Consider sequential algorithm selecting a2 argmaxaes f()
and observing Yt HG. (a:) + &t. where G2 G. (D, and & N((,, for t (T)....T}. Then, this algorithm's 's expected cumulative regret is such that:
Ry E Ö(VT),
(15)
where (-) suppresses logarithmic of the (-) asymptotic rate.
Proof of Proposition 2. Starting with the assumption that foHH an observation
only provides information about HG. AA - y, missing any component of G. mapping to the null space ker ()) C u of the observation operator H. Thus, any G = G. + indistinguishable from G., for any Z A - U with range ((A) C ker(H), based on the information available in the observations. Therefore, for the optimization objective I oo G. A - to be identifiable, we restrict admissible functionals such that f(u + w) f(u), for all ker(H) and E U, and assuming f= fo o H ensures that this requirement is satisfied.
By linearity, follows that foG ~ GP ky) for a fixed bounded linear functional fill - I. Hence, o G. is equal in distribution scalar-valued GP h ky)) with ky A A + R given by:
(a a²) ((K
a, E E
where we implicitly identify the functional f with a unique corresponding vector in L, also denoted by f. by the Riesz representation theorem apply the operator K (a, a') [(U) to f. By Lemma 2, standard GP-TS on an objective h ~ GP(O, ky), a finite domain S C A will have Bayesian cumulative regret AT E O(TTTTT). Note that 7,, in our case, corresponds to the maximum information gain after T valued observations Yi E y C RTH. for E f11.... .),, not scalar as it would be usually assumed in GP-TS. However, the proof of Lemma 2 in Takeno al. [33, Thm. 3.1] does not depend on the particular form of the posterior mean E(h(a) (D.] or variance V(h(a) D,), long the posterior remains GP, which still holds. Lastly, analyze the information gain.
The information gain about à = fo HG. after T observations is given by the mutual information I( h). As h is a deterministic function of Gs, by the data-processing inequality [see 57, Thm. 1.8.1].
I(v1:; h) < [(y1:;
In addition, there are at most |S] < co distinct (in distribution) random variables ya. for a E S, which allows to provide generic upper bound on the growth rate of the information gain.
Mutual information invariant under permutations. Then, for large enough T, we can rearrange the observations as:
[[y1: G.) = I((v2...... TT Jaes: G.),
where vr denotes the i'th observation with input function a S, for í {1, {1,..., NT, and Eaes RT,a = T. In addition, observation Ya independent of observation ya', for a * a' E S, when conditioned on G. (and the respective inputs a, a', which we are omitting to avoid notation clutter). By the chain rule of mutual information, then have that:
T((V..... .V......... AT.a Jass; G.)
Σ I(u.. T( Vo PT.a; G.).
RES
The summand corresponds to the information gain after nr,a repeated observations the same a,
Ily ... UT.. ".k....,97....
G.) = log de + A - -"KKa,),
where K7,a- = [H a) H HTg eigenvalues with multiplicity up to IT, and the maximum eigenvalue is bounded by the trace (a) E O(m). Therefore,
TT.a lij=1 E Ronrr Xmnr.a, Thus, K7, can have at most only m distinct
log dI + x- "KT,,) (oggog(m)),)))) O(m lT,
27

---

Combining all the equations above, we get:
MyyeT; 1))
GO ΣI( Poc...al Da T G.) EO
Em m mlgg(n,,a)) 0 (((Smmlog(T)
aES
NaES
as RT,a S T. Hence, Y1T is Olog T), given that m and |S| are fixed. As result the cumulative regret is:
Ri E O(VTTT,)) O(y m0/mm///ToogT = Ö(VT),
with O suppressing logarithmic factors, which concludes the proof.
Remark 1. Despite the result above assuming that f is only a function of G(a), there is a straight- forward extension functionals of the form f :u x A R. considered in our experiments. We simply need to replace G : A + U with the operator G' a - (G(a), a) by a concatenation with the identity map a + a, which deterministic. A similar result then is possible with minor adjustments.
C.4 Approximate posterior sampling via gradient descent
We briefly review the equivalence between posterior sampling and gradient descent when training only the last (or readout) layer of neural network under (regularized) ceast-squares loss and LeCun (or Kaiming Hc) initialization in the presence of observation noise. We will mainly combine major results from the NTK and NNGP literature [12, 35, 55] into the setting of our paper. When only the last layer is trained, the feature maps of the NTK and the NNGP coincide [12, App. D], so that we can follow an NTK type of analysis of how the loss function relates to the network's 's parameters, while the distribution of the trained network determined by the NNGP kernel. For simplicity, we focus on the case of standard, fully connected, ccalar-valued neural network, noticing that this analysis is readily extensible to the neural operator case by the techniques we use for our main results.
Random feature model. When training only the last layer of neural network, we have the following model at initialization:
ho(x) = wi (xx),
(43)
where we assume wo ~ (O, 1) for the initial weights of the readout layer, with M representing the network width, and given X, <b(xr) RM represents the output of the last hidden layer of the neural network, which consists of a random feature map Φ X + RMM under the initialization scheme. Observe that the NNGP kernel is given by:
AaNsoP (2, x²) lim ([ho(z)))n=')))== lim E( M-00
(44)
M-00 M
for any x, x' E X. Note that this is the same limit we obtain if wo N(O. 1) and (zz) is scaled - by V as the NTK parameterization [11]. Hence, to simplify our derivations, we will adopt the latter in the remainder of this subsection.
Regularized least-squares estimator. Given N data points DN f0. C X R. consider the following regularized least-squares loss:
N
1
1
) w woll?. (45) 2
EN (N(w)
W ((x,)) - 22) + ||w wol?
1 TTw w yll
2 where + === (d(()))),... ((TN)) E ERM y = [yn. ........ J E RRN Wo NO I), and À a regularization factor. We note that, in practice, due to the small initialization variance of order . the initial weights wo will be elementwise very close to zero, especially for large widths M. Therefore, we omit Wo from the regularizer in Eq. 7, as their practical effect is limited, and a simple L2 regularizer is typically efficiently implemented as a weight decay term in optimization algorithms found within modern deep learning frameworks, such as PyTorch [58].
2
The loss function in Eq. 45 is convex in w and therefore admits unique minimizer WN R
which we can derive in closed form as:
V(N(w) = = T w - y) + A(w - wo)
(46)
VEN (w)
W=WN = - (Φ T + AT) -by AWO
-
28

---

For À > 0, the matrix on the left-hand side is positive-definite and therefore invertible, then:
WN ( + A + Awo).
(47)
Suppose wo (O,,,) I). Then wx y // N. EN), where:
W = E ww yl = ( + XI) - By.
(48)
and the covariance matrix given by:
ΣN : V(wN yy V(( T + +XT) ( + Awo) y
VIA($#1 + +AI) 'wol
(49)
A²(#0" T + -1) "V[ol ($4 + A1)
>² XI)-2
where we used the fact that A(Aw] = AV[w]AT for random vector w, and we also note that V[wo y] V[wo]. given that wo is sampled independently yf
Alternative derivation. Another way of deriving the expression above is via the joint distribution between w and y. Assume = + w, + for some w. ~ (O,, I) and e ~ NN(0, I2) so that
Σy = Vyy = $ + IL The joint distribution is:
W N
p T + - TT X + AI)-1 (Φ4 Σy " + AI)
AI) -
Σy
(50)
The covariance of the joint distribution is obtained from the linear relation between w.v and as:
[ AI)
0] I
+
+ I
T
0 0
T
+
+
0
Ewxy=
I
Σy
+
I
We can see that the matrix above is non-singular and positive definite. In particular, its determinant can be derived as:
(#4 T
+ AI)
0)
2 det
X
8)
det Gew(E = det
Σy
0
I
I
+ Ey T
D
A2I
0 0
det(
+ AI) det
+
D
Dy
1
4yy1
det($aT AI) 2 det
Σ T
Ey
det(# + AI) 2 det Ey)
3 det + T
1²I
-1
1,,0' T
(et(A det ΦΦ T + XI)2
0, where the inequality holds long À and a. > 0. Conditioning on y then yields:
W = (Φ + -
(51)
and:
Σ = (Φ + AI) (.E, T T + 1 T + AI)
-
AII -II
- ΦΣ, (( +
AI)
2 + AII
(52)
In contrast, even if X = 0f note that En does not correspond to the exact posterior covariance, which can be derived as:
w N
0
1
y
1
T
T
AI
(53)
EN
V
y
=I
AI)
+XI
(54)
29

---

Predictions. For the predictive equations, note that adding and subtracting T wo to the expres sion for w n yields:
W = ($ + +1I) + Awo 0 + 4 wo
+D +bD' wo)
Ww ( XI) - ( -
(55)
= wo + 10(ppp 4 + Al) -y 4'woo)
where applied the identity (I- AB) A A(I BA) -1 Hence, letting hN (x(2)) ( we have that:
hN(x) = ho(z) + ((x) 4 ( Φ + Xy ho),
(56)
where ho- :- D wo [no(=)) R.V. In the infinite-width limit, we then have that:
hx(()) = ha(a)) kw (KN AI) -y -ho),
(57)
where we set = kyneP and adopt the standard GP notation for the kernel vector ky and matrix Kv
Underestimated variance. Now considering hu k)), we have that:
(z) yy| = (x(x) + AI)
(58)
(1) |y] = =(x,)) - (K) (KKN + AI) kN(x)
+ + KN (KN AI) -
(59)
()((,,))-ki ( + AII- kx()) Xkx(z) (x))),
where the last equality follows by adding and subtracting AI from the KNv factor in the previous quadratic term. We can then see that the predictive variance is lower than the exact GP posterior predictive variance by a factor of ) (KN + AI) - The two match when À - 0, as in Lee et al. [12]. However, for the noisy case with >> O. we have this mismatch, as it can also be observed in the results of Calvo-Ordoñez al. [35]. Similarly, for the weights posterior covariance, we have that:
ΣS \T T + +XI
_ +XI- En
A(0@) + +XI)= < +XI)
(60)
X(0 + XI) '1
(A- +1)-1 I,
which holds since Φ$ is positive semidefinite and À > 0. Hence, in the following we analyze the effect of the underestimated variance the algorithm's regret.
Effect on the regret bound. We may bound the effect of the posterior variance mismatch in the regret bound of GP-TS. Let Σ, V/w.ly] represent the exact posterior covariance matrix (cf. Eq. 54) after t 2 iterations, assuming À = of. and denote the exact and the approximate posterior, respectively, as:
P. = N(W,, E,) : E.)
(61)
(62)
Correspondingly, we set:
x² E argmax f(x) xEx
(63)
E argmax he (x) aEX
(64)
30

---

assuming f(x) = 0(x) (w,,),,,. for some w. (O, The instant regret at iteration t 2 is then:
Blf(2') fxx)] [=( f(x)) Dc-1]]
/ RM JRA f(x") -(zz) dP::-(w.)) d dPP-
E
/ (f(x*) - f(x)) -(x)))
d (w)) dP. dP:-1(w..) d
(65)
E
JRM JRas
aP;-1
dP,-1 RRM / JRM f(x)) f304) d (P.
dP
E
where we applied Hölder's inequality, noting that ʃ(x*) f(xx) > 0. Therefore, if the Radon- Nikodym derivative dP is uniformly bounded, the regret bound remains the same. In the finite- width case M the density ratio between multivariate normal distributions with the same mean gives us:
aP;
dP: dP: (w)
(E)) V det (et( exp WW) ET (e([[wwi] (illwww)) W RM
(66)
As Σ, Σ, (60), the difference between the inverses E, -,-- is positive semidefinite. The maximum is then achieved at w We, yielding:
dP:
det E,
dP'
V det (E,)
Vdet
(Σ,Σ;-)
(67)
det
I 1-111 !)
-Vdet
T + x
where we applied Sylvester's determinant identity to third line, and a standard determinant identity yields the last equality. In the fiinite-widtth limit as M co, have that & o converges to K; = Vamar(1; T, )1 j= leading us to:
d/: dP,
det K,).
(68)
Recall the definition of the maximum information gain [33, 59]:
Tt
max X. x. St 2
log det (I X X-'K,).
(69)
we assume that the GP information gain 1 log det + -x) is bounded by 7. as above, we would then have that:
dP.
< exp
(70)
d.P.
which is usually unbounded term, given that 7 is non-decreasing function of t. However, for finite domain [ < co, we trivially have that " I x) given that the largest finite subset X; X X' itself. Hence, in this case, the following holds:
dP dP:
E
S exp?'x''>
(71)
which is bounded for most practical kernels. Putting all together, we have that:
Vt N, E[f(x*) - f(x))] S ri exp ?|x1
(72)
where represents the Bayesian regret when 2, maximizes a sample from the exact GP posterior, instead its approximation. Given that 7J a finite constant, the asymptotic rates for the Bayesian cumulative regret remain the same in the presence of underestimated predictive variance.
31

---

Problem with bound. An issue with the finite bound on the Badon-Nikodie derivative above can be found when contrasting the classic definition of the maximum information gain "/ in the literature (69) with the actual information gain in the algorithm, i.e., the mutual information between an exact GP and the collected observations, which can be shown to be quantified by 1
2 log det (I + x -K) X) from the definition commonly found in the literature [31, 59], it does not necessarily follow that the actual information gain is bounded after we account for multiplicities in the eigenvalues. The algorithm is in principled allowed (and likely) to make repeated choices of the same 2₁ = H. for all at some point, for a fixed a, E X, which may or may not be the optimizer x². In the simplest case, if all of the algorithm's choices made any fixed Is E N, we have that:
[59]. The issue is that, although
log det (I (I K,,) log det (I P.)
log -It 22
log det - "(x") @ x(x"))
(73)
log det 1X "|<(""))) )13)
2 log(1 + ct),
for some constant > O. Therefore, this lower bound diverges + oo, whereas S 7|x1 < oo remains bounded, leading to contradiction of the previous conclusion in Eq. 72.
Exactly matching the posterior variance. The exact weights posterior covariance E, can be matched if, besides randomizing the initial weights, we randomize the observations by adding noise ~ (O, XI) to them at training time, following randomize-then-oppimne approach [60]. Specifically, we minimize the perturbed loss:
A 1 1 w (y, + Eill + 2 llw woll?, ē AI), wo NO..I. (74) 2
(w)
The corresponding minimizer given by:
W = 4 4 + AI) (v, + + + Awo),
(75)
whose conditional mean still matches the exact weights posterior mean:
yel = ( AI) b,y,
(76)
and whose conditional covariance now satisfies:
V w yrJ V[( . + AI) (b, + Ej) Awo) y7]
(2.0. + A1) - Awol Φ + A1)
-0, + XI) -1 D 0 + +AI
(77)
> XI- +
-E,
thereby recovering the exact posterior covariance. Nevertheless, note that the original difference in posterior predictive variance according to Eq. 59 is (z, + XI) (k(),, which is typically negligible for small values of noise variance o² X. As consequence, our cumulative regret bound remains approximately valid for NOTS, despite the slight underestimation of the posterior variance.
D Experiment details
D.I Darcy flow
Darcy flow describes the flow of fluid through porous medium with the following PDE form
vu())) = g(x), x x = ==.=(0,1)" 102 u(x) = 0, E an,
where u(x:) is the flow pressure, a(x) is the permeability coefficient and g(x)) is the forcing function.
We fix g(z) and generate different solutions at random with zero Neumann boundary conditions
32

---

on the Laplacian, following the setting in Li et al. [28], as implemented by the neural operator package [47]. In particular, for this problem, we generate a search space S with |S| 1000 data points. The divergence of f is V - oك=
afn de + By where f S - I² is a vector field f - (fa. f.). The gradient Vu (8u(a.)) Bu(2.2) where u(z, : 1t + R is scalar field. Inspired by previous works [5, 50, 61], we chose the following objective functions to evaluate the functions assuming that we aim to maximize the objective function f(-):
By
1. Negative total flow rates [50]
f(u,a)= a(x)(Vu(z) n)ds an
where O is the boundary of the domain and is the outward pointing unit normal vector of the boundary. g(æ) = -a(x)) Vu(æ) the volumetric flux which describes the rate of volume flow across unit area. Therefore, the objective function measures the boundary
outflux. Since the boundary is defined on grid, n E ([-1,0) 0), [1, Oj, [o, 1]. [0, -1]} for the left, right, top and bottom boundaries. The boundary integral can be simplified
[-a(0, (0, y) + (1,, (1,y)dy + [1aa((x (s,, 0) + a(xr, (),y(r,) 1)]dr
where Ur (z, y)
Gw 8e
My (2, (x, = du
8y
2. Negative total pressure (Eq 2.1 in [51])
f(u,g) = (1lu())ll2 (11g(=)
with 3 0 is coefficient for the forcing term g(x)). With constant g(x), the objective is simplified as - 2 Jn [u(x)]ldd.
3. Negative total potential energy [5]
f(u. = - (xx)) [(=))\u(z))] dr + / s(x))(() da
This functional corresponds to the system's 's total potential energy. It balances the energy dissipated by fluid friction (the first term) against the potential energy supplied by the uniform fluid source (the second term, where context, where the underlying physical state u is already a stable solution to the Darcy PDE, minimization of this functional over the set of permeability fields E S determines the permeability field a* that requires the minimum total energy to sustain the required fluid injection (source s 1) while maintaining zero pressure at the boundary (u 0). This effectively identifies the most hydrodynamically efficient design for the given flow constraints. This functional is related to the potential power functional in Wiker al. [5] with the difference that the latter requires estimates of the velocity field, while the simplified energy calculation above only uses the pressure field u.
is assumed). In our design optimization
D.2 Shallow Water
The shallow water equation on the rotating sphere is often used to model ocean waters over the surface of the globe. This problem can be described by the following PDE [46]:
04p
+ V ((p)
in S
{0, +00}
0t
0((p))
+ (pp 3 v)
in S2 x (0, +
0t
e= 00 on $² x (0)
4 = Po:
where the input function defined as the initial condition of the state a = (vo, soon)) with the geopotential layer depth io and the discharge (u is the velocity field), g is the Coriolis force term, and
33

---

Cumulative regret
Averaged cumulative regret across
250
B0
B0
200
NOTS
Cumulat
NOTS
150
BFO
BFO
GPTS
6 10
GPTS
100
SNOTS
10
3x
0-
50
100
150
200
250
300
50
100
150
200
250
300
Iterations
Iterations
Figure 4: Cumulative regret across trials for the Darcy flow rate optimization problem with only the last linear layer of single-hidden-iee FNO trained via full-batch gradient descent for NOTS (labeled as SNOTS). All our results were averaged over 10 independent trials, and shaded areas represent +1 standard deviation.
S denotes the surface of the 2-sphere in R The output function predicts the state function at time t: (Pet (pet). For this problem, we use a search space S with [S] = 200 data points.
As the shallow water equation is usually chosen simulator of global atmospheric variables, adopt the most common data assimilation objective [52, 53] in the weather forecast literature defined
(u, a) 2 ap. B B- a,)) + 2 u, R- ())
where ap describes the prior estimate of the initial condition, u, represents the ground truth function, the background kernel B and error kernel can be computed with historical data. The objective can be defined as inverse problem which corresponds to finding the initial condition that generates the ground truth solution function uy Here we simplify the objective by not penalizing the initial condition (dropping the prior term) and assuming independence and unit variance the solution functions using an identity kernel R), the simplified objective function f(u) 1 (u u4, u ue) be used measure different initial conditions.
D.3 Noise
To simulate real-world settings, noise added the observations by computing the empirical covariance matrix of the outputs in the dataset for the corresponding PDE and then adding Gaussian noise with variance set 1% of the coordinate-wis output variance.
D.4 Algorithm settings
NOTS was implemented using the Neural Operator library [47] and run on NVIDIA H100 GPUs on CSIRO's 'ss high-perte sersh-perfrmaccce computing cluster. For each dataset, we selected the recommended settings for FNO models according to examples in the library. Parameters were randomly initial- ized using Kaiming (or He) initialization [36] for the network weights, sampling from normal distribution with variance inversely proportional
the input dimensionality of each layer, while biases were initialized to zero. For all experiments, we trained the model for 10 epochs of mini-batch stochastic gradient descent with an initial learning rate of 10-3 and cosine annealing scheduler. The regularization factor for the 1.2 penalty was set as À == 10-1, This same setting for the regularization factor was also applied to our implementation STO-NTS.
E
Additional results with single-hidden-layer model
More closely to the setting in our theoretical results, we tested Lingle-hidden-laye FNO on the Darcy flow PDE. Only the last hidden layer of the model was trained via full-batch gradient descent. The FNO was configured without any lifting layer, having only a single Fourier kernel convolution and residual connection, as the original formulation. The number of hidden channels was set to 2048 to approximate the infinite- width limit.
34

---

Cumulative regret across
Averaged cumulative regret across trials
25
10-2
B0
20
NOTS
Curmlaatio
BFO
80
Cumulative
15
GPTS
NOTS
SNOTS
10-2
GP1S
STONTS
SNOTS
50
100
150
200
250
300
100
150
200
250
300
Reeetions
Iterations
Figure 5: Cumulative regret across trials for the Darey flow total pressure optimization problem with only the last linear layer of a Single-hidden-ee FNO trained via full-batch gradient descent for NOTS (labeled as SNOTS).
The results in Figure 4 show that the algorithm with the simpler model (SNOTS) can perform well in this setting, surpassing the performance of the original NOTS. However, in the more challenging scenario imposed by the potential power problem [adapted from 5], we note that SNOTS struggles, only achieving mid-range performance when compared other baselines, as shown in Figure 5. This performance drop suggests that the complexity of the pressure optimization problem may require more accurate predictions to capture details in the output functions that might heavily influence the potential power. In general, a quadratic objective will be more sensitive to small disturbances than a linear functional, hence requiring a more elaborate model.
F Limitations and extensions
Noise. We note that, although our result in Proposition 2 assumes a well specified noise model, should be possible to show that the same holds for noise which is mub--aussian with respect to the regularization factor. The latter would allow for configuring the algorithm with any regularization factor which is at least as large as the assumed noise sub-Gaussian parameter (i.e., its variance if Gaussian distributed). However, this analysis can quite involved and out of the immediate scope of this paper. Therefore, we leave such investigation for further research.
Nonlinear functionals. We assumed bounded linear functional in Proposition 2, which should cover variety of objectives involving integrals and derivatives of the operator's output. However, this assumption may not hold for more interesting functionals, such as some objectives considered in experiments. Similar the case with noise, any Lipschitz continuous functional of the neural operator's 's output should follow a sub-Gaussian distribution [62]. Hence, the Gaussian approximation remains reasonable, though more in-depth analysis would be needed to derive the exact rate of growth for the cumulative regret in these settings.
Mult-layer models. For the theoretical analysis, we assumed single hidden layer neural net- work as the basis of our Thompson sampling algorithm. While this choice provides simple and computationally efficient framework, it may not be optimal for all applications or datasets. For instance, in some cases, deeper neural network with more layers might provide better performance due to increased capacity to capture complex patterns in the data. Extending our analysis to this setting involves extending the inductive proofs for the multi-layer NNGP [38, 54] to the case of neural operators. Such extension, however, may require transforming the operator layer's output back into function in infinite-dimensoon the possibility of a kernel limit [55]. In the ungle-hidden-layer case, such effect is avoided by operating directly with the finite-dimensional input function embedding Ar) Rda, Recently, concurrent work has explored the infinite-win width limit for multi-layer neural operators [63], but their applicability NOTS is left as subject of future work.
space, which may lead to bottleneck effect affecting
Prior misspecification. which was also considered to be infinitely wide. While this assumption greatly simplifies our analysis, more practical results may be derived by considering miith neural operators and true operator
We assumed that the true operator G, follows the same prior as our model,
35

---

which might not exactly correspond to realization of the chosen class of neural operator models. For the case of finite widths. one simple way to obtain a similar regret bound is let the width of the network grow at each Thompson sampling iteration. The approximation enror between the GP model and the finite width neural operator can potentially be bounded as (M [55]. Hence if the sequence of network widths {M]]== is such that Σ V.T < oo, similar regret bound the one in Proposition 2 should possible. Furthermore, if other forms of prior misspecificatiion need to be considered, analyzing the Bayesian cumulative regret (instead of the more usual frequentist regret), as we did, allows one to bound the resulting cumulative regret of the misspecified algorithm via the Radon-Nikodym derivative dP
the true prior P with respect to the algorithm's 's prior probability
aP
measure P. If its essential supremum
dP is bounded, then the resulting cumulative regret 2o
remains proportional to the same bound derived as if the algorithm's prior was the correct one [8].
G Broader impact
This work primarily focuses on the theoretical exploration of extending Thompson sampling to function spaces via neural operators. As such, it does not directly engage with real-world applications or present immediate societal implications. However, the potential impact of this research lies in its application. By advancing methods for function-space optimization, this work may indirectly contribute to various fields that utilize complex simulations and models, such as climate science, engi- neering, and physics. Improvements in computational efficiency and predictive power in these fields could lead to positive socictal outcomes, such as better climate modeling or engineering solutions. Nevertheless, any algorithm with powerful optimization capabilities carries ethical considerations. Its deployment in domains with safety-critical implications must be approached with care to avoid misuse unintended consequences. Researchers and practitioners should ensure transparency, faimesss, and accountability in applications potentially affecting society.
36


Reviews:
[{"review_id": "vm51XMGsQ9", "reviewer": "Reviewer_Gc9e", "summary": "The paper proposed the Neural Operator Thompson Sampling (NOTS) Algorithm to solve the optimization problem in function space. The specific optimization objective is the output of a known functional of an unknown operator. The main tool in use is the neural network Gaussian process (or conjugate kernel). More specifically, the paper uses the famous result that neural networks with a single hidden layer with infinite width acts like a Gaussian process. The authors provide the regret bound adapted from the existing result on Gaussian process Thompson sampling under the assumption of linear functional and a single-hidden-layer neural network. Besides, experimental results on 2D Darcy flow and a shallow water model are presented to show the advantage of the proposed NOTS algorithm.", "questions": "1. The good flow potential figure also has open porus in the upper right and lower right corner. Can the authors provide some explanations to it?\n\n2. Could you rephrase line 195 to line 196. It is hard to understand.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strength:\n\nThe author utilize neural operator to extend the theory to more practical problem with high dimension or infinite dimension. The structure of the paper is good, and the literature review is thorough. The list of benchmark algorithm is relatively complete.\n\nWeakness:\nSome parts of the proof and discussion lack clarity:\n   1. Proposition C.1 and Proposition 1 are not consistent. \n   2. It is unclear what neural network structure is used in the simulation of NOTS algorithm in the main paper, as the appendix specify the simulation is run under the single hidden layer NN. \n   3. The simulation is not quite reproducible as some basic simulation setup is unclear. Also, in some of the figures, the simulation has fewer iterations, which I think deserve an explanation. \n   4. In line 288, it says that 'STO-NTS thrives', but in Figure 2, STO-NTS does not have low regret. \n\nFor evaluation, it will be more convincing if the regret bounds of the baseline algorithms are also shown explicitly in a table in the experiment section.", "was_revised": true, "final_justification": "The authors have addressed my concerns during the discussion. I recommend the paper to be accepted tentatively. I am unsure about the originality of the paper since the techniques in use are not that novel.", "rebuttal": "We would like to thank the reviewer for their feedback.  Below we address the issues raised by the review.\n\n## Weaknesses / Points for Clarification\n\n**Theoretical Inconsistency.**\n\"Proposition C.1 and Proposition 1 are not consistent.\"\n\n**Our Response:**\nProposition C.1 (appendix) is a more complete restatement of Proposition 1. There we clarify assumptions on the weights initialization distribution which were unfortunately misspecified in the main text.  In both the theoretical analysis and the code implementation, the weights were initialized following Gaussian Kaiming He initialization.  This scheme applies a variance which is inversely proportional to the previous layer's width, which corresponds to a variance of 1 for the first layer's weights and $1/M$ for the subsequent layers.  We will update this in the final version of the paper.\n\n**Experimental Clarity (NN Structure).**\n\"It is unclear what neural network structure is used in the simulation of NOTS algorithm in the main paper, as the appendix specify the simulation is run under the single hidden layer NN.\"\n\n**Our Response:**\nFor the experiments in the main paper, we used a multi-layer Fourier neural operator (FNO) architecture following the default settings from the examples in the `neuraloperator` package for their respective datasets.  In the appendix, however, we demonstrated that even with a single hidden layer and training only the weights of the readout layer, as suggested by our theoretical results, the proposed framework still achieves competitive performance.  One of the main objectives of this paper was to introduce a framework which is simple to implement with existing toolboxes and yet effective to tackle optimization problems in function spaces with little need for fine tuning.\n\nFNO settings for main paper experiments:\n* Darcy flow: 16 by 16 Fourier modes (2D spatial field), 64 hidden channels;\n* Shallow water: 32 by 32 Fourier models, 64 hidden channels.\n\n**Reproducibility.**\n\"The simulation is not quite reproducible as some basic simulation setup is unclear.  Also, in some of the figures, the simulation has fewer iterations, which I think deserve an explanation.\"\n\n**Our Response:** The baselines that run for fewer iterations stopped due to the code crashing due to out-of-memory issues.  One important factor to note is that GP-based BO baselines need to keep all the data they observe to form, update and predict with their GP models.  As the inputs are quite high dimensional (e.g., the shallow water problem had 6144 input dimensions for each data point), and GPs need to keep their tensors at high numerical precision ( ideally 64-bit double-precision floats), GP-based algorithms can quickly run into memory issues.  In contrast, neural operator baselines only need to keep their weights for prediction after training.  The experiments code containing the full simulation setup will be released with the camera-ready version of the paper if accepted.  We will add these clarifications to the revised text.\n\n**Results Interpretation.**\n\"In line 288, it says that 'STO-NTS thrives', but in Figure 2, STO-NTS does not have low regret.\"\n\n**Our Response:** Indeed, in Fig. 2, STO-NTS does not achieve low regret.  However, in line 288, we were referring to Fig. 1, where it achieves lower regret when compared to the other baselines, though still not as low as NOTS.  Hence, we will revise our claims on STO-NTS's performance.  In general, noting that STO-NTS uses a default fully-connected neural network architecture, as suggested by its original paper, without further adaptations to deal with high-dimensional spatial grid data, we believe that STO-NTS's poor performance is due to architectural limitations, which then lead to a need for higher amounts of data (i.e., more iterations) for its performance to become significantly better than the other baselines.  We will revise our statements on STO-NTS.\n\n**Suggestion on Baselines.**\n\"it will be more convincing if the regret bounds of the baseline algorithms are also shown explicitly in a table in the experiment section.\"\n\n**Our Response:** We will add a table to the revised version of the paper.\n\n|Problem | RS                | BO                | NOTS              | BFO               | GP-TS              | STO-NTS            |\n|--------|-------------------|-------------------|-------------------|-------------------|-------------------|-------------------|\n| Darcy flow - rates  | 0.872 +/- 0.022 | 0.703 +/- 0.045 | **0.012 +/- 0.001** | 0.788 +/- 0.066 | 0.674 +/- 0.050 | 0.068 +/- 0.002 |\n| Darcy flow - potential | 0.456 +/- 0.012 | 0.240 +/- 0.031 | **0.099 +/- 0.060** | 0.462 +/- 0.062 | 0.130 +/- 0.021 | 0.4 +/- |\n| Darcy flow - pressure  | 0.077 +/- 0.001 | 0.047 +/- 0.001 | **0.012 +/- 0.001** | 0.078 +/- 0.006 | 0.038 +/- 0.004 | 0.068 +/- 0.002 |\n| Shallow water - inversion  | 4.632 +/- 0.876 | 1.639 +/- 0.532 | **0.134 +/- 0.043** | 3.076 +/- 0.886 | 1.942 +/- 0.502 | 2.329 +/- 0.800 |\n\n## Questions\n\n**Question 1: \"Good Flow Potential\" Figure.**\n\"The good flow potential figure also has open porus in the upper right and lower right corner.  Can the authors provide some explanations to it?\"\n\n**Our Response:**\nThe good flow potential figure was selected from 1,000 input-output pairs we have generated in the experiment;  therefore, the range of the functional was bounded by this pool and the global optimal value is not necessarily 0 which would ideally be achieved by closed pores.  The functional for this best candidate is -1.17 as opposed to the worst functional of -2.43, and indeed the flow rate of this figure was contributed by the open pores in the upper right and lower right corner.\n\n**Question 2: Rephrasing Lines 195-196.**\n\"Could you rephrase line 195 to line 196. It is hard to understand.\"\n\n**Our Response:** We thank the reviewer for pointing this out.  We will rephrase it to: ``where the expectation is over the function space prior for $f$ and all other random variables involved in the decision-making process.'' We simply meant that the expectation in Bayesian regret is a \\emph{total} expectation, involving all underlying random variables.\n\n### Formatting Concerns\n\n**Typo (Line 230).**\n\"Line 230: typo. deadline -> baseline\"\n\n**Our Response:** Thanks for noticing that.  We will fix the typo in the revision.\n\n**Unclear Reference (Line 194).**\n\"Line 194: the authors use the wording 'as previously discussed' without specifying the exact location where the result was given before.  This makes the paper hard to follow.\"\n\n**Our Response:** We unfortunately were unable to find the text the reviewer mentioned in line 194. The closest match was ``As previously described...\" in line 289, which indeed can be somewhat unclear. It should be referring to Sec. 6.2 where we briefly specify the Darcy flow problem. More complete details about that problem's settings can be found in the appendix. We will revise this and other potentially unclear references in the text and make sure we specify section numbers in the main text and the appendix when referring to them."}, {"review_id": "iCTZp11Cuw", "reviewer": "Reviewer_28oU", "summary": "The paper proposes Neural operator Thompson sampling, a novel method for applications which require optimization wrt an unknown functional operator which may be too expensive to evaluate using simulation. To do so, a neural operator is trained from a number of samples. The resulting operator is treated as the correct posterior consistent with the input data and is sampled to generate new input functions. The method uses the fact that for wide neural operators these samples behave like samples from Gaussian processes.", "questions": "Having to retrain seems like a major limitations. How do other methods compare here?", "limitations": "Yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n- introduces a novel method for Thompson sampling using Neural Operators\n- this method avoids explicit uncertainty quantification, by random reinitialization and training which acts like sampling from the posterior\n- theoretical guarantees provide convergence and regret bounds in the infinite width limit\n- good evaluation with baselines, although I am unfamiliar with this subdomain to assess whether these are the right baselines\n\n**Weaknesses**\n\n- theoretical guarantees are limited to discretized input domains\n- retraining the operator at each step from scratch is expensive\n- limited evaluation on toy examples", "was_revised": true, "final_justification": "The authors were able to answer my questions and concerns. However as I have to admit that I am not particularly knowledgable about this domein, I have decided to lower my  confidence.", "rebuttal": "We'd like to thank the reviewer for their valuable feedback. Below we address the issues raised.\n\n### Weaknesses\n\n**Theoretical Guarantees.**\n\"theoretical guarantees are limited to discretized input domains\"\n\n**Our Response:**\nBesides simplifying some aspects of the derivation of regret bounds, we attached our theoretical analysis for Proposition 2 to the setting where function-valued inputs are evaluated over discretized domains as that is usually the case with neural operator models, especially FNO.  One cannot implement truly infinite-dimensional input/output data on a computer, leading to a need for discretization or parameterization via a finite number of parameters, so that the observation data are finite dimensional.  We, highlight, however, that the result in Proposition 1 (GP limit) still holds for infinite-dimensional inputs and outputs.\n\n**Computational Cost.**\n\"retraining the operator at each step from scratch is expensive\"\n\n**Our Response:**\nAnswered below.\n\n**Evaluation Scope.**\n\"limited evaluation on toy examples\"\n\n**Our Response:**\nThe scope of our experiments is in line with similar accepted works in the machine learning literature [16,25,29].  Future investigations into data from other literature, such as in [15], are possible.\n\n### Questions\n**Question 1: Comparison of Retraining Cost.**\n\"Having to retrain seems like a major limitations.  How do other methods compare here?\"\n\n**Our Response:** The retraining of the neural network (NN) model from scratch is a common implicit requirement in other NN-based Thompson sampling baselines, such as STO-NTS and the original neural Thompson sampling.  This is to ensure we have a new posterior sample at every iteration, which is part of the algorithmic design for Thompson sampling.  If training only the readout layer of the neural network, as in our theoretical analysis, one could, however, keep the other layers fixed at an initial random initialization kept throughout iterations, similar to other random feature models, such as sparse spectrum Gaussian processes;  training thus requires just regularized linear regression. In this case, one can perform incremental rank-1 updates to the model weights probability distribution (Gijsberts & Metta, 2003) without having to retrain the whole model from scratch.  However, keeping the untrained layers fixed, also limits the range of the posterior samples, reducing the algorithm's exploration, leading to worse performance, especially in practical finite-width settings where the hidden layers might not be wide enough to get the model closer to GP behavior.\n\n### References\n* Gijsberts, A., & Metta, G. (2013).  Real-time model learning using Incremental Sparse Spectrum Gaussian Process Regression. Neural Networks, 41, 59–69.  \\url{[https://doi.org/10.1016/j.neunet.2012.08.011](https://doi.org/10.1016/j.neunet.2012.08.011)}"}, {"review_id": "Cn5Izvn43N", "reviewer": "Reviewer_9PFn", "summary": "The paper presents an extension of Thompson sampling for use in neural operators. The proposed method avoids the computation of uncertainty. It is connected to active learning via sampling more effective data points which are used to train neural surrogates to solve PDEs. Experiments are performed on Darcy flow and shallow water modelling problems, via their respective PDEs. It is claimed that the proposed method (NOTS) helps to encode the compositional structure of problems and improves performance.", "questions": "1) Can authors clarify what is the main goal/ motivation behind this work? \n\n2) Where exactly is Thompson sampling coming into play? It feels like using a term without any necessity. \n\n3) Since, the goal of the paper is not clear; hence the evaluation metric presented is also not justified. The only metric of comparison is cumulative regret. Based on equation (16) the computation of such a quantity requires the use of a*. What is used to compute this metric? \n\n4) What is the set (S) over which the optimization is performed to choose a_t? How is this chosen? \n\n5) What is the rationale behind the choice of optimization functionals? \n\n6) Please explain and discuss the results in the paper more thoroughly, it will help appreciate the method. Can you introduce Thompson sampling in Section 2?", "limitations": "Limitations have been discussed.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength: \n\n1. The relevant literature has been discussed thoroughly. \n2. The problem formulation including neural PDE solvers and Gaussian process are introduced sufficiently. \n3. Theoretical results claimed in the paper are reasonable. \n\nWeakness: \n\n1. The point of the paper is absolutely unclear. The problem and the proposed solution have not been motivated appropriately. \n2. The evaluation of the method is not convincing. \n3. One of the main components of the proposed method are the optimization functional, the have been used in an ad-hoc manner. There is no justification and they vary from one problem to another, without any context. \n\nPlease see questions for more details.", "was_revised": true, "final_justification": "Based on the authors rebuttal which has satisfied my concerns, I have updated the score", "rebuttal": "We would like to thank the reviewer for the constructive feedback and helpful criticism of our paper.  We address the main concerns in our response below.\n\n### Weaknesses\n\n**Motivation & Problem Formulation.**\n\"The point of the paper is absolutely unclear.  The problem and the proposed solution have not been motivated appropriately.\"\n\n**Evaluation.**\n\"The evaluation of the method is not convincing.\"\n\n**Optimization Functionals.**\n\"One of the main components of the proposed method are the optimization functional, the have been used in an ad-hoc manner.  There is no justification and they vary from one problem to another, without any context.\"\n\n### Questions\n\n**Question 1: Main Goal/Motivation.**\n\"Can authors clarify what is the main goal/ motivation behind this work?\"\n\n**Our Response:**\nOne of our main motivations was to propose a method to optimize functionals involving simulation outputs.  With the rise of neural operators, we now have access to powerful data-driven emulators that can accurately mimic PDE solvers.  Thompson sampling then becomes an appealing approach to include such ML models within a Bayesian optimization framework, given the results in NTK theory and related areas studying overparameterized neural networks.  In the standard neural operator setting, we can only predict the output solution given the input, but we may have no idea about which input would achieve the best objective of interest, and this question can be answered by Thompson sampling in the Bayesian optimization framework.  For instance, in our Darcy flow example in section 6.2, a good neural operator can precisely predict the output flow given the input permeability function, but we still do not know which permeability function results in the minimum sinkage (our objective), and our framework indicates the best possible input function (optimal permeability function) to achieve this objective (minimum sinkage).  Lastly, we note that, in general, the objective functional will be problem dependent, and we assume that the most expensive step is the forward simulation which the neural operator is learning to emulate, not the evaluation of the objective functional on the simulated/emulated output.\n\n**Question 2: Role of Thompson Sampling.**\n\"Where exactly is Thompson sampling coming into play?  It feels like using a term without any necessity.\"\n\n**Our Response:**\nWe show that training a neural operator from randomly initialized weights is approximately equivalent to Thompson sampling, which is a well known algorithm in the Bayesian optimization (BO) and multi-armed bandits literature, allowing us to derive theoretical performance guarantees for NOTS.  The start of our section 3.2 is reference [25], and therein the introduction gives a description of Thompson sampling that is relevant to our context.  We will include a similar introduction in our supplementary material.  A more classical introduction to Thompson sampling can be found in the works of Russo and Van Roy [42].\n\nClassic Thompson sampling, as other algorithms for BO, can be seen as trying to estimate the (global) optimum $x^* \\in \\arg\\max_{x\\in\\mathcal{X}} f(x)$ of a (black-box) function $f:\\mathcal{X}\\to\\mathbb{R}$ within a limited budget of (noisy) evaluations.  In the simple case of Bayesian linear regression surrogates $\\hat{f}_t(x) = \\mathbf{w}_t ^\\top \\phi(x)$, $\\mathbf{w} _t \\in \\mathbb{R}^M$, for a scalar-valued objective $f: \\mathcal{X} \\to \\mathbb{R}$, we assume $f(x) = \\mathbf{w} ^\\top \\phi(x)$, for some $\\mathbf{w} \\sim \\mathcal{N}(\\mathbf{0}, \\mathbf{I})$, Thompson sampling proceeds as follows, for $t \\in \\{1, \\dots, T\\}$:\n1.  Sample $\\mathbf{w}_t \\sim N( \\mathbf{\\hat{w}}, \\Sigma )$, where $\\mathbf{\\hat{w}} := \\mathbb{E}[\\mathbf{w} | D ]$, $\\Sigma := \\mathbb{V}[\\mathbf{w} | D]$;\n2.  Optimize $x_t \\in \\arg\\max_{x \\in \\mathcal{X}} \\mathbf{w}_t^\\top \\boldsymbol{\\phi}(x)$;\n3.  Observe $y_t = f(x_t) + \\epsilon_t$;\n4.  Update the observations dataset $D \\leftarrow D \\cup \\lbrace x_t, y_t \\rbrace$;\n5.  Increment $t \\leftarrow t + 1$, and repeat while $t \\leq T$;\n\nwhere $\\mathbf{\\hat{w}}$ and $\\Sigma$ represent the posterior mean vector and covariance matrix, respectively, and $D$ denotes the set of observations.  The main difference with respect to NOTS and neural Thompson sampling (NTS) algorithms is that step 1 is replaced with optimization of a regularized least-squares loss (Eq. 9 for NOTS) via gradient descent, which yields an approximate posterior sample.  In STO-NTS, the feature map $\\phi: \\mathcal{X} \\to \\mathbb{R}^M$ above also has a natural correspondence with the output of the last hidden layer of a randomly initialized neural network, which has only the readout (last) layer trained.\n\n**Question 3: Justification of Cumulative Regret.**\n\"Since, the goal of the paper is not clear;  hence the evaluation metric presented is also not justified. The only metric of comparison is cumulative regret.  Based on equation (16) the computation of such a quantity requires the use of a\\*.  What is used to compute this metric?\"\n\n**Our Response:**\nRegret is the standard choice of metric in the Thompson sampling, bandits and Bayesian optimization literature to analyze the theoretical performance of an optimization algorithm.  It is usually not possible to directly measure regret in practical applications, as we might not have access to the true optimum of the objective, as the reviewer has correctly pointed out.  However, for the purposes of this paper, and similar to other BO/bandits papers, our examples with synthetic data allow us to directly measure regret for performance comparisons, since we have access to the true noise-free functional and can evaluate it over the entire search space. In the revised version of the paper, we will clarify the role of regret and include additional non-regret plots (e.g., the maximum value observed up to each iteration), which might be more meaningful for readers less familiar with the bandits/BO literature.\n\n**Question 4: Optimization set $\\mathcal{S}$.**\n\"What is the set $\\mathcal{S}$ over which the optimization is performed to choose a\\_t?  How is this chosen?\"\n\n**Our Response:**\nThe set $\\mathcal{S}$ represents the search space which the algorithm has access to.  It is a subset of the domain $\\mathcal{A}$ of the true (unknown) operator $G*$ , domain which is shared with the neural operators $G^*$.  For instance, in the Darcy flow experiments, the domain $\\mathcal{A}$ is the set of all functions mapping the 2D unit box $[0,1]^2$ to the non-negative real line $\\mathbb{R}_+$, whereas the search space $\\mathcal{S}$ is a random finite subset of the domain, generated by sampling binary values over 2D grids within the unit square.\n\n**Question 5: Rationale for Optimization Functionals.**\n\"What is the rationale behind the choice of optimization functionals?\"\n\n**Our Response:**\nOur framework is agnostic to the choice of objective functional, which could even be, e.g., application-specific black-box code run on top of a simulator output, as long as these functionals can be (cheaply) evaluated on the output of the true operator (i.e., the simulator or real process we are trying to emulate) and on the output of the neural operator.  Given the specific modeling goal of a PDE, the objective can vary.  For example, for the Darcy flow problem, one of our objectives is to find an optimal topological design of the input function which minimizes the total sinkage (Jeong, S. et al., 2025).  We therefore adopted the objective in the literature to measure the total pressure.  Contrastingly, for the shallow water equation problem, we wonder how to trace back the initial condition of a specific output function, e.g., a hurricane, as this PDE represents a weather forecast problem (Rabier, F. et al, 1998), and we therefore formulated an inverse problem as the objective.\n\n**Question 6: Deeper Explanation of Results.**\n\"Please explain and discuss the results in the paper more thoroughly, it will help appreciate the method.  Can you introduce Thompson sampling in Section 2?\"\n\n**Our Response:**\nWith the additional page allowed for the camera-ready version of the paper, we plan to add a brief overview of Thompson sampling in the background section (Sec. 2) and a more in-depth discussion of the experimental results.  In particular, we will extend the discussion on the results of Fig. 2 and 3, highlighting that GP-TS and BO with the log-expected improvement acquisition function, both equipped with the infinite-width BNN kernel, present improved performance in comparison with the other baselines, confirming the observations of Li et al.  [40] that this choice of GP covariance function yields competitive results in high-dimensional settings."}, {"review_id": "0vBthf9ygD", "reviewer": "Reviewer_aMDd", "summary": "The paper studies the problem on black-box optimization of a function composed of a known functional with a black-box unknown operator. It proposes a method based on neural operator and Thompson sampling.", "questions": "1. It is assumed that the search space is compact. Does it mean that essentially the searched function is parameterized.\n2. In algorithm 1, I feel the notation $\\arg\\min$ is a bit unsuitable. Indeed, you can usually not gurantee training it to global optimal with neural network due to non-convex optimization.", "limitations": "It is not so clear to me that why neural network should be the solution for this problem, since it can be computationally expensive and requires lots of data.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths: The paper is overall well-written and clear. The studied problem is important and interesting. The proposed solution has theoretical guarantee. Experimental results show that it outperforms other baseline methods.\nWeaknesses: From my understanding, training a neural network usually requires lots of data and computation resource. But black-box optimization typically has limited data. Then what is the advantage of using NN instead of, e.g., Gaussian process based Thompson sampling.", "was_revised": false, "final_justification": "", "rebuttal": "We'd like to thank the reviewer for their positive feedback and insightful comments.  We address the reviewer's questions below.\n\n### Weaknesses / Limitations\n**NN vs. Gaussian Process.**\n\"what is the advantage of using NN instead of, e.g., Gaussian process based Thompson sampling.\" And states: \"It is not so clear to me that why neural network should be the solution for this problem, since it can be computationally expensive and requires lots of data.\"\n\n**Our Response:** Although GPs typically perform well for Bayesian modeling tasks with low volumes of data, we have high-dimensional data in both inputs and outputs of the model, which render the application of traditional multi-output GP models challenging.  For the shallow water PDE, for example, both inputs and outputs lie in 6144-dimensional space.  Hence, without specialized kernels to capture spatial data on both inputs and outputs, conventional models would struggle due to both the challenges of modeling high-dimensional input data and the elevated computational cost of vector-valued GP inference with very large number of outputs.  In contrast, neural operators are specially designed to deal with function-valued input and output data, typically over spatial domains.  In addition, their training cost only grows linearly with the number of data points, instead of cubically (as in GPs).  Hence, NOTS is more scalable to accommodate longer runs or extensions with batched evaluations than traditional GP solutions, though we considered only up to 300 iterations in our experiments to allow for comparisons against GP baselines.\n\n### Questions\n\n**Question 1: Compact Search Space.**\n\"It is assumed that the search space is compact.  Does it mean that essentially the searched function is parameterized?\"\n\n**Our Response:** Some sort of parameterization is required to allow a computational method to search through the space.  In our case, we simply assumed the input functions and the observations are discretized on a grid.  Compactness then simply means that the values the functions can take on the grid cells are bounded.  Note, however, that this is always the case when the search space is finite.\n\n**Question 2: Notation in Algorithm 1.**\n\"In algorithm 1, I feel the notation is a bit unsuitable.  Indeed, you can usually not guarantee training it to global optimal with neural network due to non-convex optimization.  arg min\"\n\n**Our Response:** In Algorithm 1, we use the $\\arg\\min$ notation to denote the minimizer obtained by gradient descent;  we will add a remark on this relaxed use of the notation.  Nonetheless, although optimization landscapes for neural networks are usually non-convex, plenty of recent works have shown that, in the overparameterized regime, (stochastic) gradient descent converges to global minima (Allen-Zhu et al., 2019).  In addition, in the infinite-width limit, the loss function is a convex quadratic objective with respect to the network parameters, leading to a unique global minimum.\n\n### References\n* Allen-Zhu, Z., Li, Y., & Song, Z. (2019).  A convergence theory for deep learning via over-parameterization. Proceedings of the 36th International Conference on Machine Learning (ICML 2019).\n* Jeong, S., & Lee, S. (2025). Optimal control for Darcy's equation in a heterogeneous porous media.  Applied Numerical Mathematics, 207, 303-322.\n* Rabier, F., Thépaut, J. N., & Courtier, P. (1998).  Extended assimilation and forecast experiments with a four‐dimensional variational assimilation system. Quarterly Journal of the Royal Meteorological Society, 124(550), 1861-1887."}]

Author rebuttal:

Meta-review: **Summary:** The paper proposed the Neural Operator Thompson Sampling (NOTS) to solve optimization problems in function space. The specific optimization objective is the output of a known functional of an unknown operator. The paper presents a regret bound derived from an existing result on Gaussian process Thompson sampling, assuming a linear functional and a single-hidden-layer neural network. Experimental results on 2D Darcy flow and a shallow water model are presented to show the advantage of the proposed algorithm.

**Strengths:** 
- Well-written paper with a good related work section and a clear contribution
- Novel combination of Thompson sampling with neural operators
- Theoretical guarantees in the infinite-width limit
- Good experimental evaluation. 

**Weakness:** 
- Computational cost of having to retrain the network. 

**Decision and Suggested Changes:** The reviewers agree that the paper is well-written, and its contributions merit acceptance. After carefully reading the rebuttal/discussion, I tend to agree. Please incorporate the following comments in the final version of the paper. In particular, addressing the following concerns will help strengthen the current version of the paper:
- Clarify the theoretical inconsistency and add the experimental details/regret results for baseline algorithms (in response to Rev. Gc9e)
- Quantify and explicitly mention the computational cost to retrain (Rev. 28oU)
- In order to reduce the computational cost, investigate using SGLD to generate samples without complete retraining. 
- Include an overview for TS in the background section (Rev. 9PFn)
- Include a justification as to why GPs cannot be directly used instead of neural networks (Rev. aMDd)
Decision: Accept (poster)