How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?
Authors: Jingfeng Wu
Difan Zou
Zixiang Chen
Vladimir Braverman
Quanquan Gu
Peter Bartlett
Venue: ICLR 2024 spotlight (2024)
Keywords: in-context learning
linear regression
ridge regression
Bayes optimality
Primary area: learning theory

Abstract: Transformers pretrained on diverse tasks exhibit remarkable in-context learning (ICL) capabilities, enabling them to solve unseen tasks solely based on input contexts without adjusting model parameters. In this paper, we study ICL in one of its simplest setups: pretraining a single-layer linear attention model for linear regression with a Gaussian prior. We establish a statistical task complexity bound for the attention model pretraining, showing that effective pretraining only requires a small number of independent tasks. Furthermore, we prove that the pretrained model closely matches the Bayes optimal algorithm, i.e., optimally tuned ridge regression, by achieving nearly Bayes optimal risk on unseen tasks under a fixed context length. These theoretical findings complement prior experimental research and shed light on the statistical foundations of ICL.

Published as a conference paper at ICLR 2024
How MANY PRETRAINING TASKS ARE NEEDED FOR IN-CONTEXT LEARNING OF LINEAR REGRESSION?
Jingfeng Wu
Difan Zou
Zixiang Chen
UC Berkeley noujf@berkeley.eer edu
The University of Hong Kong dzouβcs hku.hk
UCLA chenzxl (eanzxx1@ccs... ucl la edd
Vladimir Braverman Rice University vb21@rice edu
Quanquan Gu UCLA cguecs.ucla.edd ucla.
Peter L. Bartlett
Google DeepMind & UC Berkeley peter@berkeley. edu
ABSTRACT
Transformers pretrained on diverse tasks exhibit remarkable in-context learning (ICL) capabilities, enabling them to solve unseen tasks solely based on input con- texts without adjusting model parameters. In this paper, we study ICL in one of its simplest setups: pretraining a linearly parameterized single-layer linear atten- tion model for lincar regression with Gaussian prior. We establish a statistical task complexity bound for the attention model pretraining, showing that effective pretraining only requires small number of independent tasks. Furthermore, we prove that the pretrained model closely matches the Bayes optimal algorithm, i.e., optimally tuned ridge regression, by achieving nearly Bayes optimal risk on seen tasks under fixed context length. These theoretical findings complement prior experimental research and shed light on the statistical foundations of ICL.
INTRODUCTION
Transformer- baasrormmer--beddddde large language models (Vaswani al. 2017) pretrained with diverse tasks have demonstrated strong ability for in-context learning (ICI.) that is. the pretrained models can answer new queries based on few in-context demonstrations (see Brown et al. 2020 and references there- after). ICI. is one of the key abilities contributing to the success of large language models, which allows pretrained models to solve multiple downstream tasks without updating their model parame- ters. However, the statistical foundation of ICL is still in its infancy.
recent line of research aims to quantify ICL by studying transformers pretrained on the linear regression task with Gaussian prior (Garg et al. 2022, Akyürek et al. 2022 Li et al. 2023b Raventos et al. 2023). Specifically, Garg et al. (2022) ; Akyürek et al. (2022) Li al. (2023b) study the setting where transformers are pretrained in online manner using independent linear regres- sion tasks with the same Gaussian prior. They find that such a pretrained transformer can perform ICL on fresh linear regression tasks. More surprisingly, the average regression error achieved by ICL is nearly Bayes optimal, and closely matches the average regression error achieved by an opti- mally tuned ridge regression given the same amount of context data. Later, Raventos al. (2023) show that the nearly optimal ICL is achievable even if the transformer is pretrained with multiple passes of a limited number of independent linear regression tasks.
On the other hand, connection has been drawn between the forward pass of (multi-layer) Trans- formers and (multi-step) gradient descent (GD) algorithms (Akyürek [ 2022 Von Oswald al.)
2023 Bai et al., 2023] Ahn et al. [2023] Zhang et al. (023a).. offering potential ICL mechanism by simulating GD (which serves as a meta-algorithm that can realize many machine learning algorithms such empirical risk minimization). Specifically, Von Oswald et al. (2023); Akyürek et al. (2022); Bai et al. (202)) show, by construction, that multi multi-layer transformers are sufficiently expressive to implement multi-step GD algorithms. In addition, Ahn et al, (2023) Zhang et al. (2023a) prove that for the ICL of linear regression by single-layer linear attention models, a global minimizer of the (population) pretraining loss can be equivalently viewed as one-step GD with matrix stepsize.

---

Published as a conference paper at ICLR 2024
Our contributions. Motivated by the above two lines of research, in this paper. we consider ICL in the arguably simplest setting: pretraining a (restricted) single-layer linear attention model for linear regression with a Gaussian prior. Our first contribution is statistical task complexity bound for pretraining the attention model (see T'heorem 4.1). Despite that the attention model contains d² free parameters, where d is the dimension of the linear regression task and is assumed be large, our bound suggests that the attention model can be effectively pretrained with a diimension-independeee number of linear regression tasks. Our theory is consistent with the empirical observations made by Raventos et al. (023).
Our second contribution is a thorough theoretical analysis of the ICL performance of the pre- trained model (see Theorem 5.3). We compute the average linear regression error achieved by an optimally pretrained single-layer linear attention model and compare with that achieved by an op- timally tuned ridge regression. When the context length in inference is close to that in pretraining, the pretrained attention model a Bayes optimal predictor, whose error matches that of an optimally tuned ridge regression. However, when the context length in inference significantly differs from that in pretraining, the pretrained single-layer linear attention model might be suboptimal.
Besides, this paper contributes novel techniques for analyzing high-order tensors. Our major tool is an extension of the operator method developed for analyzing 4-th order tensors (i.c.. linear operators on matrices) in linear regression (Bach & Moulines 2013 Dieuleveut et al. 20177: Jain jet al. 2018 2017 You et al. 2021 Wu et al. 2022) and ReL.U regression (Wu et al. (023) to 8-th order tensors (which correspond linear maps operators). We introduce two powerful new tools, namely diagonalization and operator polynomials, to this end (see Section 6 for more discussion). We believe our techniques are of independent interest in analyzing similar problems.
RELATED WORK
Empirical results for ICL for linear regression. As mentioned carlier, our paper is motivated by set of empirical results ICL for linear regression (Garg et al. 2022 Akyürek et al. 2022
Li et al. 2023b; Raventos et al. 2023 Bai et al. 2023). Along this line, the initial work by Garg ct al. (202)) considers ICL for noiseless linear regression, where they find the ICL performance of pretrained transformers is close to ordinary least squares. Later, [Akyürek et al. (2022): Li et al. 2023b) extend their results by considering ICL for linear regression with additive noise. In this case, pretrained transformers perform ICL in Bayes optimal way, matching the performance of optimally tuned ridge regression. Recently, Bai et al. (2023) consider ICL for linear regression with mixed noise levels and demonstrate that pretrained transformers can perform algorithm selection. In all these works, transformers are pretrained by an online algorithm, seeing an independent linear regression task at each optimization step. In contrast, Raventos et al. (2023) pretrain transformers using a multi-pass algorithm over a limited number of linear regression tasks. Quite surprisingly, such pretrained transformers are still able to do ICL nearly Bayes optimally. Our results can be
viewed as theoretical justifications for the empirical findings of Garg et al. 2002); : Akyürek et al. 2022
Li al. 2023b); Raventos et al. (2023).
Attention models simulating GD. Recent works explain the ICL of transformers by their capa-
bility to simulate GD. This idea is formalized by Akyürek et al. (002)); Von Oswald et al. (003);;
Dai et al. (2023), where they show, by construction, that an attention layer is expressive enough to compute one GD step. Based on the above observations, Giannou et al. (2023); Bai et al. (2023)
show transformers can approximate programmable computers as well as general machine learning algorithms. In addition. [Li et al. (202a)) show the closeness between single-layer self-attention and GD on softmax regression under some conditions. Focusing on ICL for linear regression by single- layer linear attention models, Ahhn et al.) (2023); /Zhang et al. (2023a) prove that one global minimizer of the population ICL. loss can be equivalently viewed as step GD with a matrix stepsize. A sim- ilar result specialized to ICL for isotropic linear regression has also appeared in (Mahankali et al.) 2024). Notably, Zhang et al.) (2023a) also consider the optimization of the attention model, but their results require infinite pretraining tasks; and Bai al. (2023) also establish task complexity bounds for pretraining, but their bounds are based on uniform convergence and are therefore crude (see dis- cussions after Theorem 4. In contrast, we conduct fine-gn mee-griined analysis of the task complexity bounds for pretraining single-layer linear attention model with a simplified linear parameterization and obtain much sharper bounds.
2

---

Published as a conference paper at ICLR 2024
Additional ICL theory. In addition to the above works, there are other explanations for ICL. For an incomplete list. Li et al. (2023b) use algorithm stability to show a generalization bound for ICL. Xie et al. 20211: Wang et al. (2023) explain ICL via Bayes inference, [i et al. (2023c) show transformers can learn topic structure, Zhang et al. (2023)) explain ICI. as Bayes model averaging, and Han et al. (2023) connect ICI. kernel regression. These results are not directly comparable ours, as we focus on studying the ICL of a single-layer linear attention model for linear regression.
3 PRELIMINARIES
Linear regression with a Gaussian prior. We will use R& and y I denote the covariate and response for the regression problem. We state our results in the finite-dimensional setting but most of our results are dimension-fre free and they can be extended the case when x belongs possibly infinite dimensional Hilbert space.
Assumption 1 (A fixed size dataset). For fixed mumber of contexts 2 0, a datase of size + 1, denoted by (x, x, v) Raxd x R² x R2 R, is generated as follows:
A task parameter is generated from a Gaussian prior, B ~ (O, Id).
Conditioned on B, (x,9) is generated by x ~ NNO0, H) and y ~ (3 x, d²).
Conditioned on B, each row of (X,y) E (d+(1) is an independent copy of (x .4) E Rd+1.
Here, 262 > o² 2 O, and H 2 0 are fixed but unknown quantities that govern the population data distribution. Without loss of generality, we assume H is strictly positive definite. We will refer to y) x, and y as contexts, covariate, and response, respectively.
restricted single-layer linear attention model. We use / to denote model ICI., which takes a sequence of contexts (of an unspecified length) and covariate inputs and outputs prediction of the response, i.c.,
R x R
R (X,,,)) -=
+
We will consider a (restricted version of a) single-lyeer linear attention model, which is closely related to one-step gradient descent (GD) with matrix stepsizes as model parameters. Specifically, that the function class
based the results of Ahn al. (2023); Zhang al. (2023a), one of single-layer linear attention models (when some parameters are fixed to be zero) is equivalent to the function class of one-step GD with matrix stepsizes model parameters (see Appendix B for a proof). Therefore, we will take the latter form for simplicity and consider an ICL model parameterized as one-step GD with matrix stepsize, that is,
f(X,y,; = TX y x). r T
Rddd Rdxd
(1)
din(yy)
where I' is a de-dimensional matrix parameter to optimized, and dim(y) is the dimension of y. That is, we consider two simplifications of the usual soft-max self-attention model: we remove the nonlinearity and we replace the usual parametrization with simpler linear one (see Appendix B).
ICL risk. For model (I) with fixed parameter r, measure its ICL risk by its average regres- sion risk on an independent dataset. Specifically, for n 2 O, the ICL risk evaluated on dataset of size n is defined by
(T) B(J( y, u)² where the expectation is taken with respect to the dataset (X, y, x) generated according to As- sumption I with n contexts.
We will set N generate datasets for pretraining and M to generate datasets for inference, where M is allowed to be different from N.
3

---

Published as a conference paper at ICLR 2024
We have the following theorem characterizing useful facts about the ICL risk Special cases of
Theorem 3.1 when x² 0 have appeared in (Ahn et al., 2023 Zhang et al. 2023a). The proof
is deferred to Appendix C For two matrices A and B of the same shape, we define (A :=
tr B).
Theorem 3.1 (ICL risk). Fix N as the number of contexts for generating dataset according The following holds for the ICL risk Rx(-) defined in (2)
to Assumption 7
The minimizer of (.) unique and given by
FN = (r
0²
///
N +
-1 H
(3)
N
N
2. The minimum ICL risk is given by
tr(H) + 02 02/102
11
min - =² & Hi
3. The excess ICL risk, denoted by A (-)), is given by
Ax(T)= = Rr (T) min R, (T) H, (T HN) (TN) T)
where
Hy = E ( ( T
N
= D
I+
(4)
N
N
For simplicity, we may drop the subscript N in Nÿ and Hy without causing ambiguity.
When the size of the dataset N oo, have r - H according to Theorem 3.1 Then for a fresh regression problem with task parameter B, the attention model (T. after seeing prompt (X, y, x) of infinite length, will perform Newton step the context (X,y) y) and then the result to make linear prediction for covariate x. Since the context length is infinite, the output of a Newton step precisely recovers the task parameter B, which minimizes the prediction error. Thus the attention model (I), with a fixed parameter T achieves consistent in-context learning (Zhang al. 2023ab. When N is finite, (3) is regularized Hessian inverse, so ( performs regularized Newton step in-context - the regression risk of this algorithm will be discussed in depth later in
Section
Theorem suggests that the ICI. risk parameterized by T is convex and the optimal parameter is unique. However, since the population distribution of the dataset unknown (because s!?, o², and H are unknown) and the parameter (a x d matrix) is high-dimensional, it is not immediately clear how many independent tasks are needed to learn the optimal parameter. We will address this issue in the next section.
4 THE TASK COMPLEXITY OF PRETRAINING AN ATTENTION MODEL
Pretraining dataset. During the pretraining stage, we are provided with pretraining dataset that consists of N + independent data from each of the T independent regression tasks. Specifically, the pretraining dataset is given by
(5) where each tuple (t. Ui) is independently generated according to Assumption I with N being the number of contexts, We assume is fixed during pretraining to simplify the analysis.
Nxd
E.ERN, X ER 31 E
t= 1,
Pretraining rule. Based the pretraining dataset (5). pretrain the matrix parameter Fr by stochastic gradient descent. That is, from an initialization To. e.g., To 0. we iteratively generate (T.) by
H P--1- Tt (I 2
v1) t = 1, 1,...,TT. T,
(6)
4

---

Published as a conference paper at ICLR 2024
where (1, X4, (4) is the pretraining dataset (5). is the attention model ① and (y)) is geometrically decaying stepsize schedule (Ge et al. 2019 Wu et al. 2022 i.c.,
2
2
l = |t/ log(T)
t 1, T.
Here, 7o > is an initial stepsize that is a hyperparameter. The output of SGD is the last iterate, i.e. Tr.
Our main result in this section is the following ICL. risk bound achieved by pretraining with T independent tasks. The proof is deferred Appendix [
Theorem 4.1 (Task complexity for pretraining). Fix N 2 O. Let PT be generated by (6) with
pretraining dataset ( and stepsize schedule 7 Suppose that the initialization To commutes with where is an absolute constant and Hy is defined in (4) in
H and to S
Theorem 3 Then we have
II HHN) (To (To-ri -Txx) ) 1=1
(T
HHN,
Datt
+
+ tr((H) + o² + (HHN, (To Po--ii) r2)3)) Tott
(To-\%)3)
where the effective number of tasks and effective dimension are given by
Tatt T
Datt EΣ min (1, T T2
(8)
log(T)'
respectively, and (Ni) 121 and (Xi) j21
are the eigenvalues of H and Hy that satisfy
tr( o²
02/112
N 1=)) j21.
= >B>
N
N
In particular, when To- 0, we have
[I- HH 1i )² (r + o²) t=1
Dost
(T) HHN,
Tott
Theorem 4.1 provides statistical ICL risk bound for pretraining with T tasks, which suggests that the optimal matrix parameter D (see (3)) can be recovered by SGD pretraining if T is large enough. Focusing on (9) Theorem [.1 the first term the error of directly running gradient descent on the population ICL risk (see Theorem 3.1), which decreases at exponential rate. However, seeing only finite pretraining tasks, the population ICL risk directly minimizable by the pretraining rule, and the second term in (9) accounts for the variance caused by pretraining with data from T independent tasks rather than an infinite number. The second term is small when the effective dimension small compared the effective number of tasks (see their definitions in (8)) We remark that the initial stepsize "?o induces trade-off between the two terms, where larger initial stepsize reduces the first term but increases the second term and vice versa.
We highlight that the bounds in Theorem 4.1 do not explicitly depend on the ambient dimensionality d², allowing efficient pretraining even with large number of model parameters. Specifically, our bounds (e.g., (9)) are functions of the effective dimension (8. In the worst case, for example, when
H ==I I and T' is larger, we have Deer = d² so that the excess risk bound is OdP However, the effective dimension Dest is always no larger, and can even be much smaller, than d² depending on the spectrum of the data covariance. In contrast, the pretraining bound in (Bai et al.) (2023) is based on uniform convergence analysis (see their Theorem 21) and explicitly depends on the number of model parameters, hence is worse than ours.
To further demonstrate the power of our pretraining bounds, we present three examples in the fol- lowing corollary, which illustrate how pretraining with limited tasks minimizes ICL risk. The proof deferred to Appendix D.9
additionally assume that To- 0, 2 = 1, 2 1, tr(H)= 1, and choose stepsize "/o 1/ tr H) HH)) =1/tr(H).
Corollary 4.2 (Large stepsize). Under the setup of Theorem 4.1
5

---

Published as a conference paper at ICLR 2024
1. The uniform spectrum. X 1/s sis and X, 0 for iori> >> where and N satisfy N S s S d, then
N
Tor
8 +
Tott S s².
s
EAN
(
Tort > 8²
Tatt
2. The polynomial spectrum. If - for a > 1 and N3 = o(Tott); then
EA S T -1 N log(Tarr) + T N2
3. The exponential spectrum. If = 2 and No S o(Tort), then
N + log (Terr)
EA(TT)
Test
To summarize this section, we show that the single-layer linear attention model can be effectively pretrained with a small number of independent tasks. We note that our statistical task complexity results are under the one-step GD parameterization, where we have convex (but high-dimensional) learning problem. Under the orginal attention parameterization (see Appendix [B), the learning prob- lem is non-convex, which adds extra layer of complexity from non-convex optimization. We leave for future work extending our statistical task complexity results to the original attention parameteri- zation. Finally, also empirically verify former in Appendix [ Nevertheless, it is still unclear whether or not the pretrained model achieves good ICL performance. This will be our focus in the next section.
theory both numerically and with three-layer trans-
THE IN-CONTEXT LEARNING OF THE PRETRAINED ATTENTION MODEL
In this section, we examine the ICL performance of a pretrained single-layer linear attention model. We have already shown the model can be efficiently pretrained. So in this part, will focus the model (1 equipped with the optimal parameter (T x in (3). to simplify our discussions. Our results in this section can be extended to imperfectly pretrained parameters (T') by applying an additional triangle inequality. All proofs for results in this section can be found in Appendix E
The attention estimator: According to I and the optimally pretrained attention model cor- responds to the following linear estimator:
tr (H(H) + +02/6² N + H
X dim
(X,,,
(10)
N
Average regression risk. Given task-specific dataset (X, let g(X, y, x) be an estimator of y. We measure the average linear regression risk of g by
v, x, v) generated by Assumption
1
[ == ([l([[([,,, - v) x
(11)
where the expectation is taken with respect to y, x, and y, and is conditioned on X.
The Bayes optimal estimator. is well known that the optimal estimator for linear regression with Gaussian prior is optimally tuned ridge regression estimator (see for example Bishop & Nasrabadi 2006] Section 3.3). This is formally justified by the following proposition.
Proposition 5.1 (Optimally tuned ridge regression). Given a task-specific dataset (X,y, x,y) gen- erated by Assumption [ the following estimator minimizes the average risk 11
h h(X, = ((X X
X y, x) o² dim (y)
1 dim(y)
02/1²2
X din(y)'
(12)
x x
+
x

---

we have Published as a conference paper at ICLR 2024
It is clear that the optimal estimator
12 corresponds to ridge regression estimator with regular-
ization parameter o² /v/2/ /dim(y).
Based on the analysis of ridge regression in (Tsigler & Bartlett 2023). can obtain the following bound on the average regression risk induced by the optimally tuned ridge regression.
Corollary 5.2 (Average risk of ridge regression, corollary of (Tsigler & Bartlett (023)). Con- sider the average risk defined in (TT). Assume that the signal-to to-iae ratio upper bounded, i.e., (H) S o². Then for the optimally tuned ridge regression estimator 12 least 1 e (MMM over the randomness in X, it holds that
with probability at
o² 02/4²2
o2 22 _ min {uM, de), where μM
M
where M = dim(y) refers to the number of independent data in yX
We remark that the attention estimator (10 is not the Bayes optimal estimator 12 However, will show that the average risk induced by the attention estimator (1O) can be close to that of the Bayes optimal estimator (12) in suitable regimes. In this way, we can view the attention estimator (0) as good "statistical shortcut" to the Bayes optimal estimator (12). thus achiving good ICL
performance.
Based on Theorem we have the following bounds on the average risk for the attention model. Theorem 5.3 (Average risk of the pretrained attention model). Consider the average risk defined in (TT). Assume that the signal-to-nois ratio is upper bounded, i.e., trHr(H) S o². Then for the attention estimator [T0, have
A PN
A; PM
EC(f:X) - 0² ~ 2² Σ min {AM, A1) + &² &H - HNN)² Σ min
min
where μM = ()) and un =o2/(b2N).
Theorem 5.3 provides an average risk bound for the optimally pretrained attention model. The first term in the bound in Theorem [.3 matches the bound in Corollary [5.2] When the context length in pretraining and inference is close, i.e., when M ~ N, the second term in the bound is higher-order, so the average risk bound of the attention model matches that of the optimally tuned ridge regression. In this case, the pretrained attention model achieves optimal ICL.
When M and N are not close, the attention model induces larger average risk compared ridge regression. We provide the following three examples to illustrate the gap in their performance. Corollary 5.4 (Examples). Under the setups of Corollary 5.2 Iad Theorem
5.3 additionally assume
that 0² ~ 1, 2 ~ 1, tr(H)- ~ 1, and M < ///cfoo for some constant >
1. The uniform spectrum. When X, 1/s for S and X; 0 i > have
C(1; - o² = min
s
with probability at least - e -Q(M);
M
EC(;:X) 02 ~ min
s
if < Mors> N2/M.
M
The polynomial spectrum. When X; i
for a > 1, we have
C(h; 0² = M - -1
with probability at least e eR(M);
E 8² = -
3. The exponential spectrum. When \; 2
C(b; X) 0² log M
with probability least -R(M))
M
EC(F:X)= o² log N
M
To conclude this section, we show that the pretrained model attains Bayes optimal ICL when the inference context length is close the pretraining context length. However, when the context length is very different in pretraining and in inference, the ICL of the pretrained single-layer linear attention might be suboptimal.
7

---

Published as a conference paper at ICLR 2024
6 TECHNIQUE OVERVIEW
In this section, we explain the proof of Theorem 4.1 Our techniques are motivated by the operator method developed for analyzing 4-th order tensors (i.e., linear operators matrices) arising in linear regression (Bach & Moulines 2013] Dieuleveut et al. 2017 Jain et al. 2018 2017 Zou et al., [021] need deal with 8-th order tensors that require two new tools, namely, diagonalization and operator polynomials, which will be discussed later in this section. For simplicity, we write IN and Hy as I and H. respectively.
Wu et al. 2022) and ReLU regression (Wu et al. 2023). However,
We start with evaluating (6) and get
P2 = P - xxx, -Tr) (xxTx.)
(Nxxxx)
- 7E, t= 1=....., T,
where Ee is mean random matrix given by
1 x Yr
) (xXxxx
x x1)
yext (xxT)x)²
E = x,X, r
N
Define sequence of (random) linear maps on matrices,
(+xxTx)) 1stsT. P.o : A Xx (xTx1
YA gdxd
Then we can re-write the recursion as
- P: (T+-1 - 1=
The (random) linear recursion allows us to track Fr, which serves as the basis of the operator method. From on, will heavily use tensor notations. We refer the readers to Appendix D.1 for a brief overview of tensors (especially PSD operators).
Bias-variance decomposition.
Solving the recursion of T, yields
T r
II 9 (To -
r
7 II
o
k=t+1
Taking outer product and expectation, we have
T II 1 (To
T Σ
T
82
AT = E(T -
-r")
II
P/ o
E
t=1
1=1
k=t+1
T
82
T
T
02
22
II
P o (T r
+2
Σ
7/
II
91 o
B
k=t+1
Br
Cr
where Ar. By, and Cr all PSD operators matrices (i.e., 4-th order tensors). Then we decompose the ICL risk (see Thcorem 3.1 into a bias error and a variance error:
EAN(Tr) == (H, Ar H) s 2(H, By H) + 2(H, Cr H).
In what follows, we focus on explaining the analysis of the variance error (H, Cr o H).
Operator recursion. The variance operator Cr can be equivalently defined through the following operator recursion (see Appendix D.2 for more details):
Co 0, C. = S o C.-1 + ?,N, t= 1,.....,7
(13)
where N = 62
and St is linear map operators (i.e., an 8-th order tensor) given by: for
any O E dRdxd 82
f100== :- -- - 7t (H + 0 + ..M.0...
8

---

Published as a conference paper at ICLR 2024
with M, & being given by
M : [xx]]( <:=B <<=((>xT))) (
y )))
82
Appendix D.3 includes several bounds about these operators; among them the following is crucial: for any PSD operator O, Mc O H) where s(1) (H. H
where e > is an absolute constant.
Key idea 1: diagonalization. The operator recursion 13 involves 8-th order tensors Pe that hard to compute. A critical observation is that the variance bound only depends on the results of Cr applied on diagonal matrices (assuming that H diagonal, which can be made without loss of generality). More importantly, when restricting the relevant operators to diagonal matrices (instead of all matrices), the 8-th order tensors S can be bounded by simpler 8-th order tensors F, plus diagonal operators. Specifically, based 133) we can show that (see Appendix (.4)
é, I 9 o c.11 on (H. C- o H) ((
tr(H) + + o²)80)
(14)
where C. refers restricted to diagonal matrices and 9, is linear map on operators given by:
0 = 7i (H & 11) 1)00 HH & I) (I
H) O (I ()) 7 H O H
We remark that Wu et al. (2023) has used the diagonalization idea with matrices for dealing with non-commutable matrices. In comparison, here we use the diagonalization idea with operators for dealing with high-order tensors.
Key idea 2: operator polynomials. To solve the operator recursion in need to know how the 8-th order tensor G. interacts with operator (11), To this end, we introduce powerful tool called operator polynomials. Specifically, we define operator monomials and their "multiplication" follows:
S = ( H', S sti = S((+)), ijEN.
One can verify that the multiplication "e" distributes with the usual addition "+", therefore can define polynomials of operators. We prove the following key equations that connect operator polynomials with how the 8-th order tensor H interacts with operator (11) (see Appendix D.5
o = (s(0)- - ) S(U),
o S = II (S ) s()) k=1
In addition, we note that the operator polynomials are all diagonal operators that contain only d2 degrees of freedom (unlike general operators that contain d degrees of freedom), thus we can compute them via relatively simple algebraic rules (see Appendix D.5) D.5))
Variance and bias error. Up to now, we have introduced diagonalization to simplify the operator recursion and operafor polynomials to compute the simplified operator recursion. 'he remaining efforts are to analyze the variance error following the methods introduced by Zou al. (202I); Wu ct al. (002) (see Appendix D.6.. The analysis of the bias error is more involved; it is presented in Appendix D.7
CONCLUSION
This paper studies the in-context learning single-layer linear attention model for linear regres- sion with Gaussian prior. We prove a statistical task complexity bound for the pretraining of the attention model, where we develop new tools for operator methods. In addition, we compare the average linear regression risk obtained by pretrained attention model with that obtained by op- timally tuned ridge regression, which clarifies the effectiveness of in- cn-tontext learning. Our theories complement experimental results in prior works.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENT
We thank the anonymous reviewers and area chairs for their helpful comments. DZ acknowledges the support from NSFC 62306252. ZC and QG are supported in part by the NSF grants IIS-1906160 I and the Sloan Research Fellowship. VB is partially supported by National Science Foundation Awards 2244899 and 2333887, the ONR award N000142312733 the Ministry of Trade, Industry and Energy (MOTIE) and Korea Institute for Advancement of Technology (KIAT) through the International Cooperative R&D program. JW and PB are supported in part by NSF grants DMS- 2023505 and AMS-2031883 and Simons Foundation award 814639. The views and conclusions contained in this paper are those of the authors and should not be interpreted as representing any funding agencies.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2022.
Francis Bach and Eric Moulines. Non-strongly-comm convex smooth stochastic approximation with con- vergence rate o o1/n)). Advances in neural information processing systems, 26:773-781, 2013..
Yu Bai, Fan Chen, Huan Wang. Caiming Xiong, and Song Mci. Transformers statisticians: Provable in-context learning with in-context algorithm selection. In Thirty--eventh Conference Neural Information Processing Systems, 2023.
Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning, vol- ume 4. Springer, 2006.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neurul information processing systems, 13:1877-1901, 2020.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can GPT learn in-context? language models secretly perform gradient descent as meta-optimizers, In Findings of the Association for Computational Linguistics: ACL 2023, 2023.
Aymeric Dieuleveut, Nicolas Flammarion, and Francis Bach. Harder, better, faster, stronger con- vergence rates for least-squares regression. The Journal of Machine Learning Research, 18(1): 3520-3570, 2017.
Shivam Garg. Dimitris Tsipras. Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 15:30583-30598, 2022.
Rong Ge, Sham M Kakade, Rahul Kidambi, and Praneeth Netrapalli. The step decay schedule: A near optimal, geometrically decaying learning rate procedure for least squares, Advances in neural information processing systems, 32, 2019.
Angeliki Giannou, Shashank Rajput, Jy- Yo-y Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 11398- 1442. PMLR, 23-29 Jul 2023.
Chi Han, Ziqi Wang, Han Zhao, and Heng Ji. In-context learning of large language models explained as kernel regression. arXiv preprint arXiv:2305. 10t.:2265....66. 2023.
Prateek Jain, Praneeth Netrapalli, Sham M Kakade, Rahul Kidambi, and Aaron Sidford. Paralleliz- ing stochastic gradient descent for least squares regression: Mini- Mini-batching averaging, and model misspecification. The Journal of Machine I Learning Research, 18(1):8258-837) 2017.
10

---

Published as a conference paper at ICLR 2024
Prateek Jain, Sham M Kakade, Rahul Kidambi, Praneeth Netrapalli, Venkata Krishna Pillutla, and Aaron Sidford. A markov chain theory approach to characterizing the minimax optimality of stochastic gradient descent (for least squares). In 37th IARCS Annual Conference on Foundations of Software Technology and Theoretical Computer Science (FSTTCS 2017). Schloss Dagstuhl- Leibniz-7Zentrum fuer Informatik, 2018.
Shuai Li, Zhao Song, Yu Xia, Tong Yu, and Tianyi Zhou. The closeness of in-context learning and weight shifting for softmax regression. arXiv preprint arXiv:2304. 132.6 2023a.
Yingcong Li. Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, pp. 19565-19594, PMLR, 2023b.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. 23. JMLR. org, 2023c.
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. In The Twelfth International Conference on Learning Representations, 2024.
Allan Raventos, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
James R Schott. Matrix analysis for statistics. John Wiley & Sons, 2016.
George AF Seber. A matrix handbook for statisticians. John Wiley & Sons, 2008.
Alexander Tsigler and Peter L. Bartlett. Benign overfitting in ridge regression. Journal of Machine Learning Research, 24(123):1-76, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. International Conference on Machine Learning, pp. 35151-35174. PMLR, 2023.
Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang. Large lan- guage models are latent variable models: Explaining and finding good demonstrations for in- context learning. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Jingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu, and Sham M. Kakade. Last iterate risk bounds of SGD with decaying stepsize for overparameterized linear regression. The 39th International Conference Machine Learning, 2022.
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Sham M Kakade. Finite-sample analysis of learning high-dimensional single ReLU neuron. The 40th International Conference on Machine Learning, 2023.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit Bayesian inference. In International Conference on Learning Representa- tions, 2021.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprini urXiv:2306.09922 2023a.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? Bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305. /9420, 2023b.
Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, and Sham Kakade. Benign overfitting of constant-stepsize SGD for linear regression. In Conference on Learning Theory, pp. 4633- 4635. PMLR, 2021.
11

---

Specifically, we Published as a conference paper at ICLR 2024
2.0
ICL
2,0
ICL
Ridge Regression
Ridge Regression
OLS
OLS
M5E
144
1.4
1.2
1.2
10a
10a
10a
10a
10a
10a
10a
10a
10a
10a
Number of Pretraining Tasks
Number of Pretraining Tasks
(a) Exponential decay spectrum, A; 2 2
(b) Polynomial decay spectrum,
Task complexity of ICL (of the one-step GD model), ridge regression, and OLS. The context length is MM N 200. The ambient dimension = 100. We observe that as the number of pretraining tasks increases, one-step GD achieves smaller MSE and becomes closer to the Bayes algorithm, ridge regression. This is consistent with our theory.
Figure I:
A EXPRIMENTS
In this section, we conduct experiments on the one-step GD model (1) and three-layer transformer.
A.1 THE ONE-STEP GD MODEL
Data generation. We follow the generation process outlined in Assumption sample (X,, Vn N as independent copies of (x, y). where
NOO H),
N o²), B v I2))
We treat the first N data points (Xn: bn) as the context examples, XN+1 as the covariate, and UN+1 the response.
Base experiment setup. We configure the base experiment with the following parameters:
100, N = 2d, o 1, 3 = 1, = diag(2 2- ........"9
We sample fresh sequence (Xn. Vn N+ for each task. We train the ICL model using online SGD (see (6)) with a geometrically decaying stepsize schedule defined in 77). We run online SGD for 10 steps. The default initial learning rate is set 0.1. For evaluation,
consider in-context sample size M = N 200 and compare against benchmark algorithms such as optimally tuned ridge regression (Theorem 5.1 and Ordinary Least Square (OLS). We conduct series of experiments by varying parts of this base experiment setup, that is, the experimental setups are identical to this base experiment setup unless noted otherwise.
The effeec of the number of pretraining tasks. To examine the pretraining task complexity, we vary the number of pretraining tasks in the base setup in the range [10', 102 10 10 109] In addition the exponentially decaying spectrum considered in the base setup, we consider polynomially decaying spectrum with A; i-2 For different spectrums X; = = and A; 2- initial learning rates were optimally tuned from the set (0.005, 0.01, 0.05, 0.1, 0.5), resulting in
an optimal rate of 0.1 for both. Results presented in Figure We observe that the ICL error decreases as the number of pretraining tasks increases.
The effect of the ambient dimension. To examine the pretraining task complexity, we vary the ambient dimension in the base setup in the range of d {10, 20, 50, 100). We also consider polynomial decay spectrum with X; = i 2 Results are presented in Figure We observe that the ICI. performance is relatively unaffected by the ambient dimension d.
12

---

Published as a conference paper at ICLR 2024
2.2
2.0
Ridge Regression
2,0
Ridge Regression
OLS
OLS
1.8
1.8
1.6
MSF
1.6
144
1.1
1.2
1.2
10
20
50 100
10
20
50
100
Dimension (a) Exponential decay spectrum, A 2
Dimension
(b) Polynomial decay spectrum, A;
Figure 2:
The effect of the ambient dimension for ICL (of sne GD), ridge regression, and OLS.
The context length is M = N = 200. The number of pretraining tasks is 109 for ICL. We observe that when the spectrum of the data covariance H decays relatively fast, for example, X; ~ 2- and 2; ~ 2 the performances of the three considered algorithms are not sensitive to the ambient dimension. This is consistent with our theory.
6
3.5
ICL
ICL
Ridge Regression
5-
Ridge Regression
3.0
2.5
MSE
2.0
3
1.5
1.0
1-
20
160
1
40
160
Inference Incontext Samples
Inference Incontext Samples
(a) Exponential decay spectrum, A = 2
(b) Polynomial decay spectrum, A; =
Figure 3: The effect of the number of context examples during inference for ICL (of one-step GD) and ridge regression. The number of context examples during pretraining N 40. The ambient dimension is d = 20. The MSE of OLS is significantly worse than ICI. and ridge regression when M S N 20, so we ignore OLS in this plot for better visualization. We observe that the ICL. achieves a similar MSE to ridge regression when M is close to N. However, the gap becomes larger when M is much smaller than N. This is consistent with our theory.
The effect of the number of context examples during inference. We modify the base experiment setup with d 20, NN 40. We then examine the effect of the number of context examples during inference by varying M. Similarly, we also consider polynomial decay spectrum with A; = i-2 Results are presented in Figure 3 We observe that when M is close N, the number of context examples during pretraining, the ICL risk of one-step GD is close to that of optimally tuned ridge regression. However, the gap becomes larger when M much smaller than N. This consistent with our theory.
The effect of model misspecification. The base experiment setup assumes well-specified data. We now investigate three misspecification scenarios:
1. Replacing the label generation process from y ~ NV x, 02) to y ~ B x + uniformi-c.c c), where we set c = v3 to maintain the noise variance.
Replacing the label generation process from y ~ NV x, 0²) to ~ Boo~~N(6) ((simmodd (x))))(((xx)))))))) o2).
3. Replacing the label generation process from y ~ NN x to 120g~N((3))) x)),
13

---

Published as a conference paper at ICLR 2024
10
Base
Sigmoid
Square
8
Uniform
MSE
6
4
2
5
10
25
50
100
200
400
800
Inference Incontext Samples
Figure 4: The effect of data misspecification for the ICL of one-step GD. The base setup, ~ N(B x, 000 with o² 1, is well-specified. We then consider three misspecification scenarios. [ v3). Sigmoid: y ~ ((siimood x),o²), Square: y ~
Uniform: y B x N(B o²), We observe that the type of misspecification affects the ICL performance. In particular, the ICL performance declines less when the ground-truth model closer to a linear model.
Results are shown in Figure 4 case is close to the Gaussian noise case in the base setup. However, when the mean of y is not linearly related the base setup. The performance deterioration depends the type of misspecification, with y ((3 x²),, ,²)) showing the most significant decline.
We observe that the ICL of one-step GD in the uniform noise
as the latter two cases, the ICL of one-step GD is significantly worse than
A.2 A TIREE-LAYER TRANSFORMER
We conduct experiments the task complexity for training transformer. We adopt the code by Bai et al. 22023 We consider a three-layer transformer (GPT model) with heads. We follow the generation process outlined in Assumption Specifically, we sample (Xns Bn) M copies of (x, y), where
N+1 as independent
X x~ A(O, H), y x~N((3xx,0)) B N00 II.).
We treat the first N data points (xn: ) as the context examples, XN+1 as the covariate, and gN+1 as the response. We configure the experiments with
d 20, N 2d, a 0.5, & = H = diag(1,2
For each task, will sample 64 i.i.d. sequences of (Xn; yn) W+1 The model is trained with Adam with learning rate of 0.0001. We set the number of context examples during inference to be M N. The results are presented Figure [5] Similarly the one-step GD model, we also observe that the ICI. error decreases as the number of pretraining tasks increases, approaching the performance of the Bayes optimal algorithm, ridge regression.
B SINGLE-LAVE LINEAR ATTENTION AND ONE-STEP GD
Results in this part largely follow from (Ahn al. here for completeness.
2023 Zhang al. 2023a We include them
Denote the prompt by
Z =
X y
E R (d+1))(x+1))
14

---

Published as a conference paper at ICLR 2024
0.9
ICL Ridge Regression
0.8
OLS
0.7
MSE
0.6
0.5
0.4
2
27
200 2
2
2 21 2 210
Number of Pretraining Tasks
Figure 5: ICL of three-layer transformer. The linear regression tasks are generated according to Assumption with = 20, N 2d, standard deviation a = 0.5, scaling factor v 1, and polynomial decay spectrum X; i-4, We fix the number of context examples during inference to .M = N. We observe that the ICL error decreases as the number of pretraining tasks increases, approaching the performance of the Bayes optimal algorithm, ridge regression.
Denote the query, key, and value parameters by
Q,K, V E ((++1)x((+1),
Then the single-layer attention with residue connection outputs
(Z) (KZ) E [R(++))(x++),
Z + (VZ)
"
The prediction is the bottom right entry of the above matrix, that is
1 Z + (VZ) (KZ)
n
d+1,n+1
ea-1
T
Z +
VZZ Q KZ en+1
= 0 + edl1
vzz Q KZent1
n
1 (e V)
X x xxx y x
xx
X y y
Q
K
Our key assumption is that the bottom left 1 d block in V is fixed to be zero and the bottom left 1 x d block in Q K is fixed to be zero, that is, we assume that
0 )) QK
W 0
where e E R and W E ddxd are relevant free parameters. Then we have
X X + T
T
xTy)
Wx
i - (0,
y x X
y y
XWx
((WT x
x
n
which recovers one-step GD when replace wW by П, i.e. the update formula in 1
15

---

Published as a conference paper at ICLR 2024
C
POPULATION ICL RISK
Lemma C.I. Suppose that the rows in X E Nx are generated independently
X|i ~ NO(, H), i= 1,...,M. N.
Then for every PSD matrix A, it holds that
E x XAX Ntr + N(N + 1)HAH.
Proof of Lemma C.1 This is by direct computing.
B x XAX
E. Σ
X A E [xxx))
NExx T Axx + N(N 1)HAH
-(tt(HA))H+ + 2HAH) + N(N - 1)HAH
-err(HA)H + N(N + +JHHAH.
This completes the proof.
We are ready to present the proof Theorem 3.1
Proof of Theorem
3.1 Let be the task parameter and let
: x 3, : y x,3.
Then from Assumption 1 we have
X~ (O, H).
X ~ ~(0
B~N10, E~N(0,07) c~N(0,0'IN)
Bringing this into
we have
2
R2 (T)
Ix y,
-
and
E
x r
X X 3
+ x
T r
x e
x xT -
N
x T
I
x
x
/
+
[xxx
x T rx T c)² 0
N I-
N2
602
Exx
S2
x
X
o
32 82
LE
eTe x
EE
+ 0²
N2
62
H.
E
I- r 1
x x
(1
+ N
H. No THI
+0²
N
H. J --
1-
I-r N
I
--
PHI
+02.
(15)
N
Next, we compute the matrix in
15 that involves I, that is
T
E
r
x
x
THI
N
N
2/2
(PH
N NT
N2
2
N(N + 1)H²) + 0²
[TH
H PT
by Lemma C.1
N2
N
[T
+
4
16

---

Published as a conference paper at ICLR 2024
(T H
HN (TT) T + - -
where the last equality is because N
HHH by Here, we define
1 (
2 tr H 0²
22
N +
Hy 10
x
Fw == HH
Bringing this back to (5).. we have
RT) = H [T N
H (T - ") + H. - 7) +
is clear that
min RN(()) = ( H, I- r HN (Tx) )) +o?
and
RT) min Rx(+) = H (T- r2) rTx) T)
We now compute min Rx(() (-) as follows:
min R-(())
H. I HN (Tw) / +
H, 1 HH to²
(J'HH N
HN - H² +
(tr(H)
/
+
tr((H) + 02/0²2
H
: 2
N H +
N
N
N
I+ (N +1HH = wt ((tr(H)) + + /2/1/))1
+ (
((tr(H)+00/0))
+o²,
which completes the proof.
D THE TASK COMPLEXITY FOR PRETRAINING AN ATTENTION MODEL
D.1 PRELIMINARIES OF OPERATOR METHODS
Tensor product. We use & denote the tensor product or Kronecker product. For convenience, we follow the tensor product convention used by Bach & Moulines (2033;; Diculeveut et al. (2017); Jain al. 2018 2017); Zou et al. (02T); Wu al. (2022) for analyzing SGD.
Definition 1 (Tensor product). For matrices and B of any shape, B @ A is an operator on matrices of appropriate shape. Specifically, for matrix X of an appropriate shape, define
(B A) X = AXB.
is clear that B A is linear operator. For simplicity, we also write
@ A A.
We introduce a few facts about linear operators matrices. Fact D.1. For matrices A, B, C, and D of an appropriate shape, it holds that
(D C) B A) (D @ (CA).
Proof. For matrix X of an appropriate shape, we have
D & B @
A (D & C) (AXB)
= CAXBD
(D B s (CA) o X,
which verifies the claim.
17

---

Published as a conference paper at ICLR 2024
PSD operators. A key notion in our analysis is that of PSD operators, which map PSD matrix to another PSD matrix. Definition (PSD operator). For linear operator on matrices
O R dxd
+ gdxd
say O is PSD operator, if
OoAzO, for every A 0. Definition (Operator order). For two linear operators on matrices
O1, O2 R Rdxd
- R
we say
0 W O2
O2 O1 is PSD operator.
D.2 BIAS- VARIANA DECOMPOSITION
SGD iterates.
Fix the current iterate index as t 2 1. Recall that
L
XXTY) y,
(
)7
RT; Ye, = xxx, (r -r-r") r
+ +
dr
N
where I" is defined in (3 and
[ X, ye ( N
yext ( xxxxx)²
E
X1X r
(16)
The next lemma shows that E has zero mean and hence behaves like a "noise". Lemma D.2. For random matrix =. defined in (16), it holds that B|3,] = 0.
Proof. This is because
XX T (
N xxy ((XTT) -XV N ( '7
ES = E
N xTy) 8) ( (xx)x)=
Exx T 3 (
-II'E
HH 0,
H2
where H is defined in and holds that
H = 6
We complete the proof.
We can now write the SGD update as
r, P Ør
R(T,-1
-- Xe,
xxxxx) ( -m 1= 1....,?? T,
r x (T--1- rr- r")
where (w))"1 is a stepsize schedule defined by (7)
Define
A, = r, - r
then have
1 ( -,, T xx,Tx)
A, = A-- - 78x,x, Ar-1
18

---

Published as a conference paper at ICLR 2024
Bias-variance decomposition. Define
P ddd Rddd
1 N xx!x1) (
A + A - xxIx, A
is clear that P. is linear map on matrices. Then we have
A P1 o 0 - - 1
Solving the recursion, have
0 7 II
9 o
k=H+1
Taking outer product and expectation, we have
AT == EA 02
T
= E
P oo
Σ 7
II
Pg
P,OEE
(=1
1=1
k=t+1
T
82
T
T
82
= 21E
II
P, An
+
2E
1/.
11
2x o E
1=1
k=1+1
=: 2Br + 2CT,
where we define
T
82
By := E
II 1=1
P, o Ao
(17)
CT == = E
II 0z k=t+1
(18)
Therefore, we can bound the average risk by
ER, (Tr) min (-)-)) = EE H, (Tr - I') -(T))
by Theorem 3.1
HH, Ar o H)
< H. By H) + 2(H Cr H)
The above gives the bias Marrinai decomposition of the risk.
Operators and operator maps. ces:
Define the following three linear operators on symmetric matri-
M = E(xx 002
(19)
C: ( x 8) (+xxx)) 62
(20)
N = E 02
(21)
It is easy to verify that all three operators PSD operators, that is, PSD matrix is mapped to another PSD matrix.
Define the following SGD map linear operators:
Pt (Rdxd 602 + dxdd 022
0 o
(22)
+ O...
19

---

Published as a conference paper at ICLR 2024
Similarly, define GD map linear operators:
g: (I) (Ruxd)& 82
O - O - 7a
(23)
H 0 H
When the context clear, we also use and y and ignore the subscript stepsize Th. When the context is clear, we also write
9(0)= - 9000 o0,
The following lemma explains the reason call these two maps SGD and GD maps, respectively. Lemma D.3 (GD and SGD maps). We have the following properties of the GD and SGD maps defined in (23) and (22) respectively.
1. I and I are both linear maps over the space of matrix operators, i.e., for every pair of matrix operators O₂ and O₂ and every scalar a E R.
9(0; + + (0)) + at (O2), gP + =(00)))))))=
2. For every matrix P of an appropriate shape, it holds that
( pp82) (P - HHHH) 82
and that
(? x )))
42
1 (1 xTxy)
) - E P 7xx P
N
which corresponds to a single (population) GD and SGD steps on matrix P, respectively.
3. As a consequence of the first two conclusions, holds that & (O) and (O) are both PSD operators if O is given by
O = E x P
where P is of an appropriate shape and is possibly random.
4. It holds that
g 2 () 0.
Proof. The first conclusion clear by the definitions of 23 and 22
The second conclusion also follows from the definitions of check that
23 and
22
For example, we can
(P = p - (H &H 82 a
(H&l)
+ (I& H) 882 (I H)
H 2 H 82
p222 (HPH) & P P (HPH)) + +22 (HPH)
(P HPH 82
The third conclusion follows from the first two conclusions.
The last conclusion is clear by the definitions of 23 and (22).
Bias iterate. Using the SGD map
22 can we-weite 17 recursively
B A 82 (To r
(24)
B, J, o B
t 1,...,, T.
20

---

Published as a conference paper at ICLR 2024
Variance iterates. Let consider the variance iterate defined in and is independent of 3x for & 2 t 1, we have
18
Since E has mean
T
T
82
Cr
7t
II
Rx o
E
k=1+11
82
E
k=1+1
Using the SGD map (22 and the noise operator (1). we can re-write above recursively as
Co
(25)
C. + N,N, t= l=,...., T.
D.3 SOME OPERATOR BOUNDS
Lemma D.4. Suppose that z (0, Ia), then
1. For every E Rd,
E(z, u) v)² 3/lu|/ ||v|22.
2. For every V, W Rd.
Eiz, u) (2, (2, w)² < 15 u5 1Iv12
||w12
3. For every u, v, W X.
Ez,, u) v) (z, (2, x)²
105
||
|
Proof. These inequalities can be proved by using Gaussian moment tensor equations in Section 20.5.2 in (Seber 2008) and Section 11.6 in (Schoti 2016). Specifically, for the fourth moment, we have
B(z, (2 Ez uu z z Tv
tr(uu" )tr(vv]] + 2tr(uu vv
||u|| 11v11 + v(u,v)?
3 3||u| ||v122
For the sixth moment, we have
E(z, u) (z, v) (z. w)²
Ez uu z vv z- ww z
=tr(uu ttr(vv (rr(ww] ) + 2tr (ereuuu)trrv'' Jtr(vvT ww
2tr(vv )ur ww + 2tr(ww trr(uu vv + Str(uu T wwww)
||u|| ||vl2 ||w]22 + 2 w) /l w)2
+ 2||w 12 [u, v\² + 8(u, v) w)
< 15| |5]|u|| ||v1 ||
For the eighth moment, we have
E(z, u)² (z, v)² (2, Ex un n a vv 2 z ww in x xxTx Tz
=tr(um T )tr[vv (tr(ww (xr(xx])
+ (s((truuu)) )tr(vv ww T xx + trv trr(uu Tww Twwxxx
21

---

Published as a conference paper at ICLR 2024
+ tr((ww (e(ww))tr(uu (rr(uu 'vv"xx xxx) ) + tr(xx trr(uu "vv ww ))
+ 4 (tr(uu' vv "trrr(wwwT (xx)) + tr(uu' T ww ")rr(v''' xxx)
+ +tr(uu' I xx "x"))r((vv (wrrv)""w))
+ 2 (tr(uu )tr(vv )tr(ww" xxx) + tr(uBu )tr(ww )tr(vv xx
+ tr(uu )xx(xx jtr(vvT ww + tr(vv (e(vvV)tte (tr(ww Jtr[uuuxxxx )
tr(vv ttr(xx (tr(uu ww + tr(ww )tr((xx Itr(uu (wr[uuuuTTT)) vvT) "owl
+ 16 (truuu vv ww xx x tr(uu T vv xx Txx www) + tr(uu" www 'vv vvTxx!)
=/lulB3 1|v12 || ||x1
+ ++ (2,, w) (w. (v, x) + ( ww (u, x)
+ +//wl/B (1, + (u, (u, ()(v, w))
+
w)
+ 2(11u|1
|w ||u| ||x|2 w)²
+ +|vl|| || |v1 3
(u,
+ +/ (), w)
(u. (v, + (u. w (v, (v
< 105||u|
|v12 [ 12 ||x|
We have completed the proof.
Lemma D.5 (Upper bound on M). Consider M defined in 19
For every PSD matrix A, we have
Mo A = E[xx Axx
1 A)H.
Proof. This follows from Lemma D.4
Lemma D.6. Consider defined in 20
For every PSD matrix A, we have
((>xx)))
T A (
(A, LoA =E E.
( A)?
Proof. By definition, have
((XX'))
T
A N
( 2
(A, Lo A) =F
E
N
X
X + Ell
N
X x,,1 + +BI|XTel11 ell
(26)
Next, we bound each of the two terms separately.
22

---

Published as a conference paper at ICLR 2024
Bound E
X X A
We have
A
E
X X FE(XXXXXXX XAX
I X XAXX x)²
by Lemma D.4
=300*E(A A, xx x)
=30 E E(A, xxx x,xxxxx)
XXXX xxX?)
3r0
++ 3-distinet
Σ 2-distinet
Σ
f(i,i,ke).
(27)
4-distinet
I-distinet
where we define
f(t.j,k, := X,X x A,
X xxx?)
E x Axy X; Xj X Axe
X xt)
In 27 ,. we group j. k., by their number of distinct indexes (i.e., the number of distinct random variables). We now bound the sum of the terms in each group separately.
There are no more than 2 terms that have 4 distinct random variables and each such term can be bounded by
f(1,2,3,4) (H² A) A ( A).
So we have
Σ
((.J,Kh N H A)
4-distinct
There are no more than 3 N terms that have distinct random variables. Due to the i.i.d.- ness, may assume xı appears twice and X2 and X3 appear once in such 3-distinct term without loss of generality. Due to the symmetry of (i,,j, k E), there are essentially two situations.
1. If two x 's appear in the same inner product, such 3-distinct term can be bounded by
f(1,1,2.3) E,A, X X1
X
Sex]xxxx
= (A, H E x x x Ax,]
< A, H
3 A) tr (H)
H. A) H A) 3tr H
2. two xi 's appear in different inner products, such a 3-distinct term can be bounded by
f(1,2,1,3) = E Axe xxxx xx Axs X1 X3
= [x AHx))
< 3 A)2
< 3tr HH A)
Therefore, we can upper bound the sum all 3-distinct terms by
Σ fli.j.k, < N 3tr (H, A) A).
3 distinct
There are no more than 2 terms that have 2 distinct random variables. Due to the i.i.d.nness, may assume X1 appears twice and x₂ appears twice in such a 2-distinct term without loss of generality. Due to the symmetricity of fi. Je, E), there are essentially two situations.
23

---

Published as a conference paper at ICLR 2024
1. If two x1 's appear in the same inner product, such 2-distinct term can be bounded by
f(1,1,222 - E(A, X1x x1x x (A, SxxxxxX] - A)tr()) -tr(HH) 2 (H, A)
2. If two x, 's appear in different inner products, such 2-distinct term can be bounded by
1(1,2,1,2) E[x] Ax2 x X
Ax2 X1 x2] E[x] X2X X1 Ax2x; Ax1]
S HH, X2X, Ax2x, A)
3E [x, Hx AHAx,]
S 9tr (H, AHA)
S (tr(H)) (H. A)
Therefore, we can upper bound the sum of all 2-distinct terms by
Σ ((,,,,,)) N² 9tr (H) 2 (H,
2-distinet
There N terms that have only distinct random variable and each such term can be bounded by
((.1,11)))([x([]6 Ax)]
S (H) A)
So we have
Σ
f(i,j,k, (H)
1 L-distinet
Applying these bounds to
we get
E E
x,311
S 3rd
( A)²
+ + (tr A) (H2, A)
+ +2
9tr H A)2
105
H A)2
tr(H)
A) 2 A) + tr(H)²
S N3
H A)
H.
(H. A}3
N
N2
tr(H) H
2
$ N"
HH², A)
H.
A)
N
tr(H)
2
3 N4
I I H
A
(28)
Bound on E
x
We have
A
E
X
A = - 2)²
$ <3R(0PI, XAXT)²
=30"B(A, x x
-30 1E(A) xxx, (A, xxxJ)
24

---

Published as a conference paper at ICLR 2024
30 (E x + +N(N - XIX
S 330 (3 (3N
H) 2 + N(N- 1)),
53
H 2)
0² 12
N
H
A
(29)
N
Putting things together.
Bring 28 and
29 to 26
we obtain
(A, Lo A) S N4
No
tr (H)
A
2
+ +3
W +
(1))))))))))
/0//22
N 1 A 2
N
tr(H) + +02 /2//62
2
8..39 36 I
H
)
8.. 3 (H, A)²
We have completed the proof.
Lemma D.7 (Upper bound on L). Consider defined in (20..
For every PSD matrix A, we have
Co = E 1 x y
(
x A N xxxx
(xXTs)
38 3 H
Proof. We only need to show that for every PSD matrices A and B, it holds that 3
B LoA) 58 38 H A) H)
This is because:
1 (
1 (XXTy)
(
1
(B, Lo A)
= E B.
A
N
x y
N
E
XTy y
B N
(XTT)
(
A
N
x y)
1
N
((*XT>)
T
1
5
(XTY)
B XTy
E
x
(B, Cb V(A CoA)
8- (H, ().
where the last inequality is by Lemma D.6
Lemma D.8. For every PSD matrix A, we have
62
xxXT))
yx
o
(V
Proof. First, notice that
1 N xxy) 3) ) 82
yx
AA
E N xxTy) A A(x ( x
25

---

Published as a conference paper at ICLR 2024
For the first factor, we take expectation with respect to X and (i.e., conditional on 3) to get
T
E xxTx) A x
E
e)
A
x
3
N x
1
1
EX
XAX
x3
BE
X
AXE
N2
N
B
A)
H+
N
HAH
B
H.
N
N
N
Similarly, we compute the expectation of the second factor with respect to x and (i.e., conditional on (3 to get
2yyxxx T = R(xxxx B -xxx +)) ) xx
Exx (3 xx + EEexxx
(H, )H + 2H33 H + o'H
3 (38TH5 + +
Therefore, we have
(xTT) )
82
E
yx
x
A
N
T
E2
x
A
xTx xxxx
N
AH
A)
N+
(H,A) (H,)) N
E
33
HAH
B B
(3 HB + o²) H
N
N
AH
+
(H,
3
H+
HAH
B + 30
EB H/
N
N
N
(H, A)
N
H A) N
+ RB
H
HAH
HAH)3.02 1..0+ +02
o²
H
N
N
(H. A)
N+1
HH, A) N
98 9r tr
H+
HAH
+ 30 tr(H)
2) (H)
N
N
A)
N+1
H. A
+ 2 tr
+
HAH
-0-
N
N
N
tr (H)
N
t
(H,
A)
(H
A)
A)
N
N
N
(
N+1
AH
2 A)+
H A) H
N
N
tt) 2222
N +
9
A)
N + H
t + o
N+1
HH A +
H+
= HH, A) This completes the proof.
tr (H)+
Lemma D.9. For every PSD matrix A, we have
E
XX
p
L
x
(+xxT)))
o A 38 38. ())))).
26

---

Published as a conference paper at ICLR 2024
Proof. By definition, have
xx1x) (>xxx>) ) oA
82
E
xx r
N
Exx T
xxyy)
Nx's)
1
A N
xxxxx)
1 (FxXT)
X'xx
N
N
Exx (C A)
< HHH ( H - 3" (H, A) IHHT),,
where the last inequality is due Lemma D.7
Recall from Theorem 3.1 that
(T*) ²
(r( + 0
N 1
r
It
H
H
which implies that
H. F" = tH((TTT)))
I tr(HH.
Bringing this back, we complete the proof.
Lemma D.10 (Upper bound on NV. Consider N defined in 21
For every PSD matrix A, we have
Ne A EEAST
3 (16-37 +
Proof. Note that
(A + BLBBA + BB) 2((AXA BXB]).
So we have
No A
E2 A
(M x
1
1 )
E xX T
M XT y)
xy
AA
M
M
N xTy y) 2)(3))
N (NxTx))
02
1 N xxxTy) )
02
1 2E
xx T
o +2
1/X
A
16. ( H A
tr HHH + 18 ( (r + H
by Lemmas
D.8 and
D.9
(16 37 18) (H) 2 (H
which completes the proof.
Lemma D.11 (Lower bounds on M and L). For M defined in 19 and L defined in 20 we have
M H.H,
and
L H.H.
Proof. For every PSD matrix A, we have
(M - H HHH) o A = Exx HxX HAH
E(xxx A)A(xxX HH)
27

---

Published as a conference paper at ICLR 2024
where the second equality is because
Exx T H.
Similarly, for every PSD matrix A, we have
1 ( x A
1 (xXT)
(c - H H) o A
B
x'y y)
x y)
HAH
N
(xxTT)
T
A ( ((xxTs)
(rxTx)²
=
x y
H
N
=
where the second equality is because
(x)y)
1 xTy
( N
1 H.
E
N
Lemma D.12 (Composition of PSD operators), For every PSD operator O, it holds that
H 02 MM000138.3 3 (H, O H)S(),
where S()) a PSD operator defined by
8(1) H H.
As a direct consequence of the lower bound, we have
9.0z 90 00
Proof. For the upper bound, let us consider an arbitrary PSD matrix A. We have
McDaLeA A < 88 3 (H, A)M O e H
by Lemma D.7
38.37 (H, [H, A) (H, Oo H)H
by Lemma D.5
= 3 H, O H) A,
by the definition S(1)
which verifies the upper bound.
The lower bound is direct consequence of Lemma D.11
D.4 DIAGONALIZATION
Without loss of generality, assume that H is diagonal. Let ID be the set of PSD diagonal matrices. For PSD operator O, define its diagonalization by
① D
(30)
D diag
When the context is clear, also write
diag(O):= :=
Lemma D.13 (Diagnoalization of operators). We have the following properties of diagonalization.
1. For every pair of operators O₁ and O2 and every sealar a E R, it holds that diag(O; + 02) diag(O1) + diag O2). diag(uO,) = adiag (O))
2. For two operators O₂ and O, such that 0 < O2, holds that diag (O)) - (aaag(O2)
28

---

Published as a conference paper at ICLR 2024
3. For every operator O, it holds that
diag (0))
Proof. It should be clear. We only prove the last claim.
Let K be a PSD diagonal matrix. By 23 we have
(O) K = O o K (HO (HK) + O (KH)H) HO (HKH)
Now taking a diagonal on both sides and using that K is also diagonal, we obtain that
| (0) K} diag(O K - (diiaddaHHH (HK)) + diag(O o (KH)H)
HOO (HKH)H)
Ö K (Ho
(KH)))
+H² (HKH)H
(()) o
which implies that
[iag(((0))=0(()
Bias and variance error under operator diagonalization. Since both H and H are diagonal matrices, we have
HH, By H) (H, By H),
(H, Cj o (H, Č o H)
which motivates us to control only the diagonalized bias and variance iterates. We next establish recursions about the diagonalized bias and variance iterates, respectively.
Diagonalization of the bias iterates. Consider the bias iterates given by
24 By definition of ,
in (2 and 5 in (33, we have
B. -, B,
by 24
9.0 B. + 7.M o B.-1oL. H2 BB_ 82
by 22 and
1 9 o B, + 7 M o 1
since B is PSD
3 3 B 3" (H, B, H)
by Lemma D.12
9, o B. 22 37 HH B_ H)
where the last equality is because both H and H are diagonal. Next, taking diagonal on both sides and using Lemma D.13 we have
B, h diag g. BB-1) + + 8 8.37 (H, B H)S
by Lemma D.13
g, o B, + 72
8. 3. 8.37
(H, B H))(1)
by Lemma D.13
(31)
where
B = diag (To
We have obtained a recursion about the diagonalized bias iterates.
Diagonalization of the variance iterates. Similarly, let us treat the variance iterates given by 25 By repeating the argument for the bias iterate, we have
C = 94 o C. ?,N
g o C. + + 88.2..30 8.33 (H, é. H))(1) +r.N.
29

---

Published as a conference paper at ICLR 2024
Using Lemma D.10 have
N (16-13 (16-3' ((H) + +07)00)
So we have
- 8-33 3 (H, H) 2 3 (8) + o2)s())
Similar to the treatment to the bias iterate, take diagonalization both sides and apply Lemma D.13 then we have
Ć. m 1 o c n 8-
8-37 (H, C- H)S0) + (16 3 (H)
(32)
where
Ć - -
We have established the recursion about the diagonalized variance iterates.
Monotonicity and contractivity of l on diagonal PSD operators. Finally, we introduce the following important lemma, which shows that is monotone when applied diagonal operators.
Lemma D.14 (Diagonalization of J). We have the following about the & defined in 23
1. For every diagonal operator D and every diagonal matrix K, it holds that
( Do K 2,HD (HK) + H2D o (H²
2. Suppose that
0 0<75
2tr (H)
then & is increasing map the diagonal operators. That is, for every pair of diagonal operators such that
D1 D2.
we have
3 (D2).
3. Suppose that
0<75 (tr(H)tr(H)
then 9 is a contractive map on the diagonal operators. That is, for every diagonal PSD operator
D = 0,
we have
((D) 3 D.
Proof. The first claim is clear from the definitions:
((D) o k= Dok - (HD (HK) Do (KH)H)
+ HT²HD (НKH)
DoK HD (HK) + HD (H*)).
For showing the second claim, notice that, by the linearity of E we only need to verify that for every diagonal PSD operator D, holds that
((D) - 0.
By definition, we only need to show that for every diagonal PSD matrix K, it holds that
((D) - 0.
We lower bound the left-hand side using the first conclusion:
((D) K D K HD (AK) -HH2 (H'K)
30

---

Published as a conference paper at ICLR 2024
2 Doo 110K-22100 2,HD (HK)
m DoK - (
= (1 - 2))r(
Similarly, prove the last claim by showing that
(D) K = Do K 22HD (HK) + + H (H')
Doc 22HD o (HK) + (H)))(H)) o (HK)
3 Do K -HH (HK)
3 K.
We have completed the proof.
D.5
OPERATOR POLYNOMIALS
In this section, we develop several useful new tools for computing the diagonal bias and variance iterates, (31 and (22
Operator polynomials. We first introduce operator polynomials.
Definition (Operator monomials). Define sequence of operator monomials:
((r) (H', H' tEN.
That is, for every i E N and for every symmetric matrix K,
g(t) K == (H. K H'.
Denote the set of all operator monomials by
$ (8 :1 EJ.
Definition (Operator polynomials). Let be a multiplication operation on S, defined by
S(i) s(i) = s(it)) i,j E.
Let "4" be the canonical operator addition operation. Let " distribute over "+" in the canonical manner, i.e.,
(80) + S(K) = S(1 8 + StD (kk) s(+ s(++k)
It is straightforward to verify that S(0) is the identity element under "e", 0 E ddd xd is the zero element under "4" We define a of operator polynomials by
(8 1787)-
Σ
(
IEN, R4
When the context is clear, also "II to refer to sequence of multiplication operations among the operator polynomials, e.g.,
(S II (S10) - Su)( = (S - ) (S(O) S .... (.) -1)))) 1=1
where (m))k=1 refers a sequence of positive stepsize.
The following lemma allows represent the composition of E over operator monomials as operator polynomials.
Lemma D.15 (Operator polynomials). We have the following results regarding the composition of operator monomials and other operators.
31

---

Published as a conference paper at ICLR 2024
1. For t 2 0,
(H 69 I) s o (H& & = (T H) H) o
&T H) = ((++)
2. For t 0,
H o s( H = s((+))
3. For I 0,
('()0)) = (S - - 22 s
4. For t 2 0,
II 9h (s 1=1
II (s - x s()) k=1
Proof. We now prove each claim respectively.
1. We consider symmetric matrix K and notice that
( (
H c)
(H @ I) S() (KH)
(H' ( 1) o H
(H' KH)H**
(H K)
g(++1) o K.
Similarly, we have
(I o S() (I H) o K = s((+) o K.
These verify the first claim.
2. We consider symmetric matrix K and notice that
(H Ø H) o sit (ń @ H) K
(H H) S
(HKH)
(H'', HKH @ H) H'
2 *2
g((+)) K,
which verifies the second claim.
3. Using the firs two claims and 23 we have
((S()) = S(t) H
(10)) (I+H)
7 ((t) 8
= (t)) (,5(++1) 2
(8(0) - 7)(1)) s(t)
Recursively applying the above equation and using the operator polynomials notation (see Definition 5). we get
(8(1)) = 5(1)
((((1)) (S(O) -2) +2
s(1)
(8(1)) (8 (1)) s(1)
32

---

Published as a conference paper at ICLR 2024
('(s())= = -5:1) 2* S()).
This verifies the third claim.
4. The fourth claim be verified similarly the third claim.
have completed the proof.
Computing operator polynomials. We now introduce method to compute operator polynomi-
als.
Notice that we only need to deal with diagonal PSD operators. Since diagonal PSD matrix has degrees of freedom, which can be equivalently represented by a d-dimensional (non-negative vector. Similarly, a diagonal operator has d x d degrees of freedom and thus can be equivalently represented linear map on d-dimensional (non-negative) vectors.
Define matrixization operation as
mat Rd Raxd
ki
mat(k)
ka
Then the operator monomial on diagonal PSD matrices can be equivalently written as
S(1) : D D
(33)
mv + H mat mat(v))H' = mat h of h T k).
where "0" refers to Hadamard product (i.e., entry-wise product) and h and are the diagonals of H and H. respectively, that is,
H1
h
(34)
H
Had
This viewpoint allows to compute operator polynomials. In particular, we can prove the following results.
Lemma D.16. When restricted as a diagonal operator, we have the following
1. For everyt 2 and every Rd.
(15()
s(L) mat = mat =
((J--hh"] T) o " ((hit)))
where J refers to the "all-one" matrix, that is,
= 11
2. For every 2 and every Rd
II (s - s(()) mmat(v) mat
(J- =0on1) v
Proof. By Definition
5 have
(S(0) - x)(1)
((1)
(-7) ggt+1
b
k=0
33

---

Published as a conference paper at ICLR 2024
Now using 33 we have
mat ma -581)) )
Σ
(1)
(->) mat(vv
k
k=0
(-7) mat 1 (hokt(
(hoktz) [k]
mat
Σ k) (-
(-r)" "n oott n (n ( T
k
mat
((J Th ) c (b)
which verifies the first claim. The second claim can be verified in the same way.
D.6 VARIANCE ERROR ANALYSIS
We first show crude variance upper bound.
Lemma D.17 (A crude variance bound). Suppose that
70 S 16 (Hr)
Then for 32 we have
č 3 s t2 0,
where
c c: (32 3 + 36) (v tr (r + +
Proof. We prove the claim by induction. For = 0, the claim holds since
C 0 S
Now suppose that
Let us compute è, by 32
č, x 9 o c 12 8- (H, ć o H)S + 716 3 tr)) + +0))
9 9 0 + 8. 3" ć_ H) +7 5 (())
by the definition of c
2
3 (((o)(()) + cnor? 8 3 (H, ( () + 2 s(u)
2
by the induction hypothesis
3 9 (cros()) + 7 +220 (1))
+ n 2
by the definition of ((0) and the choice of 7o
c)) + ((1)
since 9, is linear
cyo (s - (()(11) + t,
by Lemma D.15
cyo (70(0(00 20) + 722012) + c
by Definition since S - S(1)
3 CYo (S(00) - 71(11) + ((7
3 (0))
since 7 Y0
This completes the induction.
We next show sharper variance bound.
34

---

Published as a conference paper at ICLR 2024
Lemma D.18 (A sharp bound on the variance iterate). Suppose that
70 S
16- (r)
For every entry-wise non-negative vector v E R2, we have
Cr maa cmat ( (f(Tohi))
(hhi (ni))o-1)vv]
where
k 2
f(x)
x -
II j=
1
0 0<I<1,
2
2
8=
and is applied on matrix h entry-wise.
Proof. We first use Lemma D.17 to simplify the recursion in 32
C. 3 3, o + 8 3 (H, č, o H))()) + 716 3 + (8) tr((H)) + +0))8(1)
3 9, 1 c7o7 8 37 (H, ( H)S()) + 505
by Lemma D.17 and the definition of
1.
by the definition of ( and the choice of yo
We can unroll the above recursion using the monotonicity of on diagonal operators by Lemma Then have
D33
3
Co
+
II k=t+1
Sh
o ((1)
T
T
((1)
by Lemma D.3 and Co 02
1t
t=1
k=t+1
T
'''
II
(S - ) s()))
by Lemma D.15 15
[=1
k=1+1
Consider an arbitrary non-negative vector
VERd
v
and use Lemma D.16 then we have
T
o 3
(S
- 71s1))022
02 (es(i) mat
t=1
k=1+1
02 (hh v (J hhhh
emat
k=t+1
(J - (J-xhh²") @ (hh )u}
cmat
k=1+1
where the last inequality is because, by our choice of 7o the following holds in entry-wise:
0 < J ha SJ.
Let
K T:=T/lo([T) L= log(T) and recall the stepsize schedule (7). then for the non-negative vector v, have
Cr mat
35

---

Published as a conference paper at ICLR 2024
T
T
3 cmat
2
II
(J
hh
(hh)) )) v
(35)
k=
O(K
i)
cmat
70
Σ J
70 2
hh
2
L-1
oK
II
10
hh T
(hi
2
j=ビ+1
cmat
L-1 Yo
J
J
[J
7o
hh
@K
L-1
J
T0
hh
OK
2
2
2
1=
j=t+1
cmat
(f(1ohi))
① (hh
T) 0 )
where
1-1 [xx =
K
-
1
0 0<<<1,
2
2
21
t=
j=f+1
and is applied on matrix 7hhh? entry-wise.
The following lemma is an adaptation of Lemma C.3 in Wu et al. 2022 Lemma D.19. Consider a scalar function
L-1
K
L-1
K
[(x) =
0 0<<<l.
2
2
2)
t=
j=t+1
Then
8 K
<< f(x) S min
2K*
0 x 1
We are ready to show our final variance error upper bound. Theorem D.20 (Variance error bound). Suppose that
To S
16 (H)
Then we have
8c I
(H. Cr H)
min {1, x2
K
where
c: c32 37 + ( + )
K (T),
(di) are the eigenvalues of H. and (N:) i21 are the eigenvalues of H, that is
N, A tr (H)
+ o² ///
N N+1 >)) j21. 1.
N
Proof. Let compute variance error bound using Lemma D.18
H, Cr o H) = (H, ćr o H)
= (h), čr mat(h))
s mat mh) mat (1(zhhiT) @ (hh (nhT) )01))) )n)) (il)
by Lemma D.18
= ch (Soohiī) @ (hñ oo- h.
36

---

Published as a conference paper at ICLR 2024
By Lemma D.19 have
s f(zhhi") @ (hh)) (hh") )
S min { k
2K (zohh") 02
@ (hh) 6-1
8 min
{ (hiT))
K k2
K
where "min" and "s" are taken entrywise. So the variance error can be bounded by
H. Cp o H)
v ch
(f (f(rhh))
o (hh (hh")
) h
8 h min (hhit) K2 x hhh h
K
Sc
(A)
Aa)
min
1 Aedyy
xR3 (
K
8c
(A1
Aa)
E, min
K2 x23Bddyfi
K
Aedyy
8c
K K2 rards Srds
min
K
8c Σ
min {1, K2 k23g2?13))
K
where
tr H
+ 0²/1²2 //22 N
N + 1 X1)
A =
We have completed the proof.
D.7 BIAS ERROR ANALYSIS
Throughout this section, denote the bias error at the t-th iterate by
b = H, B, H = (H, B; o H).
(36)
where H (hence also H) is assumed to be diagonal and B3 admits the recursion 31p
D.7.1 MONSTANT-STEPSIEE CASE
Since the stepsize schedule ? is epoch-wise constant, begin our bias analysis by consid- ering constant-stepsize cases, where the stepsize is denoted by > 0. In this case, (31) reduces
B, x 4 By + 22 1 121, 1, where e = -7.3.
(37)
Unrolling 337 we have
R-1
B, x nd
B + 2
b 1-1 1
c1
B
=
Bo + 721
by (s(0)) 7)(1)) 11 S(1) n 2 1. by Lemma D.15 (38)
37

---

Published as a conference paper at ICLR 2024
Lemma D.21 (Controlled blow-up of bias error). Consider (77 If
1
? 201 2citr(H)rr trr
then for every n 2 0, it holds that
bn S (1+ (e+teerr(H)tr(((
Proof.
We prove the claim by induction. The claim clearly holds when n = 0. Now suppose that
by S (1 + (2)))))((())rr(()))()
t O...., n
For n, we have
n-1
B, g0 o Bu + 7 1
by (((00)) - -7(11) 2(nn1-1) S())
by
1=0
n-1
3 BB + cc
by (S(0)) 7)(11) 02(n-1-1) ((1)
by Lemma D.14
1=0
1 B + 72 260
(S ) 1(()
S(1)
where the last inequality is by the induction hypothesis and 1/2c, trH)), Next, consider an arbitrary non-negative vector v E IRd, by Lemma D.16 we have
B, mat
3 Bg mat(v} + 72 260
Σ (S ) 22(n-1-t) s(L) mat(vv
B (va + + mat
JJ 7hh 02(n-111)
c hhh )
1=0
by Lemma D.16
A Bg (v) + 722 2bymat
Σ (J hh o(n-1111) T 1=0
ohh"" hh v
since s -hhhh s J. entrywise
= Bu mat(v) + 7o, 26pmat (J (J- [J- On )
3 Bo mat(v) + yc, 2b
since 0 S J- - (J whh )on s J, entrywise
Then we have
bn (H, B, H)
(H, B, mat (h))
S (H, Bo mat (b)) + yeu 2ho (H, HH, mat(h))
bn ve, 2bo H, H)
< bn + 7c 26 tr (H)
(1 + 2c tr (H))
which completes the induction.
Lemma D.22 (A bound the sum of the bias error). Suppose that
? 20 tr tr)
38

---

Published as a conference paper at ICLR 2024
Suppose that
Bo = (To - -
and that To commutes with H. Then for every n 2 1, we have
Σ
n-1
(1 (I (I (1- 2n
(To - P7)).
Proof. By Lemma D.14 we have
9 B1 ol
B,_ -l 1,1--yHBBB11 H -HH (H))
1 B 1--1I-- -
2yHBB -
H + (H) H.
Using the above and 337. have
(I, B, 1)
(I, 9 B 1) 2 c/b-1 (I, 8 1)
by
$ (I, B - 2 B o H) + (x) B.1 o H
+72a tr (H)rr(HH))
(I, B.
276, (
S (I, 1 o - 7b-11
since S (r
Performing telescope sum, we have
Σ be 7 (1I, Bu II (I, Bn . 11).
We now derive a lower bound for Br. By Lemma D.12 and (24 we have
B f o
B
by the definition in 24
ス 6 B t 121. 1.
by Lemma D.12
Performing diagonalization using Lemma D.13 have
B, 1 B,. 121. 1.
Solving the recursion, we have
B - 1 o 1B
= g o (To - r 82
since both To and I' commute with by the definition of of in
-(1 (II- " "(To- 82
Putting these together, we have
b: 5 (I, By 1) - -(, Br 11)
1=0
< (tr((To - -))))) tr tr((1I HH (To - -)))))
I (-H)HH) 2n (To - (
which completes the proof.
Lemma D.23 (A decreasing bound on bias error). Suppose that
6c1 tr (H))trH)
Suppose that
B. (To r 32
and that To commutes with H. Then for every n > 0, we have
bn < max 1}
(I, (To - 2)
39

---

Published as a conference paper at ICLR 2024
Proof. We prove the claim by induction. For 0, have
bo (H. (To r HT, )
S (H) (I, (To 17)
(I, (To r-))
Now, suppose that
b maxft, 1}
(I. (To - ()))) t= t= ...., 0.........7 1.
For dn, considering an arbitrary non-negan non-negative vector v E Rd we have
Bn o mat(v))
n-1
(Ba o mat[[v] + +20
Σ bi ((8(0)-7(())
2( 22(=-1-1) swi)
o mva
t=0
-B)(Bn) mat(v) + +
Σ b.mat
T02001112))))) (hh T v
where the inequality is by 38 and the equality is by Lemma D.16. We will bound the second term
in two parts, En/311 and 1=n/2 separately. For the first half of the summation, we have
n/2-1
bymat {( {(J
12())1-1) R (hh)) 0
)
n/2-1
on (hh) )r} bamat ((J
since J --hh < J, entrywise
n/2-1
bamat
J ]
since (1 (1-x)" 11/(2z)) 0 < x <
ny
t=
n/2
b. mat(v}
nry
1
II (I-7HH)
(To 1'))) ny
mat(v)
by Lemma D.22
7
3 -1, (To - ))))) (
For the second half of the summation, we have
bamat {
) c (hh
02(n-1-1)
vv)
2 - (I, (To - r
Σ mat (J- 22(n-1-1) ( o (hn)7)
n-1
)v} }
ny
t=n/2
by the induction hypothesis
2 (1, (To - r
R-1
Σ mat { oon-1-t) (hh ((hit)y) v)
}
ny
t=n/2
since 0 JJ J-J- --7h" T < J. entrywise
2 [I, (To - mat/ mat { (J (J --hhī son/21)
ny
40

---

Published as a conference paper at ICLR 2024
2 (I, (To - -mat))) since J - (J - -yhh" ) < J. entrywise ny
0(n/2)
Bringing these two bounds back, we have
B, mat(v}
1 (I, (To -(-))))) 1 _ + ny
2 (I, (To rT)3) -
(B.) mat(v) + 2
7 mat (v]
ny
-ga( o mat [v) + 3yeemat(v
(I, (To - r
82 mat/v) + 3ycmmat(v
" (To (To-
(I. (T. (*)))
ny
where the last equality is by the definition of in 23
Based on the above, we have
bn = (H, B3 o H)
H. Bn mat(h))
H. (I (I- (H- (TT (To (To-T*) -rr) mat + 3rc (H, mat(h)) 1 II, (To rr))) ny -111 HI- HH, (Tu r + 370 H, H) (I, (To - 17)))
since To and I' both commute with
2ny I, (To (To- rm) + (x)rtr(
[I, (To 1
my (I, (To r*)),
since 7 ((
which completes our induction.
D.7.2 DECAYING-STEPSIIE CASE
We first show crude bound on the bias iterate.
Lemma D.24 (A crude bound). Consider the bias iterate
24
Suppose that
601 tr (eitr(H)ta (H)
Suppose that
Bo (To ) 82
and that To commutes with H. Then for every 2 K, have
b,
lo H H (To (To- fo-r"y -T*
K%0
Proof. Let
K = T/lo(TT))
L= logT).
According to in the first epoch, i.e., = 1, ... K, the stepsize is constant, ie., No- Therefore,
we can apply Lemmas D.2T and D.23 and obtain
bx < min {
(1 (1+ 2c (H) bo. Kmo
(I, (To - -r")3}}
< 22min (H. (To - -T - -T))
1
(I, (To -r))) }
K
41

---

Published as a conference paper at ICLR 2024
2 min
(Hi, (To - ))))) K'0
(I, (To - -r93}
2
min
HH,
Kno" J. (To - -r-)²
- 2
K70 lo:k + HA Hasoo (To - )
Next, recall that the stepsize schedule is epoch-wise constant, therefore we can recursively apply Lemma D.21 for epoch 2, 3...... ..... L. Suppose 2 K belongs to the L -th epoch, then we have
L' II 1=1
+ 201 Tu 28
tr (H)tr(ff)) bk
+
201 no (H)((H))(HH) bx
2bg-
We complete the proof by bringing the upper bound on bg.
Theorem D.25 (Sharp bias bound), Consider the bias iterate 24
Suppose that
1
? 6.8 (H)rr))
Suppose that
Bo (To r
and that To commutes with H. We have
II (I HH) (Fo
bt < HH.
+ +8.38 37 40
losk + H Ha:oo
(To 1m))) (To-)))))
1
Σ
min {1, 12
Koj
Proof. From
31
we have
t
Unrolling the recursion, we have
T
T-1
T
Br
II
B
4
2b
II
o S(1)
1=0
k=t+1
T-1
9
Bo
01
72
II (8 ))S(1) 02 s()) by Lemma D.15
which implies that
by H.. By
H
T-1
T II k=
<
H,
S
Bo o H
+ C1 Σ
7 bbb H.
(8 n2 +2 sws) o A
H.
g
B H
[=1
2--1
H. mat (
T II
+ C1
72 26,
- (J 0²
@ (hhT) ) h (n))
by Lemma D.16
k=t+1
42

---

Published as a conference paper at ICLR 2024
T
H
9
Bo
o
H
T-1
T
+
7t
b
(J-"" 02 @ (hh))) h.
-
(39)
1=
k=t+1
For the first term in
39
using the assumption that To commutes with H and the definition of o in
23
we have
T
T II (1- (To HH)
(To -rr) @2
H
m g
Bo
H
H,
H
1=1
t=1
HI, II (I n Hi) -rr)
rr) )
(40)
For the second term, we will bound Er= and Σi= separately. For the first part of the sum, we have
b. T
II (J
m 1J-hhhh 1o2 の (hń]) h
t=0
k=1+1
K
2K-1
S
bbhh" T
II (J
7 hh ) 1020 ① (hh h
since 1-hhh" S J, entrywise
t=0
R=
%
Σ bt hT
((J
70 hh
62K
① (hh h.
stepsize is epoch-wise constant
Next, notice that
Kx Kv - (1 (({{{11 x)) S
for <<< 1,
(1 - 2K
K2
11
So have
70 hh J
② (hh S min
4
(hh T (0-2)
J + (hh
2
K2
min
4
(hh
T) @(-1)
hh }
K2
where "min" and "S" are entrywise. Then we have
hT ((1 (J- 70 hh
62K @ (hh
(J1-220
h < hT
min
K2 (hh (h 0(-1) hn ii
2
min {1. K2 22222 12323)
K2 2
where the last inequality by the same argument in the proof of Theorem D.20 Bringing this back, we have
n h T
II (J
ya hh
)
02 c
h (hh
kmt+1
K
bt h T ((J
62K
2
hh
(hh"))
h
10
2
1=0
K
4
Σ bt
min {1, K² {2222373)
K
43

---

Published as a conference paper at ICLR 2024
- [ (I- HH)
1
(To 2
4
min [1, K² x2
by Lemma D.22
70
K
1
4
(uu + +2K Ko H (To - -
min 11, K {3
70
K
1
< 8
Iock + Hino Hisoo
(To- (T)))")
min {1, 1310
(41)
Koo
K
3d
For the second part of the sum in 39 have
II (J Amr+1
whh 02 o
(hh") h
<
II (J m-hh")
2 (hh") (hhi))
since J - < entrywise
t=k
k=(+1
T II Am1++
I
Lo:k + Ho Haooo (To To-ry r²
Σ t=
(J (J- wh ) o (hh h,
K
where the last inequality is by Lemma D.24 Notice that the sum in the above display is equivalent to the sum we encountered when analyzing the variance error (see (35) Lemma (D.8)), with the only difference being that, here, the sum starts from the second epoch. Therefore, by repeating the arguments made Lemma D.17 and Theorem D.20 (replacing "o with 70/2), we have
T
Σ
II k=++1
(J -8-70 hhhhhh 32 (hh (hi7))
8
Σ min {1, x3
K
t=k
i
Bringing this back, we have
T
Σ
II k=++1
(J "-hh"
02
(hh)T)
t=
1 Kruu
Io:k
H Ho
(To rT72)
8
Σ
min {1, 12001377)
(42)
K
Finally, putting (40). (1 and (42 in 39
we have
br <
Hi. [I
(I
Hil
(To rr)
1=1
+ + 37 40
Io + H H
(To - (o-r)))
1
Σ
min {1, K2781223).
Kyo
K i
which completes the proof.
D.8 PROOF OF THEOREM
Proof of Theorem
4.1
It follows from Theorems D.20 and D.25
D.9 PROOF OF COROLLARY
Proof of Corollary
4.2
Under the assumptions, we have
tr(H) + a //22
N
N
So we have
tr
//2
+
(
F =
H
N
N
44

---

Published as a conference paper at ICLR 2024
and
33 = 22>) (
tr(H) + 02/102
N N
N X1)
>u max
N A1 s
i 8"
\j No
> e,
where we define
min 120: N
The excess risk
p9 contains two terms.
The first term can be bounded by
Error:
HHN, N
[[[-~HHw)]
1=1
tr
To HHN (Fi)²
Σ
Tarmdidi ded (1
1
t/g A
N
1
-
Σ
N2
e
N
ise
i26
Σ TTarrod \ + Σ
2Tuuro N.
(43)
iS
x222
The second term is
Dott
Dott
Error = tr)) + +02)
~
(44)
Tors
Tors
Define
k: (i.5): A123 2 Torr
(1.5): s SE,A1 2 To U(u00)022 e, 2 Tore }
then
Σ min min {1, (Torrodzdz))
Dott
=| (T
Σ (2x2g)2
(45)
(1J))
The uniform spectrum. So we have that xn = 0 for j > and
Here, we assume that X = 1/8 for i < 8 and A = 0 for i >> 8, and
N Sssdd.
1
1
for j S s,
Ày 2 A max
M
N
sN
Therefore
tr(H) N
45

---

Published as a conference paper at ICLR 2024
and
1
= N.
tr(H)
By 43 we have
N
Error S 86.
2I 7N
8²
= -2- 1o N
8
N
Tort S 5²
Ns
Tett > 8²
Tott
By 45
we have
Datt = Σ min {1, (Tarrioddds)
1
2mmn
8² min
1, Tott?0 s2N
= s² min
1, (Tett g2
T 8²
Tort S ;²
Tort > s²
So by 44
we have
Tott
Tott S s²;
Dott ~
82
Error ~
Tott
32
Tatt > 8.
Tott"
In sum, we have
EA(TT) = Error + Errory
Tott
N
Tort S 82;
82
8²
2A
s²
Ns
s²
+
Tett > 2.
Tort
Tott
Tort
The polynomial spectrum.
Here, we assume X; = =i-a for > 1. Then
=
and
20
is và;
O-NN-1 j j> No.
Therefore
tr(fi) = Σ 2, =1,
and
70
1
ti(
46

---

Published as a conference paper at ICLR 2024
By 43
we have
Error S
Σ
27
Xa +
Σ
TTarro v N2N
Σ
xi 2 N
iSN
i
Σ
2o 30
Σ
2T 20 3 20
e
No TooiToo
Σ
2oa 2 N
i> i>(aaro/N) A
2Tarto i-
Σ
Harrys -2 2Tatt Not
N2
N
2T 20
iS
N Ni
Σ
N
1-2n
N
Tatr
Tecrio / (c
N
N
N
N
1-
N
To
Tott70
N
Tott
where the last inequality is because
70 =1
and the assumption
N3 o(Tott).
The first part in 45
Torr }
N
K
(i.5):
+ (i.))
0, (1:)):>0,AM))
Torr
Ku.D:ssN N=,
is² < (Tatt
(i.)): Ni, S
(Teerio
11
N
(Tet770)
(Te)rYo/N)
Σ
j2
15JSN1
Nà <jS (SS(Torr
Tott"o0
Tote N
(T +
log
N
The sum in the second part in
(45) is
Σ (2 (OOE
Σ
(2 +
Σ
(2r2g)2
3 (Tar)o)))))))
N2N (TTTTan
Σ
(112))2
+
Σ
(A,2//N2²
ieNt (Teno)))
juNt (Teroo/N Σ
Σ
1-2a
i 20 2aa--2N N-2
iSN a (Te
j>N Ai (T))
Σ j j-
Σ
2a + Σ
22
Σ
i 20
JSN
J>N
47

---

Published as a conference paper at ICLR 2024
5-m (TTo(3300)) 1-20
JSNU
- max {I, ((Tatroo/N)+15) 1-2a
+
j 20 N
N>Na
Σ j
2 (Ta
1 +
E
N-2N-2
JSN
1-20
Σ
j No
Nà <JS(Farro/NN)
1-2a
(Tott
1-20
To
N-2
N
1-20
Yo N
Σ
j-NN N-2
Nà <JS(Taro/N))
(Tatt
1-2a + (Tatt70)
1-20 N- (Tatt
1-20 N-2 loog
Torr
(Tatt 1-2a
+ (Tott"o)
1-20 N- log log
To
N
So the effective dimension 45 is
Datt |K + (Tatr7o)2
Σ (2:21)²
(K))
Tote70"
Tacr
(Taerro)) +
log
N
N
(To (7)
1-2u
+ (Tart 7o)
1-20
N log
(Tor
N
To N
Taet N
(T (Tat
log
Therefore 44 is
Dott Torr
Error ~
Tott
(Tot 7o
x T- (Test%)) +
log
N
N
To (1+ N log(Tor:)),
where the last inequality is because
To x 1
and the assumption
No o(Tot)).
Putting the two error terms together, we have
EA (T)))) & Error Error
m N T (1 N l
Tott
T
(1 (1+ N log(Tatt) + T 2a 22-22)
48

---

Published as a conference paper at ICLR 2024
The exponential spectrum.
Here, we assume X; = 2- Then
8 = =log(N),
and
2 2j
S log(N); log(N).
2 3 N
Therefore
tr(H) = Σ ر3 = 1,
and
70
1.
tr(H)
By
43 we have
Error
A +
2o x A,N
i51
Tou
31 2-
-Tuur702 2/ 2-22 2-
iSlog(N)
(N)
34 2-1
iSlog(N)
Tourrn222 2T Toott7022 7022-1
N2
E
N
No
log( MoggeeeeoogleTrrrooooo
Σ
2Ta N 2-12222 2-2N
c 20 8 2-
2Tet [ (c-+))dt
N
iSlog(N)
2-2i
N2
N N
2T
2T
N2
Tott
where the last inequality is because
% ~ 1
and the assumption
N oTatt))
The first part in 45 is
111 (1:00)))>> E, 2
N Tose
}
K
(i. (i.1): A 2 Tors
(1.5) :j < log(N)) : 2j log(Test%))
+ |K ((.0)):;>>0o/ log( S log
Tatr70
N
Σ (log(Teer?oo) - 22) +
- 5)
(S)
(og(N)
log(Tatt7o) (N) + (Tg2(TaerTo//N))
49

---

Published as a conference paper at ICLR 2024
The sum in the second part in (45) is
Σ (A:2g)²
Σ
(Ar2,)2+
E
(2:2,)²
j>log( M)
Σ
2 2(1+2)
Σ
2 2-((tt))
+
( 20)
Σ 2 2-
4j
Σ
2 2-
+
Σ
2 2
2 22
(Ter
{>log(N)
Σ
2 2-1 (Tesf? -2241 (222))))-224)
jSlog(N)
Σ
2-2
Σ
2-2
Σ
2 2-2N
2 2-22
12 (Tae
2 (T) Plog(N)+ +
Σ
2 2 (Tatt70/NN)
log(N))i<log(Turro/
E
2 2-21 N-2
12 (oat
(Tatt log log(N) N (Tatt 70/N) 1og(Tatr70/N) + N (TTT
(T) (o (log(N) + log log % 70/ No))
So the effective dimension 45
Dott [ (Tott
Σ (Me2,)²
(1.))K
log(Tate70) log( (N) log² (T (oN) (log(N) log(Tate NN))
log(Ter:70)) log( N) log² (TerrTo/N).
Therefore is
Dett
Error
Tart
Tell [log(TTet log(N) (ooe))+oooTTTaete (Toerro/N)
T log (Test),
where the last inequality is because
7a m
and the assumption
No o(Tot)).
Putting the two error together, we have
EA(T)) $ Error; + Error
N2
T-1 log (Tatt)
Tott
N2 + log" logr(Terr)
Tott
We have completed the proof.
50

---

Published as a conference paper at ICLR 2024
E A COMPARISON BETWEEN THE PRETRAINED ATTENTION MODEL AND OPTIMAL RIDGE REGRESSION
E.1 PROOF OF PROPOSITION 5.1
Proof of Proposition 5.7 We start with
We have
[h;; X) E (h y, x) v)² 2 X
(h y, -ank)--Evil E[y y. x)) 2 X] + (,, y)² x].
where the second term is independent of h. Therefore, the minimizer of must be h(X, y, x) = Ely y. x].
Recall from Assumption that
y ~ N 3, o²), B,
so we have
y(X, = =E[x 3|X, y, x
(E|B|X, y),
By Bayes' theorem, we have
P(y X, ()P
y) Ply X, (B) aB
Recall from Assumption that
MXX3, I²I). B~N(0, I))
so we know P) y) must be a Gaussian distribution and that which implies that (because the mean of Gaussian random variable maximizes its density)
y) x IP(y) ()
|ly
1311
||
exp
exp
2
2//²
yyXp113
||
E| arg min
22
2
(x x + 02/ -1 -1x
Putting everything together, we obtain that
4(X, y) = (1113) v), x)
+2 -T,, x).
which concludes the proof.
E.2 PROOF OF COROLLARY
Proof of Corollary 5.2 Let G be the sampled task parameter and f be the ridge estimator in 12 that is.
e xx x + 0 yTy.
By Assumption
I
we have
yi N Xi, 0²), ~~~ ~(, -~ B~,,0,,''
which allows us to apply the upper and lower bound for ridge regression in Tsigler & Bartlett (2023). then we have that, with probability at least (((M) over the randomness of X, holds that
o 0²/1²2 Eisk. Ao
Esigo || - - =
||
1 + ||3 Ha
M
51

---

Published as a conference paper at ICLR 2024
o²
k- +
( + Eish
M
Σ
12
M
13
where Esign refers taking expectation over the sign flipping randomness of [
and
o² 0²/102 + Eisk d
km min k: k
M
where e > 1 is an absolute constant. Now, taking the expectation over the Gaussian prior of 3, we have
X(;; - 0² = 3
0²/ ? + >isa' do
2
2/2
+
2
Σ
x
M
X
M
+
K* +
Σ A
M
02//022 + Σ
022
iin
12
Denote
a²| 112 + Σ bi* A 0²/10²
=
M
M
then we have
== : As 2 X}
so we have
C(h: 0² ~
2
+ 2//
Σ
x +
k* +
Σ
x
A
M
X2
isk
i>
X2
min
X
+
Σ
min
A4
M
X2
13
min
di }
+
min
1,
Arp
X2
X2
X²
min
+ min
}
2
Σ minfă, Ai}.
This completes the proof.
E.3 PROOF OF THEOREM 5.3
Proof of Theorem 5.3
Consider the attention estimator
10 and its induced average risk
11
we
have
EL(/:X) = E
x,
I M
X y) -
RM (Ti))
Therefore, we can apply Theorem
3.1 and obtain
EC(ÁN::) = RM (TT)
= (H, (T - TiM H
H (Ti - TimT) min RM
(HHM, M, (ry - rw)² 2) - min Rar(.).
For the second term, we have min RM((.) 2
52

---

Published as a conference paper at ICLR 2024
+M (M+ vr - (0/00))++/0)) H)
-
ut ((tr(H)++ /0/1/) ////////------- (M (M+ 1M+ (2(tr(H) + o + w1)
(tr((H)+ +22 + M1) +
tr
tr (H) + o
= 2M E / +
20p² Σ min {dM, Aa }
where we define
tr(H) + o² //²
o² \//²
AM
M
M
For the first term, note that
Fly - F
tr(H) + 0² /112
N ++1
-1
(r((H) + 02 1/222
M +
-1
N H
+
N
M
M
1
(xr(H) + +02//// /²/
M
N
tr(H) + 0²/ 2//662
N +
tr tr(H) + 0² //2
M + 1 M
-1
I
H
I
H
N
N
M
= -AM - 2N)
ANII+
N H
M +
H
N
M
So the first term can be bounded by
(HHM, (TN - --i))
2/ H2 tr( + 02/0/2
1+
M +
H).
(Ti rir)2)
M
M
w (Aar 2N)²
H2 I
M+1 H
(INI+
N H))
(SuI
M
H
M
N
M
(AM 2N)² tr H2 1H2(NNI+ (ANI+ + H)
(0))"(((n)) ()MI=H) -
(AAr - ÅN)² 2 Σg x? min
1
1
x2
X
Asr
X
~ v (AN åN)² 2 ΣΣ min
A
A
min
Am
Putting these two bounds together completes the proof.
E.4 PROOF COROLLARY 5.4
Proof of Corollary 5.4 Under the assumptions we have
UMM
M
We first compute ridge regression based Corollary
5.2
53

---

Published as a conference paper at ICLR 2024
The uniform case. When X; = 1/s for i and X; 0 for i >> 8, we have
[(h; X) 0²
min {uM, Xa }
min 1
M
min 1, M
The polynomial case.
When X; i- for a > we have
[(: 0² min {ur, d1
Σ min 1
i-
.
M
The exponential case.
When A, = 2-,, we have
[(:; o² ~ 2 Σ min (uA 2z}
1
min
2
.M
(Mg
M
We next compute the average risk of the attention model based on Theorem 5.3
Notice that
- (MM-MNN2 ~
M if < N/c for some constant >
M
N
The uniform case. When A;- = 1/8 for i S and A 0 for i >> 8, we have
EC(;:X) d² ~ 2 Σ min {War, Xi}
+ 2 (uM uN) 2 Σ min
A
A
min
1,
s
+ M
min
-
N22
min { M 1}
M
N
min
8 M
min (5) M
M
M
Y=1
82
8 8 M < N/c:
M
min ( 1
8
8
M M<s
M
M
N²
M <N/c <s
sM
$ M < N/c;
M
8
M < < N/c
M
N2
M < N/c<5.
SM'
54

---

Published as a conference paper at ICLR 2024
So when < M, or > N2/M, we have
{ (M' 1
EL(f:X) - o² ~ min
The polynomial case.
When A; = for a > have
X o² ~ 2/² Σ min {MAE, )o}
+ ( NN)
Σ min
A 1 min {
d (MM"
}}
X f
+ M
Σ
min
{ a N22 2a)
min {iMM, 11
N2
1 M
~
min
min
M
M
N2 M
Σ
Σ
M M
M
M
Mb c
N
~ M + .
N M
+ M
2 -
The exponential case.
When A, = 2 we have
EL([:X)- - 0² R 2 Σ min {NAS, 2z
+ 26 (HM - mx)² Σ min
Ax
A
min
MN
MM
log(M)
+ M Σ min {2
2
min
{2'MM,
'M,
1}
M
log(M)
min
2-
2
N2
2
1
min
2
1 M
M
M
M
log(M)
2
1
Σ
2 1 M
2-
M
M
M
(M
log( lgg(M)<+Slog(N)
2 .22
M 2
S>oog( N)
log(M)
log( N)
+ M
M
M
M
log(N)
M
We have completed our calculation.
55


Reviews:
[{"review_id": "t8hIow4NyO", "reviewer": "Reviewer_ozh2", "summary": "The paper investigates a single-layer linear attention model, which is equivalent to a linear model with parameters derived from a one-matrix-step gradient descent from the origin. The authors make Gaussian data generating assumptions and consider the population ICL risk, identifying the optimal step size. They demonstrate that applying gradient descent to the step size parameterization can lead to an excess risk characterized by an exponential decay plus a 1/T-like decay. The authors also compare this with the Bayes optimal estimator, analyzing their respective risks.", "questions": "- Could you provide more justification for the pretraining algorithm assumed in equation (6)?\n- Could you explain the choice of step size in Theorem 4.1?\n- Why is it reasonable to assume that the initialization Γ0 commutes with H in Theorem 4.1?\n- Could you provide more insight into why H can be assumed to be diagonal WLOG?\n- Would it be possible to include some numerical results in the main body to support the theoretical findings?\n- Could you clarify the term \"number of contexts\"? It reads like the number of in-context samples, but I guess you mean the number of sequences/datasets.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper's originality lies in analyzing the single-layer linear attention model with its one-step gradient descent parameterization. The authors' identification of the optimal step size under Gaussian data generating assumptions is a valuable contribution. The quality of the paper is evident in the rigorous mathematical analysis.", "weaknesses": "The paper's clarity could be improved. \n\n- The assumption of the pretraining algorithm in equation (6) seems arguable, as the equivalence of the function classes does not necessarily imply identical behavior under different parameterizations in gradient descent. \n- The paper is heavily reliant on mathematical notation and could benefit from more intuitive explanations to aid understanding. \n- The choice of step size in Theorem 4.1 is not adequately justified, and the assumption that the initialization Γ0 commutes with H is not clearly motivated. \n- The assumption that H can be diagonalized without loss of generality (WLOG) is also not sufficiently explained. \n- The paper could be enhanced by including some numerical results in the main body. \n- The terminology used is occasionally confusing, such as the use of \"number of contexts\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qVGZfHJ7BU", "reviewer": "Reviewer_HbBb", "summary": "The paper dives into the statistical understanding of pre-training of a single linear layer attention on in-context examples generated with Gaussian data and linear regression with a Gaussian prior.  The important theoretical contributions can be summarized as follows.\n\n(a) The authors first show the optimal solution and its excess risk in terms of the pre-training sequence lengths and the covariance matrix underlying the Gaussian data.\n\n(b) The authors then show the behavior of the excess risk with training steps via SGD, where each step uses a randomly sampled sequence from the underlying distribution.\n\n(c) The authors further show that when evaluation sequence lengths match to the training sequence lengths, the model matches the Bayes optimal solution. However, discrepancies between the two can lead to sub-optimalty. \n\nOverall, the paper takes an important step toward statistical understanding of the dependence of pre-training data and in-context abilities of attention models.", "questions": "Please see my questions in the previous section.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The main strength of the paper lies in its clinical approach to relate the existing literature on ridge regression with Gaussian prior to the statistical understanding of linear attention pre-training. Furthermore, the authors introduce novel techniques like operator polynomials to solve order-8 tensors that show up in the risk analysis of SGD training, which might be of independent interest to the community. \n\nThe advantages of the theoretical framework can be summarized as follows. First, a single linear attention model reaches the optimal linear regression solution with SGD. The framework gives statistical convergence bounds with dependence on training sequence lengths and the data covariance matrix. Secondly, one can pinpoint the gaps between evaluation and training based on the discrepancies in data properties e.g. sequence length. Overall, the paper will be an important addition to the theory community.", "weaknesses": "Overall, the paper doesn't have many pitfalls and issues as is.\n\nI have a question about the theoretical setup. Linear attentions used in practice have a query $Q$, key $K$, and value $V$ matrix. However, the authors use structural modifications in $Q$, $K$, and $V$ to represent the formulation with a single matrix $\\Gamma$ and compute SGD convergence of $\\Gamma$. Without the modification, I believe the optimal solution can be shown to be a $3$-matrix factorization of the optimal solution $\\Gamma^*$ given in theorem 1. But what will be the statistical bounds of training these $3$ matrices (or any pair among the $3$ matrices)? Can the authors discuss whether it is answerable from their theoretical framework and if not, the difficulties one might face to solve?\n\nFurthermore, how will the theory change when instead of using training sequences of a single length, we use randomly sampled training sequences with varying lengths? How will the excess risk in theorem 5.3 change then? \n\nThe authors also conduct a few experiments in the appendix on a real-world transformer to verify their theoretical claims. It would be interesting to empirically check the training time convergence with different singular value behaviors (as they pick in corollary 4.2) and observe differences in convergence and ICL performance throughout training.\n\nFinally, I haven't looked deeply into the proof, since it is extremely long to read through. But at a glance, the paper seems to be utilizing similar proof techniques that prior works have used for ridge regression with Gaussian prior. Hence, I still recommend strong acceptance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "18Hs2MtTK7", "reviewer": "Reviewer_m56U", "summary": "This paper introduces a statistical task complexity bound for a one-layer linear attention model in solving fixed-length linear regression problems. The authors investigate the required number of independent sequences or tasks necessary for pretraining the model. They conclude that only a small number of sequences are needed compared to the model parameters, and this is independent of the input dimension. Additionally, the paper theoretically demonstrates a performance drop when the number of in-context examples during evaluation differs from the training phase.", "questions": "1. Is it possible to ensure that the first term of Equation (9) will consistently reach zero? Given the exponentially decreasing learning rate $\\gamma_t$, this relationship does not seem immediately apparent.\n2. In Theorem 4.1, the learning rate $\\gamma_0$ is upper bounded. Could the authors elaborate on the reasoning behind this? Is it connected to the implicit convergence rate of gradient descent in linear problems? Additionally, what are the considerations and trade-offs in selecting an appropriate initial learning rate $\\gamma_0$?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-organized, with clear contributions and promising results in the statistical analysis of in-context learning for linear regression problems using a one-layer linear attention model.\n2. The theoretical results provided in the paper address the sample complexity of learning linear regression problems, highlighting that it is approximately $O(1/T)$. This indicates that the number of tasks required for learning the model is independent of both the model size and the input dimension ($d$).\n3. The paper underscores the significance of context length during testing, proving that optimal prediction is attainable only when the number of context examples matches that of the pre-training phase.", "weaknesses": "1. While the theoretical analysis presented is commendable and contributes to our understanding, the paper could be greatly enhanced by including more empirical results. Given that implementing linear regression over a transformer model is feasible and has been achieved in prior work, specific implementations that could improve this paper include:\n\n    - Demonstrating through examples that training models on linear tasks of varying dimensions yields similar performance, supporting the paper's claim that task complexity is independent of the task dimension $d$.\n    - Conducting experiments to show that a pre-trained model performs poorly when there is a significant discrepancy between the in-context length during testing and training.\n2. The paper’s assertion that task complexity is independent of $d$ is counterintuitive. It is typically expected that the number of tasks required would be on the order of $O(d^2/T)$ since the task distribution or covariance, encompassing at least $d^2$ parameters, needs to be retrieved. The authors are encouraged to provide additional clarification on this aspect of independence. Furthermore, the experiments presented seem to contradict the paper’s statement, as $T$ appears to be exponentially larger than $d^2$, which doesn’t align with the claim that \"$T$ could be much smaller than $d^2$\".\n\n3. Despite Equation (9) not explicitly showing dependence on $d$, it seems that $d$ is implicitly involved in the matrix operations, which should be addressed for clarity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "icdODKUceu", "reviewer": "Reviewer_dEpy", "summary": "Using a modified version of a single-layer linear attention model, the authors derive a dimension-independent complexity bound which suggests that efficient pretraining is possible even with a large number of model parameters for effective in-context learning. In the process, the authors demonstrate novel techniques for analyzing higher-order tensors that may be independently applicable.", "questions": "1. Could the authors confirm if the restrictions are significantly restrictive? It would appear significant at face value. I certainly don't discount the fact that the work aims to provides theoretical grounding for similar observations in literature.\n2. In Theorem 4.1, how is the inner product between matrices defined?\n3. It would appear that the theorems do not necessarily say anything about misspecification, except that the excess risk is controlled via terms in Equation 8. In a sense, the Gaussian assumptions are used for derivations so anything in the Gaussian family would qualify for the bound. Is that the correct assessment?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The authors provide a timely contribution to an important phenomenon of in-content learning.\n- I am positive that this work provides a good theoretical checkpoint for future work to build on top of the provided analysis. \n- The text is fairly well-presented and the research community will appreciate it as such.\n- I also specially highlight that the authors have made proper effort to delineate the assumptions behind the theoretical results.", "weaknesses": "- The obvious shortcomings of the theoretical results come from the nature of assumptions that deviate from standard practice - the choice of linear attention and a restricting the structure of Q/K/V matrices. Also see Question 1.\n- While the community is focused on Transformers, I am wondering if the analysis also holds for a different linear parametrization of the function $f$. In the broader context, if similar results hold for another parametrization, then attention would not appear that unique of a function. As an example, imagine a different parametrization that also leads to a dimension-free bound. In the broader context, such a result would then not distinguish what attention brings to the table, when in practice we have ample evidence that other parametrizations don't carry as flexible and generalizable inductive biases.\n- Regarding the distributional assumptions of the fixed sized dataset in Assumption 1, it would be great to have experiments where the model is misspecified. Please correct me if I missed, but I don't think I see such an experiment in Appendix A. Misspecification is really important, since for all practical purposes, our models are misspecified, i.e. the data does not really come from the distribution we assume to be. If Transformers can still achieved a good decay of the empirical risk as the number of pretraining tasks increase, it would certainly be a unique characteristic. Also see Question 2.\n- I would strongly recommend highlighting the experiments and moving them up to the main text. Perhaps Section 6 could be compressed a little to accommodate.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)