# CausalLM is not optimal for in-context learning

Nan Ding Tomer Levinboim Jialin Wu
Sebastian Goodman Radu Soricut
Google Research
<EMAIL_ADDRESS>

###### Abstract

Recent empirical evidence indicates that transformer based in-context learning
performs better when using a prefix language model (prefixLM), in which in-
context samples can all attend to each other, compared to causal language
models (causalLM), which uses auto-regressive attention that prohibits in-
context samples to attend to future samples. While this result is intuitive,
it is not understood from a theoretical perspective. In this paper we take a
theoretical approach and analyze the convergence behavior of prefixLM and
causalLM under a certain parameter construction. Our analysis shows that both
LM types converge to their stationary points at a linear rate, but that while
prefixLM converges to the optimal solution of linear regression, causalLM
convergence dynamics follows that of an online gradient descent algorithm,
which is not guaranteed to be optimal even as the number of samples grows
infinitely. We supplement our theoretical claims with empirical experiments
over synthetic and real tasks and using various types of transformers. Our
experiments verify that causalLM consistently underperforms prefixLM in all
settings.

CausalLM is not optimal for in-context learning

Nan Ding Tomer Levinboim Jialin Wu Sebastian Goodman Radu Soricut Google
Research<EMAIL_ADDRESS>

##  1 Introduction

Transformer-based models Vaswani et al. (2017) have become the default
foundational model for various machine learning applications such as natural
language processing Devlin et al. (2018); Brown et al. (2020); Chowdhery et
al. (2022) and computer vision Dosovitskiy et al. (2020). Beyond their
traditional usage in machine learning applications, it has recently been
discovered that pretraining large transformers on a vast amounts of data leads
them to develop a striking ability referred to as In-context learning (ICL)
Brown et al. (2020). Specifically, once such pretraining is complete, these
models are able to solve new tasks at inference time (without changing their
parameters) by simply ingesting a short sequence (prefix) of labeled examples
from a task and then computing a prediction for a query example.

The ICL capability was first demonstrated by GPT-3 Brown et al. (2020), where
a causalLM, which is a Transformer decoder with auto-regressive attention
masks, was used as the main model architecture. However, follow up work
empirically found that restricting the auto-regressive masks on the entire
sequence is too prohibitive and therefore proposed the so-called prefixLM
Raffel et al. (2020) model which allows full attention within the prefix
examples. For example, in the UL-2 paper Tay et al. (2022) the authors
combined prefixLM with various types of losses and achieved better performance
on various natural language processing tasks.

Beyond the few empirical results in those and related papers, there is yet no
theoretical explanation that can explain the different ICL behavior of
prefixLM and causalLM. Indeed, theoretical studies of ICL are difficult due to
the complicated non-linearity of the (ordinary) transformer architecture.
However, recent work Von Oswald et al. (2023) focusing on ICL of linear
regression was able to show that a specifically designed parameter
construction of a one-layer Linear Self-Attention (LSA) transformer can
simulate a single step of gradient descent by using the in-context examples as
training data. Moreover, a different recent study Zhang et al. (2023) used
gradient flow to prove that a randomly initialized LSA-transformer indeed
converges to such a construction during training.

The study in this paper continues the above line of work and investigate the
convergence properties of ICL for both prefixLM and causalLM multi-layer LSA-
transformers in the linear regression setting. We summarizes our contributions
as follows:

  * •

We first provide a clear formal proof of the relationship between a multi-
layer LSA and multi-step gradient descents in linear regression.

  * •

We then show that both causalLM and prefixLM based multi-layer LSA-
transformers converge to their respective stationary points with with linear
rates of convergence. Moreover, we show that the stationary point of prefixLM
corresponds to the optimal least square solution of the linear regression
problem, while the stationary points of causalLM correspond to the iterative
weights of an online gradient descent algorithm with non-decaying step sizes.
Importantly, the stationary points obtained by causalLM may not become optimal
even as the number of in-context examples increases, a fact that indicates
that causalLM is not optimal for in-context learning.

  * •

Finally, we verify the above theoretical insights by conducting experiments
with LSA-transformers as well as ordinary softmax attention based transformers
on various synthetic tasks including linear regression, non-linear regression,
and multiclass classifications. In addition, we also compare causalLM and
prefixLM in large language models by finetuning PaLM2 Google et al. (2023)
checkpoints on FLAN tasks Chung et al. (2022) and evaluating them on MMLU
Hendrycks et al. (2020) 5-shot tasks. Our experimental results support our
theoretical findings and consistently show the superiority of prefixLM over
causalLM on ICL.

##  2 Background

We begin by reviewing a few types of transformer attention and in-context
learning (ICL), as well as a specific transformer construction for linear
regression ICL by Von Oswald et al. (2023) which our theories will be based
on. The discussions of other related work are deferred to Section 6.

###  2.1 Transformers: SSA, LSA, causalLM, and prefixLM

Given a sequence of input vectors
$\operatorname{\mathbf{Z}}=\left(\operatorname{\mathbf{z}}_{1},\ldots,\operatorname{\mathbf{z}}_{n}\right)$,
the output of standard Softmax Self-Attention (SSA) layer is

$\displaystyle\operatorname{\mathbf{z}}_{j}\leftarrow\operatorname{\mathbf{z}}_{j}+\operatorname{\mathbf{P}}\operatorname{\mathbf{V}}\operatorname{\mathbf{Z}}\text{softmax}(\operatorname{\mathbf{Z}}^{\top}\operatorname{\mathbf{K}}^{\top}\operatorname{\mathbf{Q}}\operatorname{\mathbf{z}}_{j}),$

where
$\operatorname{\mathbf{P}},\operatorname{\mathbf{V}},\operatorname{\mathbf{K}},\operatorname{\mathbf{Q}}$
respectively corresponds to the output projection, value transformation, key
transformation and query transformation.

Since the softmax attention of standard transformers is non-linear, its
theoretical analysis becomes complicated even for a single layer. For this
reason, theoretical approaches to analyze transformers have often resorted to
the Linear Self-Attention (LSA) layer Von Oswald et al. (2023); Zhang et al.
(2023), which simply drops the softmax function from the attention,

$\displaystyle\operatorname{\mathbf{z}}_{j}\leftarrow$
$\displaystyle\operatorname{\mathbf{z}}_{j}+\operatorname{\mathbf{P}}\operatorname{\mathbf{V}}\operatorname{\mathbf{Z}}(\operatorname{\mathbf{Z}}^{\top}\operatorname{\mathbf{K}}^{\top}\operatorname{\mathbf{Q}}\operatorname{\mathbf{z}}_{j})$
$\displaystyle=$
$\displaystyle\operatorname{\mathbf{z}}_{j}+\operatorname{\mathbf{P}}\operatorname{\mathbf{V}}\sum_{i=1}^{n}\operatorname{\mathbf{z}}_{i}\left(\operatorname{\mathbf{z}}_{i}^{\top}\operatorname{\mathbf{K}}^{\top}\operatorname{\mathbf{Q}}\operatorname{\mathbf{z}}_{j}\right).$
(1)

Furthermore, since each input $\operatorname{\mathbf{z}}_{j}$ can attend to
all positions $j\in\\{1\ldots n\\}$, this form of attention is categorized as
full (or bidirectional) attention, and is typically used in the transformer
encoder.

On the other hand, a (linear) transformer decoder uses the _auto-regressive_
attention

$\displaystyle\operatorname{\mathbf{z}}_{j}\leftarrow\operatorname{\mathbf{z}}_{j}+\operatorname{\mathbf{P}}\operatorname{\mathbf{V}}\sum_{i=1}^{j}\operatorname{\mathbf{z}}_{i}\left(\operatorname{\mathbf{z}}_{i}^{\top}\operatorname{\mathbf{K}}^{\top}\operatorname{\mathbf{Q}}\operatorname{\mathbf{z}}_{j}\right).$
(2)

which restricts each token $\operatorname{\mathbf{z}}_{j}$ to attend only to
previous positions (and itself) from $\\{1\ldots j\\}$. This restriction is
due to the role of the decoder as a causal language model (causalLM) which
predicts the next token in the context of the previously generated ones.

The original transformer involves both a full attention based encoder and an
auto-regressive attention based decoder. However, prominent NLP research has
often chosen either encoder-only (e.g. BERT Devlin et al. (2018)) or decoder-
only (e.g. GPT Brown et al. (2020), PaLM Chowdhery et al. (2022)) models
according to the task at hand. This is partially for the purpose of halving
the parameter sizes.

Another version of attention, between full and auto-regressive, followed from
the observation that some tasks can benefit from a prefix sequence such as
context or prompt. That is, the input sequence $\operatorname{\mathbf{Z}}$ is
composed of $n^{\prime}$ prefix tokens
$(\operatorname{\mathbf{z}}_{1},\ldots,\operatorname{\mathbf{z}}_{n^{\prime}})$
configured for the task, while the tokens
$(\operatorname{\mathbf{z}}_{n^{\prime}+1},\ldots,\operatorname{\mathbf{z}}_{n})$
represent the sample. Specifically, prefixLM Raffel et al. (2020) suggests the
following attention (in its LSA version):

$\displaystyle\operatorname{\mathbf{z}}_{j}\leftarrow\operatorname{\mathbf{z}}_{j}+\operatorname{\mathbf{P}}\operatorname{\mathbf{V}}\sum_{i=1}^{\max(j,n^{\prime})}\operatorname{\mathbf{z}}_{i}\left(\operatorname{\mathbf{z}}_{i}^{\top}\operatorname{\mathbf{K}}^{\top}\operatorname{\mathbf{Q}}\operatorname{\mathbf{z}}_{j}\right),$

where $\max(j,n^{\prime})$ ensures each prefix token
$\operatorname{\mathbf{z}}_{j}$ with $j<n^{\prime}$ can attend to all prefix
tokens.

###  2.2 In-context learning

A formal framework of in-context learning has been described in various
existing literature such as Garg et al. (2022); Zhang et al. (2023). Here, we
briefly review the problem setting and introduce notation that will be used
across the paper.

In-context learning refers to the ability of models to produce context-driven
predictions at inference time. That is, at inference time, a model is fed with
a sequence consisting of input-label pairs and a query input
$(\operatorname{\mathbf{x}}_{1},y_{1},\ldots,\operatorname{\mathbf{x}}_{n},y_{n},\operatorname{\mathbf{x}}_{query})$
and its goal is to predict the label $y_{query}$ of
$\operatorname{\mathbf{x}}_{query}$ using the context examples
$(\operatorname{\mathbf{x}}_{1},y_{1},\ldots,\operatorname{\mathbf{x}}_{n},y_{n})$
(specifically, without changing the model parameters).

###  2.3 Linear regression in-context learners

Linear regression is a classical machine learning problem. Given a set of
input-label pairs $\left(\operatorname{\mathbf{x}}_{i},y_{i}\right)$, the goal
is to find an optimal weight vector $\operatorname{\mathbf{w}}$ that minimizes
the l2-loss:

$L(\operatorname{\mathbf{w}})=\frac{1}{2n}\sum_{i=1}^{n}\|\operatorname{\mathbf{w}}\operatorname{\mathbf{x}}_{i}-y_{i}\|_{2}^{2}.$

The gradient of the loss is
$\nabla_{\operatorname{\mathbf{w}}}L=\frac{1}{n}\sum_{i=1}^{n}(\operatorname{\mathbf{w}}\operatorname{\mathbf{x}}_{i}-y_{i})\operatorname{\mathbf{x}}_{i}^{\top}$,
and a gradient descent algorithm with step size $\eta$ follows the update
rule:

$\displaystyle\operatorname{\mathbf{w}}^{(l)}=$
$\displaystyle\operatorname{\mathbf{w}}^{(l-1)}+\frac{\eta}{n}\sum_{i=1}^{n}(y_{i}-\operatorname{\mathbf{w}}^{(l-1)}\operatorname{\mathbf{x}}_{i})\operatorname{\mathbf{x}}_{i}^{\top}.$
(3)

Using linear regression as a lens to study in-context learning was first
proposed in Garg et al. (2022), where the authors laid out an approach for
training transformers to in-context learn a class of simple predictors,
including linear regression. However, no theoretical study is provided. More
recently, and most relevant to our work, Von Oswald et al. (2023) proposed a
succinct construction that demonstrates how a single LSA layer can effectively
implement a single gradient descent step. According to their setup the input
is formulated as

$\displaystyle\operatorname{\mathbf{Z}}=(\operatorname{\mathbf{z}}_{1}^{(0)},\ldots,\operatorname{\mathbf{z}}_{n}^{(0)}),\;\text{where}\;\operatorname{\mathbf{z}}_{j}^{(0)}=\begin{pmatrix}\operatorname{\mathbf{x}}_{j}\\\
y_{j}\end{pmatrix}$ (4)

and the parameter matrices of (1) are set as:

$\displaystyle\operatorname{\mathbf{K}}$
$\displaystyle=\operatorname{\mathbf{Q}}=\begin{pmatrix}\operatorname{\mathbf{I}}_{d\times
d}&{\bf 0}\\\ 0&0\end{pmatrix},\operatorname{\mathbf{V}}=\begin{pmatrix}{\bf
0}_{d\times d}&{\bf 0}\\\ \operatorname{\mathbf{w}}^{(0)}&-1\end{pmatrix},$
$\displaystyle\operatorname{\mathbf{P}}$
$\displaystyle=\frac{\eta}{n}\operatorname{\mathbf{I}},$ (5)

where $\operatorname{\mathbf{w}}^{(0)}$ is an initial weight vector. Von
Oswald et al. (2023) then showed that this configuration results in an update
of their so-called transformed target $y_{j}\leftarrow
y_{j}+\eta\left(\nabla_{\operatorname{\mathbf{w}}^{(0)}}L\right)\operatorname{\mathbf{x}}_{j}$,
and that this target update is equivalent to the one performed by a single-
step gradient descent of linear regression.

Although the construction of Von Oswald et al. (2023) connected LSA-based ICL
to the gradient descent of linear regression, the "transformed target" view
seems unnatural***The traditional ML formulation updates the weight vector or
the model prediction, while the groundtruth target remains fixed. to work
with. Moreover, their extension from single-layer to multi-layer LSA is
unfortunately unclear.

##  3 Multi-layer in-context learner

In this section, we provide a formal proof that a multi-layer LSA under the
construction of Von Oswald et al. (2023) progresses identically to multi-step
gradient descent.

Instead of the "transformed target" view, the following proposition explicitly
connects the GD weights of (3) to the outputs of the multi-layer LSA under the
constructions of $\operatorname{\mathbf{K}}$, $\operatorname{\mathbf{Q}}$,
$\operatorname{\mathbf{P}}$ and $\operatorname{\mathbf{V}}$ in (5). Note that
we keep $\operatorname{\mathbf{w}}^{(0)}=0$ in the proposition because it
simplifies the equations and makes the outputs more meaningful. However, such
specification is not mandatory, and we provide general propositions, for
arbitrary $\operatorname{\mathbf{w}}^{(0)}$, in Appendix B.

###### Proposition 1

For a multi-layer LSA satisfying the construction (5) and with
$\operatorname{\mathbf{w}}^{(0)}=0$, if its input $\operatorname{\mathbf{Z}}$
is formatted as (4), then its $l$-th layer output is
$\operatorname{\mathbf{z}}_{j}^{(l)}=(\operatorname{\mathbf{x}}_{j}^{\top},\delta_{j}^{(l)})^{\top}$,
where
$\delta_{j}^{(l)}=y_{j}-\operatorname{\mathbf{w}}^{(l)}\operatorname{\mathbf{x}}_{j}$
and $\operatorname{\mathbf{w}}^{(l)}$ is the $l$-th updated weight from the
gradient descents update rule in (3).

_Proof Sketch:_ Plugging in $\operatorname{\mathbf{K}}$,
$\operatorname{\mathbf{Q}}$, $\operatorname{\mathbf{P}}$ and
$\operatorname{\mathbf{V}}$ of (5) with $\operatorname{\mathbf{w}}^{(0)}=0$
and
$\operatorname{\mathbf{z}}_{j}^{(l)}=(\operatorname{\mathbf{x}}_{j}^{\top},\delta_{j}^{(l)})^{\top}$
into (1), we obtain that for all $l>0$,

$\displaystyle\begin{pmatrix}\operatorname{\mathbf{x}}_{j}\\\
\delta_{j}^{(l)}\end{pmatrix}=\begin{pmatrix}\operatorname{\mathbf{x}}_{j}\\\
\delta_{j}^{(l-1)}\end{pmatrix}-\frac{\eta}{n}\sum_{i=1}^{n}\begin{pmatrix}{\bf
0}\\\
\delta_{i}^{(l-1)}\end{pmatrix}\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j}.$

Since $\operatorname{\mathbf{z}}_{j}$ never changes its first $d$-dimension
corresponding to $\operatorname{\mathbf{x}}_{j}$, we can simplify it and focus
only on $\delta_{j}^{(l)}$, which is the last output coordinate of the $j$-th
LSA-layer,

$\displaystyle\delta_{j}^{(l)}=$
$\displaystyle\;\delta_{j}^{(l-1)}-\frac{\eta}{n}\sum_{i=1}^{n}\delta_{i}^{(l-1)}\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j},$
(6)

with $\delta_{j}^{(0)}=y_{j}$. Defining
$\tilde{y}_{j}^{(l)}=y_{j}-\delta_{j}^{(l)}$ and rearranging (6), we obtain
$\tilde{y}_{j}^{(0)}=0$ and $\forall l>0$:

$\displaystyle\tilde{y}_{j}^{(l)}=$
$\displaystyle\;\tilde{y}_{j}^{(l-1)}+\frac{\eta}{n}\sum_{i=1}^{n}(y_{i}-\tilde{y}_{i}^{(l-1)})\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j}.$
(7)

Finally, using (7) and the fact that
$\tilde{y}_{j}^{(0)}=0=\operatorname{\mathbf{w}}^{(0)}\operatorname{\mathbf{x}}_{j}$,
it can be proved by induction that $\forall
l:\tilde{y}_{j}^{(l)}=\operatorname{\mathbf{w}}^{(l)}\operatorname{\mathbf{x}}_{j}$.
A complete proof is provided in Appendix A.

To summarize, the newly introduced variable $\tilde{y}_{j}^{(l)}$ is exactly
the prediction of the $l$-th gradient descent weights
$\operatorname{\mathbf{w}}^{(l)}$ for $\operatorname{\mathbf{x}}_{j}$ , and
$\delta_{j}^{(l)}$ is the difference between the true label $y_{j}$ and the
predicted $\tilde{y}_{j}^{(l)}$. Therefore, $\tilde{y}_{j}^{(l)}$ serves as a
bridge to connect the LSA output $\delta_{j}^{(l)}$ and the GD weight
$\operatorname{\mathbf{w}}^{(l)}$.

So far, we have dealt with the behavior of LSA layers with full attention. In
what follows, we move on to the practical setting of in-context learning,
where the input contains not only $n$ in-context (training) examples in the
format of (4), but also an additional (test) query
$\operatorname{\mathbf{z}}_{query}^{(0)}=(\operatorname{\mathbf{x}}_{query}^{\top},0)^{\top}$.
In particular, we will focus on the two common ICL variants: prefixLM and
causalLM, each with a different type of attention.

###  3.1 PrefixLM ICL

A prefixLM ICL treats the in-context examples $\operatorname{\mathbf{Z}}$ as
the prefix and uses full attention on the first $n$ positions, so that they
can each freely attend to each other. The last query vector
$\operatorname{\mathbf{z}}_{query}$ can also attend to any example in
$\operatorname{\mathbf{Z}}$, but cannot attend to itself†††This is because the
query does not contain a meaningful label. Attending to itself would cause it
to include its last-dim input as a label, which would contaminate the
resulting multi-layer prediction. This observation was not considered in Von
Oswald et al. (2023).. As a result, the updates of the prefixLM-ICL under the
same construction follow (6), with the outputs of the $l$-th layer being,

$\displaystyle\delta_{j}^{(l)}$
$\displaystyle=y_{j}-\tilde{y}_{j}^{(l)}=y_{j}-\operatorname{\mathbf{w}}^{(l)}\operatorname{\mathbf{x}}_{j},$
$\displaystyle\text{and}\;\;\delta_{query}^{(l)}$
$\displaystyle=-\tilde{y}_{query}^{(l)}=-\operatorname{\mathbf{w}}^{(l)}\operatorname{\mathbf{x}}_{query},$

where the initial $\tilde{y}_{j}^{(0)}=\tilde{y}_{query}^{(0)}=0$.

Intuitively, the dynamics of the prefixLM ICL is as follows: all
$\tilde{y}_{j}^{(l)}$ starts as 0 at $l=0$, and gradually approach to the true
label $y_{j}$ as $l$ increases, so that the difference (also as the output)
$\delta_{j}^{(l)}$ gradually approaches to 0. At the same time,
$\delta_{query}^{(l)}$ starts at 0, and gradually approaches to $-y_{query}$,
the negation of the query label. Figure 1 provides an illustration of these
dynamics.

Figure 1: The inputs/outputs of a multi-layer in-context learner. We omitted
$\operatorname{\mathbf{x}}_{j}$ and $\operatorname{\mathbf{x}}_{query}$ since
they are unchanged.

###  3.2 CausalLM ICL

A causalLM applies auto-regressive attention throughout the entire sequence.
Therefore, plugging the same $\operatorname{\mathbf{K}}$,
$\operatorname{\mathbf{Q}}$, $\operatorname{\mathbf{P}}$,
$\operatorname{\mathbf{V}}$ into (2), the update rules of (6) and (7) become:

$\displaystyle\delta_{j}^{(l)}=$
$\displaystyle\;\delta_{j}^{(l-1)}-\frac{\eta}{n}\sum_{i=1}^{j}\delta_{i}^{(l-1)}\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j},$
(8) $\displaystyle\tilde{y}_{j}^{(l)}=$
$\displaystyle\;\tilde{y}_{j}^{(l-1)}+\frac{\eta}{n}\sum_{i=1}^{j}(y_{i}-\tilde{y}_{i}^{(l-1)})\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j}$
(9) ‡‡‡There is another way of update which changes $\eta/n$ to $\eta/j$ for
the $j$-th example. We provide more details in Appendix C and show it performs
worse than the main version in (8).

with $\delta_{j}^{(l)}=y_{j}-\tilde{y}_{j}^{(l)}$. Moreover, since different
$\delta_{j},\tilde{y}_{j}$ are exposed to different ranges of inputs, there is
no uniform $\operatorname{\mathbf{w}}$ as in (3) that is associated with all
$\tilde{y}_{j}$. Instead, if we define $\operatorname{\mathbf{w}}_{j}$ for
each different position $j$ with $\operatorname{\mathbf{w}}_{j}^{(0)}=0$ and

$\displaystyle\operatorname{\mathbf{w}}_{j}^{(l)}=\operatorname{\mathbf{w}}_{j}^{(l-1)}+\frac{\eta}{n}\sum_{i=1}^{j}(y_{i}-\operatorname{\mathbf{w}}_{i}^{(l-1)}\operatorname{\mathbf{x}}_{i})\operatorname{\mathbf{x}}_{i}^{\top}$
(10)

then we have the following proposition:

###### Proposition 2

For a multi-layer causalLM-LSA satisfying (5) with
$\operatorname{\mathbf{w}}^{(0)}=0$, if its input $\operatorname{\mathbf{Z}}$
is formatted as (4), then its $l$-th layer output is
$\operatorname{\mathbf{z}}_{j}^{(l)}=(\operatorname{\mathbf{x}}_{j}^{\top},\delta_{j}^{(l)})^{\top}$,
where
$\delta_{j}^{(l)}=y_{j}-\operatorname{\mathbf{w}}_{j}^{(l)}\operatorname{\mathbf{x}}_{j}$
and $\operatorname{\mathbf{w}}_{j}^{(l)}$ follow (10).

The proof of Proposition 2 is provided in Appendix A. Similar to prefixLM-ICL,
causalLM-ICL also has
$\tilde{y}_{j}^{(l)}=\operatorname{\mathbf{w}}_{j}^{(l)}\operatorname{\mathbf{x}}_{j}$,
and

$\displaystyle\delta_{query}^{(l)}$
$\displaystyle=-\tilde{y}^{(l)}_{query}=-\operatorname{\mathbf{w}}_{n}^{(l)}\operatorname{\mathbf{x}}_{query}.$

Since prefixLM-ICL and causalLM-ICL are associated with different update
rules: $\operatorname{\mathbf{w}}^{(l)}$ follows (3) and
$\operatorname{\mathbf{w}}_{j}^{(l)}$ follows (10), their outcomes are
significantly different. In the next section, we look closely into their
convergence properties, and show why prefixLM is superior to causalLM for in-
context learning.

##  4 Convergence of the multi-layer in-context learners

In this section, we prove that both multi-layer prefixLM and causalLM converge
to their respective stationary points with increasing layers in linear rates.
In addition, we show that the stationary point of prefixLM corresponds to the
least-square solution of the linear regression problem, while the ones
corresponding to causalLM are equivalent to the iterative weights of online
gradient descent of linear regression.

###  4.1 Convergence of the prefixLM ICL

The fact that a multi-layer prefixLM computation exactly follows the update
rule of $\operatorname{\mathbf{w}}^{(l)}$ as in (3), implies that the layer
outputs of prefixLM have the same dynamics of multi-step gradient descent on a
linear regression problem. The convergence properties of such dynamics are
well-known, and are stated in the following proposition:

###### Proposition 3

If $\operatorname{\mathbf{w}}^{(l)}$ follows the iterative updates of (3),
then there exists a stationary point $\operatorname{\mathbf{w}}^{*}$ with
coefficients satisfying:

$\displaystyle\operatorname{\mathbf{y}}\operatorname{\mathbf{x}}^{\top}=\operatorname{\mathbf{w}}^{*}\operatorname{\mathbf{X}}\operatorname{\mathbf{X}}^{\top},$

where $\operatorname{\mathbf{y}}=(y_{1},\ldots,y_{n})$ and
$\operatorname{\mathbf{X}}=(\operatorname{\mathbf{x}}_{1},\ldots,\operatorname{\mathbf{x}}_{n})$.
Furthermore, the iterative weights $\operatorname{\mathbf{w}}^{(l)}$ converge
to $\operatorname{\mathbf{w}}^{*}$ with a linear rate of convergence:

$\displaystyle\operatorname{\mathbf{w}}^{(l)}-\operatorname{\mathbf{w}}^{*}=(\operatorname{\mathbf{w}}^{(l-1)}-\operatorname{\mathbf{w}}^{*})(\operatorname{\mathbf{I}}-\frac{\eta}{n}\operatorname{\mathbf{X}}\operatorname{\mathbf{X}}^{\top}).$

That is, Proposition 3 holds for the multi-layer prefixLM, so that the same
exact $\operatorname{\mathbf{w}}^{*}$ is also the stationary point of
prefixLM, to which it converges in a linear rate. Furthermore this stationary
point is exactly the (optimal) least square solution of the linear regression
problem.

###  4.2 Convergence of the causalLM ICL

Following the update rule of (10), we can view a multi-layer causalLM as
implicitly maintaining different weight vectors
$\operatorname{\mathbf{w}}_{j}$ for each position $j$. In what follows, we
show that: (a) Each such position $j$ has its own stationary point
$\operatorname{\mathbf{w}}_{j}^{*}$, which appears to be different from the
global optimal point $\operatorname{\mathbf{w}}^{*}$ of linear regression; (b)
even when the number of in-context samples $n$ grows to infinity, convergence
to $\operatorname{\mathbf{w}}^{*}$ is not guaranteed.

Specifically, in Appendix A we provide a proof for the following proposition:

###### Proposition 4

If
$\operatorname{\mathbf{w}}_{j}^{(l)}=\sum_{i=1}^{j}a_{i,j}^{(l)}\operatorname{\mathbf{x}}_{i}^{\top}$
follows the iterative updates of (10), then

$\displaystyle a_{i,j}^{(l)}=a_{i,i}^{(l)}\equiv a_{i}^{(l)}\;\;\;\;\forall
j\geq i,$

and there exists stationary points
$\operatorname{\mathbf{w}}_{j}^{*}=\sum_{i=1}^{j}a_{i}^{*}\operatorname{\mathbf{x}}_{i}^{\top}$
(for $j\in 1,\ldots,n$) with coefficients from
$\operatorname{\mathbf{a}}^{*}=(a^{*}_{1},\ldots,a^{*}_{n})$ that satisfy
$\operatorname{\mathbf{y}}=\operatorname{\mathbf{a}}^{*}\operatorname{\mathbf{T}}$,
where

$\displaystyle\operatorname{\mathbf{T}}=\begin{pmatrix}\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{1}&\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{2}&\cdots&\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{n}\\\
0&\operatorname{\mathbf{x}}_{2}^{\top}\operatorname{\mathbf{x}}_{2}&\cdots&\operatorname{\mathbf{x}}_{2}^{\top}\operatorname{\mathbf{x}}_{n}\\\
\vdots&\vdots&\ddots&\vdots\\\
0&0&\cdots&\operatorname{\mathbf{x}}_{n}^{\top}\operatorname{\mathbf{x}}_{n}\end{pmatrix}.$

Furthermore, the coefficients $\operatorname{\mathbf{a}}^{(l)}$ converges to
the stationary point $\operatorname{\mathbf{a}}^{*}$ with linear rate of
convergence:

$\displaystyle\operatorname{\mathbf{a}}^{(l)}-\operatorname{\mathbf{a}}^{*}=(\operatorname{\mathbf{a}}^{(l-1)}-\operatorname{\mathbf{a}}^{*})(\operatorname{\mathbf{I}}-\frac{\eta}{n}\operatorname{\mathbf{T}}).$

This proposition implies that the stationary points
$\operatorname{\mathbf{w}}_{j}^{*}$ of causalLM-ICL are different from
$\operatorname{\mathbf{w}}^{*}$, the least square solution of linear
regression. However, a natural question is: if $j$ increases, would
$\operatorname{\mathbf{w}}_{j}^{*}$ ultimately converge to the optimal
solution?

To answer this question, the next proposition shows that the stationary points
$\operatorname{\mathbf{w}}_{j}^{*}$ follow an online gradient descent
algorithm, whose loss and gradient at the $j$-th step is,

$\displaystyle L_{j}(\operatorname{\mathbf{w}}_{j})$
$\displaystyle=\frac{1}{2}(\operatorname{\mathbf{w}}_{j}\operatorname{\mathbf{x}}_{j+1}-y_{j+1})^{2},$
$\displaystyle\nabla_{\operatorname{\mathbf{w}}_{j}}L_{j}(\operatorname{\mathbf{w}}_{j})$
$\displaystyle=(\operatorname{\mathbf{w}}_{j}\operatorname{\mathbf{x}}_{j+1}-y_{j+1})\operatorname{\mathbf{x}}_{j+1}^{\top}.$

###### Proposition 5

Assuming that $\operatorname{\mathbf{w}}_{j}^{*}$ is the stationary points
obtained in Proposition 4, then

$\displaystyle\operatorname{\mathbf{w}}_{j+1}^{*}=\operatorname{\mathbf{w}}_{j}^{*}-\frac{1}{\|\operatorname{\mathbf{x}}_{j+1}\|_{2}^{2}}\nabla_{\operatorname{\mathbf{w}}_{j}^{*}}L_{j}(\operatorname{\mathbf{w}}_{j}^{*}).$
(11)

The proof of Proposition 5 is provided in Appendix A. Note that online
gradient descent is known to converge to an optimal solution only with a
decaying step size $j^{-\nu}$ for $\nu>0$ Jentzen and Von Wurstemberger
(2020). Since the step size of (11) does not decay, we conclude that causalLM
may not converge to $\operatorname{\mathbf{w}}^{*}$ even with increasing
layers and increasing number of in-context examples. More concretely, as for
the case of in-context learning, where the number of in-context examples $n$
is limited, convergence to the optimal solution
$\operatorname{\mathbf{w}}^{*}$ cannot be achieved by causalLM-ICL.

##  5 Numerical experiments

Our experiments mainly contain three parts.

  * •

We first use LSA-transformers on linear regression problems to directly verify
our theorems. We show that despite that the training error of causalLM and
prefixLM both decays in linear rates, the test error of causalLM is
significantly larger which indicates its stationary solution is not optimal.

  * •

Secondly, we use ordinary softmax transformers on a few synthetic problems
including linear regression, nonlinear regression and multiclass
classification. We show that our theoretical insights generalize to other
types of tasks, in the sense that prefixLM still outperforms causalLM on ICL
in all these cases.

  * •

Lastly, we conduct a few practical ICL experiments based on large language
models.

Figure 2: Left/Middle: the MSE errors on in-context examples and query
examples of multi-layer LSA-based prefixLM/causalLM-ICLs with 40 in-context
training examples. Right: the query MSE errors of the stationary points of the
causalLM-ICL with up to 300 in-context examples.

###  5.1 LSA-transformers on linear regression

In order to directly verify our theorems from Section 4, we first study in-
context learning on linear regression problem with the LSA transformer of (5).
Each of the input sequence contains 40 in-context examples and 200 queries,
and each query attends to all the in-context examples but does not attend to
each other. See Appendix D for an illustration. The data input
$\operatorname{\mathbf{x}}_{i}$ of the sequence is sampled from
$\operatorname{\mathcal{U}}(-1,1)^{16}$. Each sequence is associated with a
single weight vector $\operatorname{\mathbf{w}}$ that is sampled from
$\operatorname{\mathcal{N}}(0,\operatorname{\mathbf{I}})$, and the labels are
computed as $y_{i}=\operatorname{\mathbf{w}}\operatorname{\mathbf{x}}_{i}$.
Assuming the prediction of each layer is $\tilde{y}_{i}^{(l)}$, we evaluate
the mean square error $\|y_{i}-\tilde{y}_{i}^{(l)}\|_{2}^{2}$ on both in-
context examples as well as query examples across different layers $l$.

The results are plotted in Figure 2 left (for prefixLM) and middle (for
causalLM). Our results are averaged over 64 randomly generated sequences. As
we can see, although both prefixLM and causalLM has a linear rate of
convergence (with respect to the number of layers) on the in-context examples,
the query errors of causalLM stuck at above the $10^{-1}$ level, while the
query error of prefixLM decays in the same linear rate as its training error.

Furthermore, in Figure 2 right, we plot the query errors of the stationary
points of causalLM-ICL with increasing number of in-context examples up to
300. Although causalLM-ICL is able to eventually converge to optimal solution
when $\mu_{x}=0$, it takes more than 100 examples to reach below $10^{-2}$.
The convergence is even worse as we vary the input distribution
$\operatorname{\mathbf{x}}\sim\operatorname{\mathcal{U}}(-1,1)^{d}+\mu_{x}$
with increasing $\mu_{x}\in\left\\{0,1,2,3\right\\}$, which demonstrates that
causalLM-ICL is not optimal for few-shot ICL.

###  5.2 Standard transformers on synthetic tasks

Previous experiments provided a proof of concept verification of the
propositions from Section 4. Next we verify if a standard softmax transformer-
based prefixLM and causalLM ICL exhibit similar differences on various types
of synthetic tasks including linear regression, non-linear regression and
multiclass classification.

All three tasks used 16-dim inputs with
$\operatorname{\mathbf{x}}\sim\operatorname{\mathcal{U}}(-1,1)^{16}$ and
$\operatorname{\mathbf{w}}\sim\operatorname{\mathcal{N}}(0,\operatorname{\mathbf{I}})$.
For non-linear regression, we apply a sigmoid activation on the logit such
that $y=\text{sigmoid}(\operatorname{\mathbf{w}}\operatorname{\mathbf{x}})$;
and for multiclass classification, we randomly generate three
$\operatorname{\mathbf{w}}_{c}\sim\operatorname{\mathcal{N}}(0,\operatorname{\mathbf{I}})$,
and assign labels based on
$y=\operatorname*{\mathop{\mathrm{argmax}}}_{c}\left\\{\operatorname{\mathbf{w}}_{c}\operatorname{\mathbf{x}}\right\\}$.

We trained a few 24-layer transformers containing 128 hidden units and 2
heads. Besides of the comparisons of prefixLM and causalLM, we also compare
the transformers with or without sharing layers (SL vs UL). In particular, the
sharing-layer transformer can be considered a recurrent system Dehghani et al.
(2018) where the dynamic is continual along the layers and a stationary point
may exist given infinite number of layers, which makes it closer to our
constructed LSA.

The ICL training dataset contains 64,000 training sequences. Each sequence
contains 40 in-context examples and 20 queries, where queries are independent
of each other similar to Section 5.1. The transformers are trained with batch
size 64 for 100 epochs. More details of the hyper-parameters of the
experiments are provided in Appendix D.

We evaluate the ICL performance using 64 holdout test sequences and report the
test errors on the query examples. The results are summarized in Table 1. We
find that both prefixLM-SL and prefixLM-UL significantly outperform their
counterparts of causalLM in all cases. As a side note, transformer-SL appears
to outperform transformer-UL in the classification tasks, which indicates the
overfitting problem of the latter due to over-parameterization. In addition,
we also add probes at the output of each SL-transformer layer to visualize the
test errors of intermediate layers in Figure 3. Comparing Figure 3 and Figure
2 (left/middle) reveals some similarities. Although the test query errors of
causalLM decay in roughly the same rate as the ones of prefixLM in earlier
layers, the decays become much slower in later layers possibly due to the
nature of its non-optimal stationary points. These results suggest that the
title argument of the paper also holds beyond LSA-based transformers and
linear regression.

| LR | N-LR | MC
---|---|---|---
PrefixLM-SL | 8.6e-3 | 1.5e-4 | 24.1
CausalLM-SL | 1.9e-1 | 2.7e-3 | 27.0
PrefixLM-UL | 2.5e-3 | 9.0e-5 | 27.6
CausalLM-UL | 1.6e-2 | 2.9e-3 | 32.1
Table 1: The test query errors of the unshared-layer (UL) and sharing-layer
(SL) transformer-ICLs on linear regression (LR), non-linear regression (NLR),
and multiclass classification (MC) tasks.

Figure 3: The test query errors of the 24-layer SL-transformers based
prefixLM/causalLM-ICLs on linear regression (left), non-linear regression
(middle), and multiclass classification (right).

###  5.3 ICL on large language models

In order to compare the ICL performance of causalLM and prefixLM on large
language models, we conducted experiments based on the PaLM2 family Google et
al. (2023). In particular, we finetuned three sizes of PaLM2 language models:
Gecko, Otter and Unicorn§§§https://blog.google/technology/ai/google-palm-2-ai-
large-language-model/ on FLAN Chung et al. (2022), a finetuning procedure
containing a mixture of over 1000 tasks. We used the default recipe for FLAN-
PaLM2 finetuning Google et al. (2023); Chung et al. (2022) and evaluate the
ICL capability of the finetuned models on the Massive Multi-task Language
Understanding (5-shot MMLU) tasks Hendrycks et al. (2020).

From Table 2, we find that in all three sizes of the PaLM2 models, prefixLM
outperforms causalLM for more than 3% accuracies averaged over the 57 MMLU
tasks. This result empirically verifies our conjectures on highly non-linear
cases, that the autoregressive nature of causalLM restricts the mutual
attentions between the in-context examples which yields suboptimal result.

| Gecko | Otter | Unicorn
---|---|---|---
PrefixLM | 46.6 | 64.8 | 81.4
CausalLM | 43.3 | 61.0 | 78.0
Table 2: The average test query accuracies on 5-shot MMLU tasks with FLAN-
finetuned PaLM2-Gecko/Otter/Unicorn prefixLM/causalLM checkpoints. Google et
al. (2023) reported a similar averaged accuracy of 81.2 on Unicorn-Prefix.

##  6 Related Work

Ever since GPT-3 Brown et al. (2020) exhibited its in-context learning
abilities in various language inference and translation tasks, there has been
tremendous interest in understanding the mechanics behind In-Context Learning
(ICL). Currently, there are two main camps of thought that try to explain ICL:
(1) the representation camp, which views ICL behavior as a topic model that
extracts relevant memories based on the topic of the context Xie et al.
(2021); Min et al. (2022) \- these works support this view with the findings
that in-context learner sometimes behaved similarly even when the label of the
training examples were permuted Min et al. (2022). (2) the algorithmic camp,
which holds that LLMs learns to implement a learning algorithm Garg et al.
(2022); Akyürek et al. (2022); Von Oswald et al. (2023) and then run it during
ICL - these works usually propose a construction of the transformer parameters
and show that it can solve certain simple tasks (e.g. linear regression), then
empirically verify that transformers track the behavior of the algorithm of
interest.

Moreover, recent studies of large-scale data and model Wei et al. (2023)
discovered that large language models seem to exhibit certain emergent
behavior, where, ICL is memory-based on small-to-medium sized models or data,
but becomes more algorithm-based on larger model and data. For example, Wei et
al. (2023) showed that a large language model is able to respond accordingly
to the flipped label in in-context examples, opposing the findings of Min et
al. (2022).

Since most ICL applications only involve few shots of context examples, it
seems reasonable to conjecture that the memory of a _deep_ representation and
a _shallow_ predictor algorithm may co-exist in contributing the in-context
learning capabilities. Since the representation learning of large language
models have been universally acknowledged, it is more interesting to
investigate how transformer learns to in-context learn shallow predictors with
few-shot examples.

Focusing on work from the algorithmic camp, we note that Garg et al. (2022)
were the first to suggest using linear regression to study in-context
learning. The authors empirically found that a 12-layer transformer is able to
achieve similar results as a least-square solver on a 20-dim linear regression
problem with around 20 in-context examples. Beyond linear regression, they
also found that transformers can in-context learn a few other classes of
shallow predictors, including two-layer Relu networks.

Probably the first formal theoretical investigation of the linear regression
in-context learners is Akyürek et al. (2022). They first showed that a
transformer layer can approximately conduct four basic operations: mov, mul,
div, aff. They then cleverly combined these four operations and showed that a
gradient descent step of linear regression can be implemented with a 4-head
8-layer transformer with $O(d)$ hidden units, where $d$ is the dimension of
the inputs $\operatorname{\mathbf{x}}$. Despite their novel construction, the
result itself provides only a loose upper bound on the model size (or depth)
that is required for simulating linear regression within a transformer - for
example, Von Oswald et al. (2023) reported that a 2 or 5-layer transformer
already achieves significantly better results than a single-step gradient
descent for linear regression.

Because of the significant discrepancy between the construction of Akyürek et
al. (2022) and the empirical results, the one-layer LSA construction of Von
Oswald et al. (2023) appears to be more appealing and matches the experimental
results better. Moreover, a most recent work by Zhang et al. (2023) used
gradient flow to prove that by initializing
$\operatorname{\mathbf{w}}^{(0)}=0$, such matrix constructions can indeed be
learned by an LSA transformer. This is why our paper follows this construction
and studies its multi-layer convergence properties with different types of
attention (prefixLM vs causalLM).

##  7 Conclusion

In this paper, we analyzed the convergence properties of two types of widely-
used transformer-based language models (causalLM and prefixLM), during in-
context learning. Using a simplified LSA attention in a linear regression
setting, we proved that both LM types converge to their stationary points in
linear rates, but that their stationary points have significantly different
properties. In particular, the stationary points of prefixLM coincides with
the optimal least square solution; while the ones of causalLM is equivalent to
the weights of an online learning system, that is not guaranteed to converge
to the optimal solution. Our experiments verified the above theoretical
results, and also empirically extend the findings to general transformer on
non-linear regression as well as classification tasks. Finally, we compare
causalLM and prefixLM on a few large language models and find that prefixLM
also consistently wins over causalLM in practical few-shot tasks.

## Acknowledgements

We thank Xinhua Zhang for helpful discussions and the references about online
learning.

## References

  * Akyürek et al. (2022) Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. 2022\.  What learning algorithm is in-context learning? investigations with linear models.  _arXiv preprint arXiv:2211.15661_.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901.
  * Chowdhery et al. (2022) Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. 2022.  Palm: Scaling language modeling with pathways.  _arXiv preprint arXiv:2204.02311_.
  * Chung et al. (2022) Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. 2022.  Scaling instruction-finetuned language models.  _arXiv preprint arXiv:2210.11416_.
  * Dehghani et al. (2018) Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Łukasz Kaiser. 2018.  Universal transformers.  _arXiv preprint arXiv:1807.03819_.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2018.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_.
  * Dosovitskiy et al. (2020) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. 2020.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_.
  * Garg et al. (2022) Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. 2022.  What can transformers learn in-context? a case study of simple function classes.  _Advances in Neural Information Processing Systems_ , 35:30583–30598.
  * Google et al. (2023) Google, Rohan Anil, Andrew M. Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, Eric Chu, Jonathan H. Clark, Laurent El Shafey, Yanping Huang, Kathy Meier-Hellstern, Gaurav Mishra, Erica Moreira, Mark Omernick, Kevin Robinson, Sebastian Ruder, Yi Tay, Kefan Xiao, Yuanzhong Xu, Yujing Zhang, Gustavo Hernandez Abrego, Junwhan Ahn, Jacob Austin, Paul Barham, Jan Botha, James Bradbury, Siddhartha Brahma, Kevin Brooks, Michele Catasta, Yong Cheng, Colin Cherry, Christopher A. Choquette-Choo, Aakanksha Chowdhery, Clément Crepy, Shachi Dave, Mostafa Dehghani, Sunipa Dev, Jacob Devlin, Mark Díaz, Nan Du, Ethan Dyer, Vlad Feinberg, Fangxiaoyu Feng, Vlad Fienber, Markus Freitag, Xavier Garcia, Sebastian Gehrmann, Lucas Gonzalez, Guy Gur-Ari, Steven Hand, Hadi Hashemi, Le Hou, Joshua Howland, Andrea Hu, Jeffrey Hui, Jeremy Hurwitz, Michael Isard, Abe Ittycheriah, Matthew Jagielski, Wenhao Jia, Kathleen Kenealy, Maxim Krikun, Sneha Kudugunta, Chang Lan, Katherine Lee, Benjamin Lee, Eric Li, Music Li, Wei Li, YaGuang Li, Jian Li, Hyeontaek Lim, Hanzhao Lin, Zhongtao Liu, Frederick Liu, Marcello Maggioni, Aroma Mahendru, Joshua Maynez, Vedant Misra, Maysam Moussalem, Zachary Nado, John Nham, Eric Ni, Andrew Nystrom, Alicia Parrish, Marie Pellat, Martin Polacek, Alex Polozov, Reiner Pope, Siyuan Qiao, Emily Reif, Bryan Richter, Parker Riley, Alex Castro Ros, Aurko Roy, Brennan Saeta, Rajkumar Samuel, Renee Shelby, Ambrose Slone, Daniel Smilkov, David R. So, Daniel Sohn, Simon Tokumine, Dasha Valter, Vijay Vasudevan, Kiran Vodrahalli, Xuezhi Wang, Pidong Wang, Zirui Wang, Tao Wang, John Wieting, Yuhuai Wu, Kelvin Xu, Yunhan Xu, Linting Xue, Pengcheng Yin, Jiahui Yu, Qiao Zhang, Steven Zheng, Ce Zheng, Weikang Zhou, Denny Zhou, Slav Petrov, and Yonghui Wu. 2023.  Palm 2 technical report.
  * Hendrycks et al. (2020) Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020.  Measuring massive multitask language understanding.  _arXiv preprint arXiv:2009.03300_.
  * Jentzen and Von Wurstemberger (2020) Arnulf Jentzen and Philippe Von Wurstemberger. 2020.  Lower error bounds for the stochastic gradient descent optimization algorithm: Sharp convergence rates for slowly and fast decaying learning rates.  _Journal of Complexity_ , 57:101438.
  * Min et al. (2022) Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. 2022.  Rethinking the role of demonstrations: What makes in-context learning work?  _arXiv preprint arXiv:2202.12837_.
  * Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2020.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _The Journal of Machine Learning Research_ , 21(1):5485–5551.
  * Tay et al. (2022) Yi Tay, Mostafa Dehghani, Vinh Q Tran, Xavier Garcia, Dara Bahri, Tal Schuster, Huaixiu Steven Zheng, Neil Houlsby, and Donald Metzler. 2022.  Unifying language learning paradigms.  _arXiv preprint arXiv:2205.05131_.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In _Advances in Neural Information Processing Systems_ , volume 30.
  * Von Oswald et al. (2023) Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. 2023.  Transformers learn in-context by gradient descent.  In _International Conference on Machine Learning_ , pages 35151–35174. PMLR.
  * Wei et al. (2023) Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. 2023.  Larger language models do in-context learning differently.  _arXiv preprint arXiv:2303.03846_.
  * Xie et al. (2021) Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. 2021.  An explanation of in-context learning as implicit bayesian inference.  _arXiv preprint arXiv:2111.02080_.
  * Zhang et al. (2023) Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. 2023.  Trained transformers learn linear models in-context.  _arXiv preprint arXiv:2306.09927_.

##  Appendix A Proofs

In this section, we provide proofs of the propositions introduced in Section 3
and Section 4.

#### Proposition 1

For a multi-layer LSA satisfying (5) with $\operatorname{\mathbf{w}}^{(0)}=0$,
if its input $\operatorname{\mathbf{Z}}$ is formatted as (4), then its $l$-th
layer output is
$\operatorname{\mathbf{z}}_{j}^{(l)}=(\operatorname{\mathbf{x}}_{j}^{\top},\delta_{j}^{(l)})^{\top}$,
where
$\delta_{j}^{(l)}=y_{j}-\operatorname{\mathbf{w}}^{(l)}\operatorname{\mathbf{x}}_{j}$
and $\operatorname{\mathbf{w}}^{(l)}$ is the weight from the $l$-th step
gradient descents as in (3).

#### Proof:

Plugging in $\operatorname{\mathbf{K}}$, $\operatorname{\mathbf{Q}}$,
$\operatorname{\mathbf{P}}$ and $\operatorname{\mathbf{V}}$ of (5) with
$\operatorname{\mathbf{w}}^{(0)}=0$ into (1), we have

$\displaystyle\begin{pmatrix}\operatorname{\mathbf{x}}_{j}\\\
\delta_{j}^{(l)}\end{pmatrix}=\begin{pmatrix}\operatorname{\mathbf{x}}_{j}\\\
\delta_{j}^{(l-1)}\end{pmatrix}+\frac{\eta}{n}\begin{pmatrix}{\bf 0}_{d\times
d}&{\bf 0}\\\ 0&-1\end{pmatrix}\cdot$
$\displaystyle\;\;\left(\sum_{i=1}^{n}\begin{pmatrix}\operatorname{\mathbf{x}}_{i}\\\
\delta_{i}^{(l-1)}\end{pmatrix}(\operatorname{\mathbf{x}}_{i}^{\top},\delta_{i}^{(l-1)})\begin{pmatrix}\operatorname{\mathbf{I}}_{d\times
d}&{\bf 0}\\\ 0&0\end{pmatrix}\begin{pmatrix}\operatorname{\mathbf{x}}_{j}\\\
\delta_{j}^{(l-1)}\end{pmatrix}\right)$
$\displaystyle\qquad\;\;=\begin{pmatrix}\operatorname{\mathbf{x}}_{j}\\\
\delta_{j}^{(l-1)}\end{pmatrix}-\frac{\eta}{n}\sum_{i=1}^{n}\begin{pmatrix}{\bf
0}\\\
\delta_{i}^{(l-1)}\end{pmatrix}\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j}.$

It is easy to see that $\operatorname{\mathbf{z}}_{j}$ never changes its first
$d$-dimension corresponding to $\operatorname{\mathbf{x}}_{j}$. Therefore, we
can simplify the above equation and focus only on the last coordinate
$\delta_{j}^{(l)}$, where we have

$\displaystyle\delta_{j}^{(l)}=$
$\displaystyle\;\delta_{j}^{(l-1)}-\frac{\eta}{n}\sum_{i=1}^{n}\delta_{i}^{(l-1)}\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j},$
(12)

with $\delta_{j}^{(0)}=y_{j}$. Defining
$\tilde{y}_{j}^{(l)}=y_{j}-\delta_{j}^{(l)}$ and rearranging (12), we obtain
$\tilde{y}_{j}^{(0)}=0$ and

$\displaystyle\tilde{y}_{j}^{(l)}=$
$\displaystyle\;\tilde{y}_{j}^{(l-1)}+\frac{\eta}{n}\sum_{i=1}^{n}(y_{i}-\tilde{y}_{i}^{(l-1)})\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j}.$
(13)

Next we prove
$\tilde{y}_{j}^{(l)}=\operatorname{\mathbf{w}}^{(l)}\operatorname{\mathbf{x}}_{j}$
by induction. Since $\operatorname{\mathbf{w}}^{(0)}=0$, it is clear that
$\tilde{y}_{j}^{(0)}=\operatorname{\mathbf{w}}^{(0)}\operatorname{\mathbf{x}}_{j}=0$
for all $j$.

If
$\tilde{y}_{j}^{(l-1)}=\operatorname{\mathbf{w}}^{(l-1)}\operatorname{\mathbf{x}}_{j}$
for all $j$, then

$\displaystyle\tilde{y}_{j}^{(l)}=$
$\displaystyle\;\tilde{y}_{j}^{(l-1)}+\frac{\eta}{n}\sum_{i=1}^{n}(y_{i}-\tilde{y}_{i}^{(l-1)})\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j}$
$\displaystyle=$
$\displaystyle\operatorname{\mathbf{w}}^{(l-1)}\operatorname{\mathbf{x}}_{j}+\frac{\eta}{n}\sum_{i=1}^{n}(y_{i}-\operatorname{\mathbf{w}}^{(l-1)}\operatorname{\mathbf{x}}_{i})\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j}$
$\displaystyle=$
$\displaystyle\left(\operatorname{\mathbf{w}}^{(l-1)}+\sum_{i=1}^{n}(y_{i}-\operatorname{\mathbf{w}}^{(l-1)}\operatorname{\mathbf{x}}_{i})\operatorname{\mathbf{x}}_{i}^{\top}\right)\operatorname{\mathbf{x}}_{j}$
$\displaystyle=$
$\displaystyle\operatorname{\mathbf{w}}^{(l)}\operatorname{\mathbf{x}}_{j}.$

$\square$

#### Proposition 2

For a multi-layer causalLM-LSA satisfying (5) with
$\operatorname{\mathbf{w}}^{(0)}=0$, if its input $\operatorname{\mathbf{Z}}$
is formatted as (4), then its $l$-th layer output is
$\operatorname{\mathbf{z}}_{j}^{(l)}=(\operatorname{\mathbf{x}}_{j}^{\top},\delta_{j}^{(l)})^{\top}$,
where
$\delta_{j}^{(l)}=y_{j}-\operatorname{\mathbf{w}}_{j}^{(l)}\operatorname{\mathbf{x}}_{j}$
and $\operatorname{\mathbf{w}}_{j}^{(l)}$ follow (10).

#### Proof:

Plugging in $\operatorname{\mathbf{K}}$, $\operatorname{\mathbf{Q}}$,
$\operatorname{\mathbf{P}}$ and $\operatorname{\mathbf{V}}$ of (5) with
$\operatorname{\mathbf{w}}^{(0)}=0$ into (2), we have

$\displaystyle\delta_{j}^{(l)}=$
$\displaystyle\;\delta_{j}^{(l-1)}-\frac{\eta}{n}\sum_{i=1}^{j}\delta_{i}^{(l-1)}\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j}$
$\displaystyle\tilde{y}_{j}^{(l)}=$
$\displaystyle\;\tilde{y}_{j}^{(l-1)}+\frac{\eta}{n}\sum_{i=1}^{j}(y_{i}-\tilde{y}_{i}^{(l-1)})\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j}$

with $\tilde{y}_{j}^{(l)}=y_{j}-\delta_{j}^{(l)}$. Next we prove
$\tilde{y}_{j}^{(l)}=\operatorname{\mathbf{w}}_{j}^{(l)}\operatorname{\mathbf{x}}_{j}$
by induction. Since $\operatorname{\mathbf{w}}_{j}^{(0)}=0$, it is clear that
$\tilde{y}_{j}^{(0)}=\operatorname{\mathbf{w}}_{j}^{(0)}\operatorname{\mathbf{x}}_{j}=0$
for all $j$.

If
$\tilde{y}_{j}^{(l-1)}=\operatorname{\mathbf{w}}_{j}^{(l-1)}\operatorname{\mathbf{x}}_{j}$
for all $j$, then

$\displaystyle\tilde{y}_{j}^{(l)}=$
$\displaystyle\;\tilde{y}_{j}^{(l-1)}+\frac{\eta}{j}\sum_{i=1}^{n}(y_{i}-\tilde{y}_{i}^{(l-1)})\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j}$
$\displaystyle=$
$\displaystyle\operatorname{\mathbf{w}}_{j}^{(l-1)}\operatorname{\mathbf{x}}_{j}+\frac{\eta}{n}\sum_{i=1}^{j}(y_{i}-\operatorname{\mathbf{w}}_{i}^{(l-1)}\operatorname{\mathbf{x}}_{i})\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j}$
$\displaystyle=$
$\displaystyle\left(\operatorname{\mathbf{w}}_{j}^{(l-1)}+\sum_{i=1}^{n}(y_{i}-\operatorname{\mathbf{w}}_{i}^{(l-1)}\operatorname{\mathbf{x}}_{i})\operatorname{\mathbf{x}}_{i}^{\top}\right)\operatorname{\mathbf{x}}_{j}$
$\displaystyle=$
$\displaystyle\operatorname{\mathbf{w}}_{j}^{(l)}\operatorname{\mathbf{x}}_{j}.$

$\square$

#### Proposition 3

If $\operatorname{\mathbf{w}}^{(l)}$ follows the iterative updates of (3),
then there exists a stationary point $\operatorname{\mathbf{w}}^{*}$ with
coefficients satisfying:

$\displaystyle\operatorname{\mathbf{y}}\operatorname{\mathbf{x}}^{\top}=\operatorname{\mathbf{w}}^{*}\operatorname{\mathbf{X}}\operatorname{\mathbf{X}}^{\top},$

where $\operatorname{\mathbf{y}}=(y_{1},\ldots,y_{n})$ and
$\operatorname{\mathbf{X}}=(\operatorname{\mathbf{x}}_{1},\ldots,\operatorname{\mathbf{x}}_{n})$.
Furthermore, the iterative weights $\operatorname{\mathbf{w}}^{(l)}$ converges
to the stationary point $\operatorname{\mathbf{w}}^{*}$ with linear rate of
convergence:

$\displaystyle\operatorname{\mathbf{w}}^{(l)}-\operatorname{\mathbf{w}}^{*}=(\operatorname{\mathbf{w}}^{(l-1)}-\operatorname{\mathbf{w}}^{*})(\operatorname{\mathbf{I}}-\frac{\eta}{n}\operatorname{\mathbf{X}}\operatorname{\mathbf{X}}^{\top}).$

#### Proof:

From (3), we have

$\displaystyle\operatorname{\mathbf{w}}^{(l)}=$
$\displaystyle\;\operatorname{\mathbf{w}}^{(l-1)}+\frac{\eta}{n}\underbrace{\sum_{i=1}^{n}(y_{i}-\operatorname{\mathbf{w}}^{(l-1)}\operatorname{\mathbf{x}}_{i})\operatorname{\mathbf{x}}_{i}^{\top}}_{(*)}.$

The stationary point must satisfy $(*)=0$. Written in vectorized form, we have

$\displaystyle\operatorname{\mathbf{y}}\operatorname{\mathbf{X}}^{\top}=\operatorname{\mathbf{w}}^{*}\operatorname{\mathbf{X}}\operatorname{\mathbf{X}}^{\top}.$
(14)

Now plugging (14) back to (3), we have

$\displaystyle\operatorname{\mathbf{w}}^{(l)}=\operatorname{\mathbf{w}}^{(l-1)}+\frac{\eta}{n}\left(\operatorname{\mathbf{w}}^{*}\operatorname{\mathbf{X}}\operatorname{\mathbf{X}}^{\top}-\operatorname{\mathbf{a}}\operatorname{\mathbf{w}}^{(l-1)}\operatorname{\mathbf{X}}\operatorname{\mathbf{X}}^{\top}\right),$

which can be reorganized to

$\displaystyle\operatorname{\mathbf{w}}^{(l)}-\operatorname{\mathbf{w}}^{*}=(\operatorname{\mathbf{w}}^{(l-1)}-\operatorname{\mathbf{w}}^{*})(\operatorname{\mathbf{I}}-\frac{\eta}{n}\operatorname{\mathbf{X}}\operatorname{\mathbf{X}}^{\top}).$

$\square$

#### Proposition 4

If
$\operatorname{\mathbf{w}}_{j}^{(l)}=\sum_{i=1}^{j}a_{i,j}^{(l)}\operatorname{\mathbf{x}}_{i}^{\top}$
follows the iterative updates of (10), then

$\displaystyle a_{i,j}^{(l)}=a_{i,i}^{(l)}\equiv a_{i}^{(l)}\;\;\;\;\forall
j\geq i,$

and there exists stationary points
$\operatorname{\mathbf{w}}_{j}^{*}=\sum_{i=1}^{j}a_{i}^{*}\operatorname{\mathbf{x}}_{i}^{\top}$
(for $j\in 1,\ldots,n$) with coefficients from
$\operatorname{\mathbf{a}}^{*}=(a^{*}_{1},\ldots,a^{*}_{n})$ that satisfy
$\operatorname{\mathbf{y}}=\operatorname{\mathbf{a}}^{*}\operatorname{\mathbf{T}}$,
where

$\displaystyle\operatorname{\mathbf{T}}=\begin{pmatrix}\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{1}&\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{2}&\cdots&\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{n}\\\
0&\operatorname{\mathbf{x}}_{2}^{\top}\operatorname{\mathbf{x}}_{2}&\cdots&\operatorname{\mathbf{x}}_{2}^{\top}\operatorname{\mathbf{x}}_{n}\\\
\vdots&\vdots&\ddots&\vdots\\\
0&0&\cdots&\operatorname{\mathbf{x}}_{n}^{\top}\operatorname{\mathbf{x}}_{n}\end{pmatrix}.$

Furthermore, the coefficients $\operatorname{\mathbf{a}}^{(l)}$ converges to
the stationary point $\operatorname{\mathbf{a}}^{*}$ with linear rate of
convergence:

$\displaystyle\operatorname{\mathbf{a}}^{(l)}-\operatorname{\mathbf{a}}^{*}=(\operatorname{\mathbf{a}}^{(l-1)}-\operatorname{\mathbf{a}}^{*})(\operatorname{\mathbf{I}}-\frac{\eta}{n}\operatorname{\mathbf{T}}).$

#### Proof:

First notice that according to (10), we have

$\displaystyle\operatorname{\mathbf{w}}_{j}^{(l)}=$
$\displaystyle\;\operatorname{\mathbf{w}}_{j}^{(l-1)}+\frac{\eta}{n}\sum_{i=1}^{j}(y_{i}-\operatorname{\mathbf{w}}_{i}^{(l-1)}\operatorname{\mathbf{x}}_{i})\operatorname{\mathbf{x}}_{i}^{\top}$
$\displaystyle=$
$\displaystyle\sum_{i=1}^{j}\left(a_{i,j}^{(l-1)}+\frac{\eta}{n}(y_{i}-\operatorname{\mathbf{w}}_{i}^{(l-1)}\operatorname{\mathbf{x}}_{i})\right)\operatorname{\mathbf{x}}_{i}^{\top}$

or

$\displaystyle
a_{i,j}^{(l)}=a_{i,j}^{(l-1)}+\frac{\eta}{n}(y_{i}-\operatorname{\mathbf{w}}_{i}^{(l-1)}\operatorname{\mathbf{x}}_{i})\;\;\;\;\forall
j\geq i.$

Since $a_{i,j}^{(0)}=0$, and the above update is the same for any $j$ given
any $i$, then it is obvious by induction that

$\displaystyle a_{i,j}^{(l)}=a_{i,i}^{(l)}\equiv a_{i}^{(l)}\;\;\;\;\forall
j\geq i.$

Therefore, we can simplify
$\operatorname{\mathbf{w}}_{j}^{(l)}=\sum_{i=1}^{j}a_{i}^{(l)}\operatorname{\mathbf{x}}_{i}^{\top}$.

Now plugging into (10), we have

$\displaystyle\sum_{i=1}^{j}a_{i}^{(l)}\operatorname{\mathbf{x}}_{i}^{\top}$
$\displaystyle=$
$\displaystyle\sum_{i=1}^{j}a_{i}^{(l-1)}\operatorname{\mathbf{x}}_{i}^{\top}+\frac{\eta}{n}\sum_{i=1}^{j}(y_{i}-\sum_{k=1}^{i}a_{k}^{(l-1)}\operatorname{\mathbf{x}}_{k}^{\top}\operatorname{\mathbf{x}}_{i})\operatorname{\mathbf{x}}_{i}^{\top}$
$\displaystyle=$
$\displaystyle\sum_{i=1}^{j}\left(a_{i}^{(l-1)}+y_{i}-\frac{\eta}{n}\left(\sum_{k=1}^{i}a_{k}^{(l-1)}\operatorname{\mathbf{x}}_{k}^{\top}\operatorname{\mathbf{x}}_{i}\right)\right)\operatorname{\mathbf{x}}_{i}^{\top},$

which is equivalent to

$\displaystyle
a_{i}^{(l)}=a_{i}^{(l-1)}+\frac{\eta}{n}\underbrace{\left(y_{i}-\sum_{k=1}^{i}a_{k}^{(l-1)}\operatorname{\mathbf{x}}_{k}^{\top}\operatorname{\mathbf{x}}_{i}\right)}_{(*)}.$
(15)

The stationary point satisfy $(*)=0$, which gives

$\displaystyle y_{1}=$ $\displaystyle
a_{1}^{*}\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{1}$
$\displaystyle y_{2}=$ $\displaystyle
a_{1}^{*}\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{2}+a_{2}^{*}\operatorname{\mathbf{x}}_{2}^{\top}\operatorname{\mathbf{x}}_{2}$
$\displaystyle\ldots$ $\displaystyle y_{n}=$ $\displaystyle
a_{1}^{*}\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{n}+\ldots+a_{n}^{*}\operatorname{\mathbf{x}}_{n}^{\top}\operatorname{\mathbf{x}}_{n},$

or in the vectorized form
$\operatorname{\mathbf{y}}=\operatorname{\mathbf{a}}^{*}\operatorname{\mathbf{T}}$,
where

$\displaystyle\operatorname{\mathbf{T}}=\begin{pmatrix}\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{1}&\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{2}&\cdots&\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{n}\\\
0&\operatorname{\mathbf{x}}_{2}^{\top}\operatorname{\mathbf{x}}_{2}&\cdots&\operatorname{\mathbf{x}}_{2}^{\top}\operatorname{\mathbf{x}}_{n}\\\
\vdots&\vdots&\ddots&\vdots\\\
0&0&\cdots&\operatorname{\mathbf{x}}_{n}^{\top}\operatorname{\mathbf{x}}_{n}\end{pmatrix}.$

Now plugging in
$\operatorname{\mathbf{y}}=\operatorname{\mathbf{a}}^{*}\operatorname{\mathbf{T}}$
back to (15) and vectorize it, yields

$\displaystyle\operatorname{\mathbf{a}}^{(l)}=\operatorname{\mathbf{a}}^{(l-1)}+\frac{\eta}{n}\left(\operatorname{\mathbf{a}}^{*}\operatorname{\mathbf{T}}-\operatorname{\mathbf{a}}\operatorname{\mathbf{T}}\right),$

which can be reorganized to

$\displaystyle\operatorname{\mathbf{a}}^{(l)}-\operatorname{\mathbf{a}}^{*}=(\operatorname{\mathbf{a}}^{(l-1)}-\operatorname{\mathbf{a}}^{*})(\operatorname{\mathbf{I}}-\frac{\eta}{n}\operatorname{\mathbf{T}}).$

$\square$

#### Proposition 5

Assuming that $\operatorname{\mathbf{w}}_{j}^{*}$ is the stationary points
obtained in Proposition 4, then

$\displaystyle\operatorname{\mathbf{w}}_{j+1}^{*}=\operatorname{\mathbf{w}}_{j}^{*}-\frac{1}{\|\operatorname{\mathbf{x}}_{j+1}\|_{2}^{2}}\nabla_{\operatorname{\mathbf{w}}_{j}^{*}}L_{j}(\operatorname{\mathbf{w}}_{j}^{*}).$

#### Proof:

Recall the online learning system with a sequence of data-label pairs
$(\operatorname{\mathbf{x}}_{j},y_{j})$ has the following online loss and its
gradient at the $j$-th step,

$\displaystyle L_{j}(\operatorname{\mathbf{w}}_{j})$
$\displaystyle=\frac{1}{2}(\operatorname{\mathbf{w}}_{j}\operatorname{\mathbf{x}}_{j+1}-y_{j+1})^{2},$
$\displaystyle\nabla_{\operatorname{\mathbf{w}}_{j}}L_{j}(\operatorname{\mathbf{w}}_{j})$
$\displaystyle=(\operatorname{\mathbf{w}}_{j}\operatorname{\mathbf{x}}_{j+1}-y_{j+1})\operatorname{\mathbf{x}}_{j+1}^{\top}.$

According to Proposition 4, we have
$\operatorname{\mathbf{y}}=\operatorname{\mathbf{a}}^{*}\operatorname{\mathbf{T}}$,
which gives

$\displaystyle y_{j+1}=$ $\displaystyle
a_{1}^{*}\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{j+1}+\ldots+a_{j}^{*}\operatorname{\mathbf{x}}_{j}^{\top}\operatorname{\mathbf{x}}_{j+1}$
$\displaystyle\;+a_{j+1}^{*}\operatorname{\mathbf{x}}_{j+1}^{\top}\operatorname{\mathbf{x}}_{j+1}$
$\displaystyle=$
$\displaystyle\operatorname{\mathbf{w}}_{j}^{*}\operatorname{\mathbf{x}}_{j+1}+a_{j+1}^{*}\operatorname{\mathbf{x}}_{j+1}^{\top}\operatorname{\mathbf{x}}_{j+1}$
(16)

where the last equation is due to
$\operatorname{\mathbf{w}}_{j}^{*}=\sum_{i=1}^{j}a_{i}^{*}\operatorname{\mathbf{x}}_{i}^{\top}$.

Since
$\operatorname{\mathbf{w}}_{j}^{*}=\sum_{i=1}^{j}a_{i}^{*}\operatorname{\mathbf{x}}_{i}^{\top}$,
we have

$\displaystyle\operatorname{\mathbf{w}}_{j+1}^{*}=$
$\displaystyle\operatorname{\mathbf{w}}_{j}^{*}+a_{j+1}^{*}\operatorname{\mathbf{x}}_{j+1}^{\top}$
$\displaystyle=$
$\displaystyle\operatorname{\mathbf{w}}_{j}^{*}+\frac{1}{\|\operatorname{\mathbf{x}}_{j+1}\|_{2}^{2}}\left(a_{j+1}^{*}\operatorname{\mathbf{x}}_{j+1}^{\top}\operatorname{\mathbf{x}}_{j+1}\right)\operatorname{\mathbf{x}}_{j+1}^{\top}$
$\displaystyle=$
$\displaystyle\operatorname{\mathbf{w}}_{j}^{*}-\frac{1}{\|\operatorname{\mathbf{x}}_{j+1}\|_{2}^{2}}(\operatorname{\mathbf{w}}_{j}^{*}\operatorname{\mathbf{x}}_{j+1}-y_{j+1})\operatorname{\mathbf{x}}_{j+1}^{\top}$
$\displaystyle=$
$\displaystyle\operatorname{\mathbf{w}}_{j}^{*}-\frac{1}{\|\operatorname{\mathbf{x}}_{j+1}\|_{2}^{2}}\nabla_{\operatorname{\mathbf{w}}_{j}^{*}}L_{j}(\operatorname{\mathbf{w}}_{j}^{*})$

where the third equation is because of (16). $\square$

##  Appendix B Multi-layer LSA construction with non-zero w(0)

In this section, we introduce the proposition that connects a multi-layer LSA
following the construction of (5) but with non-zero
$\operatorname{\mathbf{w}}^{(0)}$ and the multi-step gradient descents of
linear regression.

###### Proposition 6

For a multi-layer LSA satisfying the construction (5), if its input
$\operatorname{\mathbf{Z}}$ is formatted as (4), then its $l$-th layer output
is
$\operatorname{\mathbf{z}}_{j}^{(l)}=(\operatorname{\mathbf{x}}_{j}^{\top},\delta_{j}^{(l)})^{\top}$,
where
$\delta_{j}^{(l)}=y_{j}-(\operatorname{\mathbf{w}}^{(l)}-\operatorname{\mathbf{w}}^{(0)})\operatorname{\mathbf{x}}_{j}$
and $\operatorname{\mathbf{w}}^{(l)}$ is the $l$-th updated weight from the
gradient descents update rule in (3).

#### Proof:

Plugging in $\operatorname{\mathbf{K}}$, $\operatorname{\mathbf{Q}}$,
$\operatorname{\mathbf{P}}$ and $\operatorname{\mathbf{V}}$ of (5) into (1),
we have

$\displaystyle\delta_{j}^{(l)}=$
$\displaystyle\;\delta_{j}^{(l-1)}-\frac{\eta}{n}\sum_{i=1}^{n}\left(\delta_{i}^{(l-1)}-\operatorname{\mathbf{w}}^{(0)}\operatorname{\mathbf{x}}_{i}\right)\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j},$
(17)

with $\delta_{j}^{(0)}=y_{j}$. Defining
$\tilde{y}_{j}^{(l)}=y_{j}-\delta_{j}^{(l)}+\operatorname{\mathbf{w}}^{(0)}\operatorname{\mathbf{x}}_{j}$
and rearranging the (17), we obtain $\tilde{y}_{j}^{(0)}=0$ and

$\displaystyle\tilde{y}_{j}^{(l)}=$
$\displaystyle\;\tilde{y}_{j}^{(l-1)}+\frac{\eta}{n}\sum_{i=1}^{n}(y_{i}-\tilde{y}_{i}^{(l-1)})\operatorname{\mathbf{x}}_{i}^{\top}\operatorname{\mathbf{x}}_{j}.$

Then it is easy to prove
$\tilde{y}_{j}^{(l)}=\operatorname{\mathbf{w}}^{(l)}\operatorname{\mathbf{x}}_{j}$
by induction, similar to the proof of Proposition 1. $\square$

##  Appendix C CausalLM with attention-length-based coefficients

Since there are $j$ terms in the summation of (10), another reasonable update
for causalLM would be

$\displaystyle\operatorname{\mathbf{w}}_{j}^{(l)}=$
$\displaystyle\;\operatorname{\mathbf{w}}_{j}^{(l-1)}+\frac{\eta}{j}\sum_{i=1}^{j}(y_{i}-\operatorname{\mathbf{w}}_{i}^{(l-1)}\operatorname{\mathbf{x}}_{i})\operatorname{\mathbf{x}}_{i}^{\top},$
(18)

which we call causalLM2. For causalLM2, we have the following proposition.

###### Proposition 7

If
$\operatorname{\mathbf{w}}_{j}^{(l)}=\sum_{i=1}^{j}a_{i,j}^{(l)}\operatorname{\mathbf{x}}_{i}^{\top}$
follows the iterative updates of (18), then

$\displaystyle a_{i,j}^{(l)}\equiv\frac{1}{j}a_{i}^{(l)}\;\;\;\;\forall j\geq
i,$

and there exists stationary points
$\operatorname{\mathbf{w}}_{j}^{*}=\frac{1}{j}\sum_{i=1}^{j}a_{i}^{*}\operatorname{\mathbf{x}}_{i}^{\top}$
(for $j\in 1,\ldots,n$) with coefficients from
$\operatorname{\mathbf{a}}^{*}=(a^{*}_{1},\ldots,a^{*}_{n})$ that satisfy
$\operatorname{\mathbf{y}}=\operatorname{\mathbf{a}}^{*}\operatorname{\mathbf{S}}$,
where

$\displaystyle\operatorname{\mathbf{S}}=\begin{pmatrix}\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{1}&\frac{1}{2}\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{2}&\cdots&\frac{1}{n}\operatorname{\mathbf{x}}_{1}^{\top}\operatorname{\mathbf{x}}_{n}\\\
0&\frac{1}{2}\operatorname{\mathbf{x}}_{2}^{\top}\operatorname{\mathbf{x}}_{2}&\cdots&\frac{1}{n}\operatorname{\mathbf{x}}_{2}^{\top}\operatorname{\mathbf{x}}_{n}\\\
\vdots&\vdots&\ddots&\vdots\\\
0&0&\cdots&\frac{1}{n}\operatorname{\mathbf{x}}_{n}^{\top}\operatorname{\mathbf{x}}_{n}\end{pmatrix}.$

Furthermore, the coefficients $\operatorname{\mathbf{a}}^{(l)}$ converges to
the stationary point $\operatorname{\mathbf{a}}^{*}$ with the following rate
of convergence:

$\displaystyle\operatorname{\mathbf{a}}^{(l)}-\operatorname{\mathbf{a}}^{*}=(\operatorname{\mathbf{a}}^{(l-1)}-\operatorname{\mathbf{a}}^{*})(\operatorname{\mathbf{I}}-\eta\operatorname{\mathbf{S}}).$

The condition number $\kappa(\operatorname{\mathbf{S}})$ is about $n/2$
greater than the one of $\kappa(\operatorname{\mathbf{T}})$, which makes
causalLM2 converge much slower than causalLM.

One can also prove that the stationary point of causalLM2 corresponds to the
following online system with online loss and gradient at the $j$-th step,

$\displaystyle L_{j}(\tilde{\operatorname{\mathbf{w}}}_{j})$
$\displaystyle=\frac{1}{2}(\tilde{\operatorname{\mathbf{w}}}_{j}\operatorname{\mathbf{x}}_{j+1}-y_{j+1})^{2},$
$\displaystyle\nabla_{\tilde{\operatorname{\mathbf{w}}}_{j}}L_{j}(\tilde{\operatorname{\mathbf{w}}}_{j})$
$\displaystyle=(\tilde{\operatorname{\mathbf{w}}}_{j}\operatorname{\mathbf{x}}_{j+1}-y_{j+1})\operatorname{\mathbf{x}}_{j+1}^{\top},$

where
$\tilde{\operatorname{\mathbf{w}}}=\frac{j}{j+1}\operatorname{\mathbf{w}}$.

###### Proposition 8

Assuming that $\operatorname{\mathbf{w}}_{j}^{*}$ is the stationary points
obtained in Proposition 4, then

$\displaystyle\operatorname{\mathbf{w}}_{j+1}^{*}=\tilde{\operatorname{\mathbf{w}}}_{j}^{*}-\frac{1}{\|\operatorname{\mathbf{x}}_{j+1}\|_{2}^{2}}\nabla_{\tilde{\operatorname{\mathbf{w}}}_{j}^{*}}L_{j}(\tilde{\operatorname{\mathbf{w}}}_{j}^{*}).$

Since the step does not have $j^{-\nu}$ ($\nu>0$) decay, such online system is
not guaranteed to converge, therefore suffers the same problem as the original
causalLM in Section 3.2.

In Figure 4, we plot the query MSE error of the stationary points of
causalLM2-ICL with increasing number of in-context examples. We can see that
the online system corresponding to causalLM2-ICL converges even slower than
the ones of causalLM-ICL in Figure 2 right.

Figure 4: The test error on the stationary point of the causalLM2-ICL with up
to 300 in-context examples.

##  Appendix D Additional experimental details and results

###  D.1 Experiment settings for Section 5.1

In order to directly verify the theorem, we used the constructed LSA-based
transformer, with
$\operatorname{\mathbf{K}}=\operatorname{\mathbf{Q}}=\begin{pmatrix}\operatorname{\mathbf{I}}_{d\times
d}&{\bf 0}\\\ 0&0\end{pmatrix}$,
$\operatorname{\mathbf{V}}=\begin{pmatrix}{\bf 0}_{d\times d}&{\bf 0}\\\
0&-1\end{pmatrix}$ and
$\operatorname{\mathbf{P}}=\frac{\eta}{n}\operatorname{\mathbf{I}}$. Although
not a trained transformer, it was recently proved in Zhang et al. (2023) that
a randomly initialized LSA-transformer does converge to such a construction.
In addition, we did an ablation test of
$\eta=\left\\{0.1,0.2,0.4,0.8,1.6,3.2\right\\}$ and chose $\eta=1.6$ as it
converges the fastest without any divergence problems.

We randomly generated 64 sequences for ICL evaluation. For each sequence, we
put the first 40 examples as the in-context examples and the last 200 examples
as the query examples. The queries are independent of each other without
attention. See Figure 5 for an illustration of the transformer attention mask.
Such multi-query design is for training efficiency purpose only and is
equivalent to 200 sequences with the same $\operatorname{\mathbf{w}}$ and
input examples $\operatorname{\mathbf{x}}_{i}$, but different one query per
sequence.

Figure 5: The illustration of the attention mask. Green arrows represent the
attentions between in-context examples. The dashed arrows only applies for
prefixLM. Red arrows represent the attentions from queries to in-context
examples. The query examples should not attend to themselves because the
inputs do not contain labels.

###  D.2 Experiment settings for Section 5.2

In order to verify that our theorems can be qualitatively applied beyond LSA
and linear regression, we conducted a few experiments over various synthetic
tasks using regular transformers. In particular, we based our code from the
repository of Akyürek et al. (2022)¶¶¶https://github.com/google-
research/google-research/tree/master/incontext and applied their default
training hyperparameters of the code. We used a transformer of 24 layers with
128 hidden units and 2 heads. The FFN intermediate size is $4\times 128=512$.
The learning schedule is based on cosine decay with base learning rate 1e-4,
for 100 epochs.