How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations
Authors: Tianyu Guo
Wei Hu
Song Mei
Huan Wang
Caiming Xiong
Silvio Savarese
Yu Bai
Venue: ICLR 2024 poster (2024)
Keywords: in-context learning
transformers
representation learning
learning theory
mechanistic understanding
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: Transformers can learn in-context on function classes with a representation, with empirical mechanisms that align well with the efficient theoretical constructions.

Abstract: While large language models based on the transformer architecture have demonstrated remarkable in-context learning (ICL) capabilities, understandings of such capabilities are still in an early stage, where existing theory and mechanistic understanding focus mostly on simple scenarios such as learning simple function classes. This paper takes initial steps on understanding ICL in more complex scenarios, by studying learning with \emph{representations}. Concretely, we construct synthetic in-context learning problems with a compositional structure, where the label depends on the input through a possibly complex but \emph{fixed} representation function, composed with a linear function that \emph{differs} in each instance. By construction, the optimal ICL algorithm first transforms the inputs by the representation function, and then performs linear ICL on top of the transformed dataset. We show theoretically the existence of transformers that approximately implement such algorithms with mild depth and size.  Empirically, we find trained transformers consistently achieve near-optimal ICL performance in this setting, and exhibit the desired dissection where lower layers transforms the dataset and upper layers perform linear ICL. Through extensive probing and a new pasting experiment, we further reveal several mechanisms within the trained transformers, such as concrete copying behaviors on both the inputs and the representations, linear ICL capability of the upper layers alone, and a post-ICL representation selection mechanism in a harder mixture setting. These observed mechanisms align well with our theory and may shed light on how transformers perform ICL in more realistic scenarios.

Published as a conference paper at ICLR 2024
HoW Do TRANSFORMERS LEARN IN-CONTEXT BEYOND SIMPLE FUNCTIONS? A CASE STUDY ON LEARNING WITH REPRESENTATIONS
Tianyu Guo Wei Hu Song Mei*
Huan Wang Caiming Xiong Silvio Savarese Yu Bai UC Berkeley University of Michigan "Salesforce AI Research t clanys_gu@@berke edu
ABSTRACT
While large language models based on the transformer architecture have demon- strated remarkable in-context learning (ICL) capabilities, understandings of such capabilities are still in an early stage, where existing theory and mechanistic un- derstanding focus mostly on simple scenarios such as learning simple function classes. This paper takes initial steps on understanding ICL. more complex scenarios, by studying learning with representations. Concretely, we construct synthetic in-context learning problems with a compositional structure, where the label depends on the input through possibly complex but fixed representation function (which we instantiate as multi-layer MLPs), composed with linear function that differs in each instance. By construction, the optimal ICL algorithm first transforms the inputs by the representation function, and then performs linear ICL on top of the transformed dataset. We show theoretically the existence of transformers that approximately implement such algorithms with mild depth and size. Empirically, we find trained transformers consistently achieve near-optimal ICL performance in this setting, and exhibit the desired dissection where lower layers transforms the dataset and upper layers perform linear ICL. Through exten- sive probing and a new pasting experiment, we further reveal several mechanisms within the trained transformers, such as concrete copying behaviors on both the inputs and the representations, linear ICL capability of the upper layers alone, and a post-ICI. representation selection mechanism observed mechanisms align well with our theory and may shed light on how trans- formers perform ICL in more realistic scenarios.
harder mixture setting. These
INTRODUCTION
Large language models based on the transformer architecture have demonstrated remarkable in- context learning (ICL.) capabilities (Brown al., 2020), where they solve newly encountered tasks when prompted with only a few training examples, without any parameter update to the model. Recent state-of-the-arr models further achieve impressive performance in context on sophisticated real-world tasks (OpenAI, 2023; Bubeck et al., 2023; Touvron et al., 2023). Such remarkable ca- pabilities call for better understandings, which recent work tackles from various angles (Xie al., 2021: Chan et al., 2022; Razeghi et al., 2022; Min et al., 2022; Olsson et al., 2022; Wei et al., 2023).
recent surge of work investigates ICL theoretically amenable setting where the context con- sists of real-valued (input, label) pairs generated from a certain function class. They find that trans- formers can learn many function classes in context, such as linear functions, shallow neural net- works, and decision trees (Garg et al., 2022; Akyürek et al., 2022; Li et al., 2023a), and further studies provide theoretical justification on how transformers can implement and learn various learn- ing algorithms in-context such ridge regression (Akyürek et al., 2022), gradient descent (von Oswald et al., 2022; Dai et al., 2022; Zhang et al., 2023a; Ahn et al., 2023), algorithm selection (Bai et al., 2023), and Bayes model averaging (Zhang al., 2023b), to name few. Despite the progress, an insufficiency of this line is that the settings and results may not actually resemble ICL in real- world scenarios- For example, ICI. in linear function classes well understood in theory with

---

Published as a conference paper at ICLR 2024
(a) Illustration of setting and theory
(b) ICL risks
(c) Linear probes
1.0
ICL proslem
x
Teratrrer
0.6
0.8
meen squardddrd
Our setting
varyling
sauve
0.6
Transfomeer Mechanism
Compute @ 0'()"")
Linear (CL (0
0.2
0.0
40
0
6
10
12
In-context esamples
Jayer
Figure I: An illustration of our setting and results. (a) We consider ICL problems with a fixed representation composed with changing linear functions, and we construct transformers that first compute the representations and then performs linear ICL. (b,c) Empirically, learned transformers can perform near-optimal ICL in this setting, and exhibit mechanisms that align with our theory (detailed setups in Section 5.1).
efficient transformer constructions (Bai al., 2023), and transformers indeed learn them well em- pirically (Garg et al., 2022); however, such linear functions in the raw input may fail to capture real-world scenarios where prior knowledge can often aid learning.
This paper takes initial steps towards addressing this by studying ICL in the setting of learning with representations, more complex and perhaps more realistic setting than existing ones. We construct synthetic ICL tasks where labels depend on inputs through a fixed representation function composed with varying linear function. We instantiate the representation as shallow neural networks (MLPs), and consider both a supervised learning setting (with input-label pairs) and a dynamical systems setting (with inputs only) for the in-context data. Our contributions can be summarized follows.
Theoretically, we construct transformers that implement in-context ridge regression on the rep- resentations (which includes the Bayes-optima algorithm) for both learning settings (Section 4). Our transformer constructions admit mild sizes, and can predict at every token using decoder architecture, (non-trivially) generalizing existing efficient constructions that predict at the last token only using an encoder architecture.
Empirically, using L-layer MLPs as representations, we find that trained small transformers consistently achieve near-optimal ICL. risk in both learning settings (Section 5 & Figure lb).
Using linear probing techniques, we identify evidence for various mechanisms in the trained transformers. Our high-level finding is that the lower layers transforms the data by the repre- sentation and prepares it into a certain format, and the upper layers perform linear ICL on top of the transformed data (Figure Ic), with often clear dissection between these two modules, consistent with our theory. See Figure Ia for pictorial illustration.
We further observe several lower-level behaviors using linear probes that align well with our (and existing) theoretical constructions, such as copying (of both the input and the representations) where which tokens are being copied are precisely identifiable (Section 5.2), and a post-ICL representation selection mechanism in a harder setting (Section 5.1.1 & Appendix E).
We perform new pasting experiment and find that the upper layers within the trained trans- former can perform nearly-optimal linear ICL in stronger evidence that the upper module alone can be strong linear ICL learner.
(Section 5.1), which provides
2 RELATED WORK
In-context learning The in-context learning (ICL) capabilities of pretrained transformers have gained significant attention since first demonstrated with GPT-3 (Brown al., 2020). Subsequent empirical studies have investigated the capabilities and limitations of ICL in large language models (Liu et al., 2021; Min et al., 2021a;b; Lu et al., 2021; Zhao et al., 2021; Rubin et al., 2021; Razeghi et al., 2022; Elhage et al., 2021; Kirsch et al., 2022: Wei et al., 2023).
A line of recent work investigates why and how pretrained transformers perform ICL from a the- oretical perspective (Garg et al., 2022; Li et al., 2023a; von Oswald et al., 2022; Akyürek et al., 2022: Xie et al., 2021; Bai et al., 2023: Zhang et al., 2023a;b; Ahn et al., 2023; Raventós et al., 2023). In particular, Xie et al. (2021) proposed a Bayesian inference framework explaining ICL. Garg et al. (2022) showed transformers could be trained from seratch for ICL of simple function classes. Other studies found transformers can implement ICL through in-context gradient descent
2

---

Published as a conference paper at ICLR 2024
(von Oswald et al., 2022; Akyürek al., 2022) and in-context algorithm selection (Bai al., 2023). Zhang et al. (2023a) studied the training dynamics of a single attention layer on linear ICL tasks. Li et al. (2023b) used the ICL framework to explain chain-of-thoughh reasoning (Wei et al., 2022). Our work builds on and extends the work of (Garg et al., 2022; Akyürek et al., 2022; von Oswald al., 2022; Bai et al., 2023), where we study the more challenging setting of ICL with a represen- tation funcciin, and also provide new efficient ICL constructions for predicting at every token using a decoder transformer, opposed predicting only at the last token in most of these work.
In-weights learning versus in-context learning Recent work has investigated when transformers learn fixed input-label mapping versus when they perform ICL (Chan et al., 2022; Wei al., 2023; Bietti et al., 2023). Chan et al. (2022) refer to learning a fixed input-label mapping from the pre- training data "in-weights learning" (IWL), in contrast with ICL. Our problem setting assumes the pre-training data admits a fixed representation function, which should be learned by IWL. In this perspective, unlike these existing works where IWL and ICL are typically treated as competing mechanisms, we study model in which IWL (computing the fixed representation by transformer weights) and ICI. (learning the changing linear function context) occur simultaneously.
Mechanistic understanding and probing techniques A line of work focuses on developing tech- niques for understanding the mechanisms of neural networks, in partieular transformers (Alain & Bengio, 2016; Geiger et al., 2021; Meng et al., 2022; von Oswald et al., 2022; Akyürek et al., 2022; Wang et al., 2022; Räuker et al., 2023). We adopted the linear probing technique of (Alain & Bengio, 2016) in a token-wise fashion for interpreting the ICL mechanisms of transformers. Beyond prob- ing, more convincing mechanistic interpretations may require advanced approaches such as causal intervention (Geiger et al., 2021; Vig et al., 2020; Wang et al., 2022); Our pasting experiment has similar interventional flavor that we feed input sequences (ICI. instances) from another distribu- tion directly (through trainable embedding layer) to the upper module of transformer.
PRELIMINARIES
Transformers
We consider -ouuence----seuence functions applied to N input vectors (h, {h,) C R Puia in Dhid dimensions, which write compactly as an input matrix H [h...... hv] R xN where each h; is a column of H (also a token).
We use standard L-layer decoder-only (autoregressive) transformer, which consists of L. consec- utive blocks each with a masked self-attention layer (henceforth "attention layer") followed by an MLP layer. Each attention layer computes
Attno = H EM (V H) x a (MSK @ (KQTH)TTK,AH) E RDDN
where = [(Qm, Km, V..) C PhiaxDDii JmE[M] are the (query, key, value) matrices, M the number of heads, MSK E NxN is the decoder mask matrix with MSK;j 1ii < j}. and of the activation function which is typically chosen as the (column- wise) softmax: [o(A) = softmax(a,) R- for A [ai...., aN] E NNN, Each MLP layer computes
MLP W. HH) = H W2o
where W (1.2) E Phaa x Paid are the weight matrices, and ((t) max (t, 0} is the ReLU activation.
We use TF to denote a transformer, and typically use H = TF THH to denote its output on H.
In-context learning We consider in-context learning (ICL) on regression problems, where each ICL instance is specified by dataset D = {(x,, u)) u))eem
iid ~ P. with (x,, W.) E RRd x R, and the model is required to accurately predict given all past observations D:-1 == ((x, H; )
and the test input x,. Each instance D - D'ú) is drawn from a different data distribution P = P()) Accurate prediction requires learning P in-context from the past observations D._ (i.e. the context); merely memorizing any fixed p(1) is not enough. This is main challenge of in-context learning.
3

---

Published as a conference paper at ICLR 2024
We consider using transformers to do ICL, where we feed sequence of length 2N into the trans- former TF using the following input format:
X1
XN
h2N]
=
1/1
UN
E E P
(I)
Pf
P1
PN
P
where P P E IR PBiaaade are fixed positional encoding vectors consisting of zero paddings. fol- lowed by non-zero entries containing information about the position index and indicator of being an z--token (1 p and p²); see (12) for our concrete choice. We refer to each odd token h as as an r-token (also the x,-token), and each even token b2; as a y-token (also the ys-token).
After obtaining the transformer output H TF(H), for every index E [N], we extract the prediction Di from the output token at position x5 Vo (hf))+1. Feeding input (1) into the transformer simultaneously computes V5 < TF(x1, : 1i x,) for all i E [N]. Denote
the parameters of transformers as
In addition to the above setting. we also consider a dynamical system setting with D = {x where the transformer predicts X, from the preceding inputs Xgi-1- See Section 4.2 for details.
4 IN-CONTEXT LEARNING WITH REPRESENTATIONS
4.1 SUPERVISED LEARNING WITH REPRESENTATION
We begin by considering ICL regression problems with representation, where labels depend the input through linear functions of a fixed representation function. Formally, let 4* gd + E be fixed representation function. We generate each in-context data distribution P ww by sampling a linear function w ~ N(O, IDD) from a Gaussian prior, and then generate the ICL instance D {(x,, (eENN Pw by a linear model on 4 with coefficient w and noise level O:
Vi (w, (xx))) azi,
iid Pa.. iid N(O, 1), i E [N].
Note that all D's share the same representation d", but each admits unique linear function w.
The representation function 4* can in principle be chosen arbitrarily. As a canonical and flexible choice for both our theory and experiments, we choose 4* to be standard L-layer MLP:
x- x*(x)= 0² (B o (Bi.-1- ... (Bix)) -.x)) Bi E Dxd (B;) R
(3)
where D is the hidden and output dimension, and a² is the activation function (applied entry-wise) which we choose to be the leaky ReLU (+) = (t) == max {1, pt} with slope p E (0, 1).
Theory As th is fixed and the w is changing in model (2), by construction, good ICL algorithm should compute the representations (Φr (x:)}; and perform linear ICL on the transformed dataset {(0" (x,), v;)); learn w. We consider the following class of & "--ridge estimators:
= arg min 2(6-1) E ((w, (x(x)) vz) + 1
(Φ* Ridge)
and we understand W #" A O. In words, w; d 'A performs ridge regression on the transformed dataset ( Vo Vi for all e [N]. By standard calculations, the Bayes-optimal predictor*
for y: given (D,-1,x;) is exactly the ridge predictor y;
= (W!" ** (x;:) o² /
We show that there exists a transformer that can approximately implement (ii-Ridge) in-context at every token i [N]. The proof can be found in Appendix B. Theorem (Transformer can implement T - Ridge). For any representation function 40 of form (3), any X > 0, B3. B.,, By > 0, < B& Bw.B /2, letting == 1 / there exists a transformer TF with L + (K log B)) layers, heads, Dnid = 2D + d 10 such that the following holds.
For any dataset D such that (x:)(2 < Ba. lysl < By and the corresponding input H E T Daaax2N of format (I), we have
'There is no information leakage, as the "prefix" property of decoder transformers Gú; haz-1 (TF (1)(211)))22-1 ensures that h. (and thus Dh) only depends on (D,-1, xz). "The predictor Di (D,,1,, x;) that minimizes the posterior square loss El 21)
4

---

Published as a conference paper at ICLR 2024
(a) The first (L + 2) layers TF transforms x; to the representation 4* (xx) each token, and copies them into the succeeding token:
TF(LL+2) (H) = (x((x,)
O (x1)
('(xN)
Φ' (x(xN)
0
Un
(4)
Pi
P
PM
PN
where p P only differ from P P the dimension of the zero paddings.
(b) For every index E [NT. the transformer output H TF(H) contains prediction T [124-1] that is close to the (Ridge predictor: T ( (x:), W, A)
The transformer construction in Theorem consists of two "modules": The lower layers computes the representations and prepares the transformed dataset (x*((x), 9)); into form (4). In particular, each (x,) appears both the i-th z-token and is also copied into the succeeding y token. The upper layers perform linear ICL (ridge regression) top of the transformed dataset. We will test whether such mechanisms align with trained transformers in reality in our experiments (Section 5.1).
Proof techniques The proof of Theorem I builds upon (1) implementing the MLP 4* by trans- formers (Lemma B.3), and (2) an efficient construction of in-context ridge regression (Theorem B.5), which to our knowledge is the first efficient construction for predicting at every token using de- coder transformers. The latter requires several new construction techniques such copying layer (Lemma B.1), and an efficient implementation all tokens simultaneously using decoder transformer (Proposition B.4). These extend the related constructions of von Oswald et al. (2022); Bai et al. (2023) who only consider predicting at the last token using encoder transformer, and could be of independent interest.
N parallel in-context gradient descent algorithms
In addition, the bounds on the number of layers, heads, and Dhia in Theorem can imply a sample complexity guarantee for pprai)trannnn:: transformer with E-excess risk (on the same ICL instance distribution) over the one constructed in Theorem 1 can be found in Õ( + x)) + d)² training instances, by the generalization analysis of Bai et al. (2023, Theorem 20). We remark that the constructions in Theorem 2 choose ō as the normalized ReLU instead of softmax, following (Bai et al., 2023) and in resonance with recent empirical studies (Wortsman et al., 2023).
4.2 DYNAMICAL SYSTEM WITH REPRESENTATION
As a variant of model (2), we additionally consider a (nonlinear) dynamical system setting with data D (x1..... .N)) where each X;+1 depends on the k preceding inputs (x for some & 2 1 through linear function on top of a fixed representation function $* Compared to the supervised learning setting in Section 4. this setting better resembles some aspects of natural language, where the next token in general depends on several preceding tokens.
Formally, k 1 denote the number of input tokens that the next token depends on, and $ Rekd RD denotes representation function. Each ICL instance 'D = [x:)ee[] is generated as
First sample P = Pw where W E Dxd sampled from a Gaussian prior: Wa, iid
follows:
N(O, 72). Then sample the initial input x1 Ps and let
X T Φ5 ('xxxxxxx; x;)) + az; 2, ndd N(O, E -1,,
(5)
where understand xj O₂ for j I O. We choose O² be the same L-layer MLP as in (3), except that the first weight matrix has size B E IR Pxkd to be consistent with the dimension of the augmented input X, := [ : x.]. We remark that (5) substantially generalizes the setting of Li et al. (2023a) which only considers linear dynamical systems (equivalent to Φ- = id), task arguably much easier for transformers to learn in context.
As x; acts both inputs and labels in model (5), we use the following input format for transformers:
XN
R PaiaXN
P1
Pv
where p; pm d-4; 4; 1; 12; 13) and extract prediction š.t1 from the i-th output token.
5

---

Published as a conference paper at ICLR 2024
Theory Similar above, consider the ridge predictor for the dynamical system setting
arg min 2(-1) WT (x(x)) - x5++ |i (& (dddRidee-Dyn)
We understand W = Opxd, and let 2 == max |W dاll for any W E Dxd
Again, (0- (idgeede)DDyn) gives Bayes-optimal predictor (W; (x.) at X A=
The following result shows that (fd-Ridge-Dyn) can also be implemented efficiently by trans- former. The proof can be found in Appendix C.2.
Theorem 2 (Transformer
implement Φ+ -Ridee for dynamical system). For the dynamical sys- tem setting where the L-layer representation function 4" IR + _R" takes (3), but otherwise same settings as Theorem 1, there exists a transformer TF with L+2 Or log(Bg. Bu/B//))) layers, max {3d, 5} heads, and Duid = max (2 + DDd d) + such that the following holds.
For any dataset D such that || (x,)ll2 < Ba, |> S By, and (odge--Dyn) for all E [M] and corresponding input H IK Puaax of format (6), have
B (cf. ($*-
(a) The first transformer layer copies the k previous inputs into the current token, and computes the first layer to, (B,xx,))) within :
X1-4+1
XN-K+1
(H) Attn (H)
X1
XN
P
PvJ
X1
XN
P
Pv
To
(BixN)
TF (H) MLP(1)
(Attn (())
=
XI
XN
(8)
P
PN
(b) The first (L + 1) layers of TF transforms each to T6 (X,), and copies the preceding repre-
sentation 4h (Xt-1) onto the same token to form the (input, label) pair (Φ (X;-1). x;):
T (X1)
(x(x)
$ (XNN
Oa
0g
Oa
( (H)
4
(X1)
4- (NN-1)
(9)
X1
X2
XN
P1
P2
PN
Above, P. Pa. Pi only differs from ps in the dimension of the zero paddings.
(c) For every index i E [N), the transformer output H -TF(H) contains prediction R
[hi]aa that is close to the -Ridgee predictor: || - (W. >> ( (x))ll S
To best knowledge, Theorem 2 provides the first transformer construction for learning nonlinear dynamical systems in context. Similar as for Theorem 1, the bounds on the transformer size here imply guarantees E excess risk within Ö(L + k:) + D)d)2e-2) (pre-)training instances.
In terms of the mechanisms, compared with Theorem the main differences in Theorem
are (1) the additional copying step (7) within the first layer, where the previous (A - 1) to- kens [xi-&+ij.... x;-1] are copied onto the x; token, to prepare for computing the intermediate output (9), where relevant information (for preparing for linear ICI.) has form (x;-1); x (x.)) and gathered in the x-tokens, different from (4) where the relevant infor- mation is [ (x;); v.l. gathered in the y-token. We will test these in experiments (Section 5.2).
x' (x;): (2)
EXPERIMENTS
We now empirically investigate trained transformers under the two settings considered in Section 4.1 4.2. both cases, we choose the representation function x to be a normalized version of the L-layer MLP (3): (*(x) = Tx (x)) (x) l2, where कx takes form (3), with weight matrices (B) [ee()) sampled random coolumn/row--orrhoooooo matrices and held fixed in each experiment,

---

Published as a conference paper at ICLR 2024
(a) Varying noise level
(b) Varying rep hidden dimension
(c) Varying depth of rep
Twndfome ape .dge
Tandfrmer viidee
Tardfomer
0.8
0.8
loss
0.6
0.6
0.6
0.4
0.4
0.4
0.0
0.0
0.0
20
40
0
100
20
30
40
10
20
30
40
contott
context
conteet
Figure 2: Test ICL risk for learning with representations. Each plot modifies single problem parameter from the base setting (L, o) = (2, 20, 0.1). Dotted lines plot the Bayes-optimal risks for each setting respectively.
(a) Probe 4* (xx,)) at x, tokens
(b) Probe @' (x.) ) tokens
(c) Probe y;
at x; tokens
mean sqaree emer
10
10
0
6
6
10
12
0
2
6
8
10
12
layer
Figure 3: Probing errors for the learning with representation setting. Each setting modifies one or two problem parameters from the hase setting (I, D) (2, 20, 0.1). Note that the orange curve corresponds to the same setting (and thus the transformer) across (a,b,c), well the red curve.
and slope = 0.01. We test L E {1, 2, 3, 3, hidden dimension D D (5, 20, 80). and noise level {0, 0.1, All experiments P. N(O, L2), 72 1, 20, and N 41.
We use a small architecture within the GPT-2 family with 12 layers, heads, and Daia = 256, following (Garg et al., 2022; Li et al., 2023a; Bai al., 2023). The (pre)-training objective for the transformer (for the supervised learning setting) is the average prediction risk at all tokens:
ming Bw, Pur
1 22
Σ
(Jo, - y:)"
(10)
where Te.: is extracted from the (2i 1)-th output token of TFo (H) (cf. Section 3). The objective for the dynamical system setting is defined similarly. Additional experimental details can found in Appendix D, and ablation studies (c.g. along the training trajectory; ef. Figure 9) in Appendix F,
5.1 SUPERVISED LEARNING WITH REPRESENTATION
We first test ICL with supervised learning data as in Section 4.1, where for each configuration of (L, D, a) (which induces 4*) we train transformer on ICL data distribution (2) and evaluate ICL
on the same distribution. Note that Figure Ic & 1b plots the results for (I, D, = (2, 20, 0.1).
ICL performance Figure reports the test risk across various settings, where we observe that trained transformers can consistently match the Bayes-optimal ridge predictor. This extends existing results which show that linear functions (without representation) can be learned near-optimally in- context by transformers (Garg et al., 2022; Akyürek et al., 2022), adding our model (2) to this list of (empirically) petrll-op learnable function classes. Among the complexity measures (L, D, o), observe that the noise level a and hidden dimension D of the representation (Figure 2a & 2b) appears have larger effect on the (nearly Bayes-optimal) risk than the depth I. (Figure 2c).
Mechanisms via linear probing We conduct probing experiments to further understand the mech- anisms of the trained transformers. In accordance with the theoretical construction in Theorem 1, main question here is: Does the trained transformer perform the following in order:
1. Computes o (x;) at s; tokens;
2. Copies them onto the following token and obtains dataset {T* (x,), ve in the form of (4);
3. Performs linear ICI top of ( (x;), v:),
7

---

Published as a conference paper at ICLR 2024
(a) Illustration of the pasting experiment
(b) Linear ICL in TF_ _upper via pasting
Original Model
9: X
Pasting
x=x
1.2
Trupeend to ombed
by
1.0
IeJesen sace cary Toan, Liaee cottral sigee
0.0
8
TF Upper Layers
Jame
TF Upper Layers
[Frozen]
Trainable linear/TF ambed.
0.2
TF Lower Layers
(0(,0)) Linear Modal
0.0
30
(6,,))01 - Lincar 0" Model
Figure 4: (a) Illustration of our pasting experiment, which examines the linear ICL capability of the upper module of a trained transformer. (b) Pasting results for the upper module of trained transformer in setting (L, D) (3, 20, 0.1). "TT SrF_upper+... with different embeddings. It achieves nearly optimal linear ICL risk (in 20 dimension with noise 0.1), using I-layer transformer embedding. and also non- non-trivial performance using linear and linear copy embeddings.
corresponn to feeding the upper module of trained transformer
While such internal mechanisms are in general difficult to quantify exactly, we adapt the linear prob- ing (Alain & Bengio, 2016) technique to the transformer setting identify evidence. Linear probing allows us to test whether intermediate layer outputs (tokens) (h, (+) Jee[ra] (E denotes the layer) and
(h. [eg[ia] "contains" various quantities of interest, by linearly regressing these quantities (as the y) on the intermediate tokens (as the x), pooled over the token index
[N]. For example, gressing 4= (x:) on h (e.(() tests whether the x; token after the é-th layer "contains" 4 (x:), where a smaller error indicates better containment. See Appendix D.I for further setups of linear probing,
Figure reports the errors of three linear probes across all 12 layers: The representation 4* (xx) in
the x; tokens and y tokens, and the optimal ridge prediction U/ o0 . in the x; tokens. Observe that the probing errors for the representation decrease through lower layers and then increase through upper layers (Figure 3a & 3b), whereas probing errors for the ridge prediction monotonically decrease through the layers (Figure 3c), aligning with our construction that the transformer first computes the representations and then performs ICL on top of the representation. Also note that deeper represen- tations take more layers to compute (Figure 3a). Further, the representation shows up later in the y- y-tokens (layers 5-6) than in the r-tokens (layers 1,3,4.5), consistent with the copying mechanism, albeit the copying appears to be lossy (probe errors are higher at y-tokens).
Finally, observe that the separation between the lower and upper modules seems to be strong in certain runs-For example, the red transformer (L 4, 0.1) computes the representation at layer 5, copies them onto y-tokens at layer 6, and starts to perform iterative ICL. from layer 7, which aligns fairly well with our theoretical constructions at high level.
Investigating upper module via pasting To further investigate upper module, we test whether it indeed strong ICL learner on its own without relying on the lower module, which would provide stronger evidence that the upper module performs linear ICL. However, key challenge here is that it is unclear how to feed raw inputs directly into the upper module, as they supposedly only admit input formats emitted from the lower module-thh part we wanted to exclude in the first place.
We address this by conducting a pasting experiment, where we feed D-dimensional linear ICL problems (y: = (w'x without representation) with input format (1) directly to the upper module of the transformer trained on representation h², by adding a trainable embedding layer in between; Figure 4a for an illustration of the pasting approach. This trainable embedding layer itself needs to be shallow without much ICL power we test the following three choices: (1) Linear embedding:
W[x; 0) and hy = W [O yıl: (2) Linear-copy embedding, where the tokens are instead h W xx y.). motivated by the format (4); (3) One-layer transformer embedding TF, which computes H TF (H). See Appendix D.2 for further setups of pasting.
h.
Figure 4b shows the pasting results on a trained transformer on (L, D.o) = (3, 20, 0.1) (an ablation in Figure 10b), where we dissect the lower and upper modules layer 4 suggested by the probing curve (Figure 3a green). Perhaps surprisingly, the upper module of the transformer can indeed perform nearly optimal linear ICL without representation when we use the one-layer transformer
8

---

Published as a conference paper at ICLR 2024
(a) Risk
(b) Probe past inputs at x, tokens
(c) Probe d ('(xi-5) at x. tokens
Traoforner ozt -ddge
10"1
0.8
O
20.0=01
squared error
20.0- 0.
10-1
10 05
oii 4)
0.4
10"3
10-2
0.0
40
10
4
6
8
1.0
12
10
12
in-context examp los
layer
Figure 5: ICL risks and probing errors for the dynamical system setting. (a) Each curve modifies problem parameters from the base setting (k:, L, o) = (3, 2, 80, 0.5). (b,c) Results are with the same base setting.
embedding. Note that a (freshly trained) single-layer transformer itself performs badly, achieving about the trivial test risk 1.01, which is expected due to our specific input format (1). This suggests that the majority of the ICL is indeed carried by the upper module, with the one-layer transformer embedding not doing much ICL. itself. Also note that the linear-copy and linear embeddings also yield reasonable (though suboptimal) performance, with linear-copy performing slightly better.
5.1.1 EXTENSION: MIXTURE OF MULTIPLE REPRESENTATIONS
We aditionally investigate an harder scenario in which there exists multiple possible representation functions (d)) and the ICL data distribution is mixture of the K distributions of form (2) each induced by $ (equivalent to using the concatenated representation क" = ......., Φi with group 1-sparse prior w E KD). We find that transformers still approach Bayes-optimal risks, though less so compared with the transformers sometimes implement the post-ICL algorithm selection mechanism identified in Bai et al. (2023), depending on the setting. Details are deferred to Appendix E due to the space limit.
setting. Using linear probes, we find that
5.2 DYNAMICAL SYSTEMS
We now study the dynamical systems setting Section 4.2 using the same approaches as in Sec- tion 5.1. Figure 5a shows that transformers can still consistently achieve nearly Bayes-optimal ICI. risk. An ablation of the risks and probing errors in alternative settings can found in Appendix F.2.
Probing copying mechanisms ration phase, where the transformer construction in Theorem 2 performs copying twice:
The main mechanistic question we ask here is about the data prepa-
i) A copying of [x; (x; ;x,, onto the x; token as in (7), to prepare for the computation of x* (xx):: As copying may not be distinguishable from the consequent matrix multiplication step [x [x- x; + [i
1;- x:). we probe instead the result -i--
after matrix multiplication, where Bū,-; E Rxd denotes the block within Bj hitting X,-j-
ii) A second copying of 4 1x; onto the x; token to obtain (9), after { (xx)); are computed.
We probe one transformer trained the dynamical systems problem with k (so that the useful preceding inputs are X;-1 and x;- and find that the transformer indeed performs the two conjec- tured copyings. Figure 5b demonstrates copying onto the current token, where the copying of x;-1 happens earlier (at layer 3) and is slightly more accurate than that of x;-2 (at layer 4), as expected. Further observe that layer 4 (which we recall contains an attention layer and an MLP layer) have seemingly also implemented the (unnormalized) MLP representation $* (xx)) = o (Bia, (B;x)), though the probing error for the actual representation x (x,) T /xxx)//11* (xx) continues to drop layer 4-6 (Figure 5c). Figure Sc further demonstrates copying ii), where & (x;-1) are in- deed copied to the i-th token, whereas by sharp contrast xx (xxxxx)) for k > 2 are not copied at all into the x; token, aligning with our conjectured intermediate output format (9).
A one-layer transformer does not have much ICL power using input format (1) Xs and are stored in separate tokens there, which makes "one-layer" mechanisms such as gradient descent (von Oswald et al., 2022; Akyürek al., 2022; Bai al., 2023) unlikely be implementable; see Appendix D.3 for discussion.

---

Published as a conference paper at ICLR 2024
6 CONCLUSION
This paper presents theoretical and mechanistic studies on the in-context learning ability of trans- formers on learning tasks involving representation functions, where we give efficient transformer constructions for linear ICL on top of representations for the supervised learning and dynamical system setting, and empirically confirm the existence of various high-level mechanisms in trained transformers. We believe our work opens up the investigation of ICL beyond simple function classes, and suggests open questions such as further investigations of the mechanisms of the linear ICL mod- ules, and theory for ICL in more complex function classes. One limitation our work that the setting still consists of synthetic data with idealistic representation functions; performing similar studies on more real-world data would be an important direction for future work.
ACKNOWLEDGMENT
WH acknowledges support from the Google Research Scholar program. S. Mei supported by NSF DMS-2210827, CCF-2315725, NSF CAREER DMS- DMS-2339904, and Amazon Research Award.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn implement preconditioned gradient descent for in-context learning. arXiv preprint arXiv:2306.0020 002977 2023,
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm in-context learning? investigations with linear models. arXiv preprint 1..Xiv:221115660 2022.
Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes. arXiv preprint 1rXiv:1610.01644 2016.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers statisticians: Prov- able in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.04637, 2023.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. arXiv preprint 2xXiv:2300.00800, 2023.
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.
Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends@ in Machine Learming, 8(3-4):231-355 2015.
Sébastien Bubeck, Varun Chandrasckaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Ka- mar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303. 2023.
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learn- ing in transformers. Advances in Neural Information Processing Systems, 35:18878 188888888888, 2022.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Zhifang Sui, and Furu Wei. Why can gpt learn in- context? language models secretly perform gradient descent as meta optimizers, arXiv preprint arXiv:2212.10559, 2022.
N Elhage, N Nanda. C Olsson. T Henighan, N Joseph, B Mann, A Askell, Y Bai. A Chen, T Conerlly, ct al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021.
Shivam Garg, Dimitris Tsipras. Percy S Liang, and Gregory Valiant. What transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
10

---

Published as a conference paper at ICLR 2024
Atticus Geiger, Hanson Lu, Thomas leard, and Christopher Potts. Causal abstractions of neural networks. Advances in Neural Information Processing Systems, 34:9574-9588 2021.
Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose in-context learning by meta-learning transformers. arXiv preprint urXiv:2212.04458 2022.
Yingcong Li, M Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers al- gorithms: Generalization and implicit model selection in in-context learning. arXiv preprint arXiv:2301.07067 2023a.
Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting chain-of-thought A study compositional in-context learning of mlps. arXiv preprint arXiv:2305. 18869, 2023b.
Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for gpt-3? arXiv preprint arXiv:2 068044 2021.
Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. arXiv preprint arXiv:2104.08786 2021.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems, 35:17359-17372, 2022.
Sewon Min, Mike Lewis, Hannaneh Hajishirzi, and Luke Zeutlemoye. Noisy channel language model prompting for few-shot text classification. arXiv preprint arXiv:2108.0411 2021a.
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. Metaicl: Learning to learn in context. arXiv preprint arXiv:21 15943, 2021b.
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 2tt.:22202./8..77 2022.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 2022.
OpenAl. Gpt-4 technical report. arXiv preprint arXiv:3333.08774, 2023.
Tilman Räuker, Anson Ho, Stephen Casper, and Dylan Hadfield-Mena Toward transparent ai: A survey interpreting the inner structures of deep neural networks. In 2023 IEEE Conference Secure and Trustworthy Machine Learning (SaTML), pp. 464-483. IEEE, 2023.
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. arXiv preprint arX1v:230..1506,, 2023.
Yasaman Razeghi, Robert L Logan IV, Matt Gardner, and Sameer Singh. Impact of pretraining term frequencies on few-shot reasoning. arXiv preprint arXiv2200.002066 2022.
Ohad Rubin, Jonathan Herzig, and Jonathan Berant. Learning to retrieve prompts for in-context learning, arXiv preprint arXiv:2112.08633 2021.
Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian. A study relu and softmax in transformer. arXiv preprint arXiv:2302.0666)) 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint drXiv:2307.09288 2023.
Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shieber. Investigating gender bias in language models using causal mediation analysis. Advances in neural information processing systems, 33:12388-12401 2020.

---

Published as a conference paper at ICLR 2024
Johannes Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov. and Max Vladymyrov. Transformers lear in-context by gradient descent. arXiv preprint arXiv:2212.07677 2022.
Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Inter- pretability in the wild: a circuit for indirect object identification in gpt-2 small. arXiv preprint arXiv:2211.00595 2022.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thougle prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824 24837, 2022.
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.0380
2023,
Mitchell Wortsman, Jachoon Lee, Justin Gilmer, and Simon Korblith. Replacing softmax with relu in vision transformers. arXiv preprint arXiv:239.008886 2023.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint 1eXiv:2//1.02000 2021.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett, Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.099) 2023a.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305. 10.55.9440 19420, 2023b.
Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving few-shot performance of language models. In International Conference on Machine Learning, pp. 12697-12706. PMLR, 2021.
A TECHNICAL TOOLS
The following convergence result for minimizing smooth and strongly convex function is standard from the convex optimization literature, e.g. by adapting the learning rate Bubeck (2015, Theorem 3.10) from 1) 1/B any 1) 1/B.
Proposition A.1 (Gradient descent for smooth and strongly convex functions). Suppose L Rd + is a-strongly convex and 3-smooth for some < S B. Then, the gradient descent iterates w 1+1
GD = W L ((w()) GD) with leurning rate S 1/B and initialization w E R& satisfies for
any 2
WOD exp t) ||W&D w'11.-
where w == arg min L(w) is the minimizer of L.
B PROOFS FOR SECTION 4.1
Throughout the rest of this and next section, consider transformer architectures defined in Sec- tion 3 where we choose to be the (entry-wise) ReLU activation normalized by sequence length, following (Bai et al., 2023): For all A N and i, [N].
[a(A) = 3 (Ass),
(11)
where we recall ott) max (t, 0} denotes the standard ReLU. This activation is similar as the softmax in that, for every (query index) j. the resulting attention weights (}o(AA)) sels is approxi- mately a probability distribution in typical scenarios, in the sense that they are non-negative and sum to O(1) when each A; O(1). We remark that transformers with (normalized) ReI.U activation is
12

---

Published as a conference paper at ICLR 2024
recently shown to achieve comparable performance with softmax in larger-scale tasks (Shen et al., 2023; Wortsman et al., 2023).
With activation chosen as (1I), a (decoder-only) attention layer H (H) with (Qm: Km, V [m]me|M] takes the following form in vector notation:
h: h; + Σ
Σ
((Qmh; Kmbj)) Vmhjj.
Recall our input format (1):
X1
0
XN
0
H
0 11
0
UN
E R
X2N
Pi
P
PN
P
We will use (hx (e(N) and (he, h.") sejN interchangeably to denote the tokens in (I), where h == h and h. := h2g Similarly, we will use (Pī KE[N] and (P*)RE(2N] interchangably to denote the positional encoding vectors in (I), where P2i-1 == p and P2i == p.P Unless otherwise specified, we typically reserve use i as (query, key) indices within [N] and k, as (query, key) indices within [2N].
We use the following positional encoding vectors for all [N]:
P [O - d-9: 1; 2: 1; (2: - 1) (2i 1);;;;;;;;1)
(12)
P 0 Dmmaa-d-gi (2 (
Note that Pş contains [1; k; A:3 A:) for all k: E [2N]; P P contains [i; indicator of being an x-token, and the product of the indicator and i.
B.1 USEFUL TRANSFORMER CONSTRUCTIONS
Lemma B.1 (Copying by single attention head). There exists single-head attention layer O
(Q, K, V) C Pmie Drid that copies each x, into the next token for every input of the form (I), i.e.
X1
Xi
XN
XN
Attno (
V1
0
UN
R Dmn X2N
P
P
PV
Proof. By assumption of the positional encoding vectors, can define matrices Q,K R Drid D Daad such that for all k, l E [2N].
Qh kk 0 Op al.
Kh, [-1;28+ -e 26: 0 Piua a].
-2:
This gives that for all S k.
a((Qha, Kh₂))
= =o(-k + 2 (2e 2) kel22 + 20)) a (k((1 - - 6- 1)3) 1: = 1}.
Further defining V such that Vha = |x;; and Vh = 0, we have for every A E [2.V] that
Σ Fo ((Qhs. Vhe
8Sk
-i(f== k 1 1}. ( if odd); 0] [x/e/2 1fc 1 - and is odd).
By the residual structure of the attention layer, the above exactly gives the desired copying behavior, where every x; the odd token H is copied the next token.
13

---

Published as a conference paper at ICLR 2024
Bw, By 0, there exists attention layer = {Qmm Km mm) Vm]]me(A]] with M = heads such that the following holds. For any input sequence H E R Puiaxx X2N that takes form
Lemma B.2 (Linear prediction layer). For any Bu,
h [x:; 0; W:: P.). h
[x6 Oa; PK]
with |x:|22 S BR, luil S By, and |wll2 S B it gives output Ating (H) = E Pmaxx2N with
h h2: (x: w.: where D (x,, we)
for all i [N].
Proof. Let R = max [B Bw: By). Define matrices (Qm: Km, V m m==
w;
Xj
O2
22
K,hj = K,,
Vihe
l
R
2j+1
D -d-
0
-2R1
0g Dnid -d-1
Q2hr
R
K2h = K1hh-
2j
+1
.Vhh,
0
for i,j E [N] and [2N]. For every i [N], we then have
2 2i 1
Σ
0 ((Qmmi h Kmhh)) Vnhe
2i
m=1
t=1
1
w X + R(- 1)) - 4 + +))] [0a: 1; 0 Diiu-d-l
21
Ø=1
1-1
0 w Xj + 24 77-22++21 - 1)) - or ( 17(--2 + 2ÿ + 1))] (0,; 2j: Op (0ma-a--1) d-1)
j=1
2i 11 w. x 00;; 10; 1;a = 00, w. x,; 0 D
By the residual structure of an attention layer, the above shows the desired result.
Lemma B.3 (Implementing MLP representation by transformers). Fix any MLP representation function 4- of the form (3), suppose Duid 2 max (2D, + + 10}, where D is the hidden di- mension within the MLP (3). Then there exists a transformer TF, with (L + 1) layers and 5 heads that exactly implements 4 a token-wise fashion, i.e. for any input of form (I),
(x*((x,)
('(xx)
= (H)
y1
UN P
Df
P
P
where P7 D. differs from P. d only in the dimension of their zero paddings.
Proof. Recall that x* op (B; op o. (Bfx) ...). We first show how to implement single MLP layer x + a ( ) by an (MLP- (til--Atention) structure.
Consider any input token h = [x4 0; p at an x-location. Define matrices W Wg E RR Daia x Dand such that
Bix; B x;
o (
W =
(W,h;)) =
0 (-
0
0₂
O₂ op
W20 (Wh =
o (Bix))
po
Bix;)
=
(B
0
14

---

Published as a conference paper at ICLR 2024
Therefore, the MLP layer (W1, W2) outputs
o (Bfxi)
h W. w. (H)), h (20(W11))
0
P
and does not change the y-tokens.
We next define an attention layer that "moves" ap (B)x; to the beginning of the token, and removes
x, Define three attention heads 0 (Qm, K Km V. Vm)me(3] as follows:
k²
k
Q (122,3 hu
K 112,33
k1 (1 is odd}
1
Top )
X, Op
Od (Bi
Vihj
0a
V2hj
Vah
0
The values for V 1,2,3,, h are defined automatically by the same operations over the h tokens (which does not matter to the proof, as we see shortly). For any { & and m E [3].
( Tg. kmh,))
okk( + 1(k odd))) 1(E k. is odd].
Therefore, these three attention heads are only active iff the query token k 2i odd (i.e. being an z-token) and l = k = 2i 1. At such tokens, the three value matrices (combined with the residual structure of attention) would further remove the x, part, and move (Bfx:) to the beginning of the token, i.e.
To (B 0
h = [Attno (H)],
P
and h. = h. Additionally, we now add two more attention heads into 0 to move all from entry + to D and leaves the æ-tokens unchanged.
Repeating the above argument L times, we obtain structure with five heads in each attention layer that exactly implements the & in a token-wise fashion. This structure can be rewritten as an (L++))-lyee -laee transformer by appending an identity { (Attention, MLP layer (with zero weights) (before, after) the structure respectively, which completes the proof.
-MLP- Atttentttooon)
B.2 IN-CONTEXT RIDGE REGRESSION BY DECODER TRANSFORMER
This section proves the existence of a decoder transformer that approximately implements in in-context ridge regression at every token [N] simultaneously. For simplicity, we specialize our results to the ridge regression problem; however, our construction can be directly generalized to any (gener- alized) linear models with a sufficiently smooth loss, by approximating the gradient of the loss by sum of relus (Bai et al., 2023, Section 3.5).
Denote the regularized empirical risk for ridge regression on dataset D; = {(xx) V;)) by
( (w) ==
w X, - yン)
w
(13)
2i
jml
for all i E [N]. Let W.- == arg min wERd T (w) denote the minimizer of the above risk (solution
of ridge regression) for dataset D:-1- We further understand Σ (w) = 0 and w = Let [,(w) == L. (w) denote the unregularized version of the above risk.
Proposition B.4 (Approximating a single GD step by single attention layer). For any n 0 and any Br. B.., By > 0, there exists an attention layer = {QQms Km, Vm))me(M) with M = heads such that the following holds. For any input sequence H IR. Paaxx2N that takes form h
[x;; 0: w:.p...
hy [x4 ys: Oa: 0
15

---

Published as a conference paper at ICLR 2024
with |xx I Br 18dl I By, and |wll2 S Bw it gives oupput Attng (H) -H R Phod x2N with
h h = (x;; O; W. W. where
W, W; 76 V VL.
i-1 with nis 26-1 1 and hi = h. for all i [M].
Proof. Let := max Bw, ByJ. By the form of the input (he)RE|2N] in (I), we can define two attention heads {(Qm; Km, V mm)m=1,22 C Pai Duid such that for all i E [N],
w:/2 22
Xj
Oa+1 xj
3R
Qu
K K
Vih; Vih =
- 3R
j
0 Daid 12d--1
R
1
0
0
RR 1
0a+1 Xj
-3R
Q2h7 Q2h
K2hă K2hı
V2hj V2h3 17-
-R
D 1Pua--dd11 2d-11
0
Further, h takes the same form h except for replacing the w:/2 location with Og and replacing the - location with (using the indicator for being an x-token within p));; K,h takes the same form as K₁ except for replacing the yj location with O.
Fixing any i E [N]. We have for all j 1
K2h)))
=o =o(w] R
-(3i 3j 1)) ([((3i - 3j 1)) w, v/'x;/2- - 4/-
and for all j < i,
a a( h K - ((A,) K2h;))
=o(w, / + R(3i - 3j - 1)) - (R(3i 3j -1) = 2.]x 1
Above, we have used w x, y, < 3R/2, w x2/2| I R/2, and the fact that o(2 + M)
(M) equals e for M 2 |a| and for M |2]
Therefore for all j < i 1, ((Qihī hi K, / + t (22h K2h
= h, K h1)) ((Q K2h)))) - -ny 00.00011xii111 D 124-1]
=-n(w,'x, x2/2 12 vj) [0,+1i 0D 22d-11
and similarly for all j S a ((Q2h K,h ) V2h. + +o (,h, V2h
=-n(w; x, /2) 12j < -i- 1) | 0 2d111
Summing the above over all key tokens & [2i 1], we obtain the combined output of the two heads query token 2i 1 (i.e. the i-th xttoken):
2i-1
1
o ((Qm h2. Kx) ha
2i
1=1
m=1
Σ
Σ
2i ( Kn) Vnhg
2i
j=1 m= 1,2
m=
1-1
7 w x2 Uz) + Σ
Si-1
Out1iX,; D 2d
2i
1=1
i
0441:
2i
(14)
16

---

Published as a conference paper at ICLR 2024
It is straightforward to see that, repeating the same operation at query token 2I (i.e. the i-th y-token)
would output 0 Dria since the query vector Q, contains [Oz; 0] instead of [w,/2; - | as Q.h.
We now define one more attention head (Q3: K3, Va) C RDD such that for all i: [2N]. E [N],
- 1/2
Od+1 nyww
(1 - /2
Qahk =
Kah
Vahh
Vzhj = 0 O
1 1/2
Dnid 2d-1
0
For any E k; we have
k
a(QQ3h Kzhe)) = =o(k)) k(1
0
lfe k}.
2
2
Therefore, for query token
2i the attention head outputs
k
1
((Q Kzhe) me - - 2
-1
lfe k3 Vnhe
(15)
i-1
i-1
V h
V h
[0a+1: w,, 0 O 2d 2d-
2k
2i
straightforward to see that the same attention head query token outputs O as the value vector Vzh;- = Vzh is zero.
Combining (14) and (15), letting the full attention layer 0 == {(Qm, Km, Vm))m=12,5 we have (H) H. where for all [N].
2i
h h2i-1 h2-
((Qm Kmhb)) V he
Xi
0a+1
X;
0
i-
0
+
V L, (w,) Aw:)
W
2i
7h L (w:)
0 Drid 2d-1
i-1 where 7) = 21-1 Ws. and h2 h. This finishes the proof.
Theorem B.5 (In-context ridge regression by decoder-only transformer). For any X 2 O, Br, Bio By 0 with 1 + B / A. and
B..Bvo/. 2. let Dnid 2d + 10, then there ists an L-layer transformer TFo with M = 3 heads and hidden dimension Duid, where
3x log BB, ((=))] 22
(16)
such that the following holds. On any input matrix H form (1) such that problem (13) has bounded inputs and solution: for all E [N]
||x:l12 S.B Br,
198l < By,
| < Bw/2,
(17)
TFa approximately implements the ridge regression algorithm (minimizer of risk (13)) at every token E [N]; The prediction h. = (T (H) (1))41122211 satisfies
W X4
SE
(18)
Proof. The proof consists of two steps.
Step 1 We analyze the convergence rate of gradient descent ΣP11 simultaneously for all 2 S i < N, each with learning rate πソ: 3 implemented in Proposition B.4.
Fix 2 < i < N. Consider the ridge risk T defined in (13), which is convex quadratic function
that is A-sronggyy convex and Xuaa (x7_ X 1)) + X < B² + A 3 smooth over Rd Recall 3/ 1 + /
17

---

Published as a conference paper at ICLR 2024
Consider the following gradient descent algorithm D Initialize w.. 0, and for every t w W w 17: VL (w().
(19)
1-1 with 7 2--1 Taking 7) ;= 2/3, we have ni E (3/(3 1/3). and thus 7; [2/(3k), 1/k:.
By standard convergence results for strongly convex and smooth functions (Proposition A.1), we have for all 2 1 that
|| 11 I exp (-mpAt)))) |/ - will2 -ep (-m |
Further, taking the number of steps as
B Br
T : 3x 3x
2
so that 1/AT/2 2 2/(k)) 3x log( B Bw/ (2:) log( B (2 we have
!! w S exp (-1x11/2) w S B,B. 2
28
B
8
(20)
B
Step 2 We construct a (T + 2a-layer transformer TFo by concatenating the copying layer Lemma B. T identical gradient descent layers as constructed in Proposition B.4, and the linear prediction layer in Lemma B.2. Note that the transformer is attention only (all MLP layers being zero), and the number of heads within all layers is at most 3.
The copying layer ensures that the output format is compatible with the input format required in Proposition B.4, which in uurn ensures that the gradient descent layers implement (19) simulta- neously for all 1 S i < N (wf 0 is not updated at token i 1). Therefore, the final linear prediction layer ensures that, the output matrix H TF, (H) contains the following prediction every i E [N]:
T = h d+1 xx),
which satisfies
M -w xx
(w: w/)xx)) S (E/B2) B₁ E.
This finishes the proof.
B.3
PROOF OF THEOREM 1
The result follows directly by concatenating the following two transformer constructions:
The MLP implementation module in Lemma B.3, which has (L+ + 1)-layers, 5 heads, and trans- forms every x, to 4* (x;) to give output matrix (4);
The in-context ridge regression module in Theorem B.5 (with inputs being { (xx)) instead of
x2) which has (w log( By Bw//e) layers, 3 heads, and outputs prediction Øi := [hā] where
0 (Φ" (I"(x;) (x:), W. ) S E, where w; the (Q"-Ridge) predictor.
Claim (4) can be seen by concatenating the (L+1)-layee MLP module with the first layer in the ridge regression module (Theorem B.5), which copies the (*(x;) in each token to the same location in the succeeding y token.
Further, the hidden dimension requirements are Daad 2 max (2D, D + d 10) for the first module and Dbia 2 2D + 10 for the second module, which is satisfied at our precondition D = 2D + d + 10. This finishes the proof.
L
C PROOFS FOR SECTION 4.2
Recall our input format (6) for the dynamical system setting:
X
E R DalaxN
P1
PN
our choice of the positional encoding vectors P. = 00 Daa d-4; 2;;;; for all i [N], and that understand x; 0 for all i 0.
18

---

Published as a conference paper at ICLR 2024
C.I USEFUL TRANSFORMER CONSTRUCTIONS
Lemma C.I (Copying for dynamical systems), Suppose Dnid 2 kd + 4. For any k E [N], there
exists (k + 1)-head atlention layer 0 = (QQm, Km, Vm) me(km) D Dma such that for
every input H of the form (6), we have
X1-k+1
X,-k+1
XN-k+1
H Attnn
E PaiaNN
(21)
X1
X
XN
P
PN
where P; only differs from pi in the dimension of the zero paddings. In words, Attng copies the k 1 previous tokens ....; onto the i-th token.
Proof. For every 1' E [k]], define attention head (Qu. KE Vv) C R Phia Daia such that for all j Si [N]
Qc:h: = i;; -3),
11 (K' - 1 Daa- h =[-1;2 + 2(k' - 1); 1)j 1
Vwh = 10 (0)=
Note that
o( h, K h,)) o0-13 + 2(K 1)i² -i² + 2ij(K - - i(k' 1)?)
i - -iiikka 1))3) (1j=) +1].
Therefore, at output token E [N]. this attention head gives
a ((Q k, = - Veh 10:
when i - + 2 1, and zero otherwise. Combining all & heads, and defining one more head (Q K411, V V++1) to "remove" x; at its original location (similar in the proof of Lemma B.3), we have
X k+1 Xi
X4
(1+1)+1
X;
0
By the residual structure of an attention layer, we have
x - x; X;
X
x
(1-(111)
x; p:
[Attne (H)),
+
X;
X;
P
(The precondition Daid 2 D + 4 guarantees that the entries would not interfere with the non-zero entries within Pt.) This the desired result.
Lemma C.2 (Implementing MLP representation for dynamical systems). Fix any MLP representa-
tion function Φx Red - RDD of the form (3), suppose Duia 2 2(k + 1)d + 3D 2d + 5. Then there exists a module MLLP-Attention-ML . -Attention-LLP) with L+ (Attention-Me MMPP) blocks
(i.e. transformer layers) and 5 heads in each attention layer (this is equivalent to an (L + 2) -ayley
transformer without the initial attention layer) that implements Th" in the following fashion: For any input H of form
H
X1 P
XN
n
19

---

Published as a conference paper at ICLR 2024
where recall x;
[x- :x.] Rdd, the following holds. The first MLP layer outputs
Top (Bix))
(Bix,
MLP
(H)
X1
Di
p
The full transformer outputs
(x1x1)
T' (x2)
$ (x:)
0a
0d
02
H (H)
Op
s+ (xx)
4
-x,
(22)
XI
X2
X;
P1
P2
P;
where Di, Pi differs from P Dj only in the dimension of their zero paddings.
Proof. We first construct the first MLP layer. Consider any input token h, matrices W1, W2 IR Pra Dhid such that (below +u = [u: (u)
[;: Define
x
a ( Bix;)
+x;
(+x,)
W,h,
o (Wih,) =
(+x,)
0n
W20(W1h;) = (B(x)) p (-Bix;)
+ (xx)) + o [x(x)x)
+ o(xi) - (-4)
0
Therefore, the MLP layer (W1, W2) outputs
Top (Bix,)
h. = [MLP w. (H)), = h, W20(W,h))
X;
(23)
P
The requirement for Daid above is Daia 2 max (2D + 2(k 1)d, D + d + 5].
The rest of the proof follows by repeating the proof of Lemma B.3 (skipping the first (MLP- Attention) block), with the following modifications:
Save the x; x; location within each token, and move it into the (2D + 1: 2D + 2d) block in the final layer (instead of moving the label in Lemma B.3); this takes the same number (at most of attention heads in every layer, same in Lemma B.3.
Append one more copying layer with a single attention head (similar as the construction in Lemma C.I) to copy each x* to the (D d + 1: 2D + d) block of the next token.
The above module has structure (L 1) which can be rewritten as an MLP-Attention-MLP module with identity MLP layers. Alto- gether, the module has an MLP- (ottexAAttnnioon-NM tion heads within the above module is 5. The required hidden dimension here D > max (kd - 2D + d max (D,d) 5}. with Dria 2 max (kd, 3D + 2d) + 5 being a sufficient condition.
followed by a single attention layer
structure. The max number of atten-
Combining the above two parts, sufficient condition for Drid is Dria 2 2(k: 1)d 3D + 2d as assumed in the precondition. This finishes the proof.
Consider the following multi-output ridge regression problem:
1
A |w
W 2(i -
arg min
x,
(24)
We
20

---

Published as a conference paper at ICLR 2024
Theorem C.3 (In-context multi-output ridge regression with alternative input structure). For any 2 0, Bz, Bw By > 0 with x. := + B²// X, and E < B,Bw/2 let Duid 2 Dd + 2(D + d) + 5, then there exists L-layer transformer TFa with M = 3d heads and hidden dimension Dnia. where
L = log Bg Bw/(e)))
(25)
such that the following holds. On any input matrix
X
X2
XN
0₂
Oa
0₂
H
Op
X1
XN-1
Oa
Y1
YN-1
P1
P2
PN
(where x; E IR yi IRd) such that problem (24) has bounded inputs and solution: for all i E [N]
xxil2 S Bz,
llllllso S By,
W S B
(26)
TFg approximately implements the ridge regression algorithm (24) at every token E [N]: The prediction y; == (D))))))))))))))))))))))))) +1)) satisfies
8 (W:) Se.
(27)
Proof. Observe that the multi-output ridge regression problem (24) equivalent to d separable single-output ridge regression problems, one for each output dimension. Therefore, the proof fol- lows by directly repeating the same analysis as in Theorem B.5, with the adaptation that
Omit the copying layer since each token already admits the previous (input, label) pair;
Use (s (,, (e./://())))ale transformer with 3d heads to perform d parallel ridge regres- sion problems (each with 3 heads), using in-context gradient descent (Proposition B.4) as the internal optimization algorithm, and with slightly different input structures that can be still ac- commodated by using relu to implement the indicators. Further, by the precondition (26) and Daid 2(D + d) 5 2 Dd, we have enough empty space to store the W E Dxd within the zado-paddings in p..
Use a single-attention layer with d parallel linear prediction heads (Lemma B.2), one for each j [d]. to write prediction (Fr) into location (i, D + j) with 1(5+), - ((W?), xiii s s. Therefore,
| (WA) x
(v)), ( xx)))),>;)
max JE[d
This finishes the proof.
C.2
PROOF THEOREM 2
Proof of Theorem 2. The proof is similar as that of Theorem 1. The result follows directly by con- catenating the following three transformer modules:
The copying layer in Lemma C.I, which transforms the input to format (21), and thus verifies claim (7).
The MLP representation module in Lemma C.2, which transforms (21) to (22). Together with the above single attention layer, the module is now an (L + 1) 1l-layer transformer with heads. Claim (8) follows by the intermediate output (23) within the proof of Lemma C.2.
The in-context multi-output ridge regression construction in Theorem C.3 (with inputs be- ing {Ø (xx)) and labels being (x++1)). This TF has O(k log( Ba Bw.ee) layers, and 3d heads. It takes in input of format (22), and outputs prediction ÿ; = [h:]D41:D4 where [| (W:
0 .yTTT T & E, where W A is the ($* Ridge--DD) predictor.
21

---

Published as a conference paper at ICLR 2024
The resulting transformer has max 53 heads, and the hidden dimension requirement is Dnid
max {kd + 5, 2(k + + + 5, Dd - 2 + d) + +5). A sufficient condition is Dhid
max {2(k- + 1), D). + d) + 5, assumed in the precondition. This finishes the proof. m
D DETAILS FOR EXPERIMENTS
Architecture and training details We train 12-layer decoder model in GPT-2 family with 8 heads and hidden dimension Dnid 256, with positional encoding. We use linear read-in and read-out layer before and after the transformers respectively. both applying same affine trans- form to all tokens in the sequence and are trainable. The read-in layer maps any input vector to a Daia- -ineenimenii hidden state, and the read-out layer maps -ia-ndimeniioal hidden state 1-dimensional scalar for model (2) and to d-dimensional scalar for model (5).
Under the in-context learning with representation setting, first generate and fix the represen- tation 4* For a single ICL instance, We generate new coefficients w and N training examples {(x,, (e and test input (xN+1, $N+1). Before feeding into transformer, we re-format the sequence HICL-reps as shown in equation (28).
HICL-rep
X1:
3/1 0a-1
N
VN Oa-1
RRd22N
(28)
We use the use the Adam optimizer with a fixed learning rate 10-4, which works well for all ex- periments. We train the model for 300K steps, where each step consists of a (fresh) minibatch with batch size 64 for single representation experiments, except for the mixture settings Appendix where we train for 150K iterations, each containing K batches one for each task.
Under ICL. dynamic system setting, for a single ICL instance, we don't need to reformat the input sequence. We feed the original sequence H Dynamic [x1 :X::-..,NN] R to transformer.
Generating representations: Denote the column- orthogonal matrices with size rx (r s) as Orr, s). We sample (ni s)] through generating an r x matrix, getting its QR decompo- sition, then taking the first s columns of the Q matrix. To generate L layers MLP with hidden dimension D. We first generate L weight matrices B,, f= l. ......,, L. with B ~ (n, D) (if d D. we sample B (n, D)D and the rest independently from Unif[O(D, D)}. After calculating MLP(x, (x). we self self normalize it with |MLP(x)]|2.
All our plots show one-standard -e--stadarrdddviittion error bars, though some of those are not too visible.
D.1 DETAILS FOR LINEAR PROBING
Denote the L-th hidden state of transformers
H (
h v.(
2.(6) IN
v.()) Pmai2N for E [12].
Denote the probing target as g((x,, U3 seta) E Rdrooe for i E [N]. Denote the linear probing parameter as ((.) and w0 that belong to R. Puaaxdprobes
Denote the best linear probing model as
w (. = arg min E
Σ (w))e) T nn:(00 (1x),8113e1s) and
wo.ck)
N
w (9)
min E
(wwze) T hy,(00) - ((xx))))( (it.)))))))))
To find them, we generate 2560 ICL input sequences with length N, and obtain 12 hidden states for each input sequences. We leave 256 sequences as test sample and use the remaining samples to estimate w. =.(() and 7.) for each with ordinary least squares. We use the mean squared error to measure the probe errors. In specific, define
a.(5 (9) = -B ([((w T (12,00) - o (1xx),053)11)) with
Probe Error,
22

---

Published as a conference paper at ICLR 2024
N
Probe Biror":(00)))
=
N
Probe Error;
, (g), and
;
Probe Error
(oroor:(00() (9) (g) E.
{ (w:)
h ((x))) (Era))
with
Probe Error (.0) (g) N ==1
Probe Error,
Error. (g). (g).
v.
control to the probe emrors in the hidden layer. We normalize each probe error with E [ v) 2 We use the 256 leaved-out samples to estimate these errors. We replicate the above procedure for three times and take their mean to get the final probe errors.
When é = O, we let h; (0) h (0)
D.2 DETAILS FOR PASTING
From the single fixed representation settings above, we pick a trained transformer trained on the representation with D d 20 avoid dimension mismatch between 4 (x) and We choose
L = 3 and noise level a 0.1.
We change the data generating procedure of y from Equation (2)
Yi (w, x;) 0z2, i [N],
(29)
which corresponds to linear-ICL task. According to the results of probing Fig 3a, conjecture that transformer use the first layers to recover the representation, and implement in-context learn- ing through the 5-th to the last layers. Therefore, extract the 5 12 layers as the transformer upper layers. Then paste them with three kinds of embeddings:
1. Linear embedding W TR. Paix (D+1) with re-formateed input Hinear:
Hinear
x
Op
x
Op
D-122N
0
1/
UN
2. Linear copy embedding W E I8 Din (D+1) with re-formatted input Heopy that copies x, to y tokens in advance:
H Heopy
x
x1 yn
xN
X UN
D+1x2N
3. Transformer embedding TFF using the same input format HICL-rep with normal settings, as shown (28). We extract the 4-th layer of GPT-2 model, its complete transformer block with trainable layer norm. We use linear read-in matrix map HICI-rep to the Duia- dimension hidden state, apply one block of transformer to it to get the TF embedding H TF(H).
We apply the upper layers to the three embeddings, then use the original read-out matrix get the prediction of ye. For comparison, we also train a one-layer transformer using the input sequence HiCL-rep-
We use the same training objective as in (10). In the retraining process, we switch to task (29), fix the parameters of upper layers of the transformer, and only retrain the embedding model. The training methods are exact the same with the original transformer. We also find that using a random initialized transformer block extracting the 4-th layer of the transformer don't 't make difference to the results.
D.3 DIFFICULTY OF LINEAR ICL WITH A SINGLE-LAYER TRANSFORMER WITH SPECIFIC INPUT FORMAT
Recall the input format (I):
X1
0
XN
0
H
0
11
0
UN
LP
PT
P
PN_
23

---

Published as a conference paper at ICLR 2024
Here heuristically argue that single attention layer alone (the only part in single-layer trans- former that handles interaction across tokens) is unlikely to achieve good linear ICL performance on input format (1).
Consider a single attention head (Q, K. V). As we wish the transformer to do ICL prediction at every token, the linear estimator w; used predict is likely best stored in the x; token (the only token that can attend to all past data D and the current input x,). In this case, the attention layer needs use the following (key, value) vectors to compute good estimator w; from the data D;-1:
(Vhz, Wh5 se(s'' (Vhy, Vhy) jE(i-1])
However (apart from position information), h only contains x,, and h only contains 3/. Therefore, using the normalized ReLU activation as in Appendix B & C.2, it is unlikely that an attention layer implement even simple ICL algorithms such as step of gradient descent (von Oswald al., 2022; Akyürek et al., 2022):
W
E 71 1-1 J5
1
(w -1,)
which (importantly) involves term =y x, that is unlikely be implementable by the above attention, where each attention head at each key token can observe either x, yj but not both.
D.4 REPRODUCIBILITY
Code for our experiments is provided the following anonymous link
E
EXPERIMENTS ON MIXTURE OF MULTIPLE REPRESENTATIONS
We train transformers mixture of multiple ICL tasks. where each task admits different rep- resentation function. This setting is a representation selection problem similar as the "algorithm selection" setting of Bai et al. (2023). In specific, let K 2 denote the number of tasks. Given j. let
V1 = (w, (x,))) + aze 21 ~ N(0,1), E (N), where
(x) = o² (BB g+00) Bi.--.... ... -000) BB, x ...))...) BT E RDdd (B:())) c RxD,
The distributions for w, [xI)ie and (Br()) are same with previous setting. We generate dif- ferent D for j E [K] independently. We choose K E (3.6), 6), E (0.0..00.5). L = 3, and noise E 00,0.1,..00 0.5).
At each training step, we generate K independent minibatches, with the j-th minimatch takes the representation +; to generate (y: hegn)- Due to multiple minibatches, we shorten the number of total training steps to 150K. The other training details are the same with fixed single representation setting.
ICL performance We choose representation क from the representations that transformers are trained on. Figure 6a & Figure 6b report the test risk. We vary K E {3, 6} and noise level (0.1,0.5), We consider two baseline models.
1. The Bayes optimal algorithu:: Note that the training distribution follows the Bayesian hierar- chical model:
Unif([K]), N(O, La), w~ rL),, and yi Xi, w N((w, ,
This gives the Bayes optimal predictor
K
(i) with (....n?),, ......) =SO0=AAO {[ Σ (yn 2200)///02 ,... 11}
i =
(30)
with 07
being ridge predictor with optimal À based on (())(x)) )
hhttps://anonymouus 4open wwwp:e.w...eee.e....eeee..e..e wtt.eecce///ffeeeeeeeee//ee
24

---

Published as a conference paper at ICLR 2024
(a) Risk for K =
(b) Risk for K
1.0
0.8
0.8
Transforer
Square loss
Baseline Bayes optimal
0.6
0.6
ride
Square
0.4
0.2
0.2
0.0
0.0
10
20
30
40
0
10
20
30
40
in-context examples
in-context examples
Figure 6: ICL risks for multiple representations setting. Dotted lines plot two baseline risks. (a) The trasformer with lower risks is trained with (K, o) (3,3, 20, 0.1). The upper one is trained with L,,,, (3,3, 20,0.5) 0.5). (b) The two transformers are trained with and same settings otherwise.
2. The oracle ridge algorithm: We use the ridge predictor 7. based on (x,), (reir I)+ which is the representation for test distribution. Note that this is an (improper) algorithm that relies on knowledge of the ground truth task.
Comparable to those trained on single fixed representation, transformers consistently match the Bayes-optimal ridge predictor. As expected, the oracle ridge algorithm is better than transform- ers and the Bayes optimal algorithm and transformers. Increasing number of tasks K can slightly increase this gap. Increasing the noise level has the same effect on transformers and baseline algo- rithms.
(a) Probe (x.)
(b) Concatenated D and ()
(c) Last layer probe g(
tokers
109
tokens
meen squared error
rlokens
concatetated prediction.
10-1
bayaa ele gitss
Rep function
641
10-1
0.15
0.01
0.16
0.1.
0.12
0.01
2
6
8
10
12
layer
layer
Task
Figure 7: Probing errors for transformer trained with LK, D) = (3, 0.1). Dotted lines plot probing errors tokens.
23,20,011
(a) Probe (x,)
(b) Concatenated / and v)
(c) Last layer probe gf
meea Squares ercer
9,
trioens
tolers
0.01
0.03
10-1
meen squared
Rep function
10-1
0.82
0.01
0.02
0.02
0.01
6
10
layer
Task
Figure Probing errors for transformer trained with a 0.5
Probe setup: Similar single fixed representation setting, conduct linear probing experi- ments. We are wondering transformer implements the ICI.-larning on representations with algo- rithm selections mechanism. We identify three sets of probing targets: 4 oxxx),, gol) and n,(") All of them are intermediate values to compute the Bayes optimal estimator (30). We generate different data for different probing targets:
25

---

Published as a conference paper at ICLR 2024
1. To probe 4 (x,) and g(L) for each j, choose from the representations that transformers are trained on, then train and test our linear probing model. This is the same with the training and testing methods for probing transformers trained on single representation.
2. To probe choose concatenated probing targets Y,
1 ...... 1 and B, Ing (1) ...... .m ",""',, we generate 2560 in-context sequences for each representation. and obtain 2560 x K samples together. We use ordinary linear square on 2560 x K 256 samples to get the linear probing models. Then test them the remaining 256 samples to get the probing errors. We also repeat this process for three times and take means to get the final probing errors.
Probe representations: Take the transformer trained on K = mixture representations with noise level E {0.1, 0.53. Figure 7 show the probing errors for 0.1: Figure 7a reports the errors of probing $; j E [3]. with probing models trained on task $1- Echoing the results for transformers trained single representation, the probing errors for each representations decrease through lower layers and increase through upper layers on x tokens. The probing errors on tokens drop after x tokens, which suggests copy mechanism. Surprisingly, on x-tokens, the probing errors for all representations attain their minimum at the 3-th layer, with transformers trained on single representation achieving their minimum on 4-th layer (compare with Figure 3a).
More importantly, for both x and tokens. the probing errors for each representation are similar through lower layers, but the probing errors for the true representation &i become the lowest through the upper layers. The gap between the probing errors increases. At the last layer, the probing error for the other representations go up to match the initial input.
Probe intermediate values for computing Bayes optimal predictor: Figure 7b shows the prob- ing errors for concatenated ridge predictors g(2)) and Bayes weights v("), i.e., Y; and B, The probing errors for Y; start dropping the 4-th layer, which suggest that transformer are imple- menting ICL using each representations. Probing errors for B; have a sudden drop at the 10-th layer. Figure 7c shows the probing errors for probing x03 At (j, k)-th cell, we show the probing error of 7.5) with probing models trained on & at the x tokens of the last layer. The diagonal elements are much smaller than others. The results combined together suggest the possibility that transformer implements in-context learning with three representations and selects one at the 10-th layer.
In comparison, Figure 8 shows results of probing the same targets for transformer under = 0.5. Figure 8a differs with Figure 7b upper layers, where probing errors for different representations don't Ion't have significant differences. Figure 8b is close to Figure 7b, also suggesting the algorithm selection mechanism. Figure &c shows that the last layer encodes the information of all ridge pre- dictors ( }. which is drastically different from the results in Figure 7c.
Conjecture on two different algorithm selection mechanisms: Based on the empirical findings, we conjecture two possible mechanisms of algorithm selection in transformer: (1) For small noise level data, transformers implement "coneurrent- ICL algorithm selection", which means they concur- rently implement ICL with algorithm selection. They wouldn't 't finish algorithms that are not have good performances. (2) For large noise level data, transformers implement "post-ICL algorithm lection", which means they would finish each algorithm. Then they select and output the best one. However, we need further experimental and theoretical work inspect this conjecture.
F ABLATIONS
F.1 SUPERVISED LEARNING WITH REPRESENTATION
Probing results along training trajectory Figure 9a, Figure 9b, and Figure 9c show the probing error for 4 (xx) at x and y tokens and g "idgee at x tokens. As expected, all probe errors reduce through training steps, showing that the progress of learning 4* is consistent with the progress of the training loss. At the 2000 training steps, transformer cannot recover the representation. At the 5000 training steps, the transformer starts memorizing the representation, starting showing differences
26

---

Published as a conference paper at ICLR 2024
between lower and upper layers. From 5000 training steps 10000, the trend of probe errors varying with layers remains the same.
(a) 2000 training steps
(b) 5000 training steps
(c) 10000 training steps
Gd tokaro
tokero
de USeeTa
meansaaans onor
yeocene
meen vqgutred errOr
tnkeos
hooens
10
10
10
2
8
12
10
12
2
8
10
layer
layer
layer
Figure 9: Probing errors for transformer trained after 2000, 5000, and 10000 steps. All three plots are for the training run on (L, D, σ) = (2, 0.1).
Additional results for probing and pasting Figure 10a plots the same probing errors as in Figure 3a with (L, o. = (0, 0.1) (the green line there), except that we separate the errors of the first 4 tokens with the rest (token 5-41), but the probing training remains the same (pooled across all tokens). We observe that lower layers compute the representation pretty much the same ways, though later layers forget the representations more for the beginning tokens (1-4) than the rest tokens.
Figure 10b plots the same pasting experiment as in Figure 4b, except that for noise level o = 0.5 as opposed 0.1 therein. The message is mostly the same in Figure 4b.
(a) Probe errors per token
(b) Pasting experiment
point
1.4
uppep-lllperr T__eeper-Enebr,co ombced uppei-Innor eribed fieshg ay
point
point
1.2
probe aroor
square loss
1.0
0.8
0.6
0.4
10
0.2
2
6
8
10 12
layer
in-context examples
Figure 10: (a) Probing emrors of 6 (xx) in X; tokens evaluated per-token. (b) Pasting results for the upper module of trained transformer in setting (J,, D2 σ) = (3, 20,20,10 0.5).
F.2 DYNAMICAL SYSTEMS
Risk Figure 11 gives ablation studies for the ICL risk in the dynamical systems setting in Sec- tion 4.2. In all settings, the trained transformer achieves nearly Bayes-optimal risk. Note that the noise appears to have larger effect than the hidden dimension, or the number of input tokens.
(a) Varying noise level
(b) Varying rep hidden dim.
(c) Varying number of input tokens
1.0
1.0
Taodsmmr
fransomer ser vidge
Tardoweer
163
0.8
ter dige
61
0.8
00:50 50
0.8
0.6
0-05
0.6
0-
0.6
owon
0.2
0.2
0.2
0.0
0.0
0.0
contoxt exampies
examples
contoct exampies
Figure 11: Ablation studies for the risk for Risk for fixed rep setting. Each plot modifies single problem parameter from the base setting (k, L, D, .) 23, 20, 0.1).
27

---

Published as a conference paper at ICLR 2024
Probing Figure 12a & 12b gives ablation studies for the probing errors in the dynamical systems setting in Section 4.2, with D 20 instead of = 80 as in Figure 5b & 5c. The message is largely similar except that in Figure 12a, all past inputs and intermediate steps in 4 (X.) are simultaneously best implemented after layer 4.
(a) Probe past inputs x; tokens
(b) Probe 4 (x:-5) at x; tokens
10
(2)
mean squared error
0 -00--
1
0 (R-1
dia
9 olB; x)
(4--4
10-2
10
0
2
4
8
10
12
layer
layer
Figure 12: Ablation study for the probing errors in the dynamics setting. Here (k: L, D) 20,
different from Figure where D 80.
F.3 REPRESENTATION FUNCTIONS WITH HIGHER COMPLEXITY OR ALTERNATIVE STRUCTURES
Here we provide additional experiments when the representation function D' has a higher complex- ity or alternative structures (such as the nonlinearity o²) from the ones in the main text.
We consider three scenarios for representations: 1, deep MLPs with D = 20, L {8, 12, 16} and 0.1; 2. wide MLPs with D {100, 200, 300). L 20 and 0.1; 3. entry-wise nonlinear functions, where we replace MLP with entry-wise square, exponential, square leaky ReLU, and exponential leaky ReLU functions.
Figure 13 shows their risk curves. Both deep MLPs and other nonlinear functions match the optimal ridge regression. For wide MIPs, TF has slightly higher risk than optimal ridge regression. The risk curve of D = 100 is the lowest, with the other two on top of cach other. Figure 14 shows probing error for deep MLPs. seems that transformers would stop implement representations at the 7-th layer, with linear-ICL algorithm starts from 8-th layer.
Figure 15 shows the result of wide MLPs, we find that transformer needs one extra layer to get the representations if D = 300, which is wider than GPT2. Figure 16 shows for other entry-wise nonlinear functions. Transformer can get the representation within the first two layers, with copying happens in 4-th to 6-th layers. The gradient descent happens thereafter.
(a) Deep MI.Ps
(b) Wide MLPs
(c) Other Functions
1.0
1.0
1.0
Taastorrer
Tranaforms
0.8
0.8
opt
1-
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.0
20
40
0.0
10
20
30
In-oontet examples
Figure 13: Test ICL risk for lcarning with complex representations. MLPs, nonlinear functions.
From left to right: deep MLPs, wide
28

---

Published as a conference paper at ICLR 2024
(a) Probe 4' (x:) at x.-tokens
(b) Probe 4 (xx) at yi-tokens
(c) Probe Di
at x; tokens
109
10"
6 10
10a
10-
10-1
10-1
x 10
2
10"2
6
10
12
layer
Inyer
layer
Figure 14: Probe & (x:) and "0" A making the error 0-th layer be 1.
and tokens with deep MLPs. We normalize the probing error,
(a) Probe 4- (xx) at x,-tokens
(b) Prohe T5 (xx;;) at yi-toeens
(c) Probe V/
++* AA
at x; tokens
10"
109
moan sror
2x100
550
60-!
10a
4 10
10"1
3 10
x 10
x 10
0
6
1.0
12
2
8
10
12
layst
layer
Figure 15: Probe 4= (xxx) and y5 at x, and ys tokens with wide MLPs. We normalize the probing error, making the error at 0-th th layer to be 1.
(a) Probe T (xx; at xy-tokens
(b) Prohe 0 (x:) at yo-tokens
(c) Probe D;
-4" A
tokens
10at
eop Iealy
10a
souare Ieaky
10
12
10
6
8
10
12
layer
Inyer
layer
Figure 16: Probe (xx) and D. probing error, making the error at 0-th layer 1.
x; and zi tokens with other non-linear functions. We normalize the
29


Reviews:
[{"review_id": "xNOEPkibGm", "reviewer": "Reviewer_oQN2", "summary": "This paper focuses on understanding the internal mechanism by which a Transformer model solves an in-context learning task where the label $y$ for an instance $x$ linearly depends on a representation $\\phi^{\\star}(x)$.  A recent line of work has focused on explicitly constructing transformer models that can simulate various learning methods (e.g., gradient descent) on a training objective defined by the in-context labeled examples during a forward pass of the Transformer model. This paper extends this line of work by considering a more general data model where the final label depends on the input instance through a linear function of a representation. The paper provides explicit constructions for Transformer networks that can simulate ridge regression for 1) supervised learning with a representation, and 2) learning dynamical systems with a representation. The explicit constructions first aim to employ the underlying representation map $\\phi^{\\star}$ in the lower layers of the Transformer and then implement gradient descent in the upper layers of the Transformer.\n\nThrough experiments on synthetic datasets, the paper demonstrates that the performance of in-context learning via Transformers closely agrees with the performance of an optimal ridge predictor. Through probing analysis of the Transformer models, the authors show evidence that supports representation mapping following by label prediction aspects of their constructions.", "questions": "1) In general, the transformers are known to be universal approximators. In light of these, could authors comment on the significance of the key contributions of the paper, i.e., learning a representation map before applying gradient-descent in the representation space?\n\n2) Could the authors elaborate on using a **linear model** for their investigation on the upper module via pasting (Figure 4)?\n\n3) Currently the non-linearity in the true representation map is closely tied to the nonlinearity used in the Transformer network. Could the authors comment on generalizing this to broader nonlinearities in the true representation map? How would it increase the required number of layers?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1) The paper successfully extends the recent line of work on showing the feasibility of in-context learning via empirical risk minimization during a forward pass by considering data models where labels depend on the input feature via a representation map.\n2) The paper is well-written and explains the key contributions and techniques clearly.\n3) The empirical results (on synthetic datasets) do indicate the feasibility/presence of the explicit in-context learning mechanism hypothesized in the paper.", "weaknesses": "1) Novelty of the technical contributions is limited given prior works of similar flavor that provide the feasibility of empirical risk minimization during forward pass. One of the aspects which authors claim to be novel is that they allow for representation-based learning. However, the underlying assumption is that the representation map is a multi-layer MLP model, which Transformers should be easily able to simulate through its MLP layers. In that sense, the results in not very surprising. \n\n2) The probing analysis is done only on a synthetic setup. It would be nice to get some supporting evidence for the proposed in-context learning mechanism on a real dataset.\n\nMinor issues:\n\n1) The authors may want to formally introduce/discuss the pre-training phase (e.g., Eq (10)) which learns the representation map early in the section on preliminaries. \n\n2) In the paragraph on **In-context learning** in Section 2, $\\mathcal{D}^{(j)}$ and $\\mathbf{w}_{\\star}^{(j)}$ are not defined before their usage.", "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": "vIxJaCSeE3", "reviewer": "Reviewer_67Bi", "summary": "Results illustrating the performance of transformers in ICL tasks that necessitate some degree of representation learning are presented. The theory can be partially validated through probing experiments.", "questions": "The reviewer is open to learning about new evidences or analyses which address the points of the “Weaknesses” section above in this review.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "Clear and well-written. \n\nThis paper is one of the pioneering efforts to formalize how transformers execute ICL tasks that necessitate a degree of representation learning.", "weaknesses": "The theory only encompasses representational results by providing some settings of the parameters in a way that a transformer performs an ICL task. Given the transformer's highly expressive capability, these types of constructions are generally relatively straightforward.\n\nThere's no assurance that these theoretical constructs are truly internalized by the model during the training process. Although probing experiments gave us some confidence that, in specific instances, the theory can predict the model's behavior, these types of experiments generally don't offer robust guarantees. As a result, while the theory is logical and sometimes mirrors empirical events, it could be counterproductive to lean too heavily on these theoretical constructs. It might be necessary to carry out an analysis of training dynamics in order to theoretically determine under which conditions the model actually aligns with the theoretical constructs.", "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": "duHUCeQ4j7", "reviewer": "Reviewer_R1v9", "summary": "The goal of this paper is to theoretically and empirically understand the mechanism of in-context learning with underlying representations. Specifically, the setting considered is where there is a fixed representation function, chosen to be an MLP, and the ICL problem is to learn ridge regression on these representations. The transformer must learn this fixed representation function during pretraining and a regression hypothesis in-context. The authors theoretically show that it is possible to construct transformers that can perform ridge regression in supervised and linear dynamical system settings on fixed representations. Empirically, the paper verifies that transformers can learn to perform this type of ICL by probing for the emergence of mechanisms and values that should emerge according to theoretical construction.", "questions": "- What happens when the representation function is of a different form? If either the transformer does not have enough layers or the width, is there an approximate representation function learned on which regression is performed, or does the entire mechanism fall apart?\n- How robust is learning of the representation function in settings where the pretraining data contains spurious correlations? Can we say anything about the transformer's ability to compositionally generalize with either the representation function, regression, or both?\n- What is OLS in Fig 1b?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The paper well-written.\n- The experiments do a good job at validating the claims by probing for the relevant information.\n- The results offer valuable insights into how in-context learning, which is very relevant and timely.", "weaknesses": "- Labels in figures and the figure captions can be more clear. For example, items like \"TF_upper+1_layer_TF_embed\" in Figure 4b are not very readable.\n- Section 3 could be significantly condensed by considering theorem 2 as a generalization of theorem 1 instead of presenting them separately.\n- Section 3.1 states that the representation function can be chosen arbitrarily, but Lemma B.3 requires a specific structure and non-linearity to work.\n- Although the paper does a good job of illustrating the claimed mechanism, it does not analyze settings where the mechanism breaks.", "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": "EvDZvKKegY", "reviewer": "Reviewer_z67a", "summary": "This works studies in-context learning in transformers using synthetic data. It extends previous work, by studying composition of a *fixed* non-linear function (L-layer MLP) with a linear function that is learned in-context. This work provides a construction of a transformer that can solve this task, but also demonstrates it empirically on synthetic data. Additionally, the authors provide a mechanistic understanding of the algorithm implemented by a trained transformer.", "questions": "1. Could the authors add more details on how the non-linear functions are created? Can the authors also clarify in the introduction/abstract that the functions are L-layer MLPs?\n\n2. Is it possible to show some of these results on other families of non-linear functions? For example, what happens if the functions are polynomials or exponential functions of the input? Are there scenarios where it fails empirically?\n\n3. What happens if we increase the number of layers used to create the representation from 5 to 15. Does the model start to fail if L=15 (even if transformer has only 11) layers or does it find a good approximation to the non-linear function using just 4-5 layers?\n\n4. Results in appendix E were very interesting! As future work, it would be great to investigate how many different non-linear functions can be learnt. I would also be interested in understanding if in-context learning becomes difficult and if the model sometimes struggles to identify the right non-linear representation.\n\n5. Ruiqi et al. (https://arxiv.org/abs/2306.09927) show that in-context learning fails if the linear functions are selected to be out-of-distribution. Is this also the case here?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The results extend the the setup of Garg et al. to study in-context learning with more complex function classes. In particular, transformers can learn a composition of a *fixed* non-linear function with a linear function learnt from context. The authors in Figure 1 provide evidence that a transformer matches the optimal predictor.\n\nThe mechanistic analysis is thorough, and provides compelling evidence for the underlying 3-step mechanism. It is surprising that the mechanism is consistent across multiple training runs and adds to the results of the paper. \n\nThe results also hold when multiple different non-linear representations are used which further strengthens the main claims of the paper. I would have liked to see the results of 4.1.1 more prominently in the main paper, but I understand the authors are constrained by space. \n\nOverall, I think the results would be of interest to the community and the toy setup may be more representative of in-context learning in language models.", "weaknesses": "**Why are the constructive proofs important for understanding in-context learning in transformers?**\nI am aware that there are prior works that design transformers that are capable of in-context learning. However, I am not convinced of the importance and significance of these results. Couldn't we also find weights for other architectures (like large MLPs or LSTMs) and argue that they are capable of in-context learning. Is the existence of these model weights informative of what is learnt in practice?\n\n**Choice of non-linear functions.**  I think the authors could be more rigorous in evaluating the non-linear representations used in their setup. In particular, the non-linear representations are L-layer MLPs with the matrices being random orthogonal matrices. Do the results hold for other families of functions and does it fail to work for some other classes of functions? I think it would be helpful to clarify that the results are specific to this setup.\n\n**Is the synthetic setup an accurate toy-model to understand language models?** Like previous work, all the results are on synthetic data. It remains unclear if the toy setup is representative of in-context learning in language models. What kind real world tasks are captured by a composition of a fixed non-linear function and a linear function that is learnt from context?", "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)