In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness
Authors: Liam Collins
Advait U Parulekar
Aryan Mokhtari
sujay sanghavi
Sanjay Shakkottai
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Transformers
self-attention
in-context learning
Primary area: learning_theory

Abstract: A striking property of transformers is their ability to perform in-context learning (ICL), a machine learning framework in which the learner is presented with a novel context during inference implicitly through some data, and tasked with making a prediction in that context. As such, that learner must adapt to the context without additional training. We explore the role of *softmax* attention in an ICL setting where each context encodes a regression task. We show that an attention unit learns a window that it uses to implement a nearest-neighbors predictor adapted to the landscape of the pretraining tasks. Specifically, we show that this window widens with decreasing Lipschitzness and increasing label noise in the pretraining tasks. We also show that on low-rank, linear problems, the attention unit learns to project onto the appropriate subspace before inference. Further, we show that this adaptivity relies crucially on the softmax activation and thus cannot be replicated by the linear activation often studied in prior theoretical analyses.

In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness
Liam Collins"
Advait Parulekar'
Chandra Family Department of ECE The University of Texas at Austin liamc@utexas edu
Chandra Family Department of ECE
The University of Texas Austin
edvaitplutexas.. edu
Aryan Mokhtari
Sujay Sanghavi
Chandra Family Department of ECE The University of Texas at Austin mokhtar @austin. utexas edu
Chandra Family Department ECE
The University of Texas Austin sanghavi@mall utexas. edu
Sanjay Shakkottai
Chandra Family Department of ECE The University of Texas at Austin sanjay maakkottai@ttezen edu
Abstract
A striking property of transformers their ability to perform in-context learning (ICL), machine learning framework in which the learner is presented with novel context during inference implicitly through some data, and tasked with making a prediction in that context. As such, that learner must adapt to the context without additional training. We explore the role of softmax attention in an ICL setting where each context encodes regression task. We show that attention unit learns window that it uses to implement a nearest-neighbors predictor adapted the landscape of the pretraining tasks. Specifically, widens with decreasing Lipschitzness and increasing label noise in the pretraining tasks. We also show that on low-rank, linear problems, the attention unit learns project onto the appropriate subspace before inference. Further, we show that this adaptivity relies crucially the softmax activation and thus cannot be replicated by the linear activation often studied in prior theoretical analyses.
show that this window
Introduction
One of the most compelling behaviors of pretrained transformers is their ability to perform in-context learning (ICL) [1]: determining how solve unseen task simply by making forward pass on input context tokens. Arguably the most critical innovation enabling ICL is the self-attention mechanism [2], which maps each token in input sequence
token using information from all other tokens. key design choice in this self-attention architecture is of the activation function that controls how much "attention" token pays other tokens. Sefimax-activated self-attention (i.e. softmax attention) is most commonly, and successfully, used in practice [1, 3-6].
A natural approach to explain ICL adopted by the literature is to equate it with classical machine learning algorithms, primarily variants of gradient descent (GD). Several works have shown that
"Co-fi first authors, listed in alphabetical order.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

WV
Atte ntion Window
Aquery
Attortton We gihts
Softmax Linear
Position Relative to
Softmax
Error
Linear
Number of Context Samples
The black line denotes the target function over domain (horizontal axis). The gray dots are noisy training data, and the white dot is a query. From left to right, the Lipschitzness of the target function grows and the optimal softmax attention window (shaded blue) shrinks. Middle Row:
Figure 1: Top Row:
Attention weights - which determine the attention window - as a function of the relative position from the query for softmax and linear attention. The softmax weights adjust to the Lipschitzness. Bottom Row: ICL error versus number of context samples for the three settings. Adapting to function Lipschitzness leads softmax attention to achieve small error. Please Remark 2.1 and Appendix for further discussion and details.
when the ICL tasks are linear regressions and the activation in the attention unit is identity (referred as linear attention), transformers that implement preconditioned GD during ICL are global optima of the pretraining loss, which is the population loss on ICL tasks [7-9]. In particular, the prediction of such transformers with / linear attention layers equals the prediction of a regressor trained with / preconditioned GD steps on the context examples. However, since these analyses are limited to linear attention and tasks, they do explain the widespread success of softmax attention ICL.
More recent work [10] extends these results by showing that for general regression tasks and any attention activation that is kernel, ICL equates training kernel regressor via functional GD in the Reproducing Kernel Hilhert Space (RKHS) induced by the activation. However, this functional GD yields generalization guarantees only when the activation kernel is identical generates the labels, which does not apply to the softmax activation, as it is not a kernel. Further, like the aforementioned studies of the linear setting [7-9], this analysis only shows that pretraining leads to learning the covariate distribution, while the activation implicitly encodes the label distribution needed for accurate predictions. Thus, this line of work has not explained the very fundamental question of what softmax attention learns during pretraining that enables it to perform ICL on wide variety downstream tasks. Motivated by this gap in the literature, we ask the following question.
kemee that
How does softmax attention learn to perform ICL?
To answer this question, we study general settings which pretraining and evaluation ICL. tasks are regressions that share only Lipschitzness and label noise variance. Specifically, the rate which their ground-truth labels change along particular directions in the input space, and the variance the label noise, is similar across tasks. In such settings, we observe that softmax attention acts nearest neighbors regressor with an attention window - i.e. neighborhood of points around the query that strongly influence, or "attend to",, the prediction that adapts the pretraining tasks. Specifically, our main result is as follows:
Main Claim: Softmax attention performs ICL by calibrating its attention window the Lipschitzness and label noise variance of the pretraining tasks.
While this does not contradict the line of work showing that ICL manifests via "meta-learned" gradient-based algorithm, we show in a general setting that a simpler mechanism can explain the capabilities of widely accepted model of ICL.

---

Outline. We substantiate the above claim via two streams of analysis. To knowledge, these the first results showing that softmax attention pretrained on ICL. tasks recovers shared structure among the tasks that facilitates ICL on downstream tasks.
(1) Attention window scale adapts to Lipschitzness and noise variance Section 3. We prove that the pretraining-optimaal softmax attention estimator scales its attention window inversely with the task Lipschitzness and jointly with the noise level to optimally trade-off bias and variance in its prediction (Theorem 3.4). This requires tight upper and lower bounds on the pretraining ICL loss. While the upper bounds (Lemma C.8) hold for all L-Lipschitz tasks, the lower bounds (Lemma C.9) are more challenging and require considering specific classes of tasks. We consider two classes of generalized linear models (GLMs) MG), and obtain lower bounds via novel concentrations for particular functionals on the distribution of the attention weights for tokens distributed on the hypersphere (Corollary G.5).
(2) Attention window directions adapt to dtrection-wise Lipschitzness - Section 4. We prove that when the target function class consists of linear functions that share common low-dimensional structure, the optimal softmax attention weight matrix from pretraining projects the data onto this subspace (Theorem 4.4). In other words, softmax attention learns to zero-out the zero- Saro-Lipschitznesss directions in the ambient data space, and thereby reduces the effective dimension of ICL. We prove this via careful symmetry- argument to characterize a particular gradient of the ICL loss as positive (Lemmas H.3 and H.4).
Tightness of results. Our results highlight the importance of shared I Lipschitzness across training and test, as well as the critical role of the softmax activation, ICL. We show that softmax attention pretrained on the setting from Section 3 in-context learns any downstream task with similar Lipschitz- ness to the pretraining tasks, while changing only the Lipschitzness of the evaluation tasks degrades performance (Theorem 3.5) implying learning Lipschitzness is both sufficient and necessary for generalization. Further, to emphasize the necessity of the softmax, we show that the minimum ICL loss achievable by linear attention exceeds that achieved by pretrained softmax attention (Theorem 3.6). We verify all of these results with empirical simulations (Section 3.2 and Appendix J).
Notations. We use (upper-, lower- ccase boldface for (matrices, vectors), respectively. We denote the (identity, zero) matrix in Raddd as (Ia, dd) respectively, the set of ootumn-orthonormal matrices in Bax* as O and the (column space, 2-norm) of matrix B as (B))(B), |[B ()) respectively. We indicate the unit hypersphere in Rd by gd-1 and the uniform distribution over Sd-1 as Ud, We use asymptotic notation (O, !2) hide constants that depend only on the dimension d.
1.1 Additional Related Work
Numerous recent works have constructed transformers that can implement GD) and other machine learning algorithms during ICL [11-15), but is unclear whether pretraining leads to such transform- ers. [16] and [13] provide generalization bounds for ICL via tools from algorithmic stability and uniform concentration, respectively. [17] investigate the pretraining statistical complexity of learning Bayes-optimal predictor for ICL on linear tasks with linear attention. [18-20] study the role of the pretraining data distribution, rather than the learning model, in facilitating ICL. [21] studies the dynamics of a softmax attention unit trained with GD on ICL tasks, but this analysis considers only linear tasks and orthogonal inputs. The connection between ICL with softmax attention and non-param non-parametrie regression has been noticed by other works that analyze the ICL performance of a softmax-like kernel regressor [22] and aim to improve upon softmax attention [23-27] rather than explain what it learns during pretraining. Please see Appendix A for further discussion of the large body of related works studying the theory of transformers, ICL and kernel regression.
2 Preliminaries
In-Context Learning (ICL) regression tasks. We study ICL in the regression setting popularized by [28], wherein each task regression problem Rdd The context for task f. consists of a set of n feature vectors paired with noisy labels {x, (t) (x) (t) + G
(t) (z) jm., where f(t)_ Rd + I
generates the ground-truth labels for task and <: label noise. Given this context, the model (1)
solves the task if it accurately predicts the label of query x.t1- During pretraining, the model
3

---

observes many such tasks. Then, is evaluated on new task with context {m,, (ar)) + (1) } and query a We emphasize that the model is trained only on the pretraining tasks, not the evaluation context. Unlike traditional supervised learning, which would involve training on the context { (xx + of in order to predict f(*) ICL happens entirely in a forward pass, so there is no training using labels from f()), Our inquiry focuses on how ICL is facilitated by the softmax activation in the self-attention unit, which we introduce next.
)==1
The Softmax Attention Unit. We consider a single softmax attention head Hs.(:: ()) R(4+1x(=++) + R(++))(x+1) parameterized by 0 == (W. Wq. Wv), where W K: Wq. Wv E R(4+1)x(d+1) are known as key, query, and value weight matrices, respectively. Intuitively, for a se- quence of tokens Z = [zi,.... sn+2] E 2((+1)x(n=1) the attention layer creates a "hash map" where the key-value pairs come from key and value embeddings of the input tokens, {Wx k 2; Ww 2 Each token 2, is interpreted as a query W Q ns, and during a pass through the attention layer, this query is matched with the keys (W к æj}l to return average over the associated values {Wy z5}; with a weight determined by the quality of the match (proportional to cWx x5) (Wa si)). Specifically, Hsa (Z; 0) = [ Z) ... hSA(Znn ( , Z; 0)], where
hs 0. ΣØ=1 (Wvz,) eWa,,,) (Waz:)
Rd+1
(ATTN)
Ex ( ) (Waz,)
With slight abuse of notation, we denote ( = hs , Z; 0) when is not ambiguous. To study how this architecture enables ICL, we follow [28] to formalize ICL as regression problem. Below we define the tokenization, pretraining objective and evaluation task.
Tokenization for regression. The learning model encounters token sequences of the form
Z
11 [(x1) +1 f(x2) +
32
En f(x
Fn+
R I(d++)x(++))
(1)
6n
where the ground-truth labelling function f maps from IRd to R and belongs to some class F, each 6, mean-zero noise, and the i-th input feature vector æ; E Rd is jointly embedded in the same token with its noisy label (x;; + e& E R.. We denote this token z;. The ICL task is to accurately predict this label given the n context tokens ((r,, ++))= where f may vary across sequences. The prediction for the label of the (n+1))tth feature vector is the (d+1)-th element of hs (Zn+1) [10], denoted hs (zn+1)d++1 Ultimately, the goal is to learn weight matrices such that hs Bn+1)d+11 is likely to approximate the (n + 1)-th label on random sequence Z.
We study what softmax attention learns when its weight matrices are pre- frained using sequences of the form of (1). These sequences are randomly generated as follows:
Pretraining protocol.
f D(F), LT In+11
iid @)(#+1),
Ely..., En Lid
D)(+11)
(2)
where ((F) is distribution over functions in F, Da is distribution over T&d, and D. is distribution over R with mean zero and variance o². The token embedding sequence then constructed as in (1). Given this generative model, the pretraining loss of the parameters O = (Wo. Q. WK. Wv) is the expected squared difference between the prediction of softmax attention and the ground-truth label of the (n 1)tth input feature vector in each sequence, namely
2( := Eeffee (hsA(Zn+1)441 - f()x+1))
We next reparameterize the attention weights to make (3) more interpretable. For the last column of Wv, we show in Appendix B that any minimizer of (3) in the settings we consider must have the first d elements of this last column equal to zero. We follow [7, 9, 10] by setting the first n columns of Wy to zero. As in [10], we fix the (d+1, d+1))th element of W here as I for simplicity. In the same vein, we follow [7, 10] by setting the (d+1))-hh row and column of W k and Wa equal to zero To summarize, the reparameterized weights are:
Oaxd 0dx1
Mæ
Odx1]
MQ
0ax
W =
W
Wq =
(4)
01xd 1
01xd
0
0
0

---

where MR, MQ E Rdxd, Now, since goal is reveal properties of minimizers of the pretraining loss, rather than study the dynamics of optimizing the loss, without loss of generality we can define M M2 Mq and re-define the pretraining loss (3) as function of M. Doing so yields:
2
(f=(((())) (.) e Mæ++1
((M) -BAaaedled.
(ICL)
Σi= ex! M M=.+1
Interpretation of the pretraining loss. The loss (ICL.) clarifies how softmax attention can be inter-
preted as a nearest neighbors regressor. When x; MIn+1 proxy for the distance between æ; and xx+11 (which we formally show in Section 3 as happening under reasonable assumptions), the softmax attention prediction a convex combination of the noisy labels with weights determined by the closeness of to En+1, such that the labels of points closer to 2n++1 have larger weight. Moreover, the decay in weights on points further from #4+1 is exponential and controlled by M, which effectively defines a neighborhood, or attention window, of points around 2n+11 whose labels have non non-trivial weight. More formally, we can think of the attention window defined for query 1n++1 as the set AttnWindow(Ww... (woddoo((n++3 M) := {æ x M In++1 (2(1). As we have observed in Figure I, our key insight is that pretrained M scales this attention window with the Lipschitzness of the function class. Generally speaking, larger M entails averaging over smaller window and incurring less bias due to the function values of distant tokens in the estimate, while smaller M entails averaging over larger window, resulting in larger bias due to distant token labels, but smaller noise variance. Figure 2 further depicts this tradcoff.
Connection to non- parametric estimation and the Vadaraya-Watson Watson estimator, A nonparamet-
ric estimation technique interpolate between known values of a function is to use kernel estimator. The Nadaraya- -attoo (NW) estimator [29-31] is one such estimator, and interpolates the data as
K (() E, K(((141,22))
fNw (In+1) = E
where K(r) = 1//h for some bandwidth h. In Section B.1 we show that optimizing the pretraining loss (ICL) reduces to meta-learning the bandwidth of an NW estimator tasks. However, to our knowledge, the literature has not determined the optimal bandwidth for the kernel, there has been analysis of non-asympiotic lower bounds on the loss, which we need to characterize the optimal solution. A close work to ours [32], which considers regression on general L-Lipschitz tasks, but this analysis provides only tight upper bound on the loss.
distribution of pretraining
Remark 2.1 (Extreme cases). Consider the following two settings.
(i) Constant functions. If each of the functions the attention unit sees in pretraining is constant, as in the Left column of Figure 1, it is best to consider an infinite attention window, that is, take M Oaxa as this results a uniform averuge over all the noisy token labels.
(i)) Rapidly changing functions. If the pretraining functions change rapidly, as in the Right column of Figure attending to distant token might only corrupt the estimate at the target. For example suppose the input tokens are used to construct Voronoi cells on the surface of the hypersphere, and the label for a new token in a cell the label of the token used to construct that cell. The optimal estimator attends only to the single nearest token since this incurs error only from label noise.
Remark 2.2 (Softmax advantage). To further highlight the utility of the softmax, we compare with linear attention [7, 9, 111, whose estimator can be written as (LA(x) - -I,(((x)) + +e)) x. M æ. up to universal scaling due to the value embedding. This is again a weighted combination of labels. but one that does not allow for adapting an attention window - any sealing of M does not change the relative weights placed on each label - unlike softmax attention. Please see Figure I (Middle Row) for a comparison of the weights used in the different estimators.
3 Pretraining Learns Scale of Attention Window
One of our observations of the attention estimator Işy is that it computes nearest neighbours regression. We hypothesize that the role of pretraining to select neighbourhood within which to select tokens for use in the estimator. In this section we characterize the radius of this neighborhood.
5

---

Figure 2: From left to right, as we shrink the attention window (shaded in blue), the estimator has lower bias (the expected value of the estimate, depicted in purple, is closer to the ground-truth label, depicted by the white circle) but larger variance (shaded tan).
Definition 3.1 (Lipschitzness). A function f X + has Lipschitzness L if the smallest number satisfying J' (a) - f'(x') S Ell x x' [or all (xr, E X22
The general requirement for the function classes to which our results apply is that the class should be invariant to isometries, each function should be Lipschitz, and the function value at two points should less correlated as those points get further. These are written formally in Assumption B. 4. To be concrete, work with the following two function classes that satisfy these assumptions (this shown in Lemmas C.3 and C.7) to derive explicit bounds.
Definition 3.2 (Affine and ReLU Function Classes). The function classes FL and F are respec- tively defined as:
F = (f f(x) =1 =l T x+ W WES .b [-,,
= (f: f(x)=
= (w(w' x)+ + 12(-- w x)+ + bb, w E (., E [[-L,
((FT).D(F)) are induced by drawing Eu and b, I, 11 li Unif(l-L. L, L)) for some E > Odxd. Note that the max Lipschitzness of any function in these classes is L. and (z)+ == max( (2, 0).
Next, make the following assumption, similar to [7], on the covariate distribution.
Assumption 3.3 (Covariate Distribution). The covariate distribution satisfies D = E -E-14 U
Now we are ready to state our main theorem that characterizes minimizers of (ICL).
Theorem 3.4. Lel Assumption 3.3 hold and tasks f be drawn from (Case 1) D(FT.) or (Case 2) D(Ft). ). For (1) and ((n-d/2 I o² < O(T.2), any minimizer of the pretraining loss (ICL. satisfies? M = -uk E, where for = oz, == J44 and Jt2:
nf.²
(Case 1) Ω (A9) < |W 0 (A=)), (Case 2) Ω (A S [wKQl S 0 (A20).
Theorem 3.4 shows that optimizing the pretraining population loss in Equation (ICL) leads attention key-query parameters that scale with the Lipschitzness of the function class, well as the noise level and number of in-context samples. These bounds align with our observations from Figures and that softmax attention selects an attention window that shrinks with the function class Lipschitzness, recalling that larger W'Q results in a smaller window. Further, the dependencies the bounds on o² and n also intuitive, since larger noise should encourage wider averaging average out the noise, and larger n should encourage a smaller window since more samples makes it more likely that there samples very close to the query. To knowledge, this is the first result showing that softmax attention learns properties of the task distribution during pretraining that facilitate ICL.
Learning Lipschitzness is critical to generalization. We next give the following generalization result for downstream tasks.
Theorem 3.5. Suppose softmax attention first pretrained on tasks drawn from D(FL) and then tested on an arbitrary L-Lisschitz task, then the loss on the new task is upper bounded as L <
We further show in Appendix B that M = -KQQ E for scalar uxQ holds for a broad family of rotationally. invariant function classes.

---

15otropic
ReLU Non-Isotropic
Cos Isotropic
Cos Non-Isotropic
100
100
80
80
40
2.2
33
500
1030
1520
500
1000
1500
506 1003
1500
500
1000
1500
2003
Pretraining Iteration
Pretraining Iterations
Pretraining Iterations
Pretraining Iterations
Figure 3: Spectral norm of M during pretraining with varying L. Each plot shows results for different task and covariate distributions, with (tasks, covariates) drawn from (Left) (22)) LMiddle-Left) (D(F)) ), (Middle-Right) (D(FL ), (Right) (D(FE")) where U is non-isotropic distribution on Sd-1 (see Section 3.2 for its definition).
O ((=)) Furthermore, the new task is instead drawn from D(FD.) the loss is lower bounded as C 21 Reo for L' L and 2 (222
A64/2 for L' L.
Theorem 3.5 shows that pretraining D(FI) yields model that can in-context learn downstream tasks if and only if they have similar Lipschitzness as L. Thus, learning Lipschitzness is both sufficient and necessary for ICL. If the evaluation task Lipschitzness is much larger than that seen in pretraining, the pretrained model will give highly biased estimates. Conversely, if the evaluation Lipschitzness is much lower, the pretrained model will not optimally average the label noise.
Necessity of Softmax. To further emphasize the importance of the softmax in Theorem 3.4, we next study the performance of an analogous model with the softmax removed. We consider linear self-attention [7, 9, 11), which replaces the softmax activation with an identity operation. In particular, in the in-context regression setting study, the prediction of fxx++1) by linear attention and the corresponding pretraining loss are given by:
hLA(Xn+1) = (f(e)) + t M M#n+1,
1=1
(M(M) == Effaalaffa (hLA(En+1) -f fxx+11)2
As discussed in Remark 2.1, h (1r41) cannot adapt attention window to the problem setting. We show below that this leads it to large ICL loss when tasks are drawn from D(F))
Theorem 3.6 (Lower Bound for Linear Attention). Consider pretraining on LIa with tasks I drawn from D(Ft) and covariates drawn from Ud, Then for all M E Rdxd, LL. C1.(M) = R(L2).
This lower bound on LLA is strictly larger than the upper bound on from Theorem 3.5, up to factors in d, as long as a² < 1, which holds in all reasonable cases. Please see Appendix F for the proof.
3.1 Proof Sketch
To highlight the key insights of our analysis, in this section we consider a modification of the softmax attention that exhibits important properties the original. Note that this approximation for illustration only; the above results use the original softmax attention see Appendices C, D, E. For now, consider a function class FL = {f f(x) Lw x, w e gd-1) of linear functions.
(Temporary) modification of the softmax attention. Rather than averaging over every token with a weight that decays exponentially with distance, consider modification which wniformly averages all tokens within a distance specified by wKQ = MJ] From Lemma B.5, without loss of generality (WLOG) can consider M WKQLs. This means that, ignoring normalization, the weight assigned to f(x;)) by the true soft-max attention is e=wxll æ æ, il That is, for all æ; satisfying x| x; || < 2/ /wRq, the assigned weights within constant factor of each other. Meanwhile, for æ; satisfying || æ - x; v for c > 1, the weights are e decaying exponentially in This f())1. motivates to consider "modified softmax attention" given by bMSA (x) =; Σ,
where
1, 1|ll x, < 1/ //
7

---

The In-context Loss. The pretraining loss in Equation ICL can be decomposed
Σ (f(an+1)= 1111))111
(Σ
Σ G1 Eg
L( = E
+ Beeshilei
=iCggna (WxQ)
We first upper and lower bound each of these terms separately, starting with Lragmal (uxQQ.
Noiseless Estimator Bias. (Please see Appendix C) This term is the squared difference between an unweighted average of the token labels within a radius of æ, and the true label. Take wKQ (1)). Then for large d, most of the points 2; satisfying xc 2c, || S 1/ γक lie on the boundary of the = 1/ - This motivates us to approximate the
cap, that is, || x 2; || < 1/v/qq + || æ set of points x; satisfying the above as coming from uniform distribution over just the boundary of the cap. The center of mass of a ring of radius 1/ f/eeq embedded on the surface of a hyper-sphere, is (1/waa) from the boundary of sphere, so the squared bias is O /wxa))
Noise, (Please see Appendix D for details) Since the noise is independent across tokens, we can write Cnoise Croise(wKK) (WKQ) Σ, T which is related to the number of tokens found within 1//wxq radius of æ. In lemma G.1, we derive bounds for the measure of this region. For now, we replace the sum in the denominator with its expectation. We can bound E, = O (w as long as WKQ S n2/d.
o²
Combining the Lsignal and Lnoise terms. (Please see Appendix for details) Overall, we have C = Liignal + Looise with fisignal ((22/wxa) and Lanise = 0 ("kk0=/))) Minimizing this reveals that the optimal wxQ satisfies wxq =
3.2 Experiments
We next empirically verify our intuitions and results regarding learning the scale of the attention window. In all cases we use the Adam optimizer with one task sampled per round, use the noise distribution D. (0, o²), and run 10 trials and plot means and standard deviations over these 10 trials. Please see Appendix J for full details as well as additional results.
Ablations over L, a and n. We verify whether the relationship between the attention window scale ie. |M] and L, and n matches our bounds in Theorem 3,4 for the case when tasks are drawn from D(Ft) and the covariates are drawn from 2td, as well as whether these relationships generalize to additional function classes and covariate distributions. We train on tasks drawn from D(Ft) and D)F where FL == [ f()) cos (Lw x), E 36-1) and D(FE is induced by sampling ~ Ldd In all cases set d 5, and use (I, (1, 0.01, 20) if not ablating over these parameters, and vary only one of (L. o, n) and no other hyperparameters within each plot.
Attention window scales inversely with L. Figure shows that [M]] increases with L in various settings. In Figure 3(eeft, Middle-Leftt, tasks are drawn from D(Ft), and in Figure M Right, Right), they are drawn D(FF In Figure 3(Left, Mighle-Right), each æ; is drawn from 4dd, whereas in Figure 3(Middle-Left, Right), each x; is drawn from non-isotropic distribution i7dh on gd-1 defined as follows. First, let Sa == diag([L, .......... d) E Rddd then æ Ltdt is generated by sampling s 1//22 Although larger L implies larger V=f(()) on
~ N (Oa, Ia), then computing x = |
CMM,C(WWW average across /, it is clear that it implies larger |/V L(W) K Wo)ll || L(W] K W W))l so it is surprising that larger I. implies larger pretrained M (although is consistent with our results).
Attention window scales with 0, inversely with n. Figure 4 shows that the dependence of ||M] on o and n also aligns with Theorem 3.4. As expected, || increases slower during pretraining for larger (shown Figures 4(Left, Middle-Left), since more noise encourages more averaging over a larger window to cancel out the noise. Likewise, ||M]] increases faster during pretraining for larger n (shown in Figures 4(Middle-B Right, Right)), since larger increases the likelihood that there is highly informative sample in a small attention window. Here always the covariate distribution is tt
8

---

ReLU Varying d
Cos Varying d
ReLU Varyingg
Cos Varying n
100
200
10
0.5
500
0.0
500
1000
1500
0
500
1003
1506
1500 Pretraining Iterations
503
1030
502 1000 2500 Pretraining Iterations
2009
Pretraining Iteration:
Pretraining Iterations
Figure 4: Spectral norm of M during pretraining tasks drawn from D(F)) in Left, Middle-Right and ((F in Middle-Left, Right. Left, Middle-Left show ablations over the noise standard deviation a and Middle-Right, Right show ablations over the number of context samples n.
Learning new tasks in-context. An implication of our work is that for the function classes consider, the softmax attention estimator does not adapt to the function class beyond its Lips- chitzness. We have already seen in Figures and 4 that the growth of |[M]] during pretraining is similar across different function classes with the same Lipschitzness, as long as a and n are fixed. Here verify the conclusion from Theorem 3.5 that for fixed n and ø, the necessary and suffi- cient condition for downstream generalization, measured by small ICL error, is that the pretraining and downstream tasks have similar Lipschitzness. Figure 5 supports this conclusion. Here we set d = 5, n 200, a = 0.01 and draw each æ; i.i.d. from 4t, In Figure 5(Left, Middle-Left, Middle- Right), we train three attention units on tasks drawn from the 1-L.ipschizz affine (D(FT)), D) ReLU (D(FT)) and cosine (D(FF)) task distributions. Each plot shows the test ICL error on tasks drawn from a distribution in {D(Ftt). D(FT), D(FT")). Performance is similar regardless of the pairing of pretraining and test distributions, as the Lipschitzness is the same in all cases, demonstrating that pretraining on tasks with appropriate Lipschitzness is sufficient for generalization.
Moreover, Figure 5(Right) shows that when the Lipschitzness of the pretraining tasks does not match that of the test tasks, ICL. performance degrades sharply, even when the tasks otherwise share similar structure. Here the test task distribution is ((F(os), and the pretraining task distributions are D(F D(For). ). and ((Fi0"). The only pretraining distribution that leads to downstream generalization is D(F(l) since its Lipschitzness matches that of the downstream tasks, despite the fact that is not distribution over cosine functions, unlike the other distributions. Thus, these results lend credence to the idea that in addition being sufficient, pretraining tasks with appropriate Lipschitzness is necessary for generalization.
Softmax Attention Learns Direction of Attention Window
Thus far, we have considered distributions over tasks that treat the value of the input data in all directions within the ambient space as equally relevant to its label. However, practice the ambient dimension of the input data is often much larger than its information content the labels may change very little with many features of the data, meaning that such features are spurious. This generally true of embedded language tokens, whose embedding dimension is typically far larger than the minimum dimension required to store them (logarithmic in the vocabulary size) [1]. Motivated by this, we define notion of "direction-we classes that may depend on some directions within the ambient input data space more than others.
Lipschitzness" of function class to allow for analyzing
Definition 4.1 "Direction-wise Lipschitzness of Function Class). The Lipschitzness of a function class F with domain C RRd in the direction w E gd-1 is defined as as the largest Lipschitz constant of all functions in F over the domain X projected onto w, that is: (F, := inf (L: f(ww" LER
x) - f(ww" T x')
L/w' x x-w x' (x,2') E FJ..
Using this definition, we analyze function classes consisting of linear functions with parameters lying in subspace of Rd, follows:
Definition 4.2 (Low-rank Linear Function Class). The function class Flin defined as F == { f(x) = a TB x, a E R and ((Fii)) is induced by drawing a ~ U
where B Ox is column-wise orthonormal matrix. Since our motivation is settings with low- dimensional structure, can think of k < d. Let B1 (dx(dkk) denote a matrix whose columns

---

Test on Affine, L=1
Test on ReLU, L-1
Test on Cos, L=1
Test on Cos, 1=1
10P
10P
10
109
Trunnad lell 1-1)
Traited 111
Truined (al)
Affine
Africc
Affine
That ICL Emrorr
RelL
PulL
RoLU
Cos
Cos
Chd
10-1
10-1
10-4
10-1
Con 10
10-
10-2
10-1
10-3
250
500
1000
750
1000
750
1000
500 750 1000
Pretraining Iterations
Pretraining Iterations
Pretraliiing Iterations
Pretraining Iterations
Figure 5: Left, Middle-Left, Middle-Right The test error for softmax attention it is trained on the distributions over I-Lipschitz affine, ReLU, and cosine function (D(FT"), D(FF), ). and D(F respectively), where the error is evaluated at each pretraining iteration 5 tasks drawn from the distributions over the I-Lipschitz (affine, ReLU, cosine) function classes in (Left, Middle-Left, Middle--ighh)) respectively. Right: The test error evaluated on tasks drawn from (FF(o)) for three softmax attention trained on tasks drawn from D(F D(F8.). ), and D(F10)) respectively.
form an orthonormal basis for the subspace perpendieular to col(B), and note that the Lipschitzness of Fig in the direction w is T. w E (B) and if w E c.) Observe that any function in FB can be learned by projecting the input onto the non-zero Lipschitzness directions, i.e. col(B), then solving a k < d-dimensional regression. To formally study whether softmax attention recovers col (B), we assume the covariates are generated as follows.
Assumption 4.3 (Covariate Distribution). There are fixed constants Cu + O and -oo < Cy < 0o s.t. sampling x; Dz is equivalent to x;; = Bu, + B. V; where u; Ut and V; ~ uu-
Assumption 4.3 entails that the data is generated by latent variables u; and v, that determine label- relevant and spurious features. This may be interpreted as a continuous analogue of dictionary learning models studied in feature learning works [33, 34]. We require no finite upper bound on |cv] so the data may be dominated by spurious features.
nor [ea|
Theorem 4.4. Let B E Odlk and consider the pretraining population loss (ICL.) with f ~ D(Fin) Suppose Assumption 4.3 holds, as well as at least one of two cases: (Case 1) a 0, or (Case 2) n = 2. Then among all M E M == {M E gdxdd :M= = M .IB MB]| S 1 }. the minimizer of cズ
B
the pretraining population loss (ICL) is M = cBB T for some c E (0, 11
Theorem 4.4 shows that softmax attention can achieve dimensionality reduction during ICL any downstream task that has non- zero Lipschitzness only in col (B) by removing the Lipschhissss features while pretraining on Jiin. Removing the LischLipschitzeess features entails that the nearest neighbor prediction of pretrained softmax attention uses a neighborhood, i.c. attention window, defined strictly by projections of the input onto cB) To our knowledge, this is the first result showing that softmax attention pretrained on ICL. tasks recovers a shared low-dimensional structure among the tasks. Please see Appendix J for empirical results verifying that softmax attention indeed recovers low-dimensional structure, even for tasks consisting of (nonlinear) generalized linear models.
Conclusion
We have presented, to our knowledge, the first results showing that softmax attention learns shared structure among pretraining tasks that facilitates downstream ICL. Moreover, we have provided empirical evidence suggesting that our conclusions about what softmax attention learns during pretraining generalize to function classes beyond those considered in our analysis.
The model we use in this work is an attempt to understand a phenomenon that emerges in LLMs, which is that the output of the model can be 'primed' with some examples provided the context that resembles few-shot learning, even though they are only trained on next token prediction. Establishing a mathematical framework for this remains an interesting question. 2. We consider the output of a single layer of attention. Studying the nature of the solution when this is iterated over multiple trained layers is an interesting future prospect.
Limitations and Future Work. 1.
10

---

6
mcknowledgments
This work was supported in part by NSF Grants 2127697, 2019844, 2107037, and 2112471, ARO Grant W91 1911NN2110222 ONR Grant N00014-19-1-2566, the Machine Learning Lab (MLL) at UT Austin, and the Wireless Networking and Communications Group (WNCG) Industrial Affiliates Program.
References
[1] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neclakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[2] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez. Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[3] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1- 113, 2023.
[4] Sewon Min, Mike Lewis, Luke Zettlemoyer. and Hannaneh Hajishirzi. Metaicl: Learning to learn in context. In Proceedings of the 2022 Conference of the North American Chupter of the Association for Computational Linguistics: Human Language Technologies. Association for Computational Linguistics, 2022.
[5] Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446. 2021.
[6] Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng- Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.082. 2022.
[7] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning, 2023.
[8] 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:2007.03576, 2023.
[9] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927, 2023.
[10] Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:23/2.06522 2023.
[11] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. International Conference on Machine Learning, pages 35151-35174. PMLR, 2023.
[12] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:221 15661, 2022.
[13] Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection, 2023.

---

[14] 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:2310.17086 2023.
[15] Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dim- itris Papailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301.33196, 2023.
[16] Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, pages 19565-19594, PMLR, 2023.
[17] Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint arXiv:2310.083/) 2023.
[18] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint 1xXi:21111.00081 2021.
[19] Xinyi Wang, Wanrong Zhu, and William Yang Wang, Large language models are implicitly topic models: Explaining and finding good demonstrations for in-context learning. arXiv preprint arXiv:2301. 2023.
[20] Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305. 19420, 2023.
[21] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.052499 2023.
[22] Chi Han, Ziqi Wang, Han Zhao, and Heng Ji. In-context learning of large language models explained as kernel regression. arXiv preprint tmXiv:2305.12766 2023.
[23] Yingyi Chen, Qinghua Tao, Francesco Tonin, and Johan AK Suykens. Primal-attention: Self-attention through asymmetric kernel svd in primal representation. arXiv preprint arXiv:2305.19798 2023.
[24] Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philipp Morency, and Ruslan Salakhutdinov. Transformer dissection: unified understanding of transformer's attention via the lens of kernel. arXiv preprint arXiv: 1175, 2019.
[25] Tan Nguyen, Minh Pham, Tam Nguyen, Khai Nguyen, Stanley Osher, and Nhat Ho. Fourier- former: Transformer meets generalized fourier integral theorem. Advances in Neural Informa- tion Processing Systems, 35:29319-29333 2022.
[26] Xing Han, Tongzheng Ren, Tan Minh Nguyen, Khai Nguyen, Joydeep Ghosh, and Nhat Ho. Designing robust transformers using robust kernel density estimation. arXiv preprint arXiv:2210.05794, 2022.
[27] Yichuan Deng, Zhihang Li, and Zhao Song. Attention scheme inspired softmax regression, 2023.
[28] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35: 35:30583-30598, 2022.
[29] Kathryn A Prewitt. A distribution-fre theory of nonparametric regression. laszlo gyorfi, michael kohler, adam krzyzak, and harro walk. Journal of the American Statistical Association, 188464):1084-1084 2003.
[30] Flizbar Nadaraya. On estimating regression. Theory of Probability Its Applications, 11):141-442, 1964.
12

---

[31] Geoffrey S Watson. Smooth regression analysis. Sankhyā: The Indian Journal of Statistics, Series A, pages 359-372, 1964.
[32] Samuele Tosatto, Riad Akrour, and Jan Peters. An upper bound of the bias of nadaraya-watson kernel regression under lipschitz assumptions. Stats, 4(1):1-17, 2021.
[33] Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self-supervised contrastive learning. In International Conference on Machine Learning, pages 11112-11122 PMLR, 2021.
[34] Zhenmei Shi, Junyi Wei, and Yingyu Liang. theoretical analysis on feature learning in neural networks: Emergence from inputs and advantage over fixed features. arXiv preprint arXiv:2206.01717 2022.
[35] Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diver- sity and the emergence of non-bayesian in-context learning for regression. arXiv preprint arXiv:2306. /5063, 2023.
[36] 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 transformers. arXiv preprint arXiv:2309.05858, 2023.
[37] Damai Dai. Yutao Sun, Li Dong, Yaru Hao, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models secretly perform gradient descent as meta optimizers. arXiv preprint arXiv:2212.1055: 2022.
[38] Lingfeng Shen, Aayush Mishra, and Daniel Khashabi. Do pretrained transformers really learn in-context by gradient descent? arXiv preprint arXiv:23 10.0040 2023.
[39] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 11895, 2022.
[40] Asher Trockman and J. Zico Kolter. Mimetic initialization of self-attention layers, 2023.
[41] Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. Joma: De- mystifying multilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00535, 2023.
[42] Enric Boix-Adsera, Etai Littwin, Emmanuel Abbe, Samy Bengio, and Joshua Susskind. Trans- formers learn through gradual rank increase, 2023.
[43] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. arXiv preprint arXiv:2303.04245 2023.
[44] Samy Jelassi, Michael E. Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure, 2022.
[45] Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015. 2023.
[46] Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis. and Samet Oymak. Trans- formers as support vector machines. arXiv preprint arXiv:2308. 16898, 2023.
[47] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism, 2023.
[48] Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. arXiv preprint 20X.0:23 2023.
13

---

[49] Clayton Sanford, Daniel Hsu, and Matus Telgarsky, Representational strengths and limitations of transformers. arXiv preprint arXiv:2306.028%6 2023.
[50] 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 representations. arXiv preprint arXiv: 10. 106/6, 2023.
[51] Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and vari- able creation in self-attention mechanisms. In International Conference on Machine Learning, pages 5793-5831. PMLR, 2022.
[52] Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthyy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.10749, 2022.
[53] Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What single attention layer learn? study through the random features lens. arXiv preprint arXiv:2307.11353, 2023.
[54] Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research, 22(1):3463-3499 2021.
[55] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of gequence-to-sqqueee functions? arXiv preprint 19X. 2019.
[56] Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of trans- formers to recognize formal languages. arXiv preprint arXiv:2009. /264, 2020.
[57] Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self-attention matrices. arXiv preprint arXivv2106.03764 2021.
[58] Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 35:
2022.
2022.
[59] Zhao Song, Guangyi Xu, and Junze Yin. The expressibility of polynomial based attention scheme, 2023.
[60] Kevin Christian Wibisono and Yixin Wang. On the role of unstructured training data in transformers' in-context learning capabilities. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023.
[61] Yichuan Deng, Zhao Song, and Tianyi Zhou. Superiority of softmax: Unveiling the performance edge over linear attention. arXiv preprint arXiv:2310. 1685, 2023.
[62] James Vuckovic, Aristide Baratin, and Remi Tachet des Combes. A mathematical theory of attention, 2020.
[63] Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant of self-attention, 2020.
[64] Herbert Robbins. A remark stirling's formula. The American Mathematical Monthly, 62(1):26-29, 1955.
[65] Brian Knaeble.
Variations on the projective central http://arrivvorr/////00044 http:///00000000..oo//..0o. pdff 2015.
limit
theorem.
[66] Elliott H Lieb and Michael Loss. Analysis, volume 14. American Mathematical Soc., 2001.
[67] G.H. Hardy, J.E. Littlewood, and G. Polya. Inequalities. Cambridge Mathematical Library. Cambridge University Press, 1952.
14

---

Contents
1
Introduction
1.1 Additional Related Work
2
Preliminaries
3
Pretraining Learns Scale of Attention Window
3.1 Proof Sketch
3.2 Experiments
4
Softmax Attention Learns Direction of Attention Window
5
Conclusion
10
6
Acknowledgments
11
Additional Related Work
16
Preliminaries
16
B. Rewriting the Loss
20
C The Signal Term
21
C.I Affine functions
21
C.2 Resed functions
24
D Bounds on Noise Variance
26
Optimizing the Loss
27
E.1 Generalization Bounds
30
Lower Bound for Linear Attention
G Bounds for (r)
31
G.I Bounds on Spherical Caps
32
G.2 Bounds on (r)
H Attention Window Captures Appropriate Directions
36
H.l Proof Sketch
37
H.2 Full Proof
37
I Additional Lemmas
51
J Additional Experiments and Details
52
J.1 Low-Rank Experiments
53
15

---

A Additional Related Work
Empirical study of ICL. Several works have studied ICL linear tasks in the framework introduced by [28], and demonstrated that pretrained transformers mimic the behavior of gradient descent [11-13, 28], Newton's 's method [14], and certain algorithm selection approaches [13, 16]. [35] studied the same linear setting with the goal of understanding the role of pretraining task diversity, while [36] argued via experiments on general auto-regressive tasks that ICL implicitly constructs learning objective and optimizes it within
forward pass. Other empirical works have both directly supported [37] and contradicted [38] the hypothesis that ICL is gradient- based optimization algorithm via experiments on real ICL tasks, while [39] empirically concluded that induction heads with softmax attention are the key mechanism that enables ICL in transformers. Lastly, outside of the context of ICL., [40] noticed that the attention parameter matrices of trained transformers are often close to scaled identities in practice, consistent with our findings on the importance of learning scale to softmax attention training.
Transformer training dynamics. [21] and [41] studied the dynamics of softmax attention trained with gradient descent, but assumed orthonormal input features and either linear tasks [21] or that the softmax normalization is a fixed constant [41]. [42] proved that softmax attention with diagonal weight matrices incrementally learns features during gradient-basse training. Other work has shown that trained transformers can learn topic structure [43], spatial structure [44], visual features [45] and support vectors [46, 47J in specific settings disjoint from ICL.
Expressivity of transformers. Multiple works have shown that transformers with linear [11, 36], ReLU [13, 14, 48], and softmax [12. 15] attention expressive enough to implement general- purpose machine learning algorithms during ICL, including gradient descent. series of works have shown the existence of transformers that recover sparse functions of the input data [49-52]. [53] studied the statistical complexity the learning capabilities of attention with random weights. More broadly, [54-59] have analyzed various aspects of the expressivity of transformers.
Other studies of softmax attention. [60] hypothesized that the role of the softmax in attention is to facilitate mixture-of-cxpee softmax regression problem and analyzed the convergence of a stylized algorithm to solve it. [22] showed that in setting with ICL regression tasks a la [28], kernel regressor akin softmax attention with M equal to the inverse covariance of x converges to the Bayes posterior for a new ICL task - in this setting the conditional distribution of the label given the query and n labelled context samples polynomially with the number of context samples, but did not study what softmax attention learns during pretraining. [61] also compared softmax and linear attention, but focused softmax's greater capacity to separate data from two classes. [62] and [63] investigate the Lipschitz constant of attention rather than what attention learns.
algorithm amenable to unstructured training data. [27] formulated
Non-par con-maramttric regression. Our results imply that pretraining softmax attention reduces to the problem of meta-learning the bandwidth of a Nadaraya- Watson estimator with Gaussian kernel. However, to our knowledge, the non- non-paraan parametric regression literature has not addressed this problem. The closest work is [32], which only upper bounds the noiseless loss, and only in the limit oc, whereas our result characterizes the optimal bandwidth, which requires upper and lower bounds on the noisy loss.
B Preliminaries
We first justify our claim that the first d rows of the last column of Wy can be set to O2 for any optimal choice of parameters.
Lemma B.1. If under the function distribution, function f is equally likely as likely -f, then any optimal solution to L(Wv. W.. W Wa) in 3 satisfies Wv
Oaxd
Oax1
01xd c
16

---

Proof. For readability we write 3 wxql 2x+1 e -mn-1ll
Suppose Wr
(aaxd v) x) was optimal, then the loss can be written
Oixd
(I (f(((()) Σv x; B: 2: f(zx+1))
-= [m,,)))
But because f and -f are equally likely, and because the noise is also symmetric about 0, we can write this
27
= 2 B
Σ (
+ +6)3+ + Σ v
ai Bi f(xn+t))
27
1 + E (.mi)(aaa))
Σ -c((-f(())-)))
Σ v D; Br - (1 (-f)
We can couple the noise {{} and the data {a:} in the two summands above to write this as
[ B C) C
where = =, (2x)) - f(x) - (2, -==))==)))) = E, T a, 3, and = E, ce,
We can set B = 0 simply by setting 0ax1, and this has loss
Σ (f(xx) + B: (
C = B
1 2 (E (-A
(-A-)) = (A+ (B[A4BBC
In all of the distributions over functions consider for pretraining, f is equally likely as fs-f, without loss of generality we set all elements of Wv besides the (d + 1, + 1)-th to O. For simplicity, set the (d + 1, d 1)-th element to 1.
Assumption B.2 (Covariate Distribution). For each token first we draw i as a ~ Ud, Then æ is constructed as æ g/
Definition B.3 (Linear and 2-ReLU Function Classes). The function classes Fhn and Fi are respectively defined as:
Flin (Jw (w( Iw +b, WES W (-(-L,, L]},
F+ = [sw [w((x) = Rel.U( x) + 11 RelU((w b WES
(6)
D(F D(Fi) induced by drawing (O,, and b, I, 1,12 Unif(l- L, LI). We say that these classes are L Lipschitzz because the maximum Lipschitz constant for any function in the class is L.
Note that because Σ
1/2 x;
always, have
zx; M xn+1
= Σ 1/2 0z; 112 | 1/ 1/2 -112
1/2 Ex+1 - 1/2 x; 1/
1/
1/2 Øn+1 11
Let M' = 1/2 1/2 This means the attention estimator can be rewritten as
hsa(#) = Σ (xx)) Mæ,41
Σ [(x,)e -112 || -1--1// m; -ME-112 m 1
Eje0) M
E 1E-1- M'E a
So the attention a token #n+1 places on another æ; is related to the distance between M' 41 and Σ- 11122 x, It is natural to suppose under some symmetry conditions that M' is best chosen be a scaled identity matrix so that the attention actually relates to distance between tokens. Below discus sufficient conditions for this.
17

---

My
Figure 6: Comparison between using M and in Lemma B.5. Here we denote y := Vn+t1- Under the attention induced by M, the center of attention for is actually y', and the attention weights are depicted by the light orange shading. Under the attention induced by w, the center of attention for y and the weights are depicted by the light blue shading. Naturally, using the blue shaded attention should lead a better estimate of fy) under mild regularity conditions.
Assumption B.4. The function class F and distribution D(F) satisfy
1. [(x) [(v)] - L| x y uals Vx, E
2. E ((D)()) [f() p(x 3) Vz X2, for some monotonically increasing p.
3. For any isometry preserving the unit sphere, and f E F, we have f o s F.
Lemma B.S. Under Assumption B.4, any minimizer of Equation ICL satisfees M = WKQ E for some scalar WKQ 2 0.
Proof. Let {u,} = (E 1/2 x;}. Suppose M y * cya+1 for any > 0 for some Mn+11 Take M Vn+1 (the projection of y onto the sphere). Consider a function
Cy.11 M yn- || and k4+1
6pn-1
w: Rd - R&d satisfying (1n+11)
Cun11 Un+1 Note that this need not be linear. Let denote
rotation that sends 4 to Vnti1
We show that L(M)= > C(w). that is, it favorable not rotate ya+1- We have
[
(reemt))
E.f(y,e ll -M 112 27
((M) E (.
Eje e 11u, M v.+1
27
Eis(yee ll M v.
Es (n.)) f(yn+1) E. ESnniri
Eje e 1-11 Mi -M IP
Σ ( (1n+1))(()))) -1u, -M Mn+i
E, Dj M
Lets compare this with the loss of w. For depiction of this, please Figure
( (Fru...))) E(y,)e -w ((w) = B5 11
E, e -11,, -w(w
E, e (2,
[ [En f(y,)e 11 u. -w (w 27
= =E f(yn41)2 Essotrile))
ww Ill
18

---

2 F t.. (y,, (W !1
Σ, e v, -W )1
There three terms to compare. The first in each identical. The second is also the same:
(
Σ, (yy)e 1l 1MMMMM 2o+1 II? Ej e M, 1.-My43 IP?
2
E B5
E My )e-1 V. M 11?
y [
Eg e- 11,,,, -M Mo+i 112
Σ, f(y,)e )e- -MN W 11?
Byn-1 [.v(
Eg e----11. 11313 -OMN Vn-1 113
Σx f(O(y))) (1(())) conts Mats 112
27
-Ey, E.
rotational symmetry of (y.h. vn-1
Eje -110(1,))- comtt Maia
Σ, f(y,))e v. -mnt1 Maee 2
= Ey... [f((
rotational symmetry of {y:}
Eye-lyy 11y, -cunt1 Mee
The third takes some more work. For any choice of (y:}. let
e Mnt1 -v.
Unas (2,) (y.) e- Mn+t v. 112 E, e M/++ 1|
We see that a y (u,) (Iw) varies monotonically with yT V for all yn-1 (y:)- That is,
y. ynt! > V. Un+1 Cy (v.) u Unt (v.)
f(yn4))f( (alg))e-lle ll MM 1I?
E,e ---11 N-MM,. -My++1ll IP
Es ((Un+1)f(()))) --llu vn+1
= =Ey.grit.
ES E -11 Mots
ply Dide 11 -My
= =Eg
Σg e Us -M M.++1 11
ply 2 vr)e -1 un+1 11?
= n Ey B.
v. M MM11 + +E, e- v,
Mn
p(u Untt U Ja)am un 1. 1 (B.))
= n B Py.111... 12.... 11..))
n E Un B. 4 11r, (
P(U T y. y. v' v'(( (v.))
n Ey B .t. Jam
pLu T V. y.)ae Maen (w (e (d-(),))
n EE y A h
plutt2 y. y.)aegmer Vaen ( (l
Similarly, we have
Σ f(yn41)f(y,)) -1 -W )ll
Ef (r
Σ, -- lu, -W 1l1
19

---

By (f(yn+ (f).)] 11 -anta 1o II?
E Ya4
Σj 11,,, -uo+1 Mo+t
E-
Ply yi)e 111:: -wn Mnll 1 1127 Σ
= -
Σj lv, -yn+t Ma+i 1|2
p(y++1 e. -e--l.. 11.. -can11 Mn+1 11 MM Cyntt Doti 11?
=n (..biii--
v. - -uuti Ma+z 11? Σ,
-n.B... By ANdeumn p(yn+1 V.) De Mo-1 (v.)
Critically, for given In+1 a (v.t (y.) can be re-parameterized as
Oyngar(M,) (y.) (v.) (y. Unti) where (u is symmetric about and decreasing. Similarly, plyn+1 V.) can be ee-parameterized p(y V-) p' 1n+1) where a', p' symmetric decreasing rearrangement (that is, the set of points 2 such that p(x) > r a ball about the origin). From Lemma 1.2 we then have
Ey..1 E yar (y:)i=nn-i (1 V.)0 Unt (g.))
= Bytt Ey. {3, (ytii=nn-1 [o'(l yn+1 - 1a (1l Ynt1 -6 y. II)
Evn.1 E By.: (Yeli=(n-1) [p'(1 (11 yo (1) } Untl - y- [D]]
= -E By. (yddi=nn11 ply yd
) (.rennt Vati (.t..vv.(().)
So ((w) < C(M). Let
( (1(bh.))) Ei f(y,) 11 c Mott 112
((yy,. 1)= El
Eje² My -cants Ma+t
Observe that L(w) = Fyy,, ((Cya4z) We might well set to be such that Cy.t, is the same for all Unt1 and minimizer of q, have w(yn+1) = cyn+1 for all Yn+1 which implies cla for some c. Because the optimal M' is identity, the corresponding optimal M is E-
B.1 Rewriting the Loss
As result of this, we can take M = WKQ E and write the attention estimator as
hSA( Σ f(z:)e wxxl E 1/2 m, E 1/2 1
(8)
#a+1 1|? Σ, e uxQ E 1/3 x. -E
This allows us to make the transformation X' + Σ 1-122 X. This has the effect of making both the data covariance and the induced function class covariance equal to the identity. Essentially, WLOG we will henceforth consider Σ = I Henceforth, the estimator will be taken to be
h5A = = f(x;)e "xall 2
(9)
Σ, Exe-oxolee 2.+1
and the loss will be parameterized by WxQ as
(I (f(z;) + c.) exxollz ma+1 1|2
((WKQ)= = E.
f(x
E #a-1
Because the noise is independent of everything else, we can decompose this into two terms, signal term and a noise term as follows
(f(xn+1) f(xx)) uxqll mn 112 2
Σ
[(wKQ)= = Es,(m.))
E, e wx-lkklle 2ut1
Lima (WxQ)
20

---

21
Ge axel xa41
+ {,
ʃ (Bn+1)
ES e wkoll, ant1
We bound the first term in Appendix C and the second in Appendix D. A useful function that we bound in Lemma G.4 and Corrolary G.5 in Appendix Gs
(r)r)
Σ
x
[|
We will use this function, particularly for p 0 and 1.
C The Signal Term
The purpose of this section of the Appendix is obtain upper and lower bounds on Lignal ( Because we work with two different distributions over functions, and because the bounds depend on the distributions, we will make the distribution explicit in the argument to the function
(wigu D(F)) = =B)((=) f((x)) - E wkellll #n+1 11?
exowllll -
2.+1
As reminder, we consider the following two distributions over functions. Please see section B.1 to see why we have set the covariance of be identity.
Definition C.1 (Affine and 2-ReLU Function Classes). The function classes Ft and F are respectively defined as:
J = =f f(x) Iw x +b, WESTH1), F = [f: f(x) = ReLU(w" x)) + 12 ReLU(-w" x) +b, w E St
D(FT), D(F)) are induced by taking w ~ Ud, b,1, Unif[-L,, L.. First we have the following trivial bound (WKQ).
Lemma C.2. For all WKQ we have Esignal (WKQ) < 4L?.
f(00)) ((onn+))) 27 )²] Σn
Proof. We have Esignal (WkQ) S [Σ
for some positive (w). By Lipsehitz-
ness, f(x;) f(xn+1) S L|| x, Fn+11 2L.
C.1 Affine functions
Here consider the affine function class Full. First, that this class satisfies Assumption B.4. Lemma C.3. The affine class Full in Definition 3.2 satisfies Assumption B.4.
Proof.
1. We have (s(=) - [())] = |w (x b)1 < |w|ll x by Cauchy-Schwarz.
2. Because is independent of we have
E22 En /w|22
1
E, [f()) Ew [12 ww y +b²
y +
3. w isotropic, so ((w) is also supported by the distribution on w.
21

---

Take any Lemma C.4. For affine functions, the signal term is upper bounded
0
WRQ
+ 1
WKQ 2 dtvd
Eeigni(WKQQ ((F): 2
2
4L?
WKQ < d+vd
Proof. In the interest of readability, we will denote 2n+1 as D. Consider ä such that a Σ, Ti e 2=xxxq=
Ej ""x0"="="="" Then loss is given by E [ (xr =]] First, since is indepen- T (æ - = T Now has uniformly
dent of æ { we have E1? //W randomly chosen direction, its covariance is multiple of the identity. We have E [r((w-T) L 3. so Elaww] = 3d
E |||
_2,
L. Continuing, E (w (æ -Z)) L7 E x - -l12.
33
x' x. we have
E 2uxQ s a
EE Et
=
x
Ej 2wkqu a.
2 x! æ
E Σ Elx; x' x Σ 2 in 0 iterated expectation and symmetry
Decomposing i into an orthogonal and a parallel component, we have E æ -
E x x x --a x' xx 112 for some x' 1 x with x' = But
x x T1
x T) x' xx = E x - x 2 + E क x x1 22-x(1
E [=1 x e||| E x' x' m
= 2 - x: x))= a' xn" 0
(10)
Case 1: WKQ 2 d+vl
Consider first the term E x01 x T 2)1l E(1 Tx x² Here we have with probability
E(1 x xz)e "xxll x =x;;
1 1- I
Σ e w 20.
(
4+1
Cbn
WKQ
Corollary G.5
(11)
Ch WKQ
20nn
The other term E || x xx 2 B(x'T 2 is the component of the bias in the direction orthogonal to x.
E x 2i -o wk
(T.T a)/²
E, e- -xxxl mi
E, TT Di c -xxlll
E; E-w-e 0i 00
x TT Wi E -xxll 11 m - 1|
2
-wKq 20
Σx (x xx))²) E -2wk @
Popoviciu's 's Variance inequality
E e exxxl m; m. 11
Σ, 2 (1 xr T ax:) -2wka ae -0
E, e e- || 12)2
22

---

Σ, 24 x - 2wxa ==l"
92 (2(2ukQ)
E moog 21 =1²)²
9 (wx))
With probability 1- when wkQ + v we have
$+1
(2(2uKQQ
Cgn
2wka
Cg RQ
(12)
go g0(wkQ}²
Cg 22 1
n
Caff
BKQ
Putting together Equations 11 and 12, we have with probability
1
L
1
W
KQ
Lsigna ((F S O
3d
WKQ
The signal bias is upper bounded by 41,2 always (Lemma C.2). The overall upper-bound on the expectation is
4-1
4 1)
Luiga (wx:: ((FL)) ((F))
I?2
1
WKQ
+2)
3d
WKQ
n
Case 2: WXQ < dtvd We always have ((wKQ) S
4L from Lemma C.2.
Lemma C.5, For affine functions, the signal term is lower bounded as
Ω
L
> dt WKQ
Esignal (WKQ: D(FL)
uko (1)
2
WKQ dtyl
Proof. Similar to Equation (10), for ā = E, -2nkq
we have
1
I
Saganal (WKQi D(FLL)
E x0 x x)11
E(1 x x)
3d
Now consider the term - x x. We have
Σ - x xz
wxq wkk 2, 1|?
92 (2(w)Q)
Σe Se-wkqll -
2go (go(wKQ)
Case 1: WKQ > d+yd Here we have from Corollary G.5, with probability 1 - 1/n
[( x. 2y wwg wexg
Can
Ch 20,
DKQ
Σ e uxa mz
WKQ
20,
WKQ
With probability 1/n the lowest can have is (uigna(WuQ)) so overall we have
2
C
Erignal (WKQ)
24d
Cs WKQ
Case 2: d+vd 4 < wKQ < dtvd From Corollary G.5, with probability
4+1
(( x w.je--xxa -xxll 1|
Cin
C 2uxQ
WKQ
E -xqq x x;
20
2wxQ
20b w wkQ
$+1*
23

---

With probability 1/n 2 the lowest have is (uKq ((FQ: so overall have
2
1,2
Ch e
(xq 2
24d
Cb
4+1 KQ
Case 3: d+yd 4
WKQ. From Corollary G.5, with probability 1 -
E x T æz)e "xqll my 11?
Oone-Auka tWKU
Cb
2mkq,
e e m:
206ne-2wxe
2Cb
With probability 1/n I 2 the lowest we can have is Cuigna(WKQ) ((::::D(F))) D(FT..))
so overall we have
I
Cb
2
Lsigna( D(FT..)
-2wxe
24d
C6
Corollary C.6. Combining the above, we have
1-1
o
Exignal D(FL)
1
w
KQ
+
(13)
(wxq + 1)?
wkq
We can now perturb these bounds in the case ofthe ReL.Ubbaed function class Fī.
C.2 ReLU-based functions
Consider the function class
Fi. Re x) I-ReLU(-WT x))
: ==) S ,,,1,12 L)},
where ReLU(=) = (z)+ max 0). Consider a distributions on Fī. namely D(F)). Let D(FI) be induced by 2tth b, l1, Unif[- LJ. That is, vector w is drawn uniformly on the unit hypersphere, Then two norms are selected, 4,, l2, and the overall function is given by
fw.d (x) ReRLLU( x) lgReLU(-W x) +
so that it follows one affine rule in one halfspace, and another affine rule in the opposite halfspace. Please section B.1 1.1 to see why we have set the covariance of to be identity.
Lemma C.7. The class Ft and distribution D(F defined above satisfy Assumption B.4
Proof.
1. Each function defined as being piece-wise L-Lipschitz, and continuous, is also L-Lipschitz overall.
arccos the points and y are such that (w < y)
2. With probability (that is, they are on opposite sides of the hyperplane defining the two pieces of the ReLU). Because the bias is independent of the other parameters, have as in the proof of Lemma C.3
2 + Eww [12 x ww y/ (w #)(w' 3) 2 0] [][[[(]')]])) x) y) 2 0
E, [f(x)
3
Eww [ ww yl (w a:((w" 4) P( x)( V) < 0)
L
Eew a ww yl x ww y arccos (x y) 1. l 12
Let x = [
for any vector æ. Consider a re-parameterization of the pair (æ. y) as
Eo (ar, y) (2r y. y). Because and y are the unit sphere, this is bijection
+ x
1-8
1
1
E (x, 4) =
2
2 y
2
2
2
24

---

That is, for any a; y. --
y (Ea= (()) (a, y). The push-forward of & is also uniform, that is for æ, y satisfying x y O. (a, y) is distributed 4t xu==1 For any æ, y. let &p (x, y) (ze, ya). Then we have E (x(s) is decreasing function of 0. Finally, for 0 I 0' L z ww Yo L' ww yo zg ww V < Fg. ww Yo O. The product of two positive increasing functions is itself non-increasing.
Since we have both Ew [L x wwT T y x ww 0] and functions of we also have
arrens (x v) are increasing
2 aarcoos( y) F [12 a WW y a ww y > 0l
π
is an increasing function of x T y since Ew I xT ww T y x ww
0
2
and
2 2arcons (m
y)
20.
3. w is distributed uniformly on the hypersphere, so ((w) is also also distributed uniformly on the hypersphere for any isometry @ that preserves the origin.
Lemma C.8. The signal term is upper bounded as
L
WNG + 21)
wKQ 2 4tyd
Esignal (w D(FL))
<
4L2
LtVd < 2 2 WKQ
Proof. We have
Σ, (f(x;) - f(xx)))
wxool xi 10 112
Suigna(W@@@ = Es((mr)
E, E-wux
uxQ mi Dn 1
Ei Tbl xi - Fn Sle 1xxxxxll 20 11a 2
$ (x.)
Σ5 -x-klll 22,
< L 91 (wkk)) go (wkQ)
2
With probability 1 1. when wkQ
dtyd we have
Thn
91 (wxQ) go (o(
"kQ
C
Cp
WKQ
Cen
We always have Erignal (wxQ)= S 4I2 from Lemma C.2. So the overall upper bound is
1
4
(wxq D) <
+
WKQ
n
For WKQ deyd as before, we always have Lsignal (wKQ; D) < 4L2. Lemma C.9. The signal term is lower bounded as
Loigmnt (wxQ; D(FF) Lxigna (wKq; D(Fr)))
For any E Fi let fa denote the corresponding affine function that is equal to f in the halfspace containing æ, that is if f(x') = 1, Rel. x²) + IgReeUU(- æ²) and WL.G x' > O. then (n.))(x)) Note that fafff comes from selected from the unit sphere and b, I-L, L) exactly f
Proof. Again for readability will write En++1
x' +b.
25

---

DFF,), is actually statistically indistinguishable from sample from D(F2"), the distribution over affine functions Definition 3.2 (and the object of Lemma C.5). The error of the nonlinear estimator be written
(
(I f(xi)) - f((n)
E [x,)
ux
5:-1
where 7/
1 1--1 Let us compare the two errors due to the two functions. Let
Σ, wkol =
A = {i: (z: w) T w) < 0} denote the set of points on the opposite side to of the hyperplane defining the function.
Cuigna(wwqq D(FT))
27
= {E,
Σ (1i)): S(xo) -
27
= (m.)
Σ fa,an(a)) + Σ (f(x;) fa.an(x))) 14 - f(xx)
iEA
27
27
Ez,, { E
Σ (a;))();;) Smaa((zr)
+ EE
Σ sami))s)
iEA
YEA
27
[I fueat((02))1
-E,r, -Er,((r:) (wx)) E
Σ (2 farar(()))
+ +EE,
27
E [ (as sm)rr(())
Σ /mmur((00)7x
EE
Σ
(ans)):)
dEA
Weed
2 _f,. Σ (efff(((::))) - faaat(())
+ E fan)
Σ Vied
atr af
27
Emmmnazi)) VeA
(2n))
E (m,.)
2 f(x))))))
Σ
+ fa
PEA
Here the third equality holds because fax;) is independent of fa,))(()) if E A A.
Let (mm) f(:))11)))))) = (x,) (EiE have
Then from the above
27
Esaan - -f(e)) ([rtee))-(aa))
E (ant
2 (Uappal -) + q²,
which has minimum g Lsignal (iggna(WKkr (wKQ; ((F)))/2222 /22 completing the proof.
D Bounds on Noise Variance
In this section we obtain upper and lower bounds on the variance of the estimator due to label noise. There three relevant parameters: d, the ambient dimension of the data; WKQ, the scaling induced
26

---

by the attention layer; and n, the number of tokens. Recall that the noise term is
27
uuiss (WQ)) = (m.))
Σ
Ge wxQ æ;
#0+11 1
E, --eollllz 21 1|2
Because the are independent, this can further be simplified as
-2ww #n+1
( =0²EEx
Σ
wxol my 10
Lemma D.1. The noise term is bounded for d + v uKQ
45vdlogg
o² [1 w xxa)
o² o
S Enoisc(W(Q) SO
Proof. We have
2wxo [x 2nn1
9o (2wkQ)
Lanise L (WKQ) ==xx))=== = Σ
o²
E
e uxxll my - an 12)
9o (w(xxQ)
Using Lemma G.5, we have with probability at least
(o(2wxQ)
wak
Cn.
W
KQQ
1//
go (o(wKQ)2
71
Cn
WKQ
and similarly
(o(2wxq) go(wkQ)2
Cy
12KKQ
On
w
KQ
n
WKQ
Finally, in the worst case, we have 0 < moose (WxQ) S 1.
Finally, we show that the noise term is monotonic WKQ- Lemma D.2. Enoise(w) > Lmoise (w') - > w!
Proof. Let a; = e-w [a, Rn 1211 b, = - -w--ww 11:00 The result follows from Lemma I.3
because {a:} and (bi) satisfy a; aj - b, bj = || Fn+il < [| In+ill
E Optimizing the Loss
For the nonlinear function class F we have the following.
Theorem E.I. Suppose the functions seen in pretraining are drawn from D(FL) in Definition 3.2,
the covariates are drawn as Assumption 3.3, = 52 [.logn) d M satisfies
and n - Ω(1), then the optimal
M wkQld
(14)
where WKQ satisfies
a ((II.2) x s wxq s o (1II.2)
(15)
27

---

upper
oveal uppor bound ovenal lower bound
bias lower bound
nose
Figure 7: Left: Rough upper and lower bounds for the bias term (shaded region), along with the noise variance (gray). Right: Overall upper and lower bound for the in-context loss. The horizontal dashed line establishes an upper bound for the optimal loss, while the vertical dashed lines establish lower and upper bounds for the parameter wKQ that can attain the optimal loss.
Proof. We consider three regions in which the optimal value could potentially lie and see that only the third region is viable.
Case 1. WKQ d + Vd: In this case, the signal term lower bounds the optimal loss by Lemma C.5 as 2(1).
Case 2. WKQ 2
logn In this case, the noise term lower bounds the optimal loss. From Lemma
D.2 we know that the noise term is non-decreasing in WKQ so in the range WkQ > 2
n
is
log
lower bounded by Lnoise (WxQQ at wxQ = 2( log
which is Ω log
Case 3. d vd < wKQ S Ω
By combining Lemmas C.8, C.9, and D.1, we obtain the
log
following overall bound on the loss:
T
0² o KQ
1,2
W KQ
σ² + T?
[(wxQ) se
o²
(wxq + 1)²
WKQ
for some constants e,g e that only depend on d. In the range WKQ 2 + Vd, we have WKQ
which is o² ww
and WKQ < n, so the upper bound can be relaxed as Lmoise (WXQ) < 25
La
n
minimized "kQ
J
Here is upper bounded by 47 dI
JTZ
We note first of all
that for large enough n (as long as 52 log and n प 2())) this is lower than the lower bounds we got in Case 1 and Case 2, so this is indeed the region of global optimal solution. From
Lemma C.9 we have Laoise (WxQ) z
wko + o
La
La wko
which gives
1
o²d dl
Lnoise (WxQ) S 4et.
"
2
nI?
nL?
< <uR
WKQ
d
4e
for the upper bound, we similarly also have Cnoise (WkQ) 2
12
w
kg
which gives
wkq
dd
KQ
47
>
71
28

---

nT.2
242
WKQ
4-
d
Of course, for this to not be vacuous we need
nL
212
3
n
4-d
duta
o²
C
45 Iivgloogn
We will again hide constants that depend only d and write this
nL
T2
C1
S
log
which is true as long as n L looo
For the affine function class FLL we have the following
Theorem E.2. If the functions seen in pretraining drawn from D(FL) in Definition 3.2, and the noise variance o² and Liphscitz constant satisfies 2 [Llog" d42
and nå 2 91(1), and
the covariates are drawn as Assumption 3.3, the optimal M satisfies
M WKQLa
(16)
where WKQ satisfies
Ω ((LL3)
S WKQ O (nL²)
2d++2) ura44)
(17)
Proof. Again we work with three cases.
Case 1. WKQ S d + Vd. Again in this case we have a lower bound to the signal term of 2(11.
Case 2. WKQ 2 logn
Again we have a lower bound of Ω log
Case 3. d vd wKQ S S
Combining Lemmas C.4, C.5, D.1
log n
2
w
KQ
[(wKQ) 5a
L WKQ
0² w
KQ
L
w
RQ
L + 0²
(WKQ + 1)2
n
n
n
We will minimize the upper bound. First suppose bound. Then have
L2 G3
>
WkQ for the wxQ that minimizes the upper
2
o²
w KQ 212
((wKQ) Se
+
u
n
n
This upper bound is minimized at WKQ naiz, However, this contradicts the constraint that L2 WKQ S -r, when n -i2 2 o as we assume. we have wkQ La zz for the minimizer. This means
L2
the upper bound is no more than
L2
2w
² + L
((wKQ) 57
+
KQ
"kQ
n
n
d=4
This upper bound is minimized at WKQ
where it is upper bounded by
odd
d
Lnoise (wkQ) 5452
5122
nI?
nIa
29

---

whenever 2
6z. We that
12
((wkQ) 20
2
L
+ 0²²
KQ
c 2
L
"KQQ
KK
L2
a²d
474
2
c
S
5
nT.
KQ
nL?
1
WKQ
σ²
5c
d
for the upper bound, we similarly also have
((wxQ)= 20
L
+
0²
kK
2 Co
2
KQ
WKQ
2
n
2(d+2) a(d+4)
nL²
WKQ
5 d
c
Of course, for this to not be vacuous we need
2id+
nL2
2td
n 45 log n
5 d-d
o2
We will again hide constants that depend only on d and write this
nL
4
C2
2
log
which again is true as long as n = ( [Llog*
d+2
E.1 Generalization Bounds
We conclude this section with proof of the generalization error new L- Lipschitz task. Theorem E.3. Suppose our attention is first pretrained on tasks drawn from D(F and then tested on an arbitrary L-Lisschita task, then the loss on the new task is upper bounded as L < O L2
A3
Furthermore, if the new task is instead drawn from D(FT)), the loss is lower bounded L 2 min 1/2
A 21 A04/2 }}
Proof. We know from Theorem E.2 that ((A7) & WKQ I O(A3)). The upper bound for L(ukQ).
which is O L
" KRQ) a convex function for 2 so in any range it attains its maximum value a
uxQ
at the extreme points. We can check the cases to see that this is Omaxx L A 45/2
A43/2
L
+
A
A
A
( for large enough n.
R +
Now consider testing on new task from D(F)). The ICL loss for Ω (A9) < wxQ 2 0 is bounded below as Q( A3s and 2 Ad4/2
The implication of this is that if > L, the error scales as (I7)² rather than (7/)445 while for L' & L. the error is lower bounded by constant.
30

---

F Lower Bound for Linear Attention
In this section we prove Theorem 3.6.
Lemma F.1. Consider the function distributions D(FL) and D(Ft) described in Definition 3.2. We have LLA 2 S((L), that is, the ICL error is lower bounded as 92(
Proof. We start by decomposing the ICL loss into a bias dependent term and cenetered term. For f E FL E { F. let f denote the centered function Ex f. Let f denote the flip of f about its expected value, so f' = Ez f F. We observe that f is independent of Eg f. For linear attention, we have, for f D(FL)
(M)) (..a))..
(hLA(En41) - (zn++1))
27
n (f((,))a)a,, Mœn+1) f(Ent1))
27
Σ (f(x)) M Mon+1 + c,a, Mæ + Ey føe M MXn+1) I (1n+11
27
Σ (f(x,) M ao+1 + Gar Mæ I ((ntt)
(18)
[ (E, f) a. M Mæn+1) )
+ Byoeted
27
2 Batalale
Mæ + qæ Mœn+1) - F(xn+1)) - Bx Σ (Tx))
(19)
By symmetry, this is also equal to the same expression using f' instead of f, since f and f are distributed identically. Besides, Ry f Fr I' and is symmetric about the origin, so
CLAMM) 2 E.
(f)))) M Mæ + Ga. M Mxn+1) 8 (3 [ E 8' Σ
27
E Σ ((x)) Mæ cx.. M Mœn+1) (aan)--Faaa T((nt1)) (-(E
Σ (f(())) Mæ c,a0 M Mxn+1) (2( -B
27
Er,,
Let A = En (T( M Ønti + c M xn++1 ((x and B = E, f. Then we see that (M) 2 B(A + B)2 1 (E-A++++ B)² A2 E B2, Meanwhile, E (IE,, f)² just the variance of the signal term in ((F D(FF), which is 4. So CLA(M 2 4
O
G Bounds for (r)
The purpose of this section is to obtain upper and lower bounds on
(p(r) = Σ x x Pe rrll,] - I|?
31

---

for 0, 1/22,1 For this, will need high probability upper and lower bounds the number of points in a spherical cap under a uniform distribution over the hypersphere. Consider n points {x:) drawn uniformly from 0d-1, the uniform measure over Sa-1, the d d-dimensional hypersphere. The
measure of the e- spherical cap around æ E Sa-1, C(e, x) = {x' : x x de.
> 1 - Ef is denoted by
G.1 Bounds on Spherical Caps
Figure 8: The surface area of the purple hemisphere is used to upper bound the surface area of C(²). while the volume of the green hypersphere is used as lower bound. Points in the orange region are Si+1 St, and their count is N.+1 - N₁.
Lemma G.1. The area of the spherical cap C(c). de bounded
(2e 2)
So. S (2c 2) S (2c)
1-1 edd/1
2d
Proof. We derive lower bound follows. We replace the surface area of a spherical cap in S with a d - dimensional ball of the same boundary. Let Va denote the volume d dimensional ball (that is, Vr(r) 5 TT"), and let Aa denote the surface area of dimensional sphere (so As(a) 4mr²). It is known that
π rd, and Ad(r) p
1
Va(r) = r( + 1)
I
Then we have
Va-1 (1 (1 (1-)
Oe 2
Ad(1)
(1 (1- (1-)) I(
2/T
I'( d+1
2 (1 (1 - c))))
Lemma G.6
dd
(2e- 2) d=1
V2dx
The upper bound is similar. This time we replace the cap with the surface of hemisphere with the same boundary. We have
Ad (1 -1-(6 (1-(1-()37) (1 (1 2)
d-1
< (2e 2
T 5
2Aa(1)
32

---

We will also need upper and lower bounds on a discretized version of the incomplete gamma function. Definition G.2. Denote by 7(d, a, m) the expression 7 a: m) E== i
We have the following
Lemma G.3. For > 5, 1 S a S 2, the incomplete Gamma function is bounded as
104e-am-1/2 1/2 < 7dd, a. m) S mm=tle ann amm11/2 [id+1l 20d4r 5 r(d, a, m) x 21(++1)
m d + vd m M d + vd
a
Proof. We compare with the Gamma function
l(d+1)= E'e'd (
Note that Jod ree-ott dt = aeti So de = adti adti [(d + 1). Because the function e is uni-modal with maximum
(ae) we have from Lemma I.1
d
[ cle-ald dl = ad+1 r(d + 1)
ai
+
ai 2 j
ax
Now suppose m dtyd
Then we have
Σ re
c
t= dtv/d
dIVE
d Vd
e (d++V
II 1=
1
d+va
+ J +1
ea
d+Vd
+j
im d=va
i- d+v
8
d Vd
(d+V2)
II j=0
1
d+vd + 1
a
c
d+va
divd
dtva
d + Vd
(d+V2)
av atvl
a
d+va
d + Vd
d
2vd
e (d++V/)
S
1-e-ava/(atV
a
ae
the first inequality follows because dtvd I m. the second follows because
2d+1 2d
>
2d+i+1
the last
2d43
follows because (1+ +
vd
< and I- S 2x for S 2. Over all, we have
Σ m
2 Vd
vdd+1
d d
0
I(d 1) ad+1
d
+1
>
rde
a
ae
j
While for the upper bound we have
d
[(d 1) dd1
<
the
att dt
ae
0
Finally, we use Lemma G.3, specifically that (dd d desired result.
S ब v 2nd (2)a <
[(d+1)
I to yield the
33

---

For m dty/ we have from Lemma G.7 that m
2 aa
ve
Σ
8 e ai
mm
amm-
2 0
and
m
# -ai
S m
d+ -om-
0
G.2 Bounds on 9p(r)
Lemma G.4. Suppose {ari are drawn independently and uniformly from the unit hypersphere. For
log 2 45 n > 5, > 2.p S 2, we have gp(r) =E== x; - x -iee Prll a.
x: satisfies
(1 c
n2f
(:)
++
7 2 2"
P 2, S (r(r) & 3n
(z)
S+4
?
d
+ P 2,r)
r)
TE
r
r
2
2
with probability at least 1
Proof. For 0 S S
let N, denote the number, and S; denote the set, of points satisfying
1 1 S x. x < x; x || I (2
& Also denote by N-1 the points satisfying x. x < 0, and
let S_1 denote this set. Note that
(p(r)
Σ
m x 1e [T
1=0
| 1|-erll1)) = +
Σ
|
[Pe-rla] - 1|2
== jES++1 US
JES-1
2(i + 1)
- (N141 - N:) +
2 N_1
r
Similarly,
2i
M()) 2
2((+1) (Ni+1 No)
Note that because N; > 0.
2(1 11)
2( 22
1)
N41e 2i
(N+1 NN)e -2
1=0
And similarly,
$
2i 1
2i
2i
N N
(N - N2) -
vi= 0
0
1=0
1=0
(Ni+1 - N₂))
Σ
2i
i=)
2i y
(W41--N) Σ
34

---

1-1
2
2j
&
j++) i-j
(N N2)
2
e
2j
vi<j
r
62
j
80 2
2j
4
i-j
(NJ+1 18.-Ny N2))
2j
e
.1+ + SE
(N - Σ
2j
2j
(88-2)
::21
22
1
(N5+1 - -N²)
rr
e
2j
pp<4 4
and so
- e -
2i Σ
Notee 2(1+1) < Σ (N5+1 - N2) ( &
e 2(++)
i=0
By a Chernoff bound for Binomial random variables, we have with probability 1 - ;2:
N₂ no S Snog + v 66 Ginlognog g S 2 & Yr
and
N; no: 2 no:
4n Lologno no; S no:
Whenever
1
16 log
10 2 16llgn Vi
22nd
n
and
N
Over all we have with probability 1 - n
2(i + 1)
h(r) $
Σ
Ni+1e 21 # 2 22N_
,e -2 2r
r
0=
r-1
22(i 1 1)
S n
2e 21
+2
r
1=0
2ne 2
12 2e 2i
221 + 2ee 25
f=1
dip
2
d
2ne
7
+
2) + 2F n1
Definition G.2
2
2'
We always have for p 2
2
4+8
d
P 2.r) 2 2
2no²
7(
+
r
2
2
7
2r P 2
r
2
So at last, we have
2 5+ン
(r)
< 16n
7( 2
p 2, r)
2
35

---

We obtain lower bound in the same way.
2i 2/2
1(r) 2 (1 c 22 -2 Σ
2(1+1]
n
i
6
r
2v 2V2nd 2nd
1=0
n22
2 e e 2)
v V8eATd Td
2 (1 3 2)
n2#
2+2
7
2: 2.r)
p
$49
++4
n2
d P r) - h()) S 3n
2
(1 (1-68 c
BBe Rd
F)
?
2 2
+
7( 2
+
2
r)
with probability
// 2 1
when
N 45 2
2n
log
will be useful to simplify this bound in regimes that we are interested
Corollary G.5. Suppose } drawn independently and uniformly from the unit hypersphere.
n For Tog 2 4r n > 5, p S 2 S d, we have gp(r) Eim1
23 - x P r - = Il" satisfies
with probability 1 2n
(p(r) 0 (p(r) = θ (ne 22)
d
d+vd
r
The following bounds are known for the Gamma function.
Lemma G.6. The Gamma function satisfies
1. V2nd (2)² S <T(d + S ev 2nd (4)
2. r(++ } [(*+1)
2 1
V2+0.3
Proof.
1. Please see [64].
2. Please see [65].
Lemma G.7. The following inequality holds:
(1 1 va)) c-va 26
vil
(20)
Proof. Take the logarithm of both sides, we have that this is equivalent to
Vd 1 2
dlog 1
A Taylor series expansion of log(1 + z) demonstrates that log (1 +1) = -
sva" For
d > 1, these terms are decreasing in absolute value beyond i 2, can upper bound the log
with just the first two terms: log (1 + val
1
2d
H Attention Window Captures Appropriate Directions
In this section we prove Theorem 4.4, which entails showing that if the Lipschitzness of the function class is zero in some directions, one-layer self-attention learns to ignore these directions when the function class consists of linear functions. First, we give brief sketch of the proof.
36

---

H.1 Proof Sketch
We briefly sketch the proof of Theorem 4.4. WLOG we write M BF + B 1 G where F := B M and G B_ B_M M. Lemma H.2 leverages the rotational symmetry of F' in col(B) to show that the loss is minimized over (F, G) at (F, G) = (cB cBB1) for some constants c, c, It remains to show that L(cBB T + c B)] L(cBB whenever d' is nonzero. Intuitively, if the attention estimator incorporates the closeness of B a, and BI #n+1 into its weighting scheme via nonzero Q, this may improperly up- or down-weight fxx,), since projections of n, onto col( do not carry any information about the closeness of (x,) and f(xn-1).
Using this intuition, we show that for any fixed c' and {vi]: such that c'v. Vn+1 7 v Qv for some i, i', the attention estimator improperly up up-weights f(x1), where 1 E arg max; c'v, Vn+1 WLOG. In particular, the version of the pretraining population loss (ICL) with expectation {u:): and {e)) is reduced by reducing c'v, Vn+1- The only way to ensure all (d'v Vn+i); are equal for all instances of {vi)s is to set c' 0, so this d' must be optimal.
To show that reducing c'v] Vn+1 reduces the loss with fixed (vi).. we define a; == e pfvdvvi vata for all i [n]] and show the loss' partial derivative with respect to 0i is positive, i.c.
a Dai
(lchta == (m.)ea... Ei=
27 ))) (a u, - a un+1 + c.) "aa Σ ecch an;
> 0. (21)
This requires a careful symmetry-based argument as the expectation over {us}: cannot be evaluated in closed-form. To overcome this, we fix all u, but u, and one other uy & uz+1 with ay < 01 We show the expectation over (u1, uy) can be written as an integral over (2): E g x g of sum of the derivatives at each of the four assignments of (nj uy) V2). and show that this sum is always positive. Intuitively, any "bad" assignment for which increasing outweighed by the other assignments, which favor smaller 0r1. For example, if y1 un+1 A y2, and u1 y1 and uy = y2, we observe from (21) that increasing @1 can reduce the loss. However, the cumulative increase in the loss on the other three assignments due increasing aj is always greater.
reduces the loss is
H.2 Full Proof
We now prove Theorem 4.4 in full detail. Lemma H.1. For any gk-1 and a On such that min; a; > O, and any cas Cu E R (0). define
Ei ク= (u u) (u, uj ouT utejeuu "azaj
=(c) = E
EE Eo. )2,)
E 22 u
+ 0² (u.
EEL c uaa)
Then for any > 0, 0 & arg minogess [cc..
Proof.
We show that there exists some arbitrarily small > 0 such that J(c) < J(0) by showing
dJ(c) < O. We have
de c==
dc)
de
"dddu,turuun""" "azayas"
(u u) T
(u, - u)(u,u u u.Tu) En clau uag)3
1=1
1 ΣΣ
(u. --u-uch u.u esteea.turtum) "aattuuuaaaaaaaaa (Ean a "a)
+ 20 (u:)ett
i=1 8= ==
Setting = 0 results in
(c)(c)
202 cr EE
71
7
E (m [i u) [u, u)(u, u u.uu)aaaay
dc
1 a
c==0
37

---

20 cx
Y1
12
(u." u u-fu-u,uaaff u,,u)afaaray
Ci ass)
20 cA
(u--u)((us u)(u u- a
EE= a;)3
(22)
202
E (es) [[ uy + 1) (n u - u
i= );
20
R
11
B [u.he [(u uy + u u)! u u
EE 05)
202
12
1= 0:)3
B Beu [u uy + uz uu? u u.Ju)aaayaw)
(23)
20 c4
ag
03)
P=1
[u u. uayay]] n
n
E (usheej
20
E (uchet
[u uu u uf u
EL= a;)3
2c 4
E (ushet [u. u(u u
(24)
E1=1 03)3
20
uuuu O3 ((.))ee (u.Je
EE 12
1= 0:)3
1=1
20
n
0,a, B (ushee [u u.uuu
EE 03))
2g
2cd
a 0v
(25)
k(>in a3)3
24
(1
Σn of
h:
> )2
(26)
where (22) follows since E[u;; 0g. (23) similarly follows since odd moments of uniform random variables on the hypersphere are zero, (24) follows by the i.i.d.nnss of the u,'s, (25) follows since
E|u,u,] L and
1, and (26) follows since min; a; > 0. This completes the proof. L
Lemma H.2. Consider any B dx and resulting function class Flin Consider the training population loss L defined in (ICL), and tasks drawn from D(Ftin)) such that E, [aa LL for some Ca and les M == M k MQ be optimized over the domain Ma {M Rdxd M
M. |B MBB|2 & for any è 0. Then any
M E arg min L(M) MEM
(27)
satisfies M' = BB + B. for some oi Icil E (0, _].
Proof. Without loss of generality (WI.GG), we can decompose M BF+B, B.G.G.Gw where F= B M
and G B M.M. M. Recall that for each E in 11, x, Bu, + c,Biv, v,. Thus, for each i E [m].
38

---

have
ex Mxx+i = ex, BFx xx B1 GX.+1
e Px. eo Gxal1
ec Fxa+ign aq
(28)
where, for each i E [n]. aq == c Gx.+1, For ease of notation, denote Xn+1-
We start by expanding the square and using the linearity of the expectation to re-write the population loss as:
L(M)
Bax haeet
Σ Ei -i= (a Baxxx, x; - ( B x, ---a7 TB x ej)ex Mx+x] Mx
x Ei= ex Mx)2
= E. x, (u.)),((v[vvveee
E E)= (a un a (a)u(aann) u, a uoJueecu Extcou Fx 0,0)
Ch (x;))Fxa;)
Σ 22 Fx n2 (E gauu: Fx
+ 0 2 haw.
(u: u) (u, ww Exteuuy Fx azaj
Ex
2 Ba (u)
(2i=1 enau, px
= Eugera (,xx
ビi=1 2cw u Fx
+ a +
(29)
u Fx az )2
(M, (M,x)
WLOG write Fx R(Fx))||Fx]] for some rotation matrix R(Fx) Okok Denote (Fx)) = Fx|12. Then we have
Exignal (M.
ΣmE E (u: (u: - -u) (uy u) C. RIBX) a
=g qu.)).ve
EE eou (FF)) R(Fx)Ua,)
Σ Di=l Ex=1 (n n) R(Fx) (ny C R(Px)utt,, R(Fx))u a.F
2 (m.).v)
ER (IF)))) a(
(30)
2 (u:).ve)
En - ca C(Fx)u R(Fx)ua,)2
(R(Fx) E, R(Fx) u) (R(Fx) (((Fx)Ju, u, RRFx)
( R(PX)) ((
E Eイ iu: R(Fx) u) (u, R(Fx) T (a "asas
"azas
2 (u))
En (F "
(31)
39

---

where (30) follows since [FFx) I and (31) follows since the distribution of u, the same the distribution of (Fx) for any rotation R (FF)]. Define
E Ey=1(u (u: R(xx) u) (uj - R(xx) (u)es (Fx)u (to, "asaz
g u. (u.).(v:)
C Di=1 (Fx)j "a,)²
for any F E Rxdd We have Signnll (M) = E. [g u, v)], and Note that F' = =BB,, then R Ig and Cj (F'x) Cuc. Thus,
Ci (Fx)
g(F. v. g(
B v)
Cu
E Ej=1 (n R(Fx) 1) (u, - R(Fx) u) weeeaaiipss/ "azay
(v
EE ((F)))))) "a;)²
E EØ=1(1 (u. - u) T u - C Ci a,ay
ea Ci (Fxju, "a,)2
E EØ= (u: u R(Fx) u u u R(Fx) ( uto Ct (Fxjj a,,a
Di=1 e (i())) "a,)²
Σ (u u u R(Fx) u u "a; Σラ=1 (FF)ju "Oj
EL Ci(x)u )
M [ u u, R(xx) ule Ci "au
D (Fx)ju Yog
-= Ci a0s
22( - ((F(
(32)
E (TFe)j)) "a:
u.e C Nas
Define ú = [....
and WLOG write u; Pu; + qu. where Pu; =
Σn euCI "as
uu u; and qu, == (I equal probability, and flipping qu, does not change any exponent or as (32). Thus
uu u, Note that for any u, = Pu, + qu, uu = Pu, qu, occurs with
n
iiml (Pu, qn, C "aa
ń Eqpp, (ve
En e (1(x)u) "aa
1i=1 (2pm: + qu; qu, Jes Ci aa;
10
E (Fx)), n0c
Σ Pu, Pu, en "os
(33)
enn (Fx)) 40;
=ēu
En un Ce C wa; En ( "o.
where à = [(u...vv.) probability, so
Note that for any u, -u; occurs with equal
u G (Fxju uos
1 aa (F()ju, "aa
u. ecu ((Fx), "o
E (v.))
2
Ci,)i(F "a, Ej e (I "aj
40

---

u. Ci(x)u "a,
(iF))u, + Σj=1 el (i()ju Laj
e "a:
u; u [u. (v..
('nCi(Fx)n("aa Lj=1,x e Cx or
e (i + e Ci "as )
Ci(x)() "o:
(34)
Since 04 0 and Gxpv 0 by definition, ((xC)(E aay is monotonically increasing in u Also, f(x) = x_ is monotonically increasing for x >> 0 for all c > 0. Thus we have that
u u > 0
c C2.0i(Fii (Fx)u "a;
O1 "a,
> 0,
(.,CC(FF) on; + Ej=ije (xC)(( "ay (iF)), "os + Ej (..C((P)) "aj
(35)
and thereby ē > O. Therefore, arg max (u') Tú u. in particular u Ju > R(Fx) à whenever [FFx)u u, so (32) is strictly positive if R(Fx)u # u. Thus, for any u, such that R(Fx)u * u, g u.v) > g( Ci (FF) BJ u, v). Also, for any u, v such that R(Fx)u u, g g( v. = g (x
cu
v..)..
Next we need to account for Enoss M, x). Again writing Fx = R(Fx))u|Fx|]] and Cj (Fx) = Fx||22 and using the rotational invariance of u, obtain
20 Fx 0
Enoise (M) = o² 2x hc/ml
Co ox
2e (( R(FX) ((P)))
T =.
- c (Fxju R(
2c (F uo?
(36)
Ci (Fxjuu "a
where (36) follows using the rotational invariance u, So, returning to (29), we have
L(M)= = Ex Erignal
BF +B G, + Emoise (BF B.G. G))
D En (u; - u) T - C. (FC. "a;a)
2 Eu,v
e ( "a;)²
(2 "aa
(37)
(FF))u "a;)²
where (37) is strict if R(F)ju & u for some u, v, which is equivalent to saying that F & T c' 0}.
Next, recall that we have defined o; e Using a similar argument earlier, by the rotational
invariance of the v, 's, for any fixed x, we can write a, = c C= es where Cy (Gx) = ||Gx||2 and e₁ is the first standard basis vector.
Next, for C1, 02 2 0 and some fixed u, v, define
E E (u, T)T(u, uueuu uut.uutuuuuu Cureeuuuugerrr wo
H(u, v,cI, (2) = B
Eh u ucooeve
41

---

Ei=r e " 2e e
+ o² D Buublid
(38)
clueru/ ueerervi e:)a
and let
(C, C3,, E arg
min (.,0))00 2.02
H(,.v,,,,22)
(39)
Since H does not vary with v, have (u, (y, (C)( ()()) WLOG. In fact, H does not vary with u either, due to the rotational invariance of the u,'s So, we have (u, v), (u, v)) (CI, C3) WLOG, i.e. there single pair (Ci, C3) that minimizes c1, C2 for all Sk-1 and E gd-k-1
(u, v. C1, (2)
Thus F = B and G* satisfies |G'x|| = C% for all x, which implies, using also that M is symmetric, that G = C.
Lemma H.3. Consider any a jai,.... anl such that a₁ max, at and a1 > min, ag >> Further, let c E (0, Define
Di= (u: u) InT[u,, (u, u)e utou, uaye uaya
Hsignat u, a)) == Efu
(40)
(E eau! ua,
Then
0H siginall u, a
0.
da
Proof.
We first compute (M.o)(n.o)) 801
Using the lincarity of the expectation and the quotient rule we
obtain:
ö (u, α))
Ooa
② Ei Σj= (u: - Tu,(u, ef utu u aayay
=Equhellll
Ocxi
E poul 4
(E e
eme "a;"² EE - u) Tu, u) u-on aa + Iu u 2 "cy
(uE(u;);
you, "aa)4
R Σn Ej=(u u) (u, uje u utou "aja) (E e "a (Et cru; 4a,)²
-i=1 coul ua)
n
E u) (uj mj utruu "o uaj
[nE(
Ei'= e ",,33
2im1 Σソ
n
v
(u: u) T (uy u) u+eu? (aa
2u
Er mu "a))
Sk.
(41)
j=1
where
(u u.) (u, w)ee(u] u-u; utu u)
Suy = a,a,
EE cu, (ks))3 "xw))
Note that terms with do appear in (41). We analyze S., + S. and each S.j. j & {1, separately, and will ultimately show that each of these terms is positive. We start with the latter case as it easier to handle. For j & (1, we have
(n - u,) (u, u utu u+u. u)
St = aaayEEuuppae...
EE cu aw
42

---

(n u,) uu (u, uje utuf u+u} u)
(Er=1 "a,""
U1 (Is uu (u, u) u)e utu, putfuutuuuuupp utuJu)
+ ajjy heem
(E#=1 e'"/"ap))
u [Ig uu (u, u) ut utujur w)
a;a, (ug)re
(42)
E /u//uaa
(u) u - u. ( u 1) utul utu w)
=0 =0
E em )ay)
where the latter two terms in (42) are zero by the same argument as in (33): flipping the component of either u, or u, perpendicular to u does not change any of the values in any exponent, and each flip occurs with equal probability. Next, note that if a; a1
u uuuu u 1)e u+u u+u3 u)
u u(u I)e u+u? utu u)
(ug Mreil)
(u hrem
1 cu e0,
Er u 4 aay)
thus S.j O. Otherwise, a₁ < an by definition of an, and there must be some such ag since if not, there would be some E 4 such that a = d'.". For the case oj < on, we use a symmetry argument show that Si.j > 0.
First we define additional notations. Let U1 == (u Sewheer hem)) and for any (a, E [-1,1 define
ba-b)(u 1)
fa,(( = (e + c + Ev e )3)3
In particular, for any a E [-1, 1]. define Pa = Pu, [u u a). Since u and u, are i.i.d., we have
Fu,.u, [u, u a, u, u
P Pu..u. Iu u b, u u al PaPs for any (a, b) E [-1, Thus, by the law of total expectation we have
S fas(Urr)PPPo da db [T (aastts
St = Lo.,
0, Eo..,
(fas(Ur,)) + fe.a (01.1) da db
(43)
2
Next we show that for any instance of a, b and U1 (u,((U fo,a (u, (U1,)) is positive. We have:
fa,6(U1,, + fo.a (U1,))
(a bb, u 1)e clabbtu))))))
x
(ecenn ed + Ev a (eeba) e + Ev eu uu
20
with equality only if bor u, since uz u 5 with equality only if uj
u, and
> (6 e Σ
cuu"aay) (e + on ai
Σ
(44)
due to > and for all i'. So we have Sij > 0.
Next analyze S St In these cases cannot immediately drop the components of u and u, that are perpendicular to u. We have:
(u u,) TT)TT( (u, (2u] utu, u)
Sτ1 + St, = 0;00 B
E e ou
43

---

[u u.) u.) (u: u)edu, ut2u! u)
a Bluddee
(1 u u1 u, + u, uj utu, u)
-v== ca
(n 1 11 11 +
(uf u-, 2
a E (ug)rem
-1== e ac)3
(u u uue utu, w) (c "ai + e "a;)
(EH= aa))
(1 (1- u (e u+u," w) (e "on e uas)
+ aulB
Now we can split u u, into the product of the components of u, in the direction u and the product of their components in the perpendicular subspace as before. Doing yields
uu u uf () u+u!
u) (ecuj "an + cou mas)
S + St, = ailtuyyren
e suu"aay)
(1 u uu un
(u utu, u) (e "ai e "aa))
+ a (u
CE c
u, (Ix uu Juu utu, u/ (E (e "al coul uas) u(ss)
-a,B(u,)vet
EE e "a )3
(u. 11 of u)e utu, w (ecu uan + you uas)
Er= em "ay""a"""
(1 u, uu un MeteefuuTute utu, (e "on "ai)
+ a. (n.)ee))
E2 c coulua)3
Next, define
(u u, u utu; u) (e u "an caul u uai)
(U,,( = Eturhre
(1 u uj c( utu u) (cau) uai coul couTTuc) u(s)
Ru hem
(45)
E ب. a.)
We argue similarly as in the previous case, except that here we must include additional terms.
S + St,5
S S (omp(UU..)) (UI.)()))) da db
Bo,
[. Gop(Drr,)apo db
ai Bo..
(46)
2
where
Go,((U1,)) = (U1,) 9
(47)
We show that for any (a, b) E [-1, and any Un Ga)(UI,)) is positive, which implies that Si,1 St, is positive by (46).
First, note that if = a for any [-1, and U1 we have
(1
-a))
Ia.a (U,,) Equgeec [(as + areeca + EnEWNNI a
20
(48)
44

---

since each term inside the expectation is nonnegative, as a² < 1 and o1 > Og. Note that this implies Ga,b 2 when a = b, so WLOG we consider * a for the remainder of the proof. Now we focus on showing (61). Throughout, we will make use of the notation
du.b = c on + e
04
E
(49)
which represents the cube root of the denominator in all terms when u u and u, u b, and
Ja,b == - ab b.
Using this notation, we can rewrite
(UI,(U = e cla-b)) e 1,,00 - e Ta,00'
(la-b)e
(50)
a.b
Therefore.
(a,0(U1,1) + (U,(()
=e((a+b) e00 01
E Ya,00s
claat)) cllatb) e %0
c n2
Lz.
- = =e(atb)) 1-33 b.a
(n
To + e rapdo,))
04 (ea" + t Tapag,o)))
Note that c((++)) d > O, it remains to show that the term inside the parentheses is positive. This term be rearranged as:
an "b E Ta,dd,,)) as (E" To + c (a
= ( an) (e the + E rabda
+ 0z ye e Ta, e yab - Ta.aaa
= =(a1 - as)) nba + e radda, + (dh,o - da.n) (c nb,a - cb %a,b)
(51)
=
First we show that Ti is positive by analyzing %a, and no.a For any (a, E [-1, such that +
ob (70,b) 0b (1 ab + -
a50
(52)
with equality holding if and only if a 1. If a -1, we have Ya.b = 1 + b 1 b- 0 for all b E [1,,1] 1]. Otherwise, (52) shows that 7a b is strictly decreasing with b, so it is minimized over
b 1] at 1. When b 1, we have aa.b 1
0 for all So, 7a.6 > 0 with
cquality holding and only if a = -1 or = Note that by symmetry, this implies Nb 2 0 with
equality holding if and only if 1 or 6 1. So. can have both %.. = 0 and 78.a 0 if and
only a
or = b = However, we have a * b, so at least one of 7a.5 and %6 are strictly positive, and Tj is strictly positive (using also that 01 > as).
We next show that T2 is positive. Observe that
-
>0
+ b > a
(53)
since 0₁ og so it remains to show
b>a + e T cb Ya, b > 0.
(54)
where
Ab c To.b = e00 - ab - a + b) -co(11 ab + a -b).
(55)
We first show the forward direction, namely b 0000 ct 0.
45

---

Note that if a, Vb Ya, 0. So, if show that for any fixed a, 10 Ta, Ta,b is increasing with b as long 2 a, then will have " Ta,6 Mba eaab b is increasing, we take its partial derivative with respect to b:
for a. To show
bb (e. va. = e (1 a a) + e (1
ca - cab)
(56)
We would like to show that the RHS of (56) nonnegative. To do so, we show that its partial derivative with respect to is positive, so achieves minimum value at -1, at which point the value is positive. We have:
e (c-(c-co-1) ca -) (e no Ta.)) c
c cb)
da
06
= = - g(a)
(57)
where (2x) e c11(1 cx - c). Note that ((x) is monotonically increasing in x E [-1, to see this, observe that
0
2 dr <<<)) et ct c (x c 2
(58)
where the inequality follows since E (0, and E [-1, Therefore, since b > a, we have q(b) (a) 2 0 and 2
Be (a (e. - Ta.))) 2 0 from (57). As a result, a
86 (em - e¹ ra.b)
achieves minimum value a -1 At this point, using (56) we have
0 ob
Tha
T0,a T = 2e + th (cb c-
= 2e
This implies that the minimum value of e0 Th, - ja, over b [a, achieved b = a, and we know this value zero, so we have that " odhy ja when
To show the backward direction of (54), namely e " To Ta,6 converse, namely >> b = , Ta < 0, follows by the argument above with and swapped. Therefore, we have T2 > O as desired.
a, note that the
Lemma H.4. Consider any a = [ar,, such that ai > 02 O. Further, let e E (0, 1]. Define
2ou7 ua + 2cu] "aa
H (u, ax) = Eug.u
(ecu, "0i + eou "aa2)²
Then
HHnoise ac)
> 0
dai
Proof. We have
2au ua? + e aaa ueu uai + eouf "a2))²
Hnoise (u, α)) = Eugus
Since n 2, we have
H (u, xx)
ua 2 day y yan (e uaz)
+ 2oou Doufua? ua
= Eu, uua
dai
"on 20 en + "ai ( "aa))
= Eu
yag (ecu enuj 2(cau]ua + u "az)e u 2ou u + 2ou uaa)
4az)4
- Eu
(e an cn "az)d
46

---

cdlui utu w) (e "a1 "a2)
2azBur
(ecu yai + eaus yaz)3
colufa-uu a-uz 4) (e "a c m(a)
Define = uu
and
(e "an t 0023
da, e a1 e a2
ha, == (((+b)) e0
01
a;
13
0.b
Now, we have
N _ [ PaPP da db
3. [ (ha.6 + Papo da db
[ [ (ha. ho.oPapoxtaa # b3 da db +
1 [ [[] (hap haaaPaPbx(a b} da db
[ [ (ha, + hb.aPaPaxtaa + bj da db +
[" Na,aPa da
[ [ (ha. + ho.o)PaPoxta # b) da db +
1 2 S ha da 2 f ho.6pt db
2 [ [ (ha.b + ho.olpapotta # b) da db +
4 { [ ha da db
A. S hoppP da db
2 [
[ (ha, + ho.olPaPoxta + b) da db +
1 L {'
(ha aa + ho,oP() da db
4 -1
| [
Ha, da db
(59)
where
pd
p
Has = PaPPbhaab + hs,a) +
Na.a +
hd.6
(60)
2
2
We will show that for any (a, b) E [-1, and (Pa,p)) E [0, 1|2 Ha, positive, which implies
that N, is positive by (59). To do this, assuming ha is nonnegative for any it is sufficient to show
Hab = ha,s + hs,a + Vha,aho.s > 0,
(61)
since this implies An. a + hr.a v ha and thus, from (60).
Pe
Pb hb.b
Hab > Paalo v Maaahis. +
ha,a
+
2
2
haa
h8,6
Pa
pbl
0
(62)
Before showing (61), need confirm that Na is not negative for all a E [-1, 1]. We have
3ea (a1
az)
ha
2
(63)
d3
47

---

since each term inside the expectation is nonnegative, o1 02- Note that this implies Ha, 2 when a = b, so WLOG we consider a > b for the remainder of the proof.
Note that
@dclatbl (a) - -xi) (a) o2)²
k c3c(o+b) 02)
(64)
(a1 (2) )
Using this, we have
Ha. = ha, + ha,a + / ha,ab
22oteb e2cbtco
e2obtea 0n
20a+e0
02
01 -
a,b
(01 + a2)
d a e datbi (a1 + az)3
3 3
c 01 E a2) a2)2 + (e c a2)dd a2)
=P
+e +e-c(utb) 13 (0.b o2))
(65)
To show that Ha, positive, we need to show that positive. Without loss of generality we can consider aı 1 and a2 10 by dividing the numerator and denominator of Hnoise by a Thus, for the remainder of the proof we treat as and write := a2 for ease of notation. Using this notation we expand P follows:
P = (e - ed add (1 + a)s + (c - a)dd + a). +e =c(att)) di, 1h a) 2n
(e - cm + t a)) (1+ + (e -c a) + (1
aa2²
+ (etab)) + db a) (c a)) a)
5cc-cc + b (0 aa)
+ (etea Addy (-a 50 5a a
(0atth + 30bbcca (1+ 6a + 10a² 10a - 6a - a²)
+ +e 20a+ (1 5a + +Ta 3a -3a -7a 5a a²)
(1 a)) x
(c50a-ob+ + b
+ (EAca + eacy (-a a² 6a a² - a²)
(eacttc debtoaa (1+ 7a 7a 17a 7a 70 +
+ 2ea+ (1+ 6a + +302 360 + 33a + la + a²)
Recall that a > 0, so we need to show that the sum of the remaining terms positive. These terms can be written as polynomial in y == ca(a-b) as follows:
P(1 2 -cc
ya²
+ 1/ (-a - a² 6a a² a²)
+ 12 (1 + 7a +7 + 17 3 7a + 7 + a)
+ (1 6a +3a 2 36a 33a 1 + 6a + a
+ +y² (1 (1+ 7a + 17a3 + T + 7 + a0)
+ (-a - a² 6a at - a²)
+ 3
(66)
48

---

We know that y 36 ... 1 since a b. We also have that 1. Using these facts next show that the sum of the third and smaller-order terms in the RHS of (66) is positive.
(+) = v" (1 + ba + 33a 2 + +6o + 33a 4 + 60 + a²)
+ v² (1 + Ta 7a + 17a 7a + 7 + aaa)
y (-a a² 6a - a - aaa)
+a²
y (1 + 6a + 33a² 360 3 + 33a + 6a + a9)
+ + (1 (1+ Ta 7a 17a + Tax + 7a + a²)
+ (-a--aa - a² 6a a² - aa))
+a²
(2 + 12a + 39a 47a 39a + 12a la")
> 0
Next we show that the sum of the sixth-, fifth-, and fourth-order terms is positive, Let as a²
a5 = - a² + 6a + a* + a and a4 = 1+ 7a + 7a + 17a + 7a + 7a +a", so the sum of the sixth-, fiftth, and fourth-order terms is yaas y + yyaa Note that 32a; < a4 since a 1, and
- 4as = a + a² - -a + a a.
(7.5x + 7.50 15a² 7.5a + 7.50
7.5
(1 (1+ 7a + 7a 17a 7a + 70 + a))
7.5
44
(67)
7.5
thus 05 < 7.5 + 4ac. Also, y claab) - < 2 < 7.5 since S 1. Therefore,
yoas 8d as + 3 04 = 322 (y²a6 as yas + a.)
> 4 yaa 4yas 7.5
a1
+ as
> yd
ya0s - 4yac + a4
7.5
y 7.7.5)
> 3d (v (y'a6 06 4yas + +20g6
7/
1
7.5
62
y 46
y²
y/ + 32
7.5
1
62
> A a6
+ 32
(68)
4
7.5
> 0
(69)
where (68) follows by minimizing the terms inside the parentheses over y. Thus, we have Ho.b > which completes the proof.
Now we can finally prove Theorem 4.4. We prove slightly stronger result, formally stated as follows.
Theorem H.5. Consider any B E Odxk and the corresponding function class Juhn as defined in (4.2).
Suppose tasks are drawn from (Fllin) loss:
B and Assumption 4.3 holds. Recall the pretraining population
(
Ez (i=(((())) - f(xx+1) + ex) Mxa+1
27
C(M) = Ey.nhheepeppeeehheee
(70)
Σn= ex! Mxn+1
49

---

Consider two cases:
Case 1: a = 0, n 1. Then define Cp = 2.
Case 2: 0 > 0, n 2. Then define Cp := 1.
Then in each case, among all M M {M Raxd M M, B MB minimizer M of (70) satisfies M cBB for some E (0, C
Cp ) any
2
Proof. From Lemma H.2, we have M = BB + Emathb EmathbfBBBB. for some and some Cp (0, ) where Cp 2 in Case 1 and Cp = 1 in Case 2. Suppose that è o 0. Then it remains to show that [B, +BB
BB
We start by establishing the same notations as the proof of Lemma H.2. For each [n 1].
x; cu Bu, + C.Bivi V,. Thus, for each i E [n], we have
x Mx ex>x, BB" x B1B
= uf un geve Mnl1
E uf un-ta;
(71)
where, for each i [n)], ag Cev va4s, vn41, For ease of notation, denote Xn+1, u:= Un+11 and = epcu. Also, note that for any x, f(x,) = a " x, = cua u, and that drawing f ~ D
is equivalent to drawing ~ Da for some distribution Da over R such that EawD, [aa""] = e Using this, we have:
C BB BB)
27
Eh (G u, Gua u e)eeet way))
E cau ux;)²
E Ej= (ahee (caa" u, c Tu + ()) u, caa G))l utcuj u a,aa
EE eau u)2
(n: T)TTun u utou, a0j
+ 0² Ei= 2 " "aaajj
G i=
n e uy2 u32
Di= eau uy2
Eu,a [H(u, a]]
where = [0n, (01,..., Onl and
(u,
Li Ei
Ej= (u, - Tu w uueuuuteuuua uaya
Di E 2au! ua²
cg
0²
EE eoul "o;)² c
Dis cau, 10x))
(72)
Define a" = =I,,..., 1,,,,1]] R". We proceed by showing that for any u Sd-11 all a E R satisfy
(i) a = c'a" for some d' E IR,, then H(u, a) = H(u,a a²)
(ii) if a a c'a" for any Ri,, then H(u, a) > au,
This implies BB BB + BB > BB T), since
P Pa(fa = c'a" for some e' RR,) =1 + 2=0,
50

---

which implies that is the unique argument that achieves the minimal value of L BB B over ē IR (and this value is E,, (u, x²) x)))).
Proving (i) is trivial as it can be easily checked that H(u, a:) H(u, d'a) for all u E sd-1, , E ERt and c' R+-
Proving (ii) is more involved. Consider any a & c'a'" for any e' R.,. WLOG arg max; d;- We show that the partial derivative of (u, a)) with respect to a1 is strictly positive, which means that H (u, ax) can be reduced by reducing a by some c 0. We can repeat this argument, repeatedly reducing max; Oi each step and thereby reducing the loss, until we reach an a' satisfying a' c'o". Since loss is reduced at each step, we have that (u, a) > Huu, x*).
To show that the partial derivative of H (u, a) with respect to a₁ is strictly positive, we decompose where
(H(o)) (H.
(Haa((u.a) da:
8m
001
E E/= (n Tu((u, u) "utttn " agjj
Heignal (u, a.) :=) 2f (
(E eou, ua,))
i== 20n "a,
H Hnoise (1, (x) = 2 (u:
n ii=l e mcxx)2
By Lemma H.3, we have @Hiuuluua)) 3a1
> If 0 we are done, otherwise we have n and
01lmaeu,ar by Lemma H.4. This completes the proof. Oag
I Additional Lemmas
Lemma 1.1. Consider continuous unimodal function Then we have
Σ f(i)- max f S 1 f(t)dt < 2=0
f(1) + maxx]
1=1
Proof. Let T denote the point that achieves the maximum of f. Then we know that f() f((+] for
t < T, while f(t) 2 f([t]) for > T. This means f; f()dtt s f(i) $ S S:-1 f(t for t < [T)
and [i, dt) 2 f(1) 2 J f(t)dt for 2 [T] So
TT
Σ
f(i) + Σ f(i)
Dc
f())dt
f(t)dt
[m] [(T]
f(t)dt
f(t)dt
1=0
<
f(t)dt + max f
Similarly we have
8
R Σ f() Σ f(1)
Σ
f(1
1=1M
[2]
i+1
<
(t)dtt
f(t)dt
[
51

---

Σ [(1)dt - s(1)dt
Im
% 1()dt-
Lemma 1.2. If and g nonnegative measurable real functions, then
I [())g(z)dx s ('())" (x)dxr
where g" are the symmetric decreasing rearrangements of f and g.
Proof. Please see [66] or [67].
Lemma L.3. Suppose {a:). {b,} are sorted the same way, 0; > aj b, > by. Then we have
Egjbz
aa))
(Eab,))
Proof. Cross multiplying and expanding, have
arbi)²
<< att)
Σ a,)²
Σ Labba,bjak
atb,a,ax
1 Σ
a.biazbzat + azbzarbrat - abbraibaa <
a bbbjaaa + a3bbaaa; + aabbaaa
i.jk
a aafaa + azbzaza + abbaaa (a.brazbzak + azbzarbaaf + aabba,
Σ a.a,a (a,b) + azbb + bb a,b,be - a,byb; axb;b;)
The last of which follows from the rearrangement inequality [67].
J Additional Experiments and Details
All experiments were run in Google Colab in a CPU runtime. We used a random seed of 0 in all cases, All training was executed in Py Torch with the Adam optimizer. We tuned learning rates in {10-3, 10 2, 10 separately for linear and softmax attention, and we initialized M, and MQ by setting each to 0.001I, and tie the weights of Mæ and MQ to speed up training.
Figure 1. The upper row depicts functions, which increase in Lipschitzness from left to right. The black curve depicts the ground truth, while the gray dots depict the noisy training samples. The shaded region represents the attention window. The middle depicts the attention weights for softmax and linear attention. We remark that the softmax is able to adapt to the Lipschitzness while linear is not. The bottom row depicts the ICL and ReLU pretraining using Linear and Softmax attention. That is, at each iteration, a context is drawn from non-linean regression (defined below) consisting of a randomly phase shifted cosine function. The ICL task is to predict the function value at a randomly chosen query on the unit circle. Each point in the plot depicts the ICL error of pretrained attention unit (using softmax (blue) or linear (orange) activation) at the end of 15000 iterations with learning rate 10-3 We use d = 2 and distribution D(Fr, Here we define
function of the context length for Linear
Fu,bills cos cos (0 b)}
and distribution ((Fr.hill) induced by drawing uniformly from 1- T). We use v = 0, 1.5.6 6 for the left, middle and right plots in the bottom row, respectively.
52

---

Representation Learning Erorr
Affine Tasks
Representation Leamiin Esror
Quadratic Tasks
Represantatin Learning Enar
%
Cosine Tasiss
Linea
Test ICL Croor
Test ICL Error
Error
Pretraining teration
Pretraining Iteration
Pretraining Iteration
Figure 9: Representation learning error (M B)) and test ICL error (mean squared error) during pretraining softmax and linear attention on tasks from Left: Jull Center: F7 .nd Right: JB
Figures 3, 5. In all cases, we use an exponentially decaying learning rate schedule with factor 0.999. In Figures 3 and 5 we use initial learning rate 0.1 and in Figure 4 we use an initial learning rate 0.01. Moreover, in all cases besides those with varying n in Figure 4, we compute gradients with respect to the ICL loss evaluated on N == ly query samples per task (that is, each context input to the attention unit has n + N samples, of which n are labeled, and the other N labels are inferred). When n varies in Figure 4, we use N = 1. In Figure 5 we show smoothed test ICL errors with smoothing rate 0.01.
J.1 Low-Rank Experiments
Due to our results in Section 3 showing that softmax attention can learn an appropriate attention window scale when pretrained on nonlinear tasks, we hypothesize that it can also learn the appropriate directions during pretraining on nonlinear tasks. To test this, we consider tasks drawn from low- rank versions of affine, quadratic and cosine function classes, in particular: J := {f f(x)
B T x +2, E 5k-1), 1], FB := {f f(z) BB x) a a Sk-13 and FB := [f: f(x)
cos 4a B x) a E S Each task distribution D(F D(F)) D(FB is induced by drawing a Uh, We train M and Me with Adam with learning rate tuned separately for softmax and linear attention. We set = 10, k = 2, n 50, and o = 0.01. We draw { i.idd from a non-uniform distribution on Sd-1 for each task, and draw one task per training iteration. We draw B randomly at the start of each trial, and repeat each trial 5 times and plots means and standard deviations over the 5 trials. We capture the extent to which the learned M = M2 recovers col via the metric [B MB lla
PM,, B) == Gmin [B MB) where min (A) is the minimum singular value of A. For test error, we compute the average squared error on 500 random tasks drawn from the same distribution as the (pre)training tasks. Please see Appendix for more details.
We randomly generate B each trial by first sampling each element of B i..d. from the standard normal distribution, then take its QR decomposition to obtain B. To draw the covariates, we draw a random matrix J E ddxd by sampling each element i.i.d. from the standard normal distribution. Then, we compute J = (J 1)//2. Then we draw I, ~ (O.. Ia) and set x; ||
1/2,
Jä.
Results. Figure 9 shows that softmax attention recovers the low-rank structure when tasks are drawn from each of the three function classes, which leads to test error improving with the quality of the learned subspace. In contrast, linear attention does not learn any meaningful structure in these cases.
53

---

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]
Justification: Our abstract is consistent with our introduction. In the introduction, we point to the places in the paper in which we substantiate our claims.
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.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are 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 provide a discussion in Section
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 are encouraged to create a separate "Limitations" section 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 on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should refleet the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness.
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 't acknowledged 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.
the paper. The authors should use their best
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]
54

---

Justification:
Assumptions are specified before all theorem statements.
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
statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: 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 or not)?
Answer: [Yes]
Justification:
We provide details Sections 3.2 and
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 is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or 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 to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS 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) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is 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).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model limited in some way (e.g., registered users), but it should be possible for other researchers 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?
55

---

Answer: [Yes]
Justification: Please see supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/ aablic/guides//
Setiiddes/ooeeSSumiiiiiiioo olicy) for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is 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 NeurIPS code and data submission guidelines (https: wwaicc/gggie//eee
/ //nips.cc..
SoodeSubmissione
for more details.
The authors should provide instructions on 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. If 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 as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide details in Sections 3.2 and J.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core 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, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Please see results in Sections 3.2 and J.
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, statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, 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.
56

---

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 enror 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: Please see Section
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http://nnuuipppp
Answer: [Yes]
Justification: we have read the Code of Ethics and ensured conformity.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations 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: This is primarily analysis of an already existing algorithm.
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 or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
57

---

generate deepfakes for disinformation. On the other hand, it is not needed to point 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility 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]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses no 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 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.
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]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. 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 has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's ss creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
58

---

Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/codc/model 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create 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 to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
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 NeurlPS 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 equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsoureing 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 NeurlPS Code of Ethics and the
and locations, and we expect authors to adhere
guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
59


Reviews:
[{"review_id": "XYJApVHAXa", "reviewer": "Reviewer_3LPG", "summary": "The paper studies in-context learning (ICL) of one-layer attention-only transformers in a regression task. The paper argues that the product of query and key projection matrix is associated with the Liptchitzness of input data. The notion of attention windows is introduced based on this. The paper shows that attention windows will adapt to the Lipchitness and the noise level of data and perform dimensional reduction. Experiments are conducted to support theoretical claims.", "questions": "- The assumption in Equation 4 can be understood as the optimal weights of values $W_V$ will ignore covariates $x$ while the attention will only consider covariates to decide where to attend in the sequence while disregarding predictor $f$. I wonder if this makes sense for hidden states of LLMs. In particular, there is no separation between covariates and predictors in LLMS.  Also, why is the optimal $W_V$ in the proof of Lemma B.1 assumed to have first columns to be zeros (in between Line 577-578)?\n\n\n- I agree that $x^\\top M y$ can be understood as the distance between $x$ and $y$. It might not straightforward to change $x^\\top M y$ into $||x - y||$ like used in Appendix. Can you clarify this more on how to obtain Equation (8) in Appendix from Equation (ICL) in the main text.", "limitations": "Please see above.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper provides a rigorous approach to the problem both with proof and empirical simulations.", "weaknesses": "- Although explaining ICL is important, the setting seems to be less realistic (see question 1).\n- The approach of understanding ICL through the lens of function Lipschitzness is somewhat limiting. For example, Figure 5 demonstrates a decline in generalization error, which is primarily attributed to the function's structure. While the Lipschitz property of a function is indeed a valuable aspect in studying ICL, it doesn't provide a comprehensive explanation or solution for ICL as a whole.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback. Please find the response below, and a global response above.\n\n1. **Closeness of the setting to reality.**\nThe data model we consider in the paper is widely studied (Ahn et al. 2023, Akyürek et al. 2023, Mahankali et al. 2023, Garg et al. 2022, von Oswald et al. 2023, etc.) as a setting that demonstrates a form of in-context learning. In that sense, our work brings a common framework for understanding ICL closer to practice using nonlinear function classes and softmax activations. \nWe address the reviewer's specific points related to the closeness of the setting we study to reality below.\n\n\n    * **Assumption that the optimal $\\mathbf{W}_V$ ignores the covariates.** At a high level, it makes sense that the optimal $\\mathbf{W}_V$ ignores the covariates because the value embeddings should live in label space rather than input (covariate) space. For example, suppose each task $t$ entails translating a word in Language $A_t$ to Language $B_t$. A natural prediction of the translation of the query word is some mixture of the example translated words in Language $B_t$, weighted by the closeness of the corresponding word in Language $A_t$ to the query. Languages $A_t$ and $B_t$ might be very different, so the prediction should not include any mixture of the example words in Language $A_t$ (as the prediction must be in Language $B_t$).\n\n        Returning to our setting, we prove that the optimal $\\mathbf{W}_V$ leads to a prediction that ignores the covariates; this is not an assumption. First, the estimator does not depend on the first $d$ *rows* of $\\mathbf{W}_V$ (not *columns*, as we mistakenly said in line 135), since these only affect the top $d$ entries of the output token which does not affect the loss, which is the standard loss studied by e.g. Ahn et al. 2023, Zhang et al. 2023a, and Mahankali et al. 2023. \n        Second, and more importantly, we prove in  Lemma B.1 that the optimal value of the first $d$ elements in the $(d+1)$-th row of $\\mathbf{W}_V$ is zero. Concretely, suppose $\\textbf{W}_V = \\begin{bmatrix} \\mathbf{A}&\\mathbf{b} \\\\\\\\ \\mathbf{v}^\\top & c \\end{bmatrix}$ and let the attention weight on the $i$-th token be $\\beta_i$ (as in the paper). Then the estimator, which is the $(d+1)$-st coordinate of the output token, is $\\sum_i \\left(cf(\\mathbf{x}_i)\\beta_i + \\mathbf{v}^\\top \\mathbf{x}_i\\right)$. As mentioned above, this does not depend on $\\mathbf{A}$ or $\\mathbf{b}$ (the first $d$ rows of $\\mathbf{W}_V$ only affect the first $d$ rows of the output), and we show in  Lemma B.1 that the optimal value of $\\mathbf{v}$ is $\\mathbf{0}$. Please see the proof of Lemma B.1 for additional technical details.\n  \n    * **Assumption that the optimal  attention parameters $\\mathbf{W}_K^\\top \\mathbf{W}_Q$ ignore the predictor $f$.** Again using the translation task example, it makes sense that the attention weights should only depend on the similarities of the query word in Language $A_t$ with the example inputs in Language $A_t$, since the query lives in a different space than the example labels (which are in Language $B_t$) and thus cannot be compared with them. In the model we use, we follow prior works by setting the last rows of $\\mathbf{W}_K$ and $\\mathbf{W}_Q$ to zero; for example please see Equation (2) in Ahn et al. 2023 and Equation (11) in Mahankali et al. 2023. \n\n\n* **Clarification on how to derive $\\Vert \\mathbf{x}-\\mathbf{y}\\Vert$ from $\\mathbf{x}^\\top \\mathbf{My}$.** We would  like to clarify that we show $e^{\\mathbf{x}^\\top \\mathbf{My}}$ simplifies to $e^{-w||\\mathbf{x}-\\mathbf{y}||^2}$. To do this, we show all optimal $\\mathbf{M}$ must be of the form $\\mathbf{M}=w\\mathbf{I}_d$. Then, we complete the square, and use the fact that all covariates are on the unit sphere to cancel the squared terms in both numerator and denominator, leaving us only with exponents raised to the cross terms. Please see lines 593, 594, Equation (7), and Lemma B.5.\n\n* **Understanding in-context learning via Lipschitzness.** In the model that we consider, we show how a transformer can exploit a shared Lipschitzness in pretraining for ICL. Figure 5 shows that this is necessary and sufficient for generalization: the attention layer fails to generalize to tasks that would be considered as sharing a common structure to those it sees in pretraining. A model that is pretrained on Affine, ReLU and Cosine functions performs equally well for inference on any particular one of those tasks despite not sharing any common structure other than lipshitzness (Right), meanwhile, pretraining on tasks with a different lipshitzness leads to poor test performance. While this is one important aspect, we agree that there could be other aspects of the data and training that contribute to ICL. For instance, in Section 4 we show that attention can adapt to a ``low rank\" structure, essentially picking out specific directions along which to consider distance.", "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": "qVlen5gUnK", "reviewer": "Reviewer_hHKc", "summary": "This paper explores how softmax attention in transformer models enables in-context learning (ICL), where a model can adapt to solve new tasks using only a few input examples without additional training. The authors focus specifically on regression tasks, where the model must predict a continuous value given some input features. They show that during pretraining on a distribution of ICL regression tasks, softmax attention learns to implement a nearest neighbors predictor that is adapted to properties of the pretraining task distribution.\n\n The key insight is that softmax attention learns an \"attention window\" - a neighborhood around each input query point that determines which other input points influence the prediction. The size and shape of this attention window adapts based on properties of the pretraining tasks, specifically their Lipschitzness (how quickly function values can change) and the amount of label noise. Importantly, the authors demonstrate that learning this adapted attention window is crucial for generalization. The authors also prove that softmax attention can learn to project inputs onto a relevant low-dimensional subspace when the pretraining tasks depend only on projections of the inputs onto this subspace. \n\nTo validate their theoretical results, the authors conduct experiments on synthetic regression tasks with varying Lipschitzness, noise levels, and input dimensionality. Empirically, the authors demonstrate that softmax attention learns appropriate attention window scales across a range of nonlinear function classes, including ReLU networks and trigonometric functions.", "questions": "Would there be (not extremely involved) ways of moving towards more realistic settings ? i.e. moving to infinite width settings ? Would classification tasks with a cross-entropy loss be somewhat tractable if making the right Gaussian (concentration) assumptions ?", "limitations": "There are some obvious, if hard to tackle, limitations in this work:\n\na. Simplicity of setting: The analysis focuses on single-layer models and synthetic tasks, which may limit its immediate applicability to more complex real-world scenarios.\n\nb. Gap to large language models: While insightful, the work doesn't fully explain emergent ICL in large language models trained on natural data.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "A first strength is that the paper is not only well written but also excellently presented, which helps in understanding its mathematical content and putting it into a better light. The clarity of exposition is therefore a first great point. In terms of originality, it starts off with the fairly known/commonplace insight that there exists a connection between self-attention and Nadaraya-Watson kernel regression, thus establishing that learning the bandwidth of that estimator across multiple tasks is a necessity for ICL. In this sense and if it stopped there, the contribution wouldn't be particularly novel, as this is intuitive (if thinking of self-attention as learning a summary of the autocovariance function of data and therefore of its characteristic length) and fairly well understood in the literature already (see the cited works of Tosatto et al). One can argue - as in the paper - that Tosatto et al only give an upper bound on the bias, and this work provides a lower bound as well.\n\nHowever where the paper takes off in my opinion is when these arguments move away from purely the 'frequency cutoff / Lipschitzness' length-related realm, to then move into the *directional*, via using concentration arguments on the hypersphere. This is Theorem 4.4 which formalizes the intuition that ICL in transformers identifies low-dimensional subspaces shared by training tasks, an argument more typically found in the analysis of (non-contrastive) self supervised learning. The mathematical method of proof is elegant as well. The authors derive novel concentration inequalities for functionals of points uniformly distributed on high-dimensional spheres, and use a careful symmetry argument to show that any non-zero component in the orthogonal subspace increases the loss. Overall this represents a standout technical contribution well worthy of publication in my view.", "weaknesses": "- A small weakness in presentation is I believe that Theorem 4.4 should be emphasized, as to my knowledge this is the more novel part of the contribution. \n\n- Similarly, the theoretical guarantees are provided for relatively simple function classes (affine and ReLU-based), which may not represent the full range of tasks where ICL is effective.\n\n- The limited scope of experiments is understandable given theoretical assumptions (single-layer) but also important enough that it becomes a weakness, IMHO.\n\n- Finally, the paper could benefit from a more extensive comparison to other theoretical frameworks for understanding ICL.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback. Please find the response below, and a global response above.\n\n1. **Simple function class.** Please see Point 1 of the global rebuttal.\n\n2. **Extensions** Please see Point 3 of the global rebuttal.\n \n3. **Emphasis on Theorem 4.4.** We thank the reviewer for appreciating Theorem 4.4 as a novel, significant and technically impressive contribution. If the paper is accepted, we will allocate some of the additional space provided in the camera-ready version towards further discussing Theorem 4.4's importance in the introduction.\n\n4. **Comparison to other frameworks.**\nWe will add a section in the appendix dedicated to induction heads and the framework comparing ICL to gradient based \"mesa\" optimizers. Please see Appendix A for an extended discussion of related works. We will refine this discussion and include it in the main body given the additional space 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": "ySejPciBd4", "reviewer": "Reviewer_BJxR", "summary": "This paper analyzes how softmax attention learns to perform in-context learning (ICL) through pretraining. The authors show that softmax attention adapts its \"attention window\" based on the Lipschitzness and noise characteristics of the pretraining tasks. They provide theoretical analysis for affine and ReLU-based function classes, demonstrating that softmax attention learns an optimal trade-off between bias and variance. The paper also explores how softmax attention can recover low-dimensional structure in the input space. Experiments are conducted to validate the theoretical findings.", "questions": "- How do you think your analysis might extend to more complex, hierarchical function classes that might better represent the capabilities of large language models? How can it be generally applied when there is no neighbor data?\n- Have you considered how your findings might relate to or explain the emergent capabilities observed in large language models as they scale?", "limitations": "The authors acknowledge some limitations in Section 5, including that their model only considers the output of a single layer of attention and that establishing a mathematical framework for priming effects in LLMs remains an open challenge. These are important limitations to note, as they highlight the gap between this analysis and the full complexity of modern language models.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper addresses an important and timely topic in understanding the mechanisms behind in-context learning in transformer models. This is a crucial area of research given the widespread adoption of large language models. The paper makes a novel connection between softmax attention and nearest neighbor regression, providing an intuitive interpretation of how ICL works in this setting.\n- The analysis is clean and the main message is clear: softmax activation plays a critical role in enabling effective in-context learning. For example, Theorem 3.4 provides concrete bounds on how the attention window scales with task Lipschitzness, noise level, and context size. Also, I like Theorem 3.5, whose test data are any arbitrary L−Lipschitz task, which is pretty general ICL setting. \n- The experimental results generally support the theoretical analysis. Figure 3 nicely illustrates how the spectral norm of M (representing the inverse of the attention window size) increases with the Lipschitz constant L, aligning with the theory.", "weaknesses": "- There's a potential gap between the paper's main message and practical ICL scenarios. The analysis is based on a nearest-neighbor interpretation, but in real-world applications, LLMs often perform well even when there are no \"close\" examples in the context. \n- Some of the experimental writing and presentation could be improved. Figures 4 and 5 are particularly hard to follow without careful reading of the appendix. It would be helpful to have more detailed captions and clearer explanations in the main text.\n- The theoretical analysis is limited to relatively simple function classes (affine and ReLU-based). While this provides valuable insights, it's not clear how well these results generalize to the more complex functions learned by large language models 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 valuable feedback. Please find the response below, and a global response above.\n\n1. **No ``close\" examples, simple function classes.** Please see the global rebuttal (Points 1 and 2). To summarize, it is possible that the intermediate layers of the model can learn the notion of closeness that it needs, but this is outside the scope of our analysis of a single layer. \n \n2. **Writing improvements.** We thank the reviewer for this feedback. Should the paper be accepted, we will make these revisions in the camera-ready version with the extra space provided.\n\n3. **How can our results be extended to hierarchical function classes?** Our analysis in Section 4 begins to answer this question. When the ground-truth functions are a hierarchy consisting of a linear projection followed by a linear link function, the attention weights learn the direction-wise Lipschitzness of the class, that is, they learn this projection. Based on our intuition and experiments in Figure 10 in Appendix J, we suspect that when the functions may have nonlinear link functions (which are still preceded by a shared linear projection), the attention weights again learn the projection and implement a nearest neighbors regressor in the range of the projection with appropriate neighborhood size based on the Lipschitzness of the link functions and noise level. More complex hierarchies likely require additional layers to learn; this is an interesting direction for future work.", "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": "VEOnlWjD61", "reviewer": "Reviewer_ftLf", "summary": "This study examines the impact of softmax attention on ICL regression, advancing beyond the typical linear treatment of the topic. In particular, the authors find that 1) softmax attention leads the model to adapt to the target function's Lipschitz constant and 2) enables the model to recover low dimensional structure in the target.", "questions": "- Could your results transfer in any way to an ICL classification setting, for instance like the one in [Reddy 2024](https://arxiv.org/abs/2312.03002) among others? Your Lipschitz adaptation argument, particularly as it controls a \"window size\" over inputs, sounds obliquely related to the way softmax attention is sometimes discussed as implementing \"selection\" or \"copy\" operations over context examples in ICL classification. I'm curious if there could be a deeper connection here?\n- When you train models in practice, do you find that the QKV matrices of the final model become like the M-parameterization of the matrices you highlight in eq. 4?", "limitations": "The authors adequately address their limitations.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper is very well-written and easy to read. It's great to see work that pushes beyond the usual linear treatment of attention, and explicitly considers the role of the softmax nonlinearity. The connection to kernel smoothing is very neat, and the resulting interpretation of softmax attention as an adaptive process sensitive to the target's Lipschitz constant is very insightful. I was particularly surprised that Lipschitz adaptation is both necessary and sufficient to illustrate the role of softmax attention in ICL regression, and that swapping function classes did not matter so long as the Lipschitz constant remained the same. Very cool work overall. Well done!", "weaknesses": "I found no substantive weaknesses in the analysis, but do have some follow-up question listed below. As with any theory paper, there's always more that can be done, but I think the level of work demonstrated in this paper is more than sufficient to merit a strong accept.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback. Please find the response below, and a global response above.\n\n1. **Do trained models have the same decomposition for $\\mathbf{W}_K^\\top \\mathbf{W}_Q$ as in \n    Equation (4)?** We indeed observe this decomposition  empirically, please see the attached pdf in the global rebuttal for examples of $\\mathbf{W}_K^\\top \\mathbf{W}_Q$ trained on ReLU functions of varying ambient dimensions.\n    \n2. **Classification problems and selection.**\n    Please see our global response for some discussion of how our results may be extended to classification.\n    Regarding the selection idea: One framework that is used to interpret ICL in language is that of the ``induction head\", which allows a transformer to implement an algorithm that searches for a previous occurrence of a token and repeating its following token. Concretely, if the context provided is $\\{... \\texttt{[A]}, \\texttt{[B]}, ... \\texttt{[A']}\\}$, the network outputs $\\texttt{[B']}$ such that $\\texttt{[A]}$ is to $\\texttt{[B]}$ analogously what $\\texttt{[A']}$ is to $\\texttt{[B']}$. While this is outside the scope of our work, one of the intuitions from our work is that the soft-max attention can help inform this analogy in two ways: (1) by informing a notion of \"distance\" that constitutes a significance of match between $[\\texttt{A}]$ and $[\\texttt{A'}]$, and (2) by selecting a subspace within which to calculate attention (in case such structure exists).", "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:
We would like to thank all reviewers for their thoughtful and detailed feedback. Some reviewers raised common concerns and questions; we will address these next here.

1. **Restricted function class (Reviewers BJxR, ftLf).** We expect that our  results will hold for quite  general function classes, in particular,  those that satisfy Assumption B.4. To instantiate specific bounds in the paper we work with specific function classes, which nevertheless comprises the main technical difficulty. Specifically, we have Lemmas C.3 and C.7 that show that our classes satisfy Assumption B.4. The upper bounds are shown in Lemmas C.5 and C.9. It is necessary that the bias of the estimator should increase with an increase in the size of the attention window (to counteract the decrease in variance), and we establish this for these two function classes specifically in Lemmas C.5 and C.9. The lemmas in Appendix C relate the correlation between the function values of a sample of neighbouring points and the bias of the estimator that is built from such a random sample of neighbours, and to derive these correlation we need to work with a specific class.

2. **Interpreting ``closeness" (Reviewers BJxR, ftLf).** Our paper is a study of what a single layer of attention does; a formal exploration of the notion of attending to "close" points similar to a Nadaraya-Watson (NW) kernel (Nadaraya 1964) (which is a consequence of the particular model we study, but also an intuition that many practitioners carry). We hypothesize that with more layers the hidden token embeddings at intermediate layers reflect some learned metric that this type of estimator can exploit for a notion of closeness beyond closeness of the input tokens. The groundwork for this is laid in Section 4 which shows that in the presence of a low rank structure, the kernel projects out the invariant components of the tokens. Suppose for instance that there are two layers. Then the first layer could learn to project onto the relevant subspace by placing all of its attention on tokens in that subspace, and for the second layer, distances are computed only within this subspace.

3. **How can our results be extended to more realistic settings (Reviewer hHKc) and/or settings that capture emergent properties of transformers as they scale (Reviewer BJxR)? (cc. Reviewer 3PLG)**
The regression model for in-context learning that we consider has been studied extensively (Garg et al. 2022, Akyurek et al. 2023, Ahn et al. 2023,  Zhang et al. 2023a, von Oswald et al. 2023a, Raventos et al. 2023, Fu et al. 2023, etc.) as a setting that demonstrates the in-context learning capabilities of transformer models. Our work brings this framework closer to practice due to its consideration of softmax-activated attention and nonlinear and low-dimensional target functions. Below are possible avenues to extend our results closer to practice (though out of scope of the current paper). 

    * *Classification.* The nature of the attention estimate will not change based on the loss function, so it will again be a nearest neighbors estimate in the classification setting. Thus, we expect that the optimal attention weights will again scale with the label noise level and inversely with the Lipschitzness of the labels, for an analogous notion of Lipschitzness defined for the classification setting (e.g. inverse of minimum KL divergence between cluster distributions, or using Fourier coefficients for functions on the hypercube). Technical details need to be resolved regarding analyzing the cross entropy loss, but these should not be prohibitive.
    * *More layers.* Additional layers introduce the following difficulty: the covariates themselves change across layers, so a deeper transformer does not simply implement a kernel estimator with a different kernel. We believe the post-attention multi-layer perceptron (MLP) block and/or layer normalization has to be incorporated into the model and analysis to prevent the tokens from collapsing to a single point similar to what happens in a power iteration (a version of this phenomenon is noted in (Geshkovski et al. 2024)).
    With that being said, it is conceivable that the first $l-1$ layers of an $l$-layer transformer can be interpreted as mapping the input tokens into an appropriate space in which to do nearest-neighbors regression in the $l$-th layer. This is an interesting direction to explore in future work.

    While plausible, these extensions present significant technical challenges that we believe can be topics for separate works. Regarding *infinite width* (suggested by Reviewer hHKc), increasing the inner dimension of $\mathbf{M}_K^\top \mathbf{M}_Q$ will not change the optimal value of $\mathbf{M}=\mathbf{M}_K^\top \mathbf{M}_Q$, so our results will not change.

**References**

Geshkovski, Borjan, et al. "The emergence of clusters in self-attention dynamics." Advances in Neural Information Processing Systems 36 (2024).

Meta-review: The paper analyzes how softmax attention enables in-context learning during pre-training, focusing on regression tasks. The authors show that, during pre-training, softmax attention learns a nearest-neighbor predictor tailored to the pre-training task distribution: the attention window adapts based on the Lipschitzness and the label noise of the pre-training tasks. They demonstrate that attention learns to balance the trade-off between bias and variance, and conditions generalization.

This is primarily a theoretical paper that focuses on a simplified yet widely adopted framework for analyzing in-context learning (ICL). All reviewers agree on the originality and relevance of the contribution to a timely topic. They consider the developed theoretical analysis to provide significant and interesting insights into the understanding of how ICL works, and they all express their appreciation for the contribution.
Decision: Accept (spotlight)