# An Analysis of Attention via the Lens of Exchangeability and Latent Variable
Models

Yufeng Zhang , Boyi Liu11footnotemark: 1 , Qi Cai, Lingxiao Wang, Zhaoran Wang
equal contributionNorthwestern University;
<EMAIL_ADDRESS>University;
<EMAIL_ADDRESS>University;
<EMAIL_ADDRESS>University;
<EMAIL_ADDRESS>University;
<EMAIL_ADDRESS>
###### Abstract

With the attention mechanism, transformers achieve significant empirical
successes in natural language processing and computer vision. Despite the
intuitive understanding that transformers perform relational inference (or
“inductive reasoning”) over long sequences to produce desirable
representations, we lack a rigorous theory on how the attention mechanism
achieves it. In particular, several intriguing questions remain open: (a) What
makes a desirable representation? (b) How does the attention mechanism infer
the desirable representation within the forward pass? (c) How does a
pretraining procedure learn to infer the desirable representation through the
backward pass?

We aim to answer the three questions via the lens of exchangeability.
Specifically, we observe that, as is the case in BERT and ViT, input tokens
are often exchangeable since they already include positional encodings. The
notion of exchangeability induces a latent variable model that is invariant to
input sizes, which enables our theoretical analysis.

  * -

To answer (a) on representation, we establish the existence of a sufficient
and minimal representation of input tokens. In particular, such a
representation instantiates the posterior distribution of the latent variable
(or “concept”) given input tokens, which plays a central role in predicting
output labels and solving downstream tasks.

  * -

To answer (b) on inference, we prove that attention with the desired parameter
infers the latent posterior up to an approximation error, which is decreasing
in input sizes. In detail, we quantify how attention approximates the
conditional mean of the value given the key, which characterizes how it
performs relational inference over long sequences.

  * -

To answer (c) on learning, we prove that both supervised and self-supervised
objectives allow empirical risk minimization to learn the desired parameter up
to a generalization error, which is independent of input sizes. Particularly,
in the self-supervised setting, we identify a condition number that is pivotal
to solving downstream tasks.

Our theoretical analysis gives a complete characterization of the attention
mechanism as a “greybox” design, which unifies the handcrafted architecture
induced by the latent variable model (“whitebox”) and the learnable parameter
estimated from data (“blackbox”) with provable approximation, generalization,
and optimization guarantees.

##  1 Introduction

Transformers are the state-of-the-art architecture for a variety of tasks in
natural language processing (Vaswani et al., 2017), computer vision
(Dosovitskiy et al., 2020), and multimodal generation (Ramesh et al., 2021).
At the core of their significant empirical successes is the attention
mechanism, which is defined by a computation graph for the forward pass. In
particular, the computation graph performs a specific form of message passing
across input tokens (Bronstein et al., 2021). It is commonly believed that the
attention mechanism is capable of handling long sequences and performing
relational inference (or “inductive reasoning”), which appears to be the key
advantage of transformers. However, the intuitive understanding lacks a
quantitative justification, which leaves many intriguing questions open:

  * (a)

What makes a desirable representation? Ideally, the desirable representation
of input tokens is sufficient and minimal in the sense that it preserves all
relevant information for predicting output labels or solving downstream tasks
(sufficiency) while it neglects all irreverent information (minimality).
However, we lack a quantitative definition of sufficiency and minimality,
which requires a probabilistic model.

  * (b)

How does the attention mechanism infer the desirable representation within the
forward pass? Intuitively, the attention mechanism is defined by a computation
graph that resembles kernel smoothing or kernel regression (Shawe-Taylor et
al., 2004) for predicting the value given the key. However, we lack a formal
characterization of what function class the attention mechanism parameterizes
or approximates. Also, it remains unclear why the specific form of message
passing produces the desirable representation of input tokens, that is, one
with sufficiency and minimality.

  * (c)

How does a pretraining procedure learn the desirable representation through
the backward pass? Empirically, the pretraining procedure that minimizes
empirical risks for predicting masked tokens (Devlin et al., 2018; Dosovitskiy
et al., 2020; He et al., 2022) appears to succeed in the presence of long
sequences. However, we lack a theoretical justification of whether the
pretraining procedure with the masked objective attains a desirable estimator
that generalizes and why the generalization error does not appear to degrade
for long sequences. In particular, it remains unclear to what degree the
estimated representation facilitates solving downstream tasks.

In this paper, we answer the three questions via the lens of exchangeability.
The key observation is that, as is the case in BERT (Devlin et al., 2018) and
ViT (Dosovitskiy et al., 2020), input tokens are exchangeable since they
include positional encodings. In other words, the joint distribution of input
tokens, e.g., vector embeddings of words in a paragraph or patches in an image
with positional encodings, remains the same upon permuting their orders.
Meanwhile, the attention mechanism and entrywise feedforward neural networks
preserve the notion of exchangeability throughout all transformer layers. By
the de Finetti Theorem (de Finetti, 1937), the notion of exchangeability
induces a latent variable model that is invariant to input sizes. Unlike
classical Bayesian settings, where the latent variable model is defined across
many data points, ours is defined over input tokens within one data point (in
an “in-context” manner), which captures a fine-grained structure of
interactions as relational inductive biases (Battaglia et al., 2018). The
latent variable model enables our theoretical analysis, which is summarized in
the following:

  * -

To answer (a) on representation, we establish the existence of a sufficient
and minimal representation of input tokens based on the latent variable model,
which is induced by exchangeability. In particular, we leverage the latent
variable model to define sufficiency and minimality following the
factorization theorem and the sufficiency principle (Fisher, 1922). Moreover,
we prove that the posterior distribution of the latent variable given input
tokens is a sufficient and minimal representation, which plays a central role
in predicting output labels and solving downstream tasks. Intuitively, the
latent variable instantiates the “concept” of a paragraph or an image, which
is “summarized” over words or patches. In detail, the “summarization” process
is formalized by the mapping from input tokens to the posterior distribution
of the latent variable, that is, inferring the “concept” in a Bayesian manner
within the forward pass.

Given the answer to (a), which defines the desirable representation as the
latent posterior, it remains unclear how to parameterize or approximate the
latent posterior, which is addressed by the answer to (b).

  * -

To answer (b) on inference, we prove that the attention mechanism with the
desired parameter infers the latent posterior up to an approximation error,
which is decreasing in input sizes. In particular, we prove that a specific
parameterization of the latent posterior yields a variant of the attention
mechanism based on kernel conditional mean embedding (CME), namely the CME
attention, which infers the conditional mean of the value given the key. Here
the value and the key (or the query) are obtained from a parameterized
transformation of input tokens, where the unknown parameter requires learning.

Although the CME attention recovers the latent posterior for any input sizes,
it differs from the commonly used softmax attention by a normalization matrix.
To this end, we prove that the CME attention and the softmax attention are
equivalent at the infinite limit of input sizes by drawing a connection to
nonparametric conditional density estimation. In other words, the softmax
attention recovers the latent posterior up to an approximation error that is
decreasing in input sizes, which characterizes how it performs relational
inference over long sequences. As byproducts, we justify the necessity of
multiple attention heads in transformers and provide a causal interpretation
of the inferred representation through instrumental variables.

Given the answer to (b), which quantifies the approximation error for the
latent posterior, it remains unclear how to learn the desired parameter of the
attention mechanism, which is addressed by the answer to (c).

  * -

To answer (c) on learning, we prove that both supervised and self-supervised
objectives allow empirical risk minimization to learn the desired parameter up
to a generalization error, which is independent of input sizes. In particular,
through maximum likelihood estimation, we establish the connection between the
latent posterior and the masked objective, which is defined by the empirical
risk for predicting masked tokens.

Moreover, we prove that the global minimizer of the masked objective attains a
generalization error that is independent of input sizes, which justifies why
transformers allow long sequences. Our proof exploits the invariance and
equivariance of the attention mechanism and entrywise feedforward neural
networks, which deviates from most existing analyses of the generalization
error. Particularly, in the self-supervised setting, e.g., as in MAE (He et
al., 2022), we identify a condition number that is pivotal to solving
downstream tasks. Intuitively, the condition number quantifies the amount of
information that is transferred from the pretraining task to a new task.

Meanwhile, in the overparameterized regime, we prove that any stationary point
of the masked objective is almost globally optimal when the attention
mechanism and entrywise feedforward neural networks have sufficient expressive
power. As a result, stochastic gradient descent finds the global minimizer of
the masked objective, which generalizes as discussed above.

Combining the above analysis of the approximation, generalization, and
optimization errors in the answer to (a)-(c), we provide a complete
characterization of the attention mechanism.

Contribution: In summary, our theoretical contribution is threefold:

  * (i)

We identify a general principle for parameterizing function classes and
constructing learning objectives based on latent posterior inference, which
requires a minimal assumption of data. In contrast to classical learning
paradigms, the latent variable model is defined over input tokens within one
data point, which captures relational inductive biases.

  * (ii)

We recover the attention mechanism from a specific parameterization of latent
posterior inference based on kernel conditional mean embedding and
nonparametric conditional density estimation. In particular, we demonstrate
how the attention mechanism combines the handcrafted architecture, which is
induced by latent posterior inference, and the learnable parameter, which
determines the kernel function.

  * (iii)

We characterize the approximation, generalization, and optimization errors for
estimating the learnable parameter of the attention mechanism through
minimizing the masked objective. In particular, we prove that input sizes do
not degrade the approximation and generalization errors, which justifies why
transformers allow long sequences.

Discussion: Our theoretical analysis casts the attention mechanism as a
“greybox” approach to modeling, that is, it combines the handcrafted
architecture, which is coined by a probabilistic model over input tokens
within one data point (“whitebox”), and the learnable parameter, which is
estimated in an end-to-end manner through empirical risk minimization
(“blackbox”). It is worth mentioning that our theoretical analysis studies the
class of transformers like BERT (Devlin et al., 2018) and ViT (Dosovitskiy et
al., 2020) (“encoder-only”), which does not exploit the autoregressive
structure as in GPT (Brown et al., 2020) (“decoder-only”). On the other hand,
the general principle identified in (i) is applicable to other probabilistic
models like hidden Markov models or general graphical models over trees and
grids, which motivates other principled architectures beyond the attention
mechanism. We leave it as a future direction.

### Related Works

Transformers and Attention. The pioneering work (Vaswani et al., 2017)
proposes transformers for the first time and highlights the key role of the
attention mechanism. Subsequently, there are a vast body of works that propose
various transformer architectures and different pretraining paradigms. See,
e.g., Devlin et al. (2018); Radford et al. (2018, 2019); Dai et al. (2019);
Brown et al. (2020); Dosovitskiy et al. (2020); He et al. (2022) and the
references therein. Transformers demonstrate significant empirical successes
in natural language processing (Wolf et al., 2020), computer vision
(Dosovitskiy et al., 2020), protein structure prediction (Jumper et al.,
2021), and sequential decision making (Chen et al., 2021). Our work provides a
theoretical justification of transformers and the attention mechanism, that
is, how a latent variable model induced by exchangeability allows us to derive
transformer architectures and pretraining paradigms in a principled manner.

Analysis of Transformers and Attention. Our work is related to a recent line
of works that analyze transformers and the attention mechanism (Tsai et al.,
2019; Vuckovic et al., 2020; Hron et al., 2020; Yang, 2020; Yang and Littwin,
2021; Edelman et al., 2021; Wei et al., 2021; Xie et al., 2021; Malladi et
al., 2022; Garg et al., 2022; Zhang et al., 2022b). Specifically, Tsai et al.
(2019) demonstrate that the attention mechanism can be viewed as a kernel
smoother over input tokens. Vuckovic et al. (2020) establish the Lipschitz
continuity of transformers via the lens of interacting particle systems. Hron
et al. (2020); Yang (2020); Yang and Littwin (2021); Malladi et al. (2022)
characterize the infinite-width limit of transformers under the framework of
neural tangent kernels (Jacot et al., 2018). Among them, Malladi et al. (2022)
demonstrate that neural tangent kernels can capture the parameter update in
the fine-tuning phase. Edelman et al. (2021) prove that transformers can
represent a sparse function of input tokens and establish a sample complexity
that scales logarithmically in input sizes. Wei et al. (2021) characterize the
approximation and generalization errors for learning a Turing machine with
transformers. Xie et al. (2021) prove that transformers can infer a latent
variable (or “concept”) assuming that the data distribution is a mixture of
hidden Markov models. Garg et al. (2022) demonstrate that transformers can
learn to perform linear predictions within one data point (in an “in-context”
manner). Zhang et al. (2022b) evaluate the empirical performance of
transformers for learning equality and group operations.

Our work provides a complete characterization of the representation,
inference, and learning aspects of the attention mechanism via the lens of
exchangeability and latent variable models, which requires a minimal
assumption on the data distribution (exchangeability). Specifically, in
comparison with Tsai et al. (2019), we demonstrate that the attention
mechanism not only parameterizes nonparametric conditional density estimation
but also approximates kernel conditional mean embedding, which infers the
conditional mean of the value given the key. Moreover, we invoke the latent
variable model induced by exchangeability to justify the attention mechanism
as a specific parameterization of latent posterior inference. Meanwhile, we
leverage the latent variable model to derive the common choice of both
supervised and self-supervised objectives, e.g., the masked objective. In
comparison with Xie et al. (2021), we do not assume that the data distribution
takes a specific form (a mixture of hidden Markov models) or the latent
posterior is given a priori by a specific parameter (no learning required).
Instead, we prove that the attention mechanism is capable of instantiating
latent posterior inference up to an approximation error and the masked
objective allows us to learn to infer the latent posterior up to the
generalization and optimization errors. Also, it is worth mentioning that Xie
et al. (2021) focus on the class of transformers like GPT (Brown et al., 2020)
(“decoder-only”), while we focus on the class of transformers like BERT
(Devlin et al., 2018) and ViT (Dosovitskiy et al., 2020) (“encoder-only”). In
comparison to Edelman et al. (2021), we exploit the invariance and
equivariance of transformers and establish a generalization error that is
independent of input sizes.

Generalization of Deep Neural Networks. Our work is related to the vast body
of works that analyze the generalization error of deep neural networks. See,
e.g., Jiang et al. (2019); Valle-Pérez and Louis (2020) for a comprehensive
introduction. However, most of them do not exploit invariance and
equivariance. As a result, a direct application of such results yields a
vacuous bound as input sizes increase. On the other hand, Sokolic et al.
(2017); Sannai et al. (2021); Elesedy (2021); Zhu et al. (2021) establish a
generalization error that captures the improvement from invariance and
equivariance, which, however, is not applicable to the attention mechanism.
Our theoretical analysis of the generalization error follows the framework of
Bartlett et al. (2017), which stems from Bartlett (1996); Bartlett and
Mendelson (2002). In addition, the concurrent work (Zhang et al., 2022a)
provides a PAC-Bayes analysis of the generalization error of the attention
mechanism in the context of multiagent reinforcement learning.

Optimization of Deep Neural Networks. Our work is built on the vast body of
works that analyze the optimization error of deep neural networks (Allen-Zhu
et al., 2019a, c, b; Arora et al., 2019; Du et al., 2018, 2019; Zhang et al.,
2019; Zou et al., 2018; Zou and Gu, 2019; Allen-Zhu et al., 2019c; Cao and Gu,
2019; Li and Liang, 2018; Chizat et al., 2019; Mei et al., 2018, 2019;
Rotskoff and Vanden-Eijnden, 2018; Nguyen, 2019; Sirignano and Spiliopoulos,
2020). Most of them focus on overparameterized neural networks in the neural
tangent kernel (Jacot et al., 2018) or mean-field regime (Mei et al., 2018).
Our work analyzes the optimization error in the neural tangent kernel regime,
which is similar to Malladi et al. (2022). Meanwhile, it is worth mentioning
that our theoretical analysis of the approximation and generalization errors
is not restricted to the neural tangent kernel regime.

Invariance and Equivariance in Deep Neural Networks. Our work is related to a
recent line of works on deep neural networks with invariance or equivariance
with respect to permutations and other group operations. See, e.g., Scarselli
et al. (2008); Zaheer et al. (2017); Lee et al. (2019); Keriven and Peyré
(2019); Romero and Cordonnier (2020); Bloem-Reddy and Teh (2020); Hutchinson
et al. (2021); Satorras et al. (2021); Kossen et al. (2021) and the references
therein. Also, see Valle-Pérez and Louis (2020); Han et al. (2022) for a
detailed survey. In comparison, we exploit the latent variable model induced
by exchangeability to provide a complete characterization of the
representation, inference, and learning aspects of the attention mechanism.

##  2 Preliminary

Notations. We denote by $[L]$ the index set $\\{1,2,\ldots,L\\}$ for any
$L\in\mathbb{N}_{+}$. For any vector $v\in\mathbb{R}^{L}$, we denote by
$\mathtt{softmax}(v)=(\exp(v^{\ell})/(\sum_{\ell^{\prime}=1}^{L}\exp(v^{\ell^{\prime}})))_{\ell\in[L]}\in\mathbb{R}^{L}$
the softmax function. We denote by $\|\cdot\|_{2}$ the spectral norm, which
becomes the $\ell_{2}$-norm when it operates on a vector. We denote by
$\|\cdot\|_{\rm F}$ the Frobenius norm. For any $d\in\mathbb{N}_{+}$, we
denote by $\mathbb{S}^{d-1}=\\{x\in\mathbb{R}^{d}{\,|\,}\|x\|_{2}=1\\}$ the
$(d-1)$-dimensional unit sphere.

Reproducing Kernel Hilbert Space. Let $\mathcal{H}_{x}$ be a Hilbert space
over a domain ${\mathfrak{X}}$, which contains functions
$f:{\mathfrak{X}}\rightarrow\mathbb{R}$ and is equipped with the inner product
$\langle\cdot,\cdot\rangle_{\mathcal{H}_{x}}$. We say that $\mathcal{H}_{x}$
is a reproducing kernel Hilbert space (RKHS) with the kernel function
$\mathfrak{K}:{\mathfrak{X}}\times{\mathfrak{X}}\rightarrow\mathbb{R}$ if we
have the reproducing property that $\langle
f,\mathfrak{K}(x,\cdot)\rangle_{\mathcal{H}_{x}}=f(x)$ for any
$f\in\mathcal{H}_{x}$ and $x\in{\mathfrak{X}}.$ An RKHS $\mathcal{H}_{x}$ is
associated with a feature mapping $\phi:{\mathfrak{X}}\rightarrow\ell_{2}$
such that $\mathfrak{K}(x,x^{\prime})=\phi(x)^{\top}\phi(x^{\prime})$ for any
$x,x^{\prime}\in{\mathfrak{X}}$ (Muandet et al., 2016). Here we denote by
$\ell_{2}$ the space of all square-summable series.

Attention Mechanism. For an input sequence $X=\\{x^{\ell}\\}_{\ell\in[L]}$
with the input tokens $x^{\ell}\in\mathbb{R}^{d}$, we consider the key matrix
$K\in\mathbb{R}^{L\times d_{\mathrm{p}}}$ and the value matrix
$V\in\mathbb{R}^{L\times d}$ defined as

$\displaystyle K$
$\displaystyle=(k^{1},\ldots,k^{L})^{\top}=\bigl{(}k_{\theta}(x^{1}),\ldots,k_{\theta}(x^{L})\bigr{)}^{\top}\in\mathbb{R}^{L\times
d_{\mathrm{p}}},$ $\displaystyle V$
$\displaystyle=(v^{1},\ldots,v^{L})^{\top}=\bigl{(}v_{\theta}(x^{1}),\ldots,v_{\theta}(x^{L})\bigr{)}^{\top}\in\mathbb{R}^{L\times
d}.$

Here $k_{\theta}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d_{\mathrm{p}}}$ and
$v_{\theta}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ map the $\ell$-th input
token $x^{\ell}$ to the key $k^{\ell}$ and the value $v^{\ell}$, respectively,
where $\theta\in\Theta$ is the the learnable parameter. For any query
$q\in\mathbb{R}^{d_{\mathrm{p}}}$, we define the attention mechanism as
follows,

$\displaystyle\mathtt{attn}(q,K,V)=V^{\top}\mathtt{norm}\bigl{(}\mathfrak{K}(K,q)\bigr{)}\in\mathbb{R}^{d},$
(2.1)

where
$\mathfrak{K}:\mathbb{R}^{d_{\mathrm{p}}}\times\mathbb{R}^{d_{\mathrm{p}}}\rightarrow\mathbb{R}$
is a kernel function and we write
$\mathfrak{K}(K,q)=(\mathfrak{K}(k^{\ell},q))_{\ell\in L}\in\mathbb{R}^{L}$.
Here we denote by $\mathtt{norm}:\mathbb{R}^{L}\rightarrow\mathbb{R}^{L}$ a
normalization mapping.

A common example of the attention mechanism is the softmax attention (Vaswani
et al., 2017), where the kernel function is the exponential kernel
$\mathfrak{K}_{\exp}(q,k)=\exp(q^{\top}k/\gamma)$ with a fixed $\gamma>0$ and
the normalization mapping is the following softmax normalization,

$\displaystyle\mathtt{norm}_{\mathtt{SM}}\bigl{(}\mathfrak{K}(K,q)\bigr{)}=\bigl{(}{\bm{1}}^{\top}\mathfrak{K}(K,q)\bigr{)}^{-1}\cdot\mathfrak{K}(K,q).$

The attention mechanism in (2.1) with the exponential kernel and the softmax
normalization is the softmax attention (Vaswani et al., 2017), which takes the
following form,

$\displaystyle\mathtt{attn}_{{\mathtt{SM}}}(q,K,V)=V^{\top}\mathtt{norm}_{\mathtt{SM}}\bigl{(}\mathfrak{K}_{\exp}(K,q)\bigr{)}=\mathtt{softmax}(Kq/\gamma).$

##  3 Representation, Inference, and Estimation
via Latent Variable Model

From Exchangeability to Latent Variable Model. We consider the input sequence
$X=\\{x^{\ell}\\}_{\ell\in[L]}$, where $x^{\ell}\in\mathbb{R}^{d}$ is an input
token and $L\in\mathbb{N}_{+}$ is the sequence length. In natural language
processing (NLP), such a sequence consists of embeddings of words in a
paragraph, while in computer vision (CV), such a sequence consists of
embeddings of patches in an image.

As is the case in BERT (Devlin et al., 2018) and ViT (Dosovitskiy et al.,
2020), the input sequence is exchangeable since it includes positional
encodings. Specifically, we say that a random variable sequence
$\\{x^{\ell}\\}_{\ell\in\mathbb{N}_{+}}$ is exchangeable if and only if it
holds for any sequence length $L\in\mathbb{N}_{+}$ and any index permutation
$\pi:[L]\rightarrow[L]$ that

$\displaystyle\mathbb{P}(x^{1},x^{2},\ldots,x^{L})=\mathbb{P}(x^{\pi(1)},x^{\pi(2)},\ldots,x^{\pi(L)}).$

In other words, permuting the index order within the random variable sequence
does not affect its joint distribution. The following proposition states that
the exchangeability of a random variable sequence induces a latent variable.

######  Proposition 3.1 (de Finetti Representation Theorem (de Finetti,
1937)).

Let $\\{x^{\ell}\\}_{\ell\in\mathbb{N}_{+}}$ be an exchangeable sequence.
Then, there exists a latent variable $z$ such that for any sequence length
$L\in\mathbb{N}_{+}$,

$\displaystyle\mathbb{P}(x^{1},\ldots,x^{L})$
$\displaystyle=\int\prod_{\ell=1}^{L}\mathbb{P}(x^{\ell}{\,|\,}z)\cdot\mathbb{P}(z){\mathrm{d}}z,$
$\displaystyle\mathbb{P}(x^{\ell}{\,|\,}x^{1},\ldots,x^{\ell-1},x^{\ell+1},\ldots,x^{L})$
$\displaystyle=\int\mathbb{P}(x^{\ell}{\,|\,}z)\cdot\mathbb{P}(z{\,|\,}x^{1},\ldots,x^{\ell-1},x^{\ell+1},\ldots,x^{L}){\mathrm{d}}z,\quad\forall\ell\in[L].$

We remark that Proposition 3.1 requires an infinite-length exchangeable
sequence. Up to an approximation error, a finite-length exchangeable sequence
also induces a latent variable (Diaconis and Freedman, 1980). In what follows,
we consider the former case where the input sequence includes positional
encodings and is thus exchangeable (Devlin et al., 2018; Dosovitskiy et al.,
2020). See Figure 1 for an illustration of the exchangeability.

Figure 1: The input sequence (the raw version without positional encodings)
becomes exchangeable with positional encodings. In practice, the positional
encoding is incorporated in an additive manner (instead of concatenation).

Proposition 3.1 guarantees the existence of a latent variable, which forms the
basis of our theoretical analysis. See Figure 2(a) for an illustration.
Intuitively, the latent variable can be viewed as the “concept” of the input
sequence, which is “summarized” over words or patches. For instance, in NLP,
the latent variable instantiates the “meaning” of a paragraph, while in CV,
the latent variable instantiates the “theme” of an image. In particular, the
latent posterior $\mathbb{P}(z{\,|\,}X)$ plays a key role in solving
downstream tasks (Song et al., 2014; Xie et al., 2021), as it is a desired
representation of the input sequence $X$. See Figure 2(b) for an illustration.
In the following lemma, we prove that the latent posterior
$b_{z}(X)=\mathbb{P}(z=\cdot{\,|\,}X)$ is a minimal sufficient statistic
(Fisher, 1922).

######  Lemma 3.2 (Minimal Sufficiency of Latent Posterior).

Let $z$ be the latent variable induced by the exchangeability of the input
sequence $X$. The latent posterior $b_{z}(X)=\mathbb{P}(z=\cdot{\,|\,}X)$ is a
minimal sufficient statistic of the input sequence $X$ for the latent variable
$z$. Meanwhile, for any target variable $y$ that is independent of the input
sequence $X$ conditioning on the latent variable $z$, we assume the
invertibility of the operator ${\mathcal{T}}$ defined by

$\displaystyle({\mathcal{T}}f)(y)=\int\mathbb{P}(y{\,|\,}z)f(z){\mathrm{d}}z.$

Then, the latent posterior $b_{z}(X)=\mathbb{P}(z=\cdot{\,|\,}X)$ is a minimal
sufficient statistic of the input sequence $X$ for the target variable $y$.

###### Proof.

See §B.3 for a detailed proof. ∎

From Latent Variable Model to Learning Objectives. In what follows, we
consider the prediction task in BERT (Devlin et al., 2018) and ViT
(Dosovitskiy et al., 2020). Let $y$ be the target variable and
$X=\\{x^{\ell}\\}_{\ell\in[L]}$ is the input sequence. In particular, in self-
supervised learning (BERT), the target variable $y$ is a masked token of the
input sequence, while in supervised learning (ViT), $y$ is the unknown label
corresponding to the class encoding. We remark that in ViT, the unknown label
$y$ corresponds to the masked token in BERT, while the input class encoding
corresponds to the mask in BERT. In both cases, the concatenation
$\\{y,x^{1},\ldots,x^{L}\\}$ is treated as an exchangeable sequence since it
includes the positional encodings. By Proposition 3.1, we have

$\displaystyle\mathbb{P}(y{\,|\,}X)=\int\mathbb{P}(y{\,|\,}z)\cdot\mathbb{P}(z{\,|\,}X){\mathrm{d}}z,$
(3.1)

where $z$ is the latent variable induced by the exchangeability of $X$. See
Figure 2 for an illustration.

In what follows, we treat $y$ as a target variable that satisfies (3.1), which
specifies that $y$ is independent of the input sequence $X$ conditioning on
the latent variable $z$. By Lemma 3.2, the latent posterior $b_{z}(X)$ is a
minimal sufficient statistic of $X$ for $y$. In other words, the latent
posterior $b_{z}(X)$ is a desired representation of the input sequence $X$.
According to (3.1), the prediction of the target variable $y$ from the input
sequence $X$ (forward pass) takes two implicit steps: i) the inference of the
latent posterior $\mathbb{P}(z{\,|\,}X)$, and ii) the prediction of $y$ based
on the generative distribution $\mathbb{P}(y{\,|\,}z)$ integrated with the
latent posterior $\mathbb{P}(z{\,|\,}X)$.

(a) Prediction of a masked token $x^{\ell}$.

(b) Prediction of the target variable $y$, which can be viewed as a masked
token.

Figure 2: The forward pass for the prediction of the masked token $x^{\ell}$
and the target variable $y$. The prediction of $y$ takes two steps: i) the
inference of the latent posterior $\mathbb{P}(z{\,|\,}X)$, and ii) the
prediction of $y$ based on the generative distribution $\mathbb{P}(y{\,|\,}z)$
integrated with the latent posterior $\mathbb{P}(z{\,|\,}X)$.

To construct the learning objective, we consider the distribution of the
target variable $y$ conditioning on the input sequence $X$ and parameterize it
by $\mathbb{P}_{\theta}(y{\,|\,}X)$, where $\theta$ is the learnable
parameter. Given a dataset $\mathcal{D}_{n}=\\{(X_{i},y_{i})\\}_{i\in[n]}$,
where $X_{i}$ is the $i$-th input sequence and $y_{i}$ is the $i$-th target
variable, the maximum likelihood estimation (MLE) objective takes the
following form,

$\displaystyle\max_{\theta}\widehat{\mathbb{E}}_{(X,y)\sim\mathcal{D}_{n}}\bigl{[}\log\mathbb{P}_{\theta}(y{\,|\,}X)\bigr{]}$
(3.2)
$\displaystyle\quad=\widehat{\mathbb{E}}_{(X,y)\sim\mathcal{D}_{n}}\biggl{[}\log\int\mathbb{P}_{\theta}(y{\,|\,}z)\mathbb{P}_{\theta}(z{\,|\,}X){\mathrm{d}}z\biggr{]}.$

We define $\widehat{\mathbb{E}}_{(X,y)\sim\mathcal{D}_{n}}[\,\cdot\,]$ as the
empirical expectation with respect to the dataset $\mathcal{D}_{n}$.

###  3.1 Preliminary Finite-Dimensional Example

Latent Variable Model. We provide a finite-dimensional Gaussian-distributed
example to illustrate the latent variable model and the MLE objective.
Specifically, we consider the setting with the input sequence
$X=\\{x^{\ell}\\}_{\ell\in[L]}$ and the target variable $y$, where
$x^{\ell}\in\mathbb{R}^{d}$ and $y\in\mathbb{R}^{d}$. For the input sequence
$X$, we consider the following example of the latent variable model in (3.1),

$\displaystyle r^{\ell}=zc^{\ell}+\epsilon^{\ell},\qquad\text{where}\quad
c^{\ell}=c_{*}(x^{\ell}),\quad
r^{\ell}=r_{*}(x^{\ell}),\quad\forall\ell\in[L].$ (3.3)

Here $z\in\mathbb{R}^{d_{r}\times d_{c}}$ is the latent variable induced by
the exchangeability of the input sequence, $c^{\ell}\in\mathbb{R}^{d_{c}}$ and
$r^{\ell}\in\mathbb{R}^{d_{r}}$ are the covariate and response, respectively,
which are determined by two unknown functions
$c_{*}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d_{c}}$ and
$r_{*}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d_{r}}$, and $\epsilon^{\ell}\sim
N(0,\sigma^{2}I)$ is the noise, which is independent of $c^{\ell}$. In
practice, the covariate $c^{\ell}$ instantiates the contextual information,
while the response $r^{\ell}$ instantiates the semantic information. We
consider the prediction of the (unknown) target variable $y$ based on its
(known) input mask ${\mathtt{msk}}$. In the self-supervised setting (Devlin et
al., 2018), $y$ is a masked token of the input sequence, while
${\mathtt{msk}}$ is the positional encoding. In the supervised setting
(Dosovitskiy et al., 2020), $y$ is the label of the input sequence, while
${\mathtt{msk}}$ is the class encoding. Specifically, corresponding to (3.3),
we consider the prediction model with $y=r^{\mathtt{msk}}$, where
$r^{\mathtt{msk}}$ is the response corresponding to the covariate
$c^{\mathtt{msk}}$ of the input mask such that

$\displaystyle
y=r^{\mathtt{msk}}=zc^{\mathtt{msk}}+\epsilon,\qquad\text{where}\quad
c^{\mathtt{msk}}=c_{*}^{\mathtt{m}}({\mathtt{msk}}).$ (3.4)

Here $c_{*}^{\mathtt{m}}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d_{c}}$ is an
unknown function and $\epsilon\sim N(0,\sigma^{2}I)$ is the noise, which is
independent of $c^{\mathtt{msk}}$. For example, (3.4) holds when we consider
an exchangeable sequence $\\{x^{1},\ldots,x^{L},x^{L+1}\\}$ satisfying (3.3)
for any $\ell\in[L+1]$ with the input sequence $X=\\{x^{1},\ldots,x^{L}\\}$,
$c^{\mathtt{msk}}=c^{L+1}$, and the target variable
$y=r^{\mathtt{msk}}=r^{L+1}$. In the next section, we consider an advanced
infinite-dimensional example and show that $c^{\mathtt{msk}}$, $c^{\ell}$, and
$r^{\ell}$ correspond to the query, the key, and the value in the attention
mechanism, respectively.

Note that the regression model in (3.3) is a conditional model. Instead of
modeling the conditional distribution of $y$ given $X$ as in (3.1), we model
the conditional distribution of $y$ given $X$ and ${\mathtt{msk}}$. Recall
that $y=r^{\mathtt{msk}}$. Corresponding to (3.1), the generative distribution
takes the following form,

$\displaystyle\mathbb{P}(y{\,|\,}{\mathtt{msk}},z)\propto\exp\Bigl{(}-\big{\|}y-zc_{*}^{\mathtt{m}}({\mathtt{msk}})\big{\|}_{2}^{2}\big{/}2\sigma^{2}\Bigr{)}.$
(3.5)

We take the Gaussian distribution $N(0,\lambda I)$ as the prior of $z$. By
(3.3), the latent posterior $\mathbb{P}(z{\,|\,}X)$ is a Gaussian
distribution, which (approximately) takes the following form,

$\displaystyle\mathbb{P}(z{\,|\,}X)\propto\exp\Bigl{(}-\big{\|}z-\bar{z}(X)\big{\|}_{2}^{2}\big{/}2\iota^{2}\Bigr{)}.$
(3.6)

Here the covariance of the latent posterior is approximated by $\iota^{2}I$
and the mean $\bar{z}(X)$ of the latent posterior takes the following form,

$\displaystyle\bar{z}(X)=\mathbb{E}[z{\,|\,}X]=R^{\top}(CC^{\top}+\lambda
I)^{-1}C,$

where we define $C=(c^{1},\ldots,c^{L})^{\top}\in\mathbb{R}^{L\times d_{c}}$
and $R=(r^{1},\ldots,r^{L})^{\top}\in\mathbb{R}^{L\times d_{r}}$. Combining
(3.5) and (3.6), we obtain

$\displaystyle\mathbb{P}(y{\,|\,}{\mathtt{msk}},X)=\int\mathbb{P}(y{\,|\,}{\mathtt{msk}},z)\cdot\mathbb{P}(z{\,|\,}X){\mathrm{d}}z\propto\exp\Bigl{(}-\big{\|}y-\bar{z}(X)c_{*}^{\mathtt{m}}({\mathtt{msk}})\big{\|}_{2}^{2}\big{/}2\widetilde{\sigma}^{2}\Bigr{)},$
(3.7)

which corresponds to (3.1), Here we approximate the covariance of $y$
conditioning on $X$ and ${\mathtt{msk}}$ by $\widetilde{\sigma}^{2}I$, where
$\widetilde{\sigma}$ does not depend on $X$. We remark that (3.7) is a form of
Bayesian model averaging (Wasserman, 2000) within one data point.

Parameterization of Latent Variable Model. Recall that $c_{*}$, $r_{*}$, and
$c^{\mathtt{m}}_{*}$ in (3.3) and (3.4) are unknown. We parameterize them with
$c_{\theta}$, $r_{\theta}$, and $c^{\mathtt{m}}_{\theta}$, where
$\theta\in\Theta$ is a learnable parameter. With the ideal parameter
$\theta^{*}\in\Theta$, it holds for any $\ell\in[L]$ that

$\displaystyle c_{\theta^{*}}(x^{\ell})=c_{*}(x^{\ell})=c^{\ell},\quad
r_{\theta^{*}}(x^{\ell})=r_{*}(x^{\ell})=r^{\ell},\quad
c^{\mathtt{m}}_{\theta^{*}}({\mathtt{msk}})=c^{\mathtt{m}}_{*}({\mathtt{msk}})=c^{\mathtt{msk}}.$
(3.8)

By (3.6), we parameterize the latent posterior $\mathbb{P}(z{\,|\,}X)$ as
follows,

$\displaystyle\mathbb{P}_{\theta}(z{\,|\,}X)\propto\exp\Bigl{(}-\big{\|}z-\bar{z}_{\theta}(X)\big{\|}_{2}^{2}\big{/}2\iota^{2}\Bigr{)},$
(3.9)

where $\bar{z}_{\theta}(X)$ is calculated as follows,

$\displaystyle\bar{z}_{\theta}(X)=r_{\theta}(X)^{\top}\bigl{(}c_{\theta}(X)c_{\theta}(X)^{\top}+\lambda
I\bigr{)}^{-1}c_{\theta}(X).$

Here
$c_{\theta}(X)=(c_{\theta}(x^{1}),\ldots,c_{\theta}(x^{L}))^{\top}\in\mathbb{R}^{L\times
d_{c}}$ and
$r_{\theta}(X)=(r_{\theta}(x^{1}),\ldots,r_{\theta}(x^{L}))^{\top}\in\mathbb{R}^{L\times
d_{r}}$. By (3.5), we parameterize the generative distribution
$\mathbb{P}(y{\,|\,}{\mathtt{msk}},z)$ as follows,

$\displaystyle\mathbb{P}_{\theta}(y{\,|\,}{\mathtt{msk}},z)\propto\exp\Bigl{(}-\big{\|}y-zc^{\mathtt{m}}_{\theta}({\mathtt{msk}})\big{\|}_{2}^{2}\big{/}2\sigma^{2}\Bigr{)}.$

By (3.7), we define the conditional likelihood
$\mathbb{P}(y{\,|\,}{\mathtt{msk}},X)$ as follows,

$\displaystyle\mathbb{P}_{\theta}(y{\,|\,}{\mathtt{msk}},X)\propto\exp\Bigl{(}-\big{\|}y-\bar{z}_{\theta}(X)c_{\theta}^{\mathtt{m}}({\mathtt{msk}})\big{\|}_{2}^{2}\big{/}2\widetilde{\sigma}^{2}\Bigr{)}.$
(3.10)

Training and Testing. In the training phase, given the dataset
$\mathcal{D}_{n}=\\{(X_{i},y_{i})\\}_{i\in[n]}$, we aim to maximize the MLE
objective in (3.2). By (3.10), maximizing the MLE objective is equivalent to
minimizing the mean-squared error as follows,

$\displaystyle\min_{\theta}\widehat{\mathbb{E}}_{(X,y)\sim\mathcal{D}_{n}}\Bigl{[}\big{\|}y-\bar{z}_{\theta}(X)c^{\mathtt{m}}_{\theta}({\mathtt{msk}})\big{\|}_{2}^{2}\Bigr{]}.$
(3.11)

Note that the learnable parameter $\theta$ is estimated across different data
points in the dataset $\mathcal{D}_{n}$ through the backward pass, while the
latent variable $z$ is inferred within one data point $(X_{i},y_{i})$ through
the forward pass. We remark that by learning $\theta$, the model learns to
perform Bayesian model averaging. Suppose that we solve (3.11) and obtain the
estimator $\widehat{\theta}$. In the testing phase, given an input sequence
$X_{\dagger}$ and an input mask ${\mathtt{msk}}_{\dagger}$, we predict the
target variable $y_{\dagger}$ by maximizing the posterior of $y$,

$\displaystyle\widehat{y}=\mathop{\mathrm{argmax}}_{y}\mathbb{P}_{\widehat{\theta}}(y{\,|\,}{\mathtt{msk}}_{\dagger},X_{\dagger})=\mathbb{E}[r^{\mathtt{msk}}{\,|\,}{\mathtt{msk}}_{\dagger},X_{\dagger}]=\bar{z}_{\widehat{\theta}}(X_{\dagger})c^{\mathtt{m}}_{\widehat{\theta}}({\mathtt{msk}}_{\dagger}).$
(3.12)

We remark that the learning process for the attention mechanism involves two
aspects. In the forward pass, within one data point, we infer the latent
posterior $\mathbb{P}(z{\,|\,}X)$ to predict the target variable $y$. In the
backward pass, we estimate the learnable parameter $\theta$ across different
data points. See Figure 3 for an illustration.

Figure 3: Forward pass: within one data point $(X,y)$, we infer the latent
posterior $\mathbb{P}_{\theta}(z{\,|\,}X)$ by (3.9). We predict $y_{\dagger}$
by $\widehat{y}$ in (3.12). Backward pass: across different data points in the
dataset $\mathcal{D}_{n}$, we estimate the learnable parameter $\theta$ by
(3.11).

The finite-dimensional example illustrates a “greybox” approach to modeling,
that is, it combines the handcrafted architecture in (3.7) (“whitebox”), and
the learnable parameter in (3.11), which is estimated in an end-to-end manner
through empirical risk minimization (“blackbox”). As shown in Figure 4, the
forward pass first infers the latent variable $z$ and then utilizes the latent
variable $z$ to predict the masked token (in the self-supervised setting) or
the unknown label (in the supervised setting). Meanwhile, the backward pass
estimatess the learnable parameter. In the following section, we extend the
finite-dimensional example to the infinite-dimensional setting, which recovers
the attention mechanism. In particular, we demonstrate that the attention
mechanism infers the latent posterior within a data point. Also, we show that
the covariate corresponds to the query and key and that the response
corresponds to the value in the attention mechanism.

Figure 4: The forward and backward passes in transformers. Dotted arrows stand
for forward passes (input$\rightarrow$latent$\rightarrow$target). Solid arrows
stand for backward passes (training). Masks (grey tokens) are only used to
illustrate the self-supervised setting (yellow box).

##  4 Attention as Latent Posterior Inference

In what follows, we demonstrate how the attention mechanism performs latent
posterior inference for the latent variable model, which is induced by the
exchangeability of the input sequence. In §4.1, we extend the finite-
dimensional example in §3 to an RKHS to induce a variant of the softmax
attention, namely, the conditional mean embedding (CME) attention. In
particular, we prove that it infers the latent posterior in the forward pass.
In §4.2, we prove that the softmax attention has the same limit as the CME
attention when the sequence length goes to infinity, which implies that the
softmax attention approximately infers the latent posterior.

###  4.1 Attention as Kernel Conditional Mean Embedding

Advanced Infinite-Dimensional Example. We present an infinite-dimensional
version of the preliminary example in §3.1, which motivates us to study the
CME attention. Similarly to (3.3), we consider the following model for the
input sequence $X=\\{x^{\ell}\\}_{\ell\in[L]}$ with input token
$x^{\ell}\in\mathbb{R}^{d}$,

$\displaystyle r^{\ell}$
$\displaystyle=z\phi(c^{\ell})+\epsilon^{\ell},\qquad\text{where}\quad
c^{\ell}=c_{*}(x^{\ell}),\quad
r^{\ell}=r_{*}(x^{\ell}),\quad\forall\ell\in[L].$ (4.1)

Here $c^{\ell}\in\mathbb{R}^{d_{\mathrm{p}}}$ and $r^{\ell}\in\mathbb{R}^{d}$
are the covariate and the response, respectively, which are determined by two
unknown functions $c_{*}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d_{\mathrm{p}}}$
and $r_{*}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$,
$\phi:\mathbb{R}^{d_{\mathrm{p}}}\rightarrow\mathcal{H}_{c}$ is the feature
mapping of the RKHS $\mathcal{H}_{c}$,
$z:\mathcal{H}_{c}\rightarrow\mathbb{R}^{d}$ is a linear mapping, which is
viewed as the latent variable induced by the exchangeability of the input
sequence $X$, and $\epsilon^{\ell}\sim N(0,\sigma^{2}I)$ is the Gaussian
noise, which is independent of the covariate $c^{\ell}$. Note that
$\phi(c^{1})^{\top}\phi(c^{2})=\mathfrak{K}(c^{1},c^{2})$ for any
$c^{1},c^{2}\in\mathbb{R}^{d_{\mathrm{p}}}$, where
$\mathfrak{K}:\mathbb{R}^{d_{\mathrm{p}}}\times\mathbb{R}^{d_{\mathrm{p}}}\rightarrow\mathbb{R}$
is the kernel function of the RKHS $\mathcal{H}_{c}$. A common example is the
Gaussian radial basis function (RBF) kernel
$\mathfrak{K}_{\mathtt{RBF}}(q,k)=\exp(-\|q-k\|_{2}^{2}/2\gamma)$ with
$\gamma>0$. Similarly to (3.4), the (unknown) target variable $y$ is
determined by its (known) input mask ${\mathtt{msk}}$, which satisfies

$\displaystyle
y=r^{\mathtt{msk}}=z\phi(c^{\mathtt{msk}})+\epsilon,\qquad\text{where}\quad
c^{\mathtt{msk}}=c_{*}^{\mathtt{m}}({\mathtt{msk}}).$ (4.2)

Here we denote by $c^{\mathtt{msk}}$ and $r^{\mathtt{msk}}$ the covariate and
the response corresponding to the input ${\mathtt{msk}}$, respectively,
$c_{*}^{\mathtt{m}}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d_{\mathrm{p}}}$ is
an unknown function, and $\epsilon\sim N(0,\sigma^{2}I)$ is a Gaussian noise,
which is independent of $c^{\mathtt{msk}}$. To simplify the presentation, we
view the RKHS $\mathcal{H}_{c}$ as a vector space $\mathbb{R}^{d_{\phi}}$ with
$d_{\phi}=\infty$. Correspondingly, we view the latent variable $z$ as a
matrix in $\mathbb{R}^{d\times d_{\phi}}$. We present a rigorous
characterization of $z$ in §B.1 based on Gaussian process regression.
Similarly to (3.8), we parameterize $c_{*}$, $r_{*}$, and $c^{\mathtt{m}}_{*}$
by $c_{\theta}$, $r_{\theta}$, and $c^{\mathtt{m}}_{\theta}$, where
$\theta\in\Theta$ is a learnable parameter. Similarly to (3.12), we predict
$r^{\mathtt{msk}}$ in the forward pass via

$\displaystyle\widehat{r}$
$\displaystyle=\mathbb{E}[r^{\mathtt{msk}}{\,|\,}{\mathtt{msk}},X]$
$\displaystyle=r_{\theta}(X)^{\top}\Bigl{(}\phi\bigl{(}c_{\theta}(X)\bigr{)}\phi\bigl{(}c_{\theta}(X)\bigr{)}^{\top}+\lambda
I\Bigr{)}^{-1}\phi\bigl{(}c_{\theta}(X)\bigr{)}\phi\bigl{(}c_{\theta}^{\mathtt{m}}({\mathtt{msk}})\bigr{)}$
$\displaystyle=r_{\theta}(X)^{\top}\Bigl{(}\mathfrak{K}\bigl{(}c_{\theta}(X),c_{\theta}(X)\bigr{)}+\lambda
I\Bigr{)}^{-1}\mathfrak{K}\bigl{(}c_{\theta}(X),c_{\theta}^{\mathtt{m}}({\mathtt{msk}})\bigr{)},$
(4.3)

where we define
$\mathfrak{K}(c_{\theta}(X),c_{\theta}(X))=(\mathfrak{K}(c_{\theta}(x^{i}),c_{\theta}(x^{j})))_{i,j\in[L]}\in\mathbb{R}^{L\times
L}$,
$\mathfrak{K}(c_{\theta}(X),c_{\theta}^{\mathtt{m}}({\mathtt{msk}}))=(\mathfrak{K}(c_{\theta}(x^{\ell}),c_{\theta}^{\mathtt{m}}({\mathtt{msk}})))_{\ell\in[L]}\in\mathbb{R}^{L}$,
$\phi(c_{\theta}(X))=(\phi(c_{\theta}(x^{1})),\ldots,\phi(c_{\theta}(x^{L})))^{\top}\in\mathbb{R}^{L\times
d_{\phi}}$, and
$r_{\theta}(X)=(r_{\theta}(x^{1}),\ldots,r_{\theta}(x^{L}))^{\top}\in\mathbb{R}^{L\times
d}$. We remark that (4.1) recovers the empirical version of the kernel
conditional mean embedding of $\mathbb{P}_{\mathcal{R}{\,|\,}\mathcal{C}}$
(Song et al., 2009), where we denote by
$\mathbb{P}_{\mathcal{R}{\,|\,}\mathcal{C}}$ the conditional distribution of
$r^{\ell}$ given $c^{\ell}$ (as two random variables) within one data point,
and $\mathcal{H}_{r}=(\mathbb{R}^{d})^{*}=\mathbb{R}^{d}$ is the dual space of
$\mathbb{R}^{d}$ equipped with the Euclidean kernel
$\langle\cdot,\cdot\rangle$.

From Latent Variable Model to Attention. Recall that the attention mechanism
is defined in (2.1) with $q\in\mathbb{R}^{d_{\mathrm{p}}}$,
$K=(k^{1},\ldots,k^{L})^{\top}\in\mathbb{R}^{L\times d_{\mathrm{p}}}$, and
$V=(v^{1},\ldots,v^{L})^{\top}\in\mathbb{R}^{L\times d}$. The kernel
conditional mean embedding in (4.1) motivates us to consider the following CME
normalization,

$\displaystyle\mathtt{norm}_{\mathtt{CME}}\bigl{(}\mathfrak{K}(K,q)\bigr{)}=\bigl{(}\mathfrak{K}(K,K)+\lambda
I\bigr{)}^{-1}\mathfrak{K}(K,q),$ (4.4)

where we write
$\mathfrak{K}(K,q)=(\mathfrak{K}(k^{\ell},q))_{\ell\in[L]}\in\mathbb{R}^{L}$
and
$\mathfrak{K}(K,K)=(\mathfrak{K}(k^{i},k^{j}))_{(i,j)\in[L]\times[L]}\in\mathbb{R}^{L\times
L}$ We call the attention mechanism with the CME normalization in (4.4) the
CME attention and denote it by $\mathtt{attn}_{{\mathtt{CME}}}$. In
particular, the CME attention takes the following form,

$\displaystyle\mathtt{attn}_{\mathtt{CME}}(q,K,V)=V^{\top}\mathtt{norm}_{\mathtt{CME}}\bigl{(}\mathfrak{K}(K,q)\bigr{)}=V^{\top}\bigl{(}\mathfrak{K}(K,K)+\lambda
I\bigr{)}^{-1}\mathfrak{K}(K,q)\in\mathbb{R}^{d}.$ (4.5)

We see that the CME attention recovers (4.1) when

$\displaystyle q=c_{\theta}^{\mathtt{m}}({\mathtt{msk}}),\qquad
k^{\ell}=c_{\theta}(x^{\ell}),\qquad
v^{\ell}=r_{\theta}(x^{\ell}),\qquad\forall\ell\in[L].$ (4.6)

We remark that (4.6) establishes a connection between the latent variable
model and the attention mechanism. In other words, the covariate
$c^{\mathtt{msk}}$ of the input mask ${\mathtt{msk}}$ corresponds to the query
$q$, the covariate $c^{\ell}$ of the input token $x^{\ell}$ corresponds to the
key $k^{\ell}$ for $\ell\in[L]$, and the response $r^{\ell}$ of the input
token $x^{\ell}$ corresponds to the value $v^{\ell}$ for $\ell\in[L]$. In the
attention mechanism, we denote by
$q_{\theta}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d_{\mathrm{p}}}$,
$k_{\theta}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d_{\mathrm{p}}}$, and
$v_{\theta}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ the mappings from the
input token to the query, the key, and the value, respectively, with the
learnable parameter $\theta$. In particular, we have the following
correspondence,

$\displaystyle q_{\theta}=c_{\theta}^{\mathtt{m}},\qquad
k_{\theta}=c_{\theta},\qquad v_{\theta}=r_{\theta}.$

In an common example, we instantiate $q_{\theta}$, $k_{\theta}$, and
$v_{\theta}$ for $x\in\mathbb{R}^{d}$ as follows,

$\displaystyle q_{\theta}(x)=(W^{\mathrm{q}})^{\top}\mathtt{nn}(x;A),\qquad
k_{\theta}(x)=(W^{\mathrm{k}})^{\top}\mathtt{nn}(x;A),\qquad
v_{\theta}(x)=(W^{\mathrm{v}})^{\top}\mathtt{nn}(x;A),$ (4.7)

where $W^{\mathrm{q}},W^{\mathrm{k}}\in\mathbb{R}^{d\times d_{\mathrm{p}}}$
and $W^{\mathrm{v}}\in\mathbb{R}^{d\times d}$ are learnable parameters, and we
denote by $\mathtt{nn}(\cdot;A):\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ the
feedforward neural network with the learnable parameter $A$ and we summarize
the learnable parameter as
$\theta=(A,W^{\mathrm{q}},W^{\mathrm{k}},W^{\mathrm{v}})$. Similarly to
(3.11), the MLE objective takes the following form,

$\displaystyle\min_{\theta}\widehat{\mathbb{E}}_{(X,y)\sim\mathcal{D}_{n}}\biggl{[}\Big{\|}y-\mathtt{attn}_{{\mathtt{CME}}}\bigl{(}q_{\theta}({\mathtt{msk}}),k_{\theta}(X),v_{\theta}(X)\bigr{)}\Big{\|}_{2}^{2}\biggr{]},$
(4.8)

where we write
$k_{\theta}(X)=(k_{\theta}(x^{1}),\ldots,k_{\theta}(x^{L}))^{\top}\in\mathbb{R}^{L\times
d_{\mathrm{p}}}$ and
$v_{\theta}(X)=(v_{\theta}(x^{1}),\ldots,v_{\theta}(x^{L}))^{\top}\in\mathbb{R}^{L\times
d}$.

Limit of CME Attention with $L\rightarrow\infty$. Given an input sequence
$X=\\{x^{\ell}\\}_{\ell\in[L]}$, we consider the key-value pairs
$\\{(k^{\ell},v^{\ell})\\}_{\ell\in[L]}$ obtained from
$k^{\ell}=k_{\theta}(x^{\ell})$ and $v^{\ell}=v_{\theta}(x^{\ell})$ for a
fixed $\theta$. For notational simplicity, we denote by $\mathcal{K}$ and
$\mathcal{V}$ the random variables with the same distribution as
$(k^{\ell},v^{\ell})$ within one data point. Recall that we define the CME
attention in (4.5). In the following proposition, we prove that the CME
attention approximates the kernel conditional mean embedding of
$\mathbb{P}_{\mathcal{V}{\,|\,}\mathcal{K}}$ as $L\rightarrow\infty$. Note
that the following proposition does not depend on the latent variable model in
(4.1).

######  Proposition 4.1 (CME Attention Converges to Kernel Conditional Mean
Embedding).

Let $\mathfrak{K}$ be a positive definite kernel function. We assume that
$\\{x^{\ell}\\}_{\ell\in[L]}$ in the input sequence $X$ are independent and
identically distributed (within one data point). It holds with probability at
least $1-\delta$ that,

$\displaystyle\big{\|}\mathtt{attn}_{\mathtt{CME}}(q,K,V)-\mathbb{E}[\mathcal{V}{\,|\,}\mathcal{K}=q]\big{\|}_{2}=\mathcal{O}\bigl{(}L^{-1/2}\cdot\log(1/\delta)+\lambda\bigr{)},\qquad\forall
q\in\mathbb{R}^{d_{\mathrm{p}}}$

###### Proof.

See §B.4 for a detailed proof. ∎

Note that the CME attention $\mathtt{attn}_{\mathtt{CME}}$ is a variant of the
softmax attention (Vaswani et al., 2017) with a different normalization. In
the following section, we prove that the softmax attention has the same limit
as the CME attention when the sequence length $L$ goes to infinity.

###  4.2 Softmax Attention Infers Latent Posterior

In §4.1, we demonstrate how the latent variable model motivates the design of
the CME attention. Recall that we consider the attention mechanism in the form
of (2.1) with $q\in\mathbb{R}^{d_{\mathrm{p}}}$, $K\in\mathbb{R}^{L\times
d_{\mathrm{p}}}$, and $V\in\mathbb{R}^{L\times d}$. In practice, a common
normalization is defined as follows,

$\displaystyle\mathtt{norm}_{\mathtt{SM}}(\mathfrak{K}(K,q))=\bigl{(}\mathbf{1}_{L}^{\top}\mathfrak{K}(K,q)\bigr{)}^{-1}\cdot\mathfrak{K}(K,q),$
(4.9)

where $\mathbf{1}_{L}\in\mathbb{R}^{L}$ is the $L$-dimensional all-one vector
and we recall that
$\mathfrak{K}(K,q)=(\mathfrak{K}(k^{\ell},q))_{\ell\in[L]}\in\mathbb{R}^{L}$.
We denote by $\mathtt{attn}_{\mathtt{SM}}$ the attention mechanism with the
normalization in (4.9). When the kernel function is the exponential kernel
$\mathfrak{K}_{\exp}(q,k)=\exp(k^{\top}q/\gamma)$ for any
$q,k\in\mathbb{R}^{d_{\mathrm{p}}}$ and a fixed $\gamma>0$, the attention
mechanism in (2.1) takes the following form,

$\displaystyle\mathtt{attn}_{\mathtt{SM}}(q,K,V)=V^{\top}\mathtt{norm}_{\mathtt{SM}}\bigl{(}\mathfrak{K}_{\exp}(K,q)\bigr{)}=V^{\top}\mathtt{softmax}(Kq/\gamma),$
(4.10)

which recovers the softmax attention in Vaswani et al. (2017). In what
follows, we prove that as the sequence length $L$ goes to infinity, the
softmax attention $\mathtt{attn}_{{\mathtt{SM}}}$ has the same limit as the
CME attention $\mathtt{attn}_{\mathtt{CME}}$.

Softmax Attention Has the Same Limit as CME Attention with
$L\rightarrow\infty$. First, we demonstrate that the softmax attention in
(4.10) is a conditional kernel density estimator of
$\mathbb{P}_{\mathcal{V}{\,|\,}\mathcal{K}}$. We define the conditional kernel
density estimator (KDE) as follows,

$\displaystyle\widehat{\mathbb{P}}_{\mathcal{V}{\,|\,}\mathcal{K}}^{\mathfrak{K}}(v{\,|\,}k)\propto\frac{\sum^{L}_{\ell=1}\mathfrak{K}(k^{\ell},k)\cdot\mathfrak{K}(v^{\ell},v)}{\sum^{L}_{\ell=1}\mathfrak{K}(k^{\ell},k)}.$
(4.11)

A common choice of the kernel function is the Gaussian RBF kernel
$\mathfrak{K}_{\mathtt{RBF}}(q,k)=\exp(-\|q-k\|_{2}^{2}/2\gamma)$. We remark
that although the definition of the KDE in (4.11) involves the kernel function
$\mathfrak{K}(\cdot,\cdot)$, it is not associated with any RKHS. In what
follows, we normalize the query $q$, the key $k$, and the value $v$ so that
$q,k\in\mathbb{S}^{d_{\mathrm{p}}-1}$ and $v\in\mathbb{S}^{d-1}$, where
$\mathbb{S}^{d_{\mathrm{p}}-1}$ and $\mathbb{S}^{d-1}$ are the
$(d_{\mathrm{p}}-1)$-dimensional and $(d-1)$-dimensional unit spheres,
respectively. On the unit sphere, the exponential kernel is equivalent to the
Gaussian RBF kernel. Specifically, it holds for a given rescaling $\gamma>0$
that $\exp(q^{\top}k/\gamma)=C\cdot\exp(-\|q-k\|_{2}^{2}/2\gamma)$ for any
$q,k\in\mathbb{R}^{d_{\mathrm{p}}}$, where $C>0$ is an absolute constant. The
following proposition proves that the attention mechanism in (4.10) outputs
the conditional kernel density estimator in (4.11) and has the same limit as
the CME attention as $L\rightarrow\infty$.

######  Proposition 4.2 (Softmax Attention Converges to Kernel Conditional
Mean Embedding).

Recall that the softmax attention is defined in (4.10). It holds for any
$q\in\mathbb{S}^{d_{\mathrm{p}}-1}$ that

$\displaystyle\mathtt{attn}_{\mathtt{SM}}(q,K,V)=C\int_{\mathbb{S}^{d-1}}v\cdot\widehat{\mathbb{P}}_{\mathcal{V}{\,|\,}\mathcal{K}}^{\mathfrak{K}}(v{\,|\,}q){\mathrm{d}}v,$

where $C>0$ is an absolute constant. Meanwhile, under the condition that
$\widehat{\mathbb{P}}_{\mathcal{V}{\,|\,}\mathcal{K}}^{\mathfrak{K}}(v{\,|\,}k)\rightarrow\mathbb{P}_{\mathcal{V}{\,|\,}\mathcal{K}}(v{\,|\,}k)$
uniformly for any $k$ as $L\rightarrow\infty$, it holds for
$L\rightarrow\infty$ that

$\displaystyle\mathtt{attn}_{\mathtt{SM}}(q,K,V)\rightarrow
C\cdot\mathbb{E}[\mathcal{V}{\,|\,}\mathcal{K}=q].$

###### Proof.

See §B.5 for a detailed proof. ∎

We remark that the uniform convergence
$\widehat{\mathbb{P}}_{\mathcal{V}{\,|\,}\mathcal{K}}^{\mathfrak{K}}(v{\,|\,}k)\rightarrow\mathbb{P}_{\mathcal{V}{\,|\,}\mathcal{K}}(v{\,|\,}k)$
holds when the density of $\mathbb{P}_{\mathcal{K}}$ is bounded from below (De
Gooijer and Zerom, 2003). As shown in Propositions 4.1 and 4.2, the softmax
attention $\mathtt{attn}_{\mathtt{SM}}$ and the CME attention
$\mathtt{attn}_{\mathtt{CME}}$ have the same limit as $L\rightarrow\infty$.
Since the CME attention captures the latent posterior, which is proved in
§4.1, we conclude that the softmax attention also captures the latent
posterior approximately. Moreover, in terms of the limiting expectation
$\mathbb{E}[\mathcal{V}{\,|\,}\mathcal{K}=q]$, we highlight that it implies
the necessity of using the multiple heads and connects the attention mechanism
with causal inference. See §B.2 for a detailed discussion.

Figure 5: As shown in Propositions 4.1 and 4.2, the softmax attention
$\mathtt{attn}_{\mathtt{SM}}$ and the CME attention
$\mathtt{attn}_{\mathtt{CME}}$ have the same limit
$\mathbb{E}[\mathcal{V}{\,|\,}\mathcal{K}=q]$ as $L\rightarrow\infty$.

##  5 Excess Risk Analysis

To demonstrate the theoretical benefit of incorporating latent posterior
inference into the transformer architecture, we present a compact version of
excess risk analysis for one-layer single-head softmax attention neural
networks without skip connections. See §C for a detailed analysis of the
complete setup of the transformer architecture.

Attention Neural Network. We specify the feedforward neural network in (4.7)
as $\mathtt{nn}(x;A)={\mathtt{ReLU}}(Ax)$111Here, for ease of presentation, we
consider feedforward neural networks without bias terms., where
${\mathtt{ReLU}}(\cdot)$ is the rectified linear unit (ReLU) activation that
operates elementwise. In the rest of the paper, we consider the attention
neural networks with a final aggregation layer to allow for the proper scaling
of the outputs in the supervised setting and the transfer capability to
diverse downstream tasks in the self-supervised setting, which is discussed in
§6. We define the following function class of attention neural networks,

$\displaystyle\mathcal{F}_{\mathtt{attn}}=\Bigl{\\{}\mathtt{agg}_{\theta_{0}}\circ\mathtt{attn}_{\mathtt{SM}}\bigl{(}q_{\theta}({\mathtt{msk}}),k_{\theta}(X),v_{\theta}(X)\bigr{)}:\theta=(\theta_{0},A,W^{\mathrm{q}},W^{\mathrm{k}},W^{\mathrm{v}})\in\Theta\Bigr{\\}},$
(5.1)

where $\mathtt{agg}_{\theta_{0}}:\mathbb{R}^{d}\to\mathbb{R}^{d_{\mathrm{y}}}$
is the output layer parameterized by $\theta_{0}$ and
$\mathtt{attn}_{\mathtt{SM}}$ is the softmax attention defined in (4.10) with
the learnable parameters $(A,W^{\mathrm{q}},W^{\mathrm{k}},W^{\mathrm{v}})$
defined in (4.7). To characterize the excess risk, we specify the parameter
space as follows, which grants $\mathcal{F}_{\mathtt{attn}}$ a finite
capacity.

######  Assumption 5.1 (Parameter Space).

We assume for all
$\theta=(\theta_{0},A,W^{\mathrm{q}},W^{\mathrm{k}},W^{\mathrm{v}})\in\Theta$
that

$\displaystyle\|W^{{\mathrm{q}}}\|_{2}$ $\displaystyle\leq\omega,$
$\displaystyle\|W^{{\mathrm{k}}}\|_{2}$ $\displaystyle\leq\omega,$
$\displaystyle\|W^{{\mathrm{v}}}\|_{2}$ $\displaystyle\leq\omega,$
$\displaystyle\|A\|_{2}$ $\displaystyle\leq{\alpha},$
$\displaystyle\|W^{{\mathrm{q}}}\|_{\rm F}$ $\displaystyle\leq
R^{\mathtt{attn}},$ $\displaystyle\|W^{{\mathrm{k}}}\|_{\rm F}$
$\displaystyle\leq R^{\mathtt{attn}},$ $\displaystyle\|W^{{\mathrm{v}}}\|_{\rm
F}$ $\displaystyle\leq R^{\mathtt{attn}},$ $\displaystyle\|A\|_{\rm F}$
$\displaystyle\leq{R}^{\mathtt{nn}},$

where $\omega,\alpha,R^{\mathtt{attn}},R^{\mathtt{nn}}>0$ are absolute
constants.

Excess Risk. Following (4.8), we consider the learning objective
$\mathcal{L}((X,y),f)=\|y-f(X)\|_{2}^{2}$ for
$f\in\mathcal{F}_{\mathtt{attn}}$, where $y$ is the target variable. We make
the following assumption on the training dataset.

######  Assumption 5.2 (Data Distribution).

We assume that the training dataset
$\mathcal{D}_{n}=\\{(X_{i},y_{i})\\}_{i\in[n]}$ is independently and
identically drawn from the data distribution $\mathcal{D}$, which is supported
on the product space $\mathfrak{X}^{L}\times\mathfrak{Y}$, where

$\displaystyle\mathfrak{X}^{L}=\bigl{\\{}X\in\mathbb{R}^{L\times
d}:{\textstyle\max_{\ell\in[L]}}\|x^{\ell}\|_{2}\leq
R\bigr{\\}},\qquad\mathfrak{Y}=\bigl{\\{}y\in\mathbb{R}^{d_{\mathrm{y}}}:\|y\|_{2}\leq
1\bigr{\\}}.$ (5.2)

We consider the excess risk
$\mathcal{E}=\mathbb{E}[\mathcal{L}((X,y),\widehat{f})]-\mathbb{E}[\mathcal{L}((X,y),f^{*})]$,
where $\mathbb{E}[\,\cdot\,]$ is the population expectation over the data
distribution $\mathcal{D}$. Here $\widehat{f}\in\mathcal{F}_{\mathtt{attn}}$
is the attention neural network obtained from minimizing the empirical risk
$\widehat{\mathbb{E}}[\mathcal{L}((X,y),f)]$ in the training process, where
$\widehat{\mathbb{E}}[\,\cdot\,]$ is the empirical expectation over the
training dataset $\mathcal{D}_{n}$. Here $f^{*}(X)=\mathbb{E}[y{\,|\,}X]$ is
the regression function that we aim to approximate. In other words, $f^{*}$ is
the optimal model that minimizes the population risk
$\mathbb{E}[\mathcal{L}((X,y),f)]$.

To analyze the excess risk $\mathcal{E}$, we decompose it into three terms,

$\displaystyle\mathcal{E}$
$\displaystyle=\underbrace{\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),\widehat{f}\bigr{)}\Bigr{]}-\widehat{\mathbb{E}}\Bigl{[}\mathcal{L}\bigl{(}(X,y),\widehat{f}\bigr{)}\Bigr{]}+\widehat{\mathbb{E}}\Bigl{[}\mathcal{L}\bigl{(}(X,y),\widetilde{f}\bigr{)}\Bigr{]}-\min_{f\in\mathcal{F}_{\mathtt{attn}}}\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f\bigr{)}\Bigr{]}}_{\displaystyle\mathcal{E}_{\text{gen}}\text{:
Generalization Error}}$ (5.3)
$\displaystyle\qquad+\underbrace{\min_{f\in\mathcal{F}_{\mathtt{attn}}}\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f\bigr{)}\Bigr{]}-\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f^{*}\bigr{)}\Bigr{]}}_{\displaystyle\mathcal{E}_{\text{approx}}\text{:
Approximation
Error}}+\underbrace{\widehat{\mathbb{E}}\Bigl{[}\mathcal{L}\bigl{(}(X,y),\widehat{f}\bigr{)}\Bigr{]}-\widehat{\mathbb{E}}\Bigl{[}\mathcal{L}\bigl{(}(X,y),\widetilde{f}\bigr{)}\Bigr{]}}_{\displaystyle\mathcal{E}_{\text{opt}}\text{:
Optimization Error}}.$

where
$\widetilde{f}=\mathop{\mathrm{argmin}}_{f\in\mathcal{F}_{\mathtt{attn}}}\widehat{\mathbb{E}}[\mathcal{L}((X,y),f)]$
is the attention neural network that minimizes the empirical risk over
$\mathcal{F}_{\mathtt{attn}}$.

In §5.1-5.3, we analyze the three terms on the right-hand side of (5.3) in the
supervised setting. In §6, we extend the following analysis of the
approximation error $\mathcal{E}_{\text{approx}}$ to the self-supervised
setting.

###  5.1 Generalization Error Analysis

Recall that the softmax attention in (4.10) is instantiated via the
exponential kernel, which is equivalent to the Gaussian RBF kernel when $q$
and $k$ are on the unit sphere $\mathbb{S}^{d_{\mathrm{p}}-1}$. Also, note
that vector $\ell_{2}$-norm scales with the dimension $d_{\mathrm{p}}$ at the
rate of $\sqrt{d_{\mathrm{p}}}$. In the rest of the paper, we consider the
Gaussian RBF kernel with inputs rescaled by $1/\sqrt{d_{\mathrm{p}}}$, i.e.,

$\displaystyle\mathfrak{K}(q,k)=\exp\bigl{(}-\|q/\sqrt{d_{\mathrm{p}}}-k/\sqrt{d_{\mathrm{p}}}\|_{2}^{2}/2\bigr{)}=\exp\bigl{(}-\|q-k\|_{2}^{2}/2d_{\mathrm{p}}\bigr{)}.$
(5.4)

Under Assumption 5.1, we define $\gamma=\max\\{\alpha,\omega\\}$,
$\kappa=\max\\{R^{\mathtt{nn}},R^{\mathtt{attn}}/{\omega},{R}^{\mathtt{attn}}/{\omega}^{2}\\}$,
and $\zeta={\omega}R^{\mathtt{attn}}$. Recall that
$\mathcal{F}_{\mathtt{attn}}$ is the family of attention neural networks
defined in (5.1). We provide the following characterization of the
generalization error $\mathcal{E}_{\text{gen}}$.

######  Theorem 5.3 (Generalization Error).

Let $D=\max\\{d,d_{\mathrm{p}}\\}$. Suppose that Assumptions 5.1-5.2 hold and
$\mathtt{agg}_{\theta_{0}}$ is $1$-Lipschitz with the output range within
$\mathfrak{Y}$. Then, for any $\delta>0$, it holds with probability at least
$1-\delta$ that

$\displaystyle\mathcal{E}_{\text{gen}}=O\biggl{(}\frac{D}{\sqrt{n}}\cdot\bigl{[}\sqrt{\log(1+\gamma)}+\sqrt{\log(1+\zeta
R)}+\sqrt{\log(1+\kappa/\zeta)}\bigr{]}+\sqrt{\frac{\log(1/\delta)}{n}}\biggr{)}.$

###### Proof.

See §C for a detailed proof. ∎

An important implication of Theorem 5.3 is that the generalization error for
attention neural networks does not degrade as the sequence length $L$ goes to
infinity. It is also worth mentioning that the constants $\alpha$,
$\omega$,$R^{\mathtt{nn}}$, $R^{\mathtt{attn}}$, and $R$ play crucial roles in
the theoretical analysis of the generalization error and justify the
architecture design of the original transformer. In specific, we observe that
(i) skip connections help reducing $\alpha$, $\omega$, $R^{\mathtt{nn}}$, and
$R^{\mathtt{attn}}$, and (ii) layer normalizations help reducing $R$ when
there is multilayer composition of many attention mechanisms. See §C for a
more involved analysis of the generalization error of the complete setup of
the transformer architecture and the related discussion.

###  5.2 Approximation Error Analysis

In what follows, we characterize the approximation error in the supervised
setting. Approximation Target. We aim to approximate the regression function
$f^{*}(X)=\mathbb{E}[y{\,|\,}X]$ with the attention neural network
$f_{\theta}\in\mathcal{F}_{\mathtt{attn}}$, which is defined in (5.1). The
regression function $f^{*}(X)$ is the optimal model in the sense that it
minimizes the population risk $\mathbb{E}[\mathcal{L}((X,y),f)]$. By Lemma
3.2, when the input mask ${\mathtt{msk}}$ and the latent variable $z$ are
given, the target variable $y$ is independent of the input sequence $X$. Thus,
the regression function $f^{*}(X)$ can be decomposed as follows,

$\displaystyle f^{*}(X)=\mathbb{E}[y{\,|\,}X]$
$\displaystyle=\mathbb{E}_{z{\,|\,}X}\bigl{[}\mathbb{E}[y{\,|\,}{\mathtt{msk}},z]\bigr{]}$
$\displaystyle=\int\underbrace{\mathbb{E}[y{\,|\,}{\mathtt{msk}},z]}_{\displaystyle
g^{*}(z;{\mathtt{msk}}):\text{latent-to-
target}}\cdot\mathbb{P}(z{\,|\,}X){\mathrm{d}}z.$ (5.5)

Here the latent-to-target mapping $g^{*}(z;{\mathtt{msk}})$ can be viewed as a
decoding function, which maps the latent variable $z$ to the target variable
$y$ given the input mask ${\mathtt{msk}}$. On the other hand, the latent
posterior $\mathbb{P}(z\,|\,X)$ encodes the input sequence $X$ into the latent
variable $z$. We note that the input mask ${\mathtt{msk}}$ describes the
prediction task and is fixed throughout. For example, the input mask
${\mathtt{msk}}$ corresponds to the class encoding in the supervised setting
or the positional encoding in the self-supervised setting.

From (5.2), we see that approximating the regression function $f^{*}(X)$
involves capturing (i) the latent posterior $\mathbb{P}(z{\,|\,}X)$ and (ii)
the latent-to-target mapping $g^{*}(z;{\mathtt{msk}})$. Corresponding to (i),
the latent variable $z$ summarizes the “concept” of the input sequence $X$,
while corresponding to (ii), the target variable $y$ and the input mask
${\mathtt{msk}}$ specify the prediction task. In the following, we demonstrate
the central role of the latent-to-target mapping $g^{*}(z;{\mathtt{msk}})$,
which attention neural networks aim to approximate. Approximation Surrogate.
We define the reweighted CME attention

$\displaystyle
f^{\dagger}_{W}(X;{\mathtt{msk}})=W^{\top}\mathtt{attn}_{\mathtt{CME}}\bigl{(}q_{*}({\mathtt{msk}}),k_{*}(X),v_{*}(X)\bigr{)}$
(5.6)

as a surrogate function for approximating the regression function $f^{*}(X)$
in (5.2). Here the reweighting parameter $W\in\mathbb{R}^{d\times
d_{\mathrm{y}}}$ satisfies $\|W\|_{\rm F}<\infty$. In the sequel, we
demonstrate that the latent-to-target function contained in
$f^{\dagger}_{W}(X;{\mathtt{msk}})$ approximates the latent-to-target mapping
$g^{*}(z;{\mathtt{msk}})$, which is a key component of the regression function
$f^{*}(X)$. By (4.1) and (4.5), we have

$\displaystyle
W^{\top}\mathtt{attn}_{\mathtt{CME}}\bigl{(}q_{*}({\mathtt{msk}}),k_{*}(X),v_{*}(X)\bigr{)}$
$\displaystyle=W^{\top}\mathbb{E}[v^{\mathtt{msk}}\,|\,{\mathtt{msk}},X]$
$\displaystyle=W^{\top}\int\underbrace{\mathbb{E}[v^{\mathtt{msk}}\,|\,{\mathtt{msk}},z]}_{\displaystyle\psi(z;{\mathtt{msk}})\text{:
latent-to-value}}\cdot\mathbb{P}(z\,|\,X){\mathrm{d}}z,$ (5.7)

where $q_{*}({\mathtt{msk}})$ and $v^{\mathtt{msk}}$ replace
$c_{*}^{\mathtt{m}}({\mathtt{msk}})$ and $r^{\mathtt{msk}}$ in (4.2),
respectively. Taking (5.2) into (5.6), we obtain

$\displaystyle f^{\dagger}_{W}(X;{\mathtt{msk}})$
$\displaystyle=\int\underbrace{W^{\top}\psi(z;{\mathtt{msk}})}_{\displaystyle
g_{W}^{\dagger}(X;{\mathtt{msk}})}\cdot\mathbb{P}(z\,|\,X){\mathrm{d}}z=\mathbb{E}_{z{\,|\,}X}\bigl{[}g^{\dagger}_{W}(z;{\mathtt{msk}})\bigr{]},$
(5.8)

where $g^{\dagger}_{W}(z;{\mathtt{msk}})$ is a latent-to-target function
parameterized by $W\in\mathbb{R}^{d\times d_{\mathrm{y}}}$.

Following the infinite-dimensional counterpart of (3.6), the reweighted CME
attention captures the latent posterior $\mathbb{P}(z{\,|\,}X)$ under the
latent variable model in (4.1), where the latent prior is Gaussian. Comparing
(5.2) and (5.8), we see that the reweighted CME attention
$f_{W}^{\dagger}(X;{\mathtt{msk}})$ performs the latent-to-target decoding via
$g_{W}^{\dagger}(z;{\mathtt{msk}})$, which plays the same role as the latent-
to-target mapping $g^{*}(z;{\mathtt{msk}})$. Thus, it remains to characterize
the expressity of the function class

$\displaystyle\mathcal{G}^{\dagger}=\bigl{\\{}g^{\dagger}_{W}(z;{\mathtt{msk}})=W^{\top}\psi(z;{\mathtt{msk}}):W\in\mathbb{R}^{d\times
d_{\mathrm{y}}},\|W\|_{\rm F}<\infty\bigr{\\}}$ (5.9)

in terms of approximating the latent-to-target mapping
$g^{*}(z;{\mathtt{msk}})$ in (5.2).

To characterize the function class $\mathcal{G}^{\dagger}$ defined in (5.9),
we define the function class

$\displaystyle\mathcal{G}_{i}^{\dagger}=\\{g^{\dagger}_{W,i}(z;{\mathtt{msk}})=w_{i}^{\top}\psi(z;{\mathtt{msk}}):w_{i}\in\mathbb{R}^{d},\|w_{i}\|_{2}<\infty\bigr{\\}},$
(5.10)

which is formed by the $i$-th entry of the latent-to-target function
$g^{\dagger}_{W}(z;{\mathtt{msk}})\in\mathcal{G}^{\dagger}$. Here
$i\in[d_{\mathrm{y}}]$ and $W=[w_{1},\ldots,w_{d_{\mathrm{y}}}]^{\top}$. Note
that the function class $\mathcal{G}_{i}^{\dagger}$ is the RKHS
$\mathcal{H}_{\mathtt{LTV}}$ induced by the kernel function
$\mathfrak{K}_{{\mathtt{LTV}}}(z,z^{\prime};{\mathtt{msk}})=\psi(z;{\mathtt{msk}})^{\top}\psi(z^{\prime};{\mathtt{msk}})$,
which is a reproducing kernel. Here the latent-to-value (${\mathtt{LTV}}$)
mapping $\psi(z;{\mathtt{msk}})$ is defined in (5.2). See §E.1 for a detailed
discussion. See Figure 6 for a visualization of the construction of
$\mathcal{H}_{{\mathtt{LTV}}}$.

Figure 6: The RKHS $\mathcal{H}_{{\mathtt{LTV}}}$ induced by the latent-to-
value mapping $\psi(z;{\mathtt{msk}})$. The input mask ${\mathtt{msk}}$
describes the prediction task and determines the RKHS
$\mathcal{H}_{\mathtt{LTV}}$.

Therefore, the reweighted CME attention $f_{W}^{\dagger}(X;{\mathtt{msk}})$ in
(5.2) aims to capture the $i$-th entry $g^{*}_{i}(z;{\mathtt{msk}})$ of the
latent-to-target mapping $g^{*}(z;{\mathtt{msk}})$ within the RKHS
$\mathcal{H}_{{\mathtt{LTV}}}$. To this end, we make the following assumption
on the fundamental hardness of the recovery task.

######  Assumption 5.4 (Recovery Gap).

For any fixed input mask ${\mathtt{msk}}$, let

$\displaystyle
g^{\dagger}_{W,i}(\cdot;{\mathtt{msk}})=\Pi_{\mathcal{H}_{{\mathtt{LTV}}},\infty}\bigl{(}g^{*}_{i}(\cdot;{\mathtt{msk}})\bigr{)}=\mathop{\mathrm{argmin}}_{g_{i}(\cdot;{\mathtt{msk}})\in\mathcal{H}_{{\mathtt{LTV}}}}\bigl{\|}g^{*}_{i}(\cdot;{\mathtt{msk}})-g_{i}(\cdot;{\mathtt{msk}})\bigr{\|}_{\infty}$

be the $\ell_{\infty}$-norm projection of the $i$-th entry
$g^{*}_{i}(z;{\mathtt{msk}})$ of the latent-to-target mapping
$g^{*}(z;{\mathtt{msk}})$ onto the RKHS $\mathcal{H}_{{\mathtt{LTV}}}$. We
assume that there exists $\epsilon_{g}({\mathtt{msk}})\in[0,+\infty)$ such
that

$\displaystyle\sum_{i=1}^{d_{\mathrm{y}}}\bigl{\|}g_{i}^{*}(\cdot;{\mathtt{msk}})-g^{\dagger}_{W,i}(\cdot;{\mathtt{msk}})\bigr{\|}^{2}_{\infty}\leq\epsilon_{g}^{2}({\mathtt{msk}}).$

Here the $\ell_{\infty}$-norm is taken over the latent variable $z$.

Recall that the function class of attention neural networks
$\mathcal{F}_{\mathtt{attn}}$ is defined in (5.1). We have the following
theorem characterizing the approximation error $\mathcal{E}_{\rm approx}$
defined in (5.3).

######  Theorem 5.5 (Approximation Error).

Let
$\\{g^{\dagger}_{W,i}(z;{\mathtt{msk}})=w_{i}^{\top}\psi(z;{\mathtt{msk}})\\}_{i\in[d_{\mathrm{y}}]}$
be a function class satisfying Assumption 5.4. We define
$W=[w^{\top}_{1},\cdots,w^{\top}_{d_{\mathrm{y}}}]^{\top}$. Suppose that there
exists $f_{\theta}\in\mathcal{F}_{\mathtt{attn}}$ and
$\epsilon_{\mathtt{attn}}\in[0,+\infty)$ such that

$\displaystyle\sup_{X\in\mathfrak{X}^{L}}\Bigl{\|}f_{\theta}(X;{\mathtt{msk}})-W^{\top}\mathtt{attn}_{\mathtt{CME}}\bigl{(}q_{*}({\mathtt{msk}}),k_{*}(X),v_{*}(X)\bigr{)}\Bigr{\|}_{2}\leq\epsilon_{\mathtt{attn}},$
(5.11)

where $\mathfrak{X}^{L}$ is defined in (5.2). Then, we have

$\displaystyle\mathcal{E}_{\text{approx}}\leq
2\epsilon^{2}_{g}({\mathtt{msk}})+2\epsilon^{2}_{\mathtt{attn}}.$

###### Proof.

See §E.2 for a detailed proof. ∎

The approximation error bound in Theorem 5.5 involves the recovery gap
$\epsilon_{g}({\mathtt{msk}})$ and the surrogate approximation error
$\epsilon_{\mathtt{attn}}$. Since the latent posterior $\mathbb{P}(z{\,|\,}X)$
is captured by the reweighted CME attention, the recovery gap
$\epsilon_{g}({\mathtt{msk}})$ between the function class
$\mathcal{G}^{\dagger}$ in (5.9) and the latent-to-target mapping
$g^{*}(z;{\mathtt{msk}})$ in (5.2) plays the central role in the approximation
error bound. On the other hand, the approximation error
$\epsilon_{\mathtt{attn}}$ between attention neural networks in
$\mathcal{F}_{\mathtt{attn}}$ and the reweighted CME attention is
characterized in Proposition 4.2.

###  5.3 Optimization Error Analysis

Since the learning objective of attention neural networks is nonconvex with
respect to the parameter $\theta$, we consider the property of the stationary
points. Let
$\widehat{\theta}=(\widehat{\theta}_{0},\widehat{A},\widehat{W}^{\mathrm{q}},\widehat{W}^{\mathrm{k}},\widehat{W}^{\mathrm{v}})$
be the stationary point of the empirical risk
$\widehat{\mathbb{E}}[\mathcal{L}((X,y),f)]$, that is,

$\displaystyle\biggl{\langle}\nabla_{\theta}\widehat{\mathbb{E}}\Bigl{[}\mathcal{L}\bigl{(}(X,y),{f}_{\widehat{\theta}}\bigr{)}\Bigr{]},\theta-\widehat{\theta}\biggr{\rangle}\geq
0,\quad\forall\theta\in\Theta,$ (5.12)

which is the learnable parameter obtained in the training process, i.e.,
$\widehat{f}=f_{\widehat{\theta}}$. Recall that the regression function
$f^{*}(X)=\mathbb{E}[y{\,|\,}X]$ is the minimizer of the population risk
${\mathbb{E}}[\mathcal{L}((X,y),f)]$. We have the following proposition
characterizing the optimization error $\mathcal{E}_{\rm opt}$, which is
defined in (5.3).

######  Proposition 5.6 (Optimization Error).

Suppose that Assumption 5.2 holds. Then, it holds that

$\displaystyle\mathcal{E}_{\rm opt}\leq
4\cdot\min_{\theta\in\Theta}\widehat{\mathbb{E}}\Bigl{[}\bigl{\|}f_{\widehat{\theta}}(X)+\nabla_{\theta}f_{\widehat{\theta}}(X)^{\top}(\theta-\widehat{\theta})-f^{*}(X)\bigr{\|}_{2}\Bigr{]}.$
(5.13)

###### Proof.

See §D for a detailed proof. ∎

The right-hand side of (5.13) quantifies the expressity of the function class
defined by the local linearization,

$\displaystyle\bigl{\\{}f_{\widehat{\theta}}(X)+\nabla_{\theta}f_{\widehat{\theta}}(X)^{\top}(\theta-\widehat{\theta}):\theta\in\Theta\bigr{\\}}.$

In the neural tangent kernel (NTK) regime (Yang, 2020; Yang and Littwin, 2021;
Jacot et al., 2018), it is known that,

$\displaystyle
f^{*}(X)=f_{\widehat{\theta}}(X)+\nabla_{\theta}f_{\widehat{\theta}}(X)^{\top}(\theta-\widehat{\theta})+o(1),\quad\forall
X\in\mathbb{R}^{L\times d},$

where the $o(1)$ error captures the local linearization error in the NTK-based
analysis. As a consequence, the optimization error satisfies $\mathcal{E}_{\rm
opt}=o(1)$, that is, the stationary point $\widehat{\theta}$ is
(approximately) global optimal. Such a result shows the theoretical benefit of
incorporating feedforward neural networks in the architecture design. While
NTK-based analysis involves a random initialization in the supervised setting,
Malladi et al. (2022) provide an NTK-based analysis for the downstream
training of the transformer with a pretrained initialization in the self-
supervised setting.

##  6 From Supervised Learning to Self-Supervised Learning

An important aspect of the attention mechanism is that one can obtain a
sequence embedding by pretraining in a self-supervised manner, which gives
rise to the transfer capability for diverse downstream tasks. Self-Supervised
Learning. The attention mechanism enables embedding learning and downstream
prediction via the self-supervised learning (SSL) process as follows.

  * (${\mathtt{PT}}$)

Pretraining process: We train an attention neural network
$\widehat{f}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})=f_{\widehat{\theta}_{\mathtt{PT}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})\in\mathcal{F}_{\mathtt{PT}}$ with the
learned parameter $\widehat{\theta}_{\mathtt{PT}}$ to predict the masked token
$x^{L}\in\mathbb{R}^{d}$, which is denoted by $y_{\mathtt{PT}}$, from the
truncated input sequence $\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu=\\{x^{\ell}\\}_{\ell\in[L-1]}$ and the input mask
${\mathtt{msk}}_{\mathtt{PT}}$. Here the function class of attention neural
networks for the pretraining process is defined as follows,

$\displaystyle\mathcal{F}_{\mathtt{PT}}=\Bigl{\\{}\mathtt{agg}^{\mathtt{PT}}_{\theta}\circ\mathtt{attn}_{\mathtt{SM}}\bigl{(}q_{\theta}({\mathtt{msk}}_{\mathtt{PT}}),k_{\theta}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu),v_{\theta}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)\bigr{)}:\theta\in\Theta_{\mathtt{PT}}\Bigr{\\}},$ (6.1)

where $\mathtt{agg}^{\mathtt{PT}}_{\theta}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is
the aggregation layer. For the pretraining process, the input mask
${\mathtt{msk}}_{\mathtt{PT}}$ is the positional encoding of the masked token
$x^{L}$.

  * (${\mathtt{DS}}$)

Downstream task: We freeze the learned parameter
$\widehat{\theta}_{\mathtt{PT}}$ and train another attention neural network
$\widehat{f}_{\mathtt{DS}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})=f_{\widehat{\theta}_{\mathtt{DS}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})\in\mathcal{F}_{\mathtt{DS}}$ with the
learned parameter $\widehat{\theta}_{\mathtt{DS}}$ to predict another target
variable $y_{\mathtt{DS}}\in\mathbb{R}^{d_{\mathrm{y}}}$ from the truncated
input sequence $\mkern 1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu=\\{x^{\ell}\\}_{\ell\in[L-1]}$ and another input mask
${\mathtt{msk}}_{\mathtt{DS}}$. Here the function class of attention neural
networks for the downstream task is defined as follows,

$\displaystyle\mathcal{F}_{\mathtt{DS}}=\Bigl{\\{}\mathtt{agg}^{\mathtt{DS}}_{\theta}\circ\mathtt{attn}_{\mathtt{SM}}\bigl{(}q_{\widehat{\theta}_{\mathtt{PT}}}({\mathtt{msk}}_{\mathtt{DS}}),k_{\widehat{\theta}_{\mathtt{PT}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu),v_{\widehat{\theta}_{\mathtt{PT}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)\bigr{)}:\theta\in\Theta_{\mathtt{DS}}\Bigr{\\}},$ (6.2)

which means that the aggregation layer
$\mathtt{agg}^{\mathtt{DS}}_{\widehat{\theta}_{\mathtt{DS}}}:\mathbb{R}^{d}\to\mathbb{R}^{d_{\mathrm{y}}}$
replaces the aggregation layer
$\mathtt{agg}^{\mathtt{PT}}_{\widehat{\theta}_{\mathtt{PT}}}:\mathbb{R}^{d}\to\mathbb{R}^{d}$
obtained in the pretraining process. For the downstream task, the input mask
${\mathtt{msk}}_{\mathtt{DS}}$ is the class encoding of the target variable
$y_{\mathtt{DS}}$.

With the full input sequence $X$ replaced by the truncated input sequence
$\mkern 1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu$, the attention
neural network $\widehat{f}_{\mathtt{DS}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu;{\mathtt{msk}})$
obtained in the SSL process has the same decomposition of the excess risk as
that in (5.3). In the risk decomposition, for the SSL process, we have the
same characterization of the generalization error and the optimization error
as those in the supervised setting. When the downstream task is trained using
the same set of truncated input sequences as that in the pretraining process,
our previous analysis of the generalization error in the supervised setting is
applicable to the SSL process. On the other hand, when the downstream task is
trained using an independent set of truncated input sequences, we can modify
our previous analysis to prove that the generalization error only scales with
the complexity measure (e.g., the covering number) of the function class
$\\{\mathtt{agg}^{\mathtt{DS}}_{\theta}:\theta\in\Theta_{\mathtt{DS}}\\}$ of
aggregation layers without depending on that of the attention mechanism, as
$\widehat{\theta}_{\mathtt{PT}}$ is frozen. Also, the attention neural network
$\widehat{f}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu;{\mathtt{msk}})$
obtained in the pretraining process has the same approximation error as that
in the supervised setting. To characterize the approximation error for the SSL
process, we analyze the approximation error for the downstream task by
connecting it to the approximation error for the pretraining process.

Approximation Error. In parallel to the supervised setting, we define the
regression function and the latent-to-target mapping for the pretraining
process as follows,

$\displaystyle f_{\mathtt{PT}}^{*}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)=\mathbb{E}[y_{\mathtt{PT}}{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu],\qquad
g^{*}_{\mathtt{PT}}(z;{\mathtt{msk}}_{\mathtt{PT}})=\mathbb{E}[y_{\mathtt{PT}}{\,|\,}{\mathtt{msk}}_{\mathtt{PT}},z].$
(6.3)

Correspondingly, we defined the regression function and the latent-to-target
mapping for the downstream task as follows,

$\displaystyle f_{\mathtt{DS}}^{*}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)=\mathbb{E}[y_{\mathtt{DS}}{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu],\qquad
g_{\mathtt{DS}}^{*}(z;{\mathtt{msk}}_{\mathtt{DS}})=\mathbb{E}[y_{\mathtt{DS}}{\,|\,}{\mathtt{msk}}_{\mathtt{DS}},z].$
(6.4)

In parallel to the reweighted CME attention defined in (5.6), we defined the
surrogate functions for the pretraining process and the downstream task as
follows,

$\displaystyle f_{W_{\mathtt{PT}}}^{\dagger}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})$
$\displaystyle=W_{\mathtt{PT}}^{\top}\mathtt{attn}_{\mathtt{CME}}\bigl{(}q_{*}({\mathtt{msk}}_{\mathtt{PT}}),k_{*}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu),v_{*}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu)\bigr{)},$
$\displaystyle f_{W_{\mathtt{DS}}}^{\dagger}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})$
$\displaystyle=W_{\mathtt{DS}}^{\top}\mathtt{attn}_{\mathtt{CME}}\bigl{(}q_{*}({\mathtt{msk}}_{\mathtt{DS}}),k_{*}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu),v_{*}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu)\bigr{)},$ (6.5)

where $W_{\mathtt{PT}}\in\mathbb{R}^{d\times d}$ and
$W_{\mathtt{DS}}\in\mathbb{R}^{d\times d_{\mathrm{y}}}$ are the reweighting
parameters. We use the surrogate function to bridge the regression function
and the attention neural network, which is illustrated in Figure 7. In
parallel to the latent-to-value mapping $\psi(z;{\mathtt{msk}})$ defined in
(5.2), we define the latent-to-value mappings for the pretraining process and
the downstream task as follows,

$\displaystyle\psi_{\mathtt{PT}}(z;{\mathtt{msk}}_{\mathtt{PT}})=\mathbb{E}[v^{{\mathtt{msk}}_{\mathtt{PT}}}\,|\,{\mathtt{msk}}_{\mathtt{PT}},z],\qquad\psi_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})=\mathbb{E}[v^{{\mathtt{msk}}_{\mathtt{DS}}}\,|\,{\mathtt{msk}}_{\mathtt{DS}},z],$

where $v^{{\mathtt{msk}}_{\mathtt{PT}}}$ and
$v^{{\mathtt{msk}}_{\mathtt{DS}}}$ replace $r^{\mathtt{msk}}$ in (4.2). The
latent-to-value mappings induce the kernel functions as follows,

$\displaystyle\mathfrak{K}_{\mathtt{PT}}(z,z^{\prime};{\mathtt{msk}}_{\mathtt{PT}})$
$\displaystyle=\psi_{\mathtt{PT}}(z;{\mathtt{msk}}_{\mathtt{PT}})^{\top}\psi_{\mathtt{PT}}(z^{\prime};{\mathtt{msk}}_{\mathtt{PT}}),$
$\displaystyle\mathfrak{K}_{\mathtt{DS}}(z,z^{\prime};{\mathtt{msk}}_{\mathtt{DS}})$
$\displaystyle=\psi_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})^{\top}\psi_{\mathtt{DS}}(z^{\prime};{\mathtt{msk}}_{\mathtt{DS}}),$

which induce the RKHSs $\mathcal{H}_{\mathtt{PT}}$ and
$\mathcal{H}_{\mathtt{DS}}$. Corresponding to (5.8), we have

$\displaystyle f^{\dagger}_{W_{\mathtt{PT}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})$
$\displaystyle=\int\underbrace{W_{\mathtt{PT}}^{\top}\psi_{\mathtt{PT}}(z;{\mathtt{msk}}_{\mathtt{PT}})}_{\displaystyle
g_{W_{\mathtt{PT}}}^{\dagger}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})}\cdot\mathbb{P}(z\,|\,\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu){\mathrm{d}}z=\mathbb{E}_{z{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu}\bigl{[}g^{\dagger}_{W_{\mathtt{PT}}}(z;{\mathtt{msk}}_{\mathtt{PT}})\bigr{]},$
$\displaystyle f^{\dagger}_{W_{\mathtt{DS}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})$
$\displaystyle=\int\underbrace{W_{\mathtt{DS}}^{\top}\psi_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})}_{\displaystyle
g_{W_{\mathtt{DS}}}^{\dagger}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})}\cdot\mathbb{P}(z\,|\,\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu){\mathrm{d}}z=\mathbb{E}_{z{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu}\bigl{[}g^{\dagger}_{W_{\mathtt{DS}}}(z;{\mathtt{msk}}_{\mathtt{DS}})\bigr{]}.$
(6.6)

Note that $f^{\dagger}_{W_{\mathtt{PT}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})$ and $f^{\dagger}_{W_{\mathtt{DS}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})$ share the same latent posterior since the
attention mechanism is frozen for the downstream task. By our previous
arguments following (5.8)-(5.10), it remains to characterize how the
reweighted CME attentions in (6) recover the latent-to-target mappings in
(6.3)-(6.4) within the RKHSs $\mathcal{H}_{\mathtt{PT}}$ and
$\mathcal{H}_{\mathtt{DS}}$. See Figure 7 for an illustration of the
construction of the RKHSs $\mathcal{H}_{\mathtt{PT}}$ and
$\mathcal{H}_{\mathtt{DS}}$.

Figure 7: The RKHSs $\mathcal{H}_{\mathtt{PT}}$ and
$\mathcal{H}_{\mathtt{DS}}$ induced by the latent-to-value mappings
$\psi_{\mathtt{PT}}(z;{\mathtt{msk}}_{\mathtt{PT}})$ and
$\psi_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})$, respectively. The input
masks ${\mathtt{msk}}_{\mathtt{PT}}$ and ${\mathtt{msk}}_{\mathtt{DS}}$
describe the pretraining process and the downstream task, respectively, and
determine the RKHSs correspondingly. The $\ell_{\infty}$-norm projections
$\Pi_{\mathcal{H}_{\mathtt{PT}},\infty}$ and
$\Pi_{\mathcal{H}_{\mathtt{DS}},\infty}$ are defined in Assumption 6.1.

In parallel to Assumption 5.4, we introduce the following assumption on the
fundamental hardness of approximating the latent-to-target mappings within the
RKHSs $\mathcal{H}_{\mathtt{PT}}$ and $\mathcal{H}_{\mathtt{DS}}$.

######  Assumption 6.1 (SSL Recovery Gap).

For any fixed input masks ${\mathtt{msk}}_{\mathtt{PT}}$ and
${\mathtt{msk}}_{\mathtt{DS}}$, let

$\displaystyle
g^{\dagger}_{W_{\mathtt{PT}},i}(\cdot;{\mathtt{msk}}_{\mathtt{PT}})$
$\displaystyle=\Pi_{\mathcal{H}_{{\mathtt{PT}}},\infty}\bigl{(}g^{*}_{{\mathtt{PT}},i}(\cdot;{\mathtt{msk}}_{\mathtt{PT}})\bigr{)}=\mathop{\mathrm{argmin}}_{g_{i}(\cdot;{\mathtt{msk}}_{\mathtt{PT}})\in\mathcal{H}_{{\mathtt{PT}}}}\bigl{\|}g^{*}_{{\mathtt{PT}},i}(\cdot;{\mathtt{msk}}_{\mathtt{PT}})-g_{i}(\cdot;{\mathtt{msk}}_{\mathtt{PT}})\bigr{\|}_{\infty},$
$\displaystyle
g^{\dagger}_{W_{\mathtt{DS}},i}(\cdot;{\mathtt{msk}}_{\mathtt{DS}})$
$\displaystyle=\Pi_{\mathcal{H}_{{\mathtt{DS}}},\infty}\bigl{(}g^{*}_{{\mathtt{DS}},i}(\cdot;{\mathtt{msk}}_{\mathtt{DS}})\bigr{)}=\mathop{\mathrm{argmin}}_{g_{i}(\cdot;{\mathtt{msk}}_{\mathtt{DS}})\in\mathcal{H}_{{\mathtt{DS}}}}\bigl{\|}g^{*}_{{\mathtt{DS}},i}(\cdot;{\mathtt{msk}}_{\mathtt{DS}})-g_{i}(\cdot;{\mathtt{msk}}_{\mathtt{DS}})\bigr{\|}_{\infty},$
$\displaystyle
g^{\dagger}_{W_{\mathtt{SSL}},i}(\cdot;{\mathtt{msk}}_{\mathtt{DS}})$
$\displaystyle=\Pi_{\mathcal{H}_{{\mathtt{DS}}},\infty}\bigl{(}g^{*}_{{\mathtt{PT}},i}(\cdot;{\mathtt{msk}}_{\mathtt{PT}})\bigr{)}=\mathop{\mathrm{argmin}}_{g_{i}(\cdot;{\mathtt{msk}}_{\mathtt{DS}})\in\mathcal{H}_{{\mathtt{DS}}}}\bigl{\|}g^{*}_{{\mathtt{PT}},i}(\cdot;{\mathtt{msk}}_{\mathtt{PT}})-g_{i}(\cdot;{\mathtt{msk}}_{\mathtt{DS}})\bigr{\|}_{\infty}$

be the $\ell_{\infty}$-norm projections of the $i$-th entry
$g^{*}_{{\mathtt{PT}},i}(z;{\mathtt{msk}}_{\mathtt{PT}})$ of the latent-to-
target mapping $g^{*}_{\mathtt{PT}}(z;{\mathtt{msk}}_{\mathtt{PT}})$ onto the
RKHS $\mathcal{H}_{{\mathtt{PT}}}$, the $i$-th entry
$g^{*}_{{\mathtt{DS}},i}(z;{\mathtt{msk}}_{\mathtt{DS}})$ of the latent-to-
target mapping $g^{*}_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})$ onto the
RKHS $\mathcal{H}_{{\mathtt{DS}}}$, and the $i$-th entry
$g^{*}_{{\mathtt{PT}},i}(z;{\mathtt{msk}}_{\mathtt{PT}})$ of the latent-to-
target mapping $g^{*}_{\mathtt{PT}}(z;{\mathtt{msk}}_{\mathtt{PT}})$ onto the
RKHS $\mathcal{H}_{{\mathtt{DS}}}$, respectively. We assume the following
statements hold.

  * (${\mathtt{PT}}$)

There exists $\epsilon_{g}({\mathtt{msk}}_{\mathtt{PT}})\in[0,+\infty)$ such
that

$\displaystyle\sum_{i=1}^{d}\bigl{\|}g_{{\mathtt{PT}},i}^{*}(\cdot;{\mathtt{msk}}_{\mathtt{PT}})-g^{\dagger}_{W_{\mathtt{PT}},i}(\cdot;{\mathtt{msk}}_{\mathtt{PT}})\bigr{\|}^{2}_{\infty}\leq\epsilon_{g}^{2}({\mathtt{msk}}_{\mathtt{PT}}).$
(6.7)

  * (${\mathtt{DS}}$)

There exists $\epsilon_{g}({\mathtt{msk}}_{\mathtt{DS}})\in[0,+\infty)$ such
that

$\displaystyle\sum_{i=1}^{d_{\mathrm{y}}}\bigl{\|}g_{{\mathtt{DS}},i}^{*}(\cdot;{\mathtt{msk}}_{\mathtt{DS}})-g^{\dagger}_{W_{\mathtt{DS}},i}(\cdot;{\mathtt{msk}}_{\mathtt{DS}})\bigr{\|}^{2}_{\infty}\leq\epsilon_{g}^{2}({\mathtt{msk}}_{\mathtt{DS}}).$
(6.8)

  * (${\mathtt{SSL}}$)

There exists
$\epsilon_{\mathtt{SSL}}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}})\in[0,+\infty)$
such that

$\displaystyle\sum_{i=1}^{d}\bigl{\|}g_{{\mathtt{PT}},i}^{*}(\cdot;{\mathtt{msk}}_{\mathtt{PT}})-g^{\dagger}_{W_{\mathtt{SSL}},i}(\cdot;{\mathtt{msk}}_{\mathtt{DS}})\bigr{\|}^{2}_{\infty}\leq\epsilon_{\mathtt{SSL}}^{2}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}}).$
(6.9)

Here the $\ell_{\infty}$-norms are taken over the latent variable $z$.

Intuitively, the feature $\psi_{\mathtt{PT}}(z;{\mathtt{msk}}_{\mathtt{PT}})$
is obtained in the pretraining process, while the feature
$\psi_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})$ is desired by the
downstream task. Meanwhile, (6.9) characterizes the fundamental hardness of
recovering the latent-to-target mapping
$g_{\mathtt{PT}}^{*}(z;{\mathtt{msk}}_{\mathtt{PT}})$ for the pretraining
process within the RKHS $\mathcal{H}_{\mathtt{DS}}$. Thus, the transfer error
$\epsilon_{\mathtt{SSL}}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}})$
captures the transfer capability of the sequence embedding obtained in the
pretraining process to the downstream task. In other words, when the
pretraining process is sufficiently related to the downstream task, the
transfer error
$\epsilon_{\mathtt{SSL}}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}})$
is small, which allows us to approximate the $i$-th entry of the latent-to-
target mapping $g^{*}_{\mathtt{PT}}(z;{\mathtt{msk}}_{\mathtt{PT}})$ within
the RKHS $\mathcal{H}_{\mathtt{DS}}$ up to the approximation error
$\epsilon_{\mathtt{SSL}}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}})$.

We introduce the following assumption on the condition number that
characterizes the alignment between the reweighting parameter desired by the
downstream task and the reweighting parameter obtained in the pretraining
process.

######  Assumption 6.2 (SSL Condition Number).

Let
$\\{g_{W_{\mathtt{DS}},i}^{\dagger}(z;{\mathtt{msk}}_{\mathtt{DS}})=w_{{\mathtt{DS}},i}^{\top}\psi_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})\\}_{i\in[d]}$
and
$\\{g^{\dagger}_{W_{\mathtt{SSL}},i}(z;{\mathtt{msk}}_{\mathtt{DS}})=w_{{\mathtt{SSL}},i}^{\top}\psi_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})\\}_{i\in[d_{\mathrm{y}}]}$
be the function classes satisfying (6.8) and (6.9) in Assumption 6.1,
respectively. Also, let
$W_{\mathtt{DS}}=[w_{{\mathtt{DS}},1},\ldots,w_{{\mathtt{DS}},d_{\mathrm{y}}}]^{\top}\in\mathbb{R}^{d\times
d_{\mathrm{y}}}$,
$W_{\mathtt{SSL}}=[w_{{\mathtt{SSL}},1},\ldots,w_{{\mathtt{SSL}},d}]^{\top}\in\mathbb{R}^{d\times
d}$, and222For ease of presentation, we assume that
$W_{\mathtt{SSL}}W^{\top}_{\mathtt{SSL}}\in\mathbb{R}^{d\times d}$ is
invertible. When $W_{\mathtt{SSL}}W^{\top}_{\mathtt{SSL}}$ is not invertible,
our subsequent analysis can be generalized using the pseudoinverse of
$W_{\mathtt{SSL}}W^{\top}_{\mathtt{SSL}}$.

$\displaystyle
B=W_{\mathtt{DS}}^{\top}(W_{\mathtt{SSL}}W^{\top}_{\mathtt{SSL}})^{-1}W_{\mathtt{SSL}}\in\mathbb{R}^{d_{\mathrm{y}}\times
d}.$ (6.10)

We assume that there exists $\mu\in[0,+\infty)$ such that
$\|B\|_{2}^{2}\leq\mu$.

The condition number $\mu$ plays a critical role in our subsequent analysis.
To see the intuition behind $\mu$, let $W_{\mathtt{DS}}=W_{\mathtt{SSL}}$,
which implies that $B$ is a projection matrix and $\mu=1$. Also, let the row
vectors of $W_{\mathtt{SSL}}$ be an orthonormal basis of $\mathbb{R}^{d}$,
which implies that
$W_{\mathtt{SSL}}W_{\mathtt{SSL}}^{\top}=I_{d_{\mathrm{p}}}$ and
$B=W_{\mathtt{DS}}^{\top}W_{\mathtt{SSL}}$. In this case, $B$ measures the
subspace alignment between the reweighting parameter $W_{\mathtt{DS}}$ desired
by the downstream task and the reweighting parameter $W_{\mathtt{SSL}}$
obtained in the pretraining process. In general cases where $W_{\mathtt{SSL}}$
is nonorthonormal, we have a similar interpretation through the eigenvalue
decomposition of $W_{\mathtt{SSL}}W_{\mathtt{SSL}}^{\top}$.

Recall that $\widehat{f}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})$ is the attention neural network obtained
in the pretraining process. For any $U\in\mathbb{R}^{d_{\mathrm{y}}\times d}$,
we define the following quantity that characterizes the expressity of the
function class
$\\{\mathtt{agg}^{\mathtt{DS}}_{\theta}:\theta\in\Theta_{\mathtt{DS}}\\}$ of
aggregation layers for the downstream task,

$\displaystyle\epsilon_{\mathtt{agg}}(U)$
$\displaystyle=\inf_{f_{\mathtt{DS}}\in\mathcal{F}_{\mathtt{DS}}}\sup_{\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu\in\mathfrak{X}^{L-1}}\bigl{\|}f_{\mathtt{DS}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})-U\widehat{f}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})\bigr{\|}_{2}$
$\displaystyle=\inf_{\theta\in\Theta_{\mathtt{DS}}}\sup_{\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu\in\mathfrak{X}^{L-1}}\bigl{\|}\mathtt{agg}^{\mathtt{DS}}_{\theta}\circ\mathtt{attn}_{\mathtt{SM}}\bigl{(}q_{\widehat{\theta}_{\mathtt{PT}}}({\mathtt{msk}}_{\mathtt{DS}}),k_{\widehat{\theta}_{\mathtt{PT}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu),v_{\widehat{\theta}_{\mathtt{PT}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu)\bigr{)}$
$\displaystyle\quad\qquad\qquad\qquad-U\mathtt{agg}^{\mathtt{PT}}_{\widehat{\theta}_{\mathtt{PT}}}\circ\mathtt{attn}_{\mathtt{SM}}\bigl{(}q_{\widehat{\theta}_{\mathtt{PT}}}({\mathtt{msk}}_{\mathtt{PT}}),k_{\widehat{\theta}_{\mathtt{PT}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu),v_{\widehat{\theta}_{\mathtt{PT}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu)\bigr{)}\bigr{\|}_{2},$
(6.11)

where $\mathfrak{X}^{L-1}$ is defined in (5.2). Since the attention mechanism
is frozen for the downstream task, the trainable part of the attention neural
network is the aggregation layer $\mathtt{agg}^{\mathtt{DS}}_{\theta}$. Thus,
the aggregation approximation error $\epsilon_{\mathtt{agg}}(U)$ characterizes
the expressity of the function class
$\\{\mathtt{agg}^{\mathtt{DS}}_{\theta}:\theta\in\Theta_{\mathtt{DS}}\\}$ of
aggregation layers in terms of approximating the composition of (i) the linear
transformation $U\mathtt{agg}^{\mathtt{PT}}_{\widehat{\theta}_{\mathtt{PT}}}$
of the aggregation layer obtained in the pretraining process and (ii) the
output variation induced by switching the input mask
${\mathtt{msk}}_{\mathtt{PT}}$ to another input mask
${\mathtt{msk}}_{\mathtt{DS}}$ in the attention mechanism, which is frozen. To
see the intuition behind $\epsilon_{\mathtt{agg}}(U)$, let
${\mathtt{msk}}_{\mathtt{PT}}={\mathtt{msk}}_{\mathtt{DS}}$, which implies
that $\epsilon_{\mathtt{agg}}(U)=0$ as long as the function class of
aggregation layers takes the form of
$\\{\mathtt{agg}^{\mathtt{DS}}_{\theta}=U\mathtt{agg}^{\mathtt{PT}}_{\widehat{\theta}_{\mathtt{PT}}}:\theta=U\in\mathbb{R}^{d_{\mathrm{y}}\times
d}\\}$. In this case, $\epsilon_{\mathtt{agg}}(U)$ characterizes the
compatibility between $\mathtt{agg}_{\theta}^{\mathtt{DS}}$ and
$\mathtt{agg}^{\mathtt{PT}}_{\widehat{\theta}_{\mathtt{PT}}}$ under a linear
transformation parameterized by $\theta$. In general cases where
${\mathtt{msk}}_{\mathtt{PT}}\neq{\mathtt{msk}}_{\mathtt{DS}}$,
$\epsilon_{\mathtt{agg}}(U)$ additionally characterizes the capability of
$\mathtt{agg}_{\theta}^{\mathtt{DS}}$ to capture the output variation induced
by switching the input mask.

Recall that the function class $\mathcal{F}_{\mathtt{PT}}$ of attention neural
networks for the pretraining process is defined in (6.1). Let

$\displaystyle\mathcal{E}_{\rm
approx}^{\mathtt{PT}}=\min_{f\in\mathcal{F}_{\mathtt{PT}}}\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu,y_{\mathtt{PT}}),f\bigr{)}\Bigr{]}-\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu,y_{\mathtt{PT}}),f_{\mathtt{PT}}^{*}\bigr{)}\Bigr{]}$ (6.12)

be the approximation error for the pretraining process, which is characterized
in Theorem 5.5. Recall that the function class $\mathcal{F}_{\mathtt{DS}}$ of
attention neural networks for the downstream task is defined in (6.2). For the
downstream task, the approximation error in (5.3) takes the following form,

$\displaystyle\mathcal{E}_{\rm
approx}=\min_{f\in\mathcal{F}_{\mathtt{DS}}}\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu,y_{\mathtt{DS}}),f\bigr{)}\Bigr{]}-\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu,y_{\mathtt{DS}}),f_{\mathtt{DS}}^{*}\bigr{)}\Bigr{]}.$ (6.13)

The following theorem characterizes the approximation error $\mathcal{E}_{\rm
approx}$ for the SSL process.

######  Theorem 6.3 (SSL Approximation Error).

Under Assumptions 6.1 and 6.2, it holds that

$\displaystyle\mathcal{E}_{\text{approx}}=O\Bigl{(}\mu\cdot\bigl{(}\mathcal{E}_{\rm
approx}^{{\mathtt{PT}}}+\epsilon_{\mathtt{SSL}}^{2}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}})\bigr{)}+\epsilon_{g}^{2}({\mathtt{msk}}_{\mathtt{DS}})+\epsilon_{\mathtt{agg}}^{2}(B)\Bigr{)},$

where $\mathcal{E}^{\mathtt{PT}}_{\rm approx}$,
$\epsilon_{\mathtt{SSL}}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}})$,
$\epsilon_{g}({\mathtt{msk}}_{\mathtt{DS}})$, and $\epsilon_{\mathtt{agg}}(B)$
are defined in (6.12), (6.9), (6.8), and (6), respectively, and $B$ is defined
in (6.10).

###### Proof.

See §E.3 for a detailed proof. ∎

Theorem 6.3 demonstrates that the attention neural network enables the
transfer capability to diverse downstream tasks, where the approximation error
is subsume from that in the supervised setting with a few extra error terms.
We interpret the approximation error in Theorem 6.3 as follows.

  * (i)

The condition number $\mu$ characterizes the the alignment between the
reweighting parameter $W_{\mathtt{DS}}$ desired by the downstream task and the
reweighting parameter $W_{\mathtt{SSL}}$ obtained in the pretraining process.
When $W_{\mathtt{DS}}=W_{\mathtt{SSL}}$, we have $\mu=1$.

  * (ii)

The approximation error $\mathcal{E}_{\rm approx}^{\mathtt{PT}}$ for the
pretraining process is characterized in Theorem 5.5. Specifically,
$\mathcal{E}_{\rm approx}^{\mathtt{PT}}$ involves the pretraining recovery gap
$\epsilon_{g}({\mathtt{msk}}_{\mathtt{PT}})$ defined in (6.7), which
characterizes the fundamental hardness of approximating the $i$-th entry
$g_{{\mathtt{PT}},i}^{*}(z;{\mathtt{msk}}_{\mathtt{PT}})$ of the latent-to-
target mapping $g_{\mathtt{PT}}^{*}(z;{\mathtt{msk}}_{\mathtt{PT}})$ defined
in (6.3) within the RKHS $\mathcal{H}_{\mathtt{PT}}$, and the attention
approximation error $\epsilon_{\mathtt{attn}}$ defined in (5.11), which is
characterized in Proposition 4.2.

  * (iii)

The transfer error
$\epsilon_{\mathtt{SSL}}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}})$
captures the transfer capability of the sequence embedding obtained in the
pretraining process to the downstream task. By our previous arguments
following Assumption 6.1,
$\epsilon_{\mathtt{SSL}}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}})$
is small as long as the pretraining process is sufficiently related to the
downstream task.

  * (iv)

The downstream recovery gap $\epsilon_{g}({\mathtt{msk}}_{\mathtt{DS}})$
defined in (6.8) characterizes the fundamental hardness of approximating the
$i$-th entry $g_{{\mathtt{DS}},i}^{*}(z;{\mathtt{msk}}_{\mathtt{DS}})$ of the
latent-to-target mapping $g^{*}_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})$
defined in (6.4) within the RKHS $\mathcal{H}_{\mathtt{DS}}$.

  * (v)

The aggregation approximation error $\epsilon_{\mathtt{agg}}(B)$ measures the
expressity of the function class of aggregation layers for the downstream
task. By our previous arguments following (6), $\epsilon_{\mathtt{agg}}(U)$ is
small as long as the aggregation layer $\mathtt{agg}^{\mathtt{DS}}_{\theta}$
for the downstream task can approximate the composition of the linear
transformation $B\mathtt{agg}^{\mathtt{PT}}_{\widehat{\theta}_{\mathtt{PT}}}$
of the aggregation layer obtained in the pretraining process and the variation
induced by switching the input mask ${\mathtt{msk}}_{\mathtt{PT}}$ to another
input mask ${\mathtt{msk}}_{\mathtt{DS}}$.

##  7 Conclusion

The attention mechanism in transformers demonstrates significant empirical
successes in natural language processing and computer vision, but there is a
lack of understanding about how it works and why it is effective. To this end,
we answer three questions about the attention mechanism: (i) what makes a good
representation, (ii) how the attention mechanism produces this representation
during the forward pass, and (iii) how the attention mechanism learns to
produce this representation during the backward pass. Through the lens of
exchangeability, we provide a theoretical characterization of the attention
mechanism as a combination of a “whitebox” design guided by a latent variable
model and a “blackbox” design allowing for learnable parameters. Also, we
establish the approximation, generalization, and optimization guarantees of
the attention mechanism. Several challenging questions remain open, e.g., what
is the theoretical benefit of a multilayer composition of many attention
mechanisms and how it affects the approximation guarantee.

## References

  * Allen-Zhu et al. (2019a) Allen-Zhu, Z., Li, Y. and Liang, Y. (2019a).  Learning and generalization in overparameterized neural networks, going beyond two layers.  In Neural Information Processing Systems.
  * Allen-Zhu et al. (2019b) Allen-Zhu, Z., Li, Y. and Song, Z. (2019b).  A convergence theory for deep learning via over-parameterization.  In International Conference on Machine Learning.
  * Allen-Zhu et al. (2019c) Allen-Zhu, Z., Li, Y. and Song, Z. (2019c).  On the convergence rate of training recurrent neural networks.  In Neural Information Processing Systems.
  * Arora et al. (2019) Arora, S., Du, S., Hu, W., Li, Z. and Wang, R. (2019).  Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.  In International Conference on Machine Learning.
  * Bartlett (1996) Bartlett, P. (1996).  For valid generalization the size of the weights is more important than the size of the network.  Neural Information Processing Systems.
  * Bartlett et al. (2018a) Bartlett, P., Helmbold, D. and Long, P. (2018a).  Gradient descent with identity initialization efficiently learns positive definite linear transformations by deep residual networks.  In International Conference on Machine Learning.
  * Bartlett et al. (2018b) Bartlett, P. L., Evans, S. N. and Long, P. M. (2018b).  Representing smooth functions as compositions of near-identity functions with implications for deep network optimization.  arXiv preprint arXiv:1804.05012.
  * Bartlett et al. (2017) Bartlett, P. L., Foster, D. J. and Telgarsky, M. J. (2017).  Spectrally-normalized margin bounds for neural networks.  Neural Information Processing Systems.
  * Bartlett and Mendelson (2002) Bartlett, P. L. and Mendelson, S. (2002).  Rademacher and Gaussian complexities: Risk bounds and structural results.  Journal of Machine Learning Research.
  * Battaglia et al. (2018) Battaglia, P. W., Hamrick, J. B., Bapst, V., Sanchez-Gonzalez, A., Zambaldi, V., Malinowski, M., Tacchetti, A., Raposo, D., Santoro, A., Faulkner, R. et al. (2018).  Relational inductive biases, deep learning, and graph networks.  arXiv preprint arXiv:1806.01261.
  * Bloem-Reddy and Teh (2020) Bloem-Reddy, B. and Teh, Y. W. (2020).  Probabilistic symmetries and invariant neural networks.  Journal of Machine Learning Research.
  * Bronstein et al. (2021) Bronstein, M. M., Bruna, J., Cohen, T. and Veličković, P. (2021).  Geometric deep learning: Grids, groups, graphs, geodesics, and gauges.  arXiv preprint arXiv:2104.13478.
  * Brown et al. (2020) Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A. et al. (2020).  Language models are few-shot learners.  Neural Information Processing Systems.
  * Cao and Gu (2019) Cao, Y. and Gu, Q. (2019).  Generalization bounds of stochastic gradient descent for wide and deep neural networks.  Neural Information Processing Systems.
  * Caponnetto and De Vito (2007) Caponnetto, A. and De Vito, E. (2007).  Optimal rates for the regularized least-squares algorithm.  Foundations of Computational Mathematics.
  * Chen et al. (2021) Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A. and Mordatch, I. (2021).  Decision transformer: Reinforcement learning via sequence modeling.  Neural Information Processing Systems.
  * Chizat et al. (2019) Chizat, L., Oyallon, E. and Bach, F. (2019).  On lazy training in differentiable programming.  Neural Information Processing Systems.
  * Dai et al. (2019) Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q. V. and Salakhutdinov, R. (2019).  Transformer-XL: Attentive language models beyond a fixed-length context.  arXiv preprint arXiv:1901.02860.
  * de Finetti (1937) de Finetti, B. (1937).  La prévision: ses lois logiques, ses sources subjectives.  In Annales de l’institut Henri Poincaré.
  * De Gooijer and Zerom (2003) De Gooijer, J. G. and Zerom, D. (2003).  On conditional density estimation.  Statistica Neerlandica.
  * Devlin et al. (2018) Devlin, J., Chang, M.-W., Lee, K. and Toutanova, K. (2018).  BERT: Pre-training of deep bidirectional transformers for language understanding.  arXiv preprint arXiv:1810.04805.
  * Diaconis and Freedman (1980) Diaconis, P. and Freedman, D. (1980).  Finite exchangeable sequences.  Annals of Probability.
  * Dosovitskiy et al. (2020) Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S. et al. (2020).  An image is worth 16$\times$16 words: Transformers for image recognition at scale.  arXiv preprint arXiv:2010.11929.
  * Du et al. (2019) Du, S., Lee, J., Li, H., Wang, L. and Zhai, X. (2019).  Gradient descent finds global minima of deep neural networks.  In International Conference on Machine Learning.
  * Du et al. (2018) Du, S. S., Zhai, X., Poczos, B. and Singh, A. (2018).  Gradient descent provably optimizes over-parameterized neural networks.  arXiv preprint arXiv:1810.02054.
  * Edelman et al. (2021) Edelman, B. L., Goel, S., Kakade, S. and Zhang, C. (2021).  Inductive biases and variable creation in self-attention mechanisms.  arXiv preprint arXiv:2110.10090.
  * Elesedy (2021) Elesedy, B. (2021).  Provably strict generalisation benefit for invariance in kernel methods.  Neural Information Processing Systems.
  * Fisher (1922) Fisher, R. A. (1922).  On the mathematical foundations of theoretical statistics.  Philosophical Transactions of the Royal Society A, 222 309–368.
  * Garg et al. (2022) Garg, S., Tsipras, D., Liang, P. and Valiant, G. (2022).  What can transformers learn in-context? A case study of simple function classes.  arXiv preprint arXiv:2208.01066.
  * Han et al. (2022) Han, J., Rong, Y., Xu, T. and Huang, W. (2022).  Geometrically equivariant graph neural networks: A survey.  arXiv preprint arXiv:2202.07230.
  * Hardt and Ma (2016) Hardt, M. and Ma, T. (2016).  Identity matters in deep learning.  arXiv preprint arXiv:1611.04231.
  * He et al. (2022) He, K., Chen, X., Xie, S., Li, Y., Dollár, P. and Girshick, R. (2022).  Masked autoencoders are scalable vision learners.  In Computer Vision and Pattern Recognition.
  * Hron et al. (2020) Hron, J., Bahri, Y., Sohl-Dickstein, J. and Novak, R. (2020).  Infinite attention: NNGP and NTK for deep attention networks.  In International Conference on Machine Learning.
  * Hutchinson et al. (2021) Hutchinson, M. J., Le Lan, C., Zaidi, S., Dupont, E., Teh, Y. W. and Kim, H. (2021).  Lietransformer: Equivariant self-attention for Lie groups.  In International Conference on Machine Learning.
  * Jacot et al. (2018) Jacot, A., Gabriel, F. and Hongler, C. (2018).  Neural tangent kernel: Convergence and generalization in neural networks.  Neural Information Processing Systems.
  * Jiang et al. (2019) Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D. and Bengio, S. (2019).  Fantastic generalization measures and where to find them.  arXiv preprint arXiv:1912.02178.
  * Jumper et al. (2021) Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Žídek, A., Potapenko, A. et al. (2021).  Highly accurate protein structure prediction with AlphaFold.  Nature.
  * Keriven and Peyré (2019) Keriven, N. and Peyré, G. (2019).  Universal invariant and equivariant graph neural networks.  Neural Information Processing Systems.
  * Kossen et al. (2021) Kossen, J., Band, N., Lyle, C., Gomez, A. N., Rainforth, T. and Gal, Y. (2021).  Self-attention between datapoints: Going beyond individual input-output pairs in deep learning.  Neural Information Processing Systems.
  * Lee et al. (2019) Lee, J., Lee, Y., Kim, J., Kosiorek, A., Choi, S. and Teh, Y. W. (2019).  Set transformer: A framework for attention-based permutation-invariant neural networks.  In International Conference on Machine Learning.
  * Li and Liang (2018) Li, Y. and Liang, Y. (2018).  Learning overparameterized neural networks via stochastic gradient descent on structured data.  Neural Information Processing Systems.
  * Malladi et al. (2022) Malladi, S., Wettig, A., Yu, D., Chen, D. and Arora, S. (2022).  A kernel-based view of language model fine-tuning.  arXiv preprint arXiv:2210.05643.
  * Mei et al. (2019) Mei, S., Misiakiewicz, T. and Montanari, A. (2019).  Mean-field theory of two-layers neural networks: Dimension-free bounds and kernel limit.  In Annual Conference on Learning Theory.
  * Mei et al. (2018) Mei, S., Montanari, A. and Nguyen, P.-M. (2018).  A mean field view of the landscape of two-layer neural networks.  Proceedings of the National Academy of Sciences.
  * Mohri et al. (2018) Mohri, M., Rostamizadeh, A. and Talwalkar, A. (2018).  Foundations of machine learning.  MIT press.
  * Muandet et al. (2016) Muandet, K., Fukumizu, K., Sriperumbudur, B. and Schölkopf, B. (2016).  Kernel mean embedding of distributions: A review and beyond.  arXiv preprint arXiv:1605.09522.
  * Nguyen (2019) Nguyen, P.-M. (2019).  Mean field limit of the learning dynamics of multilayer neural networks.  arXiv preprint arXiv:1902.02880.
  * Pearl (2009) Pearl, J. (2009).  Causality.  Cambridge University press.
  * Radford et al. (2018) Radford, A., Narasimhan, K., Salimans, T. and Sutskever, I. (2018).  Improving language understanding by generative pre-training.  Technical Report.
  * Radford et al. (2019) Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I. et al. (2019).  Language models are unsupervised multitask learners.  OpenAI blog.
  * Ramesh et al. (2021) Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Radford, A., Chen, M. and Sutskever, I. (2021).  Zero-shot text-to-image generation.  In International Conference on Machine Learning.
  * Romero and Cordonnier (2020) Romero, D. W. and Cordonnier, J.-B. (2020).  Group equivariant stand-alone self-attention for vision.  arXiv preprint arXiv:2010.00977.
  * Rotskoff and Vanden-Eijnden (2018) Rotskoff, G. and Vanden-Eijnden, E. (2018).  Parameters as interacting particles: Long time convergence and asymptotic error scaling of neural networks.  Neural Information Processing Systems.
  * Sannai et al. (2021) Sannai, A., Imaizumi, M. and Kawano, M. (2021).  Improved generalization bounds of group invariant/equivariant deep networks via quotient feature spaces.  In Uncertainty in Artificial Intelligence.
  * Satorras et al. (2021) Satorras, V. G., Hoogeboom, E., Fuchs, F. B., Posner, I. and Welling, M. (2021).  ${E}(n)$ equivariant normalizing flows.  arXiv preprint arXiv:2105.09016.
  * Scarselli et al. (2008) Scarselli, F., Gori, M., Tsoi, A. C., Hagenbuchner, M. and Monfardini, G. (2008).  The graph neural network model.  IEEE Transactions on Neural Networks.
  * Schulz et al. (2018) Schulz, E., Speekenbrink, M. and Krause, A. (2018).  A tutorial on Gaussian process regression: Modelling, exploring, and exploiting functions.  Journal of Mathematical Psychology.
  * Shawe-Taylor et al. (2004) Shawe-Taylor, J., Cristianini, N. et al. (2004).  Kernel methods for pattern analysis.  Cambridge University Press.
  * Sirignano and Spiliopoulos (2020) Sirignano, J. and Spiliopoulos, K. (2020).  Mean field analysis of neural networks: A central limit theorem.  Stochastic Processes and Their Applications.
  * Sokolic et al. (2017) Sokolic, J., Giryes, R., Sapiro, G. and Rodrigues, M. (2017).  Generalization error of invariant classifiers.  In Artificial Intelligence and Statistics.
  * Song et al. (2014) Song, L., Anandkumar, A., Dai, B. and Xie, B. (2014).  Nonparametric estimation of multi-view latent variable models.  In International Conference on Machine Learning.
  * Song et al. (2009) Song, L., Huang, J., Smola, A. and Fukumizu, K. (2009).  Hilbert space embeddings of conditional distributions with applications to dynamical systems.  In International Conference on Machine Learning.
  * Tsai et al. (2019) Tsai, Y.-H. H., Bai, S., Yamada, M., Morency, L.-P. and Salakhutdinov, R. (2019).  Transformer dissection: A unified understanding of transformer’s attention via the lens of kernel.  arXiv preprint arXiv:1908.11775.
  * Valle-Pérez and Louis (2020) Valle-Pérez, G. and Louis, A. A. (2020).  Generalization bounds for deep learning.  arXiv preprint arXiv:2012.04115.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł. and Polosukhin, I. (2017).  Attention is all you need.  In Neural Information Processing Systems.
  * Vuckovic et al. (2020) Vuckovic, J., Baratin, A. and Combes, R. T. d. (2020).  A mathematical theory of attention.  arXiv preprint arXiv:2007.02876.
  * Wainwright (2019) Wainwright, M. J. (2019).  High-dimensional statistics: A non-asymptotic viewpoint.  Cambridge University Press.
  * Wasserman (2000) Wasserman, L. (2000).  Bayesian model selection and model averaging.  Journal of Mathematical Psychology, 44 92–107.
  * Wei et al. (2021) Wei, C., Chen, Y. and Ma, T. (2021).  Statistically meaningful approximation: A case study on approximating Turing machines with transformers.  arXiv preprint arXiv:2107.13163.
  * Wolf et al. (2020) Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M. et al. (2020).  Transformers: State-of-the-art natural language processing.  In Empirical Methods in Natural Language Processing.
  * Xie et al. (2021) Xie, S. M., Raghunathan, A., Liang, P. and Ma, T. (2021).  An explanation of in-context learning as implicit Bayesian inference.  arXiv preprint arXiv:2111.02080.
  * Yang (2020) Yang, G. (2020).  Tensor programs II: Neural tangent kernel for any architecture.  arXiv preprint arXiv:2006.14548.
  * Yang and Littwin (2021) Yang, G. and Littwin, E. (2021).  Tensor programs IIb: Architectural universality of neural tangent kernel training dynamics.  In International Conference on Machine Learning.
  * Zaheer et al. (2017) Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B., Salakhutdinov, R. R. and Smola, A. J. (2017).  Deep sets.  Neural Information Processing Systems.
  * Zhang et al. (2022a) Zhang, F., Liu, B., Wang, K., Tan, V. Y., Yang, Z. and Wang, Z. (2022a).  Relational reasoning via set transformers: Provable efficiency and applications to MARL.  arXiv preprint arXiv:2209.09845.
  * Zhang et al. (2019) Zhang, X., Yu, Y., Wang, L. and Gu, Q. (2019).  Learning one-hidden-layer ReLU networks via gradient descent.  In International Conference on Machine Learning.
  * Zhang et al. (2022b) Zhang, Y., Backurs, A., Bubeck, S., Eldan, R., Gunasekar, S. and Wagner, T. (2022b).  Unveiling transformers with LEGO: A synthetic reasoning task.  arXiv preprint arXiv:2206.04301.
  * Zhu et al. (2021) Zhu, S., An, B. and Huang, F. (2021).  Understanding the generalization benefit of model invariance from a data perspective.  Neural Information Processing Systems.
  * Zou et al. (2018) Zou, D., Cao, Y., Zhou, D. and Gu, Q. (2018).  Stochastic gradient descent optimizes over-parameterized deep ReLU networks.  arXiv preprint arXiv:1811.08888.
  * Zou and Gu (2019) Zou, D. and Gu, Q. (2019).  An improved analysis of training over-parameterized deep neural networks.  Neural Information Processing Systems.

##  Appendix A Conditional Mean Embedding

We introduce the conditional mean embedding, which embeds the conditional
distribution to the element in the RKHS. Let $\mathcal{H}_{x}$ and
$\mathcal{H}_{y}$ be two RKHSs over spaces ${\mathfrak{X}}$ and
${\mathfrak{Y}}$ with kernels $\mathfrak{K}$ and $\mathfrak{L}$, respectively.
Let $\mathbb{P}_{\mathcal{X},\mathcal{Y}}$ be the joint distribution of two
random variables $\mathcal{X}$ and $\mathcal{Y}$ taking values in
${\mathfrak{X}}$ and ${\mathfrak{Y}}$, respectively. The conditional mean
embedding
${\mathtt{CME}}(x,\mathbb{P}_{\mathcal{X},\mathcal{Y}})\in\mathcal{H}_{y}$ of
the conditional distribution $\mathbb{P}_{\mathcal{Y}{\,|\,}\mathcal{X}}$ is
defined as

$\displaystyle{\mathtt{CME}}(x,\mathbb{P}_{\mathcal{X},\mathcal{Y}})=\mathbb{E}\bigl{[}\mathfrak{L}(\mathcal{Y},\cdot){\,\big{|}\,}\mathcal{X}=x\bigr{]}.$

Then, for any function $g\in\mathcal{H}_{y}$ and $x\in{\mathfrak{X}}$, it
holds that

$\displaystyle\mathbb{E}\bigl{[}g(\mathcal{Y}){\,\big{|}\,}\mathcal{X}=x\bigr{]}=\big{\langle}g,{\mathtt{CME}}(x,\mathbb{P}_{\mathcal{X},\mathcal{Y}})\big{\rangle}_{\mathcal{H}_{y}},$

which shows that the conditional mean embedding is a representation of the
conditional distribution in the RKHS. Correspondingly, the conditional mean
embedding operator
$\mathcal{C}_{\mathcal{Y}{\,|\,}\mathcal{X}}:\mathcal{H}_{x}\rightarrow\mathcal{H}_{y}$
is an operator such that

$\displaystyle\mathcal{C}_{\mathcal{Y}{\,|\,}\mathcal{X}}\mathfrak{K}(x,\cdot)={\mathtt{CME}}(x,\mathbb{P}_{\mathcal{X},\mathcal{Y}}),$

for any $x\in{\mathfrak{X}}$. We define the (uncentered) covariance operator
$\mathcal{C}_{\mathcal{X}\mathcal{X}}:\mathcal{H}_{x}\rightarrow\mathcal{H}_{x}$
and the cross-covariance operator
$\mathcal{C}_{\mathcal{Y}\mathcal{X}}:\mathcal{H}_{x}\rightarrow\mathcal{H}_{y}$
as follows,

$\displaystyle\mathcal{C}_{\mathcal{X}\mathcal{X}}=\mathbb{E}\bigl{[}\mathfrak{K}(\mathcal{X},\cdot)\otimes\mathfrak{K}(\mathcal{X},\cdot)\bigr{]},\quad\mathcal{C}_{\mathcal{Y}\mathcal{X}}=\mathbb{E}\bigl{[}\mathfrak{L}(\mathcal{Y},\cdot)\otimes\mathfrak{K}(\mathcal{X},\cdot)\bigr{]}.$

Here $\otimes$ is the tensor product. Then, it holds that
$\mathcal{C}_{\mathcal{Y}{\,|\,}\mathcal{X}}=\mathcal{C}_{\mathcal{Y}\mathcal{X}}\mathcal{C}_{\mathcal{X}\mathcal{X}}^{-1}$.
We consider the following regularized least squares problem,

$\displaystyle\min_{\mathcal{C}:\mathcal{H}_{x}\rightarrow\mathcal{H}_{y}}\mathcal{E}_{\lambda}(\mathcal{C})=\mathbb{E}\Bigl{[}\big{\|}\mathfrak{L}(\mathcal{Y},\cdot)-\mathcal{C}\mathfrak{K}(\mathcal{X},\cdot)\big{\|}^{2}_{\mathcal{H}_{y}}\Bigr{]}+\lambda\cdot\|\mathcal{C}\|^{2}_{\mathrm{HS}},$
(A.1)

where $\|\cdot\|_{\mathrm{HS}}$ denotes the Hilbert-Schmidt norm and
$\lambda>0$ is the regularization parameter. The solution to (A.1) is given by
$\mathcal{C}^{\lambda}_{\mathcal{Y}{\,|\,}\mathcal{X}}=\mathcal{C}_{\mathcal{Y}\mathcal{X}}(\mathcal{C}_{\mathcal{X}\mathcal{X}}+\lambda\mathcal{I})^{-1}$,
which is interpreted as a regularized conditional mean embedding operator.
Here $\mathcal{I}:\mathcal{H}_{x}\rightarrow\mathcal{H}_{x}$ is the identity
mapping. Correspondingly, we consider the following empirical least squares
problem,

$\displaystyle\widehat{\mathcal{E}}(\mathcal{U})=\frac{1}{L}\sum_{\ell=1}^{L}\big{\|}\mathfrak{L}(y^{\ell},\cdot)-\mathcal{C}\mathfrak{K}(x^{\ell},\cdot)\big{\|}^{2}_{\mathcal{H}_{y}}+\lambda\cdot\|\mathcal{C}\|^{2}_{\mathrm{HS}},$
(A.2)

where $\\{(x^{\ell},y^{\ell})\\}_{\ell\in[L]}$ are sampled from
$\mathbb{P}_{\mathcal{X},\mathcal{Y}}$. To present the solution to (A.2), we
define the sampling operator $S_{X}:\mathcal{H}_{x}\rightarrow\mathbb{R}^{L}$
and $S_{Y}:\mathcal{H}_{y}\rightarrow\mathbb{R}^{L}$ as
$S_{X}f=(f(x^{1}),\ldots,f(x^{L}))$ and $S_{Y}g=(g(y^{1}),\ldots,g(y^{L}))$
for any $f\in\mathcal{H}_{x}$ and $g\in\mathcal{H}_{y}$, respectively. We
denote by $S_{X}^{*}$ and $S_{Y}^{*}$ the adjoints of $S_{X}$ and $S_{Y}$,
respectively. Moreover, we denote by
$\widehat{\mathcal{C}}_{\mathcal{X}\mathcal{X}}=L^{-1}\sum_{\ell=1}^{L}\mathfrak{K}(x^{\ell},\cdot)\otimes\mathfrak{K}(x^{\ell},\cdot)=L^{-1}S_{X}^{*}S_{X}$
and
$\widehat{\mathcal{C}}_{\mathcal{Y}\mathcal{X}}=L^{-1}\sum_{\ell=1}^{L}\mathfrak{L}(y^{\ell},\cdot)\otimes\mathfrak{K}(x^{\ell},\cdot)=S_{Y}^{*}S_{X}$
the empirical (uncentered) covariance operator and (uncentered) cross-
covariance operator, respectively. Then, the solution to (A.2) is given by
$\widehat{\mathcal{C}}_{\mathcal{Y}{\,|\,}\mathcal{X}}^{\lambda}=\widehat{\mathcal{C}}_{\mathcal{Y}\mathcal{X}}(\widehat{\mathcal{C}}_{\mathcal{X}\mathcal{X}}+\lambda\mathcal{I})^{-1}$.
We denote by
$\widehat{\mathtt{CME}}_{\lambda}(x,\mathbb{P}_{\mathcal{X},\mathcal{Y}})\in\mathcal{H}_{y}$
the empirical conditional mean embedding. Following from the identity that
$S_{X}(S_{X}^{*}S_{X}+L\lambda\mathcal{I})^{-1}=(S_{X}S_{X}^{*}+L\lambda
I_{L})^{-1}S_{X}$, it holds that

$\displaystyle\widehat{\mathtt{CME}}_{\lambda}(x,\mathbb{P}_{\mathcal{X},\mathcal{Y}})$
$\displaystyle=\widehat{\mathcal{C}}_{\mathcal{Y}{\,|\,}\mathcal{X}}^{\lambda}\mathfrak{K}(x,\cdot)$
(A.3)
$\displaystyle=\widehat{\mathcal{C}}_{\mathcal{Y}\mathcal{X}}(\widehat{\mathcal{C}}_{\mathcal{X}\mathcal{X}}+\lambda)^{-1}\mathfrak{K}(x,\cdot)=S_{Y}^{*}(\mathfrak{K}(X,X)+L\lambda
I_{L})^{-1}\mathfrak{K}(X,x).$

Here
$\mathfrak{K}(X,X)=S_{X}S_{X}^{*}=(\mathfrak{K}(x^{i},x^{j}))_{i,j\in[L]}\in\mathbb{R}^{L\times
L}$ is the Gram matrix and
$\mathfrak{K}(X,x)=(\mathfrak{K}(x^{1},x),\ldots,\mathfrak{K}(x^{L},x))\in\mathbb{R}^{L}$.

##  Appendix B Attention Recovers Latent Posterior

###  B.1 Gaussian Process Regression

Gaussian Process Regression. We say that $f$ follows a Gaussian process (GP)
$\mathtt{GP}(\mu(\cdot),\mathfrak{K}(\cdot,\cdot))$ on $\mathbb{R}^{d}$ if for
$x^{1},\ldots,x^{L}$, $(f(x^{1}),\ldots,f(x^{L}))$ follows a Gaussian
distribution with mean $(\mu(x^{1}),\ldots,\mu(x^{L}))$ and covariance
$(\mathfrak{K}(x^{i},x^{j}))_{i,j\in[L]}$. Here $\mu(x)=\mathbb{E}[f(x)]$ is
the mean function and
$\mathfrak{K}(x,x^{\prime})=\mathbb{E}[(f(x)-\mu(x))(f(x^{\prime})-\mu(x^{\prime}))]$
is the covariance (or kernel) function. We take
$\mathtt{GP}(0,\mathfrak{K}(x,x^{\prime}))$ as the prior of $f$. Given noisy
samples $\\{(x^{\ell},y^{\ell})\\}_{\ell\in[L]}$, where
$y^{\ell}=f(x^{\ell})+\epsilon^{\ell}$ with $\epsilon^{\ell}\sim N(0,\lambda
I)$ i.i.d. Gaussian noise, the posterior over $f$ is a GP again with mean
$\mu_{p}(x)$ and covariance $\mathfrak{K}_{p}(x,x^{\prime})$ (Schulz et al.,
2018). In particular, it holds that

$\displaystyle\mu_{p}(x)$
$\displaystyle=\mathfrak{K}(x,X)(\mathfrak{K}(X,X)+\lambda I)^{-1}y,$
$\displaystyle\mathfrak{K}_{p}(x,x^{\prime})$
$\displaystyle=\mathfrak{K}(x,x^{\prime})-\mathfrak{K}(x,X)(\mathfrak{K}(X,X)+\lambda
I)^{-1}\mathfrak{K}(X,x^{\prime}).$

Here
$\mathfrak{K}(x,X)=(\mathfrak{K}(x,x^{\ell}))_{\ell\in[L]}^{\top}\in\mathbb{R}^{1\times
L}$,
$\mathfrak{K}(X,X)=(\mathfrak{K}(x^{i},x^{j}))_{i,j\in[L]}\in\mathbb{R}^{L\times
L}$, and
$\mathfrak{K}(X,x^{\prime})=(\mathfrak{K}(x^{\ell},x^{\prime}))_{\ell\in[L]}\in\mathbb{R}^{L}$.

Rigorous Characterization of Latent Variable Model. We now provide a rigorous
characterization of the advanced infinite-dimensional example of latent
variable model in §4.1. In particular, we consider the following model,

$\displaystyle r^{\ell}=f(c^{\ell})+\epsilon^{\ell},$ $\displaystyle
r^{\mathtt{msk}}=f(c^{\mathtt{msk}})+\epsilon.$ (B.1)

Here $f=(f_{1},\ldots,f_{d})$ with each
$f_{i}\sim\mathtt{GP}(0,\mathfrak{K}(\cdot,\cdot))$ for $i\in[d]$, and
$\epsilon^{\ell}$ and $\epsilon$ are independent Gaussian noise drawn from
$N(0,\lambda I)$. Then, following GP regression, we recover (4.1).

###  B.2 Implication of Limiting Expectation

Necessity of Multiple Heads. Based on the definition of the attention
mechanism $\mathtt{attn}$ in (2.1), we define the multihead attention as

$\displaystyle{\mathtt{mha}}(q,X;W)=\sum_{i=1}^{h}{\mathtt{head}}_{i}\in\mathbb{R}^{d}.$
(B.2)

Here $h\in\mathbb{N}_{+}$ is the number of heads,
$W=\\{(W_{i}^{\mathrm{q}},W_{i}^{\mathrm{k}},W_{i}^{\mathrm{v}})\\}_{i=1}^{h}$
is the (learnable) parameter, and

$\displaystyle{\mathtt{head}}_{i}=\mathtt{attn}(q,K_{i},V_{i})\in\mathbb{R}^{d},\qquad\text{where}\quad
K_{i}=XW_{i}^{\mathrm{k}},\quad V=XW_{i}^{\mathrm{v}}.$

We remark that the multihead attention defined in (B.2) is written in the
summation form, which is equivalent to the concatenation form (Vaswani et al.,
2017). In multihead attentions, we set that $d=d_{\mathrm{p}}\cdot h$, where
$h$ is the number of heads.

We remark that the limit $\mathbb{E}[\mathcal{V}{\,|\,}\mathcal{K}=q]$ shows
the necessity of using the multiple heads in the multihead attention. Note
that the key and value are calculated from the following mappings,

$\displaystyle
k^{\ell}=k_{\theta}(x^{\ell})=(W^{\mathrm{k}})^{\top}x^{\ell},\qquad
v^{\ell}=v_{\theta}(x^{\ell})=(W^{\mathrm{v}})^{\top}x^{\ell},$