CausalLM is not optimal for in-context learning
Authors: Nan Ding
Tomer Levinboim
Jialin Wu
Sebastian Goodman
Radu Soricut
Venue: ICLR 2024 poster (2024)
Keywords: in-context learning
causalLM
prefixLM
Primary area: transfer learning, meta learning, and lifelong learning

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 use 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.

Published as a conference paper at ICLR 2024
CAUSALLM IS NOT OPTIMAL FOR IN-CONTEXT LEARN- ING
Nan Ding
Tomer Levinboim Jialin Wu Sebastian Goodman
Radu Soricut
Google Research (di Ignan, tomer1, , jia inwu, seabass, rsoricut, @goog le. com
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 use auto-regressive attention that prohibits in-context samples 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 certain parameter construction. Our analysis shows that both LM types converge to their stationary points at linear rate, but that while prefixLM converges 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 all settings.
INTRODUCTION
censformer-bbssed models (Vaswani et al., 2017) have become the default foundational model for various machine learning applications such as natural language processing (Devlin al., 2018 Brown et al. 2020 Chowdhery et al. (002) and computer vision (Dosovitskiy et al. (020). Beyond their traditional usage in machine learning applications, it has recently been discovered that pretraining large transformers a vast amounts of data leads them to develop striking ability referred to as in-context learning (ICL) (Brown et al., 2020). Specifically, once such pretraining is complete, these models able to solve new tasks at inference time (without changing their parameters) by simply ingesting a short sequence (prefix) of labeled examples from task and then computing a prediction for query example.
The ICL capability was first demonstrated by GPT-3 (Brown et al. 2020), where a causalIM (a Transformer decoder with auto-regressive attention masks) was used as the main model architecture. However, follow work empirically found that restricting the auto-regressive masks on the entire sequence is too prohibitive and therefore proposed the so-called prefixLM (Raffel al. 2020b Tay et al. 2022) which allows full attention within the prefix tokens. Moreover, the latest models (such as P2LM2 Google al. 2023)) adopt a mixture of different LM objectives during pretraining to achieved state-of-art performance across a diverse set of tasks and capabilities.
However, beyond the few empirical results in those and related papers, there is yet no theoretical ex- planation that accounts for the different ICL behavior of prefixLM and causalL.. Indeed, theoretical studies of ICL are difficult due to the complicated non-linearity of the (ordinary) transformer archi- tecture. 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 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 randomly initialized I. SAA-rrnsformer indeed converges such construction during training.

---

Published as a conference paper at ICLR 2024
In this paper, continue the theoretical line of work above by investigating the convergence properties of ICL for both prefixLM and causalLM multi-layer -SA-trrasformmrss in a linear regression setting. We summarizes contributions as follows:
We first present clear, formal proof that establishes the relationship between multi-layer LSA and multi-step gradient descents linear regression.
We then show that both causalLM and prefixLM based multi-layer ASS-rransformrrs converge to their respective stationary points with linear rates of convergence. We prove that the stationary point of prefixLMM corresponds the optimal least square solution of the linear regression problem, while the stationary points of causalLM correspond to the weights obtained along the iterations of online gradient descent 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, which indicates that causalLM is not optimal for in-context learning.
Finally, we verify the above theoretical insights by conducting experiments with LSA- transformers well ordinary softmax attention based transformers on various synthetic tasks including linear and non-linear regression, and multiclass classifications. We also compare causalLM and prefixLM ICL based on LLMs including T5 (Roberts al. 2022) and PaLM2 (Google et al., 2023). as well as the multimodal model PaLI-X (Chen et al. )023 Our experimental results support our theoretical findings and consistently show the superiority of PrefixLM over causalLM on ICL for such settings.
2 BACKGROUND
We begin by reviewing a few types of transformer attention and in-context learning (ICL), as well as specific transformer construction for linear regression ICL by (Von Oswald al. 2023 which our theories will be based on. The discussions of other related work are deferred to Appendix A
2.1 TRANSFORMERS, SSA, LSA, CAUSALLM, AND PREFIXLM
Given a sequence of input vectors Z = (Zi,..... "2,,...,2y,, kn), the output of standard Softmax Self-Attentin (SSA) layer is
Zj - Zj + PV Z softmax K Qzy),
where P, V, K, Q respectively corresponds to the output projection, value transformation, key transformation and query transformation.
Since the softmax attention of standard transformers non-linear, its theoretical analysis becomes complicated even for 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,
Zj + 23 + (VZ K Q2,) = PV
24 ( K KTKQQ kz,z)))
Furthermore, since each input zj can attend to all positions [1...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
Zy Z + PV Σ
2 z K Qz.).
which restricts each token z, to attend only to previous positions (and itself) from {1 (1...]}. 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 full attention based encoder and an auto-regressive atten- tion based decoder. However, prominent NLP research has often chosen either encoder-only (e.g.
2

---

Published as a conference paper at ICLR 2024
BERT (Devlin al. 2018) or decoder-only (e.g. GPT (Brown al. 2020), PaLM (Chowdhery et al. ) models according to the task at hand. This is partially for the purpose halving the parameter sizes.
Another version of attention, between full and auto-regressive, followed from the observation that some tasks can benefit from prefix sequence such as context prompt. That is, the input sequence Z composed of n' prefix tokens (Z),..., Znd) configured for the task, while the tokens (1,'+1+- -........ Z,n) represent the sample. Specifically, prefiXLM (Raffel et al., 2020b) suggests the following attention (in its LSA version):
max(jnn))
2j 2 PV
Σ Zi (z) K (K])Z)).
where max(j, n') ensures each prefix token 2j with n' can attend to all prefix tokens.
2.2
IN-CONTEXT LEARNING
formal framework of in-context learning has been described in various existing literature such as (Garg et al. 2022 Zhang et al. 2023). Here, 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 inference time. That is, at inference time, a model is fed with a sequence consisting of input-label pairs and query input (x1, y.,.... Xny Un, Xquery) and its goal is predict the label Mquery of Xquery using the context examples (x1, 01.... Xa, Yn) (specifically, without changing the model parameters).
2.3 LINEAR REGRESSION IN-CONTEXT LEARNERS
Linear regression is classical machine learning problem. Given set of input-label pairs (xi, v:). the goal is to find an optimal weight vector w that minimizes the 12-loss:
I(w)=
2n Σ
1
x; --l12
The gradient of the loss V step size 17 follows the update rule:
Di=1 (W x;
Vii x and a gradient descent algorithm with
n
1 W (w(1-
72 Σ
1/4 (1-11xx,)) xT]
Using linear regression lens to study in-context learning first proposed in (Garg 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, theoretical study was provided. More recently, and most relevant to our work, (Von Oswald et al.) 2023) proposed a succinct construction that demonstrates how single LSA layer can effectively implement single gradient descent step. According to their setup the input is formulated as
Z
2 (0)), where 2,
X,
y/
and the parameter matrices of [ are set as:
K = Q Taxd
0
Odxe w(()) (0)
0
V
P = I,
(5)
where w0( is initial weight vector. (Von Oswald et al. results in an update of their so-called transformed target yj + Uj + 17 (V L) x,. and that this target update is equivalent the performed by single-step gradient descent of linear regression.
2023 then showed that this configuration
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 unnatura to work with. Moreover, their extension from single-layer to multi-layer L.A is unfortunately unclear.
"The traditional MI. formulation updates the weight vector the model prediction, while the groundtruth target remains fixed.
3

---

Published as a conference paper at ICLR 2024
3 MULTI-LAYER IN-CONTEXT LEARNER
In this section, we provide a formal proof that a multi-layer LSA under the construction of (Von Os progresses identically multi-step gradient descent.
wald al. 2023
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 K, Q. P and V in (5) Note that we keep w(") = 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 w in Appendix C
Proposition 1 For multi-layer LSA satisfying the construction ( and with w()) if its input Z is formatted as (4) then its I-th layer output 2 (x) 8(1) where 5(1) = y) (!) xj and
w(O is the 1-th updated weight from the gradient descents update rule in (3)
Proof Sketch: Plugging in K, Q, P and V of with w0) 0 and (1) = (x] .0 into 1 we
obtain that for all l 0,
x Xj.
Since 2, never changes its first d-dimension corresponding Xj, can simplify it and focus only on (0) which is the last output coordinate of the j-th LSA-layer,
6
6/-1)
7
x Xy.
(6)
n2
with ///0= = My- Defining Dj
(1)
= yj and rearranging 6 we obtain 1/ 0 and V/ >> O:
(1)
5/ y
11 1-11 1)
n y1 v
x Xj.
Finally, using 7 and the fact that (0) = w (0) x,, it can be proved by induction that V7 07)
w(O xj. A complete proof is provided in Appendix B
To summarize, the newly introduced variable 5/) is exactly the prediction of the 1-th gradient descent weights w()) for xj and 0 is the difference between the true label Y/ and the predicted (!))
Therefore, (()) serves as a bridge to connect the LSA output 5 and the GD weight (!)
So far, have dealt with the behavior of L.AA layers with full attention. In what follows, move on to the practical setting of in-context learning, where the input contains not only n in-context
(0) (training) examples in the format of 44) but also additional (test) query Ziuery (x (upryy T.
In particular, we will focus on the two common ICL variants: prefixI LM and causalI.M MM, each with a different type of attention.
3.1 PREFIXLM ICL
A PM ICL treats the in-context examples Z as the prefix and uses full attention on the first n positions, so that they can each freely attend each other. The last query vector Equery can also attend to any example in Z, but cannot attend to itself] As result, the updates of the PreIiXLM-ICL under the same construction follow (6). with the outputs of the 1-th layer being,
500 yj 03 yj w()) Xj.
and
(0) query Jadery w() Xguerys
'This is because the query does not contain meaningful label. Attending to itself would cause it to include its last-dim input label, which would contaminate the resulting multi-layer prediction. This observation was considered in (Von Oswald et al. 2023)
4

---

Published as a conference paper at ICLR 2024
Mqwwry
X1 m
M2 X/
9k 64
M 9g
Vn 2-4
Squary
3/ 6
91
Figure I: The inputs/outputs of multi-layer in-context learner. We omitted x, and Xquerg since they are unchanged.
where the initial v Vquery = O.
Intuitively, the dynamics of the prefiXLM ICL as follows: all 7. starts as 0 at = 0, and gradually approach to the true label yy as increases, so that the difference (also as the output) 8/ gradually approaches to O. At the same time, dquery starts at O, and gradually approaches to -Uguery- the negation of the query label. Figure I provides illustration of these dynamics.
3.2 CAUSALLM ICL
causalLM applies auto-regressive attention throughout the entire sequence. Therefore, plugging the same K, Q, P. V into ( the update rules of (6) and (7) become:
8 = 80
7
x; Xj+
(8)
11
1/ 1 +
11
11-1)
(ys
(1-1)) 1)
x
Xj
1=1
(1) F with 0/0 y, 8, Moreover, since different dy. By are exposed to different ranges of inputs, there is no uniform w as in 3 that is associated with all Øj. Instead, if we define w, for each different
position j with w (0)) and
W 1 W
(11))
17
w (1-1)
x;) x
(10)
then we have the following proposition:
Proposition 2 For multi-layer causalLM-LS.. satisfying
5 with w 0, its input Z is
formatted as then its 1-th layer output 2 (x] ((0)
where 6()) Yj w (1) x, and w (I)
follow 10
The proof of Proposition 2 is provided in Appendix has 1741) w Xj, and
Similar prefiXLM-ICL, causalL also
( query -
(1) Mqquery
(1) Xquery
In summary, causalLM-ICL and PrefixLM-ICL are associated with different update rules: w (I)
follows 10 while (1) follows Specifically, in causalLM, it can be seen that the w; (1-1)
'There is another way of update which changes 9/n2 to n/j for the j-th example. We provide more details in Appendix D and show it performs worse than the main version in (

---

Published as a conference paper at ICLR 2024
corresponding to the first positions biased due to restricted access only few data points and furthermore, that these biases are propagated to later positions by (10). In prefixLM on the other hand, each position has access all the data and single w(') can used across the entire sequence as in (3) Although Eq. (3) and Eq. (10) only hold for the structured LSA case, the profound difference between causalLM and prefiXLM stems from their architectural difference and therefore we believe extends to general transformers, as indicated by our experimental results in Section
4 CONVERGENCE OF THE MULTI-LAYER IN-CONTEXT LEARNERS
In this section, we prove that both multi-layer prefiXLM and causalL. converge to their respective stationary points with increasing layers (and with linear rates). In addition, we show that the stationary point of prefixLM corresponds to the optimal least-square solution of the linear regression problem, while the ones corresponding causalLM are equivalent the iterative weights of online gradient descent of linear regression, which are known to be sub-optimal for a limited number of examples.
4.1 CONVERGENCE OF THE PREFIXLM ICL
The fact that multi-layer prefixLM computation exactly follows the update rule of w-()) 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 w follows the iterative updates of with coefficients satisfying:
3 then there exists a stationary point w
XX X
where y = (y1...., :yn)) and X (xI,..., Xn). Furthermore, the iterative weights w(O converge to w" with linear rate of convergence:
w w w" (1-1)) (w((-))- 100--xx"" xx
That is, Proposition holds for the multi-layer prefixLM, so that the same exact w is also the stationary point of LM, which it converges in linear rate. Furthermore this stationary point exactly the (optimal) least square solution the linear regression problem.
4.2 CONVERGENCE OF THE CAUSALLM ICL
Following the update rule of (10). we view multi-layer causalLM implicitly maintaining different weight vectors w, for each position j. In what follows, we show that: (a) Each such position j has its own stationary point wf. which appears to be different from the global optimal point wt of linear regression; (b) even when the number of in-context samples grows to infinity, convergence to w' is not guaranteed.
Specifically, in Appendix B we provide proof for the following proposition:
Proposition 4 If (f = Dín Ci. x_ follows the iterative updates of (10), then
a a a. !1)
(1)
Vj 2 i,
and there exist stationary points
Ein x (for j E I,...,n) n) with coefficients from
a" = ........ ..)) that satisfy y = T, where
x X1
x
x Xn
0
x
x, Xx
x X
Furthermore, the coefficients ( converges to the stationary point a" with linear rate of convergence:
(1) a" = (1-1) a -a
7 T).

---

Published as a conference paper at ICLR 2024
This proposition implies that the stationary points wj of ca-Sal are different from w", the least square solution of linear regression. However, a natural question is: increases, would w) ultimately converge to the optimal solution?
To answer this question, the next proposition shows that the stationary points wj follow an online gradient descent algorithm, whose loss and gradient at the j-th step is,
(y(w)) ) 2
(w Xf+1 - V)+1)*
Vw, L3 (w, = (w, Xi+1 -リj+ ((ni) x5+11
Proposition 5 Assuming that w the stationary points obtained in Proposition then
V Ww Vw (w
(11)
xj+1
The proof of Proposition 5] is provided in Appendix B Note that online gradient descent is known to converge optimal solution only with decaying step size j" for > 0 (Jentzen & Von Wurstemberger 2020). Since the step size of (1T) does not decay, we conclude that causalI. may not converge 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 the optimal solution w cannot be achieved by causalLM-ICL.
NUMERICAL EXPERIMENTS
Our experiments contain three parts.
We first use LSA-transformers on linear regression problems to directly verify our theorems. In Section 5. we show that despite that the in-context example (training) error of causalLM and prefixLM both decays in linear rates, the query (test) error of causalLM is significantly larger, which indicates its stationary solution not optimal.
ordinary softmax transformers gression, nonlinear regression and multiclass classification. Section [5.2] we show that our theoretical insights generalize other tasks types (i.e., that ICL prefiXLM still outperforms M in all these cases).
Secondly,
few synthetic tasks including linear re-
Lastly, in Section [.3] we conduct LLM based ICL experiments using T5 (Roberts et al. 2022). We also provide additional experimental results PaLM2 (Google et al. 2023) as well as large multimodal models (PaLI-X, Chen et al. (2023) in Appendix E.6 and E.7
5.1 LSA-TRANSFORMERS ON LINEAR REGRESSION
In order to directly verify our theorems from Section 4 we first study in-context learning 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 E for an illustration. The data input x, of the sequence is sampled from WI (,-1,11) 1)1 Each sequence associated with a single weight vector w that sampled from (O I), and the labels are computed as y/ wx,, Assuming the prediction of each layer is 8/ we evaluate the MSE [|æ.
v! 1/ on both in-context and query examples across different layers
1/113
The results are plotted in Figure 2 left (for MM) and middle (for causalLM), Our results are averaged over 64 randomly generated sequences.
we can see, although both prefiXI.M and causalLM has lincar rate of convergence (with respect to the number of layers) on the in-context examples, the query errors of causalLM are stuck above the 10 level, while the query error of PrefixLM decays in the same linear rate its training error.
Furthermore, in Figure 2 right, we plot the query errors of the stationary points (following Proposi- tion corresponding to the outputs of infinite layers) of causaILM-ICL with increasing number of in-context examples up to 300. Although causalLM-ICL able to eventually converge to optimal solution when Paz 0, it takes more than 100 examples reach below 10-2 The convergence is
7

---

Published as a conference paper at ICLR 2024
Coneegence offuttilayer prefiuM
come gaice afoutt CRSSUUMM
Statirary parts coust.N
R-Ghtee aannie
1001
10-
10-1
10
End c5072:
20
220
225
Figure 2: Left/Middle: the MSE on in-context examples and query examples of multi-layer LSA- based prefixLM/ccuuaLLM-CCC. with 40 in-context training examples. Right: the query MSE of causalLM-ICLLs stationary points (per Proposition 4
using up 300 in-context examples.
LR
N-LR
MC
MMSL
8.6e-3
1.5e-4
24.1
CausalLM-SL
1.9e-1
2.7e-3
27.0
M-LL
2.5e-3
9.0e-5
27.6
CausaILM-UL.
1.6e-2
2.9e-3
32.1
Table The test query errors of the unshared-layer (UL) and sharing-layer (SL) transformer-ICLa on linear regression (LR), non-linear regression (NLR), and multiclass classification (MC) tasks. Both regression tasks report mean squared errors; and the MC task reports the classification error.
even worse as we vary the input distribution x ~ (((111)) + Ma with increasing us E 1,,1,2,3) which demonstrates that MuusalLM-ICLL not optimal for few-shot ICL.
5.2 STANDARD TRANSFORMERS 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 synthetic tasks including linear regression, non- non-linear regression and multiclass classification.
All three tasks used 16-dim inputs with (d(11,111)) 1)16 and w NNO0 I). For non-linear regression, we apply sigmoid activation on the logit such that sigmoid( w x): and for multiclass classifica- tion, we randomly generate three w. ~ (0, I), and assign labels based on y = argmax, {wo x}. We trained 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 UL). In particular, the sharing-layer transformer can be considered recurrent sys- tem (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 closer our constructed L.A..
The ICI. 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 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
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 []] We find that both prefixLM-SL and PrefixLM-UL significantly outperform their counterparts of causalLM in all cases, As side note, transformer-SL. appears outperform transformer- UL in the classification tasks, which indicates the overfitting problem of the latter due to over peerreeamettiiiiionn. 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 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.
8

---

Published as a conference paper at ICLR 2024
Regression
Reg
Ph Chuda LV
han Cauoll
Chuck
1b+
10-1
10
3.24
15
25
Layar
Layer
Layer
Figure 3: The test query errors of the 24-layer SL-transformers based prefiXLM/causalLM- on linear regression (left), non-linear regression (middle), and multiclass classification (right).
MMLU
BBH
Base
Large
XL
Base
Large
XL
PrefiXLM
28.8
32.0
39.5
27.4
32.2
35.8
CausalLM
28.0
26.9
30.5
24.8
29.8
32.0
Table The averaged test query accuracies on 5-shot MMLU (57 tasks) and 3-shot BBH (23 tasks) with FLAN-finetuned T5 DecoderOnly PrefsxX. checkpoints.
5.3 ICL LARGE LANGUAGE MODELS
In order to compare the ICL performance of causalLM and prefiXLM in a large language model setting, we conduct experiments using the publicly available T5 family of models (Roberts et al. 2022). Note that the existing public T checkpoints are all based on EncDec models, which are similar to prefixLM. Thus, it would unfair and unnatural to compare with causalLM by simply replacing the bidirectional attention of the encoder to the causal attention during the finetuning stage. To make a more reasonable comparison, we reran the pretraining stages of T5 on the C4 corpus (Raffel al. (020a) from random initialization point using a span corruption objective, but in the DecoderOnly setting. Moreover, for each size (from Base, Large and XI.) of the models, we pretrained two checkpoints, one with prefiXLM and the other with causalLM, each for 1M steps using the same T5 pretraining recipe. After pretraining, we use the FLAN recipe (Chung al., 2022) to finetune each checkpoint (40k steps for Base, 20k steps for Large and XL) with its pretraind attention mask and evaluate the ICL capability of the finetuned models on two benchmarks: MMLU (Hendrycks et al. 2020) and BBH (Suzgun et al. (022).
Table 2 shows that for all three sizes of T5X DecoderOnly models, the MMLU and BBH accuracies of LM surpasses that causall.M consistently and such gap widens as the size of the model becomes larger. This result empirically verifies that our conjecture generalizes to the practical case.
We supply additional empirical evidence on state-of-the-art models in Appendix E.6 and E.7
CONCLUSION
In this paper, we analyzed the convergence properties two types of widely-used transformer-based language models (causalLM and prefixLM), during in-context learning. Using 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 prefixI.M coincides with the optimal least square solution; while the ones of causalI.M . is equivalent 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 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.
5xx

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We want to specially thank Xinhua Zhang for helpful discussions and references about online learning. We also thank Yi Tay for comments regarding the PaLM2 checkpoints.
REFERENCES
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211. 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Xi Chen, Josip Djolonga, Piotr Padlewski, Basil Mustafa, Soravit Changpinyo, Jialin Wu, Car- los Riquelme Ruiz, Sebastian Goodman, Xiao Wang, Yi Tay, Siamak Shakeri, Mostafa Dehghani, Daniel M. Salz, Mario Lucic, Michael Tschannen, Arsha Nagrani, Hexiang Hu, Mandar Joshi, Bo Pang, Ceslee Montgomery, Paulina Pietrzyk, Marvin Ritter, A. Piergiovanni, Matthias Min- derer, Filip Pavetic, Austin Waters, Gang Li, Ibrahim M. Alabdulmohsin, Lucas Beyer, Julien Amelot, Kenton Lee, Andreas Steiner, Yang Li, Daniel Keysers, Anurag Arnab, Yuanzhong Xu, Keran Rong, Alexander Kolesnikov, Mojtaba Seyedhosseini, Anelia Angelova, Xiaohua Zhai, Neil Houlsby, and Radu Soricut. Pali-x: On scaling multilingual vision and language model. ArXiv, abs/2305. 18565, 2023. URL httph: api. semant i icscholar. Org/CorpusID: 258967670
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung. Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:220..03111 2022.
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. Scaling esstuucioonfinetunee language models. arXiv preprint arXiv:2210. 1416, 2022.
Mostafa Dehghami, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Eukasz Kaiser. Universal transformers. arXiv preprint arXiv: 000090 2018.
Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: /8 18.0008805 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.1 1929, 2020.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598 2022.
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 EI Shafey, Yanping Huang, Kathy Meier- Hellssee,, 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
10

---

Published as a conference paper at ICLR 2024
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. Palm 2 technical report, 2023.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300, 2020.
Arnulf' Jentzen and Philippe Von Wurstemberger. 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, 2020.
Andrej Karpathy and Li Fei-Fei. Deep visual-semantic alignments for generating image descriptions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3128-3137, 2015.
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? arXiv preprint aiviv:2202. 12837, 2022.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21 21(140):1-67, 2020a. URL
http::/ www.llr.rgggpppppepeee v21/20 074. html
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learming Research, 21(1):5485-5511, 2020b.
Adam Roberts, Hyung Won Chung, Anselm Levskaya, Gaurav Mishra, James Bradbury, Daniel Andor, Sharan Narang, Brian Lester, Colin Gaffney, Afroz Mohiuddin, Curtis Hawthorne, Aitor Lewkowycz, Alex Salcianu, Mare van Zee, Jacob Austin, Sebastian Goodman, Livio Baldini Soares, Haitang Hu, Sasha Tsvyashchenko, Aakanksha Chowdhery, Jasmijn Bastings, Jannis Bulian, Xavier Garcia, Jianmo Ni, Andrew Chen, Kathleen Kenealy, Jonathan H. Clark, Stephan Lee, Dan Garrette, James Lee- Lee-Th Colin Raffel, Noam Shazeer, Marvin Ritter, Maarten Bosma, Alexandre Passos, Jeremy Maitin-Shepard, Noah Fiedel, Mark Omernick, Brennan Saeta, Ryan Sepassi, Alexander Spiridonov, Joshua Newlan, and Andrea Gesmundo. Scaling up models and data with 5x and seqio. arXiv preprint arXiv:2203.171889 2022. URL http:::////at //arxiv org/ arg/abs/// 2203.2233..0 17189
Piyush Sharma, Nan Ding, Sebastian Goodman, and Radu Soricut. Conceptual captions: A cleaned, hypernymed. image alt-text dataset for automatic image captioning. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), Melbourne, Australia, July 2018
Mirac Suzgun, Nathan Scales, Nathanacl Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoe V Le. Ed H Chi, Denny Zhou, et al. Challenging big-bench tasks and whether thhin-tf-thought can solve them. arXiv preprint arXiv:22 2022.
Yi Tay. Mostafa Dehghani, Vinh Q Tran, Xavier Garcia, Dara Bahri, Tal Schuster, Huaixiu Steven Zheng, Neil Houlsby. and Donald Metzler. Unifying language learning paradigms, arXiv preprint arXiv:2205.05 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017.
11

---

Published as a conference paper at ICLR 2024
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference Machine Learning, pp- 35151-35174. PMLR, 2023.
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. Larger language models do in-context learning differently, arXiv preprint arXiv:2303.0888
2023.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning implicit bayesian inference. arXiv preprint arXiv:2111.0208) 2021.
Ruiqi Zhang, Spencer Frei, and Peter I Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2006.09927, 2023.
12


Reviews:
[{"review_id": "K6amDXXfof", "reviewer": "Reviewer_9q6M", "summary": "In this work, authors relate the in-context learning abilities of causalLM architectures with that of an online gradient descent solver, versus the abilities of a prefixLM architecture that show convergence to the optimal solution at a linear rate. I found the experiments well conducted and the paper's results well presented overall. I'm unsure of the problem's relevance, but find the paper interesting regardless.", "questions": "See weaknesses", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Primarily a study of ICL's limitations in standard transformer architectures, which use causal self-attention masking and hence causal language modeling objectives. prefixLMs' superiority is also described and the empirical results are interesting.", "weaknesses": "While I find the results and the paper itself interesting, as I mentioned in my summary of the paper, I'm unsure of the relevance of the problem. I think it would help improve the paper's impact if the motivations were clarified better: specifically, what models actually use a PrefixLM architecture in current literature and demonstrating or citing papers which show these models have a different qualitative behavior that a standard causalLM architecture.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YbwpL8wW7L", "reviewer": "Reviewer_RFAv", "summary": "The paper examines transformer-based in-context learning using pre-fixed language models (pre-fixedLM) and causal language models (causalLM). While pre-fixedLMs are empirically observed to outperform causalLMs, the theoretical reasons remain elusive. Through a convergence analysis, the study reveals that pre-fixedLMs optimally converge to linear regression solutions, whereas causalLMs follow online gradient descent dynamics, leading to suboptimal results. Empirical tests validate these theoretical findings.", "questions": "Q1. A simple augmentation could be used to manually feed those other possibilities of partitions to casual LM, so prefix LM and casual LM would be trained with the same number of in-context and query possibilities for a given context length. Under this setting, would prefix LM still perform better than casual LM?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Theoretical understanding of different solutions found by prefix and casual LMs are provided under the linear regression setting, which seems novel.\n2. The paper also provides an empirical study to compare the solution found by prefix and causal LMs in different tasks, which verifies the theoretical intuitions.", "weaknesses": "1. The experimental setting is limited to a given number of in-context examples which seems to naturally favor prefix LMs. Casual LMs would train the model with different numbers of in-context samples simultaneously while prefix LM using all possible in-context and query partitions with the same in-context length. Testing with fewer in-context examples could be beneficial to provide more comprehensive results.\n2. Another unfairness in the experimental setting is that the nature of casual LMs would basically reduce the number of examples with the same in-context samples compared to prefix LM (i.e. for one input sequence, prefix LMs would enumerate all possible in-context and query partitions but casual LM would only have one for a specific context length). See Q1.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GKqalPTYyN", "reviewer": "Reviewer_m1Xp", "summary": "The paper argues that in context learning with causalLM is inferior to in context learning with prefixLM. They extend the theory of Von Oswald et al. to include multi-layer Linear Self-Attention and multi-step gradient descent in the setup considered in Von Onswald et al., showing convergence result in the extended setup for both causalLM and prefixLM in context learning. Using their developed theory and experimental validation, they demonstrate that prefixLM is superior to causalLM.", "questions": "* Does prior works generally use causalLM instead of prefixLM?\n* Are the empirical results presented in this work reproducible with open models?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper follows an emerging line of papers showing the equivalence of gradient descent to in context learning in a very specific setup where the self-attention is linear, the objective is linear regression and the parameter matrices are hand-constructed.\n\nAbstracting away the limited setup, the paper does a good job extending the theory of Von Onswal et al. The theoretical argument is clear in my opinion and the evidences supporting the main thesis of the work are convincing enough. In particular, the paper does a good job arguing in details why prefixLM is better to causalLM both theoretically and empirically on more realistic setups.", "weaknesses": "My understanding is that this work aims at demonstrating that prefixLM is superior to causalLM for in context learning. While I believe they do a good job at it, I am under the impression that most projects already used prefixLM when possible. For example, InstructBLIP and Llama2 use prefixLM as far as I can understand. If this is the case that most influential language model or VLM already use prefixLM then I am unclear about the intended impact of this work.\n\nThe models used in this work are internal models that cannot be accessed by anyone outside Google. I believe that this is a big problem for open science and reproducible research. Unless the conclusions made in this work cannot be attained with available open models, which I highly doubt, I urge the authors to consider using open alternative to disseminate their results.\n\n\n\nMinor:\n* The considered theoretical setup seems contrived to me.  Since I don't expect the author to come up with a new theoretical setup, this point is not reflected in my evaluation of this paper. However, I believe that this is a weakness of this line of works and I would appreciate to be convinced why any conclusion made on this setup (LSA, linear regression and hand-crafted parameters) should necessary translate to more realistic setups.\n* Figure 1 is blurry on printed paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0PI7wwMUKe", "reviewer": "Reviewer_mzWe", "summary": "This paper analyzes the convergence behavior of in-context learning under a certain parameter construction and shows that while both LM types converge to their stationary points at a linear rate, prefixLM converges to the optimal solution of linear regression while causalLM may include a bias that may not diminish even if the number of samples grows infinitely. They also show with experiments on both synthetic and real tasks that prefixLM seems to outperform causalLM given the same amount of parameters.", "questions": "* **Q1.** (See weakness) If the reviewer's understanding is correct, then a question is whether proving such a statement is possible.\n\n* **Q2.** Is it possible to convert a pretrained CausalLM into a PrefixLM during in-context learning through fine-tuning and improve the performance of ICL?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "* **Originality.** The paper reveals the surprising finding that CausalLM may underperform PrefixLM in the task of in-context learning.\n\n* **Significance.** The paper is of significance in understanding the in-context learning mechanism.\n\n* **Clarity.** The paper is well-written and easy to read.\n\n* **Quality.** The experiments are thorough and the theory is rigorous.", "weaknesses": "* To the reviewer's understanding, this paper hasn't excluded the possibility that there exists a parameter configuration that can allow CausalLM to reach optimal solution in the linear setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)