Linear Transformers are Versatile In-Context Learners
Authors: Max Vladymyrov
Johannes von Oswald
Mark Sandler
Rong Ge
Venue: NeurIPS 2024 poster (2024)
Keywords: Linear Transformers
In-Context Learning
Noisy Linear Regression
Model Selection
Mesa-optimization
Primary area: interpretability_and_explainability
TL;DR: This paper shows that linear transformers can implicitly learn sophisticated optimization algorithms, particularly for noisy linear regression, exceeding the performance of traditional methods.

Abstract: Recent research has demonstrated that transformers, particularly linear attention models, implicitly execute gradient-descent-like algorithms on data provided in-context during their forward inference step. However, their capability in handling more complex problems remains unexplored. In this paper, we prove that each layer of a linear transformer maintains a weight vector for an implicit linear regression problem and can be interpreted as performing a variant of preconditioned gradient descent. We also investigate the use of linear transformers in a challenging scenario where the training data is corrupted with different levels of noise. Remarkably, we demonstrate that for this problem linear transformers discover an intricate and highly effective optimization algorithm, surpassing or matching in performance many reasonable baselines. We analyze this algorithm and show that it is a novel approach incorporating momentum and adaptive rescaling based on noise levels. Our findings show that even linear transformers possess the surprising ability to discover sophisticated optimization strategies.

Linear Transformers are Versatile In-Context Learners
Max Vladymyrov Google Research mxvêgoogle. com
Johannes von Oswald Google, Paradigms of Intelligence Team jvoswald@google com
Mark Sandler
Rong Ge
Google Research sandler@google. com
Duke University
ronggeEcs duke. edu
Abstract
Recent research has demonstrated that transformers, particularly linear attention models, implicitly execute gradient-descent-lie algorithms data provided in- context during their forward inference step. However, their capability in handling more complex problems remains unexplored. In this paper, we prove that each layer of a linear transformer maintains a weight vector for an implicit linear regression problem and can be interpreted as performing a variant of preconditioned gradient descent. We also investigate the use of linear transformers in challenging scenario where the training data corrupted with different levels of noise. Remarkably, we demonstrate that for this problem linear transformers discover intricate and highly effective optimization algorithm, surpassing or matching in performance many reasonable baselines. We analyze this algorithm and show that it is novel approach incorporating momentum and adaptive rescaling based on noise levels. Our findings show that even linear transformers possess the surprising ability discover sophisticated optimization strategies.
Introduction
The transformer architecture (Vaswani et al. 2017) has revolutionized the field of machine learning. driving breakthroughs across various domains and serving as a foundation for powerful models (Anill et al., 2023 Achiam et al. 2023 Team et al., 2023 Jiang et al., 2023). However, despite their widespread success, the mechanisms that drive their performance remain an active area of research.
A key component of their success is attributed to in-context learning (ICL, Brownn et al.) 2020) emergent ability of transformers to make predictions based on information provided within the input sequence itself, without explicit parameter updates.
Recently, several papers (Garg et al.) 2022] Akyürek et 2022. von Oswald et al., 2023a) have suggested that ICL might be partially explained by an implicit meta- oettaoopiiiztioonn of the transformers that happens on input context (aka optimizptinn Hubinger al. 2019). They have shown that transformers with linear self-attention layers (aka linear transformers) trained on linear regression tasks can internally implement gradient-based optimization.
Specifically, von Oswald et al. (2023a) demonstrated that linear transformers can execute iterations of an algorithm similar the gradient descent algorithm (which they call GD++), with each attention layer representing one step of the algorithm. Later, Ahn et al. (2023): Zhang et al. (2023) further characterized this behavior, showing that the learned solution form of preconditioned GD, and this solution is optimal for one-layer linear transformers.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

In this paper, we continue to study linear transformers trained on linear regression problems. We prove that each layer of every linear transformer maintains a weight vector for an underlying linear regression problem. Under some restrictions, the algorithm it runs can be interpreted as complex variant of preconditioned gradient descent with momentum- like behaviors.
While maintaining linear regression model (regardless of the data) might seem restrictive, we show that linear transformers can discover powerful optimization algorithms. As a first example, we prove that in case of GD++, the preconditioner results in a second order optimization algorithm.
Furthermore, we demonstrate that linear transformers can be trained to uncover and intricate algorithms. We modified the problem formulation to consider mixed linear regression with varying noise level=" (inspired by Bai et al. 2023). This is harder and non-trivial problem with no obvious closed-form solution, since it needs to account for various levels of noise in the input.
more powerful
Our experiments with two different noise variance distributions (uniform and categorical) demonstrate the remarkable flexibility of linear transformers. Training linear transformer in these settings leads to algorithm that outperforms GD++ well as various baselines derived from the exact closed- form solution of the ridge regression. We discover that this result holds even when training a linear transformer with diagonal weight matrices.
Through a detailed analysis, we reveal key distinctions from GD +++ including momentum- term and adaptive rescaling based on the noise levels.
Our findings contribute the growing body of research where novel, tigh-performing algorithms have been directly discovered through the reverse-engingg eevrre--ngiinerrig of transformer weights. This work expands understanding of the implicit learning capabilities of attention- based models and highlights the remarkable versatility of even simple linear transformers as in-context learners. We demonstrate that transformers have the potential to discover effective algorithms that may advance the state-of-the-art in optimization and machine learning in general.
2
Preliminaries
In this section we introduce notations for linear transformers, data, and type problems we consider.
2.1 Linear transformers and in-context learning
Given input sequence e1, .2.... en E [2++1 single head in linear self-attention layer is usually parameterized by four matrices, key WK. query Wo, value Wv and projection Wp. The output of the non-causal layer at position i is + De; where Ae; is computed
De Wp ")=( Wke))Wve,).
(1)
Equivalently, one can use parameters p WpWy and Q = W Wo. and the equation becomes
Ae Σラ= [e] Pe.)
(2)
Multiple heads (P1.Q1). (P2, Q2), ... (Pr, Qn) simply sum their effects
Ae; = E (e) Que:)
We define linear transformer multi-layer neural network composed of L linear self-attention layers parameterized by 0 = (Q): } for k == 1... H, 1... L. To isolate the core mechanisms, we consider simplified decoder-only architecture, excluding MLPs and LayerNorm components.
This architecture was also used in previous work (von Oswald et al. 2023a Ahn et al. 2023).
We consider two versions of linear transformers: FULL with the transformer parameters represented by full matrices and DIAG, where the parameters are restricted to diagonal matrices only.
Inspired by von Oswald et al. 2023a). in this paper we consider regression data as the token sequence. Each token e; ( V/) E I2 consists of feature vector F; IRd and its corresponding output
We consider a model where each sequence contains data with the same noise level, while different sequences have different noise levels.

---

E R. Additionally, we append query token ex+1 represents data. The goal of in-context learning is to predict yr for the test data 2₁. We constrain the attention to only focus the first tokens of the sequence so that it ignores the query token.
(Iz, 0) to the sequence, where E R
We use (ar vf) denote the i-th token in the transformer's output at layer 1. 'he initial layer is simply the input: (2;, yu) = (zi, For a model with parameters 0, we read out the prediction by taking the negative of the last coordinate of the final token in the last layer as ÿo
en Cn+1) = -
Let's also define the following notation to be used throughout the paper
E
Σ
a
Σ y:
X
(3:)
==1
1=1
3 = Σ d
d'
= s/al
X
Σ (w/)
2.2
Noisy regression model
As a model problem, we consider data generated from a noisy linear regression model. For each input sequence
sample ground-truth weight vector wl- N(( I), and generate data points as 3; ~ (0, J) and Y/1 = (Wr, 24) + &,, with noise & ~ (0)
Note that each sequence can have different ground-truth weight vectors wy. but every data point in the sequence shares the same w, and a, The query generated as H ~ (0, I) and 3/s = (w,, xt) (since the noise is independent, whether we include noise in Mo will only be additive constant to final objective).
We further define an ordinary least square (OLS) loss as
Lots(w) Ei= (yi - (w. xzi)
(4)
The OLS solution is w" := E- a with residuals r; == - (w" x;))
In the presence of noise or, w" in general is not equal to the ground truth wy. For known noise level or, the best estimator for wr is provided by ridge regression:
IRR (w) E (y/ - (w, xx))²² + w?|
(5)
with solution w (2) + o -1 a. Of course, in reality the variance of the noise known and has to be estimated from the data.
2.3 Fixed vs. mixed noise variance problems
We consider two different problems within the noisy linear regression framework. Fixed noise variance. In this scenario, the variance a, remains constant for all the training data. Here, the in-context loss is:
I(0)=
E w. N10.) N(O
ex) en+1) y1 ))
(6)
(0,02)
where e; = (z 31) and y4 (w,, rT) +.c. This problem was initially explored by Garg et al. (2022). Later, von Oswald et al. (2023a) have demonstrated that linear transformer (6) converges to form of gradient descent solution, which they called GD++ We define this in details later.
Mixed noise variance. In this case, the noise variance a, is drawn from some fixed distribution p(a,) ) for each sequence. The in-context learning loss becomes:
L(0)
IE
[(oo(te,,.., ent, en+1) v1 )1
((0.) (0,1
(0.
a,>p((r))
We set the actual prediction to - -ii+1. similar to von Oswald et al. 2023a), because it's easier for linear transformers to predict ye-

---

In other words, each training sequence 7 has fixed noise level Tr, but different training sequences have different noise levels sampled from specified distribution p(or). This scenario adds complexity because the model must predict w- for changing noise distribution, and the optimal solution likely would involve some sort of noise estimation. We have found that empirically, GD++ fails to model this noise variance and instead converges to a solution which can be interpreted as a single noise variance estimate across all input data.
3 Related work
In-context Learning Gradient Descent Our work builds on research that frames in-context learning as (variants of) gradient descent (Akyürek et al. [022] von Oswald et al. (023a)). For 1-layer linear transformer, several works Zhang al. (2023): Mahankali et al. (2023); Ahn al. (2023)
characterized the optimal parameters and training dynamics. More recent works extended the ideas to auto-regressive models (Li al. 2023) von Oswald et al. 2023b) and nonlinear models (Cheng et al. (222)) Fu et al. (2023) noticed that transformers perform similarly to second-order Newton methods on linear data, for which we give a plausible explanation in Theorem 5.1
In-context Learning in LLMs There are also many works that study how in-context learning works pre-trained I.L.Ss (Kossen al. [023] Wei et al., 2023] Hendel et al., 2023; [hhnn et al.
2023). Due to the complexity of such models, the exact mechanism for in-context learning is still
major open problem. Several works (Olsson et al., 2022 Chan et al. 2022 Akyürek et al. 2024)
identified induction heads crucial mechanism for simple in-context learning tasks, such as copying, token translation and pattern matching.
Other theories for training transformers Other than the setting of linear models, several other
works (Garg et al.] 2022; Tarzanagh et al.) 2023 Li et al. 2023 Huang et al., [023] Tian et al.)
2023a b) considered optimization of transformers under different data and model assumptions, Wen et al. (023)) showed that it can be difficult to interpret the "algorithm" performed by transformers without very strong restrictions.
Mixed Linear Models Several works observed that transformers can achieve good performance on mixture of linear models (Bai et al., 2023 Pathak et al. 2023 Yadlowsky et al. 2023). While these works show that transformers can implement many variants of selel-selection techniques, our result shows that linear transformers solve such problems by discovering interesting optimization algorithm with many hyperparameters tuned during the training process. Such a strategy is quite different from traditional ways of doing model selection. Transformers are also known to able to implement strong algorithms in many different setups (Guo et al. 2023 Giannou et al. 2023).
Effectiveness linear and kernel-like transformers ture is that it takes O(N2) time for sequence of length N, while for linear transformer this can be improved to ONN). Mirchandani et al. (2023) showed that even linear transformers are quite powerful for many tasks. Other works (Katharopoulos et al. 22200] Wang et al.) 2020; Schlag et al.) 2021 Choromanski et al. 2020) uses ideas similar to kernel/random features to improve the running time to almost linear while not losing much performance.
main constraint on transformer architec-
4 Linear transformers maintain linear regression model at every layer
While large, nonlinear transformers can model complex relationship, we show that linear transformers are restricted maintaining linear regression model based on the input, in the sense that the 1-th layer output is always a linear function of the input with latent (and possibly nonlinear) coefficients.

---

4.1. Suppose the outpur of a linear transformer at 1-th layer is (oh, v (1/2 122 ... (1 Un), (z), vi). then there exists matrices M', vectors u', w and scalars al such that
Theorem
at M, you.
M (w zi).
v = d'y; (w² ,3))
1+1
V/2
Note that M u², wt and a' are not linear in the input, but this still poses restrictions what the linear transformers can do. For example we show that cannot represent a quadratic function:
Theorem 4.2. Suppose the imput to linear transformer is (211.)) (2), (Fns yr where
T; (0, 1) and yi w Ii, let the I-th layer output be (21,2)) (222 v2), (5 yn) and let y' = (si .... Vnr and y (z1(1)2 (z1(1),,z2(1} (.., (..)?) (here 21( is just the first coordinate of æz). then when n > d with high probability the cosine similarity of yf and y is at most 0.1.
Theorem implies that the output of linear transformer can always be explained as linear combina- tions of input with latent weights a' and w'. The matrices M vectors u' w and numbers a' are not linear and can in fact be quite complex, which characterize below:
Lemma 4.3. In the setup of Theorem
if we let
A'
6¹
)
(c)
d'
h
(GO)) .1)) oh
Σ
P
j=1
then one can recursively compute matrices M vectors ut, wi and numbers a' for every layer using
M (1+ + )'M + b' (ww))
dt + + A'' d'b'
d = (1 d' )a' + (c,
/+| = (1 + d' w TM
with the init. condition a = 1, w 0, M = 0 =.
The updates to the parameters complicated and nonlinear, allowing linear transformers imple- ment powerful algorithms, as we will later see in Section 5 In fact, even with diagonal P and Q, they remain flexible. The updates in this case be further simplified to more familiar form:
Lemma 4,4. In the setup of Theorem
4.1 with diagonal parameters, u', w are updated as
+ (1 - A'y' I'E (a'w" w'); w 1+ -1= I''((''u" =w') D'ud.
Here A Arr,', sf 11,11',, a are matrices and numbers that depend on M u, a', w in Lemma 4.3
Note that (a'""-w")) is (proportional to) the gradient of a linear model f(w') Ei (w', x.)))² This makes the updates similar to a gradient descent with momentum:
w')
u I+ (1 3)u' V f (w") w w nu'.
Of course, the formula in Lemma 4.4 is still much more complicated with matrices in places of B and n, and also including a gradient term for the update of w.
Power of diagonal attention matrices
Although linear transformers are constrained, they can solve complex in-context learning problems. Empirically, we have found that they are able to very accurately solve linear regression with mixed noise variance (7). with final learned weights that are very diagonal heavy with some low-rank component (see Fig. Surprisingly, the final loss remains remarkably consistent even when their Q and P matrices
diagonal. Here will analyze this special case and explain its effectiveness.
5

---

Since the elements of z are permutation invariant, a diagonal parameterization reduces each attention heads to just four parameters:
:
Pk
O
(8)
0
Tys.k
It would be useful to further reparametrize the linear transformer
(3) using:
w H Pr, whey 2k=1 P Po
H
(9)
w yr
- Po
wpy= =
= k= Pyy.
This leads the following diagonal layer updates:
x 2 + w HeLE'' E'at + w wayth
x ++1 x1e w'e E'alt + w wayvia
(10)
vt +11 v + w (a sit + woytt
ait1 yt + way (a', .n/) + wbyuld..
Four variables w wyy w ye w represent information flow between the data and the labels across layers. For instance, the term controlled by wir measures information flow from x' to x(+1, w/ measures the flow from x' to y't and so forth. Since the model can always be captured by these 4 variables, having many heads does not significantly increase its representation power. When there is only one head the equation wyrn"yy wry&yy is always true, while models with more than one head do not have this limitation. However empirically even models with one head quite powerful.
5.1 GD++ and least squares solver
GD++ introduced in von Oswald al. (2023a), represents linear transformer that is trained fixed noise variance problem (6) It is a variant of diagonal linear transformer, with all the heads satisfying dy = Dynamics influenced only by wi and w leading simpler updates:
s/41 (I w E') aq
(11)
y. 1+ vf + yz (a =!)..
The update on x acts as preconditioning and the update on y performs gradient descent on the data.
While existing analysis by Ahn et al. (2023) has not yielded fast convergence rates for GD++ we show here that it actually a second-order optimization algorithm for the least squares problem (
Theorem 5.1. Given (R1,31), .... (,), (xx, 0) where E has eigenvalues in the range [v, with a condition number e = =v/pe Let u² be the optimal solution to least squares problem (4). then there exists hyperparameters for GD)+ algorithm that outputs ý with accuracy 15-(21, w²)l S c||are [ 1 Ologg log log 1/c) steps. In particular that implies there exists an l-layer linear transformer that can solve this task.
The convergence rate of Ollog log 1/c) is typically achieved only by second-order algorithms such as Newton's method.
5.2 Understanding Wyyt adaptive rescaling
If a layer only has w uy *f 0, it has rescaling effect. The amount of scaling is related to the amount of noise added in a model selection setting. The update rule for this layer is:
vi = (1 +1++0 ='y vl-
This rescales every y by factor that depends on X When vy < O, this shrinks of the output based on the norm of y in the previous layer. This is useful for the mixed noise variance problem, as ridge regression solution scales the least squares solution by a factor that depends on the noise level.
Specifically, assuming Σ A EP- = n/,, the ridge regression solution becomes w PJ n w", which is exactly scaled version of the OLS solution. Further, when noise is larger, the scaled factor smaller, which agrees with the behavior of negative wyy.
We can show that using adaptive scaling Wyy even a 2-layer linear transformer can be enough to solve simple example of categorical mixed noise variance problem a, e {01, 02) and - oo:

---

Omay 3
Ad). eval loss
109
10"
10
1-
10-2
10"1
10"1
10-10
10
Adj. oval loss
10P
Omee =4
10
Omay 5
Onee =6
Oray 7
10-1
10
10"1
10-1
Number of layers
Number of layers
Number of layers
Number of layers
GD
Diag
Full
ConstRR
AdaRR
TunedRR
Figure 1: In-context learning performance for noisy linear regression problem across models with different number of layers and maz for 0, ~ UO, mmax). ). Each marker corresponds to a separately trained model with a given number of layers. Models with diagonal attention weights (DIAG) match those with full attention weights (FULL). Models specialized on a fixed noise (GD- perform poorly, similar to a Ridge Regression solution with a constant noise (CONSTRR), Among the baselines, only tuned exact Ridge Regression solution (TUNEDRR) is comparable with linear transformers.
Theorem 5.2. Suppose the input to the transformer is (11,31), (I2, 82). Un), (I 0), where T; ~ N(,, I), yt w 24 + Es. Here E ~ N NO, o²) is the noise whose noise level a can take one of Iwo values: σı or 02- Then as goes 400, there exists a parameters for two-layer linear transformers such that the implicit w² of the linear transformer converges to the optimal ridge regression results (and the output of the linear transformer is {w² zg)). Further, the first layer only has Wyz being nonzero and the second layer only has wyy being nonzero.
5.3 Understanding Way: adapting step-sizes
The final term in the diagonal model, Way, has a more complicated effect. Since it changes only the r-coordinates, it does not have immediate effect y. To understand how it influences the y consider a simplified two-step process, where the first step only has way * 0 and the second step only has wyz * O (so the second step just doing one step of gradient descent). In this case, the first. layer will update the 2;'s as:
x! x. DiWwy
yj
n
x4 + WaySt Ui D ,,,)) } + r3)
vs + Way En""
x, Way (w" Izi) + En
(I way Ew" "ww) )a, + Eww"
There two effects of the Wry term, is multiplicative effect v;, and the other is an additive term that makes z æ-output related to the residual 7ą- The multiplicative step in I; has an unknown preconditioning effect. For simplicity assume the multiplicative term small, that is:
24 WayT; Ew": z1 Zt.
The first layer does not change y, so y! Ye and y! yt- For this set of Hı, we can write down the output on in the second layer as
n
yi 4/ + Wys
ye vi(x)) xt
Vt + Wyzz
Your Way Σ yor, Ewi Eum)
Yt + Wys + Way Σ (Tww) Zz-
7

---

33
Ad]. eval loss
0.010
0.10
0.05
0.01
0.005
0.05
0.000
0.00
0.00
0.00
0.0
0.5
1.0
Ad). loss
Omay
Omax
Onay
0.2
Omay
0.1
0.1
0.1
0.1
0.0
0.0
0.0
0.0
2
3
4
5
6
7
Variance
Varianoe
Vanance
Variance
adj evea loss
0.20
layers
5 layers
6 layers
layers
0.15
0.10
0.06
0.00
6
6
2
5
6
3 Variance
2
6.
2
3
5
6
2
3
5
Vananca
Variance
Variance
Variance
0
Variance
0
GD
Diag
Full
ConstRR
AdaRR
TunedRR
Figure 2: Per-variance profile of models behavior for uniform noise variance d, (O, daaz). ). Top
Two rows: 7-layer models with varying maa- Bottom row: models with varying numbers of layers. fixed mar 5. In-distribution noise shaded gray.
Here used the properties of residual (in particular E, Vız; Ew", and Σ, UT: r Note that (Ew") ) 2, is basically what a gradient descent step the original input should do. Therefore effectively, the two-layer network is doing gradient descent, but the step size is the product of -Wyz and (1 Way Σ,r r?). The factor (1 + Way E, r?) depends the level of noise, and when Way, wyz < 0, the effective step size is smaller when there is more noise. This is especially helpful in the model selection problem, because intuitively would like perform early-stopping (small step sizes) when the noise is high.
6 Experiments
In this section, investigate the training dynamics of linear transformers when trained with mixed noise variance problem (7). We evaluate three types of single-head linear transformer models:
FUL.... Trains full parameter matrices.
DIAG. Trains diagonal parameter matrices
TO
GD+ An even more restricted diagonal variant defined in
For each experiment, we train each linear transformer modifications with varying number of layers (1 to 7) using using Adam optimizer for 200 000 iterations with learning rate of 0.0001 and a batch size of 2 048. In some cases, especially for large number of layers, had to adjust the learning rate to prevent stability issues. We report the best result out of runs with different training seeds. We used N 20 in-context examples in D 10 dimensions. We evaluated the algorithm using 100 000 novel sequences. All the experiments were done on single H100 GPU with 80GB of VRAMM. It took on average 4-12 hours train a single algorithm, however experimenting with different weight decay parameters, better optimizer and learning rate schedule will likely reduce this number dramatically.
We use adjusted evaluation loss our main performance metric. It is calculated by subtracting the oracle loss from the predictor's loss. The oracle loss the closed-form solution of the ridge regression loss (5). assuming the noise variance d+ is known. The adjusted evaluation loss allows for direct model performance comparison across different noise variances. This important because higher noise significantly degrades the model prediction. Our adjustment does not affect the model's optimization process, since it only modifies the loss by an additive constant.
Baseline estimates. evaluated the linear transformer against closed-form solution to the ridge We estimated the noise variance a, using the following methods:
regression problem
Constant Ridge Regression (CoNSTRR). The noise variance estimated using a single scalar value for all the sequences, tuned separately for each mixed variance problem.

---

o
{1,3)
0,
(1,3,5)
Adi eval evallloss
10"
lose
0.2
109
Adi. eval loess
0.2
10"1
0.1
10"
0.1
10--
10:2
0.0
Number of layers
Noise variance
Number of leyers
Nolse variance
0r E {1,3, ( (1,3) 1,33
Ad) eval loss A/. evell lose
0.2
koyers
layers
layars
layeens
leyers
layars
0.11
5,,,,
10.0-
0.2
0.1-
0.00
6.0
6.0
2 8
6.0
Variance
Variance
Variance
Variance
Variance 0
Variance
GD
Diag
Full
ConstRR
AdaRR
TunedRR
Figure 3: In-context learning performance for noisy linear regression across models with varying
number of layers for conditional noise variance o, {1,3} and o, E f1, 5). Top: loss for models with various number of layers and per-variance profile for models with layers. Bottom: Per-variance profile of the model across different numbers of layers. In-distribution noise is shaded gray.
Adaptive Ridge Regression (ADARR). Estimate the noise variance via unbiased estimator (Cherkassyy & Ma 2003) de --a 1 (y/ - "² where ūj represents the solution to the ordinary least squares 4.. found closed-form.
Tuned Adaptive Ridge Regression (TUNEDRR) Same as above, but after the noise is estimated, we tuned two additional parameters minimize the evaluation loss: (1) a max.
threshold value for the estimated variance, (2) a multiplicative adjustment to the noise estimator. These values are tuned separately for each problem.
Notice that all the baselines above are based on ridge regression, which a closed-form, non-iterative solution. Thus, they have an algorithmic advantage over linear transformers that do not have access to matrix inversion. These baselines help us gauge the best possible performance, establishing an upper bound rather than a strictly equivalent comparison.
faithful comparison our method would be an iterative version of the ADARR that does use matrix inversion. Instead, we can use gradient descent to estimate the noise and the solution to the ridge regression. However, in practice, this gradient descent estimator converges to ADARR only after x 100 iterations. In contrast, linear transformers typically converge in fewer than 10 layers.
We consider two choices for the distribution of a;:
Uniform, o, ~ (0, Tmaz)) drawn from a uniform distribution bounded by oman. We tried multiple scenarios with dear ranging from 0 to 7.
Categorical. a, S chosen from discrete set S. We tested S {1, and S {1. 3,5,, 5).
Our approach generalizes the problem studied by Bai et al. (2023). who considered only categorical variance selection and show experiments only with two ar values.
Uniform noise variance. For the uniform noise variance, Fig. [ss[sshwwss that FULl and DIAG achieve comparable performance across different numbers of layers and different mmaz. On the other hand, +D++ converges to higher value, closely approaching the performance of the CONSTRR baseline.
As Omaz grows, linear transformers show a clear advantage over the baselines. With 4 layers, they outperform the closed-form solution ADARR for Tmas and larger. Models with more layers match or exceed the performance of TUNEDRR.
The top of Fig. offfers detailed perspective on performance of 7-layer models and the baselines. Here, we computed varinnce profiles across noise variance range from 0 to o maz + 1. We can that poor performance of GD ++ comes from its inability to estimate well across the full noise variance range. Its performance closely mirrors to CONSTRR, suggesting that GD++ +++ under the hood might also be estimating single constant variance for all the data.
ADARR perfectly estimates problems with no noise, but struggles more as noise variance increases.
TUNEDRR slightly improves estimation by incorporating omar into its tunable parameters, yet its

---

02
0.1
-0.1
-0.2
-0.2
11
-0.
123456
52 9191011
123456 91011
123456 6781011
Layer
Layer
Layer
Layer
-0.1
0.2
0.1
-00
0.0
00
0.0
10
-0.1
-0.2
10
-02
-0.1
/23456789101
23456 68911101
91011
723456
91011
129456789101
Figure 4: Weights for 4 layer linear transformer with FULL parametrization trained with categorical noise o, {1,3). Top: weights for Q' matrix, bottom: weights for fd matrix.
prediction suffers in the mid-range. FULL and DIAG demonstrate comparable performance across all noise variances. While more research is needed to definitively confirm or deny their equivalence, we believe that these models are actually not identical despite their similar performance.
At the bottom of Fig. 2 we set the noise variance to maaz = and display a per-variance profile for models with varying layers. Two-layer models for FULI. and DIAG behave akin to GD++, modeling only a single noise variance in the middle. However, the results quickly improve across the entire noise spectrum for 3 or more layers. In contrast, GD++ quickly converges a suboptimal solution.
Categorical noise variance. Fig. 3 shows a notable difference between DIAG and FULL models for categorical noise variance o E (1,3). This could stem from bad local minima, suggest a fundamental difference between the models for this problem. Interestingly, from per-variance profiling we that DIAG extrapolates better for variances not used for training, while FULL, despite its lower in-distribution еrror, performs worse on unseen variances. Fig. showws learned weights of the layer linear transformer with FULL parametrization. The weights are very diagonal heavy, potentially with some low-rank component.
For d, (1,3,5, 5}. examining the per-variann variance profile at the bottom of Fig. 3 reveals differences in their behaviors. FULL exhibits a more complex per- perivariance profile with more fluctuations than the diagonal model, suggesting greater representational capacity. Surprisingly, it did not translate to better loss results compared to DIAG.
For easy comparison, we compile the results of all methods and baselines in Table in the Appendix.
7 Conclusions
Our research reveals the surprising ability of linear transformers to tackle challenging in-context learning problems. We show that each layer maintains implicit linear regression model, akin to complex variant of preconditioned gradient descent with momentum like behavior.
Remarkably, when trained on noisy linear regression problems with unknown noise variance, linear transformers not only outperform standard baselines also uncover sophisticated optimization algorithm that incorporates noise-aware step-size adjustments and rescaling. This discovery highlights the potential of linear transformers to automatically discover novel optimization algorithms when presented with the right problems, opening exciting avenues for future research, including automated algorithm discovery using transformers and generalization to other problem domains.
While our findings demonstrate the impressive capabilities of linear transformers in learning optimiza- tion algorithms, we acknowledge limitations in our work. These include the focus on simplified linear models, analysis of primarily diagonal attention matrices, and the need for further exploration into the optimality of discovered algorithms, generalization to complex function classes, scalability with larger datasets, and applicability to more complex transformer architectures. We believe these limitations present valuable directions for future research and emphasize the need for a deeper understanding of the implicit learning mechanisms within transformer architectures.
10

---

8 Acknowledgements
The authors would like to thank Nolan Miller and Andrey Zhmoginov for their valuable suggestions and feedback throughout the development of this project. Part of this work was done while Rong Ge was visiting Google Research. Rong Ge's research is supported in part by NSF Award DMS-203 1849 and CCF-1845171 (CARBER).
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Iige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.055 0877744 2023.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn implement preconditioned gradient descent for in-context learning. arXiv preprint 0rXiv:2306600029) 2023,
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-context learning? investigations with linear models. arXiv preprini urXiv:2211.15660 2022.
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-Context language learning: Architec- tures and algorithms. arXiv preprint arXiv:2401.12975. 2024.
Rohan Anil, Andrew M Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, et al. Palm technical report. arXiv preprint 2rXiv:2305.104000 2023.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.046) 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.
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 35: 18788888-1889, 2022, 2022.
Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:2312.06522 2023.
Vladimir Cherkassky and Yunqian Ma. Comparison of model selection for regression. Neural computation, 15(7):161-1174,, 2003.
Krzysztof Choromanski. Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint arXiv:23 17086, 2023.
Shivam Garg, Dimitris Tsipras, Percy 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.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Pa- pailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301. 13199, 2023.
Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? a case study on learning with representa- tions. arXiv preprint 2rXiv:2310.10611 2023.

---

Roee Hendel, Mor Geva, and Amir Globerson. In-context learning creates task vectors, arXiv preprint arXiv:2310.15916. 2023.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
Evan Hubinger, Chris van Merwijk, Vladimir Mikulik, Joar Skalse, and Scott Garrabrant. Risks from learned optimization in advanced machine learning systems. arXiv preprint arXiv:/906.00 820,
2019.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310006825 2023.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference machine learning, pp. 5156-5165, PMLR, 2020.
Jannik Kossen, Tom Rainforth, and Yarin Gal. In-context learning in large language models learns label relationships but is not conventional learning. arXiv preprint arXiv:2307. 02.3007123735 /2375, 2023.
Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, pp. 19565-19594. PMLR, 2023.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576, 2023.
Suvir Mirchandani, Fei Xia, Pete Florence, Brian Ichter, Danny Driess, Montserrat Gonzalez Arenas, Kanishka Rao, Dorsa Sadigh, and Andy Zeng. Large language models general pattern machines. arXiv preprint arXiv:2307.0472.1 2023.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, al. In-context learning and induction heads. arXiv preprint ar arXiv:2209. 11099,1188,,,, 2022.
Reese Pathak, Rajat Sen, Weihao Kong, and Abhimanyu Das. Transformers can optimally learn regression mixture models. arXiv preprint arXiv:2311.00 103.1...0622,..1 2023.
Imanol Schlag, Kazuki Iric, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pp. 9355-9366. PMLR, 2021.
Lingfeng Shen, Aayush Mishra, and Daniel Khashabi. Do pretrained transformers really learn in-context by gradient descent? arXiv preprint 10.00:2310.085400 2023.
Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: a family highly capable multimodal models. arXiv preprint arXiv:2312. 1805, 2023.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in I-layer transformer. arXiv preprint arXiv:2305. 16380, 2023a.
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
12

---

Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference Machine Learning, pp- 35151-35174. PMLR, 2023a.
Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, et al. Uncovering mesa-optimization algorithms in tramsformers. arXiv preprint arXiv:2309.0088: 05800090088,,. 2023b.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang. and Hao Ma. Linformer:
with
linear complexity. arXiv preprint arXiv:2006. 04768, 2020.
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, al. Larger language models do in-context learning differently, arXiv preprint arXiv:2303.03886, 2023.
Kaiyue Wen, Yuchen Li, Binghin Liu, and Andrej Risteski. Transformers are uninterpretable with myopic methods: case study with bounded dyck grammars. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Steve Yadlowsky, Lyric Doshi, and Nilesh Tripuraneni. Pretraining data mixtures enable narrow model selection capabilities in transformer models. arXiv preprint arXiv:2311.00871 2023.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers leamn linear models in-context. arXiv preprini 230v:23066.9927,, 2023.
13

---

A Proofs from Sections 4 and 5
A.1 Proof of Theorem
We first give the proof for Theorem 4.1 In the process we will also prove Lemma as Theorem 4.
follows immediately from an induction based on the lemma.
Proof. We do this by induction. At 1 = 0, it's easy to check that we can set ( = 1, wv())
0, _ = I, u (0) =0 O.
Suppose this is true for some layer if the weights of layer I are (PI (1. (Pk. Q+) for & heads, at output of layer I + 1 we have:
5 ++1
)
++1
Pk
n
1 (G)
18))
(
(12)
yi
3/1
j=1
Note that the same equation true for = n + just by letting yo+ the following structure:
Let the middle matrix has
(
T
2
)
b
k
(C)) (o))
o
j=l
Then one can choose the parameters of the next layer as in Lemma 4.3
M (I A)M' + ((w')
[+1 = (I+ Ayu' a'b
d = (1 + d' + (c. (c,u²) u²)
w 1+1 = (1 + ddw (M²) Tc.
One can check that this choice satisfies 12
A.2 Proof of Lemma
This lemma is in fact corollary of Lemma 4.3 state the unknown matrices N 11', Ø::
We first give a more detailed version which explicitly
Lemma A.1. In the setup of Theorem 4.1 with diagonal parameters 9
one can recursively compute
matrices u., w using the following formula
1 = =(11 (al))) + w wee u'
+ a'w'ay (M' + a'rd(ww)) (w*) E (d'w" -w))
++1 = (1+ )'
- -wan (M TMM' +
-w²)
a' pu (M
where p = Et= P and initial conditions a = 1, w = 0, M = =T, -, =
Proof.
First, we compute the following matrix that appeared in Lemma 4.3 for the specific diagonal
case:
A' b d!
Σ
P Σ
((
(eeTT.s))) ob
(d)T
Km
=(
Why
wit a
wond )
w/ )
14

---

This implies that A' = w'a''' = waya', d = a' and d' = A. Next we rewrite a²:
y'
2 (a' - (wo', 1z) (i))((M + yzu')
Σ (a'r: + (a'w" w', a.))(( ) )x + rrad))
)zs (M' a'''' (w") ( - a' pu'.
Σ (a'w" x1) + ( Σ d'rfu'
Here the first step is by Theorem 4.1 the second step replaces y with (w", .xs) + ras the third step
uses the fact that Σ 742; = to get rid of the cross terms.
The remaining proof just substitutes the formula for of into Lemma 4.3
Now Lemma A.1
implies Lemma 4.4 immediately by setting A - (at)2p1 w EY
Is a a'w a' dlpo yo (M²)
+ a''l (w' s tuwiyy X IIo w/yz (M*) (M + a'' (w)) and Φ¹
A.3 Proof for Theorem 4.2
Proof. By Theorem 4.1 we know of = (w",xx;) for some wt. When >> d, with high probability the norm of y is on the order of ( (z))lll || ||. and the norm of y" is @(vn). Therefore we only need bound the correlation. The correlation equal
d
Ky* sn = |w1
(1)
Σ
ai)l
j=2
We know with high probability Σ 2 ( because E[x:"] 0. The second term be written (wo u) where v is vector whose coordinates are ei
0 and uj = -E x1;)::;;; for 2 S j S d, therefore with high probability ||w|| = O(vnd). Therefore, with high probability the cosine similarity is at most
(v w
1Vr", ne/lwd
1ly |1 0(1)
0(1) Im Ei 27 + Σn x1(1 23 2f)
n ww
s 00() |m| Σ si
1oo ||||||
n 1
< 00( v + Vnd
When n >> this can be made smaller than any fixed constant.
A.4 Proof for Theorem 5.1
In this section we prove Theorem 5.1 by finding hyperparameters for GD algorithm that solves least squares problems with very high accuracy. The first steps in the construction iteratively makes the data 2;'s 's better conditioned, and the last step is a single step of gradient descent. The proof is based on several lemma, first we observe that if the data very well-conditioned, then one-step gradient descent solves the problem accurately:
15

---

Lemma A.2. Given (I1, (11,31)..... (Fn: where E == I has eigenvalues between 1 and 1. Let w" == arg min 2 i=1 (3): (w, x2)) be the optimal least squares solution, then T = Ei Vidi satisfies 11 - w I 2/|
Proof. We write ys (5i w²) + ri. By the fact that w" is the optimal solution know ry's
satisfy Ei=1 rid; O. Therefore i = -i=1 Vitt = Di= (ex, w") = Ew". This implies
|/ w= || T) = Σ w c||w=||
Next show that by applying just the preconditioning step of GD)++ I matrix very quickly. Note that the Σ matrix is updated as Σ <- (I EI - -E), so an
get well-conditioned
eigenvalue of À in the original E matrix would become A(1 -2)2 The following lemma shows
that this transformation is effective in shrinking the condition number
Lemma A.3. Suppose vip 2 1.1, then there exists an universal constant e < such that choosing yU 1/3 implies
[a A(1 - -222 minxelusp X(1 >>1}² Sc.
On the other hand, 1///URRS where E 0.1, then choosing T 1/3 implies
maxxE(vsH) A(1 - 72)²
minxelp.0 A(1 7A)2 S 1 2
The first claim shows that one can reduce the condition number by constant factor in every step until it's small constant. The second claim shows that once the condition number small (1 + e), each iteration can bring it much closer to (to the order of 1 + Oc2)).
Now we prove the lemma.
Proof. First, notice that the function f(x) z(1 - 2 is monotonically nondecreasing for E [0, v) if yv 1/3 (indeed, it's derivative f'(z) = (1 - 7z) (1 3Yz) is always nonnegative). Therefore, the max is always achieved at r = V and the min is always achieved at x μ. The new ratio is therefore
v(1 11--701 ?
4/9
M(1 - 72)²² K -
"1-1/3x)²
When x 1.1 the ratio from 1.
4/9 [I-1/3x]² is always below
4/9
which is constant bounded away
(I-1/3.3)
When & = 1 + < 1.1, we can write down the RHS in terms c
v(1 72)²
4/9
+ ( /
1
)
(11 7/1)2 "i -1-1 1/3kk2
1
Note that by the careful choice of y, the RHS has the following Taylor expansion:
1
1 )
3e²
5
(1 + (1 + (1 1+
O(e*).
4
4
One can then check the RHS is always upperbounded by 26 when < 0.1.
With the two lemmas we are now ready to prove the main theorem:
Proof. By Lemma A.3 we know in (log log log 1/c) iterations, by assigning in the way of Lemma A.3 one can reduce the condition number of x to x < 11+ e/2r (we chose c/2rc here to give some slack for later analysis).
Let E' be the covariance matrix after these iterations, and v' o be the upper and lowerbound for its cigenvalues. The data I; 's are transformed to new data of Mx; for some matrix
16

---

GD
Dag
Ful
Adjusted evel loes
layers
layers
100
10
10%
layers
10"1
10"1
10-1
layers
layers
6
layers
Prediction after layers
Prediction after K layers
Prediction after K layers
Figure 5: Linear transformer models show consistent decrease in emrоr per layer when trained on data with mixed noise variance o ~ (O, 5). The error bars measure variance over 5 training seeds.
M. Let M AE-1/2 then since M' = AAT know A is matrix with singular values between v and vv. The optimal solution (w*)- M-Tw" has norm at most v V || I] Therefore by Lemma A.2 we know the one-step gradient step with w = Li=1 jr We2s satisfy
|| - (wI S (w' / // |/ The test data 2t also transformed to xz. AΣ 1/2 x,, and the algorithm outputs (rù, az)). so the error is at most /zl w" *| & (w' 1/)/////// well ( w |
By the choice of k' we can check that RHS is at most e||w" |||
A.5 Proof of the Theorem 5.2
Proof. The key observation here is that when n co, under the assumptions we have T. Therefore the ridge regression solutions converge to wy
limn- Di=1 Fyst I4o Ei VAz; and the desired output is Fg))
By the calculations before, know after the first-layer, the implicit is w wyz Σ U As long Wyr is constant, when on x is negligible compared to noise), therefore the output the second layer satisfies
we know 1 EE=(v!) (v!)² r² (as the part of that depend
2² (1 no w (1 +nooowww Wyy Waym
yisti
Therefore, as long as we choose wyz and Wyy to satisfy (1 + no no )ayz 14o2 when a = 01 or 02 (notice that these are two linear equations on Wyo and nwyyWWyy so they always have a solution), then we have lim 2² " for the two noise levels.
B More experiments
Here provide results of additional experiments that did not make it to the main text.
Fig. Sshow example unadjusted loss. Clearly, is virtually impossible compare the methods across various noise levels this way.
Fig. 7 shows per-variance profile of intermediate predictions of the network of varying depth. It appears that GD++ demonstrates behavior typical of GD-based algorithms: early iterations model higher noise (similar early stopping), gradually converging towards lower noise predictions. DIAG exhibits this patter initially, but then dramatically improves, particularly for lower noise ranges. Intriguingly, FULL displays the opposite trend, first improving low-noise predictions, followed by decline in higher noise prediction accuracy, especially in the last layer.
Finally, Table[[preeent comprehensive numerical results for experiments across various mixed noise variance models. For each model variant (represented by a column), the best-performing result is highlighted in bold.
17

---

Unadjusted eval loss
6
GD
Diag
4-
Full
ConstRR
AdaRR
2
TunedRR
0
2
4
6
Variance 0
Figure 6: Example of unadjusted loss given by directly minimizing 7 It is pretty hard to variation between comparable methods using this loss directly.
Method
Uniform o, (0, mmaz)
Categorical o, E S
0
2
6
[1.3]
[13355]
layer
GD+
1.768
1.639
1.396
1.175
1.015
0.907
0.841
0.806
1.007
0.819
DIAG
1.767
1.639
1.396
1.175
1.015
0.906
0.841
0.806
1.007
0.819
FULL
1.768
1,640
1.397
1.176
1.016
0.907
0.842
0.806
1.008
0.820
2 layers
GD++
0.341
0.295
0.243
0.265
0.347
0.366
0,440
0.530
0.305
0.427
DIAG
0.265
0.214
0.173
0.188
0.219
0.242
0.254
0.259
0.201
0.246
FULL
0.264
0.215
0.173
0.188
0.220
0.245
0.259
0.263
0.202
0.276
3 layers
GD+F
0.019
0.021
0.071
0.161
0.259
0.344
0.454
0.530
0.222
0.422
DIAG
0.013
0.015
0.048
0.087
0.109
0.118
0.121
0.123
0.098
0.119
FULL
0,012
0.015
0.049
0.075
0.101
0.117
0.124
0.127
0.076
0.113
4 layers
GD++
9.91e-05
0.014
0.066
0.160
0.258
0.344
0.454
0.530
0.222
0.422
DIAG
1.19c-04
0.006
0.024
0.041
0.050
0.059
0.065
0.073
0.043
0.062
FULL
1.63e-04
0.005
0.021
0.038
0.052
0.065
0.068
0.076
0.032
0.061
5 layers
GD
1.14e-07
0.014
0.066
0.161
0.265
0.344
0.454
0.530
0.222
0.422
DIAG
1.81e-07
0.004
0.016
0.029
0.041
0.051
0.058
0.062
0.026
0.051
FULL
1.79c-07
0.002
0.015
0.026
0.038
0.048
0.059
0.065
0.016
0.048
layers
GD+
2.37e-10
0.009
0.066
0.161
0.265
0.344
0.454
0.530
0.222
0.422
DIAG
2.57e-10
0.003
0.014
0.028
0.040
0.048
0.054
0.059
0.020
0.047
FULL
2.71e-10
0.002
0.014
0.025
0.036
0.044
0.052
0.059
0.011
0.043
layers
GD
2.65e-12 12
0.009
0.066
0.161
0.265
0.344
0.454
0.530
0.222
0.422
DIAG
2.50e-12 12
0.002
0.014
0.027
0.040
0.047
0.052
0.059
0.018
0.046
FULL
2.50e-12
0.002
0.010
0.025
0.035
0.047
0.050
0.057
0.010
0.035
Basclines
CONSTRR
0
0.009
0.066
0.161
0.265
0.365
0.454
0.530
0.222
0.422
ADARR
0
0.003
0.016
0.034
0.053
0.068
0.081
0.092
0.051
0.084
TUNEDRR
0
0.002
0.010
0.023
0.037
0.049
0.060
0.068
0.021
0.054
Table I: Adjusted evaluation loss for models with various number of layers with uniform noise
variance a, ~ (0, maa)). We highlight in bold the best results for each problem setup (i.e. each column).
18

---

GD++
Diag
Full
layers model
281 eval loss
0.75
0.50
0.25
Slayers motel
268 eval 1088
0.75
0.50
0.25
0.00
1.00
Isyers model
201 evgl l688
0.75
0.50
0.25
0.00
1.00
levers model
A0) ewel lesa
0.50
0.25
0.00
1.00
aayers model
201 86al 1e8
0.50
0.25
0.00
1.00
layers model
A/ evell cellos
0.50
0.25
0.00
1.00
lsvers model
20) aved (888
0.25
0.00
100
1300
10
Variarca
Variance 0
Variance
After layer
After 3 layers
After 5 layers
After 7 layers
After layers
After 4 layers
After 6 layers
After 8 layers
Figure 7: Layer by layer prediction quality for different models with a, (0. The error bars measure std over training seeds.
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction accurately state the paper's ' main claims: that linear transformers can learn optimization algorithms, that they maintain an implicit linear regression model each layer, and that they can handle noisy linear regression problems with varying noise levels, surpassing many reasonable baselines.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We acknowledge that the analysis is limited to linear transformers and doesn't consider the effects of non-linearities in larger models. We also mentioned in the paper that the algorithm discovered for mixed noise variance is not proven be optimal, and the analysis focuses primarily on the diagonal attention case.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
20

---

3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We provide the full set of assumptions and complete proofs for all theoretical results. The proofs are detailed and appear in the Appendix, which is clearly referenced in the main text. Each theorem and lemma is numbered, and all assumptions clearly stated.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: Yes]
We believe that we provide sufficient details on data generation, model variants, training parameters, and evaluation metrics to allow anyone to replicate the core experiments and validate the main claims.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g., in the of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
21

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No],
Justification: We talk about synthetic data only, which very easy to generate. The code, while net released at this stage, will be released after the paper's 's acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nnis..c)
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
/ httpssscccpblccgccccccooe /ttp:////uiiee//deSuuuuuiiiiii wwweeesuummsssooooooooo
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Ys]]
Justification: We provide sufficient detail on the training and test procedures. We specify the optimizer, learning rate, batch size, number iterations, and how hyperparameters were tuned (including adjustming the the learning rate for stability). The data generation process is also detailed, along with the evaluation metries and baselines used.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
22

---

Justification: Some figures (like Figs. 2 and 6) include the errorbars, however, the main figures (Figs. 1,3,5) are reported best of of training seeds. We have acknowledged that some runs have stability issues which result in models stuck in high energy region of the loss. These outliers would affect the mean significantly, since our loss is already at 10-2 level.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We disclose all the relevant information on the compute resources that we used produce the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics /ttiec/rrrrr////ppprrs...c/////p
Answer: [Yes]
Justification: Our paper's research focuses on theoretical analysis and simulated experiments linear regression, which doesn't present any obvious ethical concerns outlined in the NeurIPS Code of Ethics. The study doesn't 't involve human subjects, sensitive data, or potential real-world harms that would require further scrutiny.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
23

---

If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work focuses primarily on the theoretical analysis of linear transformers and their ability to learn optimization algorithms in simplified setting. As such, it does not directly translate into specific applications with immediate societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Our work does involve the release of any trained models or datasets, primarily focuses on theoretical analysis and simulated experiments. Therefore, safeguards for responsible release are not applicable in this context.
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
24

---

12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Our work does not utilize any existing code, data, pre-traennn pre-trained models. All experiments were conducted using synthetic data and models implemented specifically for this study.
Justification:
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name the license (c.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
This work does not introduce any new assets such as code, data, or pre-trained models. We focus theoretical analysis and use synthetically generated data for our experiments, so we are not releasing any assets alongside this paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots. if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This research does not involve any crowdsourcing or experiments with human subjects. Our study is purely analytical and based on simulated data.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
25

---

Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our research does not involve any human participants. The study is purely analytical and experimental, utilizing simulated data and not involving any interaction or data collection from human subjects. Therefore, IRB approval or considerations for participant risks not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
26


Reviews:
[{"review_id": "h0BxbMDUBT", "reviewer": "Reviewer_WFVs", "summary": "The paper tries to understand the reasons of the strong performance of Transformers. The authors study linear Transformers trained on linear regression problems. Moreover, the authors explores the problem of regression where the labels have variable noise levels.", "questions": "See above", "limitations": "NA", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The problems studied are important, especially with the wide adoption of Transformers.\n- The case of regression with variable noise is interesting \n- Theoretical analysis is presented", "weaknesses": "- Although the focus is on analyzing Transformers, the paper could have benefited from adding more empirical analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking time to look at our paper. We believe that the empirical evaluation provided in the paper is thorough and is appropriate to cover the claims and contributions presented in the paper. We would love to know what kind of empirical evaluation the reviewer has in mind that we should add?", "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": "CBlnouCXqu", "reviewer": "Reviewer_KrhZ", "summary": "This paper demonstrates that linear transformers maintain a weight vector for an implicit linear regression problem. The authors provide theoretical analysis showing that (1) linear transformers maintain a linear regression model at every block and (2) a diagonal parameterization of attention heads does not compromise the expressive power of the model. The authors support their theoretical findings with experimental validation.", "questions": "1. Comparison of attention matrices in full and diagonal parameterizations:\nGiven the similar loss patterns shown in Figures 2 and 3 for full and diagonal parameterizations, it raises questions about the final structure of the attention matrices. Specifically, is it possible that the full matrices converge to a near-diagonal structure during training? If so, what implications does this have for the choice between full and diagonal parameterizations in practice? If not, how do the differences in attention matrices reconcile with the similar performance observed?", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper extends von Oswald et al. (2023) by further investigating that each linear transformer layer maintains a weight vector that can be used for regression problems. Moreover, they show that GD++ is a second-order optimization alg for the least square problem. \n2. Experiments on three different parameterization matrices are consistent with the analysis.", "weaknesses": "1. Limited empirical validation of Section 5.3: The results in Section 5.3 suggest that the update of y occurs every two steps. However, the paper lacks empirical studies to illustrate this phenomenon, particularly in a diagonal parameterization scenario. Can the authors provide:\na) Concrete experiments demonstrating this two-step update pattern? \nb) Visualizations or quantitative analyses that show how this behavior manifests in practice?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments and valuable feedback. We address their specific points below:\n\n> Empirical validation of Section 5.3 is lacking that the update of y occurs every two steps.\n\nIndeed, since $w_{xy}$ controls the how the current $y_t^l$ prediction affects the prediction of $x_t^{l+1}$ for the next layer, it doesn't change the prediction $y_t^{l+1}$. It takes at least two layers for this effect to reach the prediction $y_t^{l+2}$. However, it is quite challenging to empirically isolate the effect of every component, since all of the elements work simultaneously at every layer. All four terms per layer are helping to improve the loss and it is not-trivial to show the exact effect of every term empirically. \n\n> Comparison of attention matrices in full and diagonal parameterizations.\n\nWe were quite surprised to see that the diagonal parametrization performs almost identical to the full one. As the reviewer has predicted, the learned attention matrices for the full case (see example in the PDF attached above) converge to the near-diagonal matrix, but with some additional structure as well. Our preliminary experiments with diagonal plus low-rank parameterizations yielded similar results to the full attention mechanism. However, given the comparable performance and interpretability of the diagonal approach, we chose to focus on this simpler diagonal model. \n\nIn practice, the choice between full and diagonal parametrizations involves several factors. The diagonal approximation is easier to interpret (only 4 terms per layer!), is faster and works just as well as the full within the settings we consider. However, it is quite possible that for different distributions of $x$, $w$ or $\\sigma$ beyond the ones that we consider in our paper, the difference between full and diagonal models would be more pronounced. We plan to investigate this further in future work.\n\nWe will make sure to include the attention weights matrices as well as the consideration above to the final version of the paper if accepted.", "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": "SH7raFpYRq", "reviewer": "Reviewer_xxDm", "summary": "In this paper, the authors study linear transformers trained on linear regression problems and prove that each layer of every linear transformer maintains a weight vector for an underlying linear regression problem. Furthermore, the authors consider the mixed linear regression problem with varying noise levels and empirically demonstrate the flexibility of linear transformers for this problem.", "questions": "What are potential practical uses of this insight?", "limitations": "n/a", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "**Overall well-presented**\n\nFor someone who is familiar with the field, the presentation is certainly very good. For others, it might not be super obvious that the linear transformer is trained on a bunch of other linear datasets with different noise levels. Finally, I didn't find the abbreviations before Section 2.2 particularly useful. They don't save a lot of space, but the reader has to go back to it. Not clear why the authors chose to use those.\n\n**Relatively Interesting Insight**\n\nI found it somewhat interesting that the transformer is doing so well in the varying noise levels scenario. But then again, might be not surprising given it had a lot of training data.", "weaknesses": "**Incremental Nature**\n\nThe authors build on top of a lot of existing literature. While this is not necessarily bad, the delta compared to previous work might be too small. Since I'm not an expert in this domain, I'll have to rely on the opinion of the other reviewers to confirm my understanding.\n\n**Unclear Baselines**\n\nI'm not sure what the baselines are for. I understand that it is useful to give an idea how well the transformer does in the task. However, the author make it look as if it was important to beat them. Furthermore, the authors claim that these methods have an advantage given that they can use matrix inversion. In my understanding, the transformer is trained on much more data which could technically used to improve the noise-level estimation for the linear models and improve their performance as well. To me it is not really clear why the authors are so defensive in this experimentation.\n\nThe authors highlight that the transformer does so much better in the varying noise levels scenario compared to the linear models. However, there was little to no effort to make the linear models good in this scenario in the first place.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer to take their time to evaluate the paper and for valuable suggestions on improving the presentation. Indeed, the linear transformer is trained on various generated sequences of noisy linear regression. We will make it more obvious in the introduction and preliminary section. We will also move the notation from the end of Section 2.1 closer to the Theorems where they are actually being used.\n\n> I found it somewhat interesting that the transformer is doing so well in the varying noise levels scenario. But then again, might be not surprising given it had a lot of training data.\n\nWhen training the model, the amount of the training data is just one factor. Another big factor to consider is the capacity of the model. Not every model can fit any given function. The original motivation for this paper came from the surprising observation that Transformers have the capacity to solve certain problems provided in-context. Previous research has demonstrated that for a single noise level this ability might be due to the transformers implementing a form of gradient-based optimization on an implicitly defined objective function. For multiple noise levels previous work (Bai et al.) relied on complicated constructions with large networks and nonlinearities. In this paper we demonstrate that even _linear_ transformers with _digaonal_ attention metrics have enough capacity to solve quite complex and non-obvious problems or noisy linear regression. \n\n> ..the transformer is trained on much more data which could technically be used to improve the noise-level estimation for the linear models and improve their performance as well. To me it is not really clear why the authors are so defensive in this experimentation.\n\nThank you for this observation. What we were trying to convey is not defensiveness, but a genuine surprise by the results that we have observed! The problem of linear regression with variable noise level, while simply stated, is practical and complex. The fact that a simple linear transformer can learn quite a sophisticated algorithm that works on par or better than many baselines that us, humans, can come up with, is quite surprising.\n\n> What are potential practical uses of this insight?\n\nWe think that there are several implications of our work. One, the algorithm discovered by the linear transformer, with its adaptive rescaling and momentum-like behavior, could be directly applied to real-world noisy linear regression problems in domains such as robust control, time series analysis, or finance. Second, our work strengthens the evidence for transformers' ability to implicitly learn sophisticated algorithms, opening up exciting possibilities for automated algorithm discovery in other machine learning tasks. \n\nIt would be interesting to see how variations in problem complexity and structure affect both the performance of the transformer and its ability to discover novel algorithms. This could involve exploring different data distributions, underlying function classes, and noise structures, ultimately leading to a deeper understanding of the factors influencing algorithm discovery in transformers and potentially uncovering a broader class of implicit algorithms with practical applications in various domains.", "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": "oJd2xRV0B9", "reviewer": "Reviewer_dAYB", "summary": "This paper proves that linear transformer layers maintain a weight vector for implicit linear regressions, including a more challenging scenario where data is corrupted with different levels of noise. In the theoretical analysis, this paper shows the intrinsic mechanism of the  gradient scent in linear attention, where implicit variants play various roles in parameter update. Besides, this paper shows that simplified linear attention with diagonal transformation also maintains a powerful GD behavior. The experiments demonstrate the flexibility of linear transformers which outperforms GD++ as well as other linear regression solutions. This research promotes the understanding of Transformer weights and implicit learning capabilities of attention-based models.", "questions": "1. How do you think about different linear attention variants in theory? For example, the most trivial one is $QK^TV$, but there are many improved versions, such as RetNet, SSM, and GLA. Do these architectures bring stronger modeling capability in theory?\n2. With linear regression layer-wisely, does your claim that linear Transformer with FFN only maintains linear regression still work? For example, in Theorem 4.2, it is easy to build a quadratic FFN to make more complex computations.", "limitations": "The limitation is discussed in the conclusion section.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The theoretical analysis in this paper is sound and insightful for understanding the behavior of linear attention.\n2. This paper broadens the analysis framework from fixed noise linear regression to flexible noise, making a step forward to more complex problems.\n3. The experiment settings can prove the mathematical formulation in this paper.", "weaknesses": "The analysis framework is still on linear regression and linear attention with simplified modifications, which is similar to previous references. The theory can not be generalized to other architectures and context settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for raising these insightful questions, which highlight the importance of generalization of our findings. While a complete exploration of these aspects is beyond the scope of our current work, we would like to provide some initial thoughts and discuss potential future directions.\n\n> How do you think about different linear attention variants in theory? For example, the most trivial one is $QK^TV$, but there are many improved versions, such as RetNet, SSM, and GLA. Do these architectures bring stronger modeling capability in theory?\n\nWhile we focused on simple attention for clarity and ease of analysis, these more advanced architectures could indeed offer stronger modeling capabilities in theory. They can implicitly capture higher-order feature interactions. These architectures often employ techniques like kernel functions or randomization, which could allow them to implicitly capture higher-order interactions and learn richer function classes. More specifically, we believe our Theorem 4.1/4.2 would generalize to the linear version of RetNet, but would not generalize to SSM and GLA (because SSM is time-varying and GLA has nonlinearities in the form of gates). Even in the setting of RetNet, the specific algorithms discovered could be more sophisticated. Exploring how our findings extend to these richer attention mechanisms is a promising direction for future work. \n\n> With linear regression layer-wisely, does your claim that linear Transformer with FFN only maintains linear regression still work? For example, in Theorem 4.2, it is easy to build a quadratic FFN to make more complex computations.\n\nIndeed, our current analysis considers linear transformers without FFNs. This simplification allows us to isolate the core behavior of the attention mechanism. With nonlinear FFNs, Theorem 4.2 no longer holds. A quadratic FFN can indeed introduce higher-order relationships, however the output of each attention layer is still a linear combination of its inputs (though the weights might be non-linear functions of the input). Exploring the complex interaction between attention and FFNs in the context of implicit optimization is crucial for future research and may require more advanced tools and techniques.", "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:
Here, we are attaching a PDF with example of learned weights for Full parametrization as requested by Reviewer KrhZ. The learned weights converge to near-diagonal matrix, which inspired us to try the diagonal parametrization.

Meta-review: This paper provides a theoretical analysis of the performance of linear transformers under in-context learning for linear regression tasks.  The authors show that each layer of the transformer maintains an implicit linear regression model and can be interpreted as conducting a preconditioned gradient descent.  The authors also show that even using diagonal attention matrices in this case the linear transformers can still offer good expressive capability.  When dealing with mixed linear regression problems with various levels of noise, the linear transformers can yield superior performance over baselines, which can be attributed to an intrinsic optimization process with noise-aware stepsize adjustments and rescaling. The paper introduces an interesting theoretical interpretation of linear transformers under in-context learning.  The rebuttal and following discussions have cleared most of the concerns raised by the reviewers and authors also put up additional results on attention matrices.  The authors are encouraged to improve the empirical investigation to make the paper stronger.
Decision: Accept (poster)