where $x^{\ell}\in\mathbb{R}^{d}$ is the input token and
$W^{\mathrm{k}}\in\mathbb{R}^{d\times d_{\mathrm{p}}}$,
$W^{\mathrm{v}}\in\mathbb{R}^{d\times d}$ are the learnable parameters. In
single-head attentions, where $h=1$ and $d_{\mathrm{p}}=d$, the matrix
$W^{\mathrm{k}}\in\mathbb{R}^{d\times d}$ is invertible with a high
probability. We denote by $\mathcal{K}$, $\mathcal{V}$, and $\mathcal{X}$ the
random variable with the same distribution as $k^{\ell}$, $v^{\ell}$, and
$x^{\ell}$, respectively. By Propositions 4.1 and 4.2, we have that

$\displaystyle\mathtt{attn}(q,K,V)\approx\mathbb{E}[\mathcal{V}{\,|\,}\mathcal{K}=q]=\mathbb{E}\bigl{[}(W^{\mathrm{v}})^{\top}\mathcal{X}\,\big{|}\,(W^{\mathrm{k}})^{\top}\mathcal{X}=q\bigr{]}=\bigl{(}(W^{\mathrm{k}})^{-1}W^{\mathrm{v}}\bigr{)}^{\top}q,$

which is a linear mapping and fails to capture the interaction query $q$ and
the input sequence $X$. Thus, the single-head attention mechanism becomes a
linear mapping in the limit when $L\rightarrow\infty$. In comparison, when the
number of heads $h>1$ and $d_{\mathrm{p}}<d$, the matrices
$W^{\mathrm{k}}\in\mathbb{R}^{d\times d}$ are not invertible. Thus, to
incorporate nonlinearity into the attention, it is necessary to use multiple
heads, where $h>1$.

Instrumental Variable. The limit also shows that the attention is effective in
the sense of causal inference. We first introduce the instrumental variable
and causal inference. Instrumental variable (IV) regression is used to
estimate the causal relationship between the input $\mathcal{X}$ and the
output $\mathcal{Y}$. Specifically, when the measurements of a pair of random
variables $(\mathcal{X},\mathcal{Y})$ is confouneded, an instrumental variable
$\mathcal{Z}$ is effective in identifying their causal relationship.
Intuitively, $\mathcal{Z}$ is an IV if it only influences $\mathcal{Y}$
through $\mathcal{X}$. More specifically, the causal assumption about the IV
is as follows.

######  Assumption B.1 (IV Model).

Let $(\mathcal{X},\mathcal{Y},\mathcal{Z})$ be a random variable on the space
${\mathfrak{X}}\times{\mathfrak{Y}}\times{\mathfrak{Z}}$ with joint
distribution $\mathbb{P}_{(X,Y,Z)}$. We assume that

  * (i)

$\mathcal{Y}=g(\mathcal{X})+\epsilon$ and
$\mathbb{E}[\epsilon{\,|\,}\mathcal{Z}]=0$.

  * (ii)

$\mathbb{P}_{\mathcal{X}{\,|\,}\mathcal{Z}}(x{\,|\,}z)$ is not constant in
$z$.

In the IV model in Assumption B.1, $\mathcal{Z}$ is the IV. Assumption B.1 (i)
is the exclusion restriction. The function $g$ is the true structural
function. The error term $\epsilon$ is the confounding noise. Assumption B.1
(ii) is the relevance condition, which ensures that $\mathcal{Z}$ is
informative. We remark that the IV model allows causal inference in a more
general contexts. To see this, when $\mathcal{X}=\mathcal{Z}$, the estimation
of $h$ reduces to standard regression of unconfounded inputs and it holds that
$h(\cdot)=\mathbb{E}[\mathcal{Y}{\,|\,}\mathcal{X}=\cdot]$. In particular, the
IV model allows for confounded input, i.e., $\mathcal{X}$ and $\epsilon$ are
not independent. By Assumption B.1, we have the following equation

$\displaystyle\mathbb{E}[\mathcal{Y}{\,|\,}\mathcal{Z}]=\mathbb{E}\bigl{[}g(\mathcal{X}){\,|\,}\mathcal{Z}\bigr{]}.$
(B.3)

The right-hand side of (B.3) provides a two-stage method for estimating the
function $g$. At the first stage, we estimate the conditional mean
$\mathcal{X}{\,|\,}\mathcal{Z}$. Then, at the second stage, we estimate the
function $g$ via regressing $\mathcal{Y}$ on $\mathcal{X}{\,|\,}\mathcal{Z}$.

For ease of presentation, for any given $q$, $K$, and $V$, we consider the
following mapping,

$\displaystyle g_{\theta}\circ\mathtt{attn}(q,K,V)\in\mathbb{R}^{d},$

where $g_{\theta}$ is a function approximation with parameter $\theta$. For
example, $g_{\theta}$ can be a deep neural network. By Proposition 4.1 and
Proposition 4.2, it holds that

$\displaystyle g_{\theta}\circ\mathtt{attn}(q,K,V)\in\mathbb{R}^{d}\approx
g_{\theta}\bigl{(}\mathbb{E}[\mathcal{V}{\,|\,}\mathcal{K}=q]\bigr{)},\quad\text{as
}L\rightarrow\infty.$

Let the target be $y$. Then, the learning problem can be formulated as
maximizing the following log-likelihood,

$\displaystyle\log\mathbb{P}_{\theta}\bigl{(}y{\,\big{|}\,}\mathbb{E}[\mathcal{V}{\,|\,}\mathcal{K}=q]\bigr{)}.$

Thus, the key $\mathcal{K}$ is the instrumental variable (Pearl, 2009), which
handles the endogeneity between $y$ and $V$. Note that the mappings from the
input sequence $X$ to $Q$, $K$, and $V$ are learned through the parameter
$\theta$. Thus, the attention is effective in relational reasoning and
identifying the instrumental variable.

###  B.3 Proof of Lemma 3.2

###### Proof.

We first prove the statement that $b(X)$ is a minimal sufficient statistic of
$X$ for $z$. The sufficiency of $b(X)$ is straightforward. We now show it is
minimal. Suppose that $U(X)$ is a sufficient statistic of $X$ for $z$. Then,
we have

$\displaystyle\mathbb{P}(z\,|\,X)=\mathbb{P}(z\,|\,U(X))\quad\text{for any
$z\in\mathcal{Z}$},$

which implies that $b(X)$ is a function of $U(X)$. Therefore, $b(X)$ is
minimal.

We now prove the second statement that $b(X)$ is a minimal sufficient
statistics of $X$ for $y$. Suppose that $U(X)$ is a sufficient statistic of
$X$ for $y$. Then, we have

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

which implies

$\displaystyle
b(X)={\mathcal{T}}^{-1}\Bigl{(}\int_{\mathcal{Z}}\mathbb{P}(y=\cdot\,|\,z)\cdot\mathbb{P}(z\,|\,U(X)){\mathrm{d}}z\Bigr{)}$

which means $b(X)$ is a function of $U(X)$. Therefore, $b(X)$ is minimal. ∎

###  B.4 Proof of Proposition 4.1

###### Proof.

According to (A.3), we see that $\mathtt{attn}_{\mathtt{CME}}$ recovers the
empirical conditional mean embedding with the Euclidean kernel function
$\mathfrak{L}(y,y^{\prime})=\langle y,y^{\prime}\rangle$. Note that we have
${\mathtt{CME}}(q,\mathbb{P}_{\mathcal{K},\mathcal{V}})=\mathbb{E}[\mathcal{V}{\,|\,}\mathcal{K}=q]$.
Recall the definition of the attention and the conditional mean embedding in
§A. It holds that

$\displaystyle\big{\|}\mathtt{attn}(q,K,V)-{\mathtt{CME}}(q,\mathbb{P}_{\mathcal{K},\mathcal{V}})\big{\|}_{2}$
$\displaystyle\quad\leq\underbrace{\big{\|}\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathfrak{K}(q,\cdot)-\mathcal{C}_{\mathcal{V}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathfrak{K}(q,\cdot)\big{\|}_{2}}_{\displaystyle\text{term
(i)}}$
$\displaystyle\qquad+\underbrace{\big{\|}\mathcal{C}_{\mathcal{V}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathfrak{K}(q,\cdot)-\mathcal{C}_{\mathcal{V}\mathcal{K}}\mathcal{C}_{\mathcal{K}\mathcal{K}}^{-1}\mathfrak{K}(q,\cdot)\big{\|}_{2}}_{\displaystyle\text{term
(ii)}}.$ (B.4)

Upper bounding term (i) of (B.4). It suffices to upper bound
$\|\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathfrak{K}(q,\cdot)-\mathcal{C}_{\mathcal{V}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\|$.
It holds that

$\displaystyle\big{\|}\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}-\mathcal{C}_{\mathcal{V}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}$
$\displaystyle\quad\leq\big{\|}\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}\bigl{(}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}-(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\bigr{)}\big{\|}+\big{\|}(\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}-\mathcal{C}_{\mathcal{V}\mathcal{K}})(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}$
$\displaystyle\quad=\big{\|}\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}-\mathcal{C}_{\mathcal{K}\mathcal{K}})(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}+\big{\|}(\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}-\mathcal{C}_{\mathcal{V}\mathcal{K}})(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}.$
(B.5)

For the first term, by the fact that
$\|\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}\|\leq\|V^{\top}V\|^{1/2}\cdot\|\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}\|^{1/2}$,
we have that

$\displaystyle\big{\|}\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}-\mathcal{C}_{\mathcal{K}\mathcal{K}})(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}$
$\displaystyle\quad\leq\|V^{\top}V\|^{1/2}\cdot\big{\|}\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}^{1/2}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1/2}\big{\|}\cdot\big{\|}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1/2}\big{\|}\cdot\big{\|}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}-\mathcal{C}_{\mathcal{K}\mathcal{K}})(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}$
$\displaystyle\quad\leq\lambda^{-1/2}\cdot\big{\|}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}-\mathcal{C}_{\mathcal{K}\mathcal{K}})(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}.$
(B.6)

Plugging (B.4) into (B.4), we have

$\displaystyle\big{\|}\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}-\mathcal{C}_{\mathcal{V}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}$
$\displaystyle\quad\leq\lambda^{-1/2}\cdot\big{\|}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}-\mathcal{C}_{\mathcal{K}\mathcal{K}})(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}+\big{\|}(\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}-\mathcal{C}_{\mathcal{V}\mathcal{K}})(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}.$
(B.7)

In what follows, we upper bound the second term on the right-hand side of
(B.4) using Lemma G.2. Let $\xi\in\mathcal{H}_{v}\otimes\mathcal{H}_{k}$
defined by

$\displaystyle\xi(k,v)=v\mathfrak{K}(k,\cdot)(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}.$

Following from
$\big{\|}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}\leq\lambda^{-1}$,
we have that

$\displaystyle\|\xi(k,v)\|=\big{\|}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}\cdot\|v\|_{\mathcal{H}_{v}}\cdot\|\phi(k)\|_{\mathcal{H}_{k}}\leq
C\cdot\lambda^{-1},$

where $C>0$ is an absolute constant. In addition, we have that

$\displaystyle\mathbb{E}\bigl{[}\|\xi(k,v)\|^{2}\bigr{]}$
$\displaystyle=\mathbb{E}\Bigl{[}\big{\|}\phi(k)^{\top}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}_{\mathcal{H}_{k}}^{2}\cdot\|v\|_{\mathcal{H}_{v}}^{2}\Bigr{]}$
$\displaystyle\leq\mathbb{E}\Bigl{[}\big{\|}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\phi(k)\big{\|}_{\mathcal{H}_{k}}^{2}\Bigr{]}$
$\displaystyle=\mathbb{E}\Bigl{[}\big{\langle}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-2}\phi(k),\phi(k)\big{\rangle}_{\mathcal{H}_{k}}\Bigr{]}.$

Using the trace operator, we have that

$\displaystyle\mathbb{E}\bigl{[}\|\xi(k,v)\|^{2}\bigr{]}$
$\displaystyle=\mathbb{E}\Bigl{[}\operatorname{tr}\bigl{(}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-2}\phi(k)\phi(k)^{\top}\bigr{)}\Bigr{]}$
$\displaystyle=\operatorname{tr}\bigl{(}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-2}\mathcal{C}_{\mathcal{K}\mathcal{K}}\bigr{)}$
$\displaystyle\leq\lambda^{-1}\operatorname{tr}\bigl{(}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathcal{C}_{\mathcal{K}\mathcal{K}}\bigr{)}=\lambda^{-1}\cdot\gamma(\lambda),$

where $\gamma(\lambda)$ is the effective dimension of
$\mathcal{C}_{\mathcal{K}\mathcal{K}}$ defined in Theorem 4.1. Applying Lemma
G.2 with $B=C\lambda^{-1}$ and $\sigma^{2}=\lambda^{-1}\cdot\gamma(\lambda)$,
we have with probability at least $1-\delta$ that

$\displaystyle\big{\|}\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}-\mathcal{C}_{\mathcal{V}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}\leq
C\cdot\biggl{(}\frac{2}{\lambda L}+\sqrt{\frac{\gamma(\lambda)}{\lambda
L}}\biggr{)}\log\frac{2}{\delta},$ (B.8)

where $C>0$ is an absolute constant. Similarly, we have with probability at
least $1-\delta$ that

$\displaystyle\big{\|}\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}-\mathcal{C}_{\mathcal{K}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}\leq
C\cdot\biggl{(}\frac{2}{\lambda L}+\sqrt{\frac{\gamma(\lambda)}{\lambda
L}}\biggr{)}\log\frac{2}{\delta}.$ (B.9)

Plugging (B.8) and (B.9) into (B.4), we have with probability at least
$1-\delta$ that

$\displaystyle\big{\|}\widehat{\mathcal{C}}_{\mathcal{V}\mathcal{K}}(\widehat{\mathcal{C}}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}-\mathcal{C}_{\mathcal{V}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\big{\|}\leq
C\cdot\lambda^{-1/2}\biggl{(}\frac{2}{\lambda
L}+\sqrt{\frac{\gamma(\lambda)}{\lambda L}}\biggr{)}\log\frac{2}{\delta}.$
(B.10)

Upper bounding term (ii) of (B.4). It holds that

$\displaystyle\big{\|}\mathcal{C}_{\mathcal{V}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathfrak{K}(q,\cdot)-\mathcal{C}_{\mathcal{V}\mathcal{K}}\mathcal{C}_{\mathcal{K}\mathcal{K}}^{-1}\mathfrak{K}(q,\cdot)\big{\|}_{2}$
$\displaystyle\quad=\big{\langle}\mathcal{C}_{\mathcal{V}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathfrak{K}(q,\cdot)-\mathcal{C}_{\mathcal{V}\mathcal{K}}\mathcal{C}_{\mathcal{K}\mathcal{K}}^{-1}\mathfrak{K}(q,\cdot),\mathcal{C}_{\mathcal{V}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathfrak{K}(q,\cdot)-\mathcal{C}_{\mathcal{V}\mathcal{K}}\mathcal{C}_{\mathcal{K}\mathcal{K}}^{-1}\mathfrak{K}(q,\cdot)\big{\rangle}_{\mathcal{H}_{k}}$
$\displaystyle\quad=\bigg{\langle}\Big{(}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathcal{C}_{\mathcal{K}\mathcal{K}}\otimes(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathcal{C}_{\mathcal{K}\mathcal{K}}-\mathcal{I}\otimes(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathcal{C}_{\mathcal{K}\mathcal{K}}$
$\displaystyle\quad\qquad(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathcal{C}_{\mathcal{K}\mathcal{K}}\otimes\mathcal{I}+\mathcal{I}\otimes\mathcal{I}\Big{)}\mathbb{E}\bigl{[}v^{\top}\bar{v}^{\top}{\,\big{|}\,}k=\cdot,\bar{k}=\dagger\bigr{]},\mathfrak{K}(q,\cdot)\otimes\mathfrak{K}(q,\dagger)\bigg{\rangle}_{\mathcal{H}_{k}\otimes\mathcal{H}_{k}}.$

Note that $\mathbb{E}[v^{\top}\bar{v}^{\top}{\,|\,}k=\cdot,\bar{k}=\dagger]$
is in the range of
$\mathcal{C}_{\mathcal{K}\mathcal{K}}\otimes\mathcal{C}_{\mathcal{K}\mathcal{K}}$
by the assumption. Let $\widetilde{\mathcal{C}}$ be in
$(\mathcal{C}_{\mathcal{K}\mathcal{K}}\otimes\mathcal{C}_{\mathcal{K}\mathcal{K}})^{-1}\mathbb{E}[v^{\top}\bar{v}^{\top}{\,|\,}k=\cdot,\bar{k}=\dagger]$.
Let $\\{\lambda_{i}\\}_{i=1}^{\infty}$ and $\\{\psi_{i}\\}_{i=1}^{\infty}$ be
the eigenvalues and eigenvectors of $\mathcal{C}_{\mathcal{K}\mathcal{K}}$.
Then, we have that

$\displaystyle\bigg{\|}\Big{(}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathcal{C}_{\mathcal{K}\mathcal{K}}\otimes(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathcal{C}_{\mathcal{K}\mathcal{K}}-\mathcal{I}\otimes(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathcal{C}_{\mathcal{K}\mathcal{K}}$
$\displaystyle\quad\qquad(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathcal{C}_{\mathcal{K}\mathcal{K}}\otimes\mathcal{I}+\mathcal{I}\otimes\mathcal{I}\Big{)}\mathbb{E}\bigl{[}v^{\top}\bar{v}^{\top}{\,\big{|}\,}k=\cdot,\bar{k}=\dagger\bigr{]}\bigg{\|}^{2}$
$\displaystyle\quad=\sum_{i,j}\biggl{(}\frac{\lambda_{i}^{2}}{\lambda_{i}+\lambda}\frac{\lambda_{j}^{2}}{\lambda_{j}+\lambda}-\frac{\lambda_{i}^{2}\lambda_{j}}{\lambda_{i}+\lambda}-\frac{\lambda_{j}^{2}\lambda_{i}}{\lambda_{j}+\lambda}+\lambda_{i}\lambda_{j}\biggr{)}^{2}\langle\psi_{i}\otimes\psi_{j},\widetilde{\mathcal{C}}\rangle_{\mathcal{H}_{k}\otimes\mathcal{H}_{k}}^{2}$
$\displaystyle\quad=\sum_{i,j}\biggl{(}\frac{\lambda_{i}\lambda_{j}\lambda^{2}}{(\lambda_{i}+\lambda)(\lambda_{j}+\lambda)}\biggr{)}^{2}\langle\psi_{i}\otimes\psi_{j},\widetilde{\mathcal{C}}\rangle_{\mathcal{H}_{k}\otimes\mathcal{H}_{k}}^{2}$
$\displaystyle\quad\leq\lambda^{4}\cdot\|\widetilde{\mathcal{C}}\|^{2}_{\mathcal{H}_{k}\otimes\mathcal{H}_{k}}.$

Thus, we have that

$\displaystyle\big{\|}\mathcal{C}_{\mathcal{V}\mathcal{K}}(\mathcal{C}_{\mathcal{K}\mathcal{K}}+\lambda\mathcal{I})^{-1}\mathfrak{K}(q,\cdot)-\mathcal{C}_{\mathcal{V}\mathcal{K}}\mathcal{C}_{\mathcal{K}\mathcal{K}}^{-1}\mathfrak{K}(q,\cdot)\big{\|}_{2}\leq
C\cdot\lambda,$ (B.11)

where $C>0$ is an absolute constant.

Finally, plugging (B.10) and (B.11) into (B.4), we complete the proof of
Proposition 4.1. ∎

###  B.5 Proof of Proposition 4.2

###### Proof.

Under the condition that
$\widehat{\mathbb{P}}_{\mathfrak{K}}(v{\,|\,}q)\rightarrow\mathbb{P}(v{\,|\,}q)$
uniformly as $L\rightarrow\infty$, we have that

$\displaystyle\int
v\widehat{\mathbb{P}}^{\mathfrak{K}}_{\mathcal{V}{\,|\,}\mathcal{K}}(v{\,|\,}q){\mathrm{d}}v\rightarrow\mathbb{E}[\mathcal{V}{\,|\,}\mathcal{K}=q],$

as $L\rightarrow\infty$. Moreover, it holds that

$\displaystyle\int
v\widehat{\mathbb{P}}^{\mathfrak{K}}_{\mathcal{V}{\,|\,}\mathcal{K}}(v{\,|\,}q){\mathrm{d}}v$
$\displaystyle=C_{0}\cdot\int_{\mathbb{S}^{d-1}}v\cdot\frac{\sum^{L}_{\ell=1}\mathfrak{K}(k^{\ell},q)\cdot\mathfrak{K}(v^{\ell},v)}{\sum^{L}_{\ell=1}\mathfrak{K}(k^{\ell},q)}{\mathrm{d}}v$
$\displaystyle=\frac{C_{0}\cdot\sum^{L}_{\ell=1}\mathfrak{K}(k^{\ell},q)\cdot\int_{\mathbb{S}^{d-1}}v\cdot\mathfrak{K}(v^{\ell},v){\mathrm{d}}v}{\sum^{L}_{\ell=1}\mathfrak{K}(k^{\ell},q)},$
(B.12)

where $\mathbb{S}^{d-1}$ is the $(d-1)$-dimensional unit sphere. It now
suffices to calculate the integration term

$\displaystyle\int_{\mathbb{S}^{d-1}}v\cdot\mathfrak{K}(v^{\ell},v){\mathrm{d}}v.$

To this end, we utilize the following lemma.

######  Lemma B.2.

Let $\mathfrak{K}(a,b)=\exp(a^{\top}b/\gamma)$ be the exponential kernel with
a fixed $\gamma>0$. It holds for all $b\in\mathbb{S}^{d-1}$ that

$\displaystyle\int_{\mathbb{S}^{d-1}}a\cdot\mathfrak{K}(a,b){\mathrm{d}}a=C_{1}\cdot
b,$

where $C_{1}>0$ is an absolute constant defined as follows,

$\displaystyle
C_{1}=\int_{\mathbb{S}^{d-1}}a^{\top}b\cdot\mathfrak{K}(a,b){\mathrm{d}}a,\quad\forall
b\in\mathbb{S}^{d-1}.$

Here we remark that $C_{1}$ is well-defined due to the symmetry of the unit
sphere $\mathbb{S}^{d-1}$.

###### Proof.

See §B.5.1 for a detailed proof. ∎

By Lemma B.2, it holds for the right-hand side of (B.5) that

$\displaystyle C_{0}\cdot
C_{1}\cdot\frac{\sum^{L}_{\ell=1}\mathfrak{K}(k^{\ell},q)\cdot
v^{\ell}}{\sum^{L}_{\ell=1}\mathfrak{K}(k^{\ell},q)}=C_{0}\cdot C_{1}\cdot
V^{\top}\mathtt{softmax}(Kq)=C_{0}\cdot
C_{1}\cdot\mathtt{attn}_{\mathtt{SM}}(q,K,V).$

Here recall that we define $K=[k^{1},\ldots,k^{L}]$ and
$V=[v^{1},\ldots,v^{\ell}]$ the keys and queries, and

$\displaystyle
C_{0}=\int_{\mathbb{S}^{d-1}}\mathfrak{K}(a^{\top}b){\mathrm{d}}a,\qquad
C_{1}=\int_{\mathbb{S}^{d-1}}a^{\top}b\cdot\mathfrak{K}(a,b){\mathrm{d}}a\quad\forall
b\in\mathbb{S}^{d-1}.$

Thus, by setting $C=C_{0}\cdot C_{1}$, we complete the proof of Lemma 4.2. ∎

####  B.5.1 Proof of Lemma B.2

###### Proof.

Let $a,b$ be two vectors in the $(d-1)$-dimensional unit sphere
$\mathbb{S}^{d-1}$. We first define the following vector,

$\displaystyle c=(a^{\top}b)\cdot b-\bigl{(}a-(a^{\top}b)\cdot
b\bigr{)}\in\mathbb{S}^{d-1}.$ (B.13)

By direct calculation, we have the following property of $c$ defined in
(B.13),

$\displaystyle c^{\top}b$
$\displaystyle=(a^{\top}b)\cdot\|b\|^{2}_{2}-a^{\top}b+(a^{\top}b)\cdot\|b\|^{2}_{2}=a^{\top}b.$
(B.14)

Thus, we obtain the following symmetry property,

$\displaystyle c=2(a^{\top}b)-a,\qquad a=2(c^{\top}b)-c.$ (B.15)

We are now ready to calculate the desired integration. Note that

$\displaystyle\int_{\mathbb{S}^{d-1}}a\cdot\exp(a^{\top}b){\mathrm{d}}a=b\cdot\int_{\mathbb{S}^{d-1}}(a^{\top}b)\exp(a^{\top}b){\mathrm{d}}a+\int_{\mathbb{S}^{d-1}}\bigl{(}a-(a^{\top}b)\cdot
b\bigr{)}\cdot\exp(a^{\top}b){\mathrm{d}}a.$ (B.16)

Meanwhile, it follows from the definition of $c$ in (B.13) and the properties
in (B.14) and (B.15) that

$\displaystyle\int_{\mathbb{S}^{d-1}}\bigl{(}a-(a^{\top}b)\cdot
b\bigr{)}\cdot\exp(a^{\top}b){\mathrm{d}}a$
$\displaystyle=-\int_{\mathbb{S}^{d-1}}\bigl{(}c-(c^{\top}b)\cdot
b\bigr{)}\cdot\exp(c^{\top}b){\mathrm{d}}a$
$\displaystyle=-\int_{\mathbb{S}^{d-1}}\bigl{(}c-(c^{\top}b)\cdot
b\bigr{)}\cdot\exp(c^{\top}b){\mathrm{d}}c,$ (B.17)

where the second equality follows from the fact that

$\displaystyle{\mathrm{d}}c=2\|b\|^{2}_{2}{\mathrm{d}}a-{\mathrm{d}}a={\mathrm{d}}a.$

Thus, by changing the variable on the right-hand side of (B.5.1), we obtain
that

$\displaystyle\int_{\mathbb{S}^{d-1}}\bigl{(}a-(a^{\top}b)\cdot
b\bigr{)}\cdot\exp(a^{\top}b){\mathrm{d}}a=0$ (B.18)

Finally, by plugging (B.18) into (B.16), we obtain that

$\displaystyle\int_{\mathbb{S}^{d-1}}a\cdot\exp(a^{\top}b){\mathrm{d}}a=b\cdot\int_{\mathbb{S}^{d-1}}(a^{\top}b)\exp(a^{\top}b){\mathrm{d}}a.$

Thus, by setting

$\displaystyle
C_{1}=\int_{\mathbb{S}^{d-1}}(a^{\top}b)\exp(a^{\top}b){\mathrm{d}}a,\quad\forall
b\in\mathbb{S}^{d-1},$

we conclude the proof of Lemma B.2. ∎

##  Appendix C Generalization Error Analysis

In this section, we analyze the generalization error for the complete setup of
the transformer architecture, which involves multiple layers, skip
connections, and multihead attention. We collect the notations used throughout
this section as follows. Notations. Let $r\geq 1$. We denote by
$\|\cdot\|_{r}$ the typical $\ell_{r}$-norm when it operates on a vector.
Also, let $s\geq 1$, and $M=(m_{1},\dots,m_{d_{2}})\in\mathbb{R}^{d_{1}\times
d_{2}}$, where $m_{i}\in\mathbb{R}^{d_{1}\times 1}$ for $i\in[d_{2}]$. We
define the matrix $(r,s)$-norm as
$\|M\|_{r,s}=(\sum_{i=1}^{d_{2}}\|m_{i}\|^{s}_{r})^{1/2}$. We define the
matrix $(r,s)$\- operator norm $\|\cdot\|_{r\to s}$ as $\|M\|_{r\to
s}=\sup_{u\in\mathbb{R}^{d_{2}\times 1}}\|Mu\|_{s}/\|u\|_{r}$ and write
$\|\cdot\|_{r}=\|\cdot\|_{r\to r}$ when the norm operates on a matrix. For
positive reals $r$ and $s$ such that $1/r+1/s=1$, we call $(r,s)$ a conjugate
pair.

###  C.1 Complete Setup of the Transformer Architecture

In what follows, we specify the complete setup of the $T$-layer transformer
parameterized by $\theta=(\theta_{0},\theta^{(0)},\dots,\theta^{(T-1)})$,
where the $t$-th layer is parameterized by $\theta^{(t)}=(A^{(t)},W^{(t)})$
and the aggregation layer is parameterized by $\theta_{0}$. We define the two-
layer feedforward network ${\mathtt{ffn}}$ with skip connection (without bias
terms) as follows,

$\displaystyle
X^{(t)}={\mathtt{ffn}}(X_{\star}^{(t)};A^{(t)})={\mathtt{ReLU}}(X_{\star}^{(t)}A^{{\mathrm{x}},(t)})A^{\sigma,(t)}+X_{\star}^{(t)},\qquad
A^{(t)}=(A^{{\mathrm{x}},(t)},A^{\sigma,(t)}),$

which is parameterized by $A^{(t)}$. Here
$A^{{\mathrm{z}},(t)}\in\mathbb{R}^{d\times
d_{\sigma}},A^{\sigma,(t)}\in\mathbb{R}^{d_{\sigma}\times d}$ and
${\mathtt{ReLU}}(\cdot)$ is the rectified linear unit (ReLU) that operates
elementwise. Corresponding to the softmax attention mechanism defined in
(4.10), we define the sequence to sequence softmax attention mechanism for
matrices as follows,

$\displaystyle\mathtt{attn}_{{\mathtt{SM}}}(Q,K,V)=\mathtt{softmax}\bigl{(}\mathfrak{K}(Q,K)\bigr{)}V\in\mathbb{R}^{L\times
d},$ (C.1)

where $Q,K\in\mathbb{R}^{L\times d_{\mathrm{p}}}$, $V\in\mathbb{R}^{L\times
d}$, and
$\mathtt{softmax}(\mathfrak{K}(Q,K))=(\mathtt{softmax}(\mathfrak{K}(K,q^{\ell})))^{\top}_{\ell\in[L]}\in\mathbb{R}^{L\times
L}$. With a slight abuse of notations, we define the sequence to sequence
counterpart of ${\mathtt{mha}}$ defined in (B.2) as follows,

$\displaystyle{\mathtt{mha}}(X;W)=\sum_{i=1}^{h}{\mathtt{head}}_{i}=\sum_{i=1}^{h}\mathtt{attn}_{\mathtt{SM}}(Q_{i},K_{i},V_{i}),$

where $Q_{i}=XW_{i}^{\mathrm{q}}$, $K_{i}=XW^{\mathrm{k}}_{i}$, and
$V_{i}=XW^{\mathrm{v}}_{i}$. Then we define the multihead attention with skip
connection as follows,

$\displaystyle\widetilde{X}_{\star}^{(t+1)}={\mathtt{mha}}(X^{(t)};W^{(t)})+X^{(t)},\qquad
W^{(t)}=\bigl{\\{}(W^{{\mathrm{q}},(t)}_{i},W^{{\mathrm{k}},(t)}_{i},W^{{\mathrm{v}},(t)}_{i})\bigr{\\}}_{i\in[h]},$

which is parameterized by $W^{(t)}$.

Empirical Function Class of Transformers. To construct the function class
$\mathcal{F}^{L}_{\mathtt{trans}}$ of transformers, we first specify the base
case
$\mathcal{F}_{\mathfrak{W}}^{L,(0)}=\\{(X^{(0)}_{\star})^{\top}\\}=\\{X^{\top}\\}$,
which is the function class formed by the input sequence $X$. Then we
recursively define the function classes as follows,

$\displaystyle\mathcal{F}^{L,(t)}_{\mathfrak{A}}$
$\displaystyle=\bigl{\\{}{\mathtt{ffn}}(X_{\star}^{(t)};A^{(t)})^{\top}:A^{(t)}\in{\mathfrak{A}}^{(t)},X_{\star}^{(t)}\in\mathcal{F}_{\mathfrak{W}}^{L,(t)}\bigr{\\}},$
$\displaystyle\mathcal{F}_{\mathfrak{W}}^{L,(t+1)}$
$\displaystyle=\bigl{\\{}{\mathtt{mha}}(X^{(t)};W^{(t)})^{\top}+(X^{(t)})^{\top}:W^{(t)}\in{\mathfrak{W}}^{(t)},X^{(t)}\in\mathcal{F}_{\mathfrak{A}}^{L,(t)}\bigr{\\}},$

where $\mathfrak{W}^{(t)}$ and $\mathfrak{A}^{(t)}$ are the parameter spaces
for the layer $t$. Then we define the function class of $T$-layer transformers
and its empirical counterpart as follows,

$\displaystyle\mathcal{F}^{L}_{\mathtt{trans}}=\\{\mathtt{agg}^{\mathtt{trans}}_{\theta_{0}}\circ
F:F\in\mathcal{F}^{L,(T)}_{\mathfrak{A}}\\},\qquad\mathcal{F}_{\mathtt{trans}}^{L,\mathcal{D}_{n}}=\Bigl{\\{}\bigl{(}f(X_{1}),\dots,f(X_{n})\bigr{)}:f\in\mathcal{F}^{L}_{\mathtt{trans}}\Bigr{\\}}.$
(C.2)

Note that the aggregation layer $\mathtt{agg}^{\mathtt{trans}}_{\theta_{0}}$
in (C.2) involves the input mask ${\mathtt{msk}}$, which describes the
prediction task. For example, in the complete setup of ViT, the aggregation
layer $\mathtt{agg}^{\mathtt{trans}}_{\theta_{0}}$ is the composition of a
function that extracts output corresponding to the input mask ${\mathtt{msk}}$
from the sequence to sequence softmax attention mechanism defined in (C.1) and
another function that computes the prediction of the target variable $y$
(e.g., by linear probing or fine tuning).

Assumptions. In parallel to (5.4), we make the following assumption on the
kernel function used in the multihead attention ${\mathtt{mha}}$.

######  Assumption C.1 (RBF Kernel).

We assume the the multihead attention ${\mathtt{mha}}$ adopts the Gaussian RBF
kernel $\mathfrak{K}(q,k)=\exp(-\|q-k\|_{2}^{2}/2\sigma^{2})$ with
$\sigma=(2d_{\mathrm{p}})^{1/s}$.

Note that the kernel function in Assumption C.1 is a more general version of
the RBF kernel in (5.4). Also, we make the following assumption on the
aggregation layer $\mathtt{agg}^{\mathtt{trans}}_{\theta_{0}}$.

######  Assumption C.2 (Aggregation Layer).

We assume that the function aggregation function
$\mathtt{agg}^{\mathtt{trans}}_{\theta_{0}}:\mathbb{R}^{L\times
d}\to\mathfrak{Y}$ is $1$-Lipschitz continuous with respect to
$\|\cdot\|_{r,\infty}$-norm. Here $\mathfrak{Y}$ is defined in (5.2).

In parallel to Assumption 5.1, we make the following assumption on the
parameter space for each layer of the complete setup of the transformer.

######  Assumption C.3 (Parameter Space).

Let the conjugate pair $(r,s)$ be given. For all $0\leq t\leq T-1$, we assume
that the parameter spaces for the $t$-th composite layer of the transformer
are

$\displaystyle{\mathfrak{A}}^{(t)}$
$\displaystyle=\Bigl{\\{}(A^{{\mathrm{x}},(t)},A^{\sigma,(t)})\in\mathbb{R}^{d\times
d_{\sigma}}\times\mathbb{R}^{d_{\sigma}\times
d}:\bigl{\|}(A^{{\mathrm{x}},(t)})^{\top}\bigr{\|}_{r}\leq{\alpha}^{{\mathrm{x}},(t)},\bigl{\|}(A^{{\mathrm{x}},(t)})^{\top}\bigr{\|}_{r,s}\leq{R}^{{\mathrm{x}},(t)},$
$\displaystyle\qquad\bigl{\|}(A^{\sigma,(t)})^{\top}\bigr{\|}_{r}\leq{\alpha}^{\sigma,(t)},\bigl{\|}(A^{\sigma,(t)})^{\top}\bigr{\|}_{r,s}\leq
R^{\sigma,(t)}\Bigr{\\}},$ (C.3) $\displaystyle\mathfrak{W}^{(t)}$
$\displaystyle=\Bigl{\\{}\bigl{\\{}(W^{{\mathrm{q}},(t)}_{i},W^{{\mathrm{k}},(t)}_{i},W^{{\mathrm{v}},(t)}_{i})\bigr{\\}}_{i\in[h]}:(W^{{\mathrm{q}},(t)}_{i},W^{{\mathrm{k}},(t)}_{i},W^{{\mathrm{v}},(t)}_{i})\in\mathbb{R}^{d\times
d_{\mathrm{p}}}\times\mathbb{R}^{d\times
d_{\mathrm{p}}}\times\mathbb{R}^{d\times d},$
$\displaystyle\qquad\bigl{\|}(W^{{\mathrm{q}},(t)}_{i})^{\top}\bigr{\|}_{r}\leq\omega^{{\mathrm{q}},(t)}_{i},\bigl{\|}(W^{{\mathrm{k}},(t)}_{i})^{\top}\bigr{\|}_{r}\leq\omega^{{\mathrm{k}},(t)}_{i},\bigl{\|}(W^{{\mathrm{v}},(t)}_{i})^{\top}\bigr{\|}_{r}\leq\omega^{{\mathrm{v}},(t)}_{i},$
(C.4)
$\displaystyle\qquad\bigl{\|}(W^{{\mathrm{q}},(t)}_{i})^{\top}\bigr{\|}_{r,s}\leq
R^{{\mathrm{q}},(t)}_{i},\bigl{\|}(W^{{\mathrm{k}},(t)}_{i})^{\top}\bigr{\|}_{r,s}\leq
R^{{\mathrm{k}},(t)}_{i},\bigl{\|}(W^{{\mathrm{v}},(t)}_{i})^{\top}\bigr{\|}_{r,s}\leq
R^{{\mathrm{v}},(t)}_{i},\forall i\in[h]\Bigr{\\}}.$

###  C.2 Generalization Error Analysis of Transformer

For ease of presentation, we define the following quantities combining the
parameter bounds across the $h$ heads in layer $t$ as follows,

$\displaystyle{\omega}^{{\mathrm{v}},(t)}$
$\displaystyle=\sum_{i=1}^{h}\omega_{i}^{{\mathrm{v}},(t)},\qquad\mkern
1.5mu\overline{\mkern-1.5mu\omega\mkern-1.5mu}\mkern
1.5mu^{{\mathrm{q}}{\mathrm{k}},(t)}=\max_{i\in[h]}\\{\omega_{i}^{{\mathrm{q}},(t)}+\omega_{i}^{{\mathrm{k}},(t)}\\},$
$\displaystyle R^{{\mathrm{v}},(t)}$
$\displaystyle=\sum_{i=1}^{h}R_{i}^{{\mathrm{v}},(t)},\qquad{R}^{{\mathrm{q}}{\mathrm{k}},(t)}=\sum_{i=1}^{h}(R_{i}^{{\mathrm{q}},(t)}+R_{i}^{{\mathrm{k}},(t)}),$
(C.5)

where $0\leq t\leq T-1$. Then we define

$\displaystyle\widetilde{\alpha}^{(t)}=1+\alpha^{{\mathrm{x}},(t)}\alpha^{\sigma,(t)},\qquad\widetilde{\omega}^{{\mathrm{v}},(t)}=1+\omega^{{\mathrm{v}},(t)},\qquad
R^{(t)}=R^{(0)}\cdot\prod_{\tau=0}^{t-1}\widetilde{\omega}^{{\mathrm{v}},(\tau)}\widetilde{\alpha}^{(\tau)},$
(C.6)

which characterize the magnitude of the intermediate input sequences. Also, we
define

$\displaystyle\kappa^{(t)}=\max\biggl{\\{}\frac{\alpha^{{\mathrm{x}},(t)}R^{\sigma,(t)}+\alpha^{\sigma,(t)}R^{{\mathrm{x}},(t)}}{\widetilde{\alpha}^{(t)}},\frac{R^{{\mathrm{v}},(t)}}{\widetilde{\omega}^{{\mathrm{v}},(t)}},\frac{{R}^{{\mathrm{q}}{\mathrm{k}},(t)}}{\mkern
1.5mu\overline{\mkern-1.5mu\omega\mkern-1.5mu}\mkern
1.5mu^{{\mathrm{q}}{\mathrm{k}},(t)}\omega^{{\mathrm{v}},(t)}}\biggr{\\}},\qquad\zeta^{(t)}=\frac{(\mkern
1.5mu\overline{\mkern-1.5mu\omega\mkern-1.5mu}\mkern
1.5mu^{{\mathrm{q}}{\mathrm{k}},(t)})^{2}R^{{\mathrm{v}},(t)}}{\widetilde{\omega}^{{\mathrm{v}},(t)}}.$
(C.7)

We further simplify the quantities defined in (C.6) and (C.7) to define

$\displaystyle\gamma=\max_{0\leq t\leq
T-1}\\{\alpha^{{\mathrm{x}},(t)}\alpha^{\sigma,(t)},\omega^{{\mathrm{v}},(t)}\\},\qquad\kappa=\max_{0\leq
t\leq T-1}\kappa^{(t)},\qquad\zeta=\max_{0\leq t\leq T-1}\zeta^{(t)}.$ (C.8)

Recall that the generalization error $\mathcal{E}_{\rm gen}$ is defined in
(5.3). We give the following characterization of generalization error of the
complete setup of the transformer.

######  Theorem C.4 (Generalization Error of Transformer).

Suppose that Assumptions C.1-C.3 and Assumption 5.2 hold. For any $\delta>0$,
it holds with probability at least $1-\delta$ that,

$\displaystyle\mathcal{E}_{\rm gen}\leq
O\biggl{(}\frac{D\sqrt{hT}}{\sqrt{n}}\cdot\Bigl{[}T\sqrt{\log(2+\gamma)}+\sqrt{T}\sqrt{\log(1+\zeta
R^{(0)})}+\sqrt{\log(1+\kappa/\zeta)}\Bigr{]}+\sqrt{\frac{\log(1/\delta)}{n}}\biggr{)},$

where $D=\max\\{d,d_{\sigma},d_{\mathrm{p}}\\}$ and $\widehat{E}[\,\cdot\,]$
is the empirical expectation over the dataset $\mathcal{D}_{n}$.

###### Proof.

See §C.4 for a detailed proof. ∎

Technical Highlight. The key difference between our analysis and that in
Edelman et al. (2021) lies in our exploitation of the benefit brought by the
invariance/equivariance property of the transformer. Specifically, due to the
invariance/equivariance property of the transformer, the dimensions of the
parameters $W^{(t)}$ and $A^{(t)}$ can be kept at constant levels regardless
of $L$. Therefore, we build our covering number of
$\mathcal{F}_{\mathtt{trans}}^{L,\mathcal{D}_{n}}$ based only on the parameter
space covering and its propagation. As a consequence, the covering number only
scales with the dimension $D$ of the parameters and the number $T$ of layers
in the transformer. In contrast to our results, the generalization error bound
presented in Edelman et al. (2021) has (i) logarithmic dependency over $L$,
and (ii) assumption that intermediate input sequences $X^{(t)}$ has low norm.

Interpretation. We interpret the generalization error in Theorem C.4 as
follows. The dependencies $O(\sqrt{\log(1/\delta)/n})$ and $O(1/\sqrt{n})$
over the size $n$ of the dataset $\mathcal{D}_{n}$ are standard in the
generalization analysis. On the other hand, the scaling $O(Dh^{1/2}T^{3/2})$
implies that the transformer needs more data to generalize as the parameter
dimensions, the number of layers, or the number of heads grow. Although our
analysis falls under the framework proposed in Bartlett et al. (2017), the
generalization error in Theorem C.4 only scales polynomially with the number
of layers $T$ of the transformer, as opposed to the typical exponential
scaling with the number of layers. Moreover, note that the dependency over
$\gamma$, $\kappa$, and $\zeta$ are all logarithmic, which implies that our
generalization error remains polynomial as long as these quantities do not
scale doubly exponentially with $D$, $h$, or $T$.

Transformer Architecture. From Theorem C.4, we see that $\alpha$, $\omega$,
$\kappa$, $\zeta$ and $R^{\mathrm{x}}$ play crucial roles in the
generalization error of the transformer. In specific, we observe that (i) skip
connections make all layers relatively close to identity maps (Bartlett et
al., 2018b, a; Hardt and Ma, 2016), which help reducing $\alpha$, $\omega$,
$\kappa$ and $\zeta$, and (ii) layer normalizations help reducing the
magnitudes $R^{{\mathrm{x}},(t)}$ of the intermediate input sequences
$X^{(t)}$.

Simplification to Theorem 5.3. Theorem C.4 characterizes the generalization
error for the complete setup of the transformer, which is a more general
result than that in Theorem 5.3. In what follows, we simplify Theorem C.4 to
obtain Theorem 5.3.

  * •

Single layer and signlehead attention: We set $h=1$, $T=1$, and $R=R^{(0)}$.

  * •

Skip connections: Since there is no skip connections, we

    * –

replace $\widetilde{\alpha}^{(t)}$ and $\widetilde{\omega}^{{\mathrm{v}},(t)}$
with ${\alpha}^{{\mathrm{x}},(t)}{\alpha}^{{\mathrm{z}},(t)}$ and
$\omega^{{\mathrm{v}},(t)}$ in (C.6), respectively, and

    * –

replace $\log(2+\gamma)$ with $\log(1+\gamma)$.

  * •

Feedforward neural network $\mathtt{nn}$: Since there is no linear
transformation at the second layer, we

    * –

remove $\alpha^{\sigma,(0)}$ and $R^{\sigma,(0)}$,

    * –

remove the intermediate output dimension $d_{\sigma}$, and

    * –

set $\alpha=\alpha^{{\mathrm{x}},(0)}$ and
$R^{\mathtt{nn}}=R^{{\mathrm{x}},(0)}$ correspondingly.

  * •

Softmax Attention $\mathtt{attn}_{\mathtt{SM}}$: Since there is only one layer
and one head, we set

    * –

$\omega=\max\\{\mkern 1.5mu\overline{\mkern-1.5mu\omega\mkern-1.5mu}\mkern
1.5mu^{{\mathrm{q}}{\mathrm{k}},(0)},\omega^{{\mathrm{v}},(0)}\\}$ and

    * –

$R^{\mathtt{attn}}=\max\\{R^{{\mathrm{q}}{\mathrm{k}},(0)},R^{{\mathrm{v}},(0)}\\}$.

  * •

Aggregation layer $\mathtt{agg}^{\mathtt{trans}}_{\theta}$: We view the
aggregation layer in (C.2) as the composition (i.e., $a\circ b$) of the
following two parts,

    * (a)

the aggregation function $\mathtt{agg}^{\mathtt{trans}}_{\theta_{0}}$ defined
in (5.1) and

    * (b)

a function taking the sequence to sequence softmax attention mechanism defined
in (C.1) as the input and gives the $d_{\mathrm{p}}$-dimensional output
corresponding to the query $q_{\theta}({\mathtt{msk}})$.

With the factors in (C.2)-(C.8) modified accordingly, we obtain Theorem 5.3.

Proof Sketch. We organize the proof of Theorem C.4 as follows.

  * (§C.3)

We introduce the framework of bounding generalization error via Rademacher
complexity, which requires a covering number bound for the empirical function
class $\mathcal{F}_{{\mathtt{trans}}}^{L,\mathcal{D}_{n}}$ of transformers.

  * (§C.4)

We provide a covering number bound for the empirical function class
$\mathcal{F}_{{\mathtt{trans}}}^{L,\mathcal{D}_{n}}$ of transformers and
sketch the proof of Theorem C.4 based on the framework introduced in §C.3.

  * (§C.5)

We give the detailed proof of the covering number bound for
$\mathcal{F}_{{\mathtt{trans}}}^{L,\mathcal{D}_{n}}$ by (i) analyzing the
covering of each ${\mathtt{mha}}$ and ${\mathtt{ffn}}$ and (ii) characterizing
the propagation of the coverings of them.

  * (§F)

For clarity of the presentation, we leave all the omitted proofs of the
intermediate lemmas to §F.

###  C.3 Fundation of Generalization

In this section, we introduce the fundamental building blocks for our analysis
of generalization error.

####  C.3.1 Rademacher Complexity

We assume that the dataset $\mathcal{D}_{n}=\\{(X_{i},y_{i})\\}_{i=1}^{n}$ is
drawn independently and identically from the data distribution $\mathcal{D}$.
We characterize the generalization error of learning within the function class
$\mathcal{F}$ via the empirical Rademacher complexity

$\displaystyle\widehat{\mathcal{R}}_{\mathcal{D}_{n}}(\mathcal{F})=\mathbb{E}_{\epsilon\in\\{-1,+1\\}^{n}}\biggl{[}\sup_{f\in\mathcal{F}}\frac{1}{n}\sum_{i=1}^{n}\epsilon_{i}\cdot\mathcal{L}\bigl{(}(X_{i},y_{i}),f\bigr{)}\biggr{]}.$
(C.9)

Here $\mathcal{L}$ is a fixed learning objective. The generalization error
$\mathcal{E}_{\rm gen}$ defined in (5.3) is connected to the empirical
Rademacher complexity through the following classic result.

######  Theorem C.5 (Mohri et al. (2018)).

Let $\mathcal{F}$ be a family of functions mapping from the data distribution
$\mathcal{D}$ to the output range $[0,1]$ through a fixed $1$-Lipschitz
objective function $\mathcal{L}$. Then, for any $\delta>0$, with probability
at least $1-\delta$ over the independent and identical draw of a dataset
$\mathcal{D}_{n}=\\{(X_{i},y_{i})\\}_{i\in[n]}$, the following holds for all
$f\in\mathcal{F}$,

$\displaystyle\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f\bigr{)}\Bigr{]}-\widehat{\mathbb{E}}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f\bigr{)}\Bigr{]}\leq
2\widehat{\mathcal{R}}_{\mathcal{D}_{n}}(\mathcal{F})+3\sqrt{\frac{\log(2/\delta)}{2n}},$

where $\widehat{E}[\,\cdot\,]$ is the empirical expectation taken over the
dataset $\mathcal{D}_{n}$ and the empirical Rademacher complexity
$\widehat{\mathcal{R}}_{\mathcal{D}_{n}}(\mathcal{F})$ is defined in (C.9).

####  C.3.2 Connection to Covering Number

In what follows, we present the connection between the Rademacher complexity
defined in (C.9) and the covering number of the empirical function class.
Given a dataset $\mathcal{D}_{n}=\\{(X_{i},y_{i})\\}_{i\in[n]}$ and a function
class $\mathcal{F}$, we define the corresponding empirical function class as
$\mathcal{F}^{\mathcal{D}_{n}}=\\{(f(X_{1}),\dots,f(X_{n})):f\in\mathcal{F}\\}$.
We define the proper covering number as follows.

######  Definition C.6 (Proper Covering Number).

Let $N({\mathcal{S}},\varepsilon,\|\cdot\|)$ denote the least cardinality of
any subset ${\mathcal{T}}\subseteq{\mathcal{S}}$ that covers ${\mathcal{S}}$
at resolution $\varepsilon$ with respect to norm $\|\cdot\|$. Formally, we
define

$\displaystyle
N\bigl{(}{\mathcal{S}},\varepsilon,\|\cdot\|\bigr{)}=\inf\Bigl{\\{}{\rm
card}({\mathcal{T}}):\sup_{S\in{\mathcal{S}}}\inf_{T\in{\mathcal{T}}}\|S-T\|\leq\varepsilon,{\mathcal{T}}\subseteq{\mathcal{S}}\Bigr{\\}},$

where ${\rm card}({\mathcal{T}})$ is the cardinality of the set
${\mathcal{T}}$.

The empirical Rademacher complexity
$\widehat{\mathcal{R}}_{\mathcal{D}_{n}}(\mathcal{F})$ defined in (C.9) can be
bounded via the covering number of the empirical function class
$\mathcal{F}^{\mathcal{D}_{n}}$. Specifically, we have the following version
of Dudley entropy integral lemma adapted to the
$\|\cdot\|_{r,\infty}$-covering.

######  Lemma C.7 (Dudley Entropy Integral for
$\|\cdot\|_{r,\infty}$-Covering).

Let $\mathcal{F}$ be a real-valued function class taking values in $[0,1]$,
and assume that $0\in\mathcal{F}$. Then we have

$\displaystyle\widehat{\mathcal{R}}_{\mathcal{D}_{n}}(\mathcal{F})\leq\inf_{\xi\in(0,1)}\biggl{(}4\xi+\frac{12}{\sqrt{n}}\int_{\xi}^{1}\sqrt{\log
N\bigl{(}\mathcal{F}^{\mathcal{D}_{n}},\varepsilon,\|\cdot\|_{r,\infty}\bigr{)}}{\mathrm{d}}\varepsilon\biggr{)}.$

###### Proof.

See §F.1 for a detailed proof. ∎

Note that Lemma C.7 considers the $\|\cdot\|_{r,\infty}$-covering, while the
standard version of Dudley entropy integral lemma (see, e.g., Mohri et al.
(2018)) considers the $\|\cdot\|_{2}$-covering (or
$\|\cdot\|_{\text{F}}$-covering). The key technical difference between Lemma
C.7 and the standard version of Dudley entropy integral lemma is that the
upper bound of the integral is now $1$ instead of $\sqrt{n}$ due to the
reduced upper bound of $\|\cdot\|_{r,\infty}$-norm compared to
$\|\cdot\|_{2}$-norm (or $\|\cdot\|_{\text{F}}$-norm).

To this end, it remains to bound the covering number of the empirical function
class $\mathcal{F}_{\mathtt{trans}}^{L,\mathcal{D}_{n}}$ induced by function
class $\mathcal{F}^{L}_{\mathtt{trans}}$ of transformers with respect to
$\|\cdot\|_{r,\infty}$-norm.

###  C.4 Generalization Error Bound

Recall that $\mathfrak{W}^{(t)}$ and $\mathfrak{A}^{(t)}$ are defined in (C.4)
and (C.3), respectively. We characterize of the covering number of the
empirical function class $\mathcal{F}^{L,\mathcal{D}_{n}}_{\mathtt{trans}}$ of
transformers in the following lemma.

######  Lemma C.8 (Covering of Transformer).

Let $D=\max\\{d,d_{\sigma},d_{\mathrm{p}}\\}$ and
$R^{(0)}=\sup_{i\in[n]}\|X_{i}^{\top}\|_{r,\infty}$. We have

$\displaystyle\log
N\bigl{(}\mathcal{F}^{L,\mathcal{D}_{n}}_{\mathtt{trans}},\varepsilon,\|\cdot\|_{r,\infty}\bigr{)}\leq(4+h)D^{2}T\cdot\log\biggl{(}1+\frac{2{R}^{{\mathtt{trans}},(T)}R^{(T)}}{\varepsilon}\biggr{)},$

where

$\displaystyle{R}^{{\mathtt{trans}},(T)}=\sum_{t=0}^{T-1}\biggl{[}\biggl{(}\frac{{R}^{{\mathtt{mha}},(t)}}{\widetilde{\rho}^{(t)}}+\frac{\alpha^{{\mathrm{x}},(t)}R^{\sigma,(t)}+\alpha^{\sigma,(t)}R^{{\mathrm{x}},(t)}}{\widetilde{\alpha}^{(t)}}\biggr{)}\cdot\prod_{\tau=t+1}^{T-1}\frac{\widetilde{\rho}^{(\tau)}}{\widetilde{\omega}^{{\mathrm{v}},(\tau)}}\biggr{]}.$
(C.10)

Here

$\displaystyle\widetilde{\rho}^{(t)}=\widetilde{\omega}^{{\mathrm{v}},(t)}+(\mkern
1.5mu\overline{\mkern-1.5mu\omega\mkern-1.5mu}\mkern
1.5mu^{{\mathrm{q}}{\mathrm{k}},(t)})^{2}\omega^{{\mathrm{v}},(t)}\cdot(R^{(t)})^{2},\qquad
R^{{\mathtt{mha}},(t)}=R^{{\mathrm{v}},(t)}+\mkern
1.5mu\overline{\mkern-1.5mu\omega\mkern-1.5mu}\mkern
1.5mu^{{\mathrm{q}}{\mathrm{k}},(t)}{R}^{{\mathrm{q}}{\mathrm{k}},(t)}\cdot(R^{(t)})^{2},$
(C.11)

where $R^{(t)}$ is defined in (C.6) and other factors are defined in (C.2).

###### Proof.

See §C.5.3 for a detailed proof. ∎

With Lemma C.8, we establish the generalization error of the transformer in
Theorem C.4 as follows.

###### Proof of Theorem C.4.

For any constant $R>0$, it holds that

$\displaystyle\int_{\xi}^{1}\sqrt{\log(1+R/\varepsilon)}{\mathrm{d}}\varepsilon$
$\displaystyle\leq\int_{\xi}^{1}\bigl{[}\sqrt{\log(1+R)}+\sqrt{\log(1+1/\varepsilon)}\bigr{]}{\mathrm{d}}\varepsilon$
$\displaystyle\leq(1-\xi)\cdot\sqrt{\log(1+R)}+\int_{\xi}^{1}1/\sqrt{\varepsilon}{\mathrm{d}}\varepsilon$
$\displaystyle=(1-\xi)\cdot\sqrt{\log(1+R)}+2-2\sqrt{\xi},$

where the first inequality follows from $\sqrt{a+b}\leq\sqrt{a}+\sqrt{b}$ and
the second inequality follows from $\log(1+a)\leq a$ for $a>-1$. By Lemma C.7,
we have

$\displaystyle\widehat{\mathcal{R}}_{\mathcal{D}_{n}}(\mathcal{F}^{L}_{\mathtt{trans}})$
$\displaystyle\quad\leq\inf_{\xi\in(0,1)}\biggl{(}4\xi+\frac{12}{\sqrt{n}}\int_{\xi}^{1}\sqrt{\log
N\bigl{(}\mathcal{F}_{\mathtt{trans}}^{L,\mathcal{D}_{n}},\varepsilon,\|\cdot\|_{r,\infty}\bigr{)}}{\mathrm{d}}\varepsilon\biggr{)}$
$\displaystyle\quad\leq
12D\sqrt{(4+h)T}\cdot\frac{2+\sqrt{\log(1+2{R}^{{\mathtt{trans}},(T)}R^{(T)})}}{\sqrt{n}}$
$\displaystyle\quad\qquad+\inf_{\xi\in(0,1)}\biggl{[}\biggl{(}4-\frac{12D\sqrt{(4+h)T}\cdot\sqrt{\log(1+2{R}^{{\mathtt{trans}},(T)}R^{(T)})}}{\sqrt{n}}\biggr{)}\cdot\xi-24D\sqrt{(4+h)T}\cdot\sqrt{\xi}\biggr{]}.$

For the infimum on the right-hand side of the above inequality, we can set
$\xi\to 0$ to obtain

$\displaystyle\widehat{\mathcal{R}}_{\mathcal{D}_{n}}(\mathcal{F}^{L}_{\mathtt{trans}})$
$\displaystyle\leq
12D\sqrt{(4+h)T}\cdot\frac{2+\sqrt{\log(1+2{R}^{{\mathtt{trans}},(T)}R^{(T)})}}{\sqrt{n}}$
$\displaystyle=O\biggl{(}D\sqrt{hT}\cdot\frac{1+\sqrt{\log(1+{R}^{{\mathtt{trans}},(T)}R^{(T)})}}{\sqrt{n}}\biggr{)},$

where the second line follows from Lemma G.4. Finally, taking the above back
into Theorem C.5, we obtain

$\displaystyle\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f\bigr{)}\big{/}4\Bigr{]}-\widehat{\mathbb{E}}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f\bigr{)}\big{/}4\Bigr{]}$
$\displaystyle\quad\leq
O\biggl{(}D\sqrt{hT}\cdot\frac{1+\sqrt{\log(1+R^{(T)}R^{{\mathtt{trans}},(T)})}}{\sqrt{n}}+\sqrt{\frac{\log(1/\delta)}{n}}\biggr{)},$
(C.12)

for any $f\in\mathcal{F}_{\mathtt{trans}}^{L}$. By simplifying the terms in
(C.4) using Lemma G.5, we obtain

$\displaystyle\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f\bigr{)}\Bigr{]}-\widehat{\mathbb{E}}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f\bigr{)}\Bigr{]}$
(C.13) $\displaystyle\quad\leq
O\biggl{(}\frac{D\sqrt{hT}}{\sqrt{n}}\cdot\Bigl{[}T\sqrt{\log(2+\gamma)}+\sqrt{T}\sqrt{\log(1+\zeta
R^{(0)})}+\sqrt{\log(1+\kappa/\zeta)}\Bigr{]}+\sqrt{\frac{\log(1/\delta)}{n}}\biggr{)},$

for any $f\in\mathcal{F}^{L}_{\mathtt{trans}}$. Recall that the regression
function $\widetilde{f}$ minimizes the empirical risk
$\widehat{\mathbb{E}}[\mathcal{L}((X,y),f)]$. Let $\mkern
1.5mu\overline{\mkern-1.5muf\mkern-1.5mu}\mkern
1.5mu=\mathop{\mathrm{argmin}}_{f\in\mathcal{F}^{L}_{\mathtt{trans}}}\mathbb{E}[\mathcal{L}((X,y),f)]$.
By the definition of the generalization error $\mathcal{E}_{\rm gen}$ in
(5.3), we have

$\displaystyle\mathcal{E}_{\rm gen}$
$\displaystyle=\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}^{L}_{\mathtt{trans}}}\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f\bigr{)}\Bigr{]}$
$\displaystyle=\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{]}+\min_{f\in\mathcal{F}^{L}_{\mathtt{trans}}}\widehat{\mathbb{E}}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f\bigr{)}\Bigr{]}-\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),\mkern
1.5mu\overline{\mkern-1.5muf\mkern-1.5mu}\mkern 1.5mu\bigr{)}\Bigr{]}$
$\displaystyle=\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{]}+\min_{f\in\mathcal{F}^{L}_{\mathtt{trans}}}\widehat{\mathbb{E}}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f\bigr{)}\Bigr{]}-\widehat{\mathbb{E}}\Bigl{[}\mathcal{L}\bigl{(}(X,y),\mkern
1.5mu\overline{\mkern-1.5muf\mkern-1.5mu}\mkern 1.5mu\bigr{)}\Bigr{]}$
$\displaystyle\quad\qquad+\widehat{\mathbb{E}}\Bigl{[}\mathcal{L}\bigl{(}(X,y),\mkern
1.5mu\overline{\mkern-1.5muf\mkern-1.5mu}\mkern
1.5mu\bigr{)}\Bigr{]}-\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),\mkern
1.5mu\overline{\mkern-1.5muf\mkern-1.5mu}\mkern 1.5mu\bigr{)}\Bigr{]}$
$\displaystyle\leq\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),\mkern
1.5mu\overline{\mkern-1.5muf\mkern-1.5mu}\mkern
1.5mu\bigr{)}\Bigr{]}-\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),\mkern
1.5mu\overline{\mkern-1.5muf\mkern-1.5mu}\mkern 1.5mu\bigr{)}\Bigr{]}$
$\displaystyle\leq
O\biggl{(}\frac{D\sqrt{hT}}{\sqrt{n}}\cdot\Bigl{[}T\sqrt{\log(2+\gamma)}+\sqrt{T}\sqrt{\log(1+\zeta
R^{(0)})}+\sqrt{\log(1+\kappa/\zeta)}\Bigr{]}+\sqrt{\frac{\log(1/\delta)}{n}}\biggr{)},$

where the last line follows from (C.13). Therefore, we conclude the proof of
Theorem C.4. ∎

###  C.5 Proof of Lemma C.8

In this section, we organize the proof of Lemma C.8 as follows.

  * (§C.5.1)

We give covering number bounds for ${\mathtt{mha}}$ and ${\mathtt{ffn}}$ as
the building blocks for covering the empirical function class
$\mathcal{F}_{\mathtt{trans}}^{\mathcal{D}_{n}}$ of transformers.

  * (§C.5.2)

We illustrate how the coverings for ${\mathtt{mha}}$ and ${\mathtt{ffn}}$
propagate thought the layers in the transformer.

  * (§C.5.3)

We combine the results in §C.5.1 and §C.5.2 to give the covering number bound
for the empirical function class
$\mathcal{F}^{L,\mathcal{D}_{n}}_{\mathtt{trans}}$ of transformers.

####  C.5.1 Block Covering

In this section, we bound the covering numbers of each ${\mathtt{mha}}$ and
${\mathtt{ffn}}$ block. First, we give the following corollary, which serves
as the fundamental tool for establishing the covering bound of the
transformer.

######  Lemma C.9.

Let the conjugate pair $(r,s)$ be given. We have

$\displaystyle\log N\Bigl{(}\bigl{\\{}M^{\top}\in\mathbb{R}^{d_{2}\times
d_{1}}:\|M^{\top}\|_{r,s}\leq
R^{m}\bigr{\\}},\varepsilon,\|\cdot\|_{r,s}\Bigr{)}\leq
d_{1}d_{2}\cdot\log\biggl{(}1+\frac{2R^{\mathrm{m}}}{\varepsilon}\biggr{)}.$

###### Proof.

See §F.1.1 for a detailed proof. ∎

Covering for FFN. In the sequel, we establish the covering number for the
residual feedforward network ${\mathtt{ffn}}$. Let

$\displaystyle{\mathfrak{A}}$
$\displaystyle=\Bigl{\\{}(A^{\mathrm{x}},A^{\sigma})\in\mathbb{R}^{d\times
d_{\sigma}}\times\mathbb{R}^{d_{\sigma}\times d}:$
$\displaystyle\qquad\bigl{\|}({A}^{\mathrm{x}})^{\top}\bigr{\|}_{r}\leq{\alpha}^{\mathrm{x}},\bigl{\|}({A}^{\mathrm{x}})^{\top}\bigr{\|}_{r,s}\leq{R}^{\mathrm{x}},\bigl{\|}({A}^{\sigma})^{\top}\bigr{\|}_{r}\leq{\alpha}^{\sigma},\bigl{\|}({A}^{\sigma})^{\top}\bigr{\|}_{r,s}\leq
R^{\sigma}\Bigr{\\}}.$

We characterize the covering number of the function class
$\\{{\mathtt{ffn}}(X_{\star};A)^{\top}:A\in\mathfrak{A}\\}$ of FFN as follows.

######  Lemma C.10 (FFN Covering).

Let positive reals
$(\alpha^{\mathrm{x}},\alpha^{\sigma},R^{\mathrm{x}},R^{\sigma},\varepsilon)$
and positive integers $(d,d_{\sigma})$ be given. Also, let
$X_{\star}\in\mathbb{R}^{L\times d}$ be given. Then we have

$\displaystyle\log
N\Bigl{(}\bigl{\\{}{\mathtt{ffn}}(X_{\star};A)^{\top}:A\in\mathfrak{A}\bigr{\\}},\varepsilon,\|\cdot\|_{r,\infty}\Bigr{)}\leq
2dd_{\sigma}\cdot\log\biggl{(}1+\frac{2(\alpha^{\mathrm{x}}R^{\sigma}+\alpha^{\sigma}R^{\mathrm{x}})\cdot\|X_{\star}^{\top}\|_{r,\infty}}{\varepsilon}\biggr{)}.$

###### Proof.

For any $A=({A}^{\mathrm{x}},{A}^{\sigma})\in\mathfrak{A}$, let
$\widehat{A}=(\widehat{A}^{\mathrm{x}},\widehat{A}^{\sigma})\in\mathfrak{A}$
be such that

$\displaystyle\bigl{\|}(A^{\mathrm{x}})^{\top}-(\widehat{A}^{\mathrm{x}})^{\top}\bigr{\|}_{r,s}$
$\displaystyle\leq\varepsilon^{\mathrm{x}},\qquad\bigl{\|}(A^{\sigma})^{\top}-(\widehat{A}^{\sigma})^{\top}\bigr{\|}_{r,s}\leq\varepsilon^{\sigma}.$

Then we have

$\displaystyle\bigl{\|}{\mathtt{ffn}}(X_{\star};A)^{\top}-{\mathtt{ffn}}(X^{\star};\widehat{A})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\quad=\Bigl{\|}\bigl{(}{\mathtt{ReLU}}(X_{\star}A^{\mathrm{x}})A^{\sigma}\bigr{)}^{\top}+X_{\star}^{\top}-\bigl{(}{\mathtt{ReLU}}(X_{\star}\widehat{A}^{\mathrm{x}})\widehat{A}^{\sigma}\bigr{)}^{\top}-X_{\star}^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\quad\leq\Bigl{\|}\bigl{(}(A^{\sigma})^{\top}-(\widehat{A}^{\sigma})^{\top}\bigr{)}{\mathtt{ReLU}}(X_{\star}A^{\mathrm{x}})^{\top}\Bigr{\|}_{r,\infty}+\Bigl{\|}(\widehat{A}^{\sigma})^{\top}\bigl{(}{\mathtt{ReLU}}(X_{\star}\widehat{A}^{\mathrm{x}})-{\mathtt{ReLU}}(X_{\star}{A}^{\mathrm{x}})\bigr{)}^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\quad\leq\big{\|}(A^{\sigma})^{\top}-(\widehat{A}^{\sigma})^{\top}\bigr{\|}_{r,s}\cdot\bigl{\|}(X_{\star}A^{\mathrm{x}})^{\top}\bigr{\|}_{r,\infty}+\bigl{\|}(\widehat{A}^{\sigma})^{\top}\bigr{\|}_{r}\cdot\bigl{\|}(\widehat{A}^{\mathrm{x}}-{A}^{\mathrm{x}})^{\top}X_{\star}^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\quad\leq\varepsilon^{\sigma}\cdot\bigl{\|}(A^{\mathrm{x}})^{\top}\bigr{\|}_{r}\cdot\|X_{\star}^{\top}\|_{r,\infty}+R^{\sigma}\cdot\bigl{\|}(\widehat{A}^{\mathrm{x}})^{\top}-({A}^{\mathrm{x}})^{\top}\bigr{\|}_{r,s}\cdot\|X_{\star}^{\top}\|_{r,\infty}\leq(\varepsilon^{\sigma}\alpha^{\mathrm{x}}+\varepsilon^{\mathrm{x}}\alpha^{\sigma})\cdot\|X_{\star}^{\top}\|_{r,\infty},$

where the second inequality and the last line follow from Lemma G.3. Thus,
setting

$\displaystyle\varepsilon^{\mathrm{x}}=\frac{R^{\mathrm{x}}\varepsilon}{(\alpha^{\mathrm{x}}R^{\sigma}+\alpha^{\sigma}R^{\mathrm{x}})\cdot\|X_{\star}^{\top}\|_{r,\infty}},\qquad\varepsilon^{\sigma}=\frac{R^{\sigma}\varepsilon}{(\alpha^{\mathrm{x}}R^{\sigma}+\alpha^{\sigma}R^{\mathrm{x}})\cdot\|X_{\star}^{\top}\|_{r,\infty}},$

we obtain
$\|{\mathtt{ffn}}(X_{\star};A)^{\top}-{\mathtt{ffn}}(X_{\star};\widehat{A})^{\top}\|_{r,\infty}\leq\varepsilon$.
Then by Corollary C.9, we have

$\displaystyle\log
N\Bigl{(}\bigl{\\{}{\mathtt{ffn}}(X_{\star};A)^{\top}:A\in\mathfrak{A}\bigr{\\}},\varepsilon,\|\cdot\|_{r,\infty}\Bigr{)}$
$\displaystyle\quad\leq\log
N\Bigl{(}\bigl{\\{}(A^{\mathrm{x}})^{\top}\in\mathbb{R}^{d_{\sigma}\times
d}:\bigl{\|}(A^{\mathrm{x}})^{\top}\bigr{\|}_{r,s}\leq
R^{\mathrm{x}}\bigr{\\}},\varepsilon^{\mathrm{x}},\|\cdot\|_{r,\infty}\Bigr{)}$
$\displaystyle\quad\qquad+\log
N\Bigl{(}\bigl{\\{}(A^{\sigma})^{\top}\in\mathbb{R}^{d\times
d_{\sigma}}:\bigl{\|}(A^{\sigma})^{\top}\bigr{\|}_{r,s}\leq
R^{\sigma}\bigr{\\}},\varepsilon^{\sigma},\|\cdot\|_{r,\infty}\Bigr{)}$
$\displaystyle\quad\leq
2dd_{\sigma}\cdot\log\biggl{(}1+\frac{2(\alpha^{\mathrm{x}}R^{\sigma}+\alpha^{\sigma}R^{\mathrm{x}})\cdot\|X_{\star}^{\top}\|_{r,\infty}}{\varepsilon}\biggr{)}.$

Therefore, we conclude the proof of Lemma C.10. ∎

Covering for MHA. We consider the set of feasible weights

$\displaystyle\mathfrak{W}$
$\displaystyle=\Bigl{\\{}\bigl{\\{}(W^{{\mathrm{q}}}_{i},W^{{\mathrm{k}}}_{i},W^{{\mathrm{v}}}_{i})\bigr{\\}}_{i\in[h]}:(W^{{\mathrm{q}}}_{i},W^{{\mathrm{k}}}_{i},W^{{\mathrm{v}}}_{i})\in\mathbb{R}^{d\times
d_{\mathrm{p}}}\times\mathbb{R}^{d\times
d_{\mathrm{p}}}\times\mathbb{R}^{d\times d},$ (C.14)
$\displaystyle\quad\qquad\bigl{\|}(W^{{\mathrm{q}}}_{i})^{\top}\bigr{\|}_{r}\leq\omega^{{\mathrm{q}}}_{i},\bigl{\|}(W^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r}\leq\omega^{{\mathrm{k}}}_{i},\bigl{\|}(W^{{\mathrm{v}}}_{i})^{\top}\bigr{\|}_{r}\leq\omega^{{\mathrm{v}}}_{i},$
$\displaystyle\quad\qquad\bigl{\|}(W^{{\mathrm{q}}}_{i})^{\top}\bigr{\|}_{r,s}\leq
R^{{\mathrm{q}}}_{i},\bigl{\|}(W^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r,s}\leq
R^{{\mathrm{k}}}_{i},\bigl{\|}(W^{{\mathrm{v}}}_{i})^{\top}\bigr{\|}_{r,s}\leq
R^{{\mathrm{v}}}_{i},\forall i\in[h]\Bigr{\\}}.$

We first give the following lemma on the Lipschitz continuity of MHA with
respect to the parameters in $\mathfrak{W}$.

######  Lemma C.11 (Lipschitz Continuous MHA in Parameter).

Let $(r,s)$ be a conjugate pair with $1\leq r<\infty$. Suppose that
$X\in\mathbb{R}^{L\times d}$ satisfies $\|X^{\top}\|_{r,\infty}\leq R$. For
any
$W=\\{(W^{{\mathrm{q}}}_{i},W^{{\mathrm{k}}}_{i},W^{{\mathrm{v}}}_{i})\\}_{i\in[h]}\in\mathfrak{W}$,
let
$\widehat{W}=\\{(\widehat{W}^{{\mathrm{q}}}_{i},\widehat{W}^{{\mathrm{k}}}_{i},\widehat{W}^{{\mathrm{v}}}_{i})\\}_{i\in[h]}\in\mathfrak{W}$
be such that for all $i\in[h]$,

$\displaystyle\bigl{\|}(W^{{\mathrm{q}}}_{i})^{\top}-(\widehat{W}^{{\mathrm{q}}}_{i})^{\top}\bigr{\|}_{r,s}\leq\varepsilon_{i}^{\mathrm{q}},\quad\bigl{\|}(W^{{\mathrm{k}}}_{i})^{\top}-(\widehat{W}^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r,s}\leq\varepsilon_{i}^{\mathrm{k}},\quad\bigl{\|}(W^{{\mathrm{v}}}_{i})^{\top}-(\widehat{W}^{{\mathrm{v}}}_{i})^{\top}\bigr{\|}_{r,s}\leq\varepsilon_{i}^{\mathrm{v}}.$

Then we have

$\displaystyle\bigl{\|}{\mathtt{mha}}(X;W)^{\top}-{\mathtt{mha}}(X;\widehat{W})^{\top}\bigr{\|}_{r,\infty}\leq
R\cdot\sum_{i=1}^{h}\varepsilon_{i}^{\mathrm{v}}+R^{3}\cdot\sum_{i=1}^{h}(\omega^{{\mathrm{q}}}_{i}+\omega^{{\mathrm{k}}}_{i})\cdot(\varepsilon_{i}^{\mathrm{q}}+\varepsilon_{i}^{\mathrm{k}}).$

###### Proof.

See §F.2.1 for a detailed proof. ∎

Then we have the following lemma characterizing the covering number of the
function class $\\{{\mathtt{mha}}(X;W)^{\top}:W\in\mathfrak{W}\\}$ of MHA.

######  Lemma C.12 (MHA Covering).

Let the positive reals $(R,\varepsilon)$ and the positive integers
$(d,d_{k},d_{\mathrm{p}})$ be given. Let $\mathfrak{W}$ be given by (C.14).
For any input sequence $X\in\mathbb{R}^{L\times d}$ such that
$\|X^{\top}\|_{r,\infty}\leq R$, we have

$\displaystyle\log
N\Bigl{(}\bigl{\\{}{\mathtt{mha}}(X;W)^{\top}:W\in\mathfrak{W}\bigr{\\}},\varepsilon,\|\cdot\|_{r,\infty}\Bigr{)}\leq(2+h)\cdot
d^{2}\cdot\log\biggl{(}1+\frac{2RR^{\mathtt{mha}}(\mathfrak{W})}{\varepsilon}\biggr{)},$

where

$\displaystyle
R^{\mathtt{mha}}(\mathfrak{W})=\sum_{i=1}^{h}R^{{\mathrm{v}}}_{i}+R^{2}\cdot\sum_{i=1}^{h}(\omega^{{\mathrm{q}}}_{i}+\omega^{{\mathrm{k}}}_{i})(R^{{\mathrm{q}}}_{i}+R^{{\mathrm{k}}}_{i}).$
(C.15)

###### Proof.

For all $i\in[h]$, we set

$\displaystyle\frac{R^{{\mathrm{q}}}_{i}}{\varepsilon_{i}^{\mathrm{q}}}=\frac{R^{{\mathrm{k}}}_{i}}{\varepsilon_{i}^{\mathrm{k}}}=\frac{R^{{\mathrm{v}}}_{i}}{\varepsilon_{i}^{\mathrm{v}}}=\frac{RR^{\mathtt{mha}}(\mathfrak{W})}{\varepsilon},$

where $R^{\mathtt{mha}}(\mathfrak{W})$ is defined in (C.15). Then by Lemma
C.11, it can be verified that

$\displaystyle\bigl{\|}{\mathtt{mha}}(X;W)^{\top}-{\mathtt{mha}}(X;\widehat{W})^{\top}\bigr{\|}_{r,\infty}\leq
R\cdot\sum_{i=1}^{h}\varepsilon_{i}^{\mathrm{v}}+R^{3}\cdot\sum_{i=1}^{h}(\omega^{{\mathrm{q}}}_{i}+\omega^{{\mathrm{k}}}_{i})\cdot(\varepsilon_{i}^{\mathrm{q}}+\varepsilon_{i}^{\mathrm{k}})=\varepsilon.$

Thus, revolking Corollary C.9, we further obtain

$\displaystyle\log
N\Bigl{(}\bigl{\\{}{\mathtt{mha}}(X;W)^{\top}:W\in\mathfrak{W}\bigr{\\}},\varepsilon,\|\cdot\|_{r,\infty}\Bigr{)}$
$\displaystyle\quad\leq\sum_{i=1}^{h}\biggl{(}\log
N\biggl{(}\Bigl{\\{}(W^{{\mathrm{q}}}_{i})^{\top}\in\mathbb{R}^{d\times
d_{\mathrm{p}}}:\bigl{\|}(W^{{\mathrm{q}}}_{i})^{\top}\bigr{\|}_{r,s}\leq
R^{{\mathrm{q}}}_{i}\Bigr{\\}},\varepsilon_{i}^{\mathrm{q}},\|\cdot\|_{r,\infty}\biggr{)}$
$\displaystyle\quad\qquad+\log
N\biggl{(}\Bigl{\\{}(W^{{\mathrm{k}}}_{i})^{\top}\in\mathbb{R}^{d\times
d_{\mathrm{p}}}:\bigl{\|}(W^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r,s}\leq
R^{{\mathrm{k}}}_{i}\Bigr{\\}},\varepsilon_{i}^{\mathrm{k}},\|\cdot\|_{r,\infty}\biggr{)}$
$\displaystyle\quad\qquad+\log
N\biggl{(}\Bigl{\\{}(W^{{\mathrm{v}}}_{i})^{\top}\in\mathbb{R}^{d\times
d}:\bigl{\|}(W^{{\mathrm{v}}}_{i})^{\top}\bigr{\|}_{r,s}\leq
R^{{\mathrm{v}}}_{i}\Bigr{\\}},\varepsilon_{i}^{\mathrm{v}},\|\cdot\|_{r,\infty}\biggr{)}\biggr{)}$
$\displaystyle\quad\leq
dd_{\mathrm{p}}\cdot\sum_{i=1}^{h}\biggl{(}\log\biggl{(}1+\frac{2R^{{\mathrm{q}}}_{i}}{\varepsilon_{i}^{\mathrm{q}}}\biggr{)}+\log\biggl{(}1+\frac{2R^{{\mathrm{k}}}_{i}}{\varepsilon_{i}^{\mathrm{k}}}\biggr{)}\biggr{)}+d^{2}\cdot\sum_{i=1}^{h}\log\biggl{(}1+\frac{2R^{{\mathrm{v}}}_{i}}{\varepsilon_{i}^{\mathrm{v}}}\biggr{)}$
$\displaystyle\quad=(2+h)\cdot
d^{2}\cdot\log\biggl{(}1+\frac{2RR^{\mathtt{mha}}(\mathfrak{W})}{\varepsilon}\biggr{)},$

where the last line follows from $d=d_{\mathrm{p}}h$. Therefore, we conclude
the proof of Lemma C.12. ∎

####  C.5.2 Covering Number Propagation

Recall that $\mathfrak{W}$ is defined in (C.14). We start with the following
lemma characterizing the Lipschitz continuity of MHA with respect to the input
sequence $X$.

######  Lemma C.13 (Lipschitz Continuous MHA in Input Sequence).

Let $(r,s)$ be a conjugate pair with $1\leq r<\infty$. Suppose that
$X,\widehat{X}\in\mathbb{R}^{L\times d}$ satisfying
$\|X^{\top}\|_{r,\infty}\leq R$ and $\|\widehat{X}^{\top}\|_{r,\infty}\leq R$
are given. Also, let
$W=\\{(W_{i}^{\mathrm{q}},W_{i}^{\mathrm{k}},W_{i}^{\mathrm{v}})\\}_{i\in[n]}\in\mathfrak{W}$
be given. Then we have

$\displaystyle\bigl{\|}{\mathtt{mha}}(X;W)^{\top}-{\mathtt{mha}}(\widehat{X};W)^{\top}\bigr{\|}_{r,\infty}\leq\rho(\mathfrak{W})\cdot\|X^{\top}-\widehat{X}^{\top}\|_{r,\infty},$

where

$\displaystyle\rho(\mathfrak{W})=\sum_{i=1}^{h}\omega^{{\mathrm{v}}}_{i}+R^{2}\cdot\sum_{i=1}^{h}(\omega^{{\mathrm{q}}}_{i}+\omega^{{\mathrm{k}}}_{i})^{2}\omega^{{\mathrm{v}}}_{i}.$
(C.16)

###### Proof.

See §F.2.2 for a detailed proof. ∎

Then we present the following lemma on the covering number of
$\mathcal{F}^{L}_{\mathtt{trans}}$.

######  Lemma C.14 (Covering Propagation).

We have

$\displaystyle\log
N\bigl{(}\mathcal{F}^{L}_{\mathtt{trans}},\varepsilon,\|\cdot\|_{r,\infty}\bigr{)}$
$\displaystyle\quad\leq\sum_{t=0}^{T-1}\sup_{\\{(W^{(\tau)},{A}^{(\tau)})\in\mathfrak{W}^{(\tau)}\times{\mathfrak{A}}^{(\tau)}\\}_{0\leq\tau\leq
t}}\Bigl{\\{}\log
N\bigl{(}\mathcal{F}^{L,(t)}_{\mathfrak{W}},\varepsilon^{(t)}_{\mathfrak{W}},\|\cdot\|_{r,\infty}\bigr{)}+\log
N\bigl{(}\mathcal{F}^{L,(t)}_{\mathfrak{A}},\varepsilon^{(t)}_{\mathfrak{A}},\|\cdot\|_{r,\infty}\bigr{)}\Bigr{\\}}.$

Here, with the convention $\prod_{\tau=T}^{T-1}\cdot\equiv 1$, the covering
resolution $\varepsilon$ is defined as

$\displaystyle\varepsilon=\sum_{t=0}^{T-1}\biggl{[}(\widetilde{\rho}^{(t)}\cdot\varepsilon^{(t)}_{\mathfrak{A}}+\varepsilon^{(t)}_{\mathfrak{W}})\cdot\prod_{\tau=t+1}^{T-1}\widetilde{\rho}^{(\tau)}\widetilde{\alpha}^{(\tau)}\biggr{]},$
(C.17)

where $\widetilde{\alpha}^{(t)}$ is defined in (C.6) and
$\widetilde{\rho}^{(t)}$ is defined in (C.11).

###### Proof.

For any $F\in\mathcal{F}^{L,(T)}_{\mathfrak{A}}$, let
$f=\mathtt{agg}^{\mathtt{trans}}_{\theta_{0}}\circ F$. There exists some
$\widehat{F}$, which also induces a corresponding
$\widehat{f}=\mathtt{agg}^{\mathtt{trans}}_{\theta_{0}}\circ\widehat{F}$, in
the $\|\cdot\|_{r,\infty}$-covering of $\mathcal{F}^{L,(T)}_{\mathfrak{A}}$
with the covering resolution $\varepsilon$ such that

$\displaystyle\Bigl{\|}\bigl{(}f(X_{1}),\dots,f(X_{n})\bigr{)}-\bigl{(}\widehat{f}(X_{1}),\dots,\widehat{f}(X_{n})\bigr{)}\Bigr{\|}_{\infty}$
$\displaystyle\quad=\max_{i\in[n]}\Bigl{|}f(X_{i})-\widehat{f}(X_{i})\Bigr{|}$
$\displaystyle\quad=\max_{i\in[n]}\Bigl{|}\mathtt{agg}^{\mathtt{trans}}_{\theta_{0}}\bigl{(}{F}(X_{i})\bigr{)}-\mathtt{agg}^{\mathtt{trans}}_{\theta_{0}}\bigl{(}\widehat{F}(X_{i})\bigr{)}\Bigr{|}\leq\max_{i\in[n]}\bigl{\|}{F}(X_{i})^{\top}-\widehat{F}(X_{i})^{\top}\bigr{\|}_{r,\infty}\leq\varepsilon,$

where the first inequality follows from $1$-Lipschitz continuity of
$\mathtt{agg}^{\mathtt{trans}}_{\theta_{0}}$. Thus, we have
$N(\mathcal{F}^{L}_{\mathtt{trans}},\varepsilon,\|\cdot\|_{\infty})\leq
N(\mathcal{F}^{L,(T)}_{\mathfrak{A}},\varepsilon,\|\cdot\|_{r,\infty})$, which
means that it suffices to bound
$N(\mathcal{F}^{L,(T)}_{\mathfrak{A}},\varepsilon,\|\cdot\|_{r,\infty})$.

Throughout this proof, we fix the parameters $\\{(W^{(t)},A^{(t)})\\}_{0\leq
t\leq T-1}$. We inductively construct the perturbed intermediate input
sequences $\widehat{X}^{(t)}$ and output sequences
$\widehat{X}_{\star}^{(t)}$. Starting from
$\widehat{X}^{(0)}_{\star}=X^{(0)}_{\star}=X$, we make the constructions as
follows,

$\displaystyle\widehat{X}^{(t)}$
$\displaystyle\in\Bigl{\\{}X:\bigl{\|}{\mathtt{ffn}}(\widehat{X}^{(t)}_{\star};A^{(t)})^{\top}-X^{\top}\bigr{\|}_{r,\infty}\leq\varepsilon^{(t)}_{\mathfrak{A}},X\in\mathcal{F}^{L,(t)}_{\mathfrak{A}}\Bigr{\\}},$
$\displaystyle\widehat{X}^{(t+1)}_{\star}$
$\displaystyle\in\Bigl{\\{}X_{\star}:\bigl{\|}{\mathtt{mha}}(\widehat{X}^{(t)};W^{(t)})^{\top}+(\widehat{X}^{(t)})^{\top}-X_{\star}^{\top}\bigr{\|}_{r,\infty}\leq\varepsilon^{(t)}_{\mathfrak{W}},X_{\star}\in\mathcal{F}_{\mathfrak{W}}^{L,(t+1)}\Bigr{\\}}.$
(C.18)

For the above inductive constructions, we have

$\displaystyle\bigl{\|}(X^{(t)})^{\top}-(\widehat{X}^{(t)})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle=\bigl{\|}{\mathtt{ffn}}({X}^{(t)}_{\star};A^{(t)})^{\top}-(\widehat{X}^{(t)})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\leq\bigl{\|}{\mathtt{ffn}}({X}^{(t)}_{\star};A^{(t)})^{\top}-{\mathtt{ffn}}(\widehat{X}^{(t)}_{\star};A^{(t)})^{\top}\bigr{\|}_{r,\infty}+\bigl{\|}{\mathtt{ffn}}(\widehat{X}^{(t)}_{\star};A^{(t)})^{\top}-(\widehat{X}^{(t)})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\leq\bigl{\|}{\mathtt{ffn}}({X}^{(t)}_{\star};A^{(t)})^{\top}-{\mathtt{ffn}}(\widehat{X}^{(t)}_{\star};A^{(t)})^{\top}\bigr{\|}_{r,\infty}+\varepsilon_{\mathfrak{A}}^{(t)},$
(C.19)

where the second inequality follows from the definition of
$\widehat{X}^{(t+1)}$ in (C.5.2). For the first term on the right-hand side of
(C.5.2), we have

$\displaystyle\bigl{\|}{\mathtt{ffn}}({X}^{(t)}_{\star};A^{(t)})^{\top}-{\mathtt{ffn}}(\widehat{X}^{(t)}_{\star};{A}^{(t)})^{\top}\bigr{\|}_{r,\infty}$
(C.20)
$\displaystyle\quad=\Bigl{\|}(A^{\sigma,(t)})^{\top}{\mathtt{ReLU}}({X}^{(t)}_{\star}{A}^{{\mathrm{x}},(t)})^{\top}+({X}^{(t)}_{\star})^{\top}-(A^{\sigma,(t)})^{\top}{\mathtt{ReLU}}(\widehat{X}^{(t)}_{\star}{A}^{{\mathrm{x}},(t)})^{\top}-(\widehat{X}^{(t)}_{\star})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\quad\leq\bigl{\|}(A^{\sigma,(t)})^{\top}\bigr{\|}_{r}\cdot\bigl{\|}{\mathtt{ReLU}}({X}^{(t)}_{\star}{A}^{{\mathrm{x}},(t)})^{\top}-{\mathtt{ReLU}}(\widehat{X}^{(t)}_{\star}{A}^{{\mathrm{x}},(t)})^{\top}\bigr{\|}_{r,\infty}+\bigl{\|}({X}^{(t)}_{\star})^{\top}-(\widehat{X}^{(t)}_{\star})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\quad\leq\Bigl{(}1+\bigl{\|}({A}^{{\mathrm{x}},(t)})^{\top}\bigr{\|}_{r}\bigl{\|}({A}^{\sigma,(t)})^{\top}\bigr{\|}_{r}\Bigr{)}\cdot\bigl{\|}({X}^{(t)}_{\star})^{\top}-(\widehat{X}^{(t)}_{\star})^{\top}\bigr{\|}_{r,\infty}\leq\widetilde{\alpha}^{(t)}\cdot\bigl{\|}({X}^{(t)}_{\star})^{\top}-(\widehat{X}^{(t)}_{\star})^{\top}\bigr{\|}_{r,\infty},$

where the first and the second inequalities follow from Lemma G.3, and the
last inequality follows from
$(A^{{\mathrm{x}},(t)},A^{\sigma,(t)})\in\mathfrak{A}$ and the definition of
$\widetilde{\alpha}^{(t)}$ in (C.6). By Lemma C.13, we have

$\displaystyle\bigl{\|}({X}^{(t+1)}_{\star})^{\top}-(\widehat{X}^{(t+1)}_{\star})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\leq\bigl{\|}{\mathtt{mha}}({X}^{(t)};A^{(t)})^{\top}+({X}^{(t)})^{\top}-{\mathtt{mha}}(\widehat{X}^{(t)};A^{(t)})^{\top}-(\widehat{X}^{(t)})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\qquad+\bigl{\|}{\mathtt{mha}}(\widehat{X}^{(t)};A^{(t)})^{\top}+(\widehat{X}^{(t)})^{\top}-(\widehat{X}^{(t)}_{\star})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\leq\bigl{(}\rho(\mathfrak{W}^{(t)})+1\bigr{)}\cdot\|(X^{(t)})^{\top}-(\widehat{X}^{(t)})^{\top}\bigr{\|}_{r,\infty}+\varepsilon^{(t)}_{\mathfrak{W}}$
$\displaystyle\leq\widetilde{\rho}^{(t)}\cdot\bigl{\|}(X^{(t)})^{\top}-(\widehat{X}^{(t)})^{\top}\bigr{\|}_{r,\infty}+\varepsilon^{(t)}_{\mathfrak{W}},$
(C.21)

where the second inequality follows from Lemma C.13, and the last inequality
follows from Lemma G.4. By taking (C.5.2) into (C.20) and (C.5.2), we further
obtain

$\displaystyle\bigl{\|}(X^{(t+1)}_{\star})^{\top}-(\widehat{X}^{(t+1)}_{\star})^{\top}\bigr{\|}_{r,\infty}\leq\widetilde{\rho}^{(t)}\widetilde{\alpha}^{(t)}\cdot\bigl{\|}(X_{\star}^{(t)})^{\top}-(\widehat{X}_{\star}^{(t)})^{\top}\|_{r,\infty}+\widetilde{\rho}^{(t)}\cdot\varepsilon^{(t)}_{\mathfrak{A}}+\varepsilon^{(t)}_{\mathfrak{W}}.$

Recursively applying the above inequality, we conclude the proof of Lemma
C.14. ∎

####  C.5.3 Proof of Transformer Covering Lemma C.8

###### Proof.

By Lemma C.14, we have

$\displaystyle\log
N\bigl{(}\mathcal{F}^{L}_{\mathtt{trans}},\varepsilon,\|\cdot\|_{r,\infty}\bigr{)}$
$\displaystyle\quad\leq\sum_{t=0}^{T-1}\sup_{\\{(W^{(\tau)},{A}^{(\tau)})\in\mathfrak{W}^{(\tau)}\times{\mathfrak{A}}^{(\tau)}\\}_{0\leq\tau\leq
t}}\Bigl{\\{}\log
N\bigl{(}\mathcal{F}^{L,(t)}_{\mathfrak{W}},\varepsilon^{(t)}_{\mathfrak{W}},\|\cdot\|_{r,\infty}\bigr{)}+\log
N\bigl{(}\mathcal{F}^{L,(t)}_{\mathfrak{A}},\varepsilon^{(t)}_{\mathfrak{A}},\|\cdot\|_{r,\infty}\bigr{)}\Bigr{\\}}.$

First, we set
$\varepsilon^{(t)}_{\mathfrak{A}}=\varepsilon^{(t)}\cdot(\alpha^{{\mathrm{x}},(t)}R^{\sigma,(t)}+\alpha^{\sigma,(t)}R^{{\mathrm{x}},(t)})/\widetilde{\omega}^{{\mathrm{v}},(t)}\widetilde{\alpha}^{(t)}$,
and $\varepsilon^{(t)}_{\mathfrak{W}}=\varepsilon^{(t)}\cdot
R^{{\mathtt{mha}},(t)}/\widetilde{\omega}^{{\mathrm{v}},(t)}$. Note that by
Lemma G.6 we have

$\displaystyle\bigl{\|}(X^{(t)})^{\top}\bigr{\|}_{r,\infty}\leq
R^{(0)}\cdot\widetilde{\alpha}^{(t)}\prod_{\tau=0}^{t-1}\widetilde{\omega}^{{\mathrm{v}},(\tau)}\widetilde{\alpha}^{(\tau)},\qquad\bigl{\|}(X^{(t)}_{\star})^{\top}\bigr{\|}_{r,\infty}\leq
R^{(0)}\cdot\prod_{\tau=0}^{t-1}\widetilde{\omega}^{{\mathrm{v}},(\tau)}\widetilde{\alpha}^{(\tau)}.$

By Lemma C.10, we have for all $0\leq t\leq T-1$,

$\displaystyle\log
N\bigl{(}\mathcal{F}^{(t)}_{\mathfrak{A}},\varepsilon^{(t)}_{\mathfrak{A}},\|\cdot\|_{r,\infty}\bigr{)}$
$\displaystyle\leq
2dd_{\sigma}\cdot\log\biggl{(}1+\frac{2(\alpha^{{\mathrm{x}},(t)}R^{\sigma,(t)}+\alpha^{\sigma,(t)}R^{{\mathrm{x}},(t)})\cdot\bigl{\|}(X_{\star}^{(t)})^{\top}\bigr{\|}_{r,\infty}}{\varepsilon^{(t)}_{\mathfrak{A}}}\biggr{)}$
(C.22) $\displaystyle\leq
2D^{2}\cdot\log\biggl{(}1+\frac{2(\alpha^{{\mathrm{x}},(t)}R^{\sigma,(t)}+\alpha^{\sigma,(t)}R^{{\mathrm{x}},(t)})R^{(0)}\cdot\prod_{\tau=0}^{t}\widetilde{\omega}^{{\mathrm{v}},(\tau)}\widetilde{\alpha}^{(\tau)}}{\varepsilon^{(t)}}\biggr{)}.$

On the other hand, by Lemma C.12, we have for all $0\leq t\leq T-1$,

$\displaystyle\log
N\bigl{(}\mathcal{F}^{L,(t)}_{\mathfrak{W}},\varepsilon^{(t)}_{\mathfrak{W}},\|\cdot\|_{r,\infty}\bigr{)}$
(C.23) $\displaystyle\quad\leq\log
N\Bigl{(}\bigl{\\{}{\mathtt{mha}}(X^{(t)};W^{(t)})^{\top}+(X^{(t)})^{\top}:W^{(t)}\in\mathfrak{W}^{(t)}\bigr{\\}},\varepsilon^{(t)}_{\mathfrak{W}},\|\cdot\|_{r,\infty}\Bigr{)}$
$\displaystyle\quad\leq\log
N\Bigl{(}\bigl{\\{}{\mathtt{mha}}(X^{(t)};W^{(t)})^{\top}:W^{(t)}\in\mathfrak{W}^{(t)}\bigr{\\}},\varepsilon^{(t)}_{\mathfrak{W}},\|\cdot\|_{r,\infty}\Bigr{)}$
$\displaystyle\quad\leq(2+h)\cdot
D^{2}\cdot\log\biggl{(}1+\frac{2R^{(t)}R^{{\mathtt{mha}},(t)}}{\varepsilon^{(t)}_{\mathfrak{W}}}\biggr{)}=(2+h)\cdot
D^{2}\cdot\log\biggl{(}1+\frac{2R^{(0)}\cdot\prod_{\tau=0}^{t}\widetilde{\omega}^{{\mathrm{v}},(\tau)}\widetilde{\alpha}^{(\tau)}}{\varepsilon^{(t)}}\biggr{)},$

where the last line follows from Lemma G.4 and the choice of
$\varepsilon^{(t)}_{\mathfrak{W}}$. It remains to choose the per-layer
resolutions $\\{\varepsilon^{(t)}\\}_{0\leq t\leq T-1}$ that satisfy (C.17),
which is

$\displaystyle\varepsilon=\sum_{t=0}^{T-1}\biggl{[}\varepsilon^{(t)}\cdot\biggl{(}\frac{R^{{\mathtt{mha}},(t)}}{\widetilde{\omega}^{{\mathrm{v}},(t)}}+\frac{\alpha^{{\mathrm{x}},(t)}R^{\sigma,(t)}+\alpha^{\sigma,(t)}R^{{\mathrm{x}},(t)}}{\widetilde{\alpha}^{(t)}}\biggr{)}\cdot\prod_{\tau=t+1}^{T-1}\widetilde{\rho}^{(\tau)}\widetilde{\alpha}^{(\tau)}\biggr{]}.$
(C.24)

For any $0\leq t\leq T-1$, we set

$\displaystyle\varepsilon^{(t)}$
$\displaystyle=\frac{\varepsilon}{{R}^{{\mathtt{trans}},(T)}\cdot\prod_{\tau=t+1}^{T-1}\widetilde{\omega}^{{\mathrm{v}},(\tau)}\widetilde{\alpha}^{(\tau)}},$

which can be verified to satisfy (C.24). Finally, combining Lemma C.14,
(C.22), (C.23) and the choices of $\\{\varepsilon^{(t)}\\}_{0\leq t\leq T-1}$,
we obtain

$\displaystyle\log
N\bigl{(}\mathcal{F}^{L}_{\mathtt{trans}},\varepsilon,\|\cdot\|_{r,\infty}\bigr{)}$
$\displaystyle\quad\leq\sum_{t=0}^{T-1}\sup_{\\{(W^{(\tau)},{A}^{(\tau)})\in\mathfrak{W}^{(\tau)}\times{\mathfrak{A}}^{(\tau)}\\}_{0\leq\tau\leq
t}}\Bigl{\\{}\log
N\bigl{(}\mathcal{F}^{L,(t)}_{\mathfrak{W}},\varepsilon^{(t)}_{\mathfrak{W}},\|\cdot\|_{r,\infty}\bigr{)}+\log
N\bigl{(}\mathcal{F}^{L,(t)}_{\mathfrak{A}},\varepsilon^{(t)}_{\mathfrak{A}},\|\cdot\|_{r,\infty}\bigr{)}\Bigr{\\}}$
$\displaystyle\quad\leq(4+h)\cdot
D^{2}T\cdot\log\biggl{(}1+\frac{2{R}^{{\mathtt{trans}},(T)}R^{(T)}}{\varepsilon}\biggr{)}.$

Therefore, we conclude the proof of Lemma C.8. ∎

##  Appendix D Optimization Error Analysis

###### Proof of Proposition 5.6.

Let
$\widehat{\mathcal{L}}(f_{\theta})=\widehat{\mathbb{E}}[\mathcal{L}((X,y),{f}_{\theta})]$.
By (5.12), it holds for the stationary point $\widehat{\theta}$ that,

$\displaystyle
0\leq\bigl{\langle}\nabla_{\theta}\widehat{\mathcal{L}}(f_{\widehat{\theta}}),\theta-\widehat{\theta}\bigr{\rangle}=\widehat{\mathbb{E}}\Bigl{[}\nabla_{f}\mathcal{L}\bigl{(}(X,y),f_{\widehat{\theta}}\bigr{)}\nabla_{\theta}f_{\widehat{\theta}}(X)^{\top}(\theta-\widehat{\theta})\Bigr{]}.$

Since the objective function $\mathcal{L}((X,y),f)=\|y-f(X)\|_{2}^{2}$ is
convex with respect to $f(X)$, we have

$\displaystyle
0\leq\widehat{\mathbb{E}}\Bigl{[}\nabla_{f}\mathcal{L}\bigl{(}(X,y),f_{\theta^{*}}\bigr{)}^{\top}(f-f_{\theta^{*}})(X)\Bigr{]},$

where
$\theta^{*}=\mathop{\mathrm{argmin}}_{\theta\in\Theta}\widehat{\mathcal{L}}(f_{\theta})$.
By definition of the objective function $\mathcal{L}((X,y),f)$, we have

$\displaystyle\Bigl{\|}\nabla_{f}\mathcal{L}\bigl{(}(X,y),f_{\widehat{\theta}}\bigr{)}\Bigr{\|}_{2}=2\bigl{\|}y-f_{\widehat{\theta}}(X)\bigr{\|}_{2}\leq
2\|y\|_{2}+2\|f_{\widehat{\theta}}(X)\bigr{\|}_{2}\leq 4,$ (D.1)

where the last inequality follows from Assumption 5.2 and that the aggregation
layer
$\mathtt{agg}_{\theta}:\mathbb{R}^{d_{\mathrm{p}}}\to\mathbb{R}^{d_{\mathrm{y}}}$
outputs within $\mathfrak{Y}$. For any $\theta\in\Theta$, it holds that,

$\displaystyle\widehat{\mathcal{L}}(f_{\widehat{\theta}})-\widehat{\mathcal{L}}(f_{\theta^{*}})$
$\displaystyle\leq\widehat{\mathbb{E}}\Bigl{[}\nabla_{f}\mathcal{L}\bigl{(}(X,y),f_{\widehat{\theta}}\bigr{)}^{\top}(f_{\widehat{\theta}}-f_{\theta^{*}})(X)\Bigr{]}$
$\displaystyle\leq\widehat{\mathbb{E}}\Bigl{[}\nabla_{f}\mathcal{L}\bigl{(}(X,y),f_{\widehat{\theta}}\bigr{)}^{\top}(f_{\widehat{\theta}}-f_{\theta^{*}})(X)\Bigr{]}+\widehat{\mathbb{E}}\Bigl{[}\nabla_{f}\mathcal{L}\bigl{(}(X,y),f_{\widehat{\theta}}\bigr{)}^{\top}\nabla_{\theta}f_{\widehat{\theta}}(x)^{\top}(\theta-\widehat{\theta})\Bigr{]}$
$\displaystyle=\widehat{\mathbb{E}}\Bigl{[}\nabla_{f}\mathcal{L}\bigl{(}(X,y),f_{\widehat{\theta}}\bigr{)}^{\top}\bigl{(}f_{\widehat{\theta}}(X)+\nabla_{\theta}f_{\widehat{\theta}}(X)^{\top}(\theta-\widehat{\theta})-f_{\theta^{*}}(X)\bigr{)}\Bigr{]}$
$\displaystyle\leq\widehat{\mathbb{E}}\biggl{[}\Bigl{\|}\nabla_{f}\mathcal{L}\bigl{(}(X,y),f_{\widehat{\theta}}\bigr{)}\Bigr{\|}_{2}\cdot\bigl{\|}f_{\widehat{\theta}}(X)+\nabla_{\theta}f_{\widehat{\theta}}(X)^{\top}(\theta-\widehat{\theta})-f_{\theta^{*}}(X)\bigr{\|}_{2}\biggr{]}$
$\displaystyle\leq
4\cdot\big{\|}f_{\widehat{\theta}}(X)+\nabla_{\theta}f_{\widehat{\theta}}(X)^{\top}(\theta-\widehat{\theta})-f_{\theta^{*}}(X)\big{\|},$
(D.2)

where the fourth line follows from the Cauchy-Schwartz inequality and the last
line follows from (D.1). Finally, since (D) holds for any $\theta\in\Theta$,
we have

$\displaystyle\widehat{\mathcal{L}}(f_{\widehat{\theta}})-\widehat{\mathcal{L}}(f_{\theta^{*}})\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_{\theta^{*}}(X)\bigr{\|}_{2}\Bigr{]}.$

Therefore, we conclude the proof of Proposition 5.6. ∎

##  Appendix E Approximation Error: Recovery from RKHS

###  E.1 Latent-to-Value RKHS

Recall that the latent-to-value mapping $\psi(z;{\mathtt{msk}})$ is defined in
(5.2), which induces the reproducing kernel function
$\mathfrak{K}_{{\mathtt{LTV}}}(z,z^{\prime};{\mathtt{msk}})=\psi(z;{\mathtt{msk}})^{\top}\psi(z^{\prime};{\mathtt{msk}})$.
Such a reproducing kernel function
$\mathfrak{K}_{{\mathtt{LTV}}}(z,z^{\prime};{\mathtt{msk}})$ induces the RKHS
as follows,

$\displaystyle\mathcal{H}_{{\mathtt{LTV}}}=\biggl{\\{}g_{\alpha}(z;{\mathtt{msk}})=\int\alpha(z^{\prime})\mathfrak{K}_{\mathtt{LTV}}(z^{\prime},z;{\mathtt{msk}}){\mathrm{d}}z^{\prime}:\|g_{\alpha}(\cdot;{\mathtt{msk}})\|_{\mathcal{H}_{\mathtt{LTV}}}<\infty\biggr{\\}},$
(E.1)

which is equipped with the inner product
$\langle\cdot,\cdot\rangle_{\mathcal{H}_{{\mathtt{LTV}}}}$. By the definition
of the kernel function
$\mathfrak{K}_{\mathtt{LTV}}(\cdot,\cdot;{\mathtt{msk}})$, we have for
$g_{\alpha}(\cdot;{\mathtt{msk}})\in\mathcal{H}_{\mathtt{LTV}}$ that,

$\displaystyle g_{\alpha}(z;{\mathtt{msk}})$
$\displaystyle=\int\alpha(z^{\prime})\mathfrak{K}_{\mathtt{LTV}}(z^{\prime},z;{\mathtt{msk}}){\mathrm{d}}z^{\prime}$
$\displaystyle=\biggl{(}\underbrace{\int\alpha(z^{\prime})\psi(z^{\prime};{\mathtt{msk}}){\mathrm{d}}z^{\prime}}_{\displaystyle
w_{\alpha}\in\mathbb{R}^{d}}\biggr{)}^{\top}\psi(z;{\mathtt{msk}})=w_{\alpha}^{\top}\psi(z;{\mathtt{msk}}).$
(E.2)

From (E.1), we see that $w_{\alpha}$ corresponds to the vector
$w_{i}\in\mathbb{R}^{d_{\mathrm{p}}}$ in the function class
$\mathcal{G}^{\dagger}_{i}$, which is defined in (5.10). On the other hand, we
have

$\displaystyle\bigl{\|}g_{\alpha}(\cdot;{\mathtt{msk}})\bigr{\|}^{2}_{\mathcal{H}_{\mathtt{LTV}}}$
$\displaystyle=\bigl{\langle}g_{\alpha}(\cdot;{\mathtt{msk}}),g_{\alpha}(\cdot;{\mathtt{msk}})\bigr{\rangle}_{\mathcal{H}_{\mathtt{LTV}}}$
$\displaystyle=\int\alpha(z^{\prime})\mathfrak{K}_{\mathtt{LTV}}(z^{\prime},z;{\mathtt{msk}})\alpha(z){\mathrm{d}}z{\mathrm{d}}z^{\prime}$
$\displaystyle=\biggl{(}\int\alpha(z^{\prime})\psi(z^{\prime};{\mathtt{msk}}){\mathrm{d}}z^{\prime}\biggr{)}^{\top}\biggl{(}\int\alpha(z)\psi(z;{\mathtt{msk}}){\mathrm{d}}z\biggr{)}=\|w_{\alpha}\|_{2}^{2},$
(E.3)

where the third equality follows from the definition of the kernel function
$\mathfrak{K}_{\mathtt{LTV}}(\cdot,\cdot;{\mathtt{msk}})$ and the last
equality follows from the definition of $w_{\alpha}$ in (E.1). Combining
(E.1), (E.1) and the definition of $\mathcal{H}_{\mathtt{LTV}}$ in (E.1), we
have

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

Thus, the function class $\mathcal{G}_{i}^{\dagger}$ corresponding to the
$i$-th entry of functions in the function class $\mathcal{G}$ is the RKHS
$\mathcal{H}_{\mathtt{LTV}}$. Here the function class $\mathcal{G}^{\dagger}$
is defined in (5.9), which contains the latent-to-target functions obtained by
the reweighted CME attention $f^{\dagger}_{W}(X;{\mathtt{msk}})$ in (5.6).

###  E.2 Supervised Learning

###### Proof of Theorem 5.5.

Let $f_{\theta}\in\mathcal{F}_{\mathtt{attn}}$ and
$\epsilon_{\mathtt{attn}}\in[0,+\infty)$ be such that (5.11) holds. By the
definition of the approximation error $\mathcal{E}_{\rm approx}$ in (5.3), we
have

$\displaystyle\mathcal{E}_{\rm approx}$
$\displaystyle=\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\leq\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f_{\theta}\bigr{)}\Bigr{]}-\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(X,y),f^{*}\bigr{)}\Bigr{]}$
$\displaystyle=\mathbb{E}\Bigl{[}\bigl{\|}f_{\theta}(X;{\mathtt{msk}})-f^{*}(X)\bigr{\|}_{2}^{2}\Bigr{]}$
$\displaystyle\leq
2\mathbb{E}\Bigl{[}\bigl{\|}f_{\theta}(X;{\mathtt{msk}})-f_{W}^{\dagger}(X;{\mathtt{msk}})\bigr{\|}_{2}^{2}\Bigr{]}+2\mathbb{E}\Bigl{[}\bigl{\|}f_{W}^{\dagger}(X;{\mathtt{msk}})-f^{*}(X)\bigr{\|}_{2}^{2}\Bigr{]}$
$\displaystyle\leq
2\epsilon^{2}_{\mathtt{attn}}+2\mathbb{E}\Bigl{[}\bigl{\|}f_{W}^{\dagger}(X;{\mathtt{msk}})-f^{*}(X)\bigr{\|}_{2}^{2}\Bigr{]},$
(E.4)

where the second line follows from $f_{\theta}\in\mathcal{F}_{\mathtt{attn}}$,
the third line follows from $\mathcal{L}((X,y),f)=\|y-f(X)\|_{2}^{2}$ and the
definition of the regression function $f^{*}(X)=\mathbb{E}[y{\,|\,}X]$, and
the last line follows from (5.11).

In what follows, we bound the gap between the regression function $f^{*}(X)$
and the reweighted CME attention in (5.6), which is used as the surrogate for
approximating $f^{*}(X)$. By (5.2), we have

$\displaystyle f^{*}(X)$
$\displaystyle=\mathbb{E}_{z{\,|\,}X}\bigl{[}g^{*}(z;{\mathtt{msk}})\bigr{]}$
$\displaystyle=\mathbb{E}_{z{\,|\,}X}\bigl{[}g_{W}^{\dagger}(z;{\mathtt{msk}})\bigr{]}+\mathbb{E}_{z{\,|\,}X}\bigl{[}g^{*}(z;{\mathtt{msk}})-g_{W}^{\dagger}(z;{\mathtt{msk}})\bigr{]}$
$\displaystyle=f_{W}^{\dagger}(X;{\mathtt{msk}})+\mathbb{E}_{z{\,|\,}X}\bigl{[}g^{*}(z;{\mathtt{msk}})-g_{W}^{\dagger}(z;{\mathtt{msk}})\bigr{]},$

where the last line follows from (5.8). Thus, it holds for the reweighted CME
attention $f_{W}^{\dagger}(X;{\mathtt{msk}})$ that

$\displaystyle\mathbb{E}\Bigl{[}\bigl{\|}f^{*}(X)-f_{W}^{\dagger}(X;{\mathtt{msk}})\bigr{\|}_{2}^{2}\Bigr{]}=\mathbb{E}\biggl{[}\Bigl{\|}\mathbb{E}_{z{\,|\,}X}\bigl{[}g^{*}(z;{\mathtt{msk}})-g^{\dagger}_{W}(z;{\mathtt{msk}})\bigr{]}\Bigr{\|}_{2}^{2}\biggr{]}.$
(E.5)

By Assumption 5.4, we have

$\displaystyle\Bigl{\|}\mathbb{E}_{z{\,|\,}X}\bigl{[}g^{*}(z;{\mathtt{msk}})-g^{\dagger}_{W}(z;{\mathtt{msk}})\bigr{]}\Bigr{\|}_{2}^{2}$
$\displaystyle=\sum_{i=1}^{d_{\mathrm{y}}}\mathbb{E}_{z{\,|\,}X}\bigl{[}g_{i}^{*}(z;{\mathtt{msk}})-g^{\dagger}_{W,i}(z;{\mathtt{msk}})\bigr{]}^{2}$
$\displaystyle\leq\sum_{i=1}^{d_{\mathrm{y}}}\bigl{\|}g_{i}^{*}(\cdot;{\mathtt{msk}})-g^{\dagger}_{W,i}(\cdot;{\mathtt{msk}})\bigr{\|}_{\infty}^{2}\leq\epsilon^{2}_{g}({\mathtt{msk}}).$
(E.6)

Taking (E.2) into (E.5), we obtain

$\displaystyle\mathbb{E}\Bigl{[}\bigl{\|}f^{*}(X)-f_{W}^{\dagger}(X;{\mathtt{msk}})\bigr{\|}_{2}^{2}\Bigr{]}\leq\mathbb{E}\bigl{[}\epsilon^{2}_{g}({\mathtt{msk}})\bigr{]}=\epsilon^{2}_{g}({\mathtt{msk}}).$
(E.7)

Finally, taking (E.7) into (E.2), we obtain

$\displaystyle\mathcal{E}_{\rm approx}\leq
2\epsilon_{\mathtt{attn}}^{2}+2\epsilon_{g}^{2}({\mathtt{msk}}),$

which concludes the proof of Theorem 5.5. ∎

###  E.3 Self-Supervised Learning

###### Proof of Theorem 6.3.

Recall that
$B=W_{\mathtt{DS}}^{\top}(W_{\mathtt{SSL}}W^{\top}_{\mathtt{SSL}})^{-1}W_{\mathtt{SSL}}$
is defined in (6.10). First, we define a key surrogate function as follows,

$\displaystyle\widetilde{f}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})=B\widehat{f}_{{\mathtt{PT}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}}).$

For the approximation error $\mathcal{E}_{\rm approx}$ defined in (6.13), we
have

$\displaystyle\mathcal{E}_{\text{approx}}$
$\displaystyle\leq\mathbb{E}\Bigl{[}\mathcal{L}\bigl{(}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu,y_{\mathtt{DS}}),f_{\theta_{\mathtt{DS}}}\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{]}$ (E.8)
$\displaystyle=\mathbb{E}\Bigl{[}\bigl{\|}f_{\theta_{\mathtt{DS}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})-f^{*}_{\mathtt{DS}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)\bigr{\|}^{2}_{2}\Bigr{]}$ $\displaystyle\leq
3\mathbb{E}\Bigl{[}\bigl{\|}f_{\theta_{\mathtt{DS}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})-\widetilde{f}_{W_{\mathtt{DS}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})\bigr{\|}^{2}_{2}\Bigr{]}+3\mathbb{E}\Bigl{[}\bigl{\|}\widetilde{f}_{W_{\mathtt{DS}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})-f^{\dagger}_{W_{\mathtt{DS}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})\bigr{\|}^{2}_{2}\Bigr{]}$
$\displaystyle\quad\qquad+3\mathbb{E}\Bigl{[}\bigl{\|}f^{\dagger}_{W_{\mathtt{DS}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})-f^{*}_{\mathtt{DS}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)\bigr{\|}^{2}_{2}\Bigr{]}$ $\displaystyle\leq
3\epsilon^{2}_{\mathtt{agg}}(B)+3\underbrace{\mathbb{E}\Bigl{[}\bigl{\|}\widetilde{f}_{W_{\mathtt{DS}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})-f^{\dagger}_{W_{\mathtt{DS}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})\bigr{\|}^{2}_{2}\Bigr{]}}_{\displaystyle\text{(i)}}+3\underbrace{\mathbb{E}\Bigl{[}\bigl{\|}f^{\dagger}_{W_{\mathtt{DS}}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})-f^{*}_{\mathtt{DS}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)\bigr{\|}^{2}_{2}\Bigr{]}}_{\displaystyle\text{(ii)}},$

where the second line follows from the definition of the regression function
$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]$ and the last line
follows from (6). In what follows, we bound the terms (i) and (ii). Bounding
(i). For any truncated input sequence $\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern 1.5mu$, it holds that,

$\displaystyle\bigl{\|}\widetilde{f}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})-f_{W_{\mathtt{DS}}}^{\dagger}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})\bigr{\|}_{2}^{2}$
$\displaystyle\quad=\Bigl{\|}B\widehat{f}_{\mathtt{PT}}(z;{\mathtt{msk}}_{\mathtt{PT}})-W^{\top}_{\mathtt{DS}}\mathbb{E}_{z{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu}\bigl{[}\psi_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})\bigr{]}\Bigr{\|}_{2}^{2}$
$\displaystyle\quad=\biggl{\|}B\Bigl{(}\widehat{f}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})-f^{*}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)\Bigr{)}+\Bigl{(}Bf^{*}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)-\mathbb{E}_{z{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu}\bigl{[}W^{\top}_{\mathtt{DS}}\psi_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})\bigr{]}\Bigr{)}\biggr{\|}_{2}^{2}$
$\displaystyle\quad\leq
2\underbrace{\Bigl{\|}B\bigl{(}\widehat{f}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})-f^{*}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)\bigr{)}\Bigr{\|}_{2}^{2}}_{\displaystyle\text{(i.i)}}$
$\displaystyle\quad\qquad+2\underbrace{\biggl{\|}\Bigl{(}Bf^{*}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)-\mathbb{E}_{z{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu}\bigl{[}W^{\top}_{\mathtt{DS}}\psi_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})\bigr{]}\Bigr{)}\biggr{\|}_{2}^{2}}_{\displaystyle\text{(i.ii)}}.$
(E.9)

In the sequel, we bound the terms (i.i) and (i.ii) separately. By Assumption
6.2, we have

$\displaystyle\text{(i.i)}\leq\|B\|^{2}_{2}\cdot\bigl{\|}\widehat{f}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})-f^{*}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)\bigr{)}\bigr{\|}_{2}^{2}\leq\mu\cdot\bigl{\|}\widehat{f}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})-f^{*}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)\bigr{)}\bigr{\|}_{2}^{2}.$ (E.10)

On the other hand, by the definition of $B$ in (6.10), we have

(i.ii) $\displaystyle=\Bigl{\|}B\mathbb{E}_{z{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu}\bigl{[}g^{*}_{\mathtt{PT}}(z;{\mathtt{msk}}_{\mathtt{PT}})-W_{\mathtt{SSL}}^{\top}\psi_{\mathtt{DS}}(z;{\mathtt{msk}}_{\mathtt{DS}})\bigr{]}\Bigr{\|}_{2}^{2}$
$\displaystyle\leq\|B\|_{2}^{2}\cdot\Bigl{\|}\mathbb{E}_{z{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu}\bigl{[}g^{*}_{\mathtt{PT}}(z;{\mathtt{msk}}_{\mathtt{PT}})-g_{W_{\mathtt{SSL}}}^{\dagger}(z;{\mathtt{msk}}_{\mathtt{DS}})\bigr{]}\Bigr{\|}_{2}^{2}$
$\displaystyle\leq\mu\cdot\sum_{i=1}^{d}\mathbb{E}_{z{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu}\bigl{[}g^{*}_{{\mathtt{PT}},i}(z;{\mathtt{msk}}_{\mathtt{PT}})-g_{W_{\mathtt{SSL}},i}^{\dagger}(z;{\mathtt{msk}}_{\mathtt{DS}})\bigr{]}^{2}$
$\displaystyle\leq\mu\cdot\sum_{i=1}^{d}\bigl{\|}g^{*}_{{\mathtt{PT}},i}(\cdot;{\mathtt{msk}}_{\mathtt{PT}})-g_{W_{\mathtt{SSL}},i}^{\dagger}(z;{\mathtt{msk}}_{\mathtt{DS}})\bigr{\|}_{\infty}^{2}\leq\mu\cdot\epsilon^{2}_{\mathtt{SSL}}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}}),$
(E.11)

where the third line follows from Assumption 6.2 and the last line follows
from Assumption 6.1. Recall that the approximation error $\mathcal{E}_{\rm
approx}^{\mathtt{PT}}$ for the pretraining process is defined in (6.12).
Taking (E.10) and (E.3) into (E.3), we obtain

(i)
$\displaystyle=\mathbb{E}\Bigl{[}\bigl{\|}\widetilde{f}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})-f_{W_{\mathtt{DS}}}^{\dagger}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{DS}})\bigr{\|}_{2}^{2}\Bigr{]}$
$\displaystyle\leq\mathbb{E}\Bigl{[}2\mu\cdot\bigl{\|}\widehat{f}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu;{\mathtt{msk}}_{\mathtt{PT}})-f^{*}_{\mathtt{PT}}(\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu)\bigr{)}\bigr{\|}_{2}^{2}+2\mu\cdot\epsilon_{\mathtt{SSL}}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}})\Bigr{]}$
$\displaystyle\leq 2\mu\cdot\mathcal{E}_{\rm
approx}^{\mathtt{PT}}+\mu\cdot\mathbb{E}\Bigl{[}\epsilon^{2}_{\mathtt{SSL}}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}})\Bigr{]}$
$\displaystyle=2\mu\cdot\bigl{(}\mathcal{E}_{\rm
approx}^{\mathtt{PT}}+\epsilon^{2}_{\mathtt{SSL}}({\mathtt{msk}}_{\mathtt{PT}},{\mathtt{msk}}_{\mathtt{DS}})\bigr{)},$
(E.12)

where the third line follows from the definition of the regression function
$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]$ for the pretraining
process. Bounding (ii). By the same arguments for (E.5), we have

$\displaystyle\text{(ii)}=\mathbb{E}\biggl{[}\Bigl{\|}\mathbb{E}_{z{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu}\bigl{[}g_{\mathtt{DS}}^{*}(z;{\mathtt{msk}}_{\mathtt{DS}})-g^{\dagger}_{W_{\mathtt{DS}}}(z;{\mathtt{msk}}_{\mathtt{DS}})\bigr{]}\Bigr{\|}_{2}^{2}\biggr{]}.$
(E.13)

By Assumption 6.1, we have

$\displaystyle\Bigl{\|}\mathbb{E}_{z{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu}\bigl{[}g_{\mathtt{DS}}^{*}(z;{\mathtt{msk}}_{\mathtt{DS}})-g^{\dagger}_{W_{\mathtt{DS}}}(z;{\mathtt{msk}}_{\mathtt{DS}})\bigr{]}\Bigr{\|}_{2}^{2}$
$\displaystyle\quad=\sum_{i=1}^{d_{\mathrm{y}}}\mathbb{E}_{z{\,|\,}\mkern
1.5mu\overline{\mkern-1.5muX\mkern-1.5mu}\mkern
1.5mu}\bigl{[}g_{{\mathtt{DS}},i}^{*}(z;{\mathtt{msk}}_{\mathtt{DS}})-g^{\dagger}_{W_{\mathtt{DS}},i}(z;{\mathtt{msk}}_{\mathtt{DS}})\bigr{]}^{2}$
$\displaystyle\quad\leq\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{\|}_{\infty}^{2}\leq\epsilon^{2}_{g}({\mathtt{msk}}_{\mathtt{DS}}).$
(E.14)

Taking (E.3) into (E.13), we obtain

$\displaystyle\text{(ii)}\leq\mathbb{E}\bigl{[}\epsilon^{2}_{g}({\mathtt{msk}}_{\mathtt{DS}})\bigr{]}=\epsilon^{2}_{g}({\mathtt{msk}}_{\mathtt{DS}}).$
(E.15)

Taking (E.3) and (E.15) into (E.8), we conclude the proof of Theorem 6.3. ∎

##  Appendix F Auxiliary Proofs for Generalization

###  F.1 Proof of Dudley Entropy Integral Lemma C.7

###### Proof.

Let $\varepsilon_{j}=2^{1-j}$. Denote by $N_{j}$ the covering of the empirical
function class $\mathcal{F}^{\mathcal{D}_{n}}$ that achieves the covering
number
$N(\mathcal{F}^{\mathcal{D}_{n}},\varepsilon_{j},\|\cdot\|_{r,\infty})$. In
other words, there exists some $\widehat{f}^{j}\in N_{j}$ such that
$\max_{i\in[n]}|f(X_{i})-\widehat{f}^{j}_{i}|\leq\varepsilon_{j}$ for any
$f\in\mathcal{F}$. We have

$\displaystyle\mathbb{E}_{\epsilon}\biggl{[}\sup_{f\in\mathcal{F}}\sum_{i=1}^{n}\epsilon_{i}\cdot
f(X_{i})\biggr{]}$
$\displaystyle\quad=\mathbb{E}_{\epsilon}\biggl{[}\sup_{f\in\mathcal{F}}\biggl{\\{}\sum_{i=1}^{n}\epsilon_{i}\cdot\bigl{(}f(X_{i})-\widehat{f}_{i}^{J}\bigr{)}+\sum_{j=1}^{J-1}\sum_{i=1}^{n}\epsilon_{i}\cdot(\widehat{f}^{j}_{i}-\widehat{f}_{i}^{j+1})-\sum_{i=1}^{n}\epsilon_{i}\cdot\widehat{f}^{1}_{i}\biggr{\\}}\biggr{]}$
$\displaystyle\quad\leq\mathbb{E}_{\epsilon}\biggl{[}\sup_{f\in\mathcal{F}}\sum_{i=1}^{n}\epsilon_{i}\cdot\bigl{(}f(X_{i})-\widehat{f}_{i}^{J}\bigr{)}\biggr{]}+\sum_{j=1}^{J-1}\mathbb{E}_{\epsilon}\biggl{[}\sup_{f\in\mathcal{F}}\sum_{i=1}^{n}\epsilon_{i}\cdot(\widehat{f}^{j}_{i}-\widehat{f}_{i}^{j+1})\biggr{]}+\mathbb{E}_{\epsilon}\biggl{[}\sup_{f\in\mathcal{F}}\sum_{i=1}^{n}\epsilon_{i}\cdot\widehat{f}^{1}_{i}\biggr{]}$
$\displaystyle\quad\leq\underbrace{\mathbb{E}_{\epsilon}\biggl{[}\sup_{f\in\mathcal{F}}\sum_{i=1}^{n}\epsilon_{i}\cdot\bigl{(}f(X_{i})-\widehat{f}_{i}^{J}\bigr{)}\biggr{]}}_{\displaystyle\text{(i)}}+\sum_{j=1}^{J-1}\mathbb{E}_{\epsilon}\biggl{[}\underbrace{\sup_{f\in\mathcal{F}}\sum_{i=1}^{n}\epsilon_{i}\cdot(\widehat{f}^{j}_{i}-\widehat{f}_{i}^{j+1})}_{\displaystyle\text{(ii)}}\biggr{]},$

where the last inequality follows from the choice of $V_{1}=\\{0^{1\times
n}\\}$. In what follows, we bound (i) and (ii) separately. Bounding (i). We
have

$\displaystyle\text{(i)}\leq\mathbb{E}_{\epsilon}\biggl{[}\sum_{i=1}^{n}|\epsilon_{i}|\biggr{]}\cdot\sup_{f\in\mathcal{F}}\max_{i\in[n]}\bigl{|}f(X_{i})-\widehat{f}_{i}^{j}\bigr{|}\leq
N\cdot\varepsilon_{J}.$ (F.1)

Bounding (ii). We have

$\displaystyle\sup_{f\in\mathcal{F},\widehat{f}^{j}\in
N_{j},\widehat{f}^{j+1}\in N_{j+1}}\|\widehat{f}^{j}-\widehat{f}^{j+1}\|_{2}$
$\displaystyle\leq\sup_{f\in\mathcal{F},\widehat{f}^{j}\in
N_{j}}\|\widehat{f}^{j}-f_{\mathcal{D}_{n}}\|_{2}+\sup_{f\in\mathcal{F},\widehat{f}^{j+1}\in
N_{j+1}}\|f_{\mathcal{D}_{n}}-\widehat{f}^{j+1}\|_{2}$
$\displaystyle\leq\sqrt{n}\cdot\sup_{f\in\mathcal{F},\widehat{f}^{j}\in
N_{j}}\|\widehat{f}^{j}-f_{\mathcal{D}_{n}}\|_{\infty}+\sqrt{n}\cdot\sup_{f\in\mathcal{F},\widehat{f}^{j+1}\in
N_{j+1}}\|f_{\mathcal{D}_{n}}-\widehat{f}^{j+1}\|_{\infty}$
$\displaystyle\leq\sqrt{n}\cdot\varepsilon_{j}+\sqrt{n}\cdot\varepsilon_{j+1}=3\sqrt{n}\cdot\varepsilon_{j+1}.$

By the Massart’s finite class lemma, we further have

$\displaystyle\mathbb{E}_{\epsilon}\biggl{[}\sup_{f\in\mathcal{F}}\sum_{i=1}^{n}\epsilon_{i}\cdot(\widehat{f}^{j}_{i}-\widehat{f}_{i}^{j+1})\biggr{]}$
$\displaystyle\leq
3\sqrt{n}\cdot\varepsilon_{j+1}\cdot\sqrt{2\log\bigl{(}|N_{j}|\cdot|N_{j+1}|\bigr{)}}\leq
6\sqrt{n}\cdot\varepsilon_{j+1}\cdot\sqrt{\log|N_{j+1}|},$ (F.2)

Combining (F.1) and (F.2), we obtain

$\displaystyle\mathbb{E}_{\epsilon}\biggl{[}\sup_{f\in\mathcal{F}}\frac{1}{n}\sum_{i=1}^{n}\epsilon_{i}\cdot
f(X_{i})\biggr{]}$
$\displaystyle\leq\varepsilon_{J}+\frac{6}{\sqrt{n}}\cdot\sum_{j=1}^{J-1}\varepsilon_{j+1}\cdot\sqrt{\log|N_{j+1}|}$
$\displaystyle\leq\varepsilon_{J}+\frac{12}{\sqrt{n}}\cdot\sum_{j=1}^{J-1}(\varepsilon_{j}-\varepsilon_{j+1})\cdot\sqrt{\log|N_{j}|}$
$\displaystyle\leq\varepsilon_{J}+\frac{12}{\sqrt{n}}\int_{\varepsilon_{J+1}}^{1}\sqrt{\log
N\bigl{(}\mathcal{F}^{\mathcal{D}_{n}},\varepsilon,\|\cdot\|_{r,\infty}\bigr{)}}{\mathrm{d}}\varepsilon$
$\displaystyle\leq 4\xi+\frac{12}{\sqrt{n}}\int_{\xi}^{1}\sqrt{\log
N\bigl{(}\mathcal{F}^{\mathcal{D}_{n}},\varepsilon,\|\cdot\|_{r,\infty}\bigr{)}}{\mathrm{d}}\varepsilon,$

where the last inequality holds for any $0<\xi<1$ and the smallest $J$ such
that $\xi\leq\varepsilon_{J}$. Therefore, we conclude the proof of Lemma C.7.
∎

####  F.1.1 Proof of Matrix Ball Covering Lemma C.9

###### Proof of Lemma C.9.

Let $M^{\top}=(m_{1},\dots,m_{d_{1}})\in\mathbb{R}^{d_{2}\times d_{1}}$, where
$m_{j}\in\mathbb{R}^{d_{1}\times 1}$, be
${\mathrm{vec}}(M)=(m_{1}^{\top},\dots,m_{d_{1}}^{\top})\in\mathbb{R}^{1\times
d_{1}d_{2}}$. We define the sectional norm for ${\mathrm{vec}}(M)$ as

$\displaystyle\bigl{\|}{\mathrm{vec}}(M)\bigr{\|}_{r(d_{2}),s(d_{1})}=\|M^{\top}\|_{r,s},$

which can be verified to be a proper norm. Thus, in Lemma G.1, setting

$\displaystyle\mathbb{B}=\mathbb{B}_{*}=\bigl{\\{}m\in\mathbb{R}^{d_{1}d_{2}}:\|m\|_{r(d_{2}),s(d_{1})}\leq
1\bigr{\\}}=\bigl{\\{}M^{\top}\in\mathbb{R}^{d_{2}\times
d_{1}}:\|M^{\top}\|_{r,s}\leq 1\bigr{\\}},$

and $\|\cdot\|=\|\cdot\|_{*}=\|\cdot\|_{r(d_{2}),s(d_{1})}$, we obtain

$\displaystyle\log N\Bigl{(}\bigl{\\{}M^{\top}\in\mathbb{R}^{d_{2}\times
d_{1}}:\|M^{\top}\|_{r,s}\leq
R^{\mathrm{m}}\bigr{\\}},\varepsilon,\|\cdot\|_{r,s}\Bigr{)}$
$\displaystyle\quad=\log
N\Bigl{(}\bigl{\\{}m\in\mathbb{R}^{d_{1}d_{2}}:\|m\|_{r(d_{2}),s(d_{1})}\leq
1\bigr{\\}},\varepsilon/R^{\mathrm{m}},\|\cdot\|_{r,s}\Bigr{)}$
$\displaystyle\quad\leq\frac{\text{vol}(2R^{\mathrm{m}}/\varepsilon\cdot\mathbb{B}_{*}+\mathbb{B})}{\text{vol}(\mathbb{B})}=d_{1}d_{2}\cdot\log\biggl{(}1+\frac{2R^{\mathrm{m}}}{\varepsilon}\biggr{)}.$

Therefore, we conclude the proof of Lemma C.9. ∎

###  F.2 Lipschitz Continuity of Multihead Attention

First, we introduce the following lemma characterizing the Lipschitz
continuity of softmax.

######  Lemma F.1 (Lipschitz Continuous Softmax).

Suppose that we have $q,\widehat{q}\in\mathbb{R}^{1\times d_{k}}$, and
$K,\widehat{K}\in\mathbb{R}^{L\times d_{p}}$. With the choice
$\sigma=(2d_{\mathrm{p}})^{1/2s}$, we have

$\displaystyle\bigl{\|}\mathtt{softmax}(q,K)-\mathtt{softmax}(q,\widehat{K})\bigr{\|}_{1}$
$\displaystyle\leq\bigl{(}\|q\|_{r}+\|K^{\top}\|_{r,\infty}\bigr{)}\cdot\|K^{\top}-\widehat{K}^{\top}\|_{r,\infty},$
(F.3)
$\displaystyle\bigl{\|}\mathtt{softmax}(q,K)-\mathtt{softmax}(\widehat{q},K)\bigr{\|}_{1}$
$\displaystyle\leq\bigl{(}\|q\|_{r}+\|K^{\top}\|_{r,\infty}\bigr{)}\cdot\|q-\widehat{q}\|_{r}.$
(F.4)

###### Proof.

Let $p=(p_{1},\dots,p_{L})=\mathtt{softmax}(q,K)\in\mathbb{R}^{1\times L}$. We
have

$\displaystyle
p_{\ell}\propto\exp\bigl{\\{}-\|q-k^{\ell}\|_{2}^{2}/2\sigma^{2}\bigr{\\}}.$

Denote $g_{\ell}=-\|q-k^{\ell}\|_{2}^{2}/2\sigma^{2}$ and
$g=(g_{1},\dots,g_{L})^{\top}$. Let the Jacobian of $p$ with respect to
$k^{\ell}$ be $J_{\ell}\in\mathbb{R}^{L\times d_{\mathrm{p}}}$, and
$P=\text{diag}(p)-p^{\top}p$. We have

$\displaystyle J_{\ell}=\frac{\partial p}{\partial k^{\ell}}=\frac{\partial
p}{\partial g}\cdot\frac{\partial g}{\partial k^{\ell}}=P\frac{\partial
g}{\partial k^{\ell}},$

where $\partial g/\partial k^{\ell}=(e_{\ell}q-E_{\ell,\ell}K)/\sigma^{2}$.
Here $E_{\ell,\ell^{\prime}}\in\mathbb{R}^{L\times L}$ is the unit matrix
whose $(\ell,\ell^{\prime})$-th entry is one. To characterize the Lipschitz
continuity of softmax with respect to $K^{\top}$ and $\|\cdot\|_{r,\infty}$,
we note that

$\displaystyle\biggl{\|}\sum_{\ell=1}^{L}J_{\ell}\Delta_{\ell}\biggr{\|}_{1}\leq\sum_{\ell=1}^{L}\|J_{\ell}\Delta_{\ell}\|_{1}\leq\sum_{\ell=1}^{L}\|J_{\ell}\|_{r\to
1}\cdot\|\Delta_{\ell}\|_{r}\leq\|\Delta\|_{r,\infty}\cdot\sum_{\ell=1}^{L}\|J_{\ell}\|_{r\to
1},$

where $\Delta=[\Delta_{1}^{\top},\dots,\Delta_{L}^{\top}]$. Thus, the
Lipschitz continuity constant is bounded by
$\sum_{\ell=1}^{L}\|J_{\ell}\|_{r\to 1}$. For any $\ell\in[L]$, we have

$\displaystyle\|J_{\ell}\|_{r\to 1}$
$\displaystyle\leq\frac{d_{\mathrm{p}}^{1-1/r}}{\sigma^{2}}\cdot\bigl{\|}P(e_{\ell}q-E_{\ell,\ell}K)\bigr{\|}_{1}$
$\displaystyle=\frac{d_{\mathrm{p}}^{1/s}}{\sigma^{2}}\cdot
p_{\ell}\cdot\bigl{\|}(e_{\ell}-p^{\top})(q-k_{\ell})\bigr{\|}_{1}=\frac{d_{\mathrm{p}}^{1/s}p_{\ell}}{\sigma^{2}}\cdot\|e_{\ell}-P^{\top}\|_{s}\cdot\|q-k_{\ell}\|_{r},$

where the first equality follows from $1/r+1/s=1$. As a consequence, we obtain

$\displaystyle\sum_{\ell\in[L]}\|J_{\ell}\|_{r\to 1}$
$\displaystyle\leq\frac{d_{\mathrm{p}}^{1/s}}{\sigma^{2}}\cdot\sum_{\ell\in[L]}p_{\ell}\cdot\|e_{\ell}-p^{\top}\|_{s}\cdot\|q-k_{\ell}\|_{r}$
$\displaystyle\leq\frac{d_{\mathrm{p}}^{1/s}}{\sigma^{2}}\cdot\bigl{(}\|q\|_{r}+\|K^{\top}\|_{r,\infty}\bigr{)}\cdot\sum_{\ell\in[L]}p_{\ell}\cdot\|e_{\ell}-p^{\top}\|_{s}.$
(F.5)

On the other hand, we have

$\displaystyle\sum_{\ell\in[L]}p_{\ell}\cdot\|e_{\ell}-p^{\top}\|_{s}$
$\displaystyle=\sum_{\ell\in[L]}\biggl{\\{}p_{\ell}\cdot\biggl{[}\sum_{\ell^{\prime}\neq
l}p_{\ell^{\prime}}^{s}+(1-p_{\ell})^{s}\biggr{]}^{1/s}\biggr{\\}}$
$\displaystyle\leq\sum_{\ell\in[L]}p_{\ell}\cdot\bigl{[}2(1-p_{\ell})^{s}\bigr{]}^{1/s}\leq
2^{1/s}.$ (F.6)

Therefore, by (F.2) and (F.2), we have for $\sigma=(2d_{\mathrm{p}})^{1/2s}$,

$\displaystyle\sum_{\ell\in[L]}\|J_{\ell}\|_{r\to
1}\leq\frac{(2d_{\mathrm{p}})^{1/s}}{\sigma^{2}}\cdot\bigl{(}\|q\|_{r}+\|K^{\top}\|_{r,\infty}\bigr{)}.$

Thus, $\mathtt{softmax}(Q,K)$ is
$(\|q\|_{r}+\|K^{\top}\|_{r,\infty})$-Lipschitz in $K^{\top}$ with respect to
$\|\cdot\|_{r,\infty}$, which proves (F.3). Since
$\mathtt{softmax}(Q,K)=\mathtt{softmax}(K,Q)$, we also have (F.4) by the same
arguments. Therefore, we conclude the proof of Lemma F.1. ∎

####  F.2.1 Proof of Parameter Lipschitz Continuity Lemma C.11

###### Proof.

For notational simplicity, we write

$\displaystyle{\mathtt{head}}_{i}=\mathtt{attn}_{\mathtt{SM}}(XW^{{\mathrm{q}}}_{i},XW^{{\mathrm{k}}}_{i},XW^{{\mathrm{v}}}_{i}),\quad\widehat{{\mathtt{head}}}_{i}=\mathtt{attn}_{\mathtt{SM}}(X\widehat{W}^{{\mathrm{q}}}_{i},X\widehat{W}^{{\mathrm{k}}}_{i},X\widehat{W}^{{\mathrm{v}}}_{i}).$

By the definition of multihehead attention in (B.2), we have

$\displaystyle\bigl{\|}{\mathtt{mha}}(X;W)^{\top}-{\mathtt{mha}}(X;\widehat{W})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\leq\biggl{\|}\sum_{i=1}^{h}{\mathtt{head}}^{\top}_{i}-\sum_{i=1}^{h}\widehat{{\mathtt{head}}}^{\top}_{i}\biggr{\|}_{r,\infty}$
$\displaystyle\leq\sum_{i=1}^{h}\bigl{\|}({\mathtt{head}}_{i}-\widehat{{\mathtt{head}}}_{i})^{\top}\bigr{\|}_{r,\infty}.$
(F.7)

By the definition of attention, we have

$\displaystyle\bigl{\|}({\mathtt{head}}_{i}-\widehat{{\mathtt{head}}}_{i})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\quad=\bigl{\|}\mathtt{attn}_{\mathtt{SM}}(XW^{{\mathrm{q}}}_{i},XW^{{\mathrm{k}}}_{i},XW^{{\mathrm{v}}}_{i})^{\top}-\mathtt{attn}_{\mathtt{SM}}(X\widehat{W}^{{\mathrm{q}}}_{i},X\widehat{W}^{{\mathrm{k}}}_{i},X\widehat{W}^{{\mathrm{v}}}_{i})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\quad=\bigl{\|}(XW^{{\mathrm{v}}}_{i})^{\top}\mathtt{softmax}(XW^{{\mathrm{q}}}_{i},XW^{{\mathrm{k}}}_{i})^{\top}-(X\widehat{W}^{{\mathrm{v}}}_{i})^{\top}\mathtt{softmax}(X\widehat{W}^{{\mathrm{q}}}_{i},X\widehat{W}^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\quad=\max_{\ell\in[L]}\bigl{\|}(XW^{{\mathrm{v}}}_{i})^{\top}\mathtt{softmax}(x^{\ell}W^{{\mathrm{q}}}_{i},XW^{{\mathrm{k}}}_{i})^{\top}-(X\widehat{W}^{{\mathrm{v}}}_{i})^{\top}\mathtt{softmax}(x^{\ell}\widehat{W}^{{\mathrm{q}}}_{i},X\widehat{W}^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r}.$
(F.8)

Note that

$\displaystyle\|x^{\ell}W^{{\mathrm{q}}}_{i}\|_{r}+\bigl{\|}(XW^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\leq\bigl{\|}(XW^{{\mathrm{q}}}_{i})^{\top}\bigr{\|}_{r,\infty}+\bigl{\|}(XW^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r,\infty}$
(F.9)
$\displaystyle\leq\Bigl{(}\bigl{\|}(W^{{\mathrm{q}}}_{i})^{\top}\bigr{\|}_{r}+\bigl{\|}(W^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r}\Bigr{)}\cdot\|X^{\top}\|_{r,\infty}\leq(\omega^{{\mathrm{q}}}_{i}+\omega^{{\mathrm{k}}}_{i})\cdot
R.$

Then for any $\ell\in[L]$, we have

$\displaystyle\bigl{\|}(XW^{{\mathrm{v}}}_{i})^{\top}\mathtt{softmax}(x^{\ell}W^{{\mathrm{q}}}_{i},XW^{{\mathrm{k}}}_{i})^{\top}-(X\widehat{W}^{{\mathrm{v}}}_{i})^{\top}\mathtt{softmax}(x^{\ell}\widehat{W}^{{\mathrm{q}}}_{i},X\widehat{W}^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r}$
$\displaystyle\quad\leq\bigl{\|}\mathtt{softmax}(x^{\ell}\widehat{W}^{{\mathrm{q}}}_{i},X\widehat{W}^{{\mathrm{k}}}_{i})X(W^{{\mathrm{v}}}_{i}-\widehat{W}^{{\mathrm{v}}}_{i})\bigr{\|}_{r}$
$\displaystyle\quad\qquad+\Bigl{\|}\bigl{(}\mathtt{softmax}(x^{\ell}W^{{\mathrm{q}}}_{i},XW^{{\mathrm{k}}}_{i})-\mathtt{softmax}(x^{\ell}\widehat{W}^{{\mathrm{q}}}_{i},X\widehat{W}^{{\mathrm{k}}}_{i})\bigr{)}X{W}^{{\mathrm{v}}}_{i}\Bigr{\|}_{r}$
$\displaystyle\quad\leq\bigl{\|}(W^{{\mathrm{v}}}_{i}-\widehat{W}^{{\mathrm{v}}}_{i})^{\top}\bigr{\|}_{r,s}\cdot\|X^{\top}\|_{r,\infty}$
$\displaystyle\quad\qquad+\bigl{\|}\mathtt{softmax}(x^{\ell}W^{{\mathrm{q}}}_{i},XW^{{\mathrm{k}}}_{i})-\mathtt{softmax}(x^{\ell}\widehat{W}^{{\mathrm{q}}}_{i},X\widehat{W}^{{\mathrm{k}}}_{i})\bigr{\|}_{1}\cdot\bigl{\|}({W}^{{\mathrm{v}}}_{i})^{\top}\bigr{\|}_{r}\cdot\|X^{\top}\|_{r,\infty}$
$\displaystyle\quad\leq
R\cdot\varepsilon_{i}^{\mathrm{v}}+\Bigl{(}\|x^{\ell}W^{{\mathrm{q}}}_{i}\|_{r}+\bigl{\|}(XW^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r,\infty}\Bigr{)}\cdot\omega^{{\mathrm{v}}}_{i}R^{2}\cdot(\varepsilon_{i}^{\mathrm{q}}+\varepsilon_{i}^{\mathrm{k}})$
$\displaystyle\quad\leq
R\cdot\varepsilon_{i}^{\mathrm{v}}+(\omega^{{\mathrm{q}}}_{i}+\omega^{{\mathrm{k}}}_{i})\omega^{{\mathrm{v}}}_{i}\cdot
R^{3}\cdot(\varepsilon_{i}^{\mathrm{q}}+\varepsilon_{i}^{\mathrm{k}}),$ (F.10)

where the second inequality follows from Lemma G.3, the third inequality
follows from Lemma F.1, and the last inequality follows from (F.9). Taking
(F.2.1) into (F.2.1) gives

$\displaystyle\bigl{\|}({\mathtt{head}}_{i}-\widehat{{\mathtt{head}}}_{i})^{\top}\bigr{\|}_{r,\infty}\leq
R\cdot\varepsilon_{i}^{\mathrm{v}}+(\omega^{{\mathrm{q}}}_{i}+\omega^{{\mathrm{k}}}_{i})\omega^{{\mathrm{v}}}_{i}\cdot
R^{3}\cdot(\varepsilon_{i}^{\mathrm{q}}+\varepsilon_{i}^{\mathrm{k}}).$ (F.11)

Finally, taking (F.11) into (F.2.1), we obtain

$\displaystyle\bigl{\|}{\mathtt{mha}}(X;W)^{\top}-{\mathtt{mha}}(X;\widehat{W})^{\top}\bigr{\|}_{r,\infty}\leq
R\cdot\sum_{i=1}^{h}\varepsilon_{i}^{\mathrm{v}}+R^{3}\cdot\sum_{i=1}^{h}(\omega^{{\mathrm{q}}}_{i}+\omega^{{\mathrm{k}}}_{i})\cdot(\varepsilon_{i}^{\mathrm{q}}+\varepsilon_{i}^{\mathrm{k}}).$

Therefore, we conclude the proof of Lemma C.11. ∎

####  F.2.2 Proof of Input Lipschitz Continuity Lemma C.13

###### Proof.

In this proof, with slight abuse of notations, we write

$\displaystyle{\mathtt{head}}_{i}=\mathtt{attn}_{\mathtt{SM}}(XW^{{\mathrm{q}}}_{i},XW^{{\mathrm{k}}}_{i},XW^{{\mathrm{v}}}_{i}),\quad\widehat{{\mathtt{head}}}_{i}=\mathtt{attn}_{\mathtt{SM}}(\widehat{X}{W}^{{\mathrm{q}}}_{i},\widehat{X}{W}^{{\mathrm{k}}}_{i},\widehat{X}{W}^{{\mathrm{v}}}_{i}).$

Similar to (F.2.1), we have

$\displaystyle\bigl{\|}{\mathtt{mha}}(X;W)^{\top}-{\mathtt{mha}}(\widehat{X};W)^{\top}\bigr{\|}_{r,\infty}\leq\sum_{i=1}^{h}\bigl{\|}({\mathtt{head}}_{i}-\widehat{{\mathtt{head}}}_{i})^{\top}\bigr{\|}_{r,\infty}.$

For any fixed $\ell\in[L]$, we have

$\displaystyle\bigl{\|}({\mathtt{head}}_{i}-\widehat{{\mathtt{head}}}_{i})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\quad=\bigl{\|}(XW^{{\mathrm{v}}}_{i})^{\top}\mathtt{softmax}(x^{\ell}W^{{\mathrm{q}}}_{i},XW^{{\mathrm{k}}}_{i})^{\top}-(\widehat{X}W^{{\mathrm{v}}}_{i})^{\top}\mathtt{softmax}(\widehat{x}^{\ell}W^{{\mathrm{q}}}_{i},\widehat{X}W^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r}$
$\displaystyle\quad\leq\bigl{\|}\mathtt{softmax}(\widehat{x}^{\ell}W^{{\mathrm{q}}}_{i},\widehat{X}W^{{\mathrm{k}}}_{i})\bigr{\|}_{1}\cdot\bigl{\|}(W^{{\mathrm{v}}}_{i})^{\top}\bigr{\|}_{r}\cdot\|X^{\top}-\widehat{X}^{\top}\|_{r,\infty}$
$\displaystyle\quad\qquad+\bigl{\|}\mathtt{softmax}(x^{\ell}W^{{\mathrm{q}}}_{i},XW^{{\mathrm{k}}}_{i})-\mathtt{softmax}(\widehat{x}^{\ell}{W}^{{\mathrm{q}}}_{i},\widehat{X}{W}^{{\mathrm{k}}}_{i})\bigr{\|}_{1}\cdot\bigl{\|}({W}^{{\mathrm{v}}}_{i})^{\top}\bigr{\|}_{r}\cdot\|X^{\top}\|_{r,\infty}$
$\displaystyle\quad\leq\bigl{\|}(W^{{\mathrm{v}}}_{i})^{\top}\bigr{\|}_{r}\cdot\|X^{\top}-\widehat{X}^{\top}\|_{r,\infty}$
$\displaystyle\quad\qquad+\bigl{\|}(W^{{\mathrm{v}}}_{i})^{\top}\bigr{\|}_{r}\cdot
R^{2}\cdot(\omega^{{\mathrm{q}}}_{i}+\omega^{{\mathrm{k}}}_{i})\cdot\Bigl{(}\|x^{\ell}W^{{\mathrm{q}}}_{i}\|_{r}+\bigl{\|}(XW^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r,\infty}\Bigr{)}$
$\displaystyle\quad\leq\omega^{{\mathrm{v}}}_{i}\cdot\bigl{[}1+R^{2}\cdot(\omega^{{\mathrm{q}}}_{i}+\omega^{{\mathrm{k}}}_{i})^{2}\bigr{]}\cdot\|X^{\top}-\widehat{X}^{\top}\|_{r,\infty},$
(F.12)

where the second inequality follows from Lemma F.1 and (F.9), the third
inequality follows from Lemma F.1, and the last inequality follows from (F.9).
Summing up (F.2.2), we obtain

$\displaystyle\bigl{\|}{\mathtt{mha}}(X;W)^{\top}-{\mathtt{mha}}(\widehat{X};W)^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle\leq\biggl{\\{}\sum_{i=1}^{h}\omega^{{\mathrm{v}}}_{i}\cdot\bigl{[}1+R^{2}\cdot(\omega^{{\mathrm{q}}}_{i}+\omega^{{\mathrm{k}}}_{i})^{2}\bigr{]}\biggr{\\}}\cdot\|X^{\top}-\widehat{X}^{\top}\|_{r,\infty}$
$\displaystyle=\biggl{[}\sum_{i=1}^{h}\omega^{{\mathrm{v}}}_{i}+R^{2}\cdot\sum_{i=1}^{h}(\omega^{{\mathrm{q}}}_{i}+\omega^{{\mathrm{k}}}_{i})^{2}\omega^{{\mathrm{v}}}_{i}\biggr{]}\cdot\|X^{\top}-\widehat{X}^{\top}\|_{r,\infty}.$

Therefore, we conclude the proof of Lemma C.13. ∎

##  Appendix G Auxiliary Lemmas

######  Lemma G.1 (Volume Ratios and Metric Entropy, Wainwright (2019)).

Consider a pair of norms $\|\cdot\|$ and $\|\cdot\|_{*}$ on $\mathbb{R}^{d}$,
and let $\mathbb{B}$ and $\mathbb{B}_{*}$ be their corresponding unit balls.
Then the $\varepsilon$-covering number of $\mathbb{B}_{*}$ in the
$\|\cdot\|$-norm obeys the bounds

$\displaystyle\varepsilon^{-d}\cdot\frac{\text{vol}(\mathbb{B}_{*})}{\text{vol}(\mathbb{B})}\leq
N\bigl{(}\mathbb{B}_{*},\varepsilon,\|\cdot\|\bigr{)}\leq\frac{\text{vol}(2/\varepsilon\cdot\mathbb{B}_{*}+\mathbb{B})}{\text{vol}(\mathbb{B})}.$

######  Lemma G.2 (Caponnetto and De Vito (2007)).

Let $(\Omega,\nu)$ be a probability space and $\xi$ be a random variable on
$\Omega$ taking value in a real separable Hilbert space $\mathcal{H}$. We
assume that there exists constants $B,\sigma>0$ such that

$\displaystyle\bigl{\|}\xi(w)\bigr{\|}_{\mathcal{H}}\leq B/2,\
\mathrm{a.s.},\quad\mathbb{E}\bigl{[}\|\xi\|_{\mathcal{H}}^{2}\bigr{]}\leq\sigma^{2}.$

Then, it holds with probability at least $1-\delta$ that

$\displaystyle\biggl{\|}L^{-1}\sum_{i=1}^{L}\xi(\omega_{i})-\mathbb{E}[\xi]\biggr{\|}\leq
2\biggl{(}\frac{B}{L}+\frac{\sigma}{\sqrt{L}}\biggr{)}\log\frac{2}{\delta}.$

######  Lemma G.3.

Let $(r,s)$ be a conjugate pair. For $M\in\mathbb{R}^{d_{1}\times d_{2}}$,
$b\in\mathbb{R}^{d_{2}\times 1}$, $B\in\mathbb{R}^{d_{2}\times d_{3}}$, we
have

$\displaystyle\|Mb\|_{r}$ $\displaystyle\leq\|M\|_{r,\infty}\cdot\|b\|_{1},$
$\displaystyle\|MB\|_{r,\infty}$
$\displaystyle\leq\|M\|_{r,s}\cdot\|B\|_{r,\infty},$
$\displaystyle\|MB\|_{r,\infty}$
$\displaystyle\leq\|M\|_{r}\cdot\|B\|_{r,\infty}.$

###### Proof.

Let $M=[m_{1}^{\top},\cdots,m_{d_{2}}^{\top}]$ and
$b=[b_{1},\cdots,b_{d_{2}}]^{\top}$. We have

$\displaystyle\|Mb\|_{r}=\biggl{\|}\sum_{j=1}^{d_{2}}b_{j}\cdot
m_{j}\biggr{\|}_{r}\leq\biggl{(}\sum_{j=1}^{d_{2}}|b_{j}|\biggr{)}\cdot\max_{j\in[d_{2}]}\|m_{j}\|_{r}=\|M\|_{r,\infty}\cdot\|b\|_{1}.$

Also, we have

$\displaystyle\|Mb\|_{r}=\biggl{\|}\sum_{j=1}^{d_{2}}b_{j}\cdot
m_{j}\biggr{\|}_{r}\leq\biggl{(}\sum_{j=1}^{d_{2}}|b_{j}|\cdot\|m_{j}\|_{r}\biggr{)}\leq\|M\|_{r,s}\cdot\|b\|_{r}.$

As a consequence, with $B=[b_{1},\cdots,b_{d_{3}}]$, we have

$\displaystyle\|MB\|_{r,\infty}=\max_{j\in[d_{3}]}\|Mb_{j}\|_{r}\leq\max_{j\in[d_{3}]}\|M\|_{r,s}\cdot\|b_{j}\|_{r}=\|M\|_{r,s}\cdot\|B\|_{r,\infty}.$

On the other hand, by the definition of matrix operator norm, we also obtain

$\displaystyle\|MB\|_{r,\infty}=\max_{j\in[d_{3}]}\|Mb_{j}\|_{r}\leq\max_{j\in[d_{3}]}\|M\|_{r}\cdot\|b_{j}\|_{r}=\|M\|_{r}\cdot\|B\|_{r,\infty}.$

Therefore, we conclude the proof of Lemma G.3. ∎

######  Lemma G.4 (Covering Coefficient Bounds).

We have for all $0\leq t\leq T-1$,

$\displaystyle 1+\rho(\mathfrak{W}^{(t)})$
$\displaystyle\leq\widetilde{\omega}^{{\mathrm{v}},(t)}+(\overline{\omega}^{{\mathrm{q}}{\mathrm{k}},(t)})^{2}\omega^{{\mathrm{v}},(t)}\cdot(R^{(t)})^{2}=\widetilde{\rho}^{(t)},$
$\displaystyle R^{\mathtt{mha}}(\mathfrak{W}^{(t)})$ $\displaystyle\leq
R^{{\mathrm{v}},(t)}+\overline{\omega}^{{\mathrm{q}}{\mathrm{k}},(t)}{R}^{{\mathrm{q}}{\mathrm{k}},(t)}\cdot(R^{(t)})^{2}=R^{{\mathtt{mha}},(t)},$

where $\widetilde{\rho}^{(t)}$ and $R^{{\mathtt{mha}},(t)}$ are defined in
(C.11).

###### Proof.

By the definition of $R^{\mathtt{mha}}(\mathfrak{W})$ in (C.15), we have

$\displaystyle R^{\mathtt{mha}}(\mathfrak{W}^{(t)})$
$\displaystyle=\sum_{i=1}^{h}R^{{\mathrm{v}},(t)}_{i}+(R^{{\mathrm{x}},(t)})^{2}\cdot\sum_{i=1}^{h}(\omega_{i}^{{\mathrm{q}},(t)}+\omega_{i}^{{\mathrm{k}},(t)})(R_{i}^{{\mathrm{q}},(t)}+R_{i}^{{\mathrm{k}},(t)})$
$\displaystyle\leq
R^{{\mathrm{v}},(t)}+(R^{(t)})^{2}\cdot\max_{i\in[h]}\\{\omega_{i}^{{\mathrm{q}},(t)}+\omega_{i}^{{\mathrm{k}},(t)}\\}\cdot\sum_{i=1}^{h}(R_{i}^{{\mathrm{q}},(t)}+R_{i}^{{\mathrm{k}},(t)})$
$\displaystyle=R^{{\mathrm{v}},(t)}+\overline{\omega}^{{\mathrm{q}}{\mathrm{k}},(t)}{R}^{{\mathrm{q}}{\mathrm{k}},(t)}\cdot(R^{(t)})^{2}=R^{{\mathtt{mha}},(t)}.$

Also, by the definition of $\rho(\mathfrak{W})$ in (C.16), we have

$\displaystyle 1+{\rho}(\mathfrak{W}^{(t)})$
$\displaystyle=1+\sum_{i=1}^{h}\omega^{{\mathrm{v}},(t)}_{i}+(R^{{\mathrm{x}},(t)})^{2}\cdot\sum_{i=1}^{h}(\omega^{{\mathrm{q}},(t)}_{i}+\omega^{{\mathrm{k}},(t)}_{i})^{2}\omega^{{\mathrm{v}},(t)}_{i}$
$\displaystyle\leq\widetilde{\omega}^{{\mathrm{v}},(t)}+(R^{(t)})^{2}\cdot\max_{i\in[h]}\\{\omega_{i}^{{\mathrm{q}},(t)}+\omega_{i}^{{\mathrm{k}},(t)}\\}^{2}\cdot\sum_{i=1}^{h}\omega^{{\mathrm{v}},(t)}_{i}$
$\displaystyle=\widetilde{\omega}^{{\mathrm{v}},(t)}+(\overline{\omega}^{{\mathrm{q}}{\mathrm{k}},(t)})^{2}\omega^{{\mathrm{v}},(t)}\cdot(R^{(t)})^{2}=\widetilde{\rho}^{(t)}.$

Therefore, we conclude the proof of Lemma G.4. ∎

######  Lemma G.5 (Simplified Covering Coefficient).

It holds that

$\displaystyle\sqrt{\log(1+R^{(T)}R^{{\mathtt{trans}},(T)})}=O\Bigl{(}T\sqrt{\log(2+\gamma)}+\sqrt{T}\sqrt{\log(1+\zeta
R^{(0)})}+\sqrt{\log(1+\kappa/\zeta)}\Bigr{)},$

where the factors on the right hand side are defined in (C.8). Here $R^{(t)}$
and $R^{{\mathtt{trans}},(T)}$ are defined in (C.6) and (C.10), respectively.

###### Proof.

By the definitions of $\kappa^{(t)}$ and $\zeta^{(t)}$ in (C.7), we have

$\displaystyle\frac{R^{{\mathtt{mha}},(t)}}{\widetilde{\rho}^{(t)}}$
$\displaystyle=\frac{R^{{\mathrm{v}},(t)}+\overline{\omega}^{{\mathrm{q}}{\mathrm{k}},(t)}{R}^{{\mathrm{q}}{\mathrm{k}},(t)}\cdot(R^{(t)})^{2}}{\widetilde{\omega}^{{\mathrm{v}},(t)}+(\overline{\omega}^{{\mathrm{q}}{\mathrm{k}},(t)})^{2}\omega^{{\mathrm{v}},(t)}\cdot(R^{(t)})^{2}}\leq\max\biggl{\\{}\frac{R^{{\mathrm{v}},(t)}}{\widetilde{\omega}^{{\mathrm{v}},(t)}},\frac{{R}^{{\mathrm{q}}{\mathrm{k}},(t)}}{\overline{\omega}^{{\mathrm{q}}{\mathrm{k}},(t)}\omega^{{\mathrm{v}},(t)}}\biggr{\\}}\leq\kappa^{(t)},$
$\displaystyle\frac{\widetilde{\rho}^{(t)}}{\widetilde{\omega}^{{\mathrm{v}},(t)}}$
$\displaystyle=1+\frac{(\overline{\omega}^{{\mathrm{q}}{\mathrm{k}},(t)})^{2}\omega^{{\mathrm{v}},(t)}}{\widetilde{\omega}^{{\mathrm{v}},(t)}}\cdot(R^{(t)})^{2}=1+\zeta^{(t)}\cdot(R^{(t)})^{2},$

which gives

$\displaystyle{R}^{{\mathtt{trans}},(T)}$
$\displaystyle=\sum_{t=0}^{T-1}\biggl{[}\biggl{(}\frac{R^{{\mathtt{mha}},(t)}}{\widetilde{\rho}^{(t)}}+\frac{\alpha^{{\mathrm{x}},(t)}R^{\sigma,(t)}+\alpha^{\sigma,(t)}R^{{\mathrm{x}},(t)}}{\widetilde{\alpha}^{(t)}}\biggr{)}\cdot\prod_{\tau=t+1}^{T-1}\frac{\widetilde{\rho}^{(\tau)}}{\widetilde{\omega}^{{\mathrm{v}},(\tau)}}\biggr{]}$
$\displaystyle\leq\sum_{t=0}^{T-1}\biggl{\\{}(\kappa^{(t)}+\kappa^{(t)})\cdot\prod_{\tau=t+1}^{T-1}\bigl{[}1+\zeta^{(t)}\cdot(R^{(t)})^{2}\bigr{]}\biggr{\\}}\leq
2\kappa\cdot\sum_{t=0}^{T-1}\biggl{\\{}\prod_{\tau=t+1}^{T-1}\bigl{[}1+\zeta\cdot(R^{(t)})^{2}\bigr{]}\biggr{\\}},$

where the last line follows from the definition of $\kappa$ and $\zeta$ in
(C.8). By the definition of $R^{(t)}$ in (C.6) and the definition of $\gamma$
in (C.8), we have

$\displaystyle
R^{(t)}=R^{(0)}\cdot\prod_{\tau=0}^{t-1}\widetilde{\omega}^{{\mathrm{v}},(\tau)}\widetilde{\alpha}^{(\tau)}\leq
R^{(0)}\cdot(1+\gamma)^{2t}.$

As a consequence, we obtain

$\displaystyle R^{{\mathtt{trans}},(T)}R^{(T)}$ $\displaystyle\leq 2\kappa
R^{(0)}\cdot(1+\gamma)^{2T}\cdot\sum_{t=0}^{T-1}\prod_{\tau=t+1}^{T-1}\bigl{[}1+\zeta
R^{(0)}\cdot(1+\gamma)^{4\tau}\bigr{]}$ $\displaystyle\leq 2\kappa
R^{(0)}\cdot(1+\gamma)^{2T}\cdot\sum_{t=0}^{T-1}\bigl{[}1+\zeta
R^{(0)}\cdot(1+\gamma)^{4T}\bigr{]}^{T-t+2}$
$\displaystyle=2\kappa\cdot\bigl{[}1+\zeta
R^{(0)}\cdot(1+\gamma)^{4T}\bigr{]}^{3}\cdot\frac{\bigl{[}1+\zeta
R^{(0)}\cdot(1+\gamma)^{4T}\bigr{]}^{T}-1}{\zeta\cdot(1+\gamma)^{2T}}$
$\displaystyle\leq\frac{2\kappa}{\zeta}\cdot\bigl{[}1+\zeta
R^{(0)}\cdot(1+\gamma)^{4T}\bigr{]}^{T+3}.$

Thus, using $\sqrt{a+b}\leq\sqrt{a}+\sqrt{b}$ and
$\log(1+ab)\leq\log(1+a)+\log(1+b)$, we further obtain

$\displaystyle\sqrt{\log(1+R^{{\mathtt{trans}},(T)}R^{(T)})}$
$\displaystyle=O\Bigl{(}\sqrt{T^{2}\log(2+\gamma)+T\log(1+\zeta
R^{(0)})+\log(1+\kappa/\zeta)}\Bigr{)}$
$\displaystyle=O\Bigl{(}T\sqrt{\log(2+\gamma)}+\sqrt{T}\sqrt{\log(1+\zeta
R^{(0)})}+\sqrt{\log(1+\kappa/\zeta)}\Bigr{)}.$

Therefore, we conclude the proof of Lemma G.5. ∎

######  Lemma G.6 (Inter-Layer Magnitude).

Suppose that $R^{(0)}=\max_{i\in[n]}\|X_{i}^{\top}\|_{r,\infty}$. We have

$\displaystyle\bigl{\|}(X^{(t)})^{\top}\bigr{\|}_{r,\infty}\leq\widetilde{\alpha}^{(t)}R^{(t)},\quad\bigl{\|}(X_{\star}^{(t)})^{\top}\bigr{\|}_{r,\infty}\leq
R^{(t)}.$

where $R^{(t)}$ is defined in (C.6).

###### Proof.

Setting $\widehat{X}=0^{L\times d}$ in (F.2.1), we have

$\displaystyle\bigl{\|}{\mathtt{mha}}(X;W)^{\top}\bigr{\|}_{r,\infty}\leq\omega^{o}\cdot\sum_{i=1}^{h}\|{\mathtt{head}}_{i}^{\top}\|_{r,\infty}.$

We have for all $i\in[h]$,

$\displaystyle\|{{\mathtt{head}}}_{i}^{\top}\|_{r,\infty}$
$\displaystyle=\bigl{\|}(X{W}^{{\mathrm{v}}}_{i})^{\top}\mathtt{softmax}(X{W}^{{\mathrm{q}}}_{i},X{W}^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r,\infty}$
$\displaystyle=\max_{l\in[L]}\bigl{\|}(X{W}^{{\mathrm{v}}}_{i})^{\top}\mathtt{softmax}(x_{l}{W}^{{\mathrm{q}}}_{i},X{W}^{{\mathrm{k}}}_{i})^{\top}\bigr{\|}_{r}$
$\displaystyle\leq\max_{l\in[L]}\bigl{\|}\mathtt{softmax}(x_{l}{W}^{{\mathrm{q}}}_{i},X{W}^{{\mathrm{k}}}_{i})\bigr{\|}_{1}\cdot\bigl{\|}(W^{{\mathrm{v}}}_{i})^{\top}\bigr{\|}_{r}\cdot\|X^{\top}\|_{r,\infty}\leq\omega^{{\mathrm{v}}}_{i}\cdot\|X^{\top}\|_{r,\infty},$

which further gives

$\displaystyle\bigl{\|}{\mathtt{mha}}(X;W)^{\top}+X^{\top}\bigr{\|}_{r,\infty}\leq\|X^{\top}\|_{r,\infty}+\biggl{(}\sum_{i=1}^{h}\omega^{{\mathrm{v}}}_{i}\biggr{)}\cdot\|X^{\top}\|_{r,\infty}=(1+\omega^{{\mathrm{v}}})\cdot\|X^{\top}\|_{r,\infty}.$

Thus, we have

$\displaystyle\bigl{\|}(X_{\star}^{(t)})^{\top}\bigr{\|}_{r,\infty}\leq(1+\omega^{{\mathrm{v}},(t)})\cdot\bigl{\|}(X^{(t)})^{\top}\bigr{\|}_{r,\infty}=\widetilde{\omega}^{{\mathrm{v}},(t)}\cdot\bigl{\|}(X^{(t)})^{\top}\bigr{\|}_{r,\infty}.$
(G.1)

On the other hand, setting $\widehat{X}_{\star}^{(t)}=0^{L\times d}$ in
(C.20), we have

$\displaystyle\bigl{\|}(X^{(t+1)})^{\top}\bigr{\|}_{r,\infty}=\bigl{\|}{\mathtt{ffn}}({X}_{\star}^{(t)};A^{(t)})^{\top}\bigr{\|}_{r,\infty}\leq\widetilde{\alpha}^{(t)}\cdot\bigl{\|}({X}_{\star}^{(t)})^{\top}\bigr{\|}_{r,\infty}.$
(G.2)

Recursively applying (G.1) and (G.2), we conclude the proof of Lemma G.6. ∎