# Continuum Attention for Neural Operators

Edoardo Calvello<EMAIL_ADDRESS>
California Institute of Technology Nikola B. Kovachki<EMAIL_ADDRESS>
NVIDIA Matthew E. Levine<EMAIL_ADDRESS>
The Broad Institute of MIT and Harvard Andrew M. Stuart<EMAIL_ADDRESS>
California Institute of Technology

###### Abstract

Transformers, and the attention mechanism in particular, have become
ubiquitous in machine learning. Their success in modeling nonlocal, long-range
correlations has led to their widespread adoption in natural language
processing, computer vision, and time-series problems. Neural operators, which
map spaces of functions into spaces of functions, are necessarily both
nonlinear and nonlocal if they are universal; it is thus natural to ask
whether the attention mechanism can be used in the design of neural operators.
Motivated by this, we study transformers in the function space setting. We
formulate attention as a map between infinite dimensional function spaces and
prove that the attention mechanism as implemented in practice is a Monte Carlo
or finite difference approximation of this operator. The function space
formulation allows for the design of transformer neural operators, a class of
architectures designed to learn mappings between function spaces, for which we
prove a universal approximation result. The prohibitive cost of applying the
attention operator to functions defined on multi-dimensional domains leads to
the need for more efficient attention-based architectures. For this reason we
also introduce a function space generalization of the patching strategy from
computer vision, and introduce a class of associated neural operators.
Numerical results, on an array of operator learning problems, demonstrate the
promise of our approaches to function space formulations of attention and
their use in neural operators.

Keywords: Operator Learning, Attention, Transformers, Discretization
Invariance, Partial Differential Equations.

##  1 Introduction

Attention was introduced in the context of neural machine translation in
Bahdanau et al. (2014) to effectively model correlations in sequential data.
Many attention mechanisms used in conjunction with recurrent networks were
subsequently developed, for example long short-term memory (Hochreiter and
Schmidhuber, 1997) and gated recurrent units (Chung et al., 2014), which were
then employed for a wide range of applications, see for example the review
Chung et al. (2014). The seminal work Vaswani et al. (2017) introduced the now
ubiquitous transformer architecture, which relies solely on the attention
mechanism and not on the recurrence of the network in order to effectively
model nonlocal, long-range correlations. In its full generality, the attention
function described in Vaswani et al. (2017) defines an operation between
sequences
$\\{u_{i}\\}_{i=1}^{N}\subset\mathbbm{R}^{d_{u}},\\{v_{i}\\}_{i=1}^{M}\subset\mathbbm{R}^{d_{v}}$
of lengths $N,M\in\mathbbm{N}$ respectively, represented as tensors
$u\in\mathbbm{R}^{N\times d_{u}},v\in\mathbbm{R}^{M\times d_{v}}$, where
learnable weights $Q\in\mathbbm{R}^{d_{K}\times
d_{u}},K\in\mathbbm{R}^{d_{K}\times d_{v}},V\in\mathbbm{R}^{d_{V}\times
d_{v}}$ parametrize

${\tt attention}(u,v):={\tt softmax}(uQ^{T}Kv^{T})vV^{T}.$ (1)

In this context, we recall that the function ${\tt
softmax}:\mathbbm{R}^{N\times M}\to\mathbbm{R}^{N\times M}$ is defined
component-wise via its action on the input as

$\bigl{(}{\tt
softmax}(w)\bigr{)}_{ij}=\frac{\exp\\{w_{ij}\\}}{\sum_{j=1}^{M}\exp\\{w_{ij}\\}},$
(2)

for any $w\in\mathbbm{R}^{N\times M}$. Therefore, in this form the attention
function defines a mapping ${\tt attention}:\mathbbm{R}^{N\times
d_{u}}\times\mathbbm{R}^{M\times d_{v}}\to\mathbbm{R}^{N\times d_{V}}$. 111We
note that in most implementations, the input of the softmax function is
rescaled by a factor of $1/\sqrt{d_{K}}$. In our context, continuum attention,
a different scaling will be used. Alternatively, we can view the output of the
attention function as a sequence $\\{{\tt
attention}(u,v)_{i}\\}_{i=1}^{N}\subset\mathbbm{R}^{d_{V}}$.

Viewed more generally, attention defines a mapping from two sequences
$\\{u_{i}\\}_{i=1}^{N}$ and $\\{v_{i}\\}_{i=1}^{M}$ into a third sequence
$\\{{\tt attention}(u,v)_{i}\\}_{i=1}^{N}$; all three sequences take values in
$\mathbbm{R}^{r}$ for various values of $r$. In this paper, we generalize this
in two ways: (a) we replace index $i$ by index $x\in D$, where $D$ is a
bounded open subset in $\mathbbm{R}^{d}$; (b) we retain a discrete index $i$
but allow the sequences to take values in Banach spaces of
$\mathbbm{R}^{r}$-valued functions defined over $D$, where $D$ is a bounded
open subset in $\mathbbm{R}^{d}$. We find the first generalization (a) most
useful in settings where $d=1$, that is for time-series. However, when
discretizing the index set $D\subset\mathbbm{R}^{d}$ with $N\coloneqq n^{d}$
points, the quadratic scaling of the attention mechanism in $N$ can be
prohibitive. This motivates generalization (b) where we retain the discrete
index but allow sequences to be function valued, corresponding to patching of
functions defined over $D\subset\mathbbm{R}^{d}$, with $d\geq 2.$

Our motivation for these generalizations is operator learning and, in
particular, design of discretization invariant neural operator architectures.
Such architectures disentangle model parameters from any particular
discretization of the input functional data; learning is focused on the
intrinsic continuum problem and not any particular discretization of it. This
allows models to obtain consistent errors when trained with different
resolution data, as well as enabling training, testing, and inference at
different resolutions. We consider the general setting where
$D\subset\mathbbm{R}^{d}$ is a bounded open set and
$\mathcal{U}=\mathcal{U}(D;\mathbbm{R}^{d_{u}})$ and
$\mathcal{Z}=\mathcal{Z}(D;\mathbbm{R}^{d_{z}})$ are separable Banach spaces
of functions. To train neural operators we consider the following standard
data scenario. The set $\\{u^{(j)},z^{(j)}\\}_{j=1}^{J}$ of input-output
observation pairs is such that $\bigl{(}u^{(j)},z^{(j)}\bigr{)}\sim\nu$ is an
i.i.d. sequence distributed according to the measure $\nu$, which is supported
on $\mathcal{U}\times\mathcal{Z}$. We let $\mu$ and $\pi$ denote the marginals
of $\nu$ on $\mathcal{U}$ and $\mathcal{Z}$ respectively, so that
$u^{(j)}\sim\mu$ and $z^{(j)}\sim\pi$. In what follows, we consider
$\mathsf{G}^{\dagger}:\mathcal{U}\to\mathcal{Z}$ to be an arbitrary map so
that

$z^{(j)}=\mathsf{G}^{\dagger}\bigl{(}u^{(j)}\bigr{)},$ (3)

noting that such maps will typically be both nonlocal (with respect to domain
$D$) and nonlinear. We assume for simplicity that the data generation
mechanism is compatible with (3) so that
$\pi=\mathsf{G}^{\dagger}_{\sharp}\mu$. Neural operators construct an
approximation of $\mathsf{G}^{\dagger}$ picked from the parametric class

$\mathsf{G}:\mathcal{U}\times\Theta\to\mathcal{Z};$

equivalently we may write $\mathsf{G}_{\theta}:\mathcal{U}\to\mathcal{Z},$ for
$\theta\in\Theta$. The set $\Theta$ is assumed to be a finite dimensional
parameter space from which a $\theta^{*}\in\Theta$ is selected so that
$\mathsf{G}(\cdot,\theta^{*})=\mathsf{G}_{\theta^{*}}\approx\mathsf{G}^{\dagger}$.
We then define a cost functional
$c:\mathcal{Z}\times\mathcal{Z}\to\mathbbm{R}$ and define $\theta^{*}$ by

$\theta^{*}={\arg\min}_{\theta\in\Theta}\mathbb{E}_{u\sim\mu}\Bigl{[}c\bigl{(}\mathsf{G}(u,\theta),\mathsf{G}^{\dagger}(u)\bigr{)}\Bigr{]}.$
(4)

In practice, we will only have access to the values of each $u^{(j)}$ and
$z^{(j)}$ at a set of discretization points of the domain $D$, which we denote
as $\\{x_{i}\\}_{i=1}^{N}$. We seek discretization invariant approaches to
operator learning which allow for sharing of learned parameters $\theta^{*}$
between different discretizations.

The choice of an architecture defining the approximation family
$\mathsf{G}_{\theta}$ is sometimes known as surrogate modeling. We use
continuum attention concepts to make this choice. The resulting transformer
based neural operator methodology that we develop here can be used to learn
solution operators to parametric PDEs, to solve PDE inverse problems, and may
also be applied to solve a class of data assimilation problems. In the context
of parametric PDEs we will explore how this algorithmic framework may be
employed in operator learning for Darcy flow and 2d-Navier Stokes problems. As
a particular case of operator learning, we will also explore a class of data
assimilation problems that involve recovering unobserved trajectories of
possibly chaotic dynamical systems. Namely, we illustrate examples involving
the Lorenz-63 dynamical system and a controlled ODE.

###  1.1 Literature Review

The attention mechanism, as introduced in Vaswani et al. (2017) and defined in
(1), has shown widespread success for a variety of tasks. In particular, given
the suitability of attention and transformers for applications involving
sequential data, there has been a large body of work extending the methodology
from natural language processing, the setting in which it was originally
introduced, to the more general context of time-series. Namely, as surveyed in
Wen et al. (2023), transformers and variants thereof have been applied to
time-series forecasting, classification and anomaly detection tasks. This
generalization has involved modifications and innovations relating to the
positional encoding and the design of the attention function itself. We refer
the reader to Wen et al. (2023) for a detailed review of the large body of
work and architectural variants of transformers that have been developed for
time-series applications.

The quadratic scaling in the input of the attention mechanism has made
transformers a computationally prohibitive architecture for applications
involving long sequences. Nonetheless, through techniques such as patching,
the attention mechanism and transformers have also enjoyed success in computer
vision, where high-resolution image data presents an increased computational
cost. In fact, in vision transformers (ViT), introduced in Dosovitskiy et al.
(2021), the input image is subdivided into patches; attention is then applied
across the shorter sequence of patches, thus decreasing the overall cost.
Given the success of the methodology in computer vision, there has been
growing interest in extending the transformer methodology to the operator
learning context. This operator setting concerns learning mappings between
infinite dimensional spaces of functions. In this case, the domains on which
the functions are defined (spatial and temporal) may inherently be a
continuum. Neural operators (Kovachki et al., 2023), generalizations of neural
networks that learn operators mapping between infinite dimensional spaces,
have been developed for this task. To this end, an integral kernel
interpretation of the attention mechanism appearing for example in Tsai et al.
(2019); Martins et al. (2020); Moreno et al. (2022); Wright and Gonzalez
(2021); Cao (2021); Kovachki et al. (2023) inspired a range of attention-based
architectures for operator learning. For example, the work of Cao (2021)
leverages the integral kernel formulation to design novel attention mechanisms
that do not require the softmax function. Inspired by the token mixing in
attention functions, in Guibas et al. (2022) the authors propose an
architecture for operator learning that involves token mixing in Fourier
space. Various other approaches have been developed. In Li et al. (2023) the
authors present a methodology relying on an encoder-decoder structure that
employs self-attention and cross-attention and uses a recurrent multi-layer
perceptron for unrolling instead of masked attention as in Vaswani et al.
(2017). The work Hao et al. (2023) introduces an architecture for operator
learning employing a normalized attention layer that handles multiple inputs.
This design, along with a learned gating mechanism for input coordinates,
allows the model to be applied to irregular meshes and multi-scale problems.
In Ovadia et al. (2023) the authors propose to use an architecture composed of
a ViT in the latent space of a U-Net (Ronneberger et al., 2015). The proposed
model takes as input the input field evaluated on a grid with associated
coordinate values and outputs an approximation of the solution to PDE inverse
problems. While it is shown that the architecture achieves state-of-the-art
accuracy, the parameters in the scheme are not invariant to the input function
resolution. In the recent work Rahman et al. (2024), the authors propose a
neural operator architecture that employs an attention mechanism that acts
across channels of the function in the latent embedding, and do so in the
context of pre-training a foundation model for PDE solution operators. We
demonstrate how such an architecture can be derived from our formulation in
Remark 16. As in our paper, various attention-related concepts are defined in
a function-space setting, but our formulations are more general; this issue is
discussed in detail at the relevant point in the paper.

The use of attention-based transformer models in the context of dynamical
systems and data assimilation problems, where spatial and temporal domains of
the systems are continuous, remains largely unexplored. For a comprehensive
review on the application of machine learning to data assimilation we refer
the reader to Cheng et al. (2023). In this context, the efficacy of
transformers has been explored for forecasting in numerical weather
prediction. In Chattopadhyay et al. (2022), a transformer module is integrated
in the latent space of a U-Net (Ronneberger et al., 2015) architecture in
order to leverage the permutation equivariance property of the transformer to
improve physical consistency and forecast accuracy. The work in Bi et al.
(2023) achieves state-of-the-art performance among end-to-end learning methods
for numerical weather prediction by applying a ViT-like architecture and
employing the shifted window methodology from Liu et al. (2021). The wide
interest in blending ML and data assimilation along with the capability of
transformer architectures provides an additional motivation for developing a
mathematical framework for attention defined on infinite-dimensional spaces of
functions.

Mathematical foundations for the methodology are only now starting to emerge,
however most analysis has been preformed in the finite-dimensional setting. In
the recent work Geshkovski et al. (2023), the authors study the attention
dynamics by viewing tokens as particles in an interacting particle system. The
formulation of the continuous-time limit of the dynamics, where the limit is
taken in “layer time” allows to investigate the emergence of clusters among
tokens. In Yun et al. (2020), the authors prove that transformers are
universal approximators of continuous permutation equivariant sequence-to-
sequence functions with compact support. More broadly, developing a firm
theoretical framework for the subject is necessary to understand the empirical
performance of the methodology and to inform the design of novel attention-
based architectures.

###  1.2 Contributions and Outline

Motivated by the success of transformers for sequential data, we set out to
formulate the attention mechanism and transformer architectures as mappings
between infinite dimensional spaces of functions. In this work we introduce a
continuum perspective on the transformer methodology. The resulting
theoretical framework enables the design of attention-based neural operator
architectures, which find a natural application to operator learning problems.
The strength of the approach of devising implementable discrete architectures
from the continuum perspective stems from the discretization invariance
property that the schemes inherit. Indeed, the resulting neural operator
architectures are designed as mappings between infinite dimensional functions
spaces that are invariant to the particular discretization, or resolution, at
which the input functions are defined. Notably, this allows for zero-shot
generalization to different function discretizations. With these insights in
view, our main contributions can be summarized as follows:

  1. 1.

Formulation of attention as a mapping between function spaces.

  2. 2.

Approximation theorem that quantifies the error between application of the
attention operator to a continuous function and the result of applying its
finite dimensional analogue to a discretization of the same function.

  3. 3.

A formulation of the transformer architecture as a neural operator, hence as a
mapping between function spaces; the resulting scheme is resolution invariant
and allows for zero-shot generalization to irregular, non-uniform meshes.

  4. 4.

A first universal approximation theorem for transformer neural operators.

  5. 5.

Formulation of patch-attention as a mapping between spaces of functions acting
on patch indices. This continuum perspective leads to the design of efficient
attention-based neural operator architectures.

  6. 6.

Numerical evidence of the competitiveness of transformer-based neural
architectures.

After introducing, in Subsection 1.3, the notation that will be used
throughout, Section 2 is focused on contributions 1 and 2, formulating
attention as acting on functions defined on a continuous domain. In Section 3
we formulate the attention mechanism as acting on a sequence of “patches” of a
function, addressing contribution 5. In Section 4 we make use of the operator
frameworks developed in Sections 2 and 3 to formulate transformer neural
operator architectures acting on function space. Thus we address 3, whilst the
last two subsections concern the algorithmic aspect of extending to patching
and address contribution 5. In Section 5 we state and prove a first universal
approximation theorem for a transformer-based neural operator, hence
addressing contribution 4. Finally, in Section 6 we explore applications of
the transformer neural operator architectures to a variety of operator
learning problems, contribution 6.

###  1.3 Notation

Throughout we denote the positive integers and non-negative integers
respectively by $\mathbbm{N}=\\{1,2,\cdots\\}$ and
$\mathbbm{Z}^{+}=\\{0,1,2,\cdots\\},$ and the notation
$\mathbbm{R}=(-\infty,\infty)$ and $\mathbbm{R}^{+}=[0,\infty)$ for the reals
and the non-negative reals. We let $\llbracket
1,n\rrbracket\subset\mathbbm{N}$ denote the set $\\{1,\ldots,n\\}$. For a set
$D$, we denote by $\bar{D}$ the closure of the set, i.e. the union of the set
itself and the set of all its limit points. We let
$\langle\cdot,\cdot\rangle,|\cdot|$ denote the Euclidean inner-product and
norm, noting that $|v|^{2}=\langle v,v\rangle.$ We may also use $|\cdot|$ to
denote the cardinality of a set, but the distinction will be clear from
context. We write
$\langle\cdot,\cdot\rangle_{\mathbbm{R}^{d}},|\cdot|_{\mathbbm{R}^{d}}$ if we
want to make explicit the dimension of the space on which the Euclidean inner
product is computed. We will also denote by $|\cdot|_{1}$ the norm on the
vector space $\ell^{1}$.

We denote by $C(D;\mathbbm{R}^{d})$ the infinite dimensional Banach space of
continuous functions mapping the set $D$ to the $d$-dimensional vector space
$\mathbbm{R}^{d}$. The space is endowed with the supremum norm. Similarly, we
denote by $L^{2}(D;\mathbbm{R}^{d})$ the infinite dimensional space of square
integrable functions mapping $D$ to $\mathbbm{R}^{d}$. The space is endowed
with the $L^{2}$ inner product, which induces the norm on the space. We will
sometimes use the shorthand notation $C(D)$ and $L^{2}(D)$ to denote
continuous functions and square integrable functions defined on the domain
$D$, respectively, when the image space is irrelevant. For integers $s\geq 0$
we denote by $C^{s}$ the space of continuously differentiable functions up to
order $s$; furthermore, for $p\in[1,\infty)$, we denote by $W^{s,p}$ the
Sobolev space of functions possessing weak derivatives up to order $s$ of
finite $L^{p}$ norm. Throughout, general vector spaces will be written using a
calligraphic font $\mathcal{U}$. We denote by
$\mathcal{L}(\mathcal{U},\mathcal{V})$ the infinite-dimensional space of
linear operators mapping the vector space $\mathcal{U}$ to the vector space
$\mathcal{V}$. Throughout, we use different fonts to highlight the difference
between operators acting on finite dimensional spaces, such as $Q$, and
operators acting on infinite dimensional spaces, such as $\mathsf{Q}$. We use
$\mathcal{F}$ to denote the Fourier transform, while $\mathcal{F}^{-1}$ will
denote the inverse Fourier transform.

We use $\mathbb{E}$ to denote expectation under the prevailing probability
measure; if we wish to make clear that measure $\pi$ is the prevailing
probability measure then we write $\mathbb{E}_{\pi}$.

##  2 Continuum Attention

In this section we formulate a continuum analogue of the attention mechanism
described in Vaswani et al. (2017). We note that in the transformer
architecture of Vaswani et al. (2017), a distinction is made between the self-
attention and cross-attention functions. The former, a particular instance of
(1) where $u=v$, can be used to produce a representation of the input that
captures intra-sequence correlations. On the other hand, cross-attention is
used to output a representation of the input $u$ that models cross-
correlations between both inputs $u$ and $v$. Indeed, in Vaswani et al. (2017)
cross-attention is employed in the decoder component of the transformer
architecture so that the input to the decoder attends to the output of the
encoder. For natural language applications in particular, the form of the
cross-attention function has allowed for sequences of arbitrary length $N$ to
be outputted from a transformer architecture while attending to sequences of
different lengths, e.g. as outputs of length $M$ from a self-attention encoder
block.

For pedagogical purposes we subdivide the presentation between the self-
attention operator and the cross-attention operator, which are treated in
Subsection 2.1 and Subsection 2.2 respectively. In Subsection 2.1.1 we define
self-attention for discrete sequences indexed on finite bounded sets. Here, we
outline the interpretation of self-attention as an expectation under a family
of carefully defined discrete probability measures, each acting on the space
of sequence indices. We show that under this formulation we recover the
standard definition of self-attention from Vaswani et al. (2017). We proceed
in Subsection 2.1.2 by formulating self-attention as an operator mapping
between spaces of functions defined on bounded uncountable sets. This
interpretation is natural, as functions may be thought of as “sequences”
indexed over an uncountable domain. As in the discrete case, we formulate the
self-attention operator as an expectation under a carefully defined
probability measure acting on the domain of the input. We obtain an
approximation result given by Theorem 6 that shows that for continuous
functions the self-attention mechanism as implemented in practice may be
viewed as a Monte Carlo approximation of the self-attention operator
described. We mimic these constructions in Subsections 2.2.1 and 2.2.2 for the
discrete and continuous formulations of cross-attention, respectively, and
obtain an analogous approximation result given by Theorem 12. As in (1),
throughout this section $Q\in\mathbbm{R}^{d_{K}\times
d_{u}},K\in\mathbbm{R}^{d_{K}\times d_{v}},V\in\mathbbm{R}^{d_{V}\times
d_{v}}$ are learnable weights that parametrize the attention operators. When
we consider self-attention we will have $d_{v}=d_{u}.$

###  2.1 Self-Attention

In the following discussion we define the self-attention operator in the
discrete setting and in the continuum setting, in Subsections 2.1.1 and 2.1.2,
respectively.

####  2.1.1 Sequences over $D^{N}\subset\mathbbm{Z}$

We begin by considering the finite, bounded set $D^{N}=\\{1,\cdots,N\\}.$ We
let $u:D^{N}\to\mathbbm{R}^{d_{u}}$ be a sequence and $j\in D^{N}$ an index so
that $u(j)\in\mathbbm{R}^{d_{u}}$. We will now focus on formulating the
attention mechanism in a manner which allows generalization to more general
domains $D^{N};$ at the end of this subsection we connect our formulation with
the more standard one used in the literature, which is specific to the current
choice of domain $D^{N}=\\{1,\cdots,N\\}.$ We define the following discrete
probability measure.

###### Definition 1

We define a probability measure on $D^{N}$, parameterized by $(u,j)$, via the
probability mass function $p(\cdot;u,j):D^{N}\to[0,1]$ defined by

$p(k;u,j)=\frac{\mathrm{exp}\Big{(}\big{\langle}Qu(j),Ku(k)\big{\rangle}_{\mathbbm{R}^{d_{\textrm{\scalebox{0.7}{$K$}}}}}\Big{)}}{\sum_{\ell\in
D^{N}}\mathrm{exp}\Big{(}\big{\langle}Qu(j),Ku(\ell)\big{\rangle}_{\mathbbm{R}^{d_{\textrm{\scalebox{0.7}{$K$}}}}}\Big{)}},$

for any $k\in D^{N}$.

We may now define the self-attention operator as an expectation over this
measure $p$ of the linear transformation $V\in\mathbbm{R}^{d_{V}\times d_{u}}$
applied to $u$.

###### Definition 2

We define the self-attention operator $\mathsf{A}^{N}$ as a mapping from a
$\mathbbm{R}^{d_{u}}-$valued sequence over $D^{N}$, $u$, into a
$\mathbbm{R}^{d_{V}}-$valued sequence over $D^{N}$, $\mathsf{A}^{N}(u)$, that
takes the form

$\mathsf{A}^{N}(u)(j):=\mathbb{E}_{k\sim p(k;u,j)}[Vu(k)],$

for any $j\in D^{N}$.

To connect these definitions with the standard formulation of transformers,
note that sequence $u$ can be reformulated as the matrix
$u\in\mathbbm{R}^{N\times d_{u}}.$ Then

$\\{uQ^{T}Ku^{T}\\}_{jk}=\big{\langle}Qu_{j},Ku_{k}\big{\rangle}_{\mathbbm{R}^{d_{\textrm{\scalebox{0.7}{$K$}}}}};$

that is, the $\\{j,k\\}-$entry of $uQ^{T}Ku^{T}\in\mathbbm{R}^{N\times N}$ is
the right-hand side of the preceding identity, where $u_{\ell}$ denotes the
$\ell$’th row of $u\in\mathbbm{R}^{N\times d_{u}}$. It is then apparent that
applying the softmax function along rows of $uQ^{T}Ku^{T}$ delivers the vector
of probabilities defined in Definition 1, indexed by $k;$ note that $j$ is a
parameter in this vector of probabilities, indicating the row of the matrix
$uQ^{T}Ku^{T}$ along which the softmax operation is applied. Finally the
attention operation from Definition 2 can be rewritten to act between matrices
as

${\mathsf{A}^{N}(u)}={\tt attention}(u,u):={\tt softmax}(uQ^{T}Ku^{T})uV^{T},$

so that ${\tt attention}(u,u)\in\mathbbm{R}^{N\times d_{V}}$.

######  Remark 3 (Sequences over $D^{N}\subset\mathbbm{Z}^{d}$)

Once the attention mechanism is formulated as in Definitions 1, 2, it is
straightforward to extend to (generalized) sequences indexed over bounded
subsets of $\mathbbm{Z}^{d};$ note, for example, that pixellated images are
naturally indexed over bounded subsets of $\mathbbm{Z}^{2}.$ Let
$D^{N}\subset\mathbbm{Z}^{d}$ possess cardinality $|D^{N}|=N.$ Let
$u:D^{N}\to\mathbbm{R}^{d_{u}}$ be a sequence and $j\in D^{N}$ an index. Then
the definition of probability on $D^{N}$, indexed by $(u,j)$, and the
resulting definition of self-attention operator $\mathsf{A}$, is exactly as in
Definitions 1, 2. This shows the power of non-standard notation we have
employed here.

####  2.1.2 Sequences over $D\subset\mathbbm{R}^{d}$

We note that the preceding two formulations of the attention mechanism view it
as a transformation defined as an expectation applied to the input sequence.
The expectation is with respect to a probability measure supported on an index
of the input sequence; furthermore the expectation itself depends on the input
sequence to which it is applied, and is parameterized by the input sequence
index, resulting in a new output sequence of the same length as the input
sequence. The formulation using an expectation results from the softmax part
of the definition of the attention mechanism; the query and key linear
operations define the expectation, and the value linear operation lifts the
output sequence to take values in a (possibly) different Euclidean space than
the input sequence. These components may be readily extended to work with
(generalized) sequences defined over open subsets of $\mathbbm{R}^{d}$. (These
would often be referred to as functions, but we call them generalized
sequences to emphasize similarity with the discrete case.)

Let $D\subseteq\mathbbm{R}^{d}$ be an open set. Let
$u:D\to\mathbbm{R}^{d_{u}}$ be a sequence (function) and $x\in D$ an index.

###### Definition 4

We define a probability measure on $D$, parameterized by $(u,x)$, via the
probability density function $p(\cdot;u,x):D\to\mathbbm{R}^{+}$ defined by

$p(y;u,x)=\frac{\mathrm{exp}\Big{(}\big{\langle}Qu(x),Ku(y)\big{\rangle}_{\mathbbm{R}^{d_{\textrm{\scalebox{0.7}{$K$}}}}}\Big{)}}{\int_{D}\mathrm{exp}\Big{(}\big{\langle}Qu(x),Ku(s)\big{\rangle}_{\mathbbm{R}^{d_{\textrm{\scalebox{0.7}{$K$}}}}}\Big{)}\>\mathrm{d}s},$

for any $y\in D$.

###### Definition 5

We define the self-attention operator $\mathsf{A}$ as a mapping from a
$\mathbbm{R}^{d_{u}}-$valued sequence (function) over $D$, $u$, into a
$\mathbbm{R}^{d_{V}}-$valued sequence (function) over $D$, $\mathsf{A}(u)$, as
follows:

$\mathsf{A}(u)(x)=\mathbb{E}_{y\sim p(y;u,x)}[Vu(y)],$

for any $x\in D$.

A connection between the discrete and continuum formulations of self-attention
is captured in the following theorem.

###### Theorem 6

The self-attention operator $\mathsf{A}$ may be viewed as a mapping
$\mathsf{A}:L^{\infty}(D;\mathbbm{R}^{d_{u}})\to
L^{\infty}(D;\mathbbm{R}^{d_{V}})$ and thus as a mapping
$\mathsf{A}:C(\bar{D};\mathbbm{R}^{d_{u}})\to C(\bar{D};\mathbbm{R}^{d_{V}})$.
Furthermore, for any compact set $B\subset C(\bar{D};\mathbbm{R}^{d_{u}})$,

$\lim_{N\to\infty}\sup_{u\in
B}\mathbb{E}\left\|\mathsf{A}(u)-\frac{\sum_{j=1}^{N}\mathrm{exp}\big{(}\langle
Qu(\cdot),Ku(y_{j})\rangle\big{)}Vu(y_{j})}{\sum_{\ell=1}^{N}\mathrm{exp}\big{(}\langle
Qu(\cdot),Ku(y_{\ell})\rangle\big{)}}\right\|_{C(\bar{D};\mathbbm{R}^{d_{V}})}=0,$

with the expectation taken over i.i.d. sequences
$\\{y_{j}\\}_{j=1}^{N}\sim\mathsf{unif}(\bar{D})$.

Proof  The proof of this result is developed in Appendix A.1.

###### Remark 7

This result connects our continuum formulation of self-attention with the
standard definition, using Monte Carlo approximation of the relevant
integrals; a similar result could be proved using finite difference
approximation on, for example, a uniform grid.

###  2.2 Cross-Attention

In the following discussion we define the cross-attention operator in the
discrete setting and in the continuum setting, in Subsections 2.2.1 and 2.2.2,
respectively.

####  2.2.1 Sequences over
$D^{N}\subset\mathbbm{Z}^{d},E^{M}\subset\mathbbm{Z}^{e}$

We begin by letting $D^{N}\subseteq\mathbbm{Z}^{d}$ and
$E^{M}\subseteq\mathbbm{Z}^{e}$ be finite sets of points with cardinalities
$N$ and $M$ respectively. Let $u:D^{N}\to\mathbbm{R}^{d_{u}}$ be a sequence
(function), $v:E^{M}\to\mathbbm{R}^{d_{v}}$ another sequence (function), and
$j\in D^{N}$, $k\in\mathbb{E}^{M}$ indices so that
$u(j)\in\mathbbm{R}^{d_{u}}$ and $v(k)\in\mathbbm{R}^{d_{v}}$. We define the
following discrete probability measure.

###### Definition 8

We define a probability measure on $D^{N}$, parameterized by $(u,v,j)$, via
the probability mass function $q(\cdot;u,v,j):D^{N}\to[0,1]$ defined by

$q(k;u,v,j)=\frac{\mathrm{exp}\Big{(}\big{\langle}Qu(j),Kv(k)\big{\rangle}_{\mathbbm{R}^{d_{\textrm{\scalebox{0.7}{$K$}}}}}\Big{)}}{\sum_{\ell\in
E^{M}}\mathrm{exp}\Big{(}\big{\langle}Qu(j),Kv(\ell)\big{\rangle}_{\mathbbm{R}^{d_{\textrm{\scalebox{0.7}{$K$}}}}}\Big{)},}$

for any $k\in E^{M}$.

We may now define the cross-attention operator as an expectation over this
measure $q$ of the linear transformation $V\in\mathbbm{R}^{d_{V}\times d_{v}}$
applied to $v$.

###### Definition 9

We define the cross-attention operator $\mathsf{C}^{N}$ as a mapping from a
$\mathbbm{R}^{d_{u}}-$valued sequence (function) over $D^{N}$, $u$, and a
$\mathbbm{R}^{d_{v}}-$valued sequence (function) over $E^{M}$, $v$, into a
$\mathbbm{R}^{d_{V}}-$valued sequence (function) over $D^{N}$,
$\mathsf{C}^{N}(u,v)$, as follows:

$\mathsf{C}^{N}(u,v)(j)=\mathbb{E}_{k\sim q(k;u,v,j)}[Vv(k)],$

for any $j\in D^{N}$.

####  2.2.2 Sequences over $D\subset\mathbbm{R}^{d},E\subset\mathbbm{R}^{e}$

The framework described in Subsection 2.2.1 may be readily extended to work
with generalized sequences (functions) defined over open subsets of
$\mathbbm{R}^{d}$. Indeed, let $D\subseteq\mathbbm{R}^{d}$ and
$E\subseteq\mathbbm{R}^{e}$ be open sets. Let $u:D\to\mathbbm{R}^{d_{u}}$ be a
sequence (function), $v:E\to\mathbbm{R}^{d_{v}}$ another sequence (function),
and $x\in D$, $y\in E$ indices. We define the following probability measure
over $D$ before providing the definition of the cross-attention operator.

###### Definition 10

We define a probability measure on $D$, parameterized by $(u,v,x)$, via the
probability density function $q(\cdot;u,v,x):D\to\mathbbm{R}^{+}$ defined by

$q(y;u,v,x)=\frac{\mathrm{exp}\Big{(}\big{\langle}Qu(x),Kv(y)\big{\rangle}_{\mathbbm{R}^{d_{\textrm{\scalebox{0.7}{$K$}}}}}\Big{)}}{\int_{E}\mathrm{exp}\Big{(}\big{\langle}Qu(x),Kv(s)\big{\rangle}_{\mathbbm{R}^{d_{\textrm{\scalebox{0.7}{$K$}}}}}\Big{)}\>\mathrm{d}s},$

for any $y\in{E}$.

###### Definition 11

We define the cross-attention operator $\mathsf{C}$ as a mapping from a
$\mathbbm{R}^{d_{u}}-$valued sequence (function) over $D$, $u$, and a
$\mathbbm{R}^{d_{v}}-$valued sequence (function) over $E$, $v$, into a
$\mathbbm{R}^{d_{V}}-$valued sequence (function) over $D$, $\mathsf{C}(u,v)$,
as follows:

$\mathsf{C}(u,v)(x)=\mathbb{E}_{y\sim q(y;u,v,x)}[Vv(y)],$

for any $x\in D$.

A connection between the discrete and continuum formulations of cross-
attention is captured in the following theorem.

###### Theorem 12

The cross-attention operator $\mathsf{C}$ may be viewed as a mapping
$\mathsf{C}:L^{\infty}(D;\mathbbm{R}^{d_{u}})\times
L^{\infty}(E;\mathbbm{R}^{d_{v}})\to L^{\infty}(D;\mathbbm{R}^{d_{V}})$, and
so as a mapping $\mathsf{C}:C(\bar{D};\mathbbm{R}^{d_{u}})\times
C(\bar{E};\mathbbm{R}^{d_{v}})\to C(\bar{D};\mathbbm{R}^{d_{V}})$.
Furthermore, for any compact set $B\subset
C(\bar{D};\mathbbm{R}^{d_{u}})\times C(\bar{E};\mathbbm{R}^{d_{v}})$

$\lim_{N\to\infty}\sup_{(u,v)\in
B}\mathbb{E}\left\|\mathsf{C}(u,v)-\frac{\sum_{j=1}^{N}\mathrm{exp}\big{(}\langle
Qu(\cdot),Kv(y_{j})\rangle\big{)}Vv(y_{j})}{\sum_{\ell=1}^{N}\mathrm{exp}\big{(}\langle
Qu(\cdot),Kv(y_{\ell})\rangle\big{)}}\right\|_{C(\bar{D};\mathbbm{R}^{d_{V}})}=0,$

with the expectation taken over i.i.d. sequences
$\\{y_{j}\\}_{j=1}^{N}\sim\mathsf{unif}(\bar{E})$.

Proof  The proof of this result is developed in Appendix A.2.

##  3 Continuum Patched Attention

Patches of a function are defined as the restriction of the function itself to
elements of a partition of the domain. We may then define patch-attention as a
mapping between spaces of functions defined on patch indices. Interest in
applying the attention mechanism for computer vision led to the development of
methodologies to overcome its prohibitive computational cost. Patching, as
employed in vision transformers (ViT), involves subdividing the data space
domain into $P\in\mathbb{N}$ “patches” and applying attention to a sequence of
flattened patches. This step drastically reduces the complexity of the
attention mechanism, which is quadratic in the input sequence length. An
analogue of the patching strategy in vision transformers can be considered in
the function space setting. Indeed, in image space, attention can be applied
across subsets of pixels (patches); in function space, attention can be
applied across the function defined on elements of a partition of the domain.
Such a generalization to functions defined on the continuum allows for the
development of architectures that are mesh-invariant. We describe the patched-
attention methodology in the continuum framework developed thus far.

Throughout this section, we let $D\subset\mathbbm{R}^{d}$ be a bounded open
set and let $D\coloneqq D_{1}\cup\ldots\cup D_{P}$ be a uniform partition of
the space $D$ so that $D_{j}\cong D^{\prime}$ are congruent for all
$j\in\llbracket 1,P\rrbracket$ for some $D^{\prime}\subset D$, where
$P\in\mathbb{N}$ represents the number of patches. We consider a function
defined on the full domain $D$, denoted as
$u\in\mathcal{U}\bigl{(}D;\mathbbm{R}^{d_{u}}\bigr{)}$. We define a mapping

$\widetilde{u}:\llbracket
1,P\rrbracket\to\mathcal{U}(D^{\prime},\mathbb{R}^{d_{u}}),$ (5)

that defines the patched version of the function
$u\in\mathcal{U}\bigl{(}D;\mathbbm{R}^{d_{u}}\bigr{)}$, so that

$\widetilde{u}(p)={\left.\kern-1.2ptu\mathchoice{\vphantom{\big{|}}}{}{}{}\right|_{D_{p}}},$
(6)

for any $p\in\llbracket 1,P\rrbracket$. We proceed to establish the
definitions of the patched self-attention operator $\mathsf{A}_{\mathsf{P}}$
in Subsection 3.1 and for completeness the patched cross-attention operator
$\mathsf{C}_{\mathsf{P}}$ in Subsection 3.2.

###  3.1 Patched Self-Attention

We begin with the definition of the patched self-attention operator
$\mathsf{A}_{\mathsf{P}}$. Indeed, we let the operators
$\mathsf{Q},\mathsf{K},\mathsf{V}$ be defined such that
$\mathsf{Q}:\mathcal{U}(D^{\prime},\mathbb{R}^{d_{u}})\to
L^{2}(D^{\prime},\mathbb{R}^{d_{K}}),\mathsf{K}:\mathcal{U}(D^{\prime},\mathbb{R}^{d_{u}})\to
L^{2}(D^{\prime},\mathbb{R}^{d_{K}}),\mathsf{V}:\mathcal{U}(D^{\prime},\mathbb{R}^{d_{u}})\to\mathcal{U}(D^{\prime},\mathbb{R}^{d_{V}})$
and let $j\in\llbracket 1,P\rrbracket$ be an index.

###### Definition 13

Define a probability measure on $\llbracket 1,P\rrbracket$, parameterized by
$(\widetilde{u},j)$, via the probability density function
$p(\cdot;\widetilde{u},j):\llbracket 1,P\rrbracket\to[0,1]$ defined by

$p(k;\widetilde{u},j)=\frac{\mathrm{exp}\Big{(}\big{\langle}\mathsf{Q}\widetilde{u}(j),\mathsf{K}\widetilde{u}(k)\big{\rangle}_{L^{2}(D^{\prime},\mathbb{R}^{d_{K}})}\Big{)}}{\sum_{\ell=1}^{P}\mathrm{exp}\Big{(}\big{\langle}\mathsf{Q}\widetilde{u}(j),\mathsf{K}\widetilde{u}(\ell)\big{\rangle}_{L^{2}(D^{\prime},\mathbb{R}^{d_{K}})}\Big{)}},$

for any $k\in\llbracket 1,P\rrbracket$.222We note that as $\mathsf{Q}$ is
linear, $\mathsf{Q}u(j)$ is shorthand notation for
$\mathsf{Q}\bigl{(}u(j)\bigr{)}$, and similarly for $\mathsf{K}$ and
$\mathsf{V}$.

###### Definition 14

The self-attention operator $\mathsf{A}_{\mathsf{P}}$ maps the function taking
a patch index to a $\mathbbm{R}^{d_{u}}-$valued function over $D^{\prime}$,
$\widetilde{u}$, into a function taking a patch index to a
$\mathbbm{R}^{d_{V}}-$valued function over $D^{\prime}$,
$\mathsf{A}_{\mathsf{P}}(\widetilde{u})$, as follows:

$\mathsf{A}_{\mathsf{P}}(\widetilde{u})(j)=\mathbb{E}_{k\sim
p(k;\widetilde{u},j)}[\mathsf{V}\widetilde{u}(k)],$

for any $j\in\llbracket 1,P\rrbracket$.

######  Remark 15 (Patched Attention in ViT)

An image may be viewed as a mapping $u:D\to\mathbbm{R}^{d_{u}}$, where
$D\subset\mathbbm{Z}^{2}$, and where $d_{u}=1$ for grey-scale and $d_{u}=3$
for RGB-valued images. In the context of vision transformers (Dosovitskiy et
al., 2021), the input to the attention mechanism is a sequence of $P$
“flattened” patches of the input image; we next outline the details to this
procedure.

Letting $N$ be the total number of discretization points in $D$ and $P$ the
number of patches, in the discrete setting each function patch
$\widetilde{u}(p):D^{\prime}\to\mathbbm{R}^{d_{u}}$ may be represented as a
$\mathbbm{R}^{N/P\times d_{u}}$ matrix. The flattening procedure of each patch
in ViT involves reshaping this matrix into a $\mathbbm{R}^{N/P\cdot d_{u}}$
vector. The full image is thus represented as a sequence of $P$ vectors of
dimension $N/P\cdot d_{u}$. Each patch vector is then linearly lifted to an
embedding space of dimension $d_{\textrm{model}}$. The embedded image, to
which attention is applied, may hence be viewed as a matrix
$\widetilde{u}\in\mathbbm{R}^{P\times d_{\textrm{model}}}$. Therefore, ViT may
be cast in the setting of the discrete interpretation of the self-attention
operator outlined in Subsection 2.1.1, where the domain on which attention is
applied is $D^{P}$. A key observation to the ensuing discussion is that the
application of a linear transformation to a “flattened” patch in ViT breaks
the mesh-invariance of the architecture; however, this operation may be viewed
as a nonlocal transformation on each patch. We leverage this insight in
Section 4 to design discretization invariant transformer neural operators that
employ patching.

######  Remark 16 (Codomain Attention)

The work of Rahman et al. (2024) employs an attention mechanism that acts
across the channels of the function
$u\in\mathcal{U}\bigl{(}D;\mathbbm{R}^{d_{\textrm{model}}}\bigr{)}$ in the
latent space of a neural operator architecture. Indeed the “codomain”
attention mechanism defined may be viewed as a particular variant of
Definition 14, where $D^{\prime}\cong D$, where
$\mathsf{A}_{\mathsf{P}}(\widetilde{u})(j)$ is defined for $j\in\llbracket
1,d_{\textrm{model}}\rrbracket$ and where $\mathsf{Q},\mathsf{K},\mathsf{V}$
are implemented as integral operators. In particular the indexing is by
channels and not by patch. For problems in computer vision, similar ideas have
been explored with discrete versions of the $\mathsf{Q},\mathsf{K},\mathsf{V}$
operators (Chen et al., 2017).

###  3.2 Patched Cross-Attention

For completeness, we outline the definition of the patched cross-attention
operator. We consider $E\subset\mathbbm{R}^{e}$ to be a bounded open set and
$E\coloneqq E_{1}\cup\ldots\cup E_{O}$ a uniform partition of $E$ so that
$E_{j}\cong E^{\prime}\cong D^{\prime}$ for all $j\in\llbracket 1,O\rrbracket$
for some $E^{\prime}\subset E$ and $O\in\mathbbm{N}$. We note that for the
following construction we require $E^{\prime}\cong D^{\prime}$, but it may
hold in general that $P\neq O$. As done for $\widetilde{u}$ we define an
operator

$\widetilde{v}:\llbracket
1,O\rrbracket\to\mathcal{U}(E^{\prime},\mathbb{R}^{d_{v}}),$ (7)

that defines the patched version of the function
$v\in\mathcal{U}\bigl{(}E;\mathbbm{R}^{d_{v}}\bigr{)}$, so that

$\widetilde{v}(o)={\left.\kern-1.2ptv\mathchoice{\vphantom{\big{|}}}{}{}{}\right|_{E_{o}}},$
(8)

for any $o\in\llbracket 1,O\rrbracket$. Furthermore, we let the operators
$\mathsf{Q},\mathsf{K},\mathsf{V}$ be defined such that
$\mathsf{Q}:\mathcal{U}(D^{\prime},\mathbb{R}^{d_{u}})\to
L^{2}(D^{\prime},\mathbb{R}^{d_{K}}),\mathsf{K}:\mathcal{U}(E^{\prime},\mathbb{R}^{d_{v}})\to
L^{2}(E^{\prime},\mathbb{R}^{d_{K}}),\mathsf{V}:\mathcal{U}(E^{\prime},\mathbb{R}^{d_{v}})\to\mathcal{U}(E^{\prime},\mathbb{R}^{d_{V}})$
and let $j\in\llbracket 1,P\rrbracket$ be an index.

###### Definition 17

We define a probability measure on $D$, parameterized by
$(\tilde{u},\tilde{v},j)$, via the probability density function
$q(\cdot;\widetilde{u},\widetilde{v},j):\llbracket 1,P\rrbracket\to[0,1]$
defined by

$q(k;\widetilde{u},\widetilde{v},j)=\frac{\mathrm{exp}\Big{(}\big{\langle}\mathsf{Q}\widetilde{u}(j),\mathsf{K}\widetilde{v}(k)\big{\rangle}_{L^{2}(E^{\prime},\mathbb{R}^{d_{K}})}\Big{)}}{\sum_{\ell=1}^{O}\mathrm{exp}\Big{(}\big{\langle}\mathsf{Q}\widetilde{u}(j),\mathsf{K}\widetilde{v}(\ell)\big{\rangle}_{L^{2}(E^{\prime},\mathbb{R}^{d_{K}})}\Big{)}},$

for any $k\in\llbracket 1,{O}\rrbracket$.

###### Definition 18

The cross-attention operator $\mathsf{C}$ maps the operator taking a patch
index to a $\mathbbm{R}^{d_{u}}-$valued function over $D^{\prime}$,
$\widetilde{u}$, and the operator taking a patch index to a
$\mathbbm{R}^{d_{v}}-$valued function over $E^{\prime}$, $\widetilde{v}$, into
an operator taking a patch index to a $\mathbbm{R}^{d_{V}}-$valued function
over $D^{\prime}$, $\mathsf{C}_{\mathsf{P}}(\widetilde{u},\widetilde{v})$, as
follows:

$\mathsf{C}_{\mathsf{P}}(\widetilde{u},\widetilde{v})(j)=\mathbb{E}_{k\sim
q(k;\widetilde{u},\widetilde{v},j)}[\mathsf{V}\widetilde{v}(k)],$

for any $j\in\llbracket 1,P\rrbracket$.

##  4 Transformer Neural Operators

In this section we turn our attention to describing how to devise transformer-
based neural operators using the continuum attention operator $\mathsf{A}$ as
defined in the function space setting in Section 2, and the continuum patched-
attention operator $\mathsf{A}_{\mathsf{P}}$ acting on function space from
Section 3. We make use of the operator frameworks developed in Sections 2 and
3 to formulate transformer neural operator architectures acting on function
space. We start in Subsection 4.1, setting-up the framework. In Subsection 4.2
we formulate an analogue of the transformer from Vaswani et al. (2017) as an
operator mapping an input function space
$\mathcal{U}\bigl{(}D;\mathbbm{R}^{d_{u}}\bigr{)}$ to the solution function
space $\mathcal{Z}\bigl{(}D;\mathbbm{R}^{d_{z}}\bigr{)}$. The resulting scheme
is invariant to the discretization of the input function and allows for zero-
shot generalization to irregular, non-uniform meshes. Motivated by the need to
develop new efficient architectures to employ attention on longer sequences,
or in our case higher resolution discretizations of functions, in Subsection
4.3 we describe a generalization of the ViT methodology to the function space
setting, inspired by the work in Dosovitskiy et al. (2021). In particular we
generalize the procedure of lifting flattened patches to the embedding
dimension, as done in ViT, to the function space setting. This is achieved by
lifting patches of functions to an embedding dimension using a nonlocal linear
operator, namely, an integral operator. Unlike ViT, the resulting neural
operator architecture is discretization invariant and allows for zero-shot
generalization to different resolutions. In Subsection 4.4 we describe a
different patching-based transformer neural operator acting on function space.
The architecture is discretization invariant and is based on
$\mathsf{Q},\mathsf{K},\mathsf{V}$, learnable parameters in the attention
mechanism, being defined as integral operators acting themselves on function
space.

###  4.1 Set-Up

In the subsections that follow we describe transformer neural operators of the
form
$\mathsf{G}:\mathcal{U}(D;\mathbbm{R}^{d_{u}})\times\Theta\to\mathcal{Z}(D;\mathbbm{R}^{d_{z}})$,
which serve as approximations to an operator
$\mathsf{G}^{\dagger}:\mathcal{U}(D;\mathbbm{R}^{d_{u}})\to\mathcal{Z}(D;\mathbbm{R}^{d_{z}})$.
The neural operators we devise have the general form

$\mathsf{G}(u,\theta)\coloneqq\Bigl{(}\mathsf{T}_{\textrm{out}}\circ\mathsf{E}_{L}\circ\mathsf{T}_{\textrm{in}}\Bigr{)}(u,\theta),$
(9)

for any $u\in\mathcal{U}(D;\mathbbm{R}^{d_{u}})$ and $\theta\in\Theta$. In
(9), the operators
$\mathsf{T}_{\textrm{in}}:\mathcal{U}(D;\mathbbm{R}^{d_{u}})\times\Theta\to\mathcal{V}$
and
$\mathsf{T}_{\textrm{out}}:\mathcal{V}\times\Theta\to\mathcal{Z}(D;\mathbbm{R}^{d_{z}})$
are defined for some appropriate embedding function space $\mathcal{V}$. We
will make explicit the definitions of the general operators
$\mathsf{T}_{\textrm{in}}$, $\mathsf{T}_{\textrm{out}}$ and the embedding
function space $\mathcal{V}$ in the context of each architecture. Throughout
this section, we drop explicit dependencies on $\theta$ for brevity unless we
wish to stress a parametric dependence; for example, we may abuse notation and
write $\mathsf{T}_{\textrm{in}}(u):=\mathsf{T}_{\textrm{in}}(u;\ \theta)$. The
operator $\mathsf{E}_{L}:\mathcal{V}\times\Theta\to\mathcal{V}$ defines the
neural operator analogue of the transformer encoder block from Vaswani et al.
(2017). Its action on the input
$v^{(0)}\coloneqq\mathsf{T}_{\textrm{in}}(u)\in\mathcal{V}$ is summarized by
the iteration

$\displaystyle v^{(l-1)}$
$\displaystyle\mapsfrom{\mathsf{W}_{1}}v^{(l-1)}+\mathsf{A}_{\textrm{MultiHead}}\big{(}v^{(l-1)}\big{)},$
(10a) $\displaystyle v^{(l-1)}$
$\displaystyle\mapsfrom\mathsf{F}_{\textrm{LayerNorm}}\big{(}v^{(l-1)}\big{)},$
(10b) $\displaystyle v^{(l-1)}$
$\displaystyle\mapsfrom{\mathsf{W}_{2}}v^{(l-1)}+\mathsf{F}_{\textrm{NN}}\big{(}v^{(l-1)}\big{)},$
(10c) $\displaystyle v^{(l)}$
$\displaystyle\mapsfrom\mathsf{F}_{\textrm{LayerNorm}}\big{(}v^{(l-1)}\big{)},$
(10d)

for $l=1,\dots,L$ layers, so that
$\mathsf{E}_{L}\bigl{(}v^{(0)}\bigr{)}=v^{(L)}\in\mathcal{V}$ 333We note that
we use $v$ to denote an arbitrary function in the space $\mathcal{V}$. This is
not to be confused with $v$ used for cross-attention in Section 2. We also
denote the output of the $\ell$’th encoder layer by $v^{(\ell)}$, not to be
confused with the notation used for the data pairs.. We note that for
notational convenience, we have suppressed dependence of the operators on the
layer $\ell$; indeed, each of the operators $\mathsf{A}_{\textrm{MultiHead}}$,
$\mathsf{W}_{1},\mathsf{W}_{2},$ $\mathsf{F}_{\textrm{NN}}$ and
$\mathsf{F}_{\textrm{LayerNorm}}$, will have different learnable
parametrizations for each encoder layer $\ell$. Before we delve into the
specific neural operator architectures, we summarize the definitions of the
elements of the transformer encoder described in (10), which are common to all
the subsequent architectures. The multi-head self-attention operator
$\mathsf{A}_{\textrm{MultiHead}}:\mathcal{V}\to\mathcal{V}$ is defined as the
composition of a linear transformation $W_{\textrm{MultiHead}}$ applied
pointwise to its input and a concatenation of the outputs of $H\in\mathbbm{N}$
self-attention operators so that

$\Bigl{(}\mathsf{A}_{\textrm{MultiHead}}(v)\Bigr{)}(x)=W_{\textrm{MultiHead}}\Bigl{(}\mathsf{A}_{\mathord{\color[rgb]{0.67,0.67,0.67}\bullet}}(v;\theta_{1})(x),\ldots,\mathsf{A}_{\mathord{\color[rgb]{0.67,0.67,0.67}\bullet}}(v;\theta_{H})(x)\Bigr{)},$
(11)

for any $v\in\mathcal{V}$. Here the notation
$\mathsf{A}_{\mathord{\color[rgb]{0.67,0.67,0.67}\bullet}}$ indicates the
ambiguity of the definition of self-attention. We note that we will make
explicit which definition of self-attention operator from the previous
sections we employ for each architecture; we will also make clear the
definitions of the parameters $\theta_{h}\in\Theta$ defining each attention
operator. The operators
$\mathsf{W}_{1},\mathsf{W}_{2}:\mathcal{V}\to\mathcal{V}$ are pointwise linear
operators. 444We note that for all experiments in Section 6 we set
$\mathsf{W}_{1},\mathsf{W}_{2}$ to be the identity operators. Next, the
operator $\mathsf{F}_{\textrm{LayerNorm}}:\mathcal{V}\to\mathcal{V}$ defines
the layer normalization and the operator
$\mathsf{F}_{\textrm{NN}}:\mathcal{V}\to\mathcal{V}$ defines a feed-forward
neural network layer. Both operators are applied pointwise to the input;
further details are provided in the context of each architecture.

A key aspect of the ensuing discussion will be the distinction between local
and nonlocal linear operators. We note that for an arbitrary function space
$\mathcal{U}\bigl{(}D;\mathbbm{R}^{r}\bigr{)}$ where
$D\subset\mathbbm{R}^{d}$, the space of linear operators
$\mathcal{L}\Bigl{(}\mathcal{U}\bigl{(}D;\mathbbm{R}^{r}\bigr{)};\mathcal{U}\bigl{(}D;\mathbbm{R}^{r^{\prime}}\bigr{)}\Bigr{)}$
includes both local and nonlocal linear operators. In the following, we will
consider pointwise linear operators in
$\mathcal{L}\Bigl{(}\mathcal{U}\bigl{(}D;\mathbbm{R}^{r}\bigr{)};\mathcal{U}\bigl{(}D;\mathbbm{R}^{r^{\prime}}\bigr{)}\Bigr{)}$
that admit a finite-dimensional representation
$W\in\mathbbm{R}^{r^{\prime}\times r}$. On the other hand, we will also
consider linear integral operators in
$\mathcal{L}\Bigl{(}\mathcal{U}\bigl{(}D;\mathbbm{R}^{r});\mathcal{U}\bigl{(}D;\mathbbm{R}^{r^{\prime}}\bigr{)}\Bigr{)}$,
a class of nonlocal operators. Namely, we consider operators $\mathsf{W}$ that
are integral operators learned from data of the form
$\mathsf{W}:\Theta_{\mathsf{W}}\to\mathcal{L}\Bigl{(}\mathcal{U}\bigl{(}D;\mathbbm{R}^{r});\mathcal{U}\bigl{(}D;\mathbbm{R}^{r^{\prime}}\bigr{)}\Bigr{)}$.

######  Definition 19 (Integral Operator $\mathsf{W}$)

The integral operator
$\mathsf{W}:\Theta_{\mathsf{W}}\to\mathcal{L}\Bigl{(}\mathcal{U}\bigl{(}D;\mathbbm{R}^{r});\mathcal{U}\bigl{(}D;\mathbbm{R}^{r^{\prime}}\bigr{)}\Bigr{)}$
is defined as the mapping given by

$\bigl{(}\mathsf{W}(\theta){u}\bigr{)}(x)\coloneqq\int_{D}{\kappa_{\theta}}(x,y){u}(y)\mathrm{d}y,\qquad\forall
u\in\mathcal{U}\bigl{(}D;\mathbbm{R}^{r}),\,\forall x\in D,$ (12)

where ${\kappa}_{\theta}:D\to\mathbbm{R}^{{r^{\prime}}\times r}$ is a neural
network parametrized by $\theta\in\Theta_{\mathsf{W}}$.

Assuming ${\kappa}_{\theta}(x,y)={\kappa}_{\theta}(x-y)$ then (12) may be
viewed as a convolution operator, which makes it possible to use the Fast
Fourier Transform (FFT) to compute (12) and to parametrize ${w}_{\theta}$ in
Fourier space. Following Li et al. (2021), this insight leads to the following
formulation of the operator given in Definition 19.

######  Definition 20 (Fourier Integral Operator $\mathsf{W}$)

We define the Fourier integral operator as

$\bigl{(}\mathsf{W}(\theta){u}\bigr{)}(x)=\mathcal{F}^{-1}\Bigl{(}R_{\mathsf{W}}(\theta)\cdot(\mathcal{F}{u})\Bigr{)}(x)\qquad\forall
u\in\mathcal{U}\bigl{(}D;\mathbbm{R}^{r}),\,\forall x\in D.$ (13)

To link to the definition (12) with translation-invariant kernel we take
$R_{\mathsf{W}}(\theta)$ to be the Fourier transform of the periodic function
${\kappa_{\theta}}:D\to\mathbbm{R}^{{r^{\prime}}\times r}$ parametrized by
$\theta\in\Theta_{\mathsf{W}}$.

For uniform discretizations $n_{1}\times\ldots\times n_{d}=N$ of the space
$D$, the Fourier transform can be replaced with the Fast Fourier Transform
(FFT). Indeed, for the Fourier transform we have
$\mathcal{F}u:\mathbbm{R}^{d}\to\mathbbm{C}^{r}$ and
$R_{\mathsf{W}}(\theta):\mathbbm{R}^{d}\to\mathbbm{C}^{r^{\prime}\times r}$.
On the other hand, for the FFT we have
$\mathcal{F}u:\mathbbm{Z}^{d}\to\mathbbm{C}^{r}$ and
$R_{\mathsf{W}}(\theta):\mathbbm{Z}^{d}\to\mathbbm{C}^{r^{\prime}\times r}$.
In practice, it is possible to select a finite-dimensional parametrization by
choosing a finite set of wave numbers

$\bigl{|}Z_{k_{\textrm{max}}}\bigr{|}:=\bigl{|}\bigl{\\{}k\in\mathbbm{Z}^{d}:|k_{j}|\leq
k_{\textrm{max},j},\,\text{for}\,j=1,\ldots,d\bigr{\\}}\bigr{|}.$

We can therefore parametrize $R_{\mathsf{W}}$ by a complex
$(2k_{\text{max},1}\times\ldots\times 2k_{\text{max},d}\times r^{\prime}\times
r)$-tensor.

These integral operator definitions will be used to define architectures for
the Vision Transformer Neural Operator in Subsection 4.3 and the Fourier
Attention Neural Operator in Subsection 4.4, but are not needed to define the
basic Transformer Neural Operator in Subsection 4.2. In the context of each
architecture, we make explicit the parametrizations of the integral operators
employed.

###  4.2 Transformer Neural Operator

Figure 1: Transformer Neural Operator.

The transformer neural operator architecture follows the general structure of
Equations 9 and 10; here, we make explicit the choices for
$\mathsf{T}_{\textrm{in}},\mathsf{T}_{\textrm{out}},\mathsf{E}_{L}$ that
instantiate it. In Figure 1 we provide a schematic representation of the full
transformer neural operator architecture. We consider model inputs
$u\in\mathcal{U}(D;\mathbbm{R}^{d_{u}})$, model embedding space
$\mathcal{V}:=\mathcal{U}(D;\mathbbm{R}^{d_{\textrm{model}}})$, and model
output space $\mathcal{Z}(D;\mathbbm{R}^{d_{z}})$, with
$D\subset\mathbbm{R}^{d}$. The function
$u\in\mathcal{U}(D;\mathbbm{R}^{d_{u}})$ is first concatenated with the
identity defined on the domain $D$, so that the resulting function
$u_{\textrm{in}}\in\mathcal{U}(D;\mathbbm{R}^{d_{u}+d})$ is defined by its
action on $x\in D$ as

$u_{\textrm{in}}(x)\coloneqq\bigl{(}u(x),x\bigr{)}.$ (14)

In practice, this step defines a positional encoding that is appended to the
input function. A linear transformation
$W_{\textrm{in}}\in\mathbbm{R}^{d_{\textrm{model}}\times(d_{u}+d)}$ is then
applied pointwise to lift the function $u_{\textrm{in}}$ into the embedding
space $\mathcal{V}$. The input operator
$\mathsf{T}_{\textrm{in}}:\mathcal{U}(D;\mathbbm{R}^{d_{u}})\to\mathcal{V}$ is
hence defined as

$\Bigl{(}\mathsf{T}_{\textrm{in}}(u)\Bigr{)}(x):=W_{\textrm{in}}u_{\textrm{in}}(x)$
(15)

for $u_{\textrm{in}}\in\mathcal{U}\bigl{(}D;\mathbbm{R}^{d_{u}+d}\bigr{)}$
defined as in (14), where $x\in D$, $u(x)\in\mathbbm{R}^{d_{u}}$, and
$W_{\textrm{in}}\in\mathbbm{R}^{d_{\textrm{model}}\times(d_{u}+d)}$. Observe
that $\mathsf{T}_{\textrm{in}}$ acts as a local, pointwise linear operator on
$u_{\textrm{in}}$.

Next, we define the details of $\mathsf{E}_{L}:\mathcal{V}\to\mathcal{V}$ in
Equation 10 by specifying a definition of $\mathsf{A}_{\textrm{MultiHead}}$,
$\mathsf{W}_{1},\mathsf{W}_{2}$, $\mathsf{F}_{\textrm{LayerNorm}}$, and
$\mathsf{F}_{\textrm{NN}}$ for the transformer neural operator setting. The
operator $\mathsf{A}_{\textrm{MultiHead}}:\mathcal{V}\to\mathcal{V}$ is the
multi-head attention operator from (11) based on the continuum self-attention
from Definition 5. Letting $H\in\mathbbm{N}$ denote the number of attention
heads, the multi-head attention operator is parametrized by the learnable
linear transformations $Q_{h}\in\mathbbm{R}^{d_{K}\times
d_{\textrm{model}}},K_{h}\in\mathbbm{R}^{d_{K}\times
d_{\textrm{model}}},V_{h}\in\mathbbm{R}^{d_{K}\times d_{\textrm{model}}}$, for
$h=1,\ldots,H$, where for implementation purposes $d_{K}$ is chosen as
$d_{K}\coloneqq d_{\textrm{model}}/H$. 555We note that $d_{\textrm{model}}$
and $H$ should be chosen so that $d_{\textrm{model}}$ is a multiple of $H$.
The multi-head attention operator is thus defined by the application of a
linear transformation
$W_{\textrm{MultiHead}}\in\mathbbm{R}^{d_{\textrm{model}}\times H\cdot d_{K}}$
to the concatenation of the outputs of $H\in\mathbbm{N}$ self-attention
operators. The operators
$\mathsf{W}_{1},\mathsf{W}_{2}:\mathcal{V}\to\mathcal{V}$ are pointwise linear
operators and hence admit finite-dimensional representations
$W_{1},W_{2}\in\mathbbm{R}^{d_{\textrm{model}}\times d_{\textrm{model}}}$, so
that they define a map

$v(x)\mapsto W_{1}v(x),$

for each $x\in D$, and similarly for $W_{2}$. The operator
$\mathsf{F}_{\textrm{LayerNorm}}:\mathcal{V}\to\mathcal{V}$ is defined such
that

$\Bigl{(}\mathsf{F}_{\textrm{LayerNorm}}({v};\gamma,\beta)(x)\Bigr{)}_{k}=\gamma_{k}\cdot\frac{\bigl{(}v(x)\bigr{)}_{k}-m\bigl{(}v(x)\bigr{)}}{\sqrt{\sigma^{2}\bigl{(}v(x)\bigr{)}+\epsilon}}+\beta_{k},$
(16)

for $k=1,\dots,d_{\textrm{model}}$, any $v\in\mathcal{V}$, and any $x\in
D\subset\mathbbm{R}^{d}$, where the notation $(\,\cdot\,)_{k}$ is used to
denote the $k$’th entry of the vector. In equation (16),
$\epsilon\in\mathbbm{R}^{+}$ is a fixed parameter,
$\gamma_{k},\beta_{k}\in\mathbbm{R}$ are learnable parameters and $m,\sigma$
are defined as

$\displaystyle m(y)$
$\displaystyle\coloneqq\frac{1}{d_{\textrm{model}}}\cdot|y|_{1},$ (17)
$\displaystyle\sigma^{2}(y)$
$\displaystyle\coloneqq\frac{1}{d_{\textrm{model}}}\sum_{k=1}^{d_{\textrm{model}}}\bigl{(}y_{k}-m(y)\bigr{)}^{2},$

for any $y\in\mathbbm{R}^{d_{\textrm{model}}}.$ The operator
$\mathsf{F}_{\textrm{NN}}:\mathcal{V}\to\mathcal{V}$ is defined such that

$\mathsf{F}_{\textrm{NN}}(v;W_{3},W_{4},b_{1},b_{2})(x)=W_{3}f\bigl{(}W_{4}v(x)+b_{1}\bigr{)}+b_{2},$
(18)

for any $v\in\mathcal{V}$ and $x\in D$, where
$W_{3},W_{4}\in\mathbbm{R}^{d_{\textrm{model}}\times d_{\textrm{model}}}$ and
$b_{1},b_{2}\in\mathbbm{R}^{d_{\textrm{model}}}$ are learnable parameters and
where $f$ is a nonlinear activation function.

Finally, we define the output operator
$\mathsf{T}_{\textrm{out}}:\mathcal{V}\to\mathcal{Z}(D;\mathbbm{R}^{d_{z}})$
as a local linear operator applied pointwise, so that

$\Bigl{(}\mathsf{T}_{\textrm{out}}({v})\Bigr{)}(x)=W_{\textrm{out}}v(x),$ (19)

for $x\in D$, $v\in\mathcal{V}$ and
$W_{\textrm{out}}\in\mathbbm{R}^{d_{z}\times d_{\textrm{model}}}$.

The composition of the operators
$\mathsf{T}_{\textrm{in}},\mathsf{E}_{L},\mathsf{T}_{\textrm{out}}$ completes
the definition of the transformer neural operator. In Appendix B.1 we outline
how this neural operator architecture is implemented in the finite-dimensional
setting, where the input function
$u\in\mathcal{U}\bigl{(}D;\mathbbm{R}^{d_{u}}\bigr{)}$ is defined on a finite
set of discretization points $\\{x_{i}\\}_{i=1}^{N}\subset D$.

###  4.3 Vision Transformer Neural Operator

We introduce an analogue of the vision transformer architecture as a neural
operator using the patching in the continuum framework developed in Section 3.
We note that in the context of the ViT from Dosovitskiy et al. (2021),
performing patch “flattening” before a local lifting to the embedding
dimension introduces a nonlocal transformation on the patches, as described in
Remark 15. This procedure as implemented in Dosovitskiy et al. (2021) violates
the desirable (for neural operators) discretization invariance property.
Hence, for a neural operator analogue, we substitute the lifting to embedding
dimension with a nonlocal integral operator. The ViT neural operator (ViT NO)
introduced is thus discretization invariant. We describe this methodology
within the framework provided by making explicit choices for
$\mathsf{T}_{\textrm{in}},\mathsf{T}_{\textrm{out}},\mathsf{E}_{L}$ that
instantiate Equations 9 and 10 appropriately. In Figure 2 we provide a
schematic for the architecture, with a graphical representation of the encoder
layer in Figure 3.

Figure 2: Vision Transformer Neural Operator.

Let $D\subset\mathbbm{R}^{d}$ be a bounded open set. Consider
$u\in\mathcal{U}(D;\mathbbm{R}^{d_{u}})$ and let $D\coloneqq
D_{1}\cup\ldots\cup D_{P}$ be a uniform partition of the space $D$ so that
$D_{j}\cong D^{\prime}$ for all $j\in\llbracket 1,P\rrbracket$ for some
$D^{\prime}\subset D$, where $P\in\mathbb{N}$ represents the number of
patches. We consider model inputs $u\in\mathcal{U}(D;\mathbbm{R}^{d_{u}})$,
model embedding space $\mathcal{V}:=\left\\{v\mid v:\llbracket
1,P\rrbracket\to\mathcal{U}\bigl{(}D^{\prime};\mathbbm{R}^{d_{\textrm{model}}}\bigr{)}\right\\}$
which is a space of functions acting on patch indices, and model output space
$\mathcal{Z}(D;\mathbbm{R}^{d_{z}})$, with domain $D\subset\mathbbm{R}^{d}$.

We begin by defining the input function
$\mathsf{T}_{\textrm{in}}:\mathcal{U}(D;\mathbbm{R}^{d_{u}})\to\mathcal{V}$.
The function $u\in\mathcal{U}(D;\mathbbm{R}^{d_{u}})$ is first concatenated
with the identity defined on the domain $D$, so that the resulting function
$u_{\textrm{in}}\in\mathcal{U}(D;\mathbbm{R}^{d_{u}+d})$ is defined as in
(14); recall that this defines a positional encoding. An operator
$\mathsf{S}_{\textrm{reshape}}$ is applied to the function
${u}_{\textrm{in}}\in\mathcal{U}(D;\mathbbm{R}^{d_{u}+d})$ that acts on the
input so that the resulting function is of the form

$\mathsf{S}_{\textrm{reshape}}({u}_{\textrm{in}}):\llbracket
1,P\rrbracket\to\mathcal{U}(D^{\prime},\mathbb{R}^{d_{u}+d}).$ (20)

We apply a nonlocal linear operator
$\mathsf{W}_{\textrm{in}}\in\mathcal{L}\Bigl{(}\mathcal{U}(D^{\prime};\mathbbm{R}^{d_{u}+d})\,;\,\mathcal{U}\bigl{(}D^{\prime};\mathbbm{R}^{d_{\textrm{model}}}\bigr{)}\Bigr{)}$
that acts on any $u:\llbracket
1,P\rrbracket\to\mathcal{U}(D^{\prime};\mathbbm{R}^{d_{u}+d})$ so that for any
$j\in\llbracket 1,P\rrbracket$

$u(j)\mapsto\mathsf{W}_{\textrm{in}}u(j).$ (21)

The above nonlocal linear operator is chosen as an integral operator as in
Definition 20 and the subsequent discussion. The composition of the operators
defined in Equations 14, 20 and 21 defines our choice of
$\mathsf{T}_{\textrm{in}}$. Namely, we define
$\mathsf{T}_{\textrm{in}}:\mathcal{U}(D;\mathbbm{R}^{d_{u}})\to\mathcal{V}$
via

$\Bigl{(}\mathsf{T}_{\textrm{in}}(u)\Bigr{)}(j)\coloneqq\mathsf{W}_{\textrm{in}}\Bigl{(}\mathsf{S}_{\textrm{reshape}}\bigl{(}{u}_{\textrm{in}}\bigr{)}(j)\Bigr{)},$
(22)

where ${u}_{\textrm{in}}\in\mathcal{U}(D;\mathbbm{R}^{d_{u}+d})$ is defined as
in (14), $j\in\llbracket 1,P\rrbracket$ and $\mathsf{W}_{\textrm{in}}$ is
defined as in (21).

Figure 3: Vision Transformer Neural Operator Encoder Layer.

Next, we define the details of $\mathsf{E}_{L}:\mathcal{V}\to\mathcal{V}$ in
Equation 10 by specifying a definition of $\mathsf{A}_{\textrm{MultiHead}}$,
$\mathsf{F}_{\textrm{LayerNorm}}$, $\mathsf{W}_{1},\mathsf{W}_{2}$, and
$\mathsf{F}_{\textrm{NN}}$ in this patched setting. In each layer the operator
$\mathsf{A}_{\textrm{MultiHead}}:\mathcal{V}\to\mathcal{V}$ is the multi-head
attention operator from (11) based on the self-attention from Definition 14.
Letting $H\in\mathbbm{N}$ denote the number of attention heads, the multi-head
attention operator is parametrized by learnable local (pointwise) linear
operators $\mathsf{Q}_{h},\mathsf{K}_{h},\mathsf{V}_{h}$ for $h=1,\ldots,H$.
In this setting, these are chosen to be linear operators applied pointwise,
they can be represented by finite dimensional linear transformations
$Q_{h}\in\mathbbm{R}^{d_{K}\times
d_{\textrm{model}}},K_{h}\in\mathbbm{R}^{d_{K}\times
d_{\textrm{model}}},V_{h}\in\mathbbm{R}^{d_{K}\times d_{\textrm{model}}}$, for
$h=1,\ldots,H$. 666We note that in the context of the Fourier attention neural
operator, presented in the next Subsection 4.4, we generalize this definition
to $\mathsf{Q}_{h},\mathsf{K}_{h},\mathsf{V}_{h}$ for $h=1,\ldots,H$ being
nonlocal linear integral operators. For implementation purposes $d_{K}$ is
chosen as $d_{K}\coloneqq d_{\textrm{model}}/H$. 777We note that
$d_{\textrm{model}}$ and $H$ should be chosen so that $d_{\textrm{model}}$ is
a multiple of $H$. The multi-head attention operator is thus defined by the
application of a local linear transformation
$W_{\textrm{MultiHead}}\in\mathbbm{R}^{d_{\textrm{model}}\times H\cdot d_{K}}$
to the concatenation of the outputs of $H\in\mathbbm{N}$ self-attention
operators. The operators
$\mathsf{W}_{1},\mathsf{W}_{2}:\mathcal{V}\to\mathcal{V}$ are pointwise linear
operators and hence admit finite-dimensional representations
$W_{1},W_{2}\in\mathbbm{R}^{d_{\textrm{model}}\times d_{\textrm{model}}}$, so
that they define a map

$v(x;p)\mapsto W_{1}v(x;p),$ (23)

for each $x\in D^{\prime}$ and $p\in\llbracket 1,P\rrbracket$, and similarly
for $W_{2}$. The operator
$\mathsf{F}_{\textrm{LayerNorm}}:\mathcal{V}\to\mathcal{V}$ is defined such
that

$\Bigl{(}\mathsf{F}_{\textrm{LayerNorm}}(v;\gamma,\beta)(x;p)\Bigr{)}_{k}=\gamma_{k}\cdot\frac{\bigl{(}v(x;p)\bigr{)}_{k}-m\bigl{(}v(x;p)\bigr{)}}{\sqrt{\sigma^{2}\bigl{(}v(x;p)\bigr{)}+\epsilon}}+\beta_{k},$
(24)

for $k=1,\dots,d_{\textrm{model}}$, any $p\in\llbracket 1,P\rrbracket$ and any
$x\in\mathbbm{R}^{d}\subset D$, where we use the notation $(\,\cdot\,)_{k}$ to
denote the $k$’th entry of the vector. In equation (24),
$\epsilon\in\mathbbm{R}^{+}$ is a fixed parameter,
$\gamma_{k},\beta_{k}\in\mathbbm{R}$ are learnable parameters and $m,\sigma$
are defined as in (17). The operator
$\mathsf{F}_{\textrm{NN}}:\mathcal{V}\to\mathcal{V}$ is defined by

$\mathsf{F}_{\textrm{NN}}(v;W_{3},W_{4},b_{1},b_{2})(x;p)=W_{3}f\bigl{(}W_{4}v(x;p)+b_{1}\bigr{)}+b_{2},$
(25)

for any $x\in D$ and any $p\in\llbracket 1,P\rrbracket$, where
$W_{3},W_{4}\in\mathbbm{R}^{d_{\textrm{model}}\times d_{\textrm{model}}}$ and
$b_{1},b_{2}\in\mathbbm{R}^{d_{\textrm{model}}}$ are learnable parameters and
$f$ is a nonlinear activation function.

Finally, we define the action of
$\mathsf{T}_{\textrm{out}}:\mathcal{V}\to\mathcal{Z}(D;\mathbbm{R}^{d_{z}})$.
This is given by first applying a reshaping operator
$\mathsf{S}^{\prime}_{\textrm{reshape}}$ to the output of the encoder, where
this map is defined as

${v}\mapsto\mathsf{S}^{\prime}_{\textrm{reshape}}\bigl{(}{v}\bigr{)}\in\mathcal{U}(D;\mathbb{R}^{d_{\textrm{model}}}).$
(26)

We note that the operator $\mathsf{S}^{\prime}_{\textrm{reshape}}$ may be
viewed as an inverse of $\mathsf{S}_{\textrm{reshape}}$. A pointwise linear
transformation $W_{\textrm{out}}\in\mathbbm{R}^{d_{z}\times
d_{\textrm{model}}}$ defined by

$v(x)\mapsto W_{\textrm{out}}{v}(x)\in\mathcal{Z}(D;\mathbbm{R}^{d_{z}}),$
(27)

for any $x\in D$, is then applied. This procedure completes the definition of
the operator
$\mathsf{T}_{\textrm{out}}:\mathcal{V}\to\mathcal{Z}(D;\mathbbm{R}^{d_{z}})$,
which is given by

$\Bigl{(}\mathsf{T}_{\textrm{out}}(v)\Bigr{)}(x)\coloneqq
W_{\textrm{out}}\Bigl{(}\mathsf{S}^{\prime}_{\textrm{reshape}}\bigl{(}v\bigr{)}(x)\Bigr{)},$
(28)

for any $v\in\mathcal{V}$ and $x\in D$.

###### Remark 21

We have found that adding a convolution in residual form helps to ameliorate
the discontinuities arising when the patching architecture is used; these are
visible in error plots and can also interfere with self-composition of maps
learned as solution operators of time-dependent PDEs. To this end define the
linear integral operator
$\mathsf{W}_{\textrm{S}}\in\mathcal{L}\Bigl{(}\mathcal{Z}\bigl{(}D;\mathbbm{R}^{d_{z}}\bigr{)}\,;\,\mathcal{Z}\bigl{(}D;\mathbbm{R}^{d_{z}}\bigr{)}\Bigr{)}$,
of the form described in Definition 20 and the subsequent discussion, and add
a last layer to the output of the architecture as defined thus far:

${z}\mapsto z+\mathsf{W}_{\rm{S}}z\in\mathcal{Z}(D;\mathbbm{R}^{d_{z}}),$ (29)

for
$z\coloneqq\Bigl{(}\mathsf{T}_{\rm{out}}\circ\mathsf{E}_{L}\circ\mathsf{T}_{\rm{in}}\Bigr{)}(u)$.

The composition of the operators
$\mathsf{T}_{\textrm{in}},\mathsf{E}_{L},\mathsf{T}_{\textrm{out}}$ completes
the definition of the ViT neural operator. In Appendix B.2 we outline how this
neural operator architecture is implemented in the finite-dimensional setting,
where the input function
$u\in\mathcal{U}\bigl{(}D;\mathbbm{R}^{d_{u}}\bigr{)}$ is defined on a finite
set of discretization points $\\{x_{i}\\}_{i=1}^{N}\subset D$.

###  4.4 Fourier Attention Neural Operator

In this subsection, we introduce a different transformer neural operator based
on patching, the Fourier attention neural operator (FANO). In this setting, to
gain additional expressivity, we replace the local linear operators
$\mathsf{Q},\mathsf{K},\mathsf{V}$ in the attention mechanism with nonlocal
integral operators. In Figure 4 we provide a schematic for the Fourier
attention neural operator architecture, with a graphical representation of the
encoder layer in Figure 5. We outline the details of the architecture in the
following discussion.

We again let $D\subseteq\mathbbm{R}^{d}$ be a bounded open set. Consider
$u\in\mathcal{U}(D;\mathbbm{R}^{d_{u}})$ and let $D\coloneqq
D_{1}\cup\ldots\cup D_{P}$ be a uniform partition of the space $D$ so that
$D_{j}\cong D^{\prime}$ for all $j\in\llbracket 1,P\rrbracket$ for some
$D^{\prime}\subset D$, where $P\in\mathbb{N}$ represents the number of
patches. We consider model inputs $u\in\mathcal{U}(D;\mathbbm{R}^{d_{u}})$,
model embedding space $\mathcal{V}:=\left\\{v\mid v:\llbracket
1,P\rrbracket\to\mathcal{U}\bigl{(}D^{\prime};\mathbbm{R}^{d_{\textrm{model}}}\bigr{)}\right\\}$
which is a space of functions acting on patch indices , and model output space
$\mathcal{Z}(D;\mathbbm{R}^{d_{z}})$, with domain $D\subset\mathbbm{R}^{d}$.

Figure 4: Fourier Attention Neural Operator.

We define $\mathsf{T}_{\textrm{in}}$ similarly to (22) in the previous
section, via a composition of Equations 14, 20 and 21; however here the linear
lifting operator is chosen to be a pointwise linear transformation
$W_{\textrm{in}}\in\mathbbm{R}^{d_{\textrm{model}}\times(d_{u}+d)}$. Namely,
we define
$\mathsf{T}_{\textrm{in}}:\mathcal{U}(D;\mathbbm{R}^{d_{u}})\to\mathcal{V}$
via

$\Bigl{(}\mathsf{T}_{\textrm{in}}(u)\Bigr{)}\bigl{(}j\bigr{)}(x)\coloneqq
W_{\textrm{in}}\Bigl{(}\mathsf{S}_{\textrm{reshape}}\bigl{(}{u}_{\textrm{in}}\bigr{)}(j)\Bigr{)}(x),$
(30)

for any $j\in\llbracket 1,P\rrbracket$ and any $x\in D$, where
${u}_{\textrm{in}}\in\mathcal{U}(D;\mathbbm{R}^{d_{u}+d})$ is defined as in
(14), $\mathsf{S}_{\textrm{reshape}}$ as in $\eqref{eq:reshape}$ and the
nonlocal operation (21) is replaced by a local pointwise one.

Next, we define the particulars of $\mathsf{E}_{L}:\mathcal{V}\to\mathcal{V}$
in Equation 10. The key distinction between the Fourier attention neural
operator that we define here and the ViT neural operator from the previous
subsection 4.3 lies in the nonlocality of operations performed within
$\mathsf{A}_{\textrm{MultiHead}}$. Here, in each layer the operator
$\mathsf{A}_{\textrm{MultiHead}}:\mathcal{V}\to\mathcal{V}$ is the multi-head
attention operator from (11) based on the self-attention operator from
Definition 14. Letting $H\in\mathbbm{N}$ denote the number of attention heads,
the multi-head attention operator is parametrized by learnable nonlocal linear
integral operators
$\mathsf{Q}_{h},\mathsf{K}_{h}\in\mathcal{L}\Bigl{(}\mathcal{U}(D^{\prime},\mathbb{R}^{d_{\textrm{model}}});L^{2}(D^{\prime},\mathbb{R}^{d_{K}})\Bigr{)}$
and
$\mathsf{V}_{h}\in\mathcal{L}\Bigl{(}\mathcal{U}(D^{\prime},\mathbb{R}^{d_{\textrm{model}}});\mathcal{U}(D^{\prime},\mathbb{R}^{d_{K}})\Bigr{)}$,
for $h=1,\ldots,H$, where for implementation purposes $d_{K}$ is chosen as
$d_{K}\coloneqq d_{\textrm{model}}/H$. The specific formulation of these
linear integral operators may be found in Definition 20 and the subsequent
discussion. The multi-head attention operator is thus defined by the
application of a pointwise linear transformation
$W_{\textrm{MultiHead}}\in\mathbbm{R}^{d_{\textrm{model}}\times H\cdot d_{K}}$
to a concatenation of the outputs of $H\in\mathbbm{N}$ self-attention
operators. The operators $\mathsf{W}_{1},\mathsf{W}_{2}$ are defined as in
(23). The layer normalization operator
$\mathsf{F}_{\textrm{LayerNorm}}:\mathcal{V}\to\mathcal{V}$ is defined as in
(24) so that it is applied to every point in every patch. Furthermore, the
operator $\mathsf{F}_{\textrm{NN}}:\mathcal{V}\to\mathcal{V}$ is defined as in
(25).

We define $\mathsf{T}_{\textrm{out}}$ identically to Equation 28. We refer to
Remark 21 for an additional operation that can be applied to ameliorate the
effect of patch discontinuities. Finally, the composition of the operators
$\mathsf{T}_{\textrm{in}},\mathsf{E}_{L},\mathsf{T}_{\textrm{out}}$ completes
the definition of the Fourier attention neural operator. In Appendix B.3 we
outline how this neural operator architecture is implemented in the finite-
dimensional setting, where the input function
$u\in\mathcal{U}\bigl{(}D;\mathbbm{R}^{d_{u}}\bigr{)}$ is defined on a finite
set of discretization points $\\{x_{i}\\}_{i=1}^{N}\subset D$.

Figure 5: Encoder Layer of the Fourier Attention Neural Operator.

##  5 Universal Approximation by Transformer Neural Operators

Universal approximation is a minimal requirement for neural operators. When
the operator to be approximated maps between spaces of functions defined over
Euclidean domains, universal approximation necessarily requires both
nonlocality and nonlinearity. In the recent paper Lanthaler et al. (2023), a
minimal architecture possessing these two properties is exhibited. In this
section, rather than trying to prove universal approximation for the various
discretization-invariant neural operators introduced in the paper, we
construct a simple canonical setting that exhibits the universality of the
attention mechanism on function space, allowing direct exploitation of the
ideas in Lanthaler et al. (2023). We build on the notation used in Subsection
4.1. In particular $\mathsf{A}$ denotes the self-attention operator from
Definition 5 and $f$ an activation function. Throughout this section, we
consider activation functions $f\in C^{\infty}(\mathbbm{R})$ which are non-
polynomial and Lipschitz continuous.

To be specific, we consider neural operators
$\mathsf{G}:\mathcal{U}\bigl{(}D;\mathbbm{R}^{d_{u}}\bigr{)}\times\Theta\to\mathcal{Z}\bigl{(}D;\mathbbm{R}^{d_{z}}\bigr{)}$
of the form

$\mathsf{G}(u;\theta)\coloneqq\Bigl{(}\mathsf{T}_{\textrm{out}}\circ\mathsf{E}\circ\mathsf{T}_{\textrm{in}}\Bigr{)}(u;\theta),$
(31)

where
$\mathsf{T}_{\textrm{in}}:\mathcal{U}\bigl{(}D;\mathbbm{R}^{d_{u}}\bigr{)}\to\mathcal{V}\bigl{(}D;\mathbbm{R}^{d_{\textrm{model}}}\bigr{)}$
and
$\mathsf{T}_{\textrm{out}}:\mathcal{V}\bigl{(}D;\mathbbm{R}^{d_{\textrm{model}}}\bigr{)}\to\mathcal{Z}\bigl{(}D;\mathbbm{R}^{d_{z}}\bigr{)}$
are defined by neural networks of the form

$\displaystyle\Bigl{(}\mathsf{T}_{\textrm{in}}(u)\Bigr{)}(x)$
$\displaystyle=R_{2}f\Bigl{(}R_{1}\bigl{(}u(x),x\bigr{)}+b_{R}\Bigr{)}+b_{R}^{\prime},$
(32) $\displaystyle\Bigl{(}\mathsf{T}_{\textrm{out}}(v)\Bigr{)}(x)$
$\displaystyle=P_{2}f\Bigl{(}P_{1}\bigl{(}v(x),x\bigr{)}+b_{P}\Bigr{)}+b_{P}^{\prime},$
(33)

where $\bigl{(}u(x),x\bigr{)}\in\mathbbm{R}^{d_{u}+d}$ and
$\bigl{(}v(x),x\bigr{)}\in\mathbbm{R}^{d_{\textrm{model}}+d}$, where
$R_{1},R_{2},P_{1},P_{2}$ are learned linear transformations of appropriate
dimensions and $b_{R},b_{R}^{\prime},b_{P},b_{P}^{\prime}$ are learned
vectors. We define the operator
$\mathsf{E}:\mathcal{V}\bigl{(}D;\mathbbm{R}^{d_{\textrm{model}}}\bigr{)}\to\mathcal{V}\bigl{(}D;\mathbbm{R}^{d_{\textrm{model}}}\bigr{)}$
as a variant of the layer defined by the iteration step in Equation 10, given
by the two-step map acting on its inputs $v\in\mathcal{V}$ as

$\displaystyle{v}(x)$ $\displaystyle\mapsfrom
W_{1}{v}(x)+\mathsf{A}\big{(}{v};Q,K,V\big{)}(x),$ (34a) $\displaystyle v(x)$
$\displaystyle\mapsfrom
W_{2}{v}(x)+W_{3}f\bigl{(}W_{4}{v}(x)+b_{1}\bigr{)}+b_{2},$ (34b)

for any $x\in D$. Note that the neural operator thus defined does not include
layer normalizations; it is thus a variant of the transformer neural operator
from Subsection 4.2. We may now apply the result of Lanthaler et al. (2023) to
show two universal approximation theorems for the resulting transformer neural
operator.

###### Theorem 22

Let $D\subset\mathbbm{R}^{d}$ be a bounded domain with Lipschitz boundary, and
fix integers $s,s^{\prime}\geq 0$. If
$\mathsf{G}^{\dagger}:C^{s}\bigl{(}\bar{D};\mathbbm{R}^{r}\bigr{)}\to
C^{s^{\prime}}\bigl{(}\bar{D};\mathbbm{R}^{r^{\prime}}\bigr{)}$ is a
continuous operator and $K\subset
C^{s}\bigl{(}\bar{D};\mathbbm{R}^{r}\bigr{)}$ a compact set, then for any
$\epsilon>0$, there exists a transformer neural operator
$\mathsf{G}(\cdot;\theta):K\subset
C^{s}\bigl{(}\bar{D};\mathbbm{R}^{r}\bigr{)}\to
C^{s^{\prime}}\bigl{(}\bar{D};\mathbbm{R}^{r^{\prime}}\bigr{)}$ so that

$\sup_{u\in
K}\left\|\mathsf{G}^{\dagger}(u)-\mathsf{G}(u;\theta)\right\|_{C^{s^{\prime}}}\leq\epsilon.$
(35)

Proof  We begin by noting that for $Q,K=0$ and $V=I$, the self-attention
mapping reduces to

$\mathsf{A}\bigl{(}v;Q,K,V\bigr{)}(\cdot)=\mathbb{E}_{y\sim
p(y:v,x)}[Vv(\cdot)]=\frac{1}{|D|}\int v(x)~{}\mathrm{d}x.$ (36)

For weights $W_{2}=0$, $W_{3}=W_{4}=I$ and $b_{2}=0$, the transformer encoder
neural operator layer (34) reduces to the mapping

$v(\cdot)\mapsto f\left(W_{1}v(\cdot)+b_{1}+\frac{1}{|D|}\int
v(x)~{}\mathrm{d}x\right).$ (37)

The existence of
$W_{1},R_{1},R_{2},P_{1},P_{2},b_{1},b_{R},b_{R}^{\prime},b_{P},b_{P}^{\prime}$
so that $\mathsf{G}(\cdot\,;\theta)$ satisfies (35) then follows from
Lanthaler et al. (2023, Theorem 2.1), which also involves the application of
the universality result for two-layer neural networks of Pinkus (1999, Theorem
4.1).

The analysis in Lanthaler et al. (2023) allows the derivation of an analogous
universal approximation theorem for functions belonging to Sobolev spaces.
This generalization is the content of the next theorem. The proof follows
easily by applying the same argument as in the proof of Theorem 22 and the
result of Lanthaler et al. (2023, Theorem 2.2).

###### Theorem 23

Let $D\subset\mathbbm{R}^{d}$ be a bounded domain with Lipschitz boundary and
fix integers $s,s^{\prime}\geq 0$, $p,p^{\prime}\in[1,\infty)$. If
$\mathsf{G}^{\dagger}:W^{s,p}\bigl{(}D;\mathbbm{R}^{r}\bigr{)}\to
W^{s^{\prime},p^{\prime}}\bigl{(}D;\mathbbm{R}^{r^{\prime}}\bigr{)}$ is a
continuous operator and $K\subset W^{s,p}\bigl{(}D;\mathbbm{R}^{r}\bigr{)}$ a
compact set of bounded functions so that $\sup_{u\in
K}\|u\|_{L^{\infty}}<\infty$, then for any $\epsilon>0$, there exists a
transformer neural operator $\mathsf{G}(\cdot;\theta):K\subset
W^{s,p}\bigl{(}D;\mathbbm{R}^{r}\bigr{)}\to
W^{s^{\prime},p^{\prime}}\bigl{(}D;\mathbbm{R}^{r^{\prime}}\bigr{)}$ so that

$\sup_{u\in
K}\left\|\mathsf{G}^{\dagger}(u)-\mathsf{G}(u;\theta)\right\|_{W^{s^{\prime},p^{\prime}}}\leq\epsilon.$
(38)

##  6 Numerical Experiments

In this section we illustrate, through numerical experiments, the capability
of the transformer neural operator architectures described in Section 4.
Throughout, we consider the supervised learning problem described by the data
model in (3) and take the viewpoint of surrogate modeling to construct
approximations of the operators $\mathsf{G}^{\dagger}$. In Subsection 6.1 we
consider problems given by dynamical systems and ordinary differential
equations, namely operators acting on functions defined on a one-dimensional
time domain. In this context, we explore the use of the transformer neural
operator for operator learning problems given by the Lorenz ‘63 dynamical
system and a controlled ODE. On the other hand, in Subsection 6.2 we consider
problems given by partial differential equations equations, namely operators
acting on functions defined on a two-dimensional spatial domain. In this
context we explore the use of the transformer neural operator and of the more
efficient ViT neural operator and Fourier attention neural operator
architectures for operator learning problems given by the Darcy flow and
Navier-Stokes equations.

######  Remark 24 (Relative Loss)

For each neural operator architecture we optimize for the parameters
$\theta\in\Theta$ according to the relative loss, i.e.

$\inf_{\theta\in\Theta}\,\frac{1}{J}\sum_{j=1}^{J}\left(\frac{\|\mathsf{G}\bigl{(}u^{(j)};\theta\bigr{)}-\mathsf{G}^{\dagger}\bigl{(}u^{(j)}\bigr{)}\|_{\mathcal{Z}}}{\|\mathsf{G}^{\dagger}\bigl{(}u^{(j)}\bigr{)}\|_{\mathcal{Z}}}\right),$
(39)

for data points $\\{u^{(j)}\\}_{j=1}^{J}$.

###  6.1 1D Time Domain

In this subsection we consider problems where attention is applied to
functions in the one-dimensional time domain. The first setting we consider is
that of the Lorenz ‘63 dynamical system (Lorenz, 1963), a simplified model for
atmospheric convection. The dynamics are given by

$\displaystyle{dx}$ $\displaystyle=\sigma(y-x){dt},$ (Lorenz ‘63)
$\displaystyle{dy}$ $\displaystyle=\bigl{(}x(\rho-z)-y\bigr{)}{dt},$
$\displaystyle{dz}$ $\displaystyle=\bigl{(}xy-\beta z\bigr{)}{dt}.$

We consider learning the operator
$\mathsf{G}^{\dagger}:C\bigl{(}[0,T];\mathbbm{R}^{d_{x}}\bigr{)}\times\mathbbm{R}^{d_{y}}\times\mathbbm{R}^{d_{z}}\to
C\bigl{(}[0,T];\mathbbm{R}^{d_{y}}\bigr{)}\times
C\bigl{(}[0,T];\mathbbm{R}^{d_{z}}\bigr{)}$ given by

$\mathsf{G}^{\dagger}:\\{x(t),y(0),z(0)\\}_{t\in{[0,T]}}\mapsto\\{y(t),z(t)\\}_{t\in{[0,T]}}.$
(40)

In this case, $\mathsf{G}^{\dagger}$ is known to exist and given by the
solution operator for the system of linear ODEs arising from the second and
third components of (Lorenz ‘63), given known first component. We note that
including the initial condition of the unobserved trajectories in the model
introduces a discrepancy in dimension within the input function to the model.
Namely, denoting by $u$ the input to the model, $u(0)\in\mathbbm{R}^{3}$ while
$u(t)\in\mathbbm{R}$ for $t\neq 0$. We overcome this issue by learning
separate embedding linear transformations $W_{\textrm{in}_{0}}$ and
$W_{\textrm{in}}$, for the initial condition and the rest of the trajectory,
respectively.In Figure 6 we display the results obtained by applying a
transformer neural operator model on a test data set of $1000$ samples with
the same discretization ($\Delta t=0.01$ and $T=2$) as the training set of
$8000$ samples on the operator learning problem described by (40). The
operator is successfully learned, and displays median relative $L^{2}$ error
of less than $1\%$, with worst case rising to just below $10\%.$

Figure 6: The panel displays the performance of the transformer neural
operator when applied to the Lorenz ‘63 operator learning problem of
recovering both unobserved $y$ and $z$ trajectories. In each column, the top
plot shows the input $x$ trajectory data while the second and third row
display the predicted against true $y$ and $z$ trajectories, respectively; the
left column concerns the testing example achieving a median relative $L^{2}$
error, and the right columns shows the test example achieving the worst error.

We also experiment with a setting for which $\mathsf{G}^{\dagger}$ is not
well-defined: we study the recovery of unobserved trajectory
$\\{y(t)\\}_{t\in{[0,T]}}$ from the observed trajectory
$\\{x(t)\\}_{t\in{[0,T]}}$ so that
$\mathsf{G}^{\dagger}:C\bigl{(}[0,T];\mathbbm{R}^{d_{x}}\bigr{)}\to
C\bigl{(}[0,T];\mathbbm{R}^{d_{y}}\bigr{)}$ and

$\mathsf{G}^{\dagger}:\\{x(t)\\}_{t\in{[0,T]}}\mapsto\\{y(t)\\}_{t\in{[0,T]}}.$

This operator would only be well-defined if $(y(0),z(0))$ were included as
inputs. We nonetheless expect accurate recovery of the hidden trajectory by
the property of synchronization (Pecora and Carroll, 1990; Hayden et al.,
2011).

We investigate an additional operator learning problem given by the setting of
the controlled ODE

$\displaystyle{dz}$ $\displaystyle=\sin(z){du},$ $\displaystyle z(0)=z_{0},$
(CDE) $\displaystyle u(t)$
$\displaystyle=\sum_{j=1}^{J}\xi_{j}\sin(\pi\eta_{j}t),$
$\displaystyle\xi_{j}\sim\mathcal{N}(0,1),\,\eta_{j}\sim U(0,J).$

In this case we aim to approximate the operator
$\mathsf{G}^{\dagger}:C\bigl{(}[0,T];\mathbbm{R}^{d_{u}}\bigr{)}\to
C\bigl{(}[0,T];\mathbbm{R}^{d_{z}}\bigr{)}$ defined by

$\mathsf{G}:\\{u(t);\xi,\eta\\}_{t\in{[0,T]}}\mapsto\\{z(t);\xi,\eta\\}_{t\in{[0,T]}}.$

For this problem we make the choice $J=10$.

The experimental settings we consider for the continuous time operator
learning problem may be viewed from the lens of data assimilation as smoothing
problems (Sanz-Alonso et al., 2023). Indeed, the mappings to be learned
represent the recovery of an unobserved trajectory $\\{z(t)\\}_{t\in{[0,T]}}$
conditioned on the observation of a whole trajectory
$\\{u(t)\\}_{t\in{[0,T]}}$.

Figure 7: The panel displays the performance in relative $L^{2}$ errors (in
log-scale) of the transformer neural operator when applied to the Lorenz ‘63
and controlled ODE operator learning test sets of different resolutions at
inference time (without retraining). All models were parametrized by
$d_{\textrm{model}}=128$ and $6$ encoder layers, and were trained with $8000$
trajectories with $T=2$, sampled according to $\Delta t=0.01$.

In Figure 7 we demonstrate the mesh-invariance property of the transformer
neural operator in the context of the operator learning problems arising from
equations (Lorenz ‘63) and (CDE). A key motivation for considering the
continuum formulation of attention and the related transformer neural operator
is the property of zero-shot generalization to different non-uniform meshes.
To demonstrate this, in Figure 8 we consider the operator learning problem of
mapping the observed $x$ trajectory of (Lorenz ‘63) to the unobserved $y$
trajectory. We display the median relative $L^{2}$ error and worst error
samples from a test set consisting of trajectories
$\\{x(t)\\}_{t\in\mathfrak{T}_{\textrm{test}}}$, where the test time index set
$\mathfrak{T}_{\textrm{test}}$ is defined as

$\mathfrak{T}_{\textrm{test}}\coloneqq\\{n\Delta
t\\}_{n=0}^{N/2}\cup\\{2n\Delta t\\}_{n=N/2+1}^{3N/4},$ (41)

for $\Delta t\coloneqq T/N$, where the model deployed is trained on
trajectories indexed at a set $\mathfrak{T}_{\textrm{train}}$ defined by

$\mathfrak{T}_{\textrm{train}}\coloneqq\\{n\Delta t\\}_{n=0}^{N}.$ (42)

In other words, the second half of the testing domain is up-sampled by a
factor of 2 compared to the training discretization. As seen qualitatively in
Figure 8 testing on the non-uniform discretization based on a model trained on
the uniform grid is successful. The use of different grids in test and train
does leads to a larger error than the one incurred when they match (seen in
Figure 7), but the errors remain small. It is our continuum formulation of
attention from Section B.1 that enables this direct generalization; it would
not be possible with standard methods.

Figure 8: The panel displays the performance in relative $L^{2}$ errors of the
transformer neural operator when applied to the Lorenz ‘63 operator learning
problem of recovering the unobserved $y$ trajectory. The results displayed
concern a model trained on trajectories indexed on the index set (42) deployed
on a test of trajectories indexed on (41). In each column, the top plot shows
the input data and the bottom compares the prediction and ground truth; the
left column shows the testing example achieving a median error, and the right
columns shows the test example achieving the worst error.

###  6.2 2D Spatial Domain

In this subsection we study two different 2D operator learning problems – flow
in a porous medium, governed by the Darcy model, and incompressible viscous
fluid mechanics governed by the Navier-Stokes equation, in particular, a
Kolmogorov flow. We apply the three architectures developed in Sections 2 and
3, concentrating on the patched architectures. In the first Subsection 6.2.1
we report the results obtained by applying the neural operators described to
an array of operator learning problems, where the domain of the functions is
two-dimensional i.e. $D\subset\mathbbm{R}^{2}$. In Subsection 6.2.2 we
describe the settings of the two Darcy flow operator learning problems
considered, along with the parametrization and training details for the
various architectures. Similarly in Subsection 6.2.3, we describe the
Kolmogorov flow operator learning problem and relevant implementation details.
Subsections 6.2.2 and 6.2.3 also contain further numerical results
illustrating the behaviour of the proposed transformer neural operators, and
discretization invariance in particular.

####  6.2.1 Results

The results presented concern the two variants of Darcy flow, with lognormal
and piecewise constant inputs, and the Kolmogorov flow. The operator learning
architectures we deploy include the three transformer-based methods introduced
in this paper, and the implementations of the Fourier neural operator from the
“NeuralOperator” library (Li et al., 2021; Kovachki et al., 2023).

Our first results compares the transformer neural operator from Subsection 4.2
with FNO on the lognormal Darcy Flow problem in a low $64\times 64$ resolution
setting. Table 1 provides a brief summary of the test relative $L^{2}$ errors
obtained. It is notable that the transformer architecture obtains performance
similar to FNO, despite not using patching, and with an order of magnitude
fewer parameters. However since higher resolution renders the architecture
from Subsection 4.2 impractical in dimensions $d\geq 2$ the remainder of our
experiments use the patched ViT and Fourier attention neural operators from
Subsections 4.3 and 4.4 respectively.

Architecture | Number of Parameters | Darcy Flow
---|---|---
|  | Lognormal Input
Transformer NO | $5.98\cdot 10^{5}$ | $5.29\cdot 10^{-3}$
FNO | $1.11\cdot 10^{7}$ | $5.74\cdot 10^{-3}$
Table 1: Median relative $L^{2}$ error for each architecture applied to a low
resolution experimental setting.

We now consider training two patched architectures on the Darcy Flow problem
with piecewise constant inputs and the Kolmogorov Flow problem; in both cases
we use a $416\times 416$ resolution setting. We compare cost versus accuracy
for the different neural operators implemented, defining cost in terms of both
number of parameters and number of FLOPS. In Table 2 we present the FLOPS for
each method; details of the derivation are provided in Appendix C. We note
that the parameter scaling for the transformer neural operator is given by
$\mathcal{O}\bigl{(}d_{\textrm{model}}^{2}\bigr{)}$, while the scaling for the
FNO and the Fourier attention neural operator is
$\mathcal{O}\bigl{(}d_{\textrm{model}}^{2}\cdot k_{\textrm{max}}\bigr{)}$,
where $k_{\textrm{max}}$ is the total number of Fourier modes used. The
parameter scaling for the ViT neural operator is given by
$\mathcal{O}\bigl{(}d_{\textrm{model}}\cdot
k_{\textrm{max}}+d_{\textrm{model}}^{2}\bigr{)}$. Further details are provided
in Appendix C.

Architecture | Scaling
---|---
FNO | $\mathcal{O}\bigl{(}d_{\textrm{model}}N\log(\sqrt{N})+Nd^{2}_{\textrm{model}}\bigr{)}$
Transformer NO | $\mathcal{O}\bigl{(}d_{\textrm{model}}N^{2}+Nd^{2}_{\textrm{model}}\bigr{)}$
ViT NO | $\mathcal{O}\Bigl{(}d_{\textrm{model}}N\bigl{(}P+\log(\sqrt{N/P})\bigr{)}+N\log(\sqrt{N})+Nd^{2}_{\textrm{model}}\Bigr{)}$
FANO | $\par\mathcal{O}\Bigl{(}d_{\textrm{model}}N\bigl{(}P+\log(\sqrt{N/P})\bigr{)}+N\log(N)+Nd^{2}_{\textrm{model}}\Bigr{)}$
Table 2: Evaluation cost (measured in FLOPS) for the three proposed
transformer neural operators compared to the Fourier neural operator (Li et
al., 2021). We give further details on the derivation of these scalings in
Appendix C.

Figures 9 and 10 display the cost-accuracy trade-off. Table 3 provides the
test relative $L^{2}$ errors values, depicted in the cost-accuracy analysis
figures, obtained with the models with the highest parameter count. Figure 9
demonstrates the parameter-efficiency of the ViT neural operator in comparison
with both the FNO and Fourier attention neural operator. Once cost is defined
in terms of FLOPS it is apparent that both patched transformer architectures
outperform the FNO; furthermore the FNO exhibits accuracy-saturation, caused
by the finite data set, whereas the patched architectures are less affected,
in the ranges we test here. This suggests that the patched transformer
architectures use the information content in the data more efficiently.

Figure 9: The two panels display the test errors (in log-scale) against number of parameters for the architectures with the channel widths of 32, 64, 96 and 128. Figure 10: The two panels display the test errors (in log-scale) against evaluation cost in FLOPS for the architectures with the channel widths of 32, 64, 96 and 128. Architecture | Darcy Flow | Darcy Flow | Kolmogorov Flow
---|---|---|---
| Lognormal Input | Piecewise Constant Input |
ViT NO | $5.41\cdot 10^{-3}$ | $1.97\cdot 10^{-2}$ | $1.59\cdot 10^{-2}$
FANO | $4.41\cdot 10^{-3}$ | $1.68\cdot 10^{-2}$ | $8.56\cdot 10^{-3}$
FNO | $7.01\cdot 10^{-3}$ | $3.31\cdot 10^{-2}$ | $2.52\cdot 10^{-2}$
Table 3: Median relative $L^{2}$ error for each architecture applied to the
different high resolution experimental settings.

####  6.2.2 Darcy Flow

We consider the linear, second-order elliptic PDE defined on the unit square

$\displaystyle-\nabla\cdot\bigl{(}a(x)\nabla u(x)\bigr{)}$
$\displaystyle=f(x),$ $\displaystyle x\in(0,1)^{2},$ (Darcy Flow)
$\displaystyle u(x)$ $\displaystyle=0,$ $\displaystyle x\in\partial(0,1)^{2},$

which is the steady-state of the 2D Darcy flow equation. The equation is
equipped with Dirichlet boundary conditions and is well-posed when the
diffusion coefficient $a\in
L^{\infty}\bigl{(}(0,1)^{2};\mathbbm{R}_{+}\bigr{)}$ and the forcing function
is $f\in L^{2}\bigl{(}(0,1)^{2};\mathbbm{R}\bigr{)}$. In this context, we will
consider learning the nonlinear operator
$\mathsf{G}^{\dagger}:L^{\infty}\bigl{(}(0,1)^{2};\mathbbm{R}_{+}\bigr{)}\to
H^{1}_{0}\bigl{(}(0,1)^{2};\mathbbm{R}\bigr{)}$ defined as

$\mathsf{G}^{\dagger}:a\mapsto u.$

We consider two different inputs. In both experimental settings the forcing
function is chosen as $f\equiv 1$.

In the first the data points $\\{a^{(j)}\\}_{j=1}^{J}$ are sampled from the
probability measure

$\mu_{1}\coloneqq({T_{1}})_{\sharp}\mathcal{N}\bigl{(}0,C\bigr{)},$

where $\mathcal{N}\bigl{(}0,C\bigr{)}$ is a Gaussian measure with covariance
operator $C$ defined as

$C=48^{2}\bigl{(}-\Delta+12^{2}I\bigr{)}^{-2},$

where the Laplacian is equipped with zero Neumann boundary conditions and
viewed as acting between spaces of spatially mean-zero functions, and
$T_{1}(\cdot)=\exp(\cdot)$ is the exponential function. We hence refer to the
data inputs $a^{(j)}$ as being lognormal. In the second setting
$\\{a^{(j)}\\}_{j=1}^{J}$ are sampled from the probability measure

$\mu_{2}\coloneqq({T_{2}})_{\sharp}\mathcal{N}\bigl{(}0,C\bigr{)},$

where the covariance operator is defined as before and

$T_{2}(x)=\begin{cases}12,\quad x\geq 0,\\\ 3,\quad x<0.\end{cases}$

We hence refer to the data inputs $a^{(j)}$ as being piecewise constant.

Figure 11: The panel displays the result of the application of the Fourier
attention neural operator on the Darcy flow experiment with lognormal
diffusion for the median and maximum relative $L^{2}$ error samples. The first
row shows the sample from the test set of the same resolution as the training
set yielding the median relative $L^{2}$ error. On the other hand, the second
row displays the sample yielding the maximum relative $L^{2}$ error. The first
column of the panel displays the input diffusion coefficient $a(x)$ of the
relevant sample. The second column shows the true solution $u(x;a)$, while the
third column displays the the predicted solution. The last column displays the
pointwise absolute error (in log-scale) between the prediction and truth.

In the context of the lognormal experimental setting, we investigate a
scenario with low resolution training samples, for which pointwise attention
and hence the transformer neural operator from Subsection 4.2 is suitable, and
a high resolution scenario where patching becomes necessary for application of
the patched-based architectures of Subsections 4.3 and 4.4.

For the low resolution setting, we train the FNO and transformer neural
operator architectures using $3600$ independent samples of resolution
$64\times 64$, we use $200$ samples for validation and $200$ samples for
testing. The Fourier neural operator is parametrized by $12$ Fourier modes in
each dimension and channel width of $128$; the model is trained using a batch
size of $8$ and learning rate of $10^{-4}$. On the other hand, the transformer
neural operator is parametrized by $128$ channels in the encoder and is
trained using a batch size of $2$ and learning rate of $10^{-3}$. On this
problem we train both neural operators using the relative $H^{1}$ loss.

For the high resolution setting, we train the FNO, ViT neural operator and
Fourier attention neural operator architectures using $3600$ independent
samples of resolution $416\times 416$; we use $200$ samples for validation and
$200$ samples for testing. The Fourier neural operator is again parametrized
by $12$ Fourier modes in each dimension and channel width of $128$; the model
is trained using a batch size of $8$ and learning rate of $10^{-4}$. Both the
ViT neural operator and Fourier attention neural operator are parametrized by
$128$ channels in the transformer encoder. The integral operator in the ViT NO
is parametrized by $12$ Fourier modes in each dimension while the integral
operators appearing in the Fourier attention neural operator are parametrized
by $9$ Fourier modes in each dimension. Both architectures employ a final
spectral convolution layer (see Remark 21) that is parametrized by $64$
Fourier modes in each dimension. Both of these neural operators are trained
using a batch size of $1$ and learning rate of $10^{-3}$. On this problem we
train all the neural operators with the relative $H^{1}$ loss. Figure 11
displays the high resolution test results obtained by applying the Fourier
attention neural operator.

Figure 12: The panel displays the result of the application of the Fourier
attention neural operator on the Darcy flow experiment with piecewise constant
diffusion for the median and maximum relative $L^{2}$ error samples. The first
row displays the sample from the test set of the same resolution as the
training set yielding the median relative $L^{2}$ error. On the other hand,
the second row displays the sample yielding the maximum relative $L^{2}$
error. The first column of the panel displays the input diffusion coefficient
$a(x)$ of the relevant sample. The second column shows the true solution
$u(x;a)$, while the third column displays the the predicted solution. The last
column displays the pointwise absolute error (in log-scale) between the
prediction and truth. Figure 13: The panel displays the performance in
relative $L^{2}$ errors (in log-scale) of the various architectures when
applied to Darcy flow with lognormal diffusion test sets of different
resolutions at inference time (without retraining). The left panel concerns
FNO and the transformer neural operator when trained on data of resolution
$64\times 64$. On the other hand, the right panel concerns the FNO, ViT neural
operator and Fourier attention neural operator architectures trained on data
of resolution $416\times 416$.

In the context of the piecewise constant experimental setting, we train the
FNO, ViT neural operator and Fourier attention neural operator architectures
using $3600$ independent samples of resolution $416\times 416$, we use $200$
samples for validation and $200$ samples for testing. The Fourier neural
operator is again parametrized by $12$ Fourier modes in each dimension and
channel width of $128$; the model is trained using a batch size of $8$ and
learning rate of $10^{-4}$. Both the ViT neural operator and Fourier attention
neural operator are parametrized by $128$ channels in the transformer encoder.
The integral operator in the ViT NO is parametrized by $12$ Fourier modes in
each dimension while the integral operators appearing in the Fourier attention
neural operator are parametrized by $9$ Fourier modes in each dimension. Both
architectures employ a final spectral convolution layer (see Remark 21) that
is parametrized by $64$ Fourier modes in each dimension. Both of these neural
operators are trained using a batch size of $1$ and learning rate of
$10^{-3}$. On this problem we train all the neural operators with the relative
$H^{1}$ loss. Figure 12 displays the high resolution test results obtained by
applying the Fourier attention neural operator.

Figure 14: The panel displays the performance in relative $L^{2}$ errors (in
log-scale) of the FNO, ViT neural operator and Fourier attention neural
operator when applied to Darcy flow with piecewise constant diffusion test
sets of different resolutions at inference time (without retraining). All
architectures are trained on data that is of resolution $416\times 416$.

In Figures 13 and 14 we display the results of applying the neural operator
architectures for the lognormal and piecewise input settings, respectively, to
test samples of different resolutions than the one used for training. The
results demonstrate the zero-shot generalization to different resolutions
capability of the transformer neural operator architectures, which does not
require retraining of the models. In the lognormal setting the FNO exhibits
invariance to discretization that is more stable to changing resolution than
are the ViT neural operator and the Fourier attention neural operator.

####  6.2.3 Kolmogorov Flow

We consider the two-dimensional Navier-Stokes equation for a viscous,
incompressible fluid,

$\displaystyle\frac{\partial u}{\partial t}+u\cdot\nabla u+\nabla p$
$\displaystyle=\nu\Delta u+\sin(ny)\hat{x},$
$\displaystyle(x,t)\in[0,2\pi]^{2}\times(0,\infty),$ (KF)
$\displaystyle\nabla\cdot u$ $\displaystyle=0,$
$\displaystyle(x,t)\in[0,2\pi]^{2}\times[0,\infty),$ $\displaystyle
u(\cdot,0)$ $\displaystyle=u_{0},$ $\displaystyle x\in[0,2\pi]^{2},$

where $u$ denotes the velocity, $p$ the pressure and $\nu$ the kinematic
viscosity. The particular choice of forcing function $\sin(ny)$ leads to a
particular example of a Kolmogorov flow. We equip the domain $[0,2\pi]$ with
periodic boundary conditions. We assume

$u_{0}\in\mathcal{U}\coloneqq\Bigl{\\{}u\in\dot{L}^{2}_{\textrm{per}}\bigl{(}[0,2\pi]^{2};\mathbbm{R}^{2}\bigr{)}:\nabla\cdot
u=0\Bigr{\\}}.$

The dot on $L^{2}$ denotes the space of spatially mean-zero functions. The
vorticity is defined as $w=(\nabla\times u)\widehat{z}$ and the stream
function $f$ as the solution to the Poisson equation $-\Delta f=w$. Existence
of the semigroup $S_{t}:\mathcal{U}\to\mathcal{U}$ is shown in Temam (2012,
Theorem 2.1). We generate the data by solving (KF) in vorticity-streamfunction
form by applying the pseudo-spectral split step method from Chandler and
Kerswell (2013). In our experimental set-up $n=4$ and $\nu=1/70$ are chosen.
Random initial conditions are sampled from the Gaussian measure
$\mathcal{N}(0,C)$ where the covariance operator is given by

$C=7^{3}\bigl{(}-\Delta+49I\bigr{)}^{-5},$

where the Laplacian is equipped with periodic boundary conditions on
$[0,2\pi]^{2}$, and viewed as acting between spaces of spatially mean-zero
functions. The input output-data pairs for the Kolmogorov flow experiment are
given by

$\Bigl{\\{}\omega\bigl{(}x,T;\omega_{0}^{(j)}\bigr{)};\omega\bigl{(}x,T+\Delta
t;\omega_{0}^{(j)}\bigr{)}\Bigr{\\}}_{j=1}^{J},$

for $\omega_{0}^{(j)}\sim\mathcal{N}(0,C)$ being i.i.d. samples. Indeed, given
any $r>0$ we aim to approximate the nonlinear operator
$\mathsf{G}^{\dagger}:H^{r}\bigl{(}[0,2\pi]^{2};\mathbbm{R}\bigr{)}\to
H^{r}\bigl{(}[0,2\pi]^{2};\mathbbm{R}\bigr{)}$ defined by

$\mathsf{G}^{\dagger}:\omega\bigl{(}x,T;\omega_{0}\bigr{)}\mapsto\omega\bigl{(}x,T+\Delta
t;\omega_{0}\bigr{)}.$

Figure 15: The panel displays the result of the application of the Fourier
attention neural operator on the Kolmogorov flow experiment for the median and
maximum relative $L^{2}$ error samples. The first row displays the sample from
the test set of the same resolution as the training set yielding the median
relative $L^{2}$ error. The second row on the other hand displays the sample
yielding the maximum relative $L^{2}$ error. The first column of the panel
displays the input vorticity $w(x,T)$ of the relevant sample. The second
column shows the absolute difference between the true vorticity $w(x,T+\Delta
t)$ and the input vorticity $w(x,T)$, while the third column displays the
absolute difference between the predicted vorticity at time $T+\Delta t$ and
the input $w(x,T)$. The last column shows the pointwise absolute error (in
log-scale) between the prediction and truth.

In our experimental setting, we set $T=11$ and $\Delta t=0.1$. We train the
FNO, ViT neural operator and Fourier attention neural operator architectures
on this problem using $9000$ independent samples of resolution $416\times
416$, $500$ samples for validation and $500$ samples for testing. The Fourier
neural operator is parametrized by $12$ Fourier modes in each dimension and
channel width of $128$; the model is trained using a batch size of $4$ and
learning rate of $10^{-3}$. Both the ViT neural operator and Fourier attention
neural operator are parametrized by $128$ channels in the transformer encoder.
The integral operator in the ViT NO is parametrized by $12$ Fourier modes in
each dimension while the integral operators appearing in the Fourier attention
neural operator are parametrized by $9$ Fourier modes in each dimension. Both
architectures employ a final spectral convolution layer (Remark 21) that is
parametrized by $64$ Fourier modes in each dimension. Both of these neural
operators are trained using a batch size of $2$ and learning rate of
$10^{-3}$. On this problem we train all the neural operators with the relative
$H^{1}$ loss. Figure 15 displays the test results using the Fourier attention
neural operators.

Figure 16: The panel displays the performance in relative $L^{2}$ errors (in
log-scale) of the FNO, ViT neural operator and Fourier attention neural
operator when applied to Kolomgorov flow test sets of different resolutions at
inference time (without retraining). All architectures are trained on data
that is of resolution $416\times 416$.

In Figure 16 we display the results of applying the three architectures to
test samples of resolutions different to the $416\times 416$ training
resolution. The results demonstrate the zero-shot generalization capability of
the three neural operator architectures to different discretizations, which
does not require retraining of the models. Again we note that FNO exhibits
invariance to discretization that is more stable to changing resolution than
the methods introduced here; but for all the neural operators it is
nonetheless clear that intrinsic properties of the continuum limit are learned
and may be transferred between discretizations.

A few additional considerations are in order. The mode truncation employed in
the Fourier neural operator make the architecture computationally efficient
but yields an over-smoothing effect that is not suitable for problems at high
resolutions, where high frequency detail is prominent. Attention-based neural
operators offer a possible solution to this issue, as demonstrated by the
performance of the transformer neural operators proposed when applied to the
PDE operator learning problems considered, which involve rough diffusion
coefficients and rough initial conditions. On the other hand, patching leads
to more efficient attention-based neural operators, but also introduces
possible discontinuities at patch-intersections. This issue has been observed
to be resolved in the large data regime. The residual spectral convolution
layers employed ameliorate the discontinuities that are visible in pointwise
error plots, but have a negligible effect on relative $L^{2}$ error in
training and testing. This property constitutes an inherent limitation of
patching and highlights the need for more suitable efficient attention
mechanisms. Furthermore, it is observed that in some experimental settings the
patching-based neural operators exhibit worse stability to changing resolution
than FNO. This is likely due to the effect of patching which may then be
further amplified for FANO due to the application of a higher number of FFT
applications to non-periodic domains. Such issues could potentially be
overcome using periodic extensions for Fourier-based parametrizations or other
kinds of parametrizations for the nonlocal operations, or could be further
improved by fine-tuning.

##  7 Conclusions

In this work we have introduced a continuum formulation of the attention
mechanism from the seminal work Vaswani et al. (2017). Our continuum
formulation can be used to design transformer neural operators. Indeed, this
continuum formulation leads to discretization invariant implementations of
attention and schemes that exhibit zero-shot generalization to different
resolutions. In the first result of its kind for transformers, the resulting
neural operator architecture mapping between infinite-dimensional spaces of
functions is shown to be a universal approximator of continuous functions and
functions of Sobolev regularity defined over a compact domain. We extend the
continuum formulation to patched attention, which makes it possible to design
efficient transformer neural operators. To this end, we introduce a neural
operator analogue of the vision transformer Dosovitskiy et al. (2021) and a
more expressive architecture, the Fourier attention neural operator. Through a
cost-accuracy analysis, we demonstrate the power of the methodology in the
context of high-resolution operator learning problems.

Acknowledgments and Disclosure of Funding

AMS is supported by a Department of Defense Vannevar Bush Faculty Fellowship;
this funding also supports EC and partly supported MEL. In addition, AMS and
EC are supported by the SciAI Center, funded by the Office of Naval Research
(ONR), under Grant Number N00014-23-1-2729. NBK is grateful to the Nvidia
Corporation for support through full-time employment. MEL is supported by
funding from the Eric and Wendy Schmidt Center at the Broad Institute of MIT
and Harvard.

We thank Théo Bourdais, Matthieu Darcy, Miguel Liu-Schiaffini, Georgia
Gkioxari, Sabera Talukder, Zihui Wu, and Zongyi Li for helpful discussions and
feedback.

##  Appendix A Proofs of Approximation Theorems

###  A.1 Proof of Self-Attention Approximation Theorem

We introduce the following auxiliary result that we will apply to prove
Theorem 6.

###### Lemma 25

Let $D\subset\mathbbm{R}^{d}$ be a bounded open set and let
$\\{y_{j}\\}_{j=1}^{N}\sim U(\bar{D})$ be an i.i.d. sequence. If $f\in
C(\bar{D}\times\bar{D})$ then, with the expectation taken over the data
$\\{y_{j}\\}_{j=1}^{N}$,

$\mathbb{E}\left\|\int_{D}f(\cdot,y)\>\mathrm{d}y-\frac{|D|}{N}\sum_{j=1}^{N}f(\cdot,y_{j})\right\|_{L^{2}}\leq|D|^{3/2}\|f\|_{C(\bar{D}\times\bar{D})}N^{-1/2}.$

Proof  Notice that, for any $x\in\bar{D}$,

$\int_{D}f(x,y)\>\mathrm{d}y=|D|\mathbb{E}_{y\sim
U(\bar{D})}\big{[}f(x,y)\big{]}.$

Hence, by expanding the square, we can observe that

$\mathbb{E}\left(\int_{D}f(x,y)\>\mathrm{d}y-\frac{|D|}{N}\sum_{j=1}^{N}f(x,y_{j})\right)^{2}\leq\frac{|D|^{2}}{N}\mathbb{E}_{y\sim
U(\bar{D})}\big{[}f(x,y)^{2}\big{]}.$

Bounding $f$ in $L^{\infty}$, and noting that it is a continuous function,
gives

$\displaystyle\mathbb{E}\left\|\int_{D}f(\cdot,y)\>\mathrm{d}y-\frac{|D|}{N}\sum_{j=1}^{N}f(\cdot,y_{j})\right\|^{2}_{L^{2}}$
$\displaystyle\leq\frac{|D|^{2}}{N}\int_{D}\mathbb{E}_{y\sim
U(\bar{D})}\big{[}f(x,y)^{2}\big{]}\>\mathrm{d}x$
$\displaystyle\leq|D|^{3}\|f\|^{2}_{C(\bar{D}\times\bar{D})}N^{-1}.$

The result follows by Jensen’s inequality.
As part of the proof of Theorem 6, we first show that the the self-attention
operator is a mapping of the form
$\mathsf{A}:L^{\infty}(D;\mathbbm{R}^{d_{u}})\to
L^{\infty}(D;\mathbbm{R}^{d_{V}})$ and hence
$\mathsf{A}:C(\bar{D};\mathbbm{R}^{d_{u}})\to C(\bar{D};\mathbbm{R}^{d_{V}})$.

###### Lemma 26

Let $u\in L^{\infty}(D;\mathbbm{R}^{d_{u}})$, then it holds that
$\mathsf{A}(u)\in L^{\infty}(D;\mathbbm{R}^{d_{V}})$. Furthermore, for $u\in
C(\bar{D};\mathbbm{R}^{d_{u}})$ it holds that $\mathsf{A}(u)\in
C(\bar{D};\mathbbm{R}^{d_{V}})$.

Proof  The result follows readily from the fact that $p$ is a probability
density function. For completeness, we show the boundedness of the
normalization constant of the function $p$. Indeed we note that,

$\displaystyle\int_{D}\mathrm{exp}\big{(}\langle
Qu(x),Ku(s)\rangle\big{)}\>\mathrm{d}s$
$\displaystyle\leq\int_{D}\mathrm{exp}\big{(}|Q||K||u(x)||u(y)|\big{)}\mathrm{d}s$
$\displaystyle\leq\int_{D}\mathrm{exp}\bigg{(}\frac{1}{2}|Q|^{2}|K|^{2}|u(x)|^{2}\bigg{)}\mathrm{exp}\bigg{(}\frac{1}{2}|u(y)|^{2}\bigg{)}\mathrm{d}s$
$\displaystyle\leq|D|\mathrm{exp}\big{(}R\|u\|_{L^{\infty}}^{2}\big{)}$

for some constant $R>0$. From definition, it is clear that

$\int_{D}p(y;u,x)\>\mathrm{d}y=1,$

hence $p$ is indeed a valid probability density function. It hence follows
that

$\|A(u)\|_{L^{\infty}}\leq|V|\|u\|_{L^{\infty}}$

hence $A:L^{\infty}(D;\mathbbm{R}^{d_{u}})\to
L^{\infty}(D;\mathbbm{R}^{d_{V}})$. Similarly, it follows that
$A:C(\bar{D};\mathbbm{R}^{d_{u}})\to C(\bar{D};\mathbbm{R}^{d_{V}})$. We note
that the continuity of $u$ is preserved by $\mathsf{A}$ under the continuity
of the inner product in its first argument and the continuity of the
exponential.
We are now ready to establish the full result of Theorem 6 which states that

$\lim_{N\to\infty}\sup_{u\in
B}\mathbb{E}\left\|\mathsf{A}(u)-\frac{\sum_{j=1}^{N}\mathrm{exp}\big{(}\langle
Qu(\cdot),Ku(y_{j})\rangle\big{)}Vu(y_{j})}{\sum_{\ell=1}^{N}\mathrm{exp}\big{(}\langle
Qu(\cdot),Ku(y_{\ell})\rangle\big{)}}\right\|_{C(\bar{D};\mathbbm{R}^{d_{V}})}=0,$

with the expectation taken over i.i.d. sequences
$\\{y_{j}\\}_{j=1}^{N}\sim\mathsf{unif}(\bar{D})$.

Proof [Proof of Theorem 6] Before commencing the proof, we first establish
useful shorthand notation. Letting $u\in B$ we define,

$f(x,y;u)\coloneqq\mathrm{exp}\big{(}\langle Qu(x),Ku(y)\rangle\big{)},\quad
g_{l}(x,y;u)\coloneqq\mathrm{exp}\big{(}\langle
Qu(x),Ku(y)\rangle\big{)}\big{(}Vu(y)\big{)}_{l}$

for all $x,y\in\bar{D}$ and $l\in\llbracket 1,d_{V}\rrbracket$. For $u\in B$
we also define

$a(x;u)\coloneqq\int_{D}f(x,y;u)\>\mathrm{d}y,\qquad
a^{(N)}(x;u)\coloneqq\frac{|D|}{N}\sum_{j=1}^{N}f(x,y_{j};u),$

and

$b_{l}(x;u)\coloneqq\int_{D}g_{l}(x,y;u)\>\mathrm{d}y,\qquad
b^{(N)}_{l}(x;u)\coloneqq\frac{|D|}{N}\sum_{j=1}^{N}g_{l}(x,y_{j};u)$

for any $x\in\bar{D}$ and any $l\in\llbracket 1,d_{V}\rrbracket$. We set
$\alpha_{l}=b_{l}/a$ and $\alpha_{l}^{(N)}=b_{l}^{(N)}/a^{(N)}$. Given this
notation, to prove the result we must show

$\lim_{N\to\infty}\sup_{u\in
B}\mathbb{E}\left\|\alpha_{l}(\cdot\,;u)-\alpha_{l}^{(N)}(\cdot\,;u)\right\|_{C(\bar{D};\mathbbm{R})}=0,$
(43)

for any $l\in\llbracket 1,d_{V}\rrbracket$. We divide the proof in the
following key steps. We first show that for fixed $u\in B$, it holds that

$\lim_{N\to\infty}\mathbb{E}\|\alpha_{l}(\cdot\,;u)-\alpha_{l}^{(N)}(\cdot\,;u)\|_{L^{2}}=0,$
(44)

for any $l\in\llbracket 1,d_{V}\rrbracket$. By using compactness and applying
the Arzelà-Ascoli theorem, we then show that

$\lim_{N\to\infty}\mathbb{E}\|\alpha_{l}(\cdot\,;u)-\alpha_{l}^{(N)}(\cdot\,;u)\|_{C(\bar{D})}=0,$
(45)

for any $l\in\llbracket 1,d_{V}\rrbracket$. We then use the compactness of $B$
and continuity of the mappings $u\mapsto\alpha_{\ell}(\cdot\,;u)$ and
$u\mapsto\alpha^{(N)}_{\ell}(\cdot\,;u)$ to deduce the result given by (43).
We now focus on establishing (44). Since $B$ is bounded, there exists a
constant $M>0$ such that,

$\sup_{u\in B}\|u\|_{L^{\infty}}\leq M.$

Therefore, we have

$\sup_{u\in
B}\max\\{\|f\|,\|g_{1}\|,\dots,\|g_{m}\|\\}\leq\max\left\\{\text{exp}\big{(}RM^{2}\big{)},M|V|\text{exp}\big{(}RM^{2}\big{)}\right\\}\coloneqq
J,$

where $\|\cdot\|=\|\cdot\|_{C(\bar{D}\times\bar{D})}$. Clearly,

$\sup_{u\in B}\max\\{\|a\|_{L^{2}},\|a^{(N)}\|_{L^{2}}\\}\leq|D|^{3/2}J$

Notice that, since $|\langle Qu(x),Ku(s)\rangle|\leq RM^{2}$, we have
$f(x,y)\geq\text{exp}\big{(}-RM^{2}\big{)}$. It follows that

$\inf_{u\in
B}\min\\{\|a\|_{L^{2}},\|a^{(N)}\|_{L^{2}}\\}\geq|D|^{3/2}\text{exp}\big{(}-RM^{2}\big{)}\coloneqq|D|^{3/2}I>0.$

Similarly,

$\sup_{u\in B}\max_{l\in\llbracket
1,d_{V}\rrbracket}\\{\|b_{l}\|_{L^{2}},\|b_{l}^{(N)}\|_{L^{2}}\\}\leq|D|^{3/2}J.$

Applying Lemma 25, for fixed $u\in B$ we find that

$\displaystyle\mathbb{E}\left\|\frac{b_{l}}{a}-\frac{b_{l}^{(N)}}{a^{(N)}}\right\|_{L^{2}}$
$\displaystyle\leq\mathbb{E}\frac{\|a^{(N)}\|_{L^{2}}\|b_{l}-b_{l}^{(N)}\|_{L^{2}}+\|b_{l}^{(N)}\|_{L^{2}}\|a-a^{(N)}\|_{L^{2}}}{\|a^{(N)}\|_{L^{2}}\|a\|_{L^{2}}}$
(46a) $\displaystyle\leq\frac{2J^{2}}{I^{2}}N^{-1/2}.$ (46b)

Setting $\alpha_{l}=b_{l}/a$ and $\alpha_{l}^{(N)}=b_{l}^{(N)}/a^{(N)}$, from
(46a) we deduce that

$\lim_{N\to\infty}\mathbb{E}\|\alpha_{l}(\cdot\,;u)-\alpha_{l}^{(N)}(\cdot\,;u)\|_{L^{2}}=0,$
(47)

for any $l\in\llbracket 1,d_{V}\rrbracket$. We now proceed to the second step
of the proof, i.e. showing that (45) holds. Using similar reasoning as before,
we notice that

$\displaystyle\inf_{u\in B}$
$\displaystyle\min\\{\|a\|_{C(\bar{D})},\|a^{(N)}\|_{C(\bar{D})}\\}\geq|D|I,$
(48a) $\displaystyle\sup_{u\in B}$ $\displaystyle\max_{l\in\llbracket
1,d_{V}\rrbracket}\\{\|b_{l}\|_{C(\bar{D})},\|b_{l}^{(N)}\|_{C(\bar{D})}\\}\leq|D|J,$
(48b)

hence the sequence $\\{\alpha_{l}^{(N)}\\}$ is uniformly bounded in $N$. Now,
for fixed $u\in B$ the sequence $\\{\alpha_{l}^{(N)}(\cdot\,;u)\\}$ is also
uniformly equicontinuous with probability 1 over the choice
$\\{y_{j}\\}_{j=1}^{N}\sim U(\bar{D})$. Indeed, we note that

$\displaystyle\left|\frac{b_{l}^{(N)}(r)}{a^{(N)}(r)}-\frac{b_{l}^{(N)}(t)}{a^{(N)}(t)}\right|$
$\displaystyle=\left|\frac{\sum_{j=1}^{N}g_{l}(r,y_{j})}{\sum_{k=1}^{N}f(r,y_{k})}-\frac{\sum_{j=1}^{N}g_{l}(t,y_{j})}{\sum_{k=1}^{N}f(t,y_{k})}\right|$
(49a)
$\displaystyle\leq\frac{1}{N^{2}I}\left|\sum_{j,n=1}^{N}g_{l}(r,y_{j})f(t,y_{n})-g_{l}(t,y_{j})f(r,y_{n})\right|$
(49b)
$\displaystyle\leq\frac{J}{N^{2}I}\sum_{j,n=1}^{N}\bigl{(}\bigl{|}g_{l}(r,y_{j})-g_{l}(t,y_{j})\bigr{|}+\bigl{|}f(t,y_{n})-f(r,y_{n})\bigr{|}\bigr{)}.$
(49c)

The result then follows from the uniform continuity of $f$ and $g_{l}$ in
their first arguments, due to the compactness of $\bar{D}$. Therefore, since
$\bar{D}\subset\mathbbm{R}^{d}$ is compact, by the Arzelà–Ascoli theorem there
exists a subsequence of indices $(N_{k})_{k\in\mathbbm{N}}$ such that
$\alpha_{l}^{(N_{k})}$ converges in $C(\bar{D})$ to some
$\alpha_{l}^{(\infty)}$. Since,

$\mathbb{E}\|\alpha_{l}^{(\infty)}-\alpha_{l}^{(N_{k})}\|_{L^{2}}\leq|D|^{1/2}\mathbb{E}\|\alpha_{l}^{(\infty)}-\alpha_{l}^{(N_{k})}\|_{C(\bar{D})},$

by uniqueness of limits, it is readily observed that
$\alpha_{l}^{(\infty)}=\alpha_{l}$. Therefore, since the limit is independent
of the subsequence, it holds that

$\lim_{N\to\infty}\mathbb{E}\|\alpha_{l}(\cdot\,;u)-\alpha_{l}^{(N)}(\cdot\,;u)\|_{C(\bar{D})}=0,$

for any $l\in\llbracket 1,d_{V}\rrbracket$. We now turn our attention to
establishing (43). By reasoning as before and by using the continuity of the
exponential and the inner product in its first argument, it is straightforward
to show that as mappings of the form $\alpha_{l}:B\to
C(\bar{D};\mathbbm{R}^{d_{V}})$ and $\alpha_{l}^{(N)}:B\to
C(\bar{D};\mathbbm{R}^{d_{V}})$ where $B$ is compact, $\alpha_{l}$ and
$\alpha_{l}^{(N)}$ are continuous and hence the sequence
$\\{\alpha_{l}^{(N)}\\}$ is uniformly equicontinuous. Therefore, we can find
moduli of continuity $\omega_{1},\omega_{2}$ such that

$\|\alpha_{l}(u)-\alpha_{l}(v)\|_{C(\bar{D})}\leq\omega_{1}\big{(}\|u-v\|_{C(\bar{D})}\big{)},\quad\|\alpha_{l}^{(N)}(u)-\alpha_{l}^{(N)}(v)\|_{C(\bar{D})}\leq\omega_{2}\big{(}\|u-v\|_{C(\bar{D})}\big{)}$

for any $u,v\in B$ and $N\in\mathbb{N}$. Fix $\epsilon>0$. Since $B$ is
compact, we can find a number $L=L(\epsilon)\in\mathbb{N}$ and functions
$\phi_{1},\dots,\phi_{L}\in B$ such that, for any $u\in B$, there exists
$j\in\llbracket 1,L\rrbracket$ such that

$\|u-\phi_{j}\|_{C(\bar{D})}<\epsilon.$

Furthermore, we can find a number $S=S(\epsilon)\in\mathbb{N}$ such that, for
any $j\in\llbracket 1,L\rrbracket$, we have, for all $N\geq S$,

$\mathbb{E}\|\alpha_{l}(\phi_{j})-\alpha_{l}^{(N)}(\phi_{j})\|_{C(\bar{D})}<\epsilon.$

It follows by triangle inequality that

$\displaystyle\mathbb{E}\|\alpha_{l}(u)-\alpha_{l}^{(N)}(u)\|_{C(\bar{D})}$
$\displaystyle\leq\mathbb{E}\|\alpha_{l}(u)-\alpha_{l}(\phi_{j})\|_{C(\bar{D})}+\mathbb{E}\|\alpha_{l}(\phi_{j})-\alpha_{l}^{(N)}(u)\|_{C(\bar{D})}$
$\displaystyle\leq|D|\omega_{1}(\epsilon)+\mathbb{E}\|\alpha_{l}(\phi_{j})-\alpha_{l}^{(N)}(\phi_{j})\|_{C(\bar{D})}+\mathbb{E}\|\alpha_{l}^{(N)}(\phi_{j})-\alpha_{l}^{(N)}(u)\|_{C(\bar{D})}$
$\displaystyle\leq|D|\omega_{1}(\epsilon)+\epsilon+|D|\omega_{2}(\epsilon).$

Therefore the result follows since the argument is uniform over all
$l\in\llbracket 1,d_{V}\rrbracket$.

###  A.2 Proof of Cross-Attention Approximation Theorem

We make straightforward modifications to Lemmas 25, 28 and to the proof of
Theorem 6 in order to establish Theorem 12.

###### Lemma 27

Let $D\subset\mathbbm{R}^{d}$ and $E\subset\mathbbm{R}^{e}$ be bounded open
sets and $f\in C(\bar{D}\times\bar{E})$. Then

$\mathbb{E}\left\|\int_{E}f(\cdot,y)\>\mathrm{d}y-\frac{|E|}{N}\sum_{j=1}^{N}f(\cdot,y_{j})\right\|_{L^{2}}\leq|D|^{1/2}|E|\|f\|_{C(\bar{D}\times\bar{E})}N^{-1/2}$

with the expectation taken over i.i.d. sequences $\\{y_{j}\\}_{j=1}^{N}\sim
U(\bar{E})$.

Proof  Notice that, for any $x\in\bar{D}$,

$\int_{E}f(x,y)\>\mathrm{d}y=|E|\mathbb{E}_{y\sim
U(\bar{E})}\big{[}f(x,y)\big{]}.$

Hence, by expanding the square, we can observe that

$\mathbb{E}\left(\int_{E}f(x,y)\>\mathrm{d}y-\frac{|E|}{N}\sum_{j=1}^{N}f(x,y_{j})\right)^{2}\leq\frac{|E|^{2}}{N}\mathbb{E}_{y\sim
U(\bar{E})}\big{[}f(x,y)^{2}\big{]}.$

Bounding $f$ in $L^{\infty}$, and noting that it is a continuous function,
gives

$\displaystyle\mathbb{E}\left\|\int_{E}f(\cdot,y)\>\mathrm{d}y-\frac{|E|}{N}\sum_{j=1}^{N}f(\cdot,y_{j})\right\|^{2}_{L^{2}}$
$\displaystyle\leq\frac{|E|^{2}}{N}\int_{D}\mathbb{E}_{y\sim
U(\bar{E})}\big{[}f(x,y)^{2}\big{]}\>\mathrm{d}x$
$\displaystyle\leq|D||E|^{2}\|f\|^{2}_{C(\bar{D}\times\bar{E})}N^{-1}.$

The result follows by Jensen’s inequality.
As part of the proof of Theorem 12, we first show that the the cross-attention
operator is a mapping of the form
$\mathsf{C}:L^{\infty}(D;\mathbbm{R}^{d_{u}})\times
L^{\infty}(E;\mathbbm{R}^{d_{v}})\to L^{\infty}(D;\mathbbm{R}^{d_{V}})$ and
hence $\mathsf{A}:C(\bar{D};\mathbbm{R}^{d_{u}})\times
C(\bar{E};\mathbbm{R}^{d_{v}})\to C(\bar{D};\mathbbm{R}^{d_{V}})$.

###### Lemma 28

Let $u\in L^{\infty}(D;\mathbbm{R}^{d_{u}})\times
L^{\infty}(E;\mathbbm{R}^{d_{v}})$, then it holds that $\mathsf{A}(u)\in
L^{\infty}(D;\mathbbm{R}^{d_{V}})$. Furthermore, for $u\in
C(\bar{D};\mathbbm{R}^{d_{u}})\times C(\bar{E};\mathbbm{R}^{d_{v}})$ it holds
that $\mathsf{A}(u)\in C(\bar{D};\mathbbm{R}^{d_{V}})$.

Proof  The result follows readily from the fact that $q$ is a probability
density function. For completeness, we show the boundedness of the
normalization constant of the function $q$. Indeed we note that,

$\displaystyle\int_{E}\mathrm{exp}\big{(}\langle
Qu(x),Kv(s)\rangle\big{)}\>\mathrm{d}s$
$\displaystyle\leq\int_{E}\mathrm{exp}\big{(}|Q||K||u(x)||v(s)|\big{)}\mathrm{d}s$
$\displaystyle\leq\int_{E}\mathrm{exp}\bigg{(}\frac{1}{2}|Q|^{2}|K|^{2}|u(x)|^{2}\bigg{)}\mathrm{exp}\bigg{(}\frac{1}{2}|v(s)|^{2}\bigg{)}\mathrm{d}s$
$\displaystyle\leq|E|\mathrm{exp}\big{(}R_{1}\|u\|_{L^{\infty}}^{2}\big{)}\mathrm{exp}\big{(}R_{2}\|v\|_{L^{\infty}}^{2}\big{)}$
$\displaystyle\leq|E|\mathrm{exp}\Big{(}R\bigl{(}\|u\|_{L^{\infty}}^{2}+\|v\|_{L^{\infty}}^{2}\bigr{)}\Big{)},$

for some constant $R\coloneqq\max\\{R_{1},R_{2}\\}$. From definition, it is
clear that

$\int_{E}q(y;u,v,x)\>\mathrm{d}y=1,$

hence $q$ is indeed a valid probability density function. It hence follows
that

$\|\mathsf{C}(u,v)\|_{L^{\infty}(D;\mathbbm{R}^{d_{V}})}\leq|V|\|v\|_{L^{\infty}(E;\mathbbm{R}^{d_{v}})}$

hence $\mathsf{C}:L^{\infty}(D;\mathbbm{R}^{d_{u}})\times
L^{\infty}(E;\mathbbm{R}^{d_{v}})\to L^{\infty}(D;\mathbbm{R}^{d_{V}})$.
Similarly, it follows that $\mathsf{C}:C(\bar{D};\mathbbm{R}^{d_{u}})\times
C(\bar{E};\mathbbm{R}^{d_{v}})\to C(\bar{D};\mathbbm{R}^{d_{V}})$. We note
that the continuity is preserved by $\mathsf{C}$ under the continuity of the
inner product in its first argument and the continuity of the exponential.
We are now ready to establish the full result of Theorem 12 which states that

$\lim_{N\to\infty}\sup_{(u,v)\in
B}\mathbb{E}\left\|\mathsf{C}(u,v)-\frac{\sum_{j=1}^{N}\mathrm{exp}\big{(}\langle
Qu(\cdot),Kv(y_{j})\rangle\big{)}Vv(y_{j})}{\sum_{\ell=1}^{N}\mathrm{exp}\big{(}\langle
Qu(\cdot),Kv(y_{\ell})\rangle\big{)}}\right\|_{C(\bar{D};\mathbbm{R}^{d_{V}})}=0,$

with the expectation taken over i.i.d. sequences
$\\{y_{j}\\}_{j=1}^{N}\sim\mathsf{unif}(\bar{E})$.

Proof [Proof of Theorem 12] Before commencing the proof, we first establish
useful shorthand notation. Letting $(u,v)\in B$ we define,

$f(x,y;u,v)\coloneqq\mathrm{exp}\big{(}\langle Qu(x),Kv(y)\rangle\big{)},\quad
g_{l}(x,y;u,v)\coloneqq\mathrm{exp}\big{(}\langle
Qu(x),Kv(y)\rangle\big{)}\big{(}Vv(y)\big{)}_{l}$

for all $x\in\bar{D},y\in\bar{E}$ and $l\in\llbracket 1,d_{V}\rrbracket$. For
$(u,v)\in B$ we also define

$a(x;u,v)\coloneqq\int_{E}f(x,y;u,v)\>\mathrm{d}y,\qquad
a^{(N)}(x;u,v)\coloneqq\frac{|E|}{N}\sum_{j=1}^{N}f(x,y_{j};u,v),$

and

$b_{l}(x;u,v)\coloneqq\int_{E}g_{l}(x,y;u,v)\>\mathrm{d}y,\qquad
b^{(N)}_{l}(x;u,v)\coloneqq\frac{|E|}{N}\sum_{j=1}^{N}g_{l}(x,y_{j};u,v)$

for any $x\in\bar{D}$ and any $l\in\llbracket 1,d_{V}\rrbracket$. We set
$\alpha_{l}=b_{l}/a$ and $\alpha_{l}^{(N)}=b_{l}^{(N)}/a^{(N)}$. Given this
notation, to prove the result we must show

$\lim_{N\to\infty}\sup_{(u,v)\in
B}\mathbb{E}\left\|\alpha_{l}(\cdot\,;u,v)-\alpha_{l}^{(N)}(\cdot\,;u,v)\right\|_{C(\bar{D};\mathbbm{R})}=0,$
(50)