# One Step of Gradient Descent is Provably the Optimal In-Context Learner with
One Layer of Linear Self-Attention

Arvind Mahankali
Stanford University
<EMAIL_ADDRESS>
Tatsunori B. Hashimoto
Stanford University
<EMAIL_ADDRESS>
Tengyu Ma
Stanford University
<EMAIL_ADDRESS>

###### Abstract

Recent works have empirically analyzed in-context learning and shown that
transformers trained on synthetic linear regression tasks can learn to
implement ridge regression, which is the Bayes-optimal predictor, given
sufficient capacity (Akyürek et al., 2023), while one-layer transformers with
linear self-attention and no MLP layer will learn to implement one step of
gradient descent (GD) on a least-squares linear regression objective (von
Oswald et al., 2022). However, the theory behind these observations remains
poorly understood. We theoretically study transformers with a single layer of
linear self-attention, trained on synthetic noisy linear regression data.
First, we mathematically show that when the covariates are drawn from a
standard Gaussian distribution, the one-layer transformer which minimizes the
pre-training loss will implement a single step of GD on the least-squares
linear regression objective. Then, we find that changing the distribution of
the covariates and weight vector to a non-isotropic Gaussian distribution has
a strong impact on the learned algorithm: the global minimizer of the pre-
training loss now implements a single step of pre-conditioned GD. However, if
only the distribution of the responses is changed, then this does not have a
large effect on the learned algorithm: even when the response comes from a
more general family of nonlinear functions, the global minimizer of the pre-
training loss still implements a single step of GD on a least-squares linear
regression objective.

##  1 Introduction

Large language models (LLMs) demonstrate the surprising ability of in-context
learning, where an LLM “learns” to solve a task by conditioning on a prompt
containing input-output exemplars (Brown et al., 2020; Lieber et al., 2021;
Radford et al., 2019; Wang and Komatsuzaki, 2021). Recent works have advanced
the understanding of in-context learning via empirical analysis (Min et al.,
2022; Wei et al., 2023; Akyürek et al., 2023; von Oswald et al., 2022; Dai et
al., 2023), but theoretical analysis remains limited (Xie et al., 2022).

A recent line of work (Garg et al., 2022; Akyürek et al., 2023; von Oswald et
al., 2022; Dai et al., 2023) empirically finds that transformers can be
trained to implement algorithms that solve linear regression problems in-
context. Specifically, in each input sequence the transformer is given a set
of in-context examples $(x_{i},y_{i})$, where
$y_{i}=w^{\top}x_{i}+\epsilon_{i}$ with a shared and hidden random coefficient
vector $w$ and random noise $\epsilon_{i}$, and a test example $x$.111In some
settings in these works, the noise is set to $0$. The transformer is then
trained to predict $y=w^{\top}x+\epsilon$, where $\epsilon$ denotes random
noise from the same distribution as $\epsilon_{i}$. These works find that the
transformer outputs a prediction $\hat{y}$ which is similar to the predictions
of existing, interpretable linear regression algorithms, such as gradient
descent (GD) or ordinary least squares, applied to the dataset consisting of
the pairs $(x_{i},y_{i})$. In particular, von Oswald et al. (2022) empirically
show that a one-layer transformer with linear self-attention and no MLP layer
will implement a single step of gradient descent when trained on such a
distribution.

Several works (e.g. Akyürek et al. (2023); Liu et al. (2023); Giannou et al.
(2023)) theoretically study the expressive power of transformers. In the
context of linear regression tasks, Akyürek et al. (2023) describe how
transformers can represent gradient descent, or Sherman-Morrison updates, and
Giannou et al. (2023) describe how transformers can represent Newton’s
algorithm for matrix inversion. However, in addition to the expressive power
of transformers, it is also of interest to understand the behavior of
transformers trained with gradient-based algorithms. Furthermore, it is still
useful to understand the behavior of models with restricted capacity—though
practical LLMs are very expressive, they need to perform many tasks
simultaneously, and therefore the capacity per problem may still be relatively
limited. Thus, motivated by von Oswald et al. (2022), we theoretically study
the global minima of the pre-training loss for one-layer transformers with
linear self-attention on the linear regression data distribution described
above.

Contributions. In this paper, we study transformers with one linear self-
attention layer, and mathematically investigate which algorithms the
transformers implement for synthetically generated linear regression datasets.
We prove that the transformer which implements a single step of gradient
descent on a least squares linear regression objective is the global minimizer
of the pre-training loss. This exactly matches the empirical findings of von
Oswald et al. (2022).

Concretely, we consider a setup similar to von Oswald et al. (2022); Akyürek
et al. (2023). The model we study is a transformer with one linear single-head
self-attention layer, which is the same model as the one empirically studied
by von Oswald et al. (2022). The training data for this transformer consist of
sequences of the form $(x_{1},y_{1},\ldots,x_{n},y_{n})$, where the $x_{i}$
are sampled from $\mathcal{N}(0,{I_{d\times d}})$ and
$y_{i}=w^{\top}x_{i}+\epsilon_{i}$, where $w$ is sampled from
$\mathcal{N}(0,{I_{d\times d}})$ once per sequence, and the $\epsilon_{i}$ are
i.i.d. Gaussian noise with variance $\sigma^{2}$. The pre-training loss is the
expected error that the transformer achieves when predicting $y=w^{\top}x$
given the test example $x$ and the context $(x_{1},y_{1},\ldots,x_{n},y_{n})$,
i.e. the pre-training loss is
$L=\mathbb{E}_{(x_{1},y_{1},\ldots,x_{n},y_{n}),x,y}[(y-\hat{y})^{2}]$, where
$\hat{y}$ is the output of the transformer given
$(x_{1},y_{1},\ldots,x_{n},y_{n})$ and $x$ as input.

We show in Section 3 that the transformer which is the global minimizer of the
pre-training loss $L$ implements one step of gradient descent on a linear
regression objective with the dataset consisting of the $(x_{i},y_{i})$. More
concretely, the transformer implements the prediction algorithm

$\displaystyle\hat{y}$ $\displaystyle=\eta\sum_{i=1}^{n}y_{i}x_{i}^{\top}x\,.$
(1)

However, one step of GD is also preferred in part due to the distribution of
the $x_{i}$. In particular, if the covariance of $x_{i}$ is no longer the
identity matrix, we show (Section 4) that the global minimum of the pre-
training loss corresponds to one step of GD, but with pre-conditioning.

Moreover, interestingly, our theory also suggests that the distribution of
$y_{i}|x_{i}$ does not play such a significant role in the algorithm learned
by the transformer. In Section 5, we study a setting where $y_{i}|x_{i}$ is
nonlinear, but satisfies some mild assumptions, such as invariance to
rotations of the distribution of the $x_{i}$. As a concrete special case, the
target function can be a neural network with any depth/width and i.i.d. random
Gaussian weights. We show in Section 5 that a one-layer transformer with
linear self-attention, which minimizes the pre-training loss, still implements
one step of GD on a linear regression objective. Intuitively, this is likely
because of the constraint imposed by the architecture, which prevents the
transformer from making use of any more complex structure in the $y_{i}$.

Concurrent Works. We discuss the closely related works of Ahn et al. (2023)
and Zhang et al. (2023) which are concurrent with and independent with our
work and were posted prior to our work on arXiv. The work of Ahn et al. (2023)
gives theoretical results very similar to ours. They study one-layer
transformers with linear self-attention with the same parameterization as von
Oswald et al. (2022), and show that with isotropic $x_{i}$, the global
minimizer of the pre-training loss corresponds to one step of gradient descent
on a linear model. They also show that for more general covariance matrices,
the global minimizer of the pre-training loss corresponds to one step of pre-
conditioned gradient descent, where the pre-conditioner matrix can be computed
in terms of the covariance of $x_{i}$.222This result is not exactly the same
as our result in Section 4, since we assume $w\sim\mathcal{N}(0,\Sigma^{-1})$
while they assume $w\sim\mathcal{N}(0,{I_{d\times d}})$.

Different from our work, Ahn et al. (2023) also show additional results for
multi-layer transformers (with linear self-attention) with residual
connections trained on linear regression data. First, they study a restricted
parameterization where in each layer, the product of the projection and value
matrices has only one nonzero entry. In this setting, for two-layer
transformers with linear self-attention, they show that the global minimizer
corresponds to two steps of pre-conditioned GD with diagonal pre-conditioner
matrices, when the data is isotropic. For linear transformers with $k$ layers,
they show that $k$ steps of pre-conditioned GD corresponds to a critical point
of the pre-training loss,333One technical point is that they show there exist
transformers representing this form of pre-conditioned GD having arbitrarily
small gradient, but not that there exists a transformer with gradient exactly
$0$ which represents this form of pre-conditioned GD. where the pre-
conditioner matrix is the inverse of the covariance matrix of the
$x_{i}$.444Here, they assume that $x_{i}\sim\mathcal{N}(0,\Sigma)$ and
$w\sim\mathcal{N}(0,\Sigma^{-1})$, which is the same assumption as our result
in Section 4 but different from their result for one-layer transformers where
$x_{i}\sim\mathcal{N}(0,\Sigma)$. Next, they study a less restrictive
parameterization where the product of the projection and value matrices can be
almost fully dense, and show that a certain critical point of the pre-training
loss for $k$-layer linear transformers corresponds to $k$ steps of a
generalized version of the GD++ algorithm, which was empirically observed by
von Oswald et al. (2022) to be the algorithm learned by $k$-layer linear
transformers.

Zhang et al. (2023) also theoretically study a setting similar to ours. They
not only show that the global minimizer of the pre-training loss implements
one step of GD (the same result as ours), but also show that a one-layer
linear transformer trained with gradient flow will converge to a global
minimizer. They also show that the transformer implements a step of pre-
conditioned GD when the $x_{i}$ are non-isotropic. They also characterize how
the training prompt lengths and test prompt length affect the test-time
prediction error of the trained transformer. Additionally, they consider the
behavior of the trained transformer under distribution shifts, as well as the
training dynamics when the covariance matrices of the $x_{i}$ in different
training prompts can be different.

One additional contribution of our work is that we also consider the case
where the target function in the pre-training data is not a linear function
(Section 5). This suggests that, compared to the distribution of the
covariates, the distribution of the responses at training time does not have
as strong of an effect on the algorithm learned by the transformer. We note
that our proof in this setting is not too different from our proof in Section
3. Zhang et al. (2023) consider the case where the $y_{i}$’s in the test time
prompt are obtained from a nonlinear target function, and consider the
performance on this prompt of the transformer trained on prompts with a linear
target function — this is different from our setting in Section 5 since we
consider the case where the training prompts themselves are obtained with a
nonlinear target function.

##  2 Setup

Our setup is similar to von Oswald et al. (2022).

One-Layer Transformer with Linear Self-Attention. A linear self-attention
layer with width $s$ consists of the following parameters: a key matrix
$W_{K}\in\mathbb{R}^{s\times s}$, a query matrix $W_{Q}\in\mathbb{R}^{s\times
s}$, and a value matrix $W_{V}\in\mathbb{R}^{s\times s}$. Given a sequence of
$T>1$ tokens $(v_{1},v_{2},\ldots,v_{T})$, the output of the linear self-
attention layer is defined to be
$({\hat{v}}_{1},{\hat{v}}_{2},\ldots,{\hat{v}}_{T})$, where for $i\in[T]$ with
$i>1$,

$\displaystyle{\hat{v}}_{i}$
$\displaystyle=\sum_{j=1}^{i-1}(W_{V}v_{j})(v_{j}^{\top}W_{K}^{\top}W_{Q}v_{i})\,,$
(2)

and ${\hat{v}}_{1}=0$. In particular, the output on the $T^{\textup{th}}$
token is

$\displaystyle{\hat{v}}_{T}$
$\displaystyle=\sum_{j=1}^{T-1}(W_{V}v_{j})(v_{j}^{\top}W_{K}^{\top}W_{Q}v_{T})\,.$
(3)

As in the theoretical construction of von Oswald et al. (2022), we do not
consider the attention score between a token $v_{i}$ and itself. Our overall
transformer is then defined to be a linear self-attention layer with key
matrix $W_{K}$, query matrix $W_{Q}$, and value matrix $W_{V}$, together with
a linear head $h\in\mathbb{R}^{s}$ which is applied to the last token. Thus,
the final output of the transformer is $h^{\top}{\hat{v}}_{T}$. We will later
instantiate this one-layer transformer with $s=d+1$, where $d$ is the
dimension of the inputs $x_{i}$. We note that this corresponds to a single
head of linear self-attention, while one could also consider multi-head self-
attention.

Linear Regression Data Distribution. The pretraining data distribution
consists of sequences $D=(x_{1},y_{1},\ldots,x_{n+1},y_{n+1})$. Here, the
exemplars $x_{i}$ are sampled i.i.d. from $\mathcal{N}(0,{I_{d\times d}})$.
Then, a weight vector $w\in\mathbb{R}^{d}$ is sampled from
$\mathcal{N}(0,{I_{d\times d}})$, freshly for each sequence. Finally, $y_{i}$
is computed as $y_{i}=w^{\top}x_{i}+\epsilon_{i}$ where
$\epsilon_{i}\sim\mathcal{N}(0,\sigma^{2})$ for some $\sigma>0$. We consider
the vector $v_{i}=\left[\begin{array}[]{c}x_{i}\\\
y_{i}\end{array}\right]\in\mathbb{R}^{d+1}$ to be a token — in other words,
the sequence $(x_{1},y_{1},\ldots,x_{n+1},y_{n+1})$ is considered to have
$n+1$ tokens (rather than $2(n+1)$ tokens). We use ${\mathcal{T}}$ to denote
the distribution of sequences defined in this way.

At both training and test time,
$(x_{1},y_{1},\ldots,x_{n},y_{n},x_{n+1},y_{n+1})$ is generated according to
the pretraining distribution ${\mathcal{T}}$, i.e. the $x_{i}$ are sampled
i.i.d. from $\mathcal{N}(0,{I_{d\times d}})$, a new weight vector
$w\in\mathbb{R}^{d}$ is also sampled from $\mathcal{N}(0,{I_{d\times d}})$,
and $y_{i}=w^{\top}x_{i}+\epsilon_{i}$ where the $\epsilon_{i}$ are sampled
i.i.d. from $\mathcal{N}(0,\sigma^{2})$. Then, the in-context learner is
presented with $x_{1},y_{1},\ldots,x_{n},y_{n},x_{n+1}$, and must predict
$y_{n+1}$. We refer to $x_{1},\ldots,x_{n}$ as the support exemplars and
$x_{n+1}$ as the query exemplar. Here, $v_{1},\ldots,v_{n}$ are defined as
above, but $v_{n+1}=\left[\begin{array}[]{c}x_{i}\\\ 0\end{array}\right]$,
following the notation of von Oswald et al. (2022).555If we were to treat
$x_{i}$ and $y_{i}$ as separate tokens, then we would need to deal with
attention scores between $y_{i}$ and $y_{j}$ for $i\neq j$, as well as
attention scores between $y_{i}$ and $x_{j}$ for $i\neq j$. Our current setup
simplifies the analysis. We note that this is not significantly different from
the standard in-context learning setting, since even though the final token
$v_{n+1}$ has $0$ as an extra coordinate, it does not provide the transformer
with any additional information about $y_{n+1}$.

Loss Function. Given a one-layer transformer with linear self-attention and
width $d+1$, with key matrix $W_{K}\in{\mathbb{R}^{(d+1)\times(d+1)}}$, query
matrix $W_{Q}\in{\mathbb{R}^{(d+1)\times(d+1)}}$, and value matrix
$W_{V}\in{\mathbb{R}^{(d+1)\times(d+1)}}$, and with a head
$h\in\mathbb{R}^{d+1}$, the loss of this transformer on our linear regression
data distribution is formally defined as

$\displaystyle L(W_{K},W_{Q},W_{V},h)$
$\displaystyle=\mathbb{E}_{D\sim{\mathcal{T}}}[(h^{\top}{\hat{v}}_{n+1}-y_{n+1})^{2}]\,,$
(4)

where as defined above, ${\hat{v}}_{n+1}$ is the output of the linear self-
attention layer on the $(n+1)^{\textup{th}}$ token, which in this case is
$\left[\begin{array}[]{c}x_{n+1}\\\ 0\end{array}\right]$.

We now rewrite the loss function and one-layer transformer in a more
convenient form. As a convenient shorthand, for any test-time sequence
$D=(x_{1},y_{1},\ldots,x_{n+1},0)$, we write
${\widetilde{D}}=(x_{1},y_{1},\ldots,x_{n},y_{n})$, i.e. the prefix of $D$
that does not include $x_{n+1}$ and $y_{n+1}$. We also define

$\displaystyle{G_{\widetilde{D}}}$
$\displaystyle=\sum_{i=1}^{n}\left[\begin{array}[]{c}x_{i}\\\
y_{i}\end{array}\right]\left[\begin{array}[]{c}x_{i}\\\
y_{i}\end{array}\right]^{\top}\,.$ (9)

With this notation, we can write the prediction obtained from the transformer
on the final token as

$\displaystyle\hat{y}_{n+1}$
$\displaystyle=h^{\top}W_{V}{G_{\widetilde{D}}}W_{K}^{\top}W_{Q}v_{n+1}\,.$
(10)

where $v_{n+1}=\left[\begin{array}[]{c}x_{n+1}\\\ 0\end{array}\right]$.
Additionally, we also define the matrix $X\in\mathbb{R}^{n\times d}$ as the
matrix whose $i^{th}$ row is the row vector $x_{i}^{\top}$, i.e.

$\displaystyle X$
$\displaystyle=\left[\begin{array}[]{ccc}\cdots&x_{1}^{\top}&\cdots\\\
\cdots&x_{2}^{\top}&\cdots\\\ \vdots&\vdots&\vdots\\\
\cdots&x_{n}^{\top}&\cdots\end{array}\right]\,,$ (15)

and we define the vector ${\vec{y}}\in\mathbb{R}^{n}$ as the vector whose
$i^{\textup{th}}$ entry is $y_{i}$, i.e.

$\displaystyle{\vec{y}}$ $\displaystyle=\left[\begin{array}[]{c}y_{1}\\\
y_{2}\\\ \vdots\\\ y_{n}\end{array}\right]\,.$ (20)

Finally, it is worth noting that we can write the loss function as

$\displaystyle L(W_{K},W_{Q},W_{V},h)$
$\displaystyle=\mathbb{E}_{D\sim{\mathcal{T}}}[(h^{\top}W_{V}{G_{\widetilde{D}}}W_{K}^{\top}W_{Q}v_{n+1}-y_{n+1})^{2}]\,.$
(21)

Thus, for $w\in\mathbb{R}^{d+1}$ and $M\in{\mathbb{R}^{(d+1)\times(d+1)}}$, if
we define

$\displaystyle L(w,M)$
$\displaystyle=\mathbb{E}_{D\sim{\mathcal{T}}}[(w^{\top}{G_{\widetilde{D}}}Mv_{n+1}-y_{n+1})^{2}]\,,$
(22)

then $L(W_{V}^{\top}h,W_{K}^{\top}W_{Q})=L(W_{K},W_{Q},W_{V},h)$. Note that we
have a slight abuse of notation, and $L$ has two different meanings depending
on the number of arguments. Finally, with the change of variables
$M=W_{K}^{\top}W_{Q}$ and $w=W_{V}^{\top}h$, we can write the prediction of
the transformer as $w^{\top}{G_{\widetilde{D}}}Mv_{n+1}$. Thus, the output of
the transformer only depends on the parameters through
$w^{\top}{G_{\widetilde{D}}}M$.

Additional Notation. For a matrix $A\in\mathbb{R}^{d\times d}$, we write
$A_{i:j,:}$ to denote the sub-matrix of $A$ that contains the rows of $A$ with
indices between $i$ and $j$ (inclusive). Similarly, we write $A_{:,i:j}$ to
denote the sub-matrix of $A$ that contains the columns of $A$ with indices
between $i$ and $j$ (inclusive). We write $A_{i:j,k:l}$ to denote the sub-
matrix of $A$ containing the entries with row indices between $i$ and $j$
(inclusive) and column indices between $k$ and $l$ (inclusive).

##  3 Main Result for Linear Models

######  Theorem 1 (Global Minimum for Linear Regression Data).

Suppose $(W_{K}^{*},W_{Q}^{*},W_{V}^{*},h^{*})$ is a global minimizer of the
loss $L$. Then, the corresponding one-layer transformer with linear self-
attention implements one step of gradient descent on a linear model with some
learning rate $\eta>0$. More concretely, given a query token
$v_{n+1}=\left[\begin{array}[]{c}x_{n+1}\\\ 0\end{array}\right]$, the
transformer outputs $\eta\sum_{i=1}^{n}y_{i}x_{i}^{\top}x_{n+1}$, where
$\eta=\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\hat{w}_{\widetilde{D}}^{\top}X^{\top}{\vec{y}}]}{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}XX^{\top}{\vec{y}}]}$.
Here given a prefix ${\widetilde{D}}$ of a test-time data sequence $D$, we let
$\hat{w}_{\widetilde{D}}$ denote the solution to ridge regression on $X$ and
${\vec{y}}$ with regularization strength $\sigma^{2}$.

One such construction is as follows. von Oswald et al. (2022) describe
essentially the same construction, but our result shows that it is a global
minimum of the loss function, while von Oswald et al. (2022) do not
theoretically study the construction aside from showing that it is equivalent
to one step of gradient descent. We define

$\displaystyle W_{K}^{*}=\begin{pmatrix}{I_{d\times d}}&0\\\
0&0\end{pmatrix}\,,W_{Q}^{*}=\begin{pmatrix}{I_{d\times d}}&0\\\
0&0\end{pmatrix}\,,W_{V}^{*}=\begin{pmatrix}0&0\\\
0&\eta\end{pmatrix}\,,h^{*}=\left[\begin{array}[]{c}0\\\
1\end{array}\right]\,.$ (25)

Here, the unique value of $\eta$ which makes this construction a global
minimum is
$\eta=\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\hat{w}_{\widetilde{D}}^{\top}X^{\top}{\vec{y}}]}{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}XX^{\top}{\vec{y}}]}$.

To see why this construction implements a single step of gradient descent on a
linear model, note that given test time inputs
$x_{1},y_{1},\ldots,x_{n},y_{n},x_{n+1}$, if we write
$v_{i}=\left[\begin{array}[]{c}x_{i}\\\ y_{i}\end{array}\right]$ for $i\leq n$
and $v_{n+1}=\left[\begin{array}[]{c}x_{n+1}\\\ 0\end{array}\right]$, then the
output of the corresponding transformer would be

$\displaystyle(h^{*})^{\top}\sum_{i=1}^{n}(W_{V}^{*}v_{i})(v_{i}^{\top}(W_{K}^{*})^{\top}W_{Q}^{*}v_{n+1})$
$\displaystyle=\begin{pmatrix}0&\eta\end{pmatrix}\sum_{i=1}^{n}\left[\begin{array}[]{c}x_{i}\\\
y_{i}\end{array}\right]\left[\begin{array}[]{c}x_{i}\\\
y_{i}\end{array}\right]^{\top}\begin{pmatrix}{I_{d\times d}}&0\\\
0&0\end{pmatrix}\left[\begin{array}[]{c}x_{n+1}\\\ 0\end{array}\right]$ (32)
$\displaystyle=\eta\sum_{i=1}^{n}y_{i}x_{i}^{\top}x_{n+1}\,.$ (33)

On the other hand, consider linear regression with total squared error as the
loss function, using the $x_{i}$ and $y_{i}$. Here, the loss function would be

$\displaystyle L(w)$
$\displaystyle=\frac{1}{2}\sum_{i=1}^{n}(w^{\top}x_{i}-y_{i})^{2}\,,$ (34)

meaning that the gradient is

$\displaystyle\nabla_{w}L(w)$
$\displaystyle=\sum_{i=1}^{n}(w^{\top}x_{i}-y_{i})x_{i}\,.$ (35)

In particular, if we initialize gradient descent at $w_{0}=0$, then after one
step of gradient descent with learning rate $\eta$, the iterate would be at
$w_{1}=\eta\sum_{i=1}^{n}y_{i}x_{i}$ — observe that the final expression in
Equation 33 is exactly $w_{1}^{\top}x_{n+1}$.

Now, we give an overview of the proof of Theorem 1. By the discussion in
Section 2, it suffices to show that
$L((W_{V}^{*})^{\top}h^{*},(W_{K}^{*})^{\top}W_{Q}^{*})$ is the global minimum
of $L(w,M)$. The first step of the proof is to rewrite the loss in a more
convenient form, getting rid of the expectation over $x_{n+1}$ and $y_{n+1}$:

######  Lemma 1.

Let $\hat{w}_{\widetilde{D}}$ be the solution to ridge regression with
regularization strength $\sigma^{2}$ on the exemplars
$(x_{1},y_{1}),\ldots,(x_{n},y_{n})$ given in a context ${\widetilde{D}}$.
Then, there exists a constant $C\geq 0$, which is independent of $w,M$, such
that

$\displaystyle L(w,M)$
$\displaystyle=C+\mathbb{E}_{D\sim{\mathcal{T}}}\|M_{:,1:d}^{\top}{G_{\widetilde{D}}}w-\hat{w}_{\widetilde{D}}\|_{2}^{2}\,.$
(36)

As discussed towards the end of Section 2, the prediction can be written as
$w^{\top}{G_{\widetilde{D}}}Mv_{n+1}$ where
$v_{n+1}=\left[\begin{array}[]{c}x_{n+1}\\\ 0\end{array}\right]$, meaning that
the effective linear predictor implemented by the transformer is the linear
function from $\mathbb{R}^{d}$ to $\mathbb{R}$ with weight vector
$M_{:,1:d}^{\top}{G_{\widetilde{D}}}w$. Thus, we can interpret Lemma 1 as
saying that the loss function encourages the effective linear predictor to
match the Bayes-optimal predictor $\hat{w}_{\widetilde{D}}$. Note that it is
not possible for the effective linear predictor of the transformer to match
$\hat{w}_{\widetilde{D}}$ exactly, since the transformer can only implement a
linear or quadratic function of the $x_{i}$, while representing
$\hat{w}_{\widetilde{D}}$ requires computing $(X^{\top}X+\sigma^{2}I)^{-1}$,
and this is a much more complex function of the $x_{i}$.

###### Proof of Lemma 1.

We can write

$\displaystyle L(w,M)$
$\displaystyle=\mathbb{E}_{D\sim{\mathcal{T}}}[(y_{n+1}-w^{\top}{G_{\widetilde{D}}}Mv_{n+1})^{2}]$
(37)
$\displaystyle=\mathbb{E}_{{\widetilde{D}},x_{n+1}}\Big{[}\mathbb{E}_{y_{n+1}}[(y_{n+1}-w^{\top}{G_{\widetilde{D}}}Mv_{n+1})^{2}\mid
x_{n+1},{\widetilde{D}}]\Big{]}\,.$ (38)

Let us simplify the inner expectation. For convenience, fix ${\widetilde{D}}$
and $x_{n+1}$, and consider the function $g:\mathbb{R}^{d}\to\mathbb{R}$ given
by

$\displaystyle g(u)$ $\displaystyle=\mathbb{E}_{y_{n+1}}[(u\cdot
x_{n+1}-y_{n+1})^{2}\mid{\widetilde{D}},x_{n+1}]\,.$ (39)

It is a well-known fact that the minimizer of $g(u)$ is given by
$\hat{w}_{\widetilde{D}}$ where
$\hat{w}_{\widetilde{D}}=(X^{\top}X+\sigma^{2}I)^{-1}X^{\top}{\vec{y}}$ is the
solution to ridge regression on $X$ and ${\vec{y}}$ with regularization
strength $\sigma^{2}$. Furthermore,

$\displaystyle
0=\nabla_{u}g(\hat{w}_{\widetilde{D}})=\mathbb{E}_{y_{n+1}}[2(\hat{w}_{\widetilde{D}}\cdot
x_{n+1}-y_{n+1})x_{n+1}\mid{\widetilde{D}},x_{n+1}]\,,$ (40)

and in particular, taking the dot product of both sides with
$u-\hat{w}_{\widetilde{D}}$ (for any vector $u\in\mathbb{R}^{d}$) gives

$\displaystyle\mathbb{E}_{y_{n+1}}[(\hat{w}_{\widetilde{D}}\cdot
x_{n+1}-y_{n+1})\cdot(u\cdot x_{n+1}-\hat{w}_{\widetilde{D}}\cdot
x_{n+1})]=0\,.$ (41)

Thus, letting $u=w^{\top}{G_{\widetilde{D}}}M_{:,1:d}$, we can simplify the
inner expectation in Equation 38:

$\displaystyle\mathbb{E}_{y_{n+1}}[($ $\displaystyle
y_{n+1}-w^{\top}{G_{\widetilde{D}}}M_{:,1:d}x_{n+1})^{2}\mid
x_{n+1},{\widetilde{D}}]$ (42)
$\displaystyle=\mathbb{E}_{y_{n+1}}[(y_{n+1}-\hat{w}_{\widetilde{D}}^{\top}x_{n+1}+\hat{w}_{\widetilde{D}}^{\top}x_{n+1}-w^{\top}{G_{\widetilde{D}}}M_{:,1:d}x_{n+1})^{2}\mid
x_{n+1},{\widetilde{D}}]$ (43)
$\displaystyle=\mathbb{E}_{y_{n+1}}[(y_{n+1}-\hat{w}_{\widetilde{D}}^{\top}x_{n+1})^{2}\mid
x_{n+1},{\widetilde{D}}]+(\hat{w}_{\widetilde{D}}^{\top}x_{n+1}-w^{\top}{G_{\widetilde{D}}}M_{:,1:d}x_{n+1})^{2}$
(44)
$\displaystyle\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,+2\cdot\mathbb{E}_{y_{n+1}}[(y_{n+1}-\hat{w}_{\widetilde{D}}^{\top}x_{n+1})(\hat{w}_{\widetilde{D}}^{\top}x_{n+1}-w^{\top}{G_{\widetilde{D}}}M_{:,1:d}x_{n+1})\mid
x_{n+1},{\widetilde{D}}]$ (45)
$\displaystyle=\mathbb{E}_{y_{n+1}}[(y_{n+1}-\hat{w}_{\widetilde{D}}^{\top}x_{n+1})^{2}\mid
x_{n+1},{\widetilde{D}}]+(\hat{w}_{\widetilde{D}}^{\top}x_{n+1}-w^{\top}{G_{\widetilde{D}}}M_{:,1:d}x_{n+1})^{2}\,.$
(By Equation 41)

We can further write the final expression as
$C_{x_{n+1},{\widetilde{D}}}+(\hat{w}_{\widetilde{D}}^{\top}x_{n+1}-w^{\top}{G_{\widetilde{D}}}M_{:,1:d}x_{n+1})^{2}$,
where $C_{x_{n+1},{\widetilde{D}}}$ is a constant that depends on $x_{n+1}$
and ${\widetilde{D}}$ but is independent of $w$ and $M$. Thus, we have

$\displaystyle L(w,M)$
$\displaystyle=\mathbb{E}_{{\widetilde{D}},x_{n+1}}[C_{x_{n+1},{\widetilde{D}}}]+\mathbb{E}_{{\widetilde{D}},x_{n+1}}[(\hat{w}_{\widetilde{D}}^{\top}x_{n+1}-w^{\top}{G_{\widetilde{D}}}M_{:,1:d}x_{n+1})^{2}]$
(46)
$\displaystyle=C+\mathbb{E}_{{\widetilde{D}}}\|\hat{w}_{\widetilde{D}}-M_{:,1:d}^{\top}{G_{\widetilde{D}}}w\|_{2}^{2}\,,$
(B.c. $x_{n+1}\sim\mathcal{N}(0,{I_{d\times d}})$)

where $C$ is a constant which is independent of $w$ and $M$. ∎

Next, the key step is to replace $\hat{w}_{\widetilde{D}}$ in the above lemma
by $\eta X^{\top}{\vec{y}}$.

######  Lemma 2.

There exists a constant $C_{1}\geq 0$ which is independent of $w,M$, such that

$\displaystyle L(w,M)$
$\displaystyle=C_{1}+\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|M_{:,1:d}^{\top}{G_{\widetilde{D}}}w-\eta
X^{\top}{\vec{y}}\|_{2}^{2}\,.$ (47)

Lemma 2 says that the loss depends entirely on how far the effective linear
predictor is from $\eta X^{\top}{\vec{y}}$. It immediately follows from this
lemma that $(W_{K},W_{Q},W_{V},h)$ is a global minimizer of the loss if and
only if the effective linear predictor of the corresponding transformer is
$\eta X^{\top}{\vec{y}}$. Thus, Theorem 1 follows almost directly from Lemma
2, and in the rest of this section, we give an outline of the proof of Lemma 2
— the detailed proofs of Theorem 1 and Lemma 2 are in Appendix A.

Proof Strategy for Lemma 2. Our overall proof strategy is to show that the
gradients of

$\displaystyle L(w,M)$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|M_{:,1:d}^{\top}{G_{\widetilde{D}}}w-\hat{w}_{\widetilde{D}}\|_{2}^{2}$
(48)

and

$\displaystyle L^{\prime}(w,M)$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|M_{:,1:d}^{\top}{G_{\widetilde{D}}}w-\eta
X^{\top}{\vec{y}}\|_{2}^{2}$ (49)

are equal at every $w,M$, from which Lemma 2 immediately follows.666This is
because $L$ and $L^{\prime}$ are defined everywhere on
$\mathbb{R}^{d+1}\times\mathbb{R}^{(d+1)\times(d+1)}$, and for any two
differentiable functions $f,g$ defined on an open connected subset
$S\subset\mathbb{R}^{k}$, if the gradients of $f$ and $g$ are identical on
$S$, then $f$ and $g$ are equal on $S$ up to an additive constant. This can be
shown using the fundamental theorem of calculus. For simplicity, we write
$A=M_{:,1:d}^{\top}$, so without loss of generality we can show that the
gradients of the following two loss functions are identical:

$\displaystyle J_{1}(A,w)$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|A{G_{\widetilde{D}}}w-\hat{w}_{\widetilde{D}}\|_{2}^{2}$
(50)

and

$\displaystyle J_{2}(A,w)$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|A{G_{\widetilde{D}}}w-\eta
X^{\top}{\vec{y}}\|_{2}^{2}\,.$ (51)

In this section, we discuss the gradients with respect to $w$ — we use the
same proof ideas to show that the gradients with respect to $A$ are the same.
We have

$\displaystyle\nabla_{w}J_{1}(A,w)$
$\displaystyle=2\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}(A{G_{\widetilde{D}}}w-\hat{w}_{\widetilde{D}})$
(52)

and

$\displaystyle\nabla_{w}J_{2}(A,w)$
$\displaystyle=2\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}(A{G_{\widetilde{D}}}w-\eta
X^{\top}{\vec{y}})\,.$ (53)

Thus, showing that these two gradients are equal for all $A,w$ reduces to
showing that for all $A$, we have

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}\hat{w}_{\widetilde{D}}=\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}X^{\top}{\vec{y}}\,.$
(54)

Recall that ${G_{\widetilde{D}}}$ is defined as

$\displaystyle{G_{\widetilde{D}}}$
$\displaystyle=\sum_{i=1}^{n}\left[\begin{array}[]{cc}x_{i}x_{i}^{\top}&y_{i}x_{i}\\\
y_{i}x_{i}^{\top}&y_{i}^{2}\end{array}\right]\,.$ (57)

Our first key observation is that for any $i\in[n]$ and any odd positive
integer $k$, $\mathbb{E}[y_{i}^{k}\mid X]=0$, since
$y_{i}=w^{\top}x_{i}+\epsilon_{i}$, and both $w^{\top}x_{i}$ and
$\epsilon_{i}$ have distributions which are symmetric around $0$. This
observation also extends to any odd-degree monomial of the $y_{i}$. Using this
observation, we can simplify the left-hand side of Equation 54 as follows. We
can first write it as

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}\hat{w}_{\widetilde{D}}$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\sum_{i=1}^{n}\left[\begin{array}[]{c}x_{i}x_{i}^{\top}(A^{\top})_{1:d,:}\hat{w}_{\widetilde{D}}+y_{i}x_{i}(A^{\top})_{d+1,:}\hat{w}_{\widetilde{D}}\\\
y_{i}x_{i}^{\top}(A^{\top})_{1:d,:}\hat{w}_{\widetilde{D}}+y_{i}^{2}(A^{\top})_{d+1,:}\hat{w}_{\widetilde{D}}\end{array}\right]\,.$
(60)

Then, since
$\hat{w}_{\widetilde{D}}=(X^{\top}X+\sigma^{2}I)^{-1}X^{\top}{\vec{y}}$, each
entry of $\hat{w}_{\widetilde{D}}$ has an odd degree in the $y_{i}$, meaning
the terms $x_{i}x_{i}^{\top}(A^{\top})_{1:d,:}\hat{w}_{\widetilde{D}}$ and
$y_{i}^{2}(A^{\top})_{d+1,:}\hat{w}_{\widetilde{D}}$ in the above equation
vanish after taking the expectation. Thus, we obtain

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}\hat{w}_{\widetilde{D}}$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\sum_{i=1}^{n}\left[\begin{array}[]{c}y_{i}x_{i}A_{:,d+1}^{\top}\hat{w}_{\widetilde{D}}\\\
y_{i}x_{i}^{\top}A_{:,1:d}^{\top}\hat{w}_{\widetilde{D}}\end{array}\right]=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\left[\begin{array}[]{c}X^{\top}{\vec{y}}A_{:,d+1}^{\top}\hat{w}_{\widetilde{D}}\\\
{\vec{y}}^{\top}XA_{:,1:d}^{\top}\hat{w}_{\widetilde{D}}\end{array}\right]\,.$
(65)

Since each entry of $\eta X^{\top}{\vec{y}}$ has an odd degree in the $y_{i}$,
in order to simplify the right-hand side of Equation 54, we can apply the same
argument but with $\hat{w}_{\widetilde{D}}$ replaced by $\eta
X^{\top}{\vec{y}}$, obtaining

$\displaystyle\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}X^{\top}{\vec{y}}$
$\displaystyle=\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\left[\begin{array}[]{c}X^{\top}{\vec{y}}A_{:,d+1}^{\top}X^{\top}{\vec{y}}\\\
{\vec{y}}^{\top}XA_{:,1:d}^{\top}X^{\top}{\vec{y}}\end{array}\right]\,.$ (68)

Thus, showing Equation 54 reduces to showing that

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\left[\begin{array}[]{c}X^{\top}{\vec{y}}A_{:,d+1}^{\top}\hat{w}_{\widetilde{D}}\\\
{\vec{y}}^{\top}XA_{:,1:d}^{\top}\hat{w}_{\widetilde{D}}\end{array}\right]$
$\displaystyle=\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\left[\begin{array}[]{c}X^{\top}{\vec{y}}A_{:,d+1}^{\top}X^{\top}{\vec{y}}\\\
{\vec{y}}^{\top}XA_{:,1:d}^{\top}X^{\top}{\vec{y}}\end{array}\right]\,.$ (73)

To show Equation 73, our key tool is Lemma 4, which follows from Lemma 3.

######  Lemma 3.

There exists a scalar $c_{1}$ such that
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]=c_{1}{I_{d\times
d}}$, and there exists a scalar $c_{2}$ such that
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}]=c_{2}{I_{d\times
d}}$.

######  Lemma 4.

If
$\eta=\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\hat{w}_{\widetilde{D}}^{\top}X^{\top}{\vec{y}}]}{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}XX^{\top}{\vec{y}}]}$,
then $\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\eta
X^{\top}{\vec{y}}{\vec{y}}^{\top}X-X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}]=0$.

Overview of Proof of Lemma 3 and Lemma 4. We give an overview of how we prove
Lemma 3 and Lemma 4 here, and defer the full proofs to Appendix A. To show
that
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]$
is a scalar multiple of the identity, we first use the fact that even when all
of the $x_{i}$ are rotated by a rotation matrix $R$, the distribution of
${\vec{y}}|X$ remains the same, since the weight vector $w$ is drawn from
$\mathcal{N}(0,{I_{d\times d}})$ which is a rotationally invariant
distribution. Thus, if we define
$M(X)=\mathbb{E}[{\vec{y}}{\vec{y}}^{\top}\mid X]$ as a function of
$X\in\mathbb{R}^{n\times d}$, then for any rotation matrix
$R\in\mathbb{R}^{d\times d}$, we have

$\displaystyle M(XR^{\top})=\mathbb{E}[{\vec{y}}{\vec{y}}^{\top}\mid
XR^{\top}]=\mathbb{E}[{\vec{y}}{\vec{y}}^{\top}\mid X]=M(X)\,,$ (74)

where the second equality is because multiplying $X$ on the right by
$R^{\top}$ corresponds to rotating each of the $x_{i}$ by $R$. Additionally,
if we rotate the $x_{i}$ by $R$, then
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]$
remains the same — this is because the distribution of the $x_{i}$ is
unchanged due to the rotational invariance of the Gaussian distribution, and
the conditional distribution $y_{i}\mid x_{i}$ is unchanged when we rotate
$x_{i}$ by $R$. This implies that

$\displaystyle\mathbb{E}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]=\mathbb{E}[X^{\top}M(X)X]=\mathbb{E}[(XR^{\top})^{\top}M(XR^{\top})XR^{\top}]\,,$
(75)

where the second equality is because, as we observed above,
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]$
remains the same when we rotate each of the $x_{i}$ by $R$. Since
$M(XR^{\top})=M(X)$, we have

$\displaystyle\mathbb{E}[(XR^{\top})^{\top}M(XR^{\top})XR^{\top}]$
$\displaystyle=R\mathbb{E}[X^{\top}M(X)X]R^{\top}\,,$ (76)

which implies that
$\mathbb{E}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]=R\mathbb{E}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]R^{\top}$
for any rotation matrix $R$, and therefore
$\mathbb{E}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]$ is a scalar multiple of the
identity matrix. To finish the proof of Lemma 3, we show that
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}]$
is a scalar multiple of the identity using essentially the same argument. To
show Lemma 4, we simply take the trace of
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\eta
X^{\top}{\vec{y}}{\vec{y}}^{\top}X-X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}]$,
and select $\eta$ so that this trace is equal to $0$.

Finishing the Proof of Lemma 2. Recall that, to show that the gradients of
$J_{1}$ and $J_{2}$ (defined in Equation 50 and Equation 51) with respect to
$w$ are equal, it suffices to show Equation 73. However, this is a direct
consequence of Lemma 4. This is because we can rewrite Equation 73 as

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\left[\begin{array}[]{c}X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}A_{:,d+1}\\\
\text{tr}(\hat{w}_{\widetilde{D}}{\vec{y}}^{\top}XA_{:,1:d}^{\top})\end{array}\right]$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\left[\begin{array}[]{c}X^{\top}{\vec{y}}{\vec{y}}^{\top}XA_{:,d+1}\\\
\text{tr}(X^{\top}{\vec{y}}{\vec{y}}^{\top}XA_{:,1:d}^{\top})\end{array}\right]\,.$
(81)

This shows that the gradients of $J_{1}$ and $J_{2}$ with respect to $w$ are
equal, and we use similar arguments to show that the gradients of $J_{1}$ and
$J_{2}$ with respect to $A$ are equal. As mentioned above, this implies that
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|M_{:,1:d}^{\top}{G_{\widetilde{D}}}w-\hat{w}_{\widetilde{D}}\|_{2}^{2}-\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|M_{:,1:d}^{\top}{G_{\widetilde{D}}}w-\eta
X^{\top}{\vec{y}}\|_{2}^{2}$ is a constant that is independent of $M$ and $w$,
as desired.

##  4 Results for Different Data Covariance Matrices

In this section, we consider the setting where the $x_{i}$’s have a covariance
that is different from the identity matrix, and we show that the loss is
minimized when the one-layer transformer implements one step of gradient
descent with preconditioning. This suggests that the distribution of the
$x_{i}$’s has a significant effect on the algorithm that the transformer
implements.

Data Distribution. Concretely, the data distribution is the same as before,
but the $x_{i}$ are sampled from $\mathcal{N}(0,\Sigma)$, where
$\Sigma\in\mathbb{R}^{d\times d}$ is a positive semi-definite (PSD) matrix.
The outputs are generated according to $y_{i}=w^{\top}x_{i}+\epsilon_{i}$,
where $w\sim\mathcal{N}(0,\Sigma^{-1})$. This can equivalently be written as
$x_{i}=\Sigma^{1/2}u_{i}$, where $u_{i}\sim\mathcal{N}(0,{I_{d\times d}})$,
and $y_{i}=(w^{\prime})^{\top}u_{i}+\epsilon_{i}$, where
$w^{\prime}\sim\mathcal{N}(0,{I_{d\times d}})$. We keep all other definitions,
such as the loss function, the same as before.

######  Theorem 2 (Global Minimum for 1-Layer 1-Head Linear Self-Attention
with Skewed Covariance).

Suppose $(W_{K}^{*},W_{Q}^{*},W_{V}^{*},h^{*})$ is a global minimizer of the
loss $L$ when the data is generated according to the distribution given in
this section. Then, the corresponding one-layer transformer implements one
step of preconditioned gradient descent, on the least-squares linear
regression objective, with preconditioner $\Sigma^{-1}$, for some learning
rate $\eta>0$. Specifically, given a query token
$v_{n+1}=\left[\begin{array}[]{c}x_{n+1}\\\ 0\end{array}\right]$, the
transformer outputs $\eta\sum_{i=1}^{n}y_{i}(\Sigma^{-1}x_{i})^{\top}x_{n+1}$,
where
$\eta=\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}X(X^{\top}X+\sigma^{2}\Sigma)^{-1}X^{\top}{\vec{y}}]}{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}X\Sigma^{-1}X^{\top}{\vec{y}}]}$.

To prove this result, we essentially perform a change of variables to reduce
this problem to the setting of the previous section — then, we directly apply
Theorem 1. The detailed proof is given in Appendix B.

##  5 Results for Nonlinear Target Functions

In this section, we extend to a setting where the target function is nonlinear
— our conditions on the target function are mild, and for instance allow the
target function to be a fully-connected neural network with arbitrary
depth/width. However, we keep the model class the same (i.e. 1-layer
transformer with linear self-attention). We find that the transformer which
minimizes the pre-training loss still implements one step of GD on the linear
regression objective (Theorem 3), even though the target function is
nonlinear. This suggests that the distribution of $y_{i}|x_{i}$ does not
affect the algorithm learned by the transformer as much as the distribution of
$x_{i}$.

Data Distribution. In this section, we consider the same setup as Section 3,
but we change the distribution of the $y_{i}$’s. We now assume
$y_{i}=f(x_{i})+\epsilon_{i}$, where
$\epsilon_{i}\sim\mathcal{N}(0,\sigma^{2})$ as before, but $f$ is drawn from a
family of nonlinear functions satisfying the following assumption:

######  Assumption 1.

We assume that the target function $f$ is drawn from a family $\mathcal{F}$,
with a probability measure $\mathbb{P}$ on $\mathcal{F}$, such that the
following conditions hold: (1) for any fixed rotation matrix
$R\in\mathbb{R}^{d\times d}$, the distribution of functions $f$ is the same as
the distribution of $f\circ R$ (where $\circ$ denotes function composition).
Moreover, the distribution of $f$ is symmetric under negation. In other words,
if $E\subset\mathcal{F}$ is measurable under $\mathbb{P}$, then
$\mathbb{P}(E)=\mathbb{P}(-E)$, where $-E=\\{-f\mid f\in E\\}$.

For example, 1 is satisfied when $f(x)$ is a fully connected neural network,
with arbitrary depth and width, where the first and last layers have i.i.d.
$\mathcal{N}(0,1)$ entries. Under this assumption, we prove the following
result:

######  Theorem 3 (Global Minimum for $1$-Layer $1$-Head Linear Self-Attention
with Nonlinear Target Function).

Suppose 1 holds, and let $(W_{K}^{*},W_{Q}^{*},W_{V}^{*},h^{*})$ be a global
minimizer of the pre-training loss. Then, the corresponding one-layer
transformer implements one step of gradient descent on the least-squares
linear regression objective, given $(x_{1},y_{1},\ldots,x_{n},y_{n})$. More
concretely, given a query token $v_{n+1}=\left[\begin{array}[]{c}x_{n+1}\\\
0\end{array}\right]$, the transformer outputs
$\eta\sum_{i=1}^{n}y_{i}x_{i}^{\top}x_{n+1}$, where
$\eta=\frac{\mathbb{E}_{\mathcal{D}}[\overline{u}_{{\widetilde{D}}}^{\top}X^{\top}{\vec{y}}]}{\mathbb{E}_{\mathcal{D}}[{\vec{y}}^{\top}XX^{\top}{\vec{y}}]}$
and
$\overline{u}_{{\widetilde{D}}}=\textup{argmin}_{u}\mathbb{E}_{x_{n+1},y_{n+1}}[(u\cdot
x_{n+1}-y_{n+1})^{2}\mid{\widetilde{D}}]$.

The result is essentially the same as that of Theorem 1 — note that the
learning rate is potentially different, as it may depend on the function
family $\mathcal{F}$. The proof is analogous to the proof of Theorem 1. First
we prove the analogue of Lemma 1, defining $L(w,M)$ as in Section 2:

######  Lemma 5.

There exists a constant $C\geq 0$ such that

$\displaystyle L(w,M)$
$\displaystyle=C+\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|M_{:,1:d}^{\top}{G_{\widetilde{D}}}w-\overline{u}_{{\widetilde{D}}}\|_{2}^{2}\,.$
(82)

Here
$\overline{u}_{{\widetilde{D}}}=\textup{argmin}_{u}\mathbb{E}_{x_{n+1},y_{n+1}}[(u\cdot
x_{n+1}-y_{n+1})^{2}\mid{\widetilde{D}}]$.

Next, in the proof of Lemma 2, we used the fact that odd-degree polynomials of
the $y_{i}$ have expectation $0$ — the corresponding lemma in our current
setting is as follows:

######  Lemma 6.

For even integers $k$ and for $i\in[n]$,
$\mathbb{E}[y_{i}^{k}\overline{u}_{{\widetilde{D}}}\mid X]=0$. This also holds
with $y_{i}^{k}$ replaced by an even-degree monomial of the $y_{i}$.

Additionally, for odd integers $k$ and for $i\in[n]$,
$\mathbb{E}[y_{i}^{k}\mid X]=0$. This also holds with $y_{i}^{k}$ replaced by
an odd-degree monomial of the $y_{i}$.

###### Proof of Lemma 6.

This follows from 1. This is because for each outcome (i.e. choice of $f$ and
$\epsilon_{1},\ldots,\epsilon_{n}$) which leads to
$(x_{1},y_{1},\ldots,x_{n},y_{n})$, the corresponding outcome
$-f,-\epsilon_{1},\ldots,-\epsilon_{n}$ which leads to
$(x_{1},-y_{1},\ldots,x_{n},-y_{n})$ is equally likely. The
$\overline{u}_{{\widetilde{D}}}$ which is obtained from the second outcome is
the negative of the $\overline{u}_{{\widetilde{D}}}$ which is obtained from
the first outcome. If $k$ is even, then $y_{i}^{k}$ is the same under both
outcomes since $k$ is even, and the average of
$y_{i}^{k}\overline{u}_{{\widetilde{D}}}$ under these two outcomes is $0$.
Additionally, if $k$ is odd, then $y_{i}^{k}$ under the second outcome is the
negative of $y_{i}^{k}$ under the first outcome, and the average of
$y_{i}^{k}$ under these two outcomes is $0$. This completes the proof of the
lemma. ∎

Next, we show the analogue of Lemma 3.

######  Lemma 7.

$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]$
and
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\overline{u}_{{\widetilde{D}}}^{\top}]$
are scalar multiples of the identity. Thus,

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\overline{u}_{{\widetilde{D}}}^{\top}]=\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]$
(83)

where
$\eta=\frac{\mathbb{E}_{\mathcal{D}}[\overline{u}_{{\widetilde{D}}}^{\top}X^{\top}{\vec{y}}]}{\mathbb{E}_{\mathcal{D}}[{\vec{y}}^{\top}XX^{\top}{\vec{y}}]}$.

The proof of Lemma 7 is nearly identical to that of Lemma 3, with 1 used where
appropriate. For completeness, we include the proof in Appendix C. We now
state the analogue of Lemma 2:

######  Lemma 8.

There exists a constant $C_{1}\geq 0$ which is independent of $w,M$, such that

$\displaystyle L(w,M)$
$\displaystyle=C+\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|M_{:,1:d}^{\top}{G_{\widetilde{D}}}w-\eta
X^{\top}{\vec{y}}\|_{2}^{2}\,,$ (84)

where
$\eta=\frac{\mathbb{E}_{\mathcal{D}}[\overline{u}_{{\widetilde{D}}}^{\top}X^{\top}{\vec{y}}]}{\mathbb{E}_{\mathcal{D}}[{\vec{y}}^{\top}XX^{\top}{\vec{y}}]}$.

Theorem 3 now follows from Lemma 8 because
$M_{:,1:d}^{\top}{G_{\widetilde{D}}}w$ is the weight vector for the effective
linear predictor implemented by the transformer. All missing proofs are in
Appendix C.

##  6 Conclusion

We theoretically study one-layer transformers with linear self-attention
trained on noisy linear regression tasks with randomly generated data. We
confirm the empirical finding of von Oswald et al. (2022) by mathematically
showing that the global minimum of the pre-training loss for one-layer
transformers with linear self-attention corresponds to one step of GD on a
least-squares linear regression objective, when the covariates are from an
isotropic Gaussian distribution. We find that when the covariates are not from
an isotropic Gaussian distribution, the global minimum of the pre-training
loss instead corresponds to pre-conditioned GD, while if the covariates are
from an isotropic Gaussian distribution and the response is obtained from a
nonlinear target function, then the global minimum of the pre-training loss
will still correspond to one step of GD on a least-squares linear regression
objective. We study single-head linear self-attention layers — it is an
interesting direction for future work to study the global minima of the pre-
training loss for a multi-head linear self-attention layer. Another
interesting direction is to study the algorithms learned by multi-layer
transformers when the response is obtained from a nonlinear target function.
We note that Ahn et al. (2023) have studied the case of multi-layer
transformers when the target function is linear. They show that for certain
restricted parameterizations of multi-layer linear transformers, the global
minima or critical points of the pre-training loss correspond to interpretable
gradient-based algorithms.

## Acknowledgments

The authors would like to thank the support from NSF IIS 2211780 and a gift
from Open Philanthropy.

## References

  * Ahn et al. [2023] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra.  Transformers learn to implement preconditioned gradient descent for in-context learning, 2023.
  * Akyürek et al. [2023] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou.  What learning algorithm is in-context learning? investigations with linear models.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=0g0X4H8yN4I.
  * Brown et al. [2020] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.  Language models are few-shot learners.  _CoRR_ , abs/2005.14165, 2020.  URL https://arxiv.org/abs/2005.14165.
  * Dai et al. [2023] Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei.  Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers, 2023.
  * Garg et al. [2022] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant.  What can transformers learn in-context? A case study of simple function classes.  In _NeurIPS_ , 2022.  URL http://papers.nips.cc/paper_files/paper/2022/hash/c529dba08a146ea8d6cf715ae8930cbe-Abstract-Conference.html.
  * Giannou et al. [2023] Angeliki Giannou, Shashank Rajput, Jy yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos.  Looped transformers as programmable computers, 2023.
  * Lieber et al. [2021] Opher Lieber, Or Sharir, Barak Lenz, and Yoav Shoham.  Jurassic-1: Technical details and evaluation.  Technical report, AI21 Labs, 2021.
  * Liu et al. [2023] Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang.  Transformers learn shortcuts to automata.  In _The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023_. OpenReview.net, 2023\.  URL https://openreview.net/pdf?id=De4FYqjFueZ.
  * Min et al. [2022] Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer.  Rethinking the role of demonstrations: What makes in-context learning work?  In _EMNLP_ , 2022.
  * Radford et al. [2019] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever.  Language models are unsupervised multitask learners.  Technical report, OpenAI, 2019.
  * von Oswald et al. [2022] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov.  Transformers learn in-context by gradient descent, 2022.
  * Wang and Komatsuzaki [2021] Ben Wang and Aran Komatsuzaki.  Gpt-j-6b: A 6 billion parameter autoregressive language model, 2021.
  * Wei et al. [2023] Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, and Tengyu Ma.  Larger language models do in-context learning differently, 2023.
  * Xie et al. [2022] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma.  An explanation of in-context learning as implicit bayesian inference.  In _The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022_. OpenReview.net, 2022.  URL https://openreview.net/forum?id=RdJVFCHjUMI.
  * Zhang et al. [2023] Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett.  Trained transformers learn linear models in-context, 2023.

##  Appendix A Missing Proofs from Section 3

###### Proof of Lemma 3.

For convenience let
$M(X)=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}{\vec{y}}^{\top}\mid
X]$ — we use this notation to make the dependence on $X$ clear. Then, the
$(i,j)$-th entry of $M(X)$ is $\mathbb{E}[(w\cdot x_{i}+\epsilon_{i})(w\cdot
x_{j}+\epsilon_{j})\mid X]$ and this is equal to $\mathbb{E}[(w\cdot
x_{i})(w\cdot x_{j})\mid X]$ for $i\neq j$, and $\mathbb{E}[(w\cdot
x_{i})(w\cdot x_{j})\mid X]+\sigma^{2}$ for $i=j$. If we perform the change of
variables $x_{i}\mapsto Rx_{i}$ for a fixed rotation matrix $R$ and all
$i\in[n]$, then $\mathbb{E}_{w}[(w\cdot Rx_{i})(w\cdot Rx_{j})\mid
X]=\mathbb{E}_{w}[(w\cdot x_{i})(w\cdot x_{j})\mid X]$ because
$w\sim\mathcal{N}(0,{I_{d\times d}})$ and $\mathcal{N}(0,{I_{d\times d}})$ is
a rotationally invariant distribution. In other words, we have
$M(XR^{\top})=M(X)$. Thus, for any rotation matrix $R$,

$\displaystyle\mathbb{E}_{X}[X^{\top}M(X)X]$
$\displaystyle=\mathbb{E}_{X}[(XR)^{\top}M(XR^{\top})(XR^{\top})]$ (By
rotational invariance of dist. of $x_{i}$)
$\displaystyle=R^{\top}\mathbb{E}_{X}[X^{\top}M(XR^{\top})X]R^{\top}$ (85)
$\displaystyle=R^{\top}\mathbb{E}_{X}[X^{\top}M(X)X]R\,.$ (Because
$M(XR^{\top})=M(X)$)

This implies that
$\mathbb{E}_{X}[X^{\top}M(X)X]=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]$
is a scalar multiple of the identity matrix.

Next, we consider
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}\mid
X]$, which we write for convenience as $J(X)$. Similarly to the above proof,
we use the observation that if we make the change of variables $x_{i}\to
Rx_{i}$, then the joint distribution of $y_{1},\ldots,y_{n},y_{n+1}$ is
unchanged due to the rotational invariance of $w$. Additionally, if we write
$\hat{w}_{\widetilde{D}}$ as
$\hat{w}_{\widetilde{D}}(x_{1},y_{1},\ldots,x_{n},y_{n})$ to emphasize the
fact that it is a function of $x_{1},y_{1},\ldots,x_{n},y_{n}$, then

$\displaystyle\hat{w}_{\widetilde{D}}(Rx_{1},y_{1},Rx_{2},y_{2},\ldots,Rx_{n},y_{n})$
$\displaystyle=R\hat{w}_{\widetilde{D}}(x_{1},y_{1},x_{2},y_{2},\ldots,x_{n},y_{n})$
(86)

because $\hat{w}_{\widetilde{D}}$ is the minimizer of
$F(w)=\|w^{\top}X-{\vec{y}}\|_{2}^{2}+\sigma^{2}\|w\|_{2}^{2}$, and if all the
$x_{i}$ are rotated by $R$, then the minimizer of $F$ will also be rotated by
$R$. Thus,

$\displaystyle
J(XR^{\top})=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}\mid
XR^{\top}]=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}R^{\top}\mid
X]=J(X)R^{\top}\,.$ (87)

We can use this to show that
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}]$
is a scalar multiple of the identity. Letting $R$ be a rotation matrix, we
obtain

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}]$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\mathbb{E}[X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}\mid
X]$ (88)
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}\mathbb{E}[{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}\mid
X]]$ (89) $\displaystyle=\mathbb{E}_{X}[X^{\top}J(X)]$ (90)
$\displaystyle=\mathbb{E}_{X}[(XR^{\top})^{\top}J(XR^{\top})]$ (By rotational
invariance of dist. of $x_{i}$)
$\displaystyle=\mathbb{E}_{X}[RX^{\top}J(X)R^{\top}]$ (B.c.
$J(XR^{\top})=J(X)R^{\top}$)
$\displaystyle=R\mathbb{E}_{X}[X^{\top}J(X)]R^{\top}$ (91)
$\displaystyle=R\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}]R^{\top}\,.$
(By Equation 90)

This implies that
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}]$
is a scalar multiple of the identity. ∎

###### Proof of Lemma 4.

By Lemma 3, we have
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]=c_{1}I$
and
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}]=c_{2}I$
for some scalars $c_{1},c_{2}$. Taking the traces of both matrices gives us

$\displaystyle
c_{1}d=\text{tr}(\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X])=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\text{tr}(X^{\top}{\vec{y}}{\vec{y}}^{\top}X)]=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}XX^{\top}{\vec{y}}]\,,$
(92)

and similarly,

$\displaystyle
c_{2}d=\text{tr}(\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}])=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\text{tr}(X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top})]=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\hat{w}_{\widetilde{D}}^{\top}X^{\top}{\vec{y}}]\,.$
(93)

By the definition of $\eta$ as
$\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\hat{w}_{\widetilde{D}}^{\top}X^{\top}{\vec{y}}]}{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}XX^{\top}{\vec{y}}]}$,
we have

$\displaystyle\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]-\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}]$
$\displaystyle=\eta c_{1}I-c_{2}I$ (94)
$\displaystyle=\eta\cdot\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}XX^{\top}{\vec{y}}]}{d}I-\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\hat{w}_{\widetilde{D}}^{\top}X^{\top}{\vec{y}}]}{d}I$
(95)
$\displaystyle=\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\hat{w}_{\widetilde{D}}^{\top}X^{\top}{\vec{y}}]}{d}I-\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\hat{w}_{\widetilde{D}}^{\top}X^{\top}{\vec{y}}]}{d}I$
(96) $\displaystyle=0\,,$ (97)

completing the proof of the lemma. ∎

###### Proof of Lemma 2.

For convenience, we write $A:=M_{:,1:d}^{\top}$. Then, we wish to show that
for all $A\in\mathbb{R}^{d\times(d+1)}$ and $w\in\mathbb{R}^{d+1}$, we have

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|A{G_{\widetilde{D}}}w-\hat{w}_{\widetilde{D}}\|_{2}^{2}=C+\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|A{G_{\widetilde{D}}}w-\eta
X^{\top}{\vec{y}}\|_{2}^{2}$ (98)

for some constant $C$ independent of $A$ and $w$. Define

$\displaystyle J_{1}(A,w)$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|A{G_{\widetilde{D}}}w-\hat{w}_{\widetilde{D}}\|_{2}^{2}$
(99)

and

$\displaystyle J_{2}(A,w)$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\|A{G_{\widetilde{D}}}w-\eta
X^{\top}{\vec{y}}\|_{2}^{2}\,.$ (100)

To show that $J_{1}(A,w)$ and $J_{2}(A,w)$ are equal up to a constant, it
suffices to show that their gradients are identical, since $J_{1}$ and $J_{2}$
are defined everywhere on
$\mathbb{R}^{d+1}\times\mathbb{R}^{(d+1)\times(d+1)}$ and by the fundamental
theorem of calculus.

Gradients With Respect to $w$. First, we analyze the gradients with respect to
$w$. We have

$\displaystyle\nabla_{w}J_{1}(A,w)$
$\displaystyle=2\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}(A{G_{\widetilde{D}}}w-\hat{w}_{\widetilde{D}})$
(101)

and

$\displaystyle\nabla_{w}J_{2}(A,w)$
$\displaystyle=2\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}(A{G_{\widetilde{D}}}w-\eta
X^{\top}{\vec{y}})$ (102)

where we use the convention that the gradient with respect to $w$ has the same
shape as $w$. To show that $\nabla_{w}J_{1}(A,w)=\nabla_{w}J_{2}(A,w)$, it
suffices to show that
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}\hat{w}_{\widetilde{D}}=\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}X^{\top}{\vec{y}}$.
Observe that we can write ${G_{\widetilde{D}}}$ as

$\displaystyle{G_{\widetilde{D}}}=\sum_{i=1}^{n}\left[\begin{array}[]{c}x_{i}\\\
y_{i}\end{array}\right]\left[\begin{array}[]{c}x_{i}\\\
y_{i}\end{array}\right]^{\top}=\sum_{i=1}^{n}\left[\begin{array}[]{cc}x_{i}x_{i}^{\top}&y_{i}x_{i}\\\
y_{i}x_{i}^{\top}&y_{i}^{2}\end{array}\right]\,.$ (109)

Given $X$, the expected value of any odd monomial of the $y_{i}$ is equal to
$0$, since $w$ and $\epsilon_{i}$ are independent of $X$ and have mean $0$.
Thus, we can ignore the blocks corresponding to $x_{i}x_{i}^{\top}$ and
$y_{i}^{2}$, since the entries of $\hat{w}_{\widetilde{D}}$ and
$X^{\top}{\vec{y}}$ contain odd monomials of the $y_{i}$ once $X$ is fixed.

We now proceed in two steps. First, we deal with the terms corresponding to
the lower left block of ${G_{\widetilde{D}}}$ and show that their respective
contributions to
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}\hat{w}_{\widetilde{D}}$
and
$\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}X^{\top}{\vec{y}}$
are equal. For
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}\hat{w}_{\widetilde{D}}$,
these terms contribute:

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\sum_{i=1}^{n}y_{i}x_{i}^{\top}(A^{\top})_{1:d,:}\hat{w}_{\widetilde{D}}$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{\vec{y}}^{\top}XA_{:,1:d}^{\top}\hat{w}_{\widetilde{D}}$
(110)
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\text{tr}(\hat{w}_{\widetilde{D}}{\vec{y}}^{\top}XA_{:,1:d}^{\top})$
(111)
$\displaystyle=\text{tr}(\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\hat{w}_{\widetilde{D}}{\vec{y}}^{\top}X]A_{:,1:d}^{\top})\,,$
(112)

while for
$\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}X^{\top}{\vec{y}}$,
these terms contribute:

$\displaystyle\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\sum_{i=1}^{n}y_{i}x_{i}^{\top}A_{:,1:d}^{\top}X^{\top}{\vec{y}}$
$\displaystyle=\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{\vec{y}}^{\top}XA_{:,1:d}^{\top}X^{\top}{\vec{y}}$
(113)
$\displaystyle=\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\text{tr}(X^{\top}{\vec{y}}{\vec{y}}^{\top}XA_{:,1:d}^{\top})$
(114)
$\displaystyle=\text{tr}(\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\eta
X^{\top}{\vec{y}}{\vec{y}}^{\top}X]A_{:,1:d}^{\top})\,,$ (115)

and since
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\hat{w}_{\widetilde{D}}{\vec{y}}^{\top}X]=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\eta
X^{\top}{\vec{y}}{\vec{y}}^{\top}X]$ by Lemma 4, these two contributions are
equal.

Second, we deal with the terms corresponding to the upper right block of
${G_{\widetilde{D}}}$. For
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}\hat{w}_{\widetilde{D}}$
these terms contribute

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\sum_{i=1}^{n}y_{i}x_{i}(A^{\top})_{(d+1),:}\hat{w}_{\widetilde{D}}$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}X^{\top}{\vec{y}}A_{:,(d+1)}^{\top}\hat{w}_{\widetilde{D}}$
(116)
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}X^{\top}{\vec{y}}\hat{w}_{\widetilde{D}}^{\top}A_{:,(d+1)}\,,$
(B.c. dot product $A_{:,(d+1)}^{\top}\hat{w}_{\widetilde{D}}$ is commutative)

while for
$\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}X^{\top}{\vec{y}}$,
these terms contribute

$\displaystyle\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\sum_{i=1}^{n}y_{i}x_{i}(A^{\top})_{(d+1),:}X^{\top}{\vec{y}}$
$\displaystyle=\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}X^{\top}{\vec{y}}A_{:,(d+1)}^{\top}X^{\top}{\vec{y}}$
(117)
$\displaystyle=\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}X^{\top}{\vec{y}}{\vec{y}}^{\top}XA_{:,(d+1)}\,,$
(B.c. dot product $A_{:,(d+1)}^{\top}X^{\top}{\vec{y}}$ is commutative)

and again these contributions are equal by Lemma 4. In summary, we have shown
that for
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}\hat{w}_{\widetilde{D}}$
and
$\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}X^{\top}{\vec{y}}$,
the contributions corresponding to the lower left and upper right blocks of
${G_{\widetilde{D}}}$ are equal, meaning that
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}\hat{w}_{\widetilde{D}}$
and
$\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}{G_{\widetilde{D}}}A^{\top}X^{\top}{\vec{y}}$
are equal. This shows that $\nabla_{w}J_{1}(A,w)$ and $\nabla_{w}J_{2}(A,w)$
are equal.

Gradients With Respect to $A$. We can compute the gradient of $J_{1}$ with
respect to $A$ as follows:

$\displaystyle\nabla_{A}J_{1}(A,w)$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\nabla_{A}\|A{G_{\widetilde{D}}}w-\hat{w}_{\widetilde{D}}\|_{2}^{2}$
(118)
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\nabla_{A}(A{G_{\widetilde{D}}}w-\hat{w}_{\widetilde{D}})^{\top}(A{G_{\widetilde{D}}}w-\hat{w}_{\widetilde{D}})$
(119)
$\displaystyle=2\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}(A{G_{\widetilde{D}}}w-\hat{w}_{\widetilde{D}})({G_{\widetilde{D}}}w)^{\top}\,,$
(120)

and we can similarly compute

$\displaystyle\nabla_{A}J_{2}(A,w)$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\nabla_{A}\|A{G_{\widetilde{D}}}w-\eta
X^{\top}{\vec{y}}\|_{2}^{2}$ (121)
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\nabla_{A}(A{G_{\widetilde{D}}}w-\eta
X^{\top}{\vec{y}})^{\top}(A{G_{\widetilde{D}}}w-\eta X^{\top}{\vec{y}})$ (122)
$\displaystyle=2\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}(A{G_{\widetilde{D}}}w-\eta
X^{\top}{\vec{y}})({G_{\widetilde{D}}}w)^{\top}\,.$ (123)

Thus, to show that the gradients with respect to $A$ are equal, it suffices to
show that
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\hat{w}_{\widetilde{D}}w^{\top}{G_{\widetilde{D}}}=\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}X^{\top}{\vec{y}}w^{\top}{G_{\widetilde{D}}}$
for all $w$.

As before, we only need to consider the lower left and upper right blocks of
${G_{\widetilde{D}}}$, since the entries of the other blocks have even powers
of the $y_{i}$. First let us consider the lower left block. The contribution
of the lower left block to
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\hat{w}_{\widetilde{D}}w^{\top}{G_{\widetilde{D}}}$
is

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\hat{w}_{\widetilde{D}}w_{(d+1)}\sum_{i=1}^{n}y_{i}x_{i}^{\top}=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\hat{w}_{\widetilde{D}}w_{(d+1)}{\vec{y}}^{\top}X=w_{(d+1)}\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\hat{w}_{\widetilde{D}}{\vec{y}}^{\top}X\,,$
(124)

while the contribution of the lower left block to
$\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}X^{\top}{\vec{y}}w^{\top}{G_{\widetilde{D}}}$
is

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\eta
X^{\top}{\vec{y}}w_{(d+1)}\sum_{i=1}^{n}y_{i}x_{i}^{\top}$ $\displaystyle=\eta
w_{(d+1)}\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}X^{\top}{\vec{y}}{\vec{y}}^{\top}X\,,$
(125)

and these two are equal by Lemma 4. The contribution of the upper right block
to
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\hat{w}_{\widetilde{D}}w^{\top}{G_{\widetilde{D}}}$
is

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\hat{w}_{\widetilde{D}}w_{1:d}^{\top}\sum_{i=1}^{n}y_{i}x_{i}=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\hat{w}_{\widetilde{D}}w_{1:d}^{\top}X^{\top}{\vec{y}}=(\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\hat{w}_{\widetilde{D}}{\vec{y}}^{\top}X)w_{1:d}\,,$
(126)

by the commutativity of the dot product $w_{1:d}^{\top}X^{\top}{\vec{y}}$,
while similarly, the contribution of the upper right block to
$\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}X^{\top}{\vec{y}}w^{\top}{G_{\widetilde{D}}}$
is

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\eta
X^{\top}{\vec{y}}w_{1:d}^{\top}\sum_{i=1}^{n}y_{i}x_{i}=(\eta\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}X^{\top}{\vec{y}}{\vec{y}}^{\top}X)w_{1:d}\,,$
(127)

and these two quantities are equal by Lemma 4. Thus, we have shown that
$\nabla_{A}J_{1}(A,w)=\nabla_{A}J_{2}(A,w)$.

Summary. We have shown that $\nabla_{w}J_{1}(A,w)=\nabla_{w}J_{2}(A,w)$, and
$\nabla_{A}J_{1}(A,w)=\nabla_{A}J_{2}(A,w)$. This completes the proof of the
lemma. ∎

###### Proof of Theorem 1.

This theorem follows from Lemma 1 and Lemma 2. To see why, recall from Section
2 that the output of the transformer on the last token $v_{n+1}$ can be
written as $w^{\top}{G_{\widetilde{D}}}Mv_{n+1}$, where $M=W_{K}^{\top}W_{Q}$
and $w=W_{V}^{\top}h$. By Lemma 2, the pre-training loss is minimized when
$M_{:,1:d}^{\top}{G_{\widetilde{D}}}w=\eta X^{\top}{\vec{y}}$ almost surely
(over the randomness of $X$ and ${\vec{y}}$), and when this holds, the output
of the transformer on the last token is

$\displaystyle
w^{\top}{G_{\widetilde{D}}}Mv_{n+1}=w^{\top}{G_{\widetilde{D}}}M_{:,1:d}x_{n+1}=(\eta
X^{\top}{\vec{y}})^{\top}x_{n+1}=\eta\sum_{i=1}^{n}y_{i}x_{i}^{\top}x_{n+1}\,,$
(128)

as desired. ∎

##  Appendix B Missing Proofs for Section 4

###### Proof of Theorem 2.

As discussed in Section 4, we can write $x_{i}=\Sigma^{1/2}u_{i}$, where
$u_{i}\sim\mathcal{N}(0,{I_{d\times d}})$. Furthermore, we can let
$U\in\mathbb{R}^{n\times d}$ be the matrix whose $i^{\textup{th}}$ row is
$u_{i}^{\top}$. Note that $X=U\Sigma^{1/2}$, since $\Sigma$ is a symmetric
positive-definite matrix. Given
${\widetilde{D}}=(x_{1},y_{1},\ldots,x_{n},y_{n})$, we can define
$\hat{w}_{{\widetilde{D}},\Sigma}=(U^{\top}U+\sigma^{2}I)^{-1}U^{\top}{\vec{y}}$
— this would be the solution to ridge regression if we were given the $u_{i}$
and $y_{i}$. Using $X=U\Sigma^{1/2}$ which implies $U=X\Sigma^{-1/2}$, we
obtain

$\displaystyle\hat{w}_{{\widetilde{D}},\Sigma}$
$\displaystyle=(\Sigma^{-1/2}X^{\top}X\Sigma^{-1/2}+\sigma^{2}I)^{-1}\Sigma^{-1/2}X^{\top}{\vec{y}}$
(129)
$\displaystyle=(\Sigma^{1/2}\Sigma^{-1/2}X^{\top}X\Sigma^{-1/2}+\sigma^{2}\Sigma^{1/2})^{-1}X^{\top}{\vec{y}}$
(130)
$\displaystyle=(X^{\top}X\Sigma^{-1/2}+\sigma^{2}\Sigma^{1/2})^{-1}X^{\top}{\vec{y}}$
(131)
$\displaystyle=\Sigma^{1/2}(X^{\top}X+\sigma^{2}\Sigma)^{-1}X^{\top}{\vec{y}}\,.$
(132)

We now change variables in order to reduce Theorem 2 to Theorem 1, using the
fact that $u_{i}$ and $y_{i}$ together are from the same distribution as the
data studied in Theorem 1. We can write the loss as

$\displaystyle L(W_{K},W_{Q},W_{V},h)$
$\displaystyle=\mathbb{E}_{D\sim{\mathcal{T}}}[(y_{n+1}-h^{\top}{\hat{v}}_{n+1})^{2}]$
(133)
$\displaystyle=\mathbb{E}_{{\widetilde{D}},x_{n+1}}\Big{[}\mathbb{E}_{y_{n+1}}[(y_{n+1}-h^{\top}{\hat{v}}_{n+1})^{2}\mid
x_{n+1},\mathcal{D}]\Big{]}\,.$ (134)

The solution to ridge regression given the $u_{i}$ and $y_{i}$ is
$\hat{w}_{{\widetilde{D}},\Sigma}$, meaning $\mathbb{E}[y_{n+1}\mid
u_{n+1},{\widetilde{D}}]=\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}$ and
therefore $\mathbb{E}[y_{n+1}\mid
x_{n+1},{\widetilde{D}}]=\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}$,
since $u_{n+1}=\Sigma^{-1/2}x_{n+1}$, which is an invertible function of
$x_{n+1}$. Thus,

$\displaystyle\mathbb{E}_{y_{n+1}}[(y_{n+1}-h^{\top}{\hat{v}}_{n+1})^{2}\mid
x_{n+1},{\widetilde{D}}]$
$\displaystyle=\mathbb{E}_{y_{n+1}}[(y_{n+1}-\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1})^{2}\mid
x_{n+1},{\widetilde{D}}]$ (135)
$\displaystyle\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,+\mathbb{E}_{y_{n+1}}[(\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}-h^{\top}{\hat{v}}_{n+1})^{2}\mid
x_{n+1},{\widetilde{D}}]$ (136)
$\displaystyle\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,+2\mathbb{E}_{y_{n+1}}[(y_{n+1}-\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1})(\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}-h^{\top}v_{n+1})\mid
x_{n+1},{\widetilde{D}}]\,.$ (137)

Now, note that
$\mathbb{E}_{y_{n+1}}[(y_{n+1}-\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1})(\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}-h^{\top}v_{n+1})\mid
x_{n+1},{\widetilde{D}}]=0$ since
$\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}-h^{\top}{\hat{v}}_{n+1}$ is
fully determined by $x_{n+1}$ and ${\widetilde{D}}$, and
$\mathbb{E}[y_{n+1}-\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}\mid
x_{n+1},{\widetilde{D}}]=0$ as mentioned above. Additionally, we can write
$\mathbb{E}_{y_{n+1}}[(y_{n+1}-\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1})^{2}\mid
x_{n+1},{\widetilde{D}}]$ as $C_{x_{n+1},{\widetilde{D}}}$ (denoting a
constant that depends only on $x_{n+1}$ and ${\widetilde{D}}$) since it is
independent of the parameters $W_{K},W_{Q},W_{V},h$. Thus, taking the
expectation over $x_{n+1}$ and ${\widetilde{D}}$, for some constant $C$ which
is independent of $W_{K},W_{Q},W_{V},h$, we have

$\displaystyle L(W_{K},W_{Q},W_{V},h)$
$\displaystyle=C+\mathbb{E}_{{\widetilde{D}},x_{n+1}}(\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}-h^{\top}{\hat{v}}_{n+1})^{2}$
(138)
$\displaystyle=C+\mathbb{E}_{{\widetilde{D}},x_{n+1}}(\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}-h^{\top}W_{V}{G_{\widetilde{D}}}W_{K}^{\top}W_{Q}v_{n+1})^{2}$
(139)
$\displaystyle=C+\mathbb{E}_{{\widetilde{D}},x_{n+1}}(\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}-h^{\top}W_{V}{G_{\widetilde{D}}}W_{K}^{\top}{W_{Q}}_{:,1:d}x_{n+1})^{2}$
(140)
$\displaystyle=C+\mathbb{E}_{{\widetilde{D}},x_{n+1}}(\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}-h^{\top}W_{V}{G_{\widetilde{D}}}W_{K}^{\top}{W_{Q}}_{:,1:d}\Sigma^{1/2}u_{n+1})^{2}\,.$
(141)

To finish the proof, observe that we can write

$\displaystyle{G_{\widetilde{D}}}$
$\displaystyle=\sum_{i=1}^{n}\left[\begin{array}[]{c}x_{i}\\\
y_{i}\end{array}\right]\left[\begin{array}[]{c}x_{i}\\\
y_{i}\end{array}\right]^{\top}$ (146)
$\displaystyle=\begin{pmatrix}\Sigma^{1/2}&0\\\
0&1\end{pmatrix}\sum_{i=1}^{n}\left[\begin{array}[]{c}u_{i}\\\
y_{i}\end{array}\right]\left[\begin{array}[]{c}u_{i}\\\
y_{i}\end{array}\right]^{\top}\begin{pmatrix}\Sigma^{1/2}&0\\\
0&1\end{pmatrix}$ (151)
$\displaystyle=H{G_{\widetilde{D}}}^{\prime}H^{\top}\,,$ (152)

where we have defined $H=\begin{pmatrix}\Sigma^{1/2}&0\\\ 0&1\end{pmatrix}$
and

$\displaystyle{G_{\widetilde{D}}}^{\prime}=\sum_{i=1}^{n}\left[\begin{array}[]{c}u_{i}\\\
y_{i}\end{array}\right]\left[\begin{array}[]{c}u_{i}\\\
y_{i}\end{array}\right]^{\top}\,.$ (157)

Now, let us make the change of variables $h=h^{\prime}$,
$W_{V}=W_{V}^{\prime}H^{-1}$, $W_{K}=W_{K}^{\prime}H^{-1}$ and
$W_{Q}=W_{Q}^{\prime}H^{-1}$. Then, we obtain

$\displaystyle L(W_{K},$ $\displaystyle W_{Q},W_{V},h)$ (158)
$\displaystyle=C+\mathbb{E}_{{\widetilde{D}},x_{n+1}}(\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}-h^{\top}W_{V}{G_{\widetilde{D}}}W_{K}^{\top}{W_{Q}}_{:,1:d}\Sigma^{1/2}u_{n+1})^{2}$
(159)
$\displaystyle=C+\mathbb{E}_{{\widetilde{D}},x_{n+1}}(\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}-(h^{\prime})^{\top}W_{V}^{\prime}H^{-1}H{G_{\widetilde{D}}}^{\prime}HH^{-1}(W_{K}^{\prime})^{\top}{W_{Q}^{\prime}}_{:,1:d}\Sigma^{-1/2}\Sigma^{1/2}u_{n+1})^{2}$
(160)
$\displaystyle=C+\mathbb{E}_{{\widetilde{D}},x_{n+1}}(\hat{w}_{{\widetilde{D}},\Sigma}^{\top}u_{n+1}-(h^{\prime})^{\top}W_{V}^{\prime}{G_{\widetilde{D}}}^{\prime}(W_{K}^{\prime})^{\top}{W_{Q}^{\prime}}_{:,1:d}u_{n+1})^{2}\,.$
(161)

In other words, this is equal to the loss obtained by the transformer
corresponding to $h^{\prime},W_{V}^{\prime},W_{K}^{\prime},W_{Q}^{\prime}$
when the data is from the distribution studied in Section 3. The above proof
also shows that the transformer corresponding to $h,W_{V},W_{K},W_{Q}$ has the
same output on $x_{n+1}$, as the transformer corresponding to
$h^{\prime},W_{V}^{\prime},W_{K}^{\prime},W_{Q}^{\prime}$ does on $u_{n+1}$.
By Theorem 1, the output of the latter transformer on $u_{n+1}$ is

$\displaystyle(\eta U^{\top}{\vec{y}})^{\top}u_{n+1}$
$\displaystyle=\eta{\vec{y}}^{\top}Uu_{n+1}$ (162)
$\displaystyle=\eta\sum_{i=1}^{n}y_{i}u_{i}^{\top}u_{n+1}$ (163)
$\displaystyle=\eta\sum_{i=1}^{n}y_{i}(\Sigma^{-1/2}x_{i})^{\top}\Sigma^{-1/2}x_{n+1}$
(164) $\displaystyle=\eta\sum_{i=1}^{n}y_{i}x_{i}^{\top}\Sigma^{-1}x_{n+1}\,,$
(165)

where the learning rate is

$\displaystyle\eta$
$\displaystyle=\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\hat{w}_{{\widetilde{D}},\Sigma}^{\top}U^{\top}{\vec{y}}]}{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}UU^{\top}{\vec{y}}]}$
(166)
$\displaystyle=\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}X(X^{\top}X+\sigma^{2}\Sigma)^{-1}\Sigma^{1/2}U^{\top}{\vec{y}}]}{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}UU^{\top}{\vec{y}}]}$
(By definition of $\hat{w}_{{\widetilde{D}},\Sigma}$)
$\displaystyle=\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}X(X^{\top}X+\sigma^{2}\Sigma)^{-1}\Sigma^{1/2}\Sigma^{-1/2}X^{\top}{\vec{y}}]}{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}X\Sigma^{-1/2}\Sigma^{-1/2}X^{\top}{\vec{y}}]}$
(B.c. $U=X\Sigma^{-1/2}$)
$\displaystyle=\frac{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}X(X^{\top}X+\sigma^{2}\Sigma)^{-1}X^{\top}{\vec{y}}]}{\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}^{\top}X\Sigma^{-1}X^{\top}{\vec{y}}]}\,.$
(167)

This completes the proof. ∎

##  Appendix C Missing Proofs from Section 5

###### Proof of Lemma 5.

We proceed similarly to the proof of Lemma 1. For convenience, fix
${\widetilde{D}}$ and consider the function

$\displaystyle g(u)$ $\displaystyle=\mathbb{E}_{x_{n+1},y_{n+1}}[(u\cdot
x_{n+1}-y_{n+1})^{2}\mid{\widetilde{D}}]\,.$ (168)

Then, we have

$\displaystyle\nabla_{u}g(u)$
$\displaystyle=\mathbb{E}_{x_{n+1},y_{n+1}}[2(u\cdot
x_{n+1}-y_{n+1})x_{n+1}\mid{\widetilde{D}}]\,.$ (169)

Therefore, if $\overline{u}_{{\widetilde{D}}}$ is the minimizer of $g(u)$
(note that $\overline{u}_{{\widetilde{D}}}$ depends on ${\widetilde{D}}$ but
not on $x_{n+1}$ and $y_{n+1}$), then for any $u\in\mathbb{R}^{d}$, we have

$\displaystyle\mathbb{E}_{x_{n+1},y_{n+1}}[(\overline{u}_{{\widetilde{D}}}\cdot
x_{n+1}-y_{n+1})(u\cdot x_{n+1}-\overline{u}_{{\widetilde{D}}}\cdot
x_{n+1})\mid{\widetilde{D}}]=(u-\overline{u}_{{\widetilde{D}}})\cdot\nabla_{u}g(\overline{u}_{{\widetilde{D}}})=0\,,$
(170)

meaning that for any $u\in\mathbb{R}^{d}$, and some $C_{{\widetilde{D}}}$
which only depends on ${\widetilde{D}}$ and is independent of $u$,

$\displaystyle\mathbb{E}_{x_{n+1},y_{n+1}}$ $\displaystyle[(u\cdot
x_{n+1}-y_{n+1})^{2}\mid{\widetilde{D}}]$ (171)
$\displaystyle=\mathbb{E}_{x_{n+1},y_{n+1}}[(u\cdot
x_{n+1}-\overline{u}_{{\widetilde{D}}}\cdot
x_{n+1}+\overline{u}_{{\widetilde{D}}}\cdot
x_{n+1}-y_{n+1})^{2}\mid{\widetilde{D}}]$ (172)
$\displaystyle=\mathbb{E}_{x_{n+1},y_{n+1}}[(u\cdot
x_{n+1}-\overline{u}_{{\widetilde{D}}}\cdot
x_{n+1})^{2}\mid{\widetilde{D}}]+\mathbb{E}_{x_{n+1},y_{n+1}}[(\overline{u}_{{\widetilde{D}}}\cdot
x_{n+1}-y_{n+1})^{2}\mid{\widetilde{D}}]$ (173)
$\displaystyle\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,+2\mathbb{E}_{x_{n+1},y_{n+1}}[(u\cdot
x_{n+1}-\overline{u}_{{\widetilde{D}}}\cdot
x_{n+1})(\overline{u}_{{\widetilde{D}}}\cdot
x_{n+1}-y_{n+1})\mid{\widetilde{D}}]$ (174)
$\displaystyle=\mathbb{E}_{x_{n+1},y_{n+1}}[(u\cdot
x_{n+1}-\overline{u}_{{\widetilde{D}}}\cdot
x_{n+1})^{2}\mid{\widetilde{D}}]+\mathbb{E}_{x_{n+1},y_{n+1}}[(\overline{u}_{{\widetilde{D}}}\cdot
x_{n+1}-y_{n+1})^{2}\mid{\widetilde{D}}]$ (By Equation 170)
$\displaystyle=\|u-\overline{u}_{{\widetilde{D}}}\|_{2}^{2}+C_{{\widetilde{D}}}\,.$
(175)

Here the last equality is because $x_{n+1}\sim\mathcal{N}(0,{I_{d\times d}})$,
and because $\mathbb{E}_{x_{n+1},y_{n+1}}[(\overline{u}_{{\widetilde{D}}}\cdot
x_{n+1}-y_{n+1})^{2}\mid{\widetilde{D}}]$ is a constant that depends on
${\widetilde{D}}$ but not on $u$. We can apply this manipulation to the loss
function:

$\displaystyle L(w,M)$
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\Big{[}\mathbb{E}_{x_{n+1},y_{n+1}}[(y_{n+1}-\hat{y}_{n+1})^{2}\mid{\widetilde{D}}]\Big{]}$
(176)
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\Big{[}\mathbb{E}_{x_{n+1},y_{n+1}}[(y_{n+1}-w^{\top}{G_{\widetilde{D}}}Mv_{n+1})^{2}\mid{\widetilde{D}}]\Big{]}$
(177)
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\Big{[}\mathbb{E}_{x_{n+1},y_{n+1}}[(y_{n+1}-w^{\top}{G_{\widetilde{D}}}M_{:,1:d}x_{n+1})^{2}\mid{\widetilde{D}}]\Big{]}$
(178)
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}\Big{[}\|\overline{u}_{{\widetilde{D}}}-M_{:,1:d}^{\top}{G_{\widetilde{D}}}w\|_{2}^{2}+C_{{\widetilde{D}}}\Big{]}$
(By Equation 175)
$\displaystyle=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[\|\overline{u}_{{\widetilde{D}}}-M_{:,1:d}^{\top}{G_{\widetilde{D}}}w\|_{2}^{2}]+C\,.$
(179)

Here, $C$ is a constant independent of $w$ and $M$. This completes the proof
of the lemma. ∎

###### Proof of Lemma 7.

We prove this using 1, imitating the proof of Lemma 3. For convenience, let
$M(X)=\mathbb{E}[{\vec{y}}{\vec{y}}^{\top}\mid X]$. Then, the $(i,j)$-th entry
of $M(X)$ is $\mathbb{E}_{f}[f(x_{i})f(x_{j})]+\sigma^{2}$ if $i=j$ and
$\mathbb{E}_{f}[f(x_{i})f(x_{j})]$ if $i\neq j$, since
$\mathbb{E}[\epsilon_{i}]=0$ and the $\epsilon_{i}$ are i.i.d. and independent
of $X$. If we perform the change of variables $x_{i}\to Rx_{i}$ for a fixed
rotation matrix $R$ and all $i\in[n]$, then by 1,
$\mathbb{E}_{f}[f(Rx_{i})f(Rx_{j})]=\mathbb{E}_{f}[f(x_{i})f(x_{j})]$, meaning
that for any rotation matrix $R$,

$\displaystyle\mathbb{E}_{X}[X^{\top}M(X)X]=\mathbb{E}_{X}[(XR^{\top})^{\top}M(XR^{\top})(XR^{\top})]=R\mathbb{E}_{X}[X^{\top}M(XR^{\top})X]R^{\top}=R\mathbb{E}_{X}[X^{\top}M(X)X]R^{\top}$
(180)

by the rotational invariance of the $x_{i}$. This implies that
$\mathbb{E}_{X}[X^{\top}M(X)X]=\mathbb{E}_{{\widetilde{D}}}[X^{\top}{\vec{y}}{\vec{y}}^{\top}X]$
is a scalar multiple of the identity matrix.

Next, we consider
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\overline{u}_{{\widetilde{D}}}^{\top}]$.
For convenience, let
$J(X)=\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[{\vec{y}}\overline{u}_{{\widetilde{D}}}^{\top}\mid
X]$. If we make the change of variables $x_{i}\to Rx_{i}$, then the joint
distribution of $y_{1},\ldots,y_{n},y_{n+1}$ does not change by 1, meaning
that $\overline{u}_{{\widetilde{D}}}$ is replaced by
$R\overline{u}_{{\widetilde{D}}}$. Thus, we can conclude that $J(X)$ is
equivariant to rotations of all the $x_{i}$ by $R$:

$\displaystyle
J(XR^{\top})=\mathbb{E}[{\vec{y}}\overline{u}_{{\widetilde{D}}}^{\top}\mid
XR^{\top}]=\mathbb{E}[{\vec{y}}\overline{u}_{{\widetilde{D}}}^{\top}R^{\top}\mid
X]=J(X)R^{\top}\,.$ (181)

Thus,

$\displaystyle\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\overline{u}_{{\widetilde{D}}}]$
$\displaystyle=\mathbb{E}_{X}[X^{\top}J(X)]$ (182)
$\displaystyle=\mathbb{E}_{X}[(XR^{\top})^{\top}J(XR^{\top})]$ (By rotational
invariance of $\mathcal{N}(0,{I_{d\times d}})$)
$\displaystyle=R\mathbb{E}_{X}[X^{\top}J(X)]R^{\top}$ (By Equation 181)
$\displaystyle=R\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\overline{u}_{{\widetilde{D}}}^{\top}]R^{\top}\,.$
(183)

Thus,
$\mathbb{E}_{{\widetilde{D}}\sim{\mathcal{T}}}[X^{\top}{\vec{y}}\overline{u}_{{\widetilde{D}}}^{\top}]$
is a scalar multiple of the identity matrix. The final statement of the lemma
follows by taking the trace of the left and right hand sides. ∎

###### Proof of Lemma 8.

This follows by the same argument as Lemma 2 — here we use Lemma 6 in order to
show that we only need to consider the lower left and upper right blocks of
${G_{\widetilde{D}}}$, and the rest of the proof follows from linear algebraic
manipulations and applying Lemma 7 (in place of Lemma 4 which was used in the
proof of Lemma 2). ∎

###### Proof of Theorem 3.

This follows directly from Lemma 8, since the effective linear predictor being
$\eta X^{\top}{\vec{y}}$ is a necessary and sufficient condition for
minimizing the pre-training loss. ∎