Learning and Transferring Sparse Contextual Bigrams with Linear Transformers
Authors: Yunwei Ren
Zixuan Wang
Jason D. Lee
Venue: NeurIPS 2024 poster (2024)
Keywords: transformer
sample complexity
optimization
Primary area: learning_theory

Abstract: Transformers have achieved significant success in natural language modeling because of their exceptional capabilities to combine contextual information and global knowledge, yet their theoretical basis remains unclear. In this paper, we first propose Sparse Contextual Bigram (SCB), a natural extension to the classical bigram model, where the generation of the next token depends on a sparse set of earlier positions determined by the last token. We investigate the training dynamics and sample complexity of learning SCB using a one-layer linear transformer with a gradient-based algorithm. We show that when trained from scratch, the training process can be split into an initial sample-intensive stage where the correlation is boosted from zero to a nontrivial value, followed by a more sample-efficient stage of further improvement. Additionally, we prove that, provided a nontrivial correlation between the downstream and pretraining tasks, finetuning from a pretrained model allows us to bypass the initial sample-intensive stage. We also empirically demonstrate that our algorithm can outperform SGD in our setting.

Learning and Transferring Sparse Contextual Bigrams with Linear Transformers
Yunwei Ren"
Zixuan Wang" Princeton University
Jason D. Lee
(yunwei.ree) ren, wangzx, jasonlee) @princeton.edd
Abstract
Transformers have excelled in natural language modeling and one reason behind this success is their exceptional ability to combine contextual informal and global knowledge. However, the theoretical basis remains unclear. In this paper, first we introduce the Sparse Contextual Bigram (SCB), a natural extension of the classical bigram model, where the next token's generation depends on a sparse set of earlier positions determined by the last token. We then analyze the training dynamics and sample complexity of learning SCB using a one-layer linear transformer with a gradient-based based algorithm. We show that when trained from scratch, the training process can be split into an initial sample-intensive stage where the correlation is boosted from stage of further improvement. Additionally, we prove that, provided a nontrivial correlation between the dow downstream and pretraining tasks, finetuning from pretrained model allows us to bypass the initial sample-intensive stage. We also empirically demonstrate that our algorithm outperform SGD in this setting and discuss its relationship with the usual softmax-based transformers.
to nontrivial value, followed by
sample-effice
Introduction
Transformers have played a central role in modern deep learning, achieving significant success across various fields, including language modeling (OpenAI, 2023), computer vision (Dosovitskiy et al., 2020), and natural sciences (Jumper al., 2021). The core of transformers is the self-attention layer (Vaswani et al., 2017), which can attend any subset of the input sequence to output weighted linear combination of the (transformed) tokens.
Several capabilities of the transformers contribute to their success in language modeling. First, they can extract contextual information from the input token sequences, which is essential some arithmetic tasks (Edelman et al., 2022: Liu et al., 2022; Nanda et al., 2023; Yao et al., 2021). In addition, transformers can memorize global in-domain knowledge (Petroni et al., 2019; Zhang et al., 2023a; Haviv et al., 2022: Carlini et al., 2021). These two abilities combined enable transformers to predict the next token based on the in-context information as well global knowledge (OpenAI, 2023) acquired during training.
To theoretically understand how transformers learn both capabilities, we propose minimalist data- generating model, the Sparse Contextual Bigram (SCB). This model builds on the classical bigram model and requires learning both contextual information and the (global) transition probabilities. Here, the next token depends on the transition matrix P and sparse set of prior tokens that is determined by the last token. particular, SCB can be represented by one-layer linear transformer simplified architecture that can serve as an abstraction for studying transformer optimization (Ahn al., 2023), which makes it suitable for theoretical analysis.
'Equal contribution.
38th Conference on Neural Information Processing Systems (NeurIPS 2024).

---

In this paper, we investigate the training dynamics and sample complexity of training a linear transformer to learn the SCB task using stochastic gradient-bass baseed algorithm. Our contributions are summarized as follows:
Data model:
We introduce the Sparse Contextual Bigram (SCB) model, a simple task that requires the model to learn both in-context and global information.
Convergence:
We prove convergence guarantees for a one-layer linear transformer trained on with the nonconvex (--reullarizd MSE loss using preconditioned projected proximal descent, given a dataset sampled from the SCB model.
Sample Complexity: Under mild conditions on the data distribution, initialization, and hyperpa- rameters, we prove that our algorithm on the sequence length T, number of states N, and the sparsity parameter Q < T. We show that the training first goes through initial sample-intensive stage which boosts the signal with poly(?) samples, followed by a more sample-efficient stage to achieve final convergence with poly(N,Q) Q) samples. We empirically verify that our gradient-basee methods converge to the ground truth with a small batch size, while unregularized stochastic gradient descent fails due to the large variance.
recover the ground-trute truth with polynomial dependence
Transfer Learning: We prove that, when there is nontrivial correlation between the pretraining and downstream tasks, we can transfer a pre-trained model to bypass the first sample intensive stage, that our algorithm converges to the ground truth of the downstream task with only poly(N. Q) samples.
1.1 Related works
Training dynamics of transformers. Several works have studied the learnability aspects of specific transformer architectures. Jelassi et al. (2022) demonstrated that Vision Transformer (ViT) (Dosovitskiy et al., 2020) trained through GD, augmented with positional-embeeddng attention matrices, can effectively capture spatial structures. Li et al. (2023) investigated the sample complexity similar VIT model. Tarzanagh et al. (2023)
necessary to achieve good generalization performance established a connection between the optimization landscape of self-attention and the formulation of hard-margin Support Vector Machine (SVM) problem that separates and selects specific optimal tokens and established global convergence under strong assumptions. Tian et al. (2023a.b) provided insights into the training dynamics of the self-attentin and MLP layers, respectively, although they did not establish convergence guarantees.
Another line of work focuses on the training dynamics of in-context learning. Mahankali et al. (2023) was among the first to introduce linear regression as an in-context learning task, while Zhang et al. (2023b) proved global convergence of gradient flow for a single-layer linear self-attention layer this task. Huang et al. (2023) provided a convergence guarantee for a one-layer transformer with softmax attention similar task where the in-context tokens drawn from specific data distribution. Chen et al. (2024) generalized the single-task linear regression task to a multi-task setting and proved the global convergence of multi-head attention architecture using gradient flow on the population loss with specific initialization. In contrast, our work focuses on the language modeling ability of transformers instead of their in-context learning ability.
Several recent works analyzed transformers from a Markov chain perspective. Bietti et al. (2024) studied the in-context bigram (phrased as induction head) from an associative memory viewpoint. Nichani et al. (2024) proved that simplified two-layer transformer can learn the induction head and generalize it to certain latent causal graphs. Edelman et al. (2024) further investigated training process on bigram and general n- n-gram tasks, and observed multi-phase dynamics. Makkuva et al. (2024) studied the loss landscape of transformers trained on sequences sampled from single Markov Chain. Our SCB model extends the classical bigram models to allow context- -dtexxt-depenent sparse attention on previous tokens.
Several works, including Tian al. (2023a); Zhang et al. (2023b); Huang et al. (2023); Tarzanagh et al. (2023); Nichani et al. (2024); Kim and Suzuki (2024), and ours, use similar reparameterization, consolidating the key and query matrices into a single matrix to simplify the dynamics of the training process. Most previous studies (Tian et al., 2023a; Zhang et al., 2023b; Huang et al., 2023; Tarzanagh et al., 2023; Nichani et al., 2024; Kim and Suzuki, 2024; Wang et al., 2024; Chen et al., 2024) uses population loss to simplify the analysis. In contrast, work goes beyond the population

---

loss to analyze the sample complexity of the stochastic gradient descent dynamics. Although Li et al. (2023) also investigated the sample complexity on a different task, their model requires a pre-trained initialization, while our model is trained from scratch.
Transfer Learning. Transfer learning (Devlin al., 2018) has gained significant attention in this deep learning era. From a theoretical perspective, several works have investigated the statistical guarantees of transfer learning from the representation learning perspective (Tripuraneni et al., 2020; Du et al., 2020; Arora et al., 2019; Hanneke et al., 2023). Recent studies on transfer learning mostly focus on linear models (Li et al., 2022; Tian and Feng, 2023; Fei and Li, 2021; Zhang et al., 2022: Ju et al., 2023; Dar and Baraniuk, 2022). For dynamics of transfer learning, Lampinen and Ganguli (2018) 8) studied the behaviors of multi-layer linear networks in a teacher-student setting, while Dhifallah and Lu (2021) analyzed single-layer perceptrons. Damian et al. (2022) showed that two-layer neural network can efficiently learn polynomials dependent on a few directions, enabling transfer learning.
To the best of our knowledge, this is the first work studying transfer learning for transformers Moreover, unlike previous works that assume a shared structure between the pretraining and downstream tasks, we only require them to have a non-trivial correlation, which is a much weaker assumption.
1.2 Outline of this paper
In Section we formalize the problem setup, including the SCB task, the transformer architecture, and the training algorithm. Section 3 consists of our main results, and we analyze the population dynamics to provide intuitions. Section 4 contains our transfer learning results. Experimental results can be found in Section 5.
2 Setup
In this section, we describe our data data-gennrating model, the one-layer linear transformer architecture, and the training algorithm.
Notations. We use [7'] to denote the set .I.,, Matrices and vectors are denoted in upper-case bold letters (A, V, A, etc.) and lower-case bold letters (a, q. etc.), respectively. For norm, II-Il denotes f2 norm and || |I denotes the Frobenius norm. Additionally, for μ E RN,, || denotes şe-norm for matrix A E ddxN for arbitrary d, which is defined as || == Tr (Adiag( A )) We use 1 (-} to denote the indicator function. We use Ō(-) hide logarithmic factors in the asymptotic notations.
2.1 Data-generating model: Sparse Contextual Bigram
The bigram model. where the next token depends only on the current one, is arguably one of the simplest language models. To learn this model, it suffices to learn the transition probabilities p e RNN where Pn.m = P[Xr+1 = X m), which is achievable through a linear model (0-layer transformer).
natural way to extend the classical bigram model is to allow the next token to depend on a context-dependent set of previous tokens. This extension model situations such generating the words after the phrase "by Theorem 3.2", which requires us to retrieve the statement of "Theorem 3.2". Here, we propose simple extension of this type, which call the Sparse Contextual Bigram (SCB). The contextual information is encoded by sparse probability vector q determined by the last token. To generate the next token, the model retrieves the tokens referenced by q and applies the transition matrix P (global knowledge) to one of them according to the distribution
Formally, data- data-genn model SCB can be described as follows. Let T be the sequence length and [N] the vocabulary. Let P E R XN be a transition matrix, with column P& being the transition probability vector of token k. Suppose that μ R. the stationary distribution of P.
Each input sequence consists of T + I tokens (xi,... xx+1). 1). i.i.d. sampled from distribution μ. The output token (label) is generated as follows. For each & E [N], there is probability vector g e RT 20
that represents the tokens the model needs to attend to when the last token XT+1 1 k. For notational
'Thi differs from the convention of the usual Markov Chain literature where the rows are the transition probability vectors. We use this convention it is compatible with notations.
3

---

simplicity, we will write x XT)), X = (ex..... exy) and Q = (q"")...,, X741 X7+1 the output token x, sampled from the distribution
((N)) E ERTN.
T = =x = Σ 5=1
P(N, =n n
(k) 41
Vn e [N]
(I)
In words, we first sample position s [7] according to (8)) and then run one step of the Markov Chain from Xy to generate Xo- Note that this model can be represented by a one-layer linear transformer (see the next subsection for details). We make the following assumptions SCB task.
Assumption 2.1 (Sparse Contextual Bigram, SCB). In the SCB task, we assume the following:
(a) (Q-sparse) For some Q < 'I' and each of q at most Q entries are nonzero.
(b) (Well-conditioned) There exists some constant C z 1 such that for every k E [M] and E [7], q, (1/(CQ), C/Q) if it is nonzero, and μa CI/(
(c) (Nontrivial transition) [|P]], || 2 ||
(d) (Long sequence) T (NQ)
Remark on condition (c). We say the transition P is trivial if the transition probability vectors are all the same, i.c., P wl", In this case, we have || (n11, , > = ||ull?. Requiring [IPI], - || 2 || rules out situations where too close to the trivial one. Also, note that for any well-conditioned μ. we have pll 2 Q(1/N).
In this work, we focus on the case where (x, Xx) are given as (one data point of) the training data with (x7+....T7+)) ii.dd sampled from μr. The SCB task can be extended to a sequence-to-sequee model: we drop x1 and append x, to get a new input sequence (x2, ..2......77 .T.11 , and then repeat the same sampling procedure to generate another token. This generates sequence (x,) where (x142. ...) are not independent, and this makes our model a true language model. We leave the study of the more complicated learning- (rom-(x,)8 task to future works.
2.2 Transformer architecture
Our learner model is a one-layer single-head linear transformer (Akyürek al., 2022: Zhang et al., 2023b; Ahn et al., 2023). A general linear transformer can be expressed as: F(x,XT=1: V, A) VE (E' AE), where E is the embedding of the input tokens and positions, and A, V are the parameters of the attention and output layers, respectively. In our setting, we only need simpler model:
F F(x, A) VX (LrA) VXa (s7-1).
where Ve E NNN and A E TxN are the trainable parameters, and a (k) denotes the k-th column of A. This model uses cross-attention (replacing the last E with ey.,)), uses only the positional embeddings together with the last token to compute the attention weights (replacing the second E with Iy), and discards the positional embeddings in the output layer (replacing the first E with X). This is equivalent to manually set certain blocks in the weight matrices to O, which a common practice in the theoretical literature simplify the analysis (Nichani al., 2024; Huang et al., 2023; Zhang et al., 2023b).
Note that Gata-generating model (1) can be represented using (2) by setting A = Q and V = P. We will show that a modified version of SGD can approximately recover this ground-truth model.
2.3 Training algorithm
We assume that the stationary distribution and certain norms of the ground-truth P and Q are known when choosing the initialization and learning rate. The goal here is to recover and Q. Our loss function is the fj-regularized MSE loss. The standard way optimize an r--reullarieed loss is to use the proximal gradient descent. We adopt this algorithm with several additional pre-conditioning and projection step to ensure some basic properties.

---

Formally, let the per-sample loss be defined
1 (xx,AT-1. A) Ellexo VX Ae -e
33)
We initialize A = //7 to have uniform attention and V = #1 to be the trivial transition. At each xx. x Br i-1
step T 2 O, we sample B, fresh samples (x(() .x
form mini-batch. The fi-regularized
mini-batch loss is defined as
Br (B,,d) ((x(0) T.I P V.A) rreg
Σ 1(x (0
(1) TT VV.
(K)
Br E=1
A=1
where 0 is parameter that controls the strength of regularization. Let V (B-1 and V ( denote the mini-batch gradients of the original w.r.t. V and A, respectively, We then define the preconditioned gradients
((Br) =
IN
IN T N (v.
diag(1/M) IN
1p
N
||
(4)
(BB)) atk)
IT
1
(v T
wk E [N].
Here, the 1/p rescaling plays a role similar to importance sampling. We multiply ((Br) with 11T/N and I up" /ll to ensure least I V I. 17A = 1 and Vμ = μ always hold throughout training. Note that we project each column of V to the affine space (r E R 1 1) instead of the probability simplex. This is sufficient for our analysis and is much easier to compute than the latter. We update the output layer using
V++1 Vr 7/ (B
(5)
where vạ > 0 is the step size. Now, consider the attention layer. Due to the existence of the fr-reuularizaion, the update rule becomes simple variant of the standard proximal gradient descent. Formally, for step size O, each & E [N] and E [T), we have
(k.) r+1
(k)
AA ((B)),
1,
(preconditioned GD step),
Pk a
(k,')
if a (R,')
a 1,44]
d,
2 d,
1,r+1
(k.")
0.
if a (k.')
[,T4]
(proximal step),
(6)
1,7+1
a r+1
(k)
= Proj
a (R."))
r+] = r+1
(R,"),
+
+(1
+1-117
1
(k.") 7+1
(projection step).
(1a=1)
For the proximal step, we will later show that no a, can ever become smaller than -d, so it suffices to consider those two cases. During the proximal step, all small a, (k) set to 0, and a is subtracted from all large coordinates. For notational simplicity, we define g == - (k)
r+1 a (²) so that we
can write the update as a r+1 this case, (6) becomes the usual projected preconditioned gradient descent and we have
(k) (k) 17gT? (k) We will choose 2 in certain stages of training. In
(k)
(k) a 7A ( (when = 0).
Our algorithm consists of three stages with different hyperparameters being used in different stages and certain rounding required between stages. The pseudocode is given in Algorithm and more details on the projectionnnormalizttion steps are provided in Appendix E and When we train the model from scratch, all three stages are used and the initialization is Vo =11 and 1, 1
Transfer learning. When doing transfer learning, the initialization will be obtained from the weights of the pre-trained model and one step of gradient update. Then, will run Algorithm from Stage 2.
5

---

Algorithm 1 Projected preconditioned (1-proximal gradient descent
Input:
(1) (0) Stationary distribution μ. initialization Vo. Ag: learning rates n .y i i [3]: threshold Anš regularization strength å; times Tí, T2, 73
Stage I: Run (5) and (6) with 7) = -7 =1y ,
(1) A= = for Ti steps;
Thresholding-ppooecctoon Ak E [x]] ((k) = =[a, [[) 2 Aa)le. a (k) (17 - 17,,
Stage 2: Run (5) and (6) with 7 A B/y =m
(2) .0 = T22 Ti steps;
Thresholding-normaliiin
Yk E [n], ((k) = 1a, 1a (1) z 201/0)
a a²)) -
á /1 (R)
Stage 3: Run (5) and (6) with 7/4 0, nj
3 y4 d 0 for T3 T2 steps:
Output: A7. V7s.
Results for training from scratch
In this section, we consider the situations where we train the model from scratch, i.e., the initialization Vo = 11" and Ag = 11 /T and discuss the ideas of the proof of the following theorem.
Theorem 3.1 (Theorem G.I), Let & > 0 be our target accuracy and Ti min(T 2 0 max(av, a 2 6(1/(QN))) We can choose the hyperparameters in Algorithm such that within poly(N. 1/E, T) steps, we have ||A 212 s & and I|V Pl & with prob- ability at least 1- and the numbers of samples used before and after Tj are poly(T. 8) and Poly(N, Q. 1/e, log T, 6). respectively.
The overall strategy is analyzing the population process and then controlling the distance between the mini-batch trajectory and the population process?. In Section 3.1, we discuss the key properties of the population process that simplify the analysis. After that, describe the dynamics of the algorithm and the signal-noise-ne (SNR) in each the three stages of Algorithm in Section 3.2-3.4.
3.1 The population process
In this subsection, we analyze the behavior of the population process and the evolution of the signal-noise ratio. More details can be found in Appendix C, where the so-called population projected process are defined and rigorously analyzed.
For case of presentation, we assume . 0 and access to the population loss L := Bl. In other words, we consider the projected preconditioned gradient descent. By Lemma B.8, the dynamics of the population process is controlled by
V++1 V- 7V (I1A122 (V u11) - (). (P- (P-117)) =11)).
Ar+1 A, nA (IIVI] -1111)) (c7 (a()) 1 T (vV.)),. -1pll())) ( (K)
One can prove via induction on r that V (resp. A) always stays on the straight line crossing #11 and (resp. / and @). In other words, there exists some time-dependent real numbers av,r, UAA. ßv.r == 1- av.r, BA. := 1 1- a',r such that V, r fv. and A, a + 11 /T. The same calculation yields the following equations that govern the dynamics of av and a:a
av
@v. = dv.r K (1 a a 1/v
0AAT+1 = -A K K (1 a) av.
where av,0 a = Kp == ||||| - ||| 2 1/N and K2 == |1 1/T R 1/2. Choose
ny = n/ Ko and 17A /7/ Kp, and we can write the above in matrix form
@v,r++ ['A.++1
av,r AA.T K
" (1 (1-ay) av)
7 - RAav)
0
+
0
(8)
2Strictly speaking. what we actually control is the distance of the mini-batch trajectory to the subspace the population process lies. This allows us to prevent the potential exponential growth of the error caused by error compounding in the analysis. For details on this technique, Appendix C.

---

Hence, in order to analyze the population process, it suffices to analyze the above 2-dimensional ODE, In what follows, when we say the signal, we usually refer to these a's or some quantities whose size is proportional to them. In particular, as one can see from (8), the size of the expected gradients is proportional to ay and/or ax²
Note that when both av, a are still small, the population dynamics of av, a are linear system with coefficient matrix 7/ O
o] and drift WT//Ke The drift term will provide a small initial signal that guides the process toward the correct direction and then the linear term will amplify this signal. Since the linear term is close to at initial and the initial signal provided by the drift term has order 1 1/T, we should expect that poly T samples are necessary to distinguish it from noises (Stage 1). After the signal becomes reasonably large, the first term will have order 1 1/poly(N, NN Q), and we can then rely on it (combined with the /1rreuularization instead of the drift term to learn the model (Stage 2).
3.2 Stage 1: boosting the signal
At initialization, we have ary = a = 0. We define Stage to be the phase until at least one of them has grown from 0 to some small 1/poly(N, Q). Note that in this stage, the mini-batch version of (8) is approximately equivalent to
1/
OV ('A,++1
0 [
ay,r 0
7 Ke
1/T
1
0
0 Enoise + Sapprox.
where Enoise and Bapprox represent the errors introduced by the difference between the mini-batch and population gradients, and the fact that we are not exactly on the population trajectory. If we had infinite amount of samples so that both Enoise and Eapprex were O, then the second term on the RHS of (9) could provide a small positive signal to a and the first term would quickly boost it to oi within Ti lTg(T) /n iterations. In order for the above analysis work, we need both Eno se and 6approx to be at least 0(1)/T small. Since, unfortunately, Enaiso does not scale with I/T, we need poly(7) samples to ensure these conditions.
We conjecture that this poly(T) dependence is unavoidable (when only a polynomial amount of computing time is available). That is because around the initialization, the only signal comes from the second term and the first term amplifies whatever the second term provides, even if it has been corrupted by the errors. It either takes poly(?") fresh samples each step reveal the signal or poly(7) steps (whence also (7ly(7) samples) for the random noises to (hopefully) cancel with each other.
3.3 Stage 2:
learning the model
We know that at the end of Stage 1, at least one of av and a is 01 = 1/poly(Q, N) large. Hence, one may expect that the signal is I/poly(Q, N) large now so that we no longer need make the noises I/T small and therefore, only PNly(N, N, Q) samples are needed. Unfortunately, this argument will not work directly, since the variance of the mini-batch gradients scales with T. Therefore, we still need (T) samples to reduce the squared f--norm of Enoise from Q(7)) to 1/ 1/poly(Q. N). To address this issue, we introduce the Ej Gi-reullarizer and use a variant of proximal gradient descent.
The idea is, while the concentration in the l2 sense is difficult, controlling the -a-error is easy as every entry of Vaua// is bounded whence subgaussian. As result, we can make the coordinate- wise difference between the population and mini-batch gradients 1/poly(N, Q) small using only Poly(N, Q) log T samples by a standard concentration argument. Moreover, we have (ef. the proof of Lemma E.3)
Q(1/2). if if # O, O(ava("))
(R)
(x -avaaa)) = KRayKP x E (0) - Hkav (q. di
(10)
if q (R) 0.
Thus, as long as av 2 1/poly (N,
2). the lor norm of the gradient noise being small is enough to create a separation between those useful entries (q, (k) * 0) and useless entries (q. 0) and ensure the fz-error of those @ useful entries is small.
We will often drop the time subscript and write ay == kv,r for simplicity. *It is possible almost explicitly to compute the covariance matrix through some tedious calculation. Intuitively, the reason it scale with T is V / has 7 entries with most of them almost uncorrelated in certain sense.
7

---

The above analysis suggests removing all small entries from the gradient will work. Now, we claim that é erregularization and proximal gradient descent naturally implement this strategy, least approximately. We believe softmax-basee attention layers also automatically implement this strategy. See Section 5 for more discussion on the relationship between our model and softmax transformers.
Note that, at the end of Stage I and after the beaghholdinn-rojeectii step - which is approximately equivalent to running one proximal step first - we know that all useful a; are at least Q(av/()= 1/ I/poly(N, Q), while all useless entries are of size O(1/TT. By our previous discussion, we know that if 2 is chosen appropriately, with poly(N, 2. log I) samples, the gradients w.r.t. those useful entries can be made approximately correct, while the gradients w.r.t. those useless entries are much
smaller than 2. Thus, after one gradient step, the absolute value of each of those useless a, is still much smaller than d. As a result, they will be set to 0 in the proximal step (and to ((ITT) after the projection step), which is equivalent to filtering out all those entries, up to small bias. Therefore, the proximal gradient updates stay close to the population trajectory, and the growth of the signals a'A, dry can be analyzed using the population dynamics.
(k)
We end Stage 2 when (av + az))/22 = 1. Similar to Stage 1, this also only takes 72 = Ō(1n)) steps. We also show that the difference between the mini-batch trajectory and the "population trajectory" can decrease to a small value (ef. Lemma E0) This allows us to decouple the error introduced by Stage 1 and the target accuracy. We defer the proof details to Appendix E.
3.4 Stage 3: final rounding and convergence
The purpose of Stage 3 is to fix minor issue regarding lay a'l- Taylor expand (8) around (a av) = (I. 1) and one will notice that although (av + a 2/2 can converge to I at a linear rate (and the approximation error also decreases exponentially fast), the convergence rate of a av is much slower, and the process will get stuck around (1 + 6, 1 - 6) for some small nonzero 6, instead of converging to (1, 1). To accelerate this process, we directly round A via normalization, which is possible only after the approximation error becomes small in Stage 2. Then we freeze A and train V to the desired accuracy. More details about this stage can be found in Appendix F.
4 Results for transfer learning
The transferability of neural networks and transformers and their benefits have been widely observed and studied in both practice and theory. is often assumed that the downstream and pretraining tasks share a common structure or representations/feauuee and these models can learn these common structures during training, and then leverage them in fine-tuning.
In this section, we offer a different perspective: as long as there is a (potentially small) nontrivial correlation between the pretraining and downstream tasks, the pretrained model can be used to provide a nonzero initial signal, allowing us to bypass the initial sample-intense
stage.
Formally, we consider the following setting. Let f be the transition matrix of the pretraining task and (P. () the transition matrix and Q-matrix of the downstream task. We still assume Assumption 2. In addition, assume P and P share the same stationary distribution μ. || @(1) [|PI] and (P. 2| || The last condition can be viewed as the transfer learning version of condition (c) of Assumption 1. Note that we allow the correlation between É and P to be small as o(1). Theorem 4.1 (informal version of Theorem H.3). Consider the above setting. Initialize = 11TT/r /T, V 0p + (1- 0)w1" for some small 0 > and run one step of gradient update on A. Then, running Algorithm 1 from Stage 2 allows us to recover (P, Q) to 8-accuracy with high probability with PNly(N, Q. 1/g,logT] T) samples.
To intuitively see why using poly(N, N. Q. 1/e, log 7') samples possible, recall from (9) that the reason we need (T) samples in Stage 1 is the signal is additive and has order (1/T), so we need the size of the noise to be at most (1//T). On the other hand, when we initialize V = 0f + (1 0) u11, we > 1/T. Then we can rely on av,, instead of the 1/T-sized additive signal, to boost the signal of a to w0(/T) in one step, which leads to a sample complexity that depends on ay instead of the 1 1/T-sized additive signal. Then, we can reuse the analysis of Stage and 3 to show that the downstream (P. Q) can be approximately recovered using poly Q. log.
have av 2 ((0((NK))
8

---

Note that unlike the case of training from scratch, when performing transfer learning, the initial approximation error || pe i.e., the distance between V and its population projection, can be much larger than the signal dry, and it might seem unreasonable to expect that we can leverage the small signal in the presence of a large approximation error. To handle this issue, we show that the influence of [|Av I| on the dynamics scales with a (* av), which is small. In addition, we also show that as long as a is bounded away from 0 and the batch size is large, the approximation error will not grow. This allows us to ignore the approximation errors the signal- stage until we enter the regime of the Stage 2 analysis.
Experiments and relationship with softmax transformers
This section contains our experimental results. We also discuss the relationship between our linear model and the usual softmax transformers.
Experiment setup We use the same shallow transformer model (2) to train on the synthetic data. The data distribution follows the SCB model (1) with randomly sampled transition matrix together with its stationary μ,, and the ground truth attention pattern Q. We choose the number of states N 3, sparsity Q = 2, and the sequence length
5000 > N. Q. We use a batch size B = 64 to run the online projected proximal gradient descent with 2 le-5 and the vanilla SGD for T 1000 iterations. Through the signal boosting stage T E (0,). we use 1)1 0.01 to accelerate the process. After 172 0.005 for further improvement. For SGD, add another set of experiments with 1/2 = 0.001 to prevent potential instability. For more details, see Appendix I.
400,
5.1 Convergence
Our experiments (cf. Fig. 1) show that after switching to proximal gradient descent after Stage I (the stage), both || Vll and ||A L1l decrease faster than SGD. The final distance to the ground truth after normalization gets close to O, and the similarity between the ground truth and parameters quickly converges close to 1. In comparison, SGD struggles to converge with the same small batch size and large learning rate, while the convergence rate too slow when smaller learning rate is used. This phenomenon verifies our theory that the variance of the original stochastic gradient will be too large for SGD converge when T' > Q, N, while proximal gradient descent with an lj regularizer can resolve this issue.
1.301
Distancc to ground truth
0.25
Distanca to agandd trudh
0.5
0.20
0.6
0.35
0.10-
0.2
0.00
0
2000
4001
600
S00
1000
2001 400
000
800
2000
Iteration
teeaaion
Convergence analysis: We plot the distance to the ground truth VV |lA 21l in different settings. After stage ends at r = 400 (when a av x 0.1), we use vanilla SGD and our proximal gradient method train the transformer. Compared with SGD, the & regularized proximal gradient descent quickly converges, and the final solution (the star) recovers the ground truth. SGD either suffers from the large gradient variance (when 1/2 is large) or a slow convergence rate (small 1y)).
Figure 1:
5.2 Relationship between our model and softmax transformers
We claim that they have our linear model and softmax transformers have qualitatively similar behaviors: there will be a sample-inten inteel-itte initial stage, and after the model and the target have a nontrivial correlation, proximal gradient descent/SGD will become much more sample efficient.
1, write a := (1), and assume the ground-truth 4 := q (1) e; (1, 0). Most of our argument below can be generalized to the general setting at least at a heuristic level. Recall that our linear model is f(x; V. a) = VXa. By softmax transformer, we mean the model for X; V.w). =Voo(w) yaa, where or is the softmax function and R is the trainable first-layer weights.
For ease of presentation, in the following, we will assume

---

Softmax Transformer
Linear Transformer
0.50
linear softmax
0.49
20
0.47
40
40-
50
50-
0.46
60
20 40 60
20 40
60
0.44
20 40
20 40 60
2000
Figure 2: Similarity between the softmax and linear attention. We train two transformers with (1) (Left) sofimax attention and (2) (Middle) linear attention layer on the SCB tasks with the same ground-truth (T 50. N 10, Q = 2). The attention pattern and the value matrix (learned transition matrix) are very similar (left two plots) and they converge to approximately the same loss (right plot).
Let / denote the (per-sample) loss. We have (x)((f(( (diag( (a..) a. ((X))
As result, the dynamics of the attention weights are controlled by
a(T+1)= a(T) 17 1- 11
T (VX) VI(((X),
in our linear model,
a do(T+1)= (T) - (diag(a.) 3) (VX)) VI() in softmax transformers.
In other words, the main difference is that there will be preconditioning matrix (diag(a) -aa T)² in the dynamics of softmax transformers.
Near initialization, i.c., when the attention pattern is still close to the uniform attention, we have (diag(a a0 = T I - ) In other words, our linear model and softmax transformers approximately equivalent up change in learning rates.
Now, suppose that there is a nontrivial correlation between ar and q = e1, say, ao,, is small constant while all other entries are O1/TT) In this case, we have (diag(a - ar)
(1-a,r,))e
O(1/T). Effectively, softmax transformers automatically adjust the learning rate according to a and roughly ignore those positions with small attention weight to stabilize the gradients. Note that this is also what reggulaariztio does in our algorithm. fact, mimicking this behavior is one the motivations of using ( rregularizattin in linear setting. We run further experiments to highlight the resemblance between softmax attention and our linear attention model (Figure 2).
Conclusion and discussion
In this paper, we propose the Sparse Contextual Bigram (SCB) model, which is a natural extension of the bigram model, that requires both contextual and global information. Then, we analyze the problem of learning a SCB model using a one-layer linear transformer and a gradient-based algorithm. We prove quantitative bounds on the convergence rate and the sample complexity. particular, we show when trained from scratch, the training process can be split into two stages, where the first stage uses lot of samples to boost the signal from zero nontrivial value, while the second stage is much more sample-efficient Then, we consider the problem in a transfer learning setting and prove that when there nontrivial correlation between the pretraining and downstream tasks, the first sample intensive stage can be bypassed.
Our data- senargennntting model and results also lead some interesting future directions. For example, can we improve the sample complexity of the first stage? What can we gain if the datapoints are sequences generated by repeatedly applying the SCB model?
Acknowledgement
JDL. acknowledges support of the NSF CCF 2002272, NSF IIS 2107304, and NSF CAREER Award 2 144994.
10

---

References
OpenAl. Gpt-4 technical report, 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghami, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x words: Transformers for image recognition scale. arXiv preprint arXiv:2010.1192 2020.
John M. Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A A Kohl, Andy Ballard, Andrew Cowie, Bernardino Camera-Paredees Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with alphafold. Nature, 596:583 589, 2021. URL https::///.cocc / api semant cschola Oarg/Corpus ID: 2.3 233959967..
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang, Inductive biases and variable creation in self-attention mechanisms. International Conference on Machine Learning, pages 5793-5831. PMLR, 2022.
Bingbin Liu, Jordan Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.10749 2022.
Neel Nanda, Lawrence Chan, Tom Liberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.02211 00521............. 2023.
Shunyu Yao, Binghui Peng. Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. arXiv preprint arXiv:2 2021.
Fabio Petroni, Tim Rocktäschel, Patrick Lewis, Anton Bakhtin, Yuxiang Wu, Alexander H Miller, and Sebastian Riedel. Language models knowledge bases? arXiv preprint arXiv: 1909.0 2019.
Chiyuan Zhang, Daphne Ippolito, Katherine Lee, Matthew Jagielski, Florian Tramèr, and Nicholas Carlini. Counterfactual memorization in neural language models. Advances in Neural Information Processing Systems, 36:39321-39362 2023a.
Adi Haviv, Ido Cohen, Jacob Gidron, Roei Schuster, Yoav Goldberg, and Mor Geva. Understanding transformer memorization recall through idioms, arXiv preprint arXiv:2210.03588. 2022.
Nicholas Carlini, Florian Tramèr, Eric Wallace, Matthew Jagielski, Ariel Herbert-Voss, Katherine Lee, Adam Roberts, Tom Brown, Dawn Song, Ülfar Erlingsson, Alina Oprea, and Colin Raffel. Extracting training data from large language models. In 30th USENIX Security Symposium (USENIX Security 21), pages 2633-2657 USENIX Association, August 2021. ISBN 978-1- 939133-24-3, URL https:tttsc.. www. usenix. wwt.//oofoeeeneeeeoooooo http://oooeereeeeeesseeeeaaaaaat usenixsecurity21/
Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). arXiv preprint arXiv:2310.01082. 2023.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances Neural Information Processing Systems, 35:37822-37836, 2022.
Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015 2023.
11

---

Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv:2308.16888 2023.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in l-layer transformer. arXiv preprint arXtv:2305.16380 2023a.
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00535, 2023b.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.0357, 2023.
Ruiqi Zhang, Spencer Frei, and Peter I. Bartlett. Trained transformers learn linear models in-context. arXiv preprint XXi::2366009927 2023b.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint 14444:240..1944, 2024.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transforeer: A memory viewpoint. Advances in Neural Information Processing Systems, 36, 2024.
Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402. 14735, 2024.
Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains. arXiv preprint arXiv:240..11004 2024.
Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: framework for principled analysis of transformers via markov chains. arXiv preprint arXi::2002.04166 2024.
Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. arXiv preprint arXiv:2-422.01258, 2024.
Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D Lee. Transformers provably learn sparse token selection while fully-connected nets cannot. arXiv preprint arXiv:2446.0688) 068430 2024.
2018. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint
Nilesh Tripuraneni, Michael Jordan, and Chi Jin. On the theory of transfer learning: The importance of task diversity. Advances in neural information processing systems, 33:7852-788 0762,, 2020.
Simon S Du, Wei Hu, Sham M Kakade, Jason D Lee, and Qi Lei. Few-shot learning via learning the representation, provably. arXiv preprint N00iv:2002.094334 2020.
Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saun- shi. A theoretical analysis of contrastive unsupervised representation Icarning. arXiv preprint 1eX2::1902.09220 2019.
Steve Hanneke, Samory Kpotufe, and Yasaman Mahdaviyeh. Limits of model selection under transfer learning. The Thirty Sixth Annual Conference on Learning Theory, pages 5781-5812. PMLR. 2023.
Sai Li, T Tony Cai, and Hongzhe Li. Transfer learning for high-dimensional linear regression: Prediction, estimation and minimax optimality. Journal of the Royal Statistical Society Series B: Statistical Methodology, 84(1):149-173, 2022.
12

---

Ye Tian and Yang Feng. Transfer learning under high high-diimnsiioal generalized linear models. Journal of the American Statistical Association, 118654)::684-2697, 2023.
Zhe Fei and Yi Estimation and inference for high dimensional generalized linear models: splitting and smoothing approach. Journal of Machine Learning Research, 22(58)11-32, 2021.
Xuhui Zhang, Jose Blanchet, Soumyadip Ghosh, and Mark S Squillante. A class of geometric structures transfer learning: Minimax bounds and optimality. In International Conference on Artificial Intelligence and Statistics, pages 3794-388 99200. PMLR, 2022.
Peizhong Ju, Sen Lin, Mark S Squillante, Yingbin Liang, and Ness Shroff. Generalization performance of transfer learning: Overparameterized and underparameterized regimes. arXiv preprint arXiv:2306. 04901, 2023.
Yehuda Dar and Richard Baraniuk, Double double descent: On generalization errors transfer learning between linear regression tasks. SIAM Journal on Mathematics of Data Science, 4(4): 1447-1472, 2022.
Andrew Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. arXiv preprint 1809. 103744 2018.
Oussama Dhifallah and Yue M Lu. Phase transitions in transfer learning for high-dimensional perceptrons. Entropy, 23(4):400, 2021.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. Conference on Learning Theory, pages 5413-5452. PMLR, 2022.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-context learning? investigations with linear models. arXiv preprint arXiv:22 1. 1,6661, 2022.
13

---

A Limitation
In this section, we briefly discuss the limitation of this work.
First, we consider one-layer single-head linear transformers with certain (blocks of the) weights merged fixed. Though this simplification are widely used in theoretical works and linear and nonlinear transformers share some training behaviors (Ahn et al., 2023), this architecture still very far away from the transformers used in practice.
We also use a non-standard training algorithm that has several manually separated stages. Some parts of the modification are made to address certain issues of linear transformers, while the other are made to simplify the analysis. would be interesting (and more challenging) to consider more natural/practical training algorithms.
Finally, for our data- Gata-generating model, we only use it to generate one next token, instead of repeatedly apply SCB on the previous generated results to obtain long sequence. In our setting, the contextual tokens are independent. While this simplifies the analysis, deviates from how natural language works.
B Probabilities, expectations, and variances
We collect in this section closed-form formulas for the probabilities of certain events, and the expectations and variances of some random vectors of interest. All proofs are deferred the end of this section.
B.1 Probabilities
Lemma B.1. For any 1 E [7] and k. n. m E [N], we have
P(x, = X; m XT+1 = k) = q; Pn.mMm + (1- (x)) Unllm-
Lemma B.2. For any # r E [T) k, n, m, E [N], we have
(k) P(x, = n JXT+1 k, x, = m, X = =) = 4s (1) pn. 41 (2) Pnn (1- -143 4s - 4, pn-
(k)
B.2 Gradient and Expectations
Lemma B.3. Suppose the last input token xT+1 and a == AXt+I- The gradients of the objective are
Vpl= (VXa - ex,) (Xa)
V /// = (VXa - exx)
Lemma B.4. For any E [N] and s,I [T], we have
diag((), 1. 1MpT.
B [ex, e = E[ex, e
$# t.
Lemma B.S. For any V E NxN with Vu l and s, [T], we have
N
B [n,, Vn,x,1 = [[[V]] u. s=r.
[||]] #
Lemma B.6. For any V e NNN with Vu μ and [T], have
V (k) (V. P) + (1 q: 11
Lemma B.7 (Expected gradients).
B || Vdiag(#)
(1-11A117) [|A mn
14

---

- (Q, A) Pdiag( pt) (1 (1- (L. (1- A) up²
uy 1 jk [IV - -11u112) (k)) + px 111111
- -uz9 (V.P)H pu (1 - q (())) 11
Lemma B.8 (Expected preconditioned gradients).
B / || (V (V- (Q. A), (P 111),
E (2) = (II - [111112)) (am))
- -
- 1
- -1mul1))) -
B.3 Deferred proofs of this section
B.3.1
Probabilities
Proof of Lemma B.I. For notational simplicity, define x-, = f.b.. .T). We com-
pute
P(xo n1 =m XT+1 k)
Σ
P(x. = n1 = m, = m XT+1 = k)
the [N|-1
Σ
F(x, X, m,x n. = k)
AE
P = m, m xT+1 = k).
By the independence assumption, have P(x, m,x_; n Xy+1 k) P(x, -,, m
P(x; m) p = n. For the first factor, we have P(x,, Xt = m. m, XT+1 = k)
e (k) Pn.m* Eser 2: (k) P Pn.m.- Therefore,
P(N, n,x, = =m XT+1 = k)
Σ
00t0 (h)
PM)PPnm Pmmm" Ee!t Pn.m, P(x, P(X = m)
mh N-I
597
0 Pn.m* + L
Σ
P Pn,m, P(x_, = m) pm.
Set
Note that for any #
Σ
Pn,n, P(x-, =m)=
Σ
P Pn.ngHim
F(x-s,- = =m_ pn-
N-
= N-2 Un
Thus,
P(x, n,X, = m XT+1 = k) =
9:
(k)
Pn. +
Σ (k)
Un
pm
(k)
Pn.mHm (1- 1, (k))
4i
Hnffm
Proof of Lemma B.2. For notational simplicity, let x-x,-1 denote the vector obtained by removing the s, coordinates from x. Then, we compute
P(X. n XT+1 k, m,X, I)
Σ P(x, x s.-r i XT++ K. = m, = 1)
T-
15

---

Σ ME[N]T-
P(xo n XT+1 =,,x, = m, = =,
A P(x m1
Σ
9s
(k) Pn.m +41 Pnd Σ
di
(k)
Pn.m
P(x-s.-t = m)
ME[N]T-2
iE(s.)}
9s Pn + qr nnd + (1 (1- 9s 9: pn.
(k)
(k)
(k) (k)
B.3.2 Expectations
Proof of Lemma B.3. For each sample X, we have
1xx XT=1, Xo; A) = 2 -leer -VXAex
and a Then we have the matrix differential:
(VXa - 'VV.a+ (VXa ex,) VVda
Therefore,
Vvl (VXa-ex. (Xa) (Xa)
V / = 1(x++1 (VVX)) (VXa exx
Proof of Lemma B.4. When $ 1, we have [e,, e, =
B({) [e,, [e [e],]
up
When = 1,
have B(k) (e,, e, Ex Mreee diag(u).
Proof of Lemma B.5. When s # we have
N
N
B Vn.x Vnn1
Vm..
Pk Vn.k
n=
n=1
N
N
Σ
Vn, = # VVV = =||
Mkl
k
n=l
When =1, we have ΣN EV2 n.x;
EN
ΣN
m& V2 n,k
= 1|VI
Proof of Lemma B.6. Recall Lemma B.1. Then, we compute
V x =
Σ
Vr.mPax, n,X, m |x+11 k)
n,m=1
N
Σ
PammPnmmmmm + (1- -q (8)
qi
Vn Um
n.m=1
nmm==
qi (k)
(V. P)
10pp+(11-1 (1- (1)
Vu.
Proof of Lemma B.7. First, we consider Vvl (VXa - ex.) (Xa),, and compute B(Xa) and B ex,, (Xa)T. Write Xa = ΣTz dzex, Then, we have
T
B(k) [(Xa)(Xa)]]=
Σ
a a, Ble
(1) (k)
16

---

(1)) e Σ a a Blex, xx
(kk)
| diag(u) +
(1 (1-lactol7) uuu
=
where the last line comes from Lemma B.4. Then, we compute
lex, (Xa) = Σ
a
a (ex, 2 =
(()
Σ
a
[P(x, = n1 = m XI+1 = k)lmmE[NN]
[=|
1=l
T
Σ
(9. Pdiag(() + (1 qi )
=
ai
1=1
( a ali)) pdiag( p) (1- (1 u
where the second line comes from Lemma B.1. Thus, for Vyl, we have
NV
BVv V E [xa
Σ mx (ex, (Xa)
|| Vd + +1 ||4 VHp -QQ.), Pdiag( /1) - -(1 -11- (Q.44),) ,) up".
Now, consider Vausl = 1(x7+1 k)(VX)' (VXa- - ex.) and compute ((k) (VX) T *(VX) and
g(k) (VX) ex. By Lemma B.5, for each s, [7]. we have
[|A1] ==t,
e (VVX) = [(Vex) Vex, Σ
B Vn,, Vnxx
|||112 # I.
In matrix form, this is
B (k) (VX) VX = (III - [1al1()) -1 1 1+11 1|al12.
Then, by Lemma B.6, for each / E [7], we have
e, B((X)) ex, 1 = [(+) xx) exxx [0o1=00 =B(KVv Xo.N;
= (k) (V. P), + (1- (1-9(t)) (..))) 11al12
In matrix for, this is
[[(V))=xx]]=< (x) (V. (1 - q(0)) |1a1112
(1)9(()
Combine these together, and we obtain
M B V / = (())) [(VX) Vx] a (4) (t) [(VX)Tex. exol
= =(IVII [1al1)) a (k) + 1||4|| - (²) (V. P) - (11-1111 g())) (1) M11l
Proof of Lemma B.8. Recall from Lemma B.7 that
B || Vdiag(H) + (1- [1A1 Vun
-(,,A), Pdiag(H)) - (1- (l, m) up"
BV / = 14 (IVI - [lall2) a(x) + ur 1|
- (k) (V. P),,, ph (1- g ( |1ael1?.
17

---

For Vv, we have
pape
pp.
(diag())) 11/p)
Iall
pM !
Pdiag(4)diag(1// (1/
||
1|
pp"
=11
||
||
In particular, note that the μM /Il terms will cancel with each other. Thus, we have
Bvyl = ||A||; f. V + || μl - (Q.A), (Q. A) pl
= || (V =11) (e. A) (P- --11)
For V Vaii we have
4 a (11V11 -11a111)
1 11 (k)
11
q (k) (V. P)
T
+ ((1
11
q (k) ||
T
=IIVI ||
(R)
7 - (1))) - (V. (V.P). P)p + +( (k)
((())
||
7
- (VV P), - |1aa1²)) ( 1
=(IVI - -1 (R)
7
B.4 Concentration
In this section, provide concentration inequalities for the gradients of the loss function. The concentration is applied on the gradient noise term
hy.r 4 (B) - -EAE!)
(BB) 1 -
AA.T
where ((8) and ((B) are the preconditioned empirical gradients computed from batch of size Here, we first consider the concentration of the original gradients:
Vv (VXa ex,) (Xa)T,
V / = 1 (x (VVX)) (VXa-exa
and then consider the concentration of the preconditioned gradients. In this paper, focus on |1-11p as the mostly used metric for the gradient matrices.
First we prove naive concentration w.r.t. any random vector y with bounded second moment with any II-II-
Lemma B.9. Fix 8, O. Let y be D-dimensional random vector with E ||>l| s G. Define
y (( B EB Yi where (y:); are i.i.d. versions of y. If
G
B
then with probability least 1 8, have lly Eyll
Proof of Lemma B.9. Assume w.L.ogg that Ey = 0. First, note that
B
B)
B B(n,y))
Σ
B
Blbail
G
B
B²
1
18

---

Hence, by the Markov inequality, we have
P 2 181 2 z) s)= =P -11100001
||> 11
82
Bg2
Thus, for fixed e, (0, 1). if choose B = G/(ōE²), then we have with probability at least 1 6, [B S..
Now we upper bound the infinity norm of the preconditioned gradients to apply concentration.
Lemma B.10. Suppose that \|ee &l O(1),V V Av, A Á + AA, where V E ENNN is transition probability matrix, and in the attention matrix Á each column ā'k) is probability vector. Moreover, ||Allf ||Av 11y S O1/TT) Then, we have
|| u - O(N),
||vect O(N)
Proof. We first consider the infinity norm of the original gradient. Recall that the gradient for V and A are
Vaail (x7++ (VX) (VXa eex.)
Vyl= (VXa ex.) (Xa)
and the preconditioned gradient is:
(( r
1N N
(v(R.)) diag ( IN 14
=
IN
N
||
((8.) alk)
1
lr
1 1
(BB)))). ((t)
We first consider the maximum absolute value in the original gradients. For V Vausl, we have
- -ex ||.
(XT+I =
= KV(X ' (ex [ = =V(X)) (VXa (VXX ( I1: +|[((T=+=
T
max x
Σ
a. (Vex,) Vex, + max sE[T]
The first term can be uouer-bounded in the following way:
I Σ
max
a (Vex, ) Vex.
= max SE[T]
a, Tex, +
(Vex,
S 1=1
1=1
[=1
< 1 +
||arrlli
max (Vex)) Vex
1=1
$ (1 + VT (12) max (Ve..) Vexl
S.1
Since V= V Ay and ||Av 11. || < O(I/T) we have m (Vex,) Vex. upper bounded by
O(1). Therefore
max S
a, (Vex.) Vex, Vex, 00
And similarly, the second term max,e(7) (Ve., Texo can be bounded by 0(1) because the infinity norm of is also upper bounded by 0(1). Therefore, we know //Vass/lee S 0(1).
Now we consider the preconditioned gradient ((B,) l:
(A))
1 1r
(B,) (1))
I||0o
((B))
4
(B)) (k))
14
T
19

---

11 (v(k.))) 1)
1
1r1 (vlE)>))
+
s SO(N)
Mk
pk
T
since 2 N for all & E [N].
We use similar technique on & (Brr) 1. First, we prove the infinity norm upper bound on the original gradient.
V|Ve /lloo
|| Vec (VXa- - ex,) (xa)"1l.
||Vee (VXa) (Xa) +
+||Vec (ex, (Xa) 1loo
max s.r(TT]
(VXa), (Xx),7+
Σ ar Vec (ex, ex
Σ aj Vec (ex, exa)
II(VXa) (Xa), +
Σ 8
(a, + +l (exxxllllll 11
Σ (ā, + +| (1
//vec Vee 1ll 0(I).
And therefore, the preconditioned gradient can also be bounded.
|| V
(B)
/||00 =
IN
1N
( (v(Ht))
diag(1/
IN
N
||
( diag(1/H)) (IN
nM
||
IN (v(..)) diag(1/m)
IN
||
s O(N)
since MR 2 N for all [N]. Now we finished the proof.
With the upper bound of the infinity norm, we have the following upper bound on the second order moments of the preconditioned gradients of A and V.
Corollary B.11. With the same setting in Lemma B.10 and || aas s ([ O(N). Moreover, ||Ax 1| 1| < O(1/T). Then, we have
B | O(TN3)
Proof. We directly upper bound || AI] and || using the upper bound on infinity norm, Since || || O(1/T), we have the infinity norm be upper bounded by
| s O(N), |/vect Vrectv/ll. /1.e s O(N)
Then, we can first bound the Frobenius norm /|1. /I f I| We have V NNN, A E RTNN so 1 = N2 [| = =(N)) || = NT |TT O(N'T).
That leads to:
||tv/[[, = (v 4 diag(H)) s 0
11,/113 s O(N3)
N
|| = 14.1. diag(r)) x 0
|| = O(N2T)
where the second inequality comes from the assumption that μ O(1NN)
20

---

Now with the upper bound of the second moments of the gradients, we begin to prove the concentration of the gradients. We first consider the first-order terms that need be bounded in the signal dynamics:
(hv, P),, (hA. 2),
KpKę
KpKę
Lemma B.12. Fix 6, O. Under Assumption 2.1, suppose || | (()), I| ((N). If B 2 ((1) max (N 2222)
N log
then with probability at least 1 - 5, we have:
(hy, P)
(hA. KpKo
8..
8
KpK@
Proof. Note ha = 3 v (t) E a thus we have the upper bound for each coordinate of the gradient error bounded by @(N). Similarly, we have the upper bound for each coordinate of the gradient error of by bounded by ((N).
Then, bound the infinity norm of (Vy((i), p)μ and (Val(i, P),i
(Vyl(1), P)μ
Σn
Vyl(mmmm Pnm
KpKQ
KpKo
No 1/ ( 11
KpKe
@(1)N3
(|) S 1.)
KpKę
(vAI().
Σ ΣN mm
KpKe
KrKe
QN A(1)]
11el
Q is Q-sparse.)
Krke
0(1)QN2
( 1.)
KpKę
Note that B V (/) - B / = -E 00 which means the two terms above have expectation 0. Since h by both averages of B gradients of single sample, we Hoeffding Inequality:
(hv, Phu
BE K K2
2eep
KpKo
N6
(h 1)p
-B& B3=K2KK'
P
26
2exp
KpKo
NºQ2
By union bound, if B max N22Q2?
No lo eKKKK
it has at least probability,
(hy,
(h 2)p
SE.
KpKę
KpKe
Then finish this section with the concentration of the second order terms || and ||y 112, which need to be bounded in the error evolution.
Lemma B.13. Fix 8 > O. Under Assumption 2.1, if y Tmpa O(N3),
Tmp4 O B @(TN²) 6e2
then with probability at least 6, we have:
| S E. | 11
21

---

@(TN2) de
Proof. Similar Lemma B.12, Bhy =, = 0. By Lemma B.9, when we pick have ||lll S 1l with probability least 1 1.
C The population projected process
In this section, we define the projection of the true SGD process onto the "space of population trajectories". Then, we derive formulas for the dynamies of projected process and the distance of the true SGD process to the space of population trajectories. All proof's - except for those short ones deferred to the end of this section.
C.1 Definition of the population projection
The main reason we analyze the population process first that on the population trajectory, both layers possess special structures. Recall that
B / / || l (V 111) (Q.
(P- -111)
P [IV 11 A
17
1r N - (v (e-
T
T
and initialize T 0 and 1-11 O. Note that for any (zr)r. if 0 and = Az + Bz.. then Z x z. for all 2 0. In other words, Zr moves only along the direction z. and therefore, can be characterized by single real number. This is exactly the same case of V u1 and A 11 /T in the population case. Hence, in the population case, V stays on the line crossing μ1 and P. and A stays on the line crossing 11 and Q.
Unfortunately, mini-batch SGD does not stay exactly on the population trajectory. We can still, however, look at the projection of SGD onto the "population trajectories". Formally, for any V.A satisfying 1 V -,y A = I and V =
μe, we define
aPP + (1 a)p - VIl
ay := argmin .E3
aQ + (1 (t) 1,1 1
0A == argmin at3
7
By setting the derivative to be 0, we can obtain the following closed-form formulas for av and d'A-
Note: Without specification, we drop the the time subscript and consider ay == av,r for similicity. Lemma C.1. For any V, A satisfying 1 V = N. 1 = NT and Vμ μ,, we have
av Kyp/Kp and a KAQ/ KQ.
where Kp ||P1 - Kyp = P) 11 Ke |1@11 1/T, KAO (A.2), 1/T. For notational simplicity, we define Bv = av, BA = AAA.
11T
V av BBHH1 and Ā Ā= a BA
Then, define Av V V and Aa = A Á so that we can decompose V Ay and A= Ā AA- By our construction, we have 11 1 Av and similarly for AĄ- We will now show that we can in fact drop u1
Lemma C.2. For any V OP (1 0) and A' eQ (1 0111/ with 0 e R, have (Av, V',,, = 0 and (Ax. A') 0. In particular, we have (Av, P) (Av. V), = and
(AA. ,), , (AA, A), =0.
Proof. Note that (IIT, Av) (M, (V V)H) = 0. Hence, (Ay. V') (Av, V =11) = 0. For (AA. A'), it suffices to note that (AA, 11 (T), = =(1 (1 (A A), 1 /T) =0.
The following lemma the basic definitions and results about the population projection.
22

---

Lemma C.3 (Definitions and basic results on the population projection). Suppose that V. A satisfy
T N. 1 v2 and V = μ. We define the following:
Kp || - || Kvp (V. P) ||all
Kv= =v=IIV[ |1
av KvP/KP. Bv - av. +
Av -V
Ke = 11211 - 1/T, KAO (A.2), - -IT,, KA = I|ll - 1/T,
aA= = KAo/Ko.
BA = BA. A +Balrly/T,
===A A A.
Moreover, by Lemma C.2, the following hold.
Kvp PV - -11p111 avKr.
Kv = || + ||Ay 11 -111112 = K2 Kp + ||Av 11.
KAy = (A. - -1/T = =axke.
K^ = || [|Аil] - 1/T Ke ||All
C.2 Dynamics of the population projected process and the approximation error
We write
V+11 V 7/ E //yl 7/ (5( B Py =: V, 7/ E E.y hyhy...
A. - - BA - 71 (b 1- E V.1) =: A. 17. E V 1A
Ar+1
where the expectations taken over the fresh samples step First, we expand the expected preconditioned gradients around the population projection. Lemma C.4 (Expanding the gradients).
/ = Kpov (αvay (e-
1 T
KKaasss + ||Av +xpasAA+1Avvll
A 1
P (RARY (P-117) av- (P- 111)
[oo - Av |||||/ (V- xl1).
A
Then, we compute the dynamics of the projected process.
Lemma C.5 (Dynamics of the population projection).
0V.r+1 av.r Ke (1 a.av)) a 1/V 1 ary
nyvay ||A
" (hv P),,
Kp
@A,7+1 = a +IAKp (1- avyad ov - A ||AiI
AA
(hA.T. 2)₂
Ke
Note that V 111 (1 and Ā = 17 arx 1TT). Hence, this also gives formulas for Vr. and À T+1-
Now, consider the dynamics of the errors.
Lemma C.6 (Dynamics of the errors).
||AA.+ = "AKKPaa 1/a ||Av I|A Il
227A ( (I-naKpa? - nA || I1Av (sshhs)
23

---

A K2
(hA.T. 2) +1A llhalll l
| = 1 - 11V
K + T +) - -nvliaali) l1vv 1l
+ 2/j 1 - 1/v [ + + - nv -mllllal))) cav hv)p
my (P. + +r² ||h 11.
Kp
C.3
Omitted proofs in this section
Proof of Lemma C.I. We compute
1
0e la a) -Il (1P+ - -V.P- - 11)),
2
=aPP V.P),, =akp || (V, P)
Set the derivative to be 0, and get av ((V,)) -11111//pp (rrl)) Similarly, we compute
11T
A 2
11T
41,
11
2 dur
a + (1
= aQ + (1-a) a)
A.
T
T
T
11T
aQ + (1 a)
A.
T
=a (1e11 1/T) - (AA.2) 1/T)
Again, the derivative be and we get a (A. م 1/T) 101211 1/T).
Proof of Lemma C.4. Recall from Lemma B.8 that
B4, K^
(V
u1T)
KAQ ((KAe 7
1
(P- μ17).
7
B = Kv (A- 1 - Kvp Q
1
T
T
First, consider the dynamics of By Lemma C.3, we can further decompose as
1
171
E BV (oēKp + [|vv]
A
N
KyKK
2
T
1r1
1T
(A- 1
= Kravv
αv
A
0
+ +11Avll
T
T
1
171
1
= =Krav
av
A
2
||Av/] A
T
T
T
(0-
1
+ KPazAA ||Avll A- 1
= Kpay (ayαx-1) 1)
T
Similarly, we can rewrite the expected preconditioned gradient of V as
B B4,
a Kg + T (V-111)
aake 7 (P- -11) (P-11T) 1| (V-+1T)
Ko + 7 + (V - (V-
(1-11²)
(anke + 7 (P-111)
+ a Ke + + dr + ||Ax|| (V-117
24

---

aake (ayar 1) (P-1P-1) -11' + av -1 (P-11')
+ +l22ke+ Ar + (V- (V-H1').
Proof of Lemma C.5. Recall that av (V. KKp and as /Ke- First, consider the dynamics of V. By Lemma C.4 and Lemma C.3, we have
@v.r+1 aV.r -1]
Kp
T/V Kp
aake (aAaV- 1) (P pll, P),,,
n av -1
(P- (P-m11.P). Pp
QV.r
Kp T
(hv,r:P) P)
|& (V (V-11² u1T.P), - KP
Kp
= =vv. + Ke (1 a(av)) aλ + I7
1 av
aaav ||Aa||-
aI (hv.r.P).
T
Kp
Similarly, for V, we have
AV + HA.T Q)
@A, = =AA, A1
KO
nA Krav (avaA 1) (aVDA- Ke
1
.0
-.A.""
T
(1A 2)₂
11A
||Av|] 1|3
A
1 N
11A
K2
KQ
1
= DA. + 11K (1 avaa av - MAaA ||Av|
11
(IA.T: Q)
Ke
Proof of Lemma C.6. First, consider the dynamics of AA. which is given by
1
M Ar nAKParaA 7A ||A
A
h
T
nA (hA.T. Ko
2)
110
1 N
T
Decompose A into Á AA, rearrange terms, and we obtain
AA.TiI AA - TAKKa ||A 11 Aa
( 1
+ NAAEA Ilvlll (e- 1
- Illlll 1
7A (HA.:: 2),, (0-
1r
-nAhA.
Kg
T
Note that À 11T/7' (1 11 /T -11T/T- asll 11 /T). Hence, we have
AA,TiI AA -naKKaaiA 777 ||Av| 1| A^
NAAA 11 1
12
10-1T 1
NAAA A e-
1
T
7A Ko
(HA
()-
1
nAAA
T
25

---

7A (KA.T. 2)p₂ (e
1r1
=(1 -1-110 (1-n - 77 [|Av/lz A1) ^^ +
nAhA-
K
7
Recall that (AA. - 11 /T), 110 - 11/ kg. and (1, ,AA,, = 0. Hence,
|AA.T-Iff [
-lAKKav (1
1IA || [| AA
(hA.T. 2),, 0
1
- nAhA
Ke
T
aA (NA.:. 2),, )
2
=(1 -1-nnAraa na I1 |ANII
Ke +1 11
Ke
227 [1nAKpaat - nA (lllll1))) (A).
2 1Aa
(HA.T. (2.
Ko
(1 nakpay-
11 (l. ||AIII - 27A (1 (1- -AKKpaa - nA mallAvll)) (AA. ha)
1l-nakaai-nallyy
n (hA.T: 2)
+ 3 ||lll
Ke
Now, consider Ay. Similar to the previous calculation, we have
Av, Av
aa + 7 Av - nv IIAll Av hyhy
av
(P 10-111) - n -l (p- -1l')
n
T
1 av - navv ||A Kr
n
(hv .. p)// (P (P- 11
1/v
T
Ke + 7 - (ttllAall) Av
11V
(IV P), (P 1--111) hv
Kp
Again, note that Ay La P ult, ||p 11p11111; -rr, and (H1T, hv) =0. Hence, we have
111T/
1 πll a 22 + + - -nVIIaalP) ||Av l
| [Al
(1- -vı a Ke + + - -n ny Ilaalp)) Av..v)),
+ 2ny
P2
(P, P. ) 2 || 1}:
Kp
D Stage 1: signal boosting
In this section, we assume both av and a are close to O. In this case, we can approximate Lemma C.5 with
1
77 (hv.Ply
0v,-+1 avvr + +10A + TKo
KO
KoKr
(hs. e)p
aA.++1 CA.r + nav - aaa
Kp 11 KpKe
We can also write this matrix form as
@v. @A.T+I
[ayr UA.T
- A|l /Ke
lav.r aA.
+ 17
-IlAvll /Kp
26

---

1/(TKe)
as I (KpK@) (hA. 2) P (KpKo).
+
-"
Suppose that ||A / and ||Av|[ (Kp both bounded by 62. Then, have
1
av.r+I + aAA.+ 2 17 26 (av.r + + + "1 TKę
(11)
P) 2) KoKr
KrKę
As long ò & 1 and we choose a sufficiently large batch size so that the second line is bounded by n/(27 2), ay++ a grows exponentially fast. Similarly, one can also bound the difference between and aA. Formally, have the following lemma.
Lemma D.I (Main result of Stage 1). Define the end of Stage 1 as
1 O(1)
T == min r 020 max 2 min
2 QN
Suppose z 1/10, nv.r = Kn and NA,r n/K, for some S [ Ko). Let B > 0 be the number fresh samples we use step r. Suppose that B 2 Ô probability 6+
PO'N
are chosen s.t. with
min(K Ko)
max P)y
(ha. (ha.2)m Kpke
(12)
KoKr
4TKe
@(1) min 112 k112
min K o.
K
VeNUT
max ||
(13)
T log TKe
Then, the following hold with probability at least 6p:
(a) T s 0(1) (og(TT /q-
(b) Throughou Stage 1, ||^Il] A/T and ||Av A/T, where
min [
ke. kr. (1)
@²N3
(c) At Ti, we have av a = O QN and lay a.l -
9(1) TKe
(d) For all t S Tú, we have || s @(1), ||]] s 0(1(Q).
The proof of this lemma is large induction argument. We will first assume the bounds on ||Aمل| and ||vll are true, so that the approximation (11) is valid, This will give an upper bound on the length of Stage Then. we show that within this many steps, the errors cannot exceed the given maximum values. Thus, the induction hypotheses are true and Lemma D.1 can be established.
Part I of the proof of Lemma D.1: Signal growth rate
Proof. Since avas $ | by definition of Stage I, we can rewrite Lemma C.5 as
3
av. + @A.t+1 2 (1 (1+3n/4) (av.r + (.A.) +
4TKQ
|
||Av|
(hv. P),
(hA, ()
mja KO
Kp
KoKr
KrKQ
3
1
(1 + 3n/4) (av.r + a.A.))
4 nay
naa 4
2
4T K
27

---

(1+5) (av,r ( + 4
where the second line comes from induction hypothesis (b) and (12). Recursively expand the RHS, and we obtain
av + UAAT 2 2012(1111
(1+ ") avv. + 2TKe 2
((1 21 - 2TKe
Since the RHS is upper bounded by by definition of stage I, and we obtain
s O(1) log(TKe)
Part II of the proof of Lemma D.1: Upper bounds on ||A||
Proof. Recall from Lemma C.6 that
Av r+1 [ I |Av.
11. 2n/ [|Av.rlly wl rrl, + ny (lhv.rlly
|| [| < | 2m || || + A AAA
By part I of the proof, Stage takes at most @(1) (og(TKK)// steps. Hence, it suffices to bound
the increase of these || in this many steps. Recall the notation A = min k Kp, induction hypothesis (b), we have for all T.
e(1)
By
@N
| = Av I| 2nv
|A r11,, || + n | /
$ Av.
[|
2 / + + / |/Avvrll.
O (() v\/T min(Kr. Ko}
0(1)A - min(K,, K3)
|| 1
tn
log(TKo)
TK log (Kk))²
|AA - AA
1,-
2n ||A..
Ila ||H...lll + ||
< || A
||
211
21
A/T / + + A
@(1)TA// Va/Tminn min (p,, Ko)
@(1) min (K?. K2)
||
+1
Kp
(TKe)
TK? log(TKo)
Since Stage at most takes Ti 0(1) (TK))/ steps, the increase of ||AA.-12 and ||A.. in S Ti are at most (since A4.0 0, A, .0 O):
0(1)NA// V\T Minn KO
0(1)A min min K...
Av AAp.11
+ +n
(Tg(
TK2 log (Tg()
e
@(1)
JATT Ki, Ko)
(()) min ( (K2 K3)
+1
Ti s A/T.
Ke
log(TKo))
TK2 Q (Tk))
@())VA/T
VATT Min(KK,,Ko
((1) min K?., K2.K2 2
|
KP
(Tg(
TK3 log (Kgy)
V/T min Kp, Kp)
@(1)A: min K? K2)
+
Ti S A/T.
Kp log(TKg)
TK2 p (oggTKoo))
Therefore, completed the induction for the error terms.
28

---

Part III of the proof of Lemma D.1: Ending state
Proof. First, consider the distance between av and aA- Similar to the part I of the proof, we rewrite Lemma C.5
@v.r+ 0A,++1 =- (1 avaa)) (av,r aA.)) +11 1
T
||A
llvlly
(hv.
(hA. Q) KpKQ
-nov
Kr " KoKp
+1
Ke
= (1 -1/2) (av,z - aA,,) (1()
Thus, whenever orvr dA,r (2(1/( it will start decrease. Since the amount of increase at each step is also upper bounded by O(1((KK)),, this implies Jav,r CA,= s O(1((K)))) The other direction
be proved in the same way.
Finally, we bound the possible amount of overshot. By the part of the proof, we can also upper bound the signal term growth
@V.r+I + +av.T+1 & (1+ 2n) (av.r aA.)) +0 n
TKę
Since uv.r drA, $ for all Ti, we have ay + + O ( at time TT.
Part IV of the proof of Lemma D.1: Upper bound on Infinity norm of and A
Here consider the upper bound of the weights V and A, which can be used in the concentration section below.
Proof. First, we upper bound the infinity norm of VT-
|V ||V, Ay -l = || + ||Av || |lvlle.
and we can upper bound ||Av F by its μ-norm:
| || Avdigg)))( ||Ay I|:
Thus we have ||o, < V + + ll s |i + 0 x by Induction hypothesis (b). And we can further bound V=
11.0 s O(1). || = |Iv, (1 av 1 ||||| 11
Therefore, we have >l., @(1). Similarly, for A we have (since || = / /l
C/O.)
|| Ι| Axloo s || ||A s ||AT|l. I|AllF 0(12)).
Part V of the proof of Lemma D.1: Concentration
Finally, we need to ensure that with high probability, all the error terms h cannot exceed the given bounds throughout z s T.. We use Lemma B.12 and Lemma B.13 to bound the concentration of the error terms.
By Lemma B.13 and union bound, have that if B, T-ITT log (TKo) then with 6, min K K2.K
probability least 1 6,/2, the following holds for all / TT:
@(1) min K 0
1/2 k1/2
min(K or Kp)
max ||A |
Q'
Vo TKQ
29

---

8(1)9, max max(N..e))N+ log
16 TKo
By Lemma B.12 and union bound, we have that B, 2 with probability at least - 6,/2, the following holds for all 1 1 TT:
then
(hvv P)μ
(hA. 2) KpKe
1 1
S 4TKº
KpKę
4TKº
And by union bound, have that with probability at least I s TI. Therefore, we conclude the proof of Lemma D.I.
Ör, all the bounds above hold for all
30

---

E Stage 2: learning the model
In this Stage 2, we use a positive A for the E-regularizzaton. In this case, can write the update rule of each a (k)
(k:))
ar
(k)
17A v
BB) ((R)
(gradient descent step),
r+l
uk
(k.)
if a (k.))
a (k.")
r+1,1
0.
a
T+1,
d,
if
T+1,1
(R.)
2 d,
S,
(proximal step),
(14)
r+1,
r+1
a r+1
(I,") + (1
r+1
(k."))
T'
(projection step).
For notational simplicity, we define
gr (k)
1 (a (R) r+1
(K)
We further define G2.- = - (Ar+1 -7
Ar) as the full gradient the matrix A. we will show that with appropriate rounding at the beginming of Stage 2, we can ensure BB (t)) using only B, x log(T) fresh samples at each step (Section E.
First,
gr (R)
E.1 Rounding and gradient denoising
Recall that the first step of the Stage 2 is rounding each (k) by setting small coordinates to and then projecting it back to the affine space the probability simplex lies. Since after Stage 1, there will be a separation between (a) with (A) and (k),, this rounding step makes all a(") with I q (k) have the same small value.
Also recall that we use E Errregularization and proximal gradients in the Stage 2. Effectively, the & f-regularization ensures those useless a²²) are always 0 (before projection). Then, similar to the first rounding step. projection will again make then have the same small value.
In this subsection, we formalize the above argument. We show that the rounding step can recover the support of q% (k), analyze its influence on a and the distance to the population subspace. Then, analyze the effect of use of the proximal gradients and show that with poly(N, Q) log(T) fresh samples at each step, we can make sure the difference between update and the population update is small with high probability.
Lemma E.1 (Separation between noise and signal). Assume that at the beginning of Stage 2, we have
Q T +VVN (NNiIAAll))
av 2 θ
Then, we can choose a threshold 0rmp = ((vv/Q) s.t. a, s Ormp iff qt (k) = 0.
Proof. Note that there exists some universal constant 0 such that a, (k) c/Q for all nonzero
(K) and μu 2 c/N for all & e [M]. Recall that the population process = ax + (1 1T Hence, for all k [N] and / [7].
q,
d S 1/T, f ( cay/e. (1)
Then, note that
T N
|
Σ
a, a, )² μk
(k)
1=1
k=1
(k) a. /
2 c N
c [
N - |
a
N
1=1
31

---

Hence, for any & E [N] and / [7], we have 1a, a = ||A vN/c. Combine these together, and we obtain
S 1 + |lxllI VNCc, if if ( 2 cav/ 12- - ||Aal|y vN/C.
a
if E q (k)
Hence, in order to get a separation,
1/T + IllAll VN/C S 1 (cav10 - -1AAll, VNIc)
2Q cT
3QVN
=
c
|| S av.
Lemma E.2 (Effect of rounding). Under the conditions of Lemma E.I, choose À as in Lemma E. and set
(k)
1 a (K) の (1fa,t) (k)
11
2 =1)) T
T
|=1
T
= aa (k) 3 (110,t) (k) 11) 7
+1 1- Σ a
(k)
1 T
16q
We have AA + a O and |All - ||lll +(1/TT
Proof. For notational simplicity, put b (k) = (k) o (1 (R) 2 21) By Lemma E. we know (R)
supported within q (k), Set ((k) Mx ((R) Then, we can write
(2) - (k) + (1-b(h) 1
T. Recall that aA KAO/ Ko where Ks = (A,0) / = Ek HX (a (al)))((k)) (())). Hence, for each k E [N]. we have
(a(k) (000.((()) (111) (B(K), .((k)) + (1- b(())
7.. g((())
1
- (aak)) (())) + (1 b(()
7
As result,
Σ M* (1) a - a + T k=1
N
0(1)
.AA +
T
Now, consider the effect of projection on the distance the population subspace. We will use subscript new indicate values after rounding and use notations such as a (k) to denote the values before rounding. Recall that || ||2 = Σ Me (|a(() - -ak)ip2 We have
|a (4)
(k) 2 nnow
b (k) (1- ((k)
1 1
a.A.www (A)
T T
(k) a..oww q (k)
(b(A)
1
a...ww
7
b (k) a...ww (12) 2 + (b()) -aA.ew
T
2 2((()) -on.mew))
(B())
(A)
aan,now
T
= (000 -(aA ((1)((1)))) -1pp:)99///p10000000 -(b(e) ..A.noww
2
7
Note that
ar) = [a())= -000f a (k) - aa -1-
- a (k) -a (2) - -
32

---

b (K) org - ax) T
Hence,
anow dbool
(k) - a 2
0(1)
T q (k) 1
q - -1b000-aa00 -aagu)/
0(1) 0(1) T
$ ((k) -Aq (k)
Thus, ||al]f - ||I O(1/TT.
As we have seen in Stage 1, the norm of V / can scale linearly with T. Hence, in order to make lgtt)rr has size (1) in terms of [1-112- However, note that all entries of V ( are bounded, and therefore are subgaussian. Hence, for each entry of V I, with O (log 7 samples, we can make sure the relative error is small with probability at least 1 - (T), By union bound, this means the Il-I| error can be made small using only log (7) samples. Note that there a separation between the signal and noise parts of B V Va I. This implies that can distinguish them using lTg(?) samples and directly remove the noise part. Formally, have the following lemma.
necessary to use poly(7) samples, which is undesirable.
Lemma E.3 (Gradient denoising). Given &
@
@² VT ). 0.1). Suppose that for any k, m, n E [N],
(1) s 1/T ift & (k), |Vnml O()), and ||Aall <c(l-vv/(VNO) and [lvll, s cll-av)
for some small constant > For the target accuracy Emmp O @N=z), we have with probability 1 6G. If we choose
Q*
CNT 6t
B, 2 C
log
ea K?
for some large constant 0. Then, for each E [N], with probability least 1 Srmp, we have
a (B+) (k) = (1 @TmPl
B
m/ 2
θ
Kv
Vr E q (k)
4;
NQ
0 (B+) 4
= 0
T
Timp
avKp
NO
=
N2
EKp
Vr $ q (k)
Proof. For s [7], recall that
( ( = = (VXa exx))
0 (k)/ MkKyas HıKvrqs- n,
First, consider the expectations. If & (k), then by our assumption, we have [Bå a gill = MyK O(Kv/(NT)). Meanwhile, for E q (k), we have
-BO anl p*Kvp(qs az) HR(Ky Kvp)as
[a<(Lv( av) (vy avkp ||ll|| VN- ||
As result, we have for any E [N] and s E (k),
1vav -
||all S c
KyKp *
Kv
- Ed un 2 c
O
=
QVN
NQ
NT
[|vi s kv))Kr
for some small constant > Then, for size each entry, we have
N Σ
s s ((ee,,) (xka|+(1 +((Ve,,) Teol Σ a,
Vn.x, Vn.xi IV S N.
n=1
33

---

Therefore, 0 az is 22 22 Tmp -subaaussin, whence 0 (BB.) is N2 / Su.-uuuuusasiiann As result, for 0x each È O, we have
a
[ Bd ml
-E2
P
2 _ s eexp
4
02
N2
Apply union bound and we get
Br
exp Iog((NT) f²Br N
P V &V 2 s 27N exp
N2
Recall that the separation between the expectations is cay Hence, it suffices to choose & = cayKp/(2N)), Then, make the failure probability at most Ómmp, we can choose B, follows:
exp [og((NN))
5²Brr N2
S B, C N²
CNT 8
=
log
d K.
for some large constant 0. Then, boost the accuracy from 1/2 Srmp. it suffices increase the batch size to C Tiog oi Kp
N²Q² log
CNT Or
Using this lemma, we can pick a = O <Kp @²N3) to sparsify our proximal gradient. Notice that our proximal gradient is biased estimate of the true preconditioned gradient, but the separation guarantees that it is possible to make the error controllable. The following lemma calculates the error each
proximal step introduces. Here we define hA.T = G (k) -B instead of h because of the bias
introduced by the proximal gradient.
Lemma E.4. Under the same setting of Lemma E.3, if the batch size
B max (C C K2
Qº log
2CNT
O(NOQ3
6- K²a
562
the noise attention score la, O(1/T) for all & (k) and all k E [N], and à = @ with probability 1 6,, the gradient error at iteration
then
Q-N3
h | == I| - E S O QN2
Proof. For notational simplicity, we drop the superscript k. The goal is to estimate the difference between E Val and ga.r by calculating the magnitude of the bias of the proximal gradient together with the concentration error.
We consider the population gradient first. Since we have (rr,r)rgg are all the same in Stage 2, we can write
a, [arly (1- 1r where br Σa. a.
Q-Q
req
and recall the projected preconditioned gradient for A.
B = (IIVI}2 -11u1() (a - -2) ((V. -
Therefore, we have
(1- bg)
+ (1IV (V. (V. 1 L
V = (BVal, + Ky
T T-Q
Now, we consider g by calculating the expression of ar+1. First, note that by Lemma E.3, we have d' = (di,.4 a) [t q}. This implies that (ar, breg all the same and the value is at most 1/T. Moreover, it also implies that it suffices to focus on [a" ++ lq. for which we have
[a" = [a'mile Alg =[arla
v B+) - -lg.
34

---

Therefore, Zieg d 1,741
b,
2
Eleq
((B+)]
Q. and
Mr.
a++ = la
17
[[(B))|
l +
1 b
Hk Σ
1
OuBr) Qdi
1
T
fk
= =ar
(1- b
17
v (B+1)
-219 +
1 b + 11
Σ
Oar
(Br) + Q. T
T T-Q
leq
Thus, write explicit update
V.l+ [B))
2
1J Σ
(Br)] 1 Q2
1
(1- b,)
g1
+ - 1
b
11
Pkk TEQ
11
n1T Q)
+ 1g [VV (B-)
|((B)) 11,
1 by + 11
Σ
Oas
(B+) 1+ 0.
M&T
Mk tEq
n/
(1 - br) lr
1 1g
(T Q)
1
Then the gradient error at step T can be decomposed into:
8a,r - B/ = (BB)) IBVal
(Concentration error)
1 1 (. (B-)
b 7
Da:
(B+) + Qd
1
d
-1
HRT
jak
n/T
(1 bg
(T Q) Ky (1 by lg
1g:
(1IV11 (V.P)
(Gradient bias error)
T-Q
T
Here the gradient bias error can be further simplified to
0 [B+)
1 b- Qd
l
(1*)
P
1/T
Eieq 0(B-) PgT
IIVI (V. P)
Q(1 by Q.) QAT - nTTKK (T)
1
(2*)
T
First, estimate the concentration error. Similar to Lemma B. we first upper bound the infinity norm of the gradient. Consider the maximum absolute value in the original gradients. For V Va we have
1 (XT+I ()(VV))XXVVVV (VXa -exo
| [ = (X)(V) (8))
I|. +||1(x7== (EXEV))TTxeelll)
max x
(Vex,, Ves,
(Vee,) Texol sE[T]
max
The first term can be upper-bounded in the following way:
T
Σ
a, (Vex, Vex,
Σ |arl Σ larl max TVex.))
XE
E=
reg
req
5 (11) max ((e Vexl
The second inequality is due to la (k) max;,1 (Vex,) "Vex upper bounded by 0(1). Therefore
< O(1/T) for r & q. Since Vnm S 0(1), we have
T
max x
ar (Vex,) Vex $ 0(1)
35

---

And similarly, the second term m [T) (Vex, ) Tex, can be bounded by O(1) because the infinity norm of V is also upper bounded by O(1). Therefore, we know // I O
Now we consider the preconditioned gradient
1
((8e))
1)
11100
1T
jk
T
1
Ir (v((k))))
i)
1
1
(v(8.)))
1)
+
s SONN)
Mk
Mkk
T
since μk 2 & for all [M]. Now since ()(+)
is Q-sparse, we have E || |
s O(ON3).
4
By Lemma B.9, when B,
@(N"0))
& with probability 4
56
(
0 8 QN2
Then, consider the gradient bias term. With the selected A O cKp
and
|o65-11
1
00
ckp for N'
@
r é q, with probability 1 - 6,/2 we have the μ-norm of first term (since Kp O(N)):
II(1))|
1
Σ
(B+)
Q1
+
ve
T
144
+0 E
S 0
s 0
+ 0
ON2
TN2
Q3/2222
ON2
and the second term can be upper-bounde by
|I(2)||| = (0 QN
QN
vT s o &
0
+ 0
T
T
7
QN2
since J/ s 2 ( (BB)) = 0(1), and [IVI S N.
Combine all three terms and by union bound, have with probability 1 - 8,
r = G (k)
BAAI
N
Mxllg (k) - B B 11 s o &
ON2
since br = O1/N).
E.2
Model aligning and the decrease of the errors
First, we show that the signal will continue to grow and approximation error will decrease. This decouples the error at the end of Stage and the final error. In particular, we show that eventually we will have ay a 2, || and [|Av 0².
For notational simplicity, define 5², = >|A^^lI, Kg and o7 [|vv|l /Kp. Recall Lemma C.5 and Lemma C.6 and that we choose 17v = /7/ Kg. 7A = Kp.. The dynamics the signals and the errors can be described using"
(hv P),
77 1 -vav avv Kg
@V. = av. + -axa)) aa
novor - -11 KrKo
$ However, cannot ensure A av since a av is not contractive toward the end of training due to the magnitude of gradient noise. This issue can he fixed by final rounding stage (See Appendix 'Here AA.. - h because each gradient step is changed to the {j l-rrguularzeed gradient. It does not change the main parts in the population process.
36

---

(ha.r.
0A.++1 = F.. +1 (1 avaaa ev naaon 11 KpKQ
and
(AA. (A KpKQ
82 A,r+1
= -1- nI-nay mmo²) 8¹ 2n/ 1 -nnay -nov
| 1
11² (hA.T. 2)
K.K
K2Ke
2y V.++1
1 a' + KoT
6,2 )
67 + 2m
1-
an
1
+ 62
(Av,hv),,
KoT
KpKe
(P. lv 3
+ +1 |1 11
K2 0 K
Kp
E.2.1 Lemmas for the dynamics
Before come to the final convergence analysis, first simplify the dynamics with some basic lemmas.
Lemma E.5 (Dynamics the errors). For the errors, we have
h
hA
82 A.7+1 S exp (-2na 62 3n
5A
KP Ke
K Kpvko
| ll VKrKo
11hvl1
82 v.11 s expp (0022000) o2 + 37
σv + 11
VKpKQ
Proof. First, we write
62 6, =(1 -1-nna? - -163) 83 -2n
-maa -mo²
KrKe
|
(hA.T. 2)
1/
+1
K2 K2
K2.Ke
5 |
nal )6%)) 2
6% 201
+3n2
Kryke
K2.Ke
|
|
$ (1 n 182) 0% + +3
ŌΛ Kr Ke
KryKe
For the first term, we have
(1-naš - 182) s exp -111)- s <exp( exp(-2nur)
Thus, for SA, we have
|
hA
22 S exp (-22aaa)) 8n. + 3n
SA + i]
KP Ke
KpvKe
37

---

Similarly, for Sy. we have
2 oyy 33
ōv 1| l1u
!1hv 11 k Q
8 V.r+1
1 11-1 aa Kot 82)
3n
vKrke
|| ||4r||,
sexp (-2po²) 24 317 óv + vKpKel
vKpKe
Lemma E.6 (Dynamics of a - av). The difference between the signals evolves as follows
(av,r+1 -AA.+1)) exp (-2n (1 ((av))) (av (av- ax)
|/ |
1 KOT
81
+ +64 6
VKpKQ
KpvKe
Proof. First, write
av.r+1 a.A.+l av.r dA.. 17 (1 (xaa)) (av axa)
Kg T navo - (hv, P)
(hA,T
- av
+ na na 10x+ + KrKQ KrKę
=: (1 17 (1 - ((aa)) (av - aay) + Tmp.
Therefore, we have
- - ( (ay aa) + Tmp)²
(11-n(1-aaav))) (av aa)² (1 (av aa Tmp Tmp?
exp (1- (øv ax))
+ 3/avv - - Tmp?.
Then, for Tmp, we compute
Ibv.rlle | VKrKo Kpvke
0.5m Tmpl I KoT
5² +6,
In particular, this implies Tmpl s 1. Thus, we have
(av, aA,++1)2 exp (-2n (1 axax)) (av- 22)) 4/Tmpl
$ exp (- (1 (( (av- - a
|| |
8n
+
KoT
VKpKo KpyKo
Lemma E.7. Suppose that both (r,, A at most 1. Then, we have
@V.r+1 CAA.+1
(1 ov CEA)
av + CEA
0V + O'A
( (82 + 64)
+
2
2
17 [yy .rl,
h
+ 2 VKrKQ 2 K KpyKo
38

---

Proof. First, we write
dv,r+1 + dA. = av.r + aA, + (1 (xar)) (av+
1 uy KΦ novo -na n (hv.r: (hA 2),
KpKe 11 KpKQ
0V, dA,T 1/TmPsig Tmprr
For the signal growth, note that asav S (a + a 2/2 & (av + aA)/2. Hence,
(1 (ravv) (ar + ara) 2 (rV + aa (2 ay aA)-
For the error terms, we have
|
(hv, P),
(A,.
|vrrlly
KpKę
KpKQ
VKpKe
Kryke
Thus, we have
av.r+1 + 2 av.r + aAr +17 (av +a
(1- av +
5% o%/
2
|nv.rlly |A..
VKpKe
Krvke
and, therefore,
av.r+1 #...++
-1 ov + (EA
(1 av + UEA
av + +
(83 + 6%)
+
2
2
[yy :ll
h
+ 2 vKpKQ
2 KPvKo
Lemma E.8. Put EA aλ and av. When leal 1/2, have
(6A,4418y,, z exp (-2n) (EΛ + 8v)2 87/8A ValEyvy 16 o
KoT + 03 + + |/h,rrly, |
+ 8ylEA + 8vl
vKrKo Koyko
Proof. Similar to the proof of the previous lemma, we write
u41 (AA,++ = av,r U'AT + 17 (1 axaav) (uv+
(hv.r.P),
(hA.T. 2) KrKQ
" 1 - av myovoo bo
KQ T
Krke
=: av. R.,r + TTTmPsig Tmpen.
For the signal term, we have
Tmpsig = (1- (1- EV)) (2- EA εV)
(2 88 8y) (5A +8v) EAEV (2 EA -v).
For the error term, we have
+ 252 25% 21|bv.rlly 2 2|1hA-1
|mparl 5 KoT
vKrKQ +
KpvKe
39

---

Combine these together, and we obtain
EA,T+ + 5V,++ = EA,T + EV, 17 -22-6,8 EΛ 1.4-64 Ev) (EΛ +Ev)
hy hy | 1
+ NYWYY E2-EA-EVE - + 2n
1
+ 80a +2
+ vKpKo kpyke
KoT
Tmp.
Thus,
(EA,T+) + + =(1 =(1
11 8Λ - (EA + +v) Tmp)² exxp (2 EΛ 8v)) (EΛ + +v) 2|Ea Tv||Tmp| Tmp?.
Note that
|
+ +82 g |y,,ily
Impl 4n EAEV KQT
vKpKe
K kevRe
Recall leal.levi 1/2. Thus,
(EA,T+1 Ev. $ exp (-2n) (EΛ + Eyy
|r
|
+ +nylEA + tvv
BΛEV
+ +62 0
KOT
VKrke
KpvKo
Ihv.elly I|aa.Ilyy
2 gEE+ + 1617²
+ +8,
A +0,
Kot
VKrKo
Kpvke
s (-2)) (E) +Ev)² + +
+ + + E6
|Iv.rlle ||
1 KoT
+ 188/88+8p
+ by
VKpKQ KpyKe
Lemma E.9. Suppose that (X+), satisfies X7+1 S se-Ax, + B for some A E (0, B 0. If
Xo 5 2B/ A. then we have X S 3B/ A for all 0. If Xo 2B/A, then we have X, 3B/ for all
т 2 6 log
XaA 3B
A
Proof: Since e 1 for A E 1J, we have X,+1 S X, - AX+/2+ B. Hence, whenever X, 2 2B/A, we will have X,+1 $ X,. Moreover, X, < 2BB/ A, we have X s 2B/A+ 3B/A. This proves the first part of the lemma.
Now, suppose that Xo 2B/A. When X, 2 3B/ A, we have
X++1 s X, -X,/2 X, AX,/6 A/6
Thus, takes at most á log analysis applies.
XoA 3B
steps to reduce X, from Xo to 3B/A. A. After that, the previous
F.2.2 Main lemma of Stage
We split the analysis of Stage into two substage. Let ca E (0,0. be a small constant. Define
Tz.1 == inf { 2 (av,r + aA,r) 2 1 1- -caj.
We call (r F2.1} stage 2.1 and {T T2,) stage 2.2. For notational simplicity, we define
||tv.rly vKpKo
|
Errsmall max
KºT
Kr
Kryke
Note that this can be made essentially arbitrarily small by choosing large enough batch size.
40

---

Lemma E.10. Suppose that the following hold at the beginning of Stage (after thresholding and projection):
(a) av, A 2 (2))
(b) 2 + oy s (6(2))2.
Let ö, be our target value for d and öv. Choose A as in Lemma E.4. Suppose that
max (2), Errsmall
Errsmall & 6,, Erissall & (a)!)))
max (av.n 2 Emall S log(1/6.)
Then, within (11((pa(2)) log(116.)))g steps, we will have 8xx, o2 s d² and av, (0.9, 1.1).
Remark. Note that our conditions a and 5(2) are much weaker that what one can obtain from Stage This allows us apply the analysis here transfer learning.
The following proof should be treated as a large induction argument though we do not explicitly write down the induction in the proof of Stage 1. In particular, will show (by induction) that the approximation errors 5A and ōy are small, so that most of the naive bounds on the entries can be transferred to A and V. In particular, IV = O(1) for all E [N] that our bounds
and
in Section E. are valid, and |a; ( O(1) for & E [N], e [N], which implies that after projection step, a, (k) O(1T)) for all (R))
Proof of the Lemma E.10
Common results for Stage 2.1 and 2.2 First, we prove some basic results that hold for both Stage 2.1 and 2.2. First, we show (by induction) that av,aa 2 (2) and 8y + 5 S 5(2) hold throughout
Stage 2. Recall from Lemma E.5 that
RA
82 A,r+1
exp (-21at))
32. 3n SA + 11
Kp Ke
Kpvke
$ eexp (22)(a 62. + +Errssmlll-
Hence, long Errsmall (0(1)(2)((a
ensure 62 2 + (6(21)2 always hold.
Stage 2.1: signal growth By Lemma E.7, we have
0V,r+1 + +.A,T
(1-11 av + +
(1 av +
av aA
(8 +7 Err
+
2
2
2
When (av + .2) Ca and av, av.as= (2), we have
" av UAA (1- uV UAA
2 2 (V Ux
Ca 2 na (2)
Hence, as long as Errsmall S a (2) (2) and ((2) < c.//,, we have
(V,++ + @'AT+ < (1 n av +aA
av aA
na
(1
a'v + a
s exp
2 2
2
2
2
Thus, stage 2.1 takes at most ((1 (2)) steps.
41

---

Stage 2.1: difference between the a's By Lemma E.6, we have
(av, (dv. a & - (ay 81/ (6)2))2+ Errsnall)
By the AM-GM inequality, we have ayav ((ax+av)/2)² (1-ca)² Therefore, 1-D 2 Ca and the above inequality can be further rewritten as
(av,r+l @A,,+1)2 & exp (-2c (av (av- + 87/ (1)12)22 +Errsmat)
Thus, by (the proof of) Lemma E.9, we have
(av,r - 2(av, a...) 4 ((6(2))2 + ErTsmall
Ca
Stage 2.2: error decrease
By Lemma E.5, we have
Ar exp (-211r2) 6 + 31 max (2(2) Errsmall Errsmall
exp (-1) 8,, + 3n max 6(22) Errsmalll
Emall
Thus, by Lemma E.9, we have
2² A.r (1) max (0(2) Err Emr $ 82
1og(1/0.)
VT 2 0
17
In other words, Stage 2.2 takes at most O log(1/6.)
steps.
Stage 2.2: stability of (ry + aA Recall from Lemma E.8 and Lemma E.6 that
(EA,r+) ++ s exp (-2n) (EΛ + 8y) 87lEA + eylEΛVy 161 E
+ +8ylEA + | (2) + Err
- - & (-2) (xAur)) (av a ara) 2 817 (6!2) Err
We wish maintain the induction hypotheses Jav a.l lev EAl 0_ for some = ((1) and
(a + a2) N c (18.) =: 1 0..
First, assume these conditions are Then, have
(av,r+1 - aAAT1) s (op())))/log((/)))) (ay- - xx) 8n (8(2) Err
Since Stage 2.2 takes most (l0g)((16)))))) steps, when the constant e is small, we have
(av,r - (dv,r-aAr)) < (av.9), - 2
- +8 (6(2) EErr
s 0(1) max (av.7 - -
-AAT% ,.8 (2)) Err =: 0_.
We can choose the parameters appropriately so that 0_ o(04). Then, when EΛ ev (-20+. -0.). we have, for some large universal constant C O,
(EΛ,111 + exp (-2n) 03 Cno Cn220 CnO,, (8 + Err
$ exp (-2 0 C + Cn0.0_
02
This establishes the induction hypotheses ay t aA.
42

---

Stage 3: Final Convergence
As mentioned last section, cannot ensure a av since O dt is not contractive toward the end of training. However, we can add a final rounding step and then continue train the model to recover the ground-truth with s-error.
First, we formally define our rounding procedure. Let > a small constant (cf. the proof of Lemma F.1). For each & [N], define
a
(kk) a;
a
2 c/0
re(r) and (+..k)
1 (k))
This A' := [[ [<[[N] is our rounded version A. For the error between A and Q. we have the following lemma.
KE[NN]
Lemma F.1 (Rounding A). Let (0(0)/T2), 0.1) be our target accuracy. Suppose that a = 1-EA for some |&Al s 0.1 and || I s 8x for some 0 < δΛ < 1/(QVN). Then, after rounding, we have
Q2
| Q1I, s 0
T2 +
In particular, achieve accuracy terms of I1- we only need | 0.1 and 8 O(E/vNQ).
Remark. In particular, this lemma implies that as long as & not too large, after rounding, the error depends solely on ||
Proof. For notational simplicity, we omit the superscript & for now. Write
Aa (1 - =a)g E Ag. a = a= avg+ (1 - av) _ T
Note that ||all & |l s 0(1) N | S 5O(NNÖA) Since all nonzero are lower bounded by (21/2),, we have, for any s q and q,
a,
2(1)
TEAl
0
(VNOA)
2(1)
Q
T
[EAl
|a,1 S
0
(N88A
T
0
Hence, we can choose small constant c 0, that
xx) + EA 1g
à = [a, [1arl 2 c/QilemTT =1-
[Aalq.
where for r E R Fg R is defined as [v,[(r e q)lerrr here. Now, consider the difference between y and à. We have
EA_ 1d=11-E4+
+ 0 (eVNNo)).
T
and therefore,
(1 8A4 8A 1 7 [Nalg
a = LTà
(1 (1-EA) x 0 (QVNOA)
(1- +
[Aal9
(1~EA)
(q = O (EALy/T (Aalg) (1 (1==0<x0//+=VNO/)
= 02 (xQ/T VVNOAA)
veNoA
Thus,
N
II, - Q11, Σ
Uk a
(R) - q (2))
= 0
+ 82
NO
72
A
43

---

Lemma F.2. Let & E (0(1)((K77)0.1) be our target accuracy. Suppose that ||4 211 1 SΛ $ 0.0læ and |l 5 ōv s 0.01 the beginning of Stage and || 1l cevKpKo for all and sufficiently small constant c. Then, we have ||V - P s & for all T 2 722 (((og(1/e)/m)
Proof. Under the condition llA 2l1 s مA,e, we have ||All s || - 221 s ÒA.- and
(A.2),
1 T Ke
1 -
Ke (Ax. 2), o (VēōA.).
aA= Ke
Recall that we only train V in Stage 3. Note that by Lemma C.3,
I|V - PIl = | P| |llll = (1 - all-v) 1l S (1 - av)² + ||Av
Hence, to get e accuracy. it suffices to have (1 - av)² 2 8/2 and | s 8/2.
First, for ev, by Lemma C.5 and fjv = 17/ Kg. we have
av, r+1 ov - + livee (1 avvaaa av + T
1- a
naav |Aal
T
(hy, P)n
T
Kp
(hV,T:P). P) KpKQ
1
or + (x) ay n
+ +
+no
Kot
Hence,
Phy
(1 - av. (1- (1 uv + n0
+ +10
KoT
Kpke
(IIV P))
1
I exp ( (-2nav) (1 - xx) + n0
+ + +1o
KoT
KpKę
s exp (-1) ar ov.r)² 2 0.1
For [| by (the proof of) Lemma E.5, we have
|| h | exp (-21a²) (/ae)))\vvvv/v ||Av. + 3n ||Av.rll, 'KpKo VKpKo
||hr||
(-1) |v. 0.1.
Thus, by Lemma E.9, have (1 ary 2 8/2 and |lAv 8/2 T (log(1/8)/n).
Corollary F.3. Let (O()) 0.1) be our target accuracy. Suppose that a E (0.9, 1.1), [|AIll, s ((v))) and || s 0.01. Then with poly(N, Q, 1/8) samples, we have with high probability that || @ and ||V P1 I after Stage 3, which takes ((log(1/e) steps.
Proof. It suffices to combine the previous two lemmas, the concentration results in Section B, and apply union bound.
G Proof of the main theorem
In this section, we combine the results from the last three sections and prove the following formal version of Theorem 3.1.
Theorem G.1. Let 0 be target accuracy and Ti min(r ( 0 max(av. A.A... ) 2 O(1((ON))) We can choose the hyperparameters in Algorithm O (Tom (2) 1og(1/s)/n) steps, we have - s 6 and I|V Pl s e with probability at least and the number of samples used before and after T are poly(T, 8) and Poly(N, Q, 1/e, log T, 6). respectively.
such that within
44

---

Proof. The results for Stage follow directly from Lemma D.I. Now, consider the results for Stage 2 and 3. First, by Corollary F.3, it suffices to make sure time 72. we have aa (0.9, 1.1), ||Алl[ ((e((VV)) and ||Av f = 0.01 (with high probability). By Lemma D.1, we know the following hold at time Ti w.h.p:
a av = (II((N)) and
||al |lAvll S O(1/T).
Therefore, by Lemma E.1, we choose the threshold be @(1/(Q2N)), then after thresholding and projection, we have a, = O(1/T) for all & ( Thus, by Lemma E.3 and Lemma E.4, with poly(2. NV 1/e, log(T)) samples, we have with high probability that a O for all $ q (8) holds and | satisfies the requirements in Lemma E.10 throughout Stage 2. Thus, by Lemma E.10
with o? ((e/VV)) at the end of Stage 2, we have with high probability that ax (0.9, 1.1), || (s<(((NN) and ||Av 0.01. When combined with Corollary F33, this completes the proof.
H Transfer learning
Lemma H.1 (Initialization). Suppose that have learned P and (P.P), 211 || @(1) [I] Let V= of (1 -1-0)) 01 for some 0 (0.1). We have
0 NKP ((0)
av θ
and ||v ((1)0
Proof. First, consider Av. Since ||Av I is the distance to a projection, we have
|lvi| s V|V (OP (1 01- )1; -
0² -|p- P|] 0(1) |lll 0(1)
For av, we compute
(V.P),, - 11112 = = -1ul11²) - 20 0(10/N.
2
In particular, this implies av 3
e NkP
Lemma H.2 (First gradient step). Let A 11 and given by Lemma H.I. Let Etpp O())(NKp)) Run one gradient step with poly Q. logTT T, 1/6mmp) samples and rj = 1, remove all entries of (k) with |a,x)| s @(0/(Nkp)) and the replace (k) with the projection (I 11 (T) Eeg a
With high probability, we have a (1 (ermp))) and ||AIl 0
Proof. The proof idea is essentially the same as Lemma E.3, though the reinitialization of the first layer allows better estimations in several places. Recall that for each E [T], we have
4 1(x7+1 (Vex.) (VX1/T- - ex...
B a m = mr kv.KvTT H4KvrQs
Also recall from Lemma C.3 that Kvp Ky and Kv az, Kr p + ||Av have
For any s E (k), we
- E un Mk avkrqs Kz ||Av|??
MRarkKp4
T
2
where the inequality comes from TyavKpga Q(av/N2/O0 > 1/T. Meanwhile, for any s [T]. we have Eö u = (1/T). as
45

---

Meanwhile, for any s E [7) we have
7 Σ exxl $ T [=1
0 00 = (Vexx) "VXI/
(V (Vxx,
V.xx)++ /V s O(N).
Thus, by some standard concentration argument similar to the one Lemma E.3, we can show that with (Nly(N, (N, Q. 1/vv 1/6rmp, log 7) samples, we can make sure with high probability,
0 0(8-) 1 = (1 pmmp)
(2)
B p0ll = + ETOP) pLayKrqs
E q
0 B d (k)
0 { 61mp
avkr 0
av
Vr q (R)
Ermp
NQ
NQ
Thus, after one gradient step with 17 = 1, we have
a, (k) (1 EImp EImp)avq (8)
Vr (k)
(k) 0 Tmp Q
Ovv
Vr (8)
a
Recall from Lemma H.1 that ay E [0(0/(NK))), 0(0). Hence, we can choose the threshold Ag to be (NK)) and SrmO $ 1/(KKp) so that
(ak) a 1(a
(k) (k) 2 do} [E[T]
= (1+
(k) 11t (2) [E(T)
Now, set
a (k) - (I (1-1111 / (1 + O (w (avg))) (1) 0 (1)
Note that this implies that after the first step, we have
Ex MK (a (x))
((0)-110 - 1/T
av(1 = (mp 11 0((T))
(1 ( )aa.
K2
KQ
and
|
> uk |a (8) -(a- +11 (1--ay 1/))
50
Tmp
ar
1
Q
T
Theorem H.3 (Main theorem for transfer learning). Let & > be our target accuracy. Consider the same setting of Lemma H. and Lemma H.2. Choose (0(1/log(//e) and Etmp O (og(1/<) Then, after one step of update A as in Lemma H.2, A and V satisfies the conditions of Lemma E.10, and therefore we can learn (P, (P.Q) to & E-accuracy using Poly Q. 1/8, 1/6) samples with probability at least within poly(N, Q, 1/e, 1/5) steps.
I Additional Experiment
For all our experiments, we use Numpy and run on a normal laptop which takes about 20 minutes. Setup. In all our experiments, we choose 7' = 5000, @ 2, N = 3. The architecture is
F(X,XT41: V, A) VX (T ) =: VXa (sr.1)
(15)
and the data model is the SCB (1) mata-generating model. The batch size is B 64 and the regularization hyperparameter is A = le-5, The total time is T = 1000 iterations where stage takes r E [0, [0,4 with learning rate 7)1 = 0.01. After r 400, we use = 0.005 for further improvement (stages 2 and 3).
Hyperparameter selection. Due to limitation of computational resources, we do experiments with N E [3, [0] for real-world batched gradient experiments, and N 500), 7 = 100000 experiments by using Gaussian noise SGD simulations based on the dynamies of Lemma C.5 and C.6. As T needs scale with N polynomially, it would be beyond computation capability to
46

---

experiment with larger N. As for other hyperparameters, à is chosen based on our theoretical results (Theorem 3.1 and G.1):
O(cKp/O N*) Lemma E.4. The batch size can be chosen from standard [64, 128, 256). while smaller batch size will lead to divergence for both SGD and regularized GD. 17 is chosen as the largest learning rate without divergence.
Besides the original parameters, we consider the approximation error after the normalization step (stage 3) in real-time. That is thresholding and normalizing the attention block
Vk E [n]. (k) = [aa (a, 2 22112)))r- (kk) - a (k) / a
and do further gradient descent to recover V. In this case, will directly use the linear regression solution on population loss for V.
Here we report in addition: (1) original signal projection on the population process trajectory a and the distance to the trajectory Av.Av. Ay. (2) The approximation error/similarity before and after normalization for both SGD and proximal gradient descent. We conclude that in all metrics proximal gradient descent performs better than the vanilla gradient descent with a small batch size (when the noise is large).
0.8
$20
0.8-
500
0.6
550
Detance to projection
560
Signal Procc Proection
0.6
11.4
0.4
0.2
0.2-
0.0
0.0
200
000
600
800
1000
200
2000
Iceration
nteration
Figure 3: Signals A, dv and the distance to population process dy. For the SGD, the distance to the population process of the attention matrix A keeps growing and dominates the signal term. That explains the failure learn the correct attention pattern, which leads to saturation of the signal. In comparison, our proximal methods dramatically help reduce the gradient noise and keep close to the population process. Though ||A eventually grows up due the bias of the gradient estimate (the original signal growth also slowed down), after normalization it can still approximately learn the correct pattern. Both ||Ay|] stay small empirically.
1.0
1.0
Simitrity (Prximal)
Sim loriss (56D)
0.6-
0.5
-... (P.P) right
0.2-
0.2
(4.5% orgmall
(P.Pi,
200
400
6000
800
1000
0
200
400
900
1000
Deration
Ieeation
Figure 4: Similarity with the ground-truth. The figure shows after Stage normalization helps further improve the solution of the proximal method. Meanwhile, with or without normalization, our proximal method always outperforms the vanilla SGD, which fails recover the ground-truth.
We tried different orders of state number N and show that E regularization is necessary and outperforms SGD when batch size is small (gradient noise is large). Due computation limitation, experiment with real batched gradient on x 20,7 T s 5000, and do SGD simulation by combining our population gradient Gaussian noise (to mimic the batch gradient noise) for N 500, T 100000.
We also corroborate our previous experiments with the new test loss plot to show the convergence of training. Note that since there are multiple global minima for the línear attention, & regularized dynamics eventually will make and Q deviate from the ground-truth while representing the same function. That is why the loss converges but the distance to the ground-trun increases after some
47

---

point, making the final normalization step essential to recover the ground-truth. Another point is that according to our theory, the regularization will eventually distort the learned pattern when trained for too many iterations. Empirically, the loss also increases little after it converges. Therefore, we must stop early and normalize before the distortion happens.
0.31
6.25-
132
330
Tonadan
Raration
Iscation
0.25
die
1.33
0.15-
2.0
553
2.0
1200
290
600
1000
1058
10:00
24
080
750
1350
Handon
1.2
3.424
3.3
8.75
823
3.0
2000
12.90
3305
2000
336
3.0
10.00
-300 2000
2540
3000
140 1000 1003
386
Figure Convergence analysis. We plot the distance to the ground-truth and the test loss for N = 3, 10, 20 (from top to bottom). It shows that when gradient noise is large, E regularized algorithm with normalization and early-stopping can almost perfectly recover the ground-truth (the star), while SGD struggles to learn the target function. Figure 3 (Appendix I) also shows when the gradient noise large, SGD never learns ground-truth @ even with normalization.
Sap
3.4
04:
150
Figure 6: Simulation with larger N and T. We simulate the SGD/C, regularized dynamics by replacing the batched noise with Gaussian noise in the dynamics formula in Lemma C.5 and C.6. The gaussian noise variance scales with the inverse of batch size. The experiments show that the conclusions drawn from the small N cases still hold in those simulations: when T 100000, N 100/500. our l regularized algorithm can recover the ground-truth since the distance to the population trajectory (Ax. Av) stays very small, while the error along SGD trajectories quickly increases with the same batch size.
48

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We accurately summarized our claims and contributions in the abstract as well as the introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings,
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer:
[Yes]
Justification:
is included in Appendix A.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't t acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
49

---

Justification:
The assumptions and conditions are included in the Setup section 2.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification:
All information is included in Section 5 and Appendix I.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully specific model and empirical evaluation, may
might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g.. large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
50

---

Answer: [Yes]
Justification:
We will upload the codes in the supplementary materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https: battps://nips: /ttp///ipsps.//eee///////eeeee.s cc/
pub public/ guides Cat//CooeSSSeee Submi ssionPolicy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: //niss.///ppp./mm cc/ http///blcccggudeeeeecc/e// Sea.ddee/oooeuuuee wwm.isssiiiiiiyy. for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer:
[Yes]
Justification: All information included Section and Appendix I.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Our experiments are for dynamics simulations. We do not include perfor- mance/accuracy results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
51

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer:
[Yes]
Justification:
All information included Section and Appendix I.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https:/ /neurieiigge ps. cc wtcp/pbllc///eeteeeeccee www.://ethhhheeeeeeeccoeccc www.eeeeeeeieieeeooeoooo.o...eee dellddinnnnn line?
Answer: [Yes]
Justification:
This work is pure theory paper without potential harmful effect.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there is a special consideration due to laws or regulations in their jurisdicioo))
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
This work is a theoretical paper without direct societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
52

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper is theoretical work without such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification:
The paper does not use existing assets.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, papers wt. thoode com/ daaaaaasaseeee has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
53

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
54


Reviews:
[{"review_id": "ZCQHCjv7J7", "reviewer": "Reviewer_yR2y", "summary": "The paper studies a natural problem that generalizes the bigram model for language generation. In the bigram model, each token $x_t$  is only dependent on the previous token $x_{t-1}$. Therefore, to learn to predict the next token, one only needs to look at the bigram frequencies $P[i,j] = \\Pr[x_t = i \\mid x_{t-1}= j]$ which is simple to learn. The authors generalize this in two steps:\n\n\n1) First there is a mapping $Q: [K] \\to [T]$ that maps each token value to a sparse distribution over the indices. The algorithm picks an index and looks at the value at this index.\n\n2) Then they use a bigram model $P$  to predict the next token based on the value chosen in step 1. \n\n\nThis is an interesting model and can capture natural tasks performed by transformers as they process long contexts. It requires the ability to learn the transition matrix $P$ as well as the mapping $Q$.  The authors try and show is that simplified linear model can also learn this efficiently in $N,Q$. This effectively means that they take one layer of the attention matrix and and remove soft-max and the non-linear activation in the final MLP layer. Finally they use the one hot shot embedding which greatly simplifies the model further. This results in a very simple form:\n$x_{T+2} \\approx VXa^{x_{T+1}}$ where $V$ is the trainable vector and $a$ is the $k^{th}$ column of a trainable matrix $A$. The key is to understand how this will be transformed using standard gradient descent and a mean squared error loss. The main result is that with the right initialization as well as a staged use of proximal gradient descent, one can learn this model using linear transformers.", "questions": "NA", "limitations": "NA", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well written and the assumptions are clearly stated. The analysis is mostly in the appendix but from what I have glanced at, it looks thorough. The bulk of the analysis is to deal with the noise introduced by the empirical gradient from having a small batch.  The model itself is a nice tractable model that still captures a nontrivial task that we observe performed by modern transformers.", "weaknesses": "Naively, one can imagine a much simpler model such as Hidden Markov Model or other sequence being capable of learning the sparse bigram model. The main interest in the paper is that a sufficiently simpler version of transformer can \"provably\" learn such a model.  The assumptions do seem necessary to make the problem sufficiently simple and hence approachable. However this means that the actual algorithm and model is quite far from what is used in practice. I would have liked to see a bit more empirical work to see if the introduction of non-linearity greatly reduces the sample complexity. In particular, I would have liked to see some empirical work showing that the actual attention matrix has patterns that are similar to what is learnt in the linear model itself.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your reviews! We'd like to answer your questions as follows.\n\n## Impact of non-linearity on the sample complexity\nWe discuss the relationship between softmax transformers and our linear model in the global response. Toward the \nbeginning of training, where the sample complexity is potentially bad, softmax transformers and our model are approximately\nequivalent, up to a change of learning rate. As a result, we do not believe that simply adding non-linearity can \ngreatly reduce the sample complexity. However, in Stage 2, softmax transformers will automatically implement something \nsimilar to what $\\ell_1$-regularization does in our setting and will enjoy the same improved sample complexity. \n\n## Attention pattern after training\nWe add plots of the learned attention patterns after training in the PDF file attached to the global response. In short, \nthey match what our theory predicts and recovers the $\\mathbf{Q}$-matrix. The pattern that linear and softmax attention learns are very similar, and converges to the solutions with similar losses.", "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": "Q3OhYX1OMs", "reviewer": "Reviewer_fyHi", "summary": "The topic of this paper is understanding transformers. Since this is an ambitious goal, the authors make some reasonable simplications. Specifically, they study the training of a **one-layer linear** transformer on synthetic datasets generated by a novel data-generating model, called **Sparse Contextual Bigram (SCB)**. SCB extends traditional bigram by also conditions the next token generation on prior earlier positions determined by the last token. Such design of SCB allows study on both contextual (last token conditioned early token dependence Q) and global (transition matrix P) knowledge. \n\nBased on this data-generation model, their analysis shows the the one-layer linear transformer training process can be divided into two stages: an initial sample-intensive stage followed by a more efficient one. In transfer learning scenarios, the first stage can be skipped if there's significant correlation between the pretraining and downstream tasks.", "questions": "- See weakness.\n- What's $I_T$ in equation 2?", "limitations": "They discuss the limitations in appendix A", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "- The paper is well-scoped. The structure is easy to follow. The authors clearly state their simplification in order to analyze the thus complicated transformer training process.\n- The proposed data-generating model, SCB, is novel to me. It extends the classic bigram in a smart way so that it encompasses both contextual and global knowledge.", "weaknesses": "- Architectural simplification. The simplification of transformer into a one-layer linear version could make the conclusions less useful or even invalid. Particularly, the softmax over the attention weights and the layernorm before the output are known to be important for the impressive empirical performance of transformer. It would be great if the authors can explain how their simplification of architecture can impact their conclusions, either theoretically or empirically.\n- Experimental setup. The experiments in Sec 5 are with a vocabulary size of N=3. Such a small vocab size is far from practical ones (~32K-250K). I doubt the conclusions on the stages would still holds across different orders of vocab size. And it is unclear to me how the hyper-parameters are chosen from line 321-line328. It would be great if the authors can explain why they use this specific setup instead of a more systematic investigation of hyper-parameters.\n- Complexity. What is the complexity of the SCB data-generation model? How does it differ from a real natural language distribution?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedbacks. We'd like to address your concerns as follows. \n\n## Meaning of $I_T$.\n$I_T$ stands for the $T$-by-$T$ identity matrix. We are using one-hot positional embeddings \nand attending only to positions here. Hence, the second $E$ is replaced by the identity matrix. \n\n## Architectural simplification (linear vs softmax)\nWe have discussed the relationship between our linear model and softmax transformers in the global response. In short, \nsoftmax transformers and our linear model are approximately equivalent up to a change of learning rate near \ninitialization, and in Stage 2, softmax transformers implicitly implement something similar to what \n$\\ell_1$-regularization does in our algorithm. In addition, analyzing the linear counterpart of the model before attacking \nthe more difficult practical models is common in the development of DL theory, and in our case, this is also partially\njustified by the empirical observations in [1]. \n\n## Architectural simplification (LayerNorm)\nIndeed, LayerNorm plays an important role in the empirical success of transformers, as the pre-softmax is usually \ninterpreted as the correlation of two tokens and one may wish to remove the effect of scaling via normalization. However,\nwe wish to stress here that in the theoretical analysis of shallow transformers, LayerNorm is not as important as it is\nin practice and is often omitted to simplify the problem (cf. [2-4]). In particular, in our setting, the token and \npositional embeddings are fixed one-hot embeddings. They are automatically orthogonal to each other and have the same \nnorm. In fact, our model is also closely related to the associative memory point of view proposed in [2], which \ninterprets (the product of) the QK matrices as a dataset and the token embeddings as the key. Indeed, if the token \nembeddings are one-hot, then effectively they are accessing the corresponding entry in the QK matrix and no normalization is needed in this case. \n\n## Experiments\nThank you for your constructive questions on the experiments. Though the experiments is to numerically verify the theoretical results when **$T\\gg N$ and the batched gradient noise is large**, we agree that a more systematic report should be added. We have added \nexperimental results for different $N$ in the pdf file attached to global response. \n\nIn this work, we focus on the regime where the sequence length $T\\gg N$ with a relative small vocabulary size $N$, which is also used in related transformer optimization works (cf. [5]). We added new experiments for small $N\\in [3,20]$, and simulation results for $N \\in \\{100, 500\\}, T = 100000$ by using Gaussian noise SGD simulations based on the dynamics of Lemma C.5 and C.6. In all cases our theoretical conclusions still hold and $\\ell_1$ regularization is necessary for this algorithm.\nAs $T$ needs to scale with $N$ polynomially, it would be beyond our computation capability to experiment directly with larger $N$. \nAs for other hyperparameters, $\\lambda$ is chosen based on our theoretical results (Theorem 3.1 and G.1): $\\lambda\\sim \\Theta(\\epsilon K_P/Q^2N^3)$ in Lemma E.4. The batch size can be chosen from standard $\\{64, 128, 256\\}$, while smaller batch size will lead to divergence for both SGD and regularized GD. $\\eta$ is chosen as the largest learning rate without divergence. \n\n\n## SCB model\nI'm not sure whether I understand your meaning of the complexity of the SCB model. The purpose of this model is to \ngeneralize the bigram model to capture a certain linguistic phenomenon (in a simplified setting), that is, the next token can depend on a small piece \nof text that is potentially far away from it, and the location of this piece depends on the recent context, instead of providing a \ncomplete modeling of the natural language. \n\n\n\n[1] Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). (2023)\n\n[2] Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, Leon Bottou. Birth of a Transformer: A Memory Viewpoint. (2023)\n\n[3] Yuchen Li, Yuanzhi Li, Andrej Risteski. How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding. (2023)\n\n[4] Ruiqi Zhang, Spencer Frei, Peter L. Bartlett. Trained Transformers Learn Linear Models In-Context. (2023)\n\n[5] Yuandong Tian, Yiping Wang, Beidi Chen, Simon Du. Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer. (2023)", "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": "ubPoflySxc", "reviewer": "Reviewer_Qmma", "summary": "This paper proposes a new data model, Sparse Contextual Bigram (SCB), to study the training dynamics and sample complexity of transformers. The paper analyzes a one-layer linear transformer trained on data generated by the SCB model using a gradient-based algorithm. They prove convergence guarantees and provide bounds on the sample complexity, showing that the training process can be split into two stages: an initial sample-intensive stage and a more sample-efficient stage. The paper also includes results on transfer learning, showing that pretraining on a related task can improve sample efficiency.", "questions": "* How do you think the results would change if a more complex transformer architecture (e.g., with multiple layers or non-linear activations) were used?\n* Can the results on transfer learning be extended to cases where the pretraining and downstream tasks have different structures or distributions?\n* Could you provide more intuition behind the proofs, either in the main text or in the appendix?", "limitations": "The authors have adequately addressed the limitations of their work in Appendix A. They acknowledge that the model and algorithm studied are simplified compared to those used in practice and that the empirical evaluation is limited to synthetic data. They also mention the need for further research to extend the results to more complex architectures and real-world datasets.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The proposed SCB model is a novel contribution that simplifies the analysis of transformer learning while still capturing essential aspects of language modeling. Especially the fact that training can be split into two stages with different data requirements and the implications for transfer learning are quite interesting.\n* Although I am not familiar enough with these type of theoretical analyses to fully verify the correctness of the proofs, I found the paper relatively easy to follow.", "weaknesses": "* The analysis focuses on a simplified linear transformer model and a specific gradient-based algorithm. The generalizability of the results to more complex architectures and optimization methods is unclear. This limitation is acknowledged in the paper, and I don't think it's reasonable to expect a theoretical analysis without any simplifications. Hence, this might not actually be a weakness, although it is a limitation.\n* The paper could benefit from more intuitive explanations of the proofs to make the theoretical results more accessible to a wider audience.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your reviews! We'd like to answer your questions as follows. \n\n## More complex architectures\nWe discuss in the global response on the relationship between our linear model and softmax transformers. At a heuristic\nlevel, they are approximately equivalent, up to a change of learning rate, in Stage 1, and in Stage 2, softmax \nautomatically implements something similar to the $\\ell_1$ regularization in our algorithm. It might be hard to extrapolate\nour results to deeper transformers rigorously. However, one can perhaps expect our message on transfer learning to hold \nat an intuitive level for deeper networks. That is, as long as pretraining can find a point that is positively \ncorrelated with the downstream task, which is easier than learning representations that can be directly used in downstream\ntasks, then the sample complexity can be greatly reduced. \n\n## Transfer learning with different architectures/distributions\nOne of the main benefits of our signal/correlation point of view on transfer learning is its flexibility: when the \narchitecture/distribution changes, the needed representations might change as well, but in most cases, one should \nexpect them to have at least a small positive correlation, and this is all we need in our theory. \n\nFor example, in our setting, we only need the transition matrices of upstream and downstream tasks to be reasonably \npositively correlated and impose no correlation constraint on the $\\mathbf{Q}$-matrices. In addition, our results can be \ngeneralized to cases where the transition matrices have different but reasonably correlated stationary distributions. \nThe key here is when initializing the model as in Theorem 4.1, if both stationary distributions are well-conditioned, then \nthe error should scale with $1/N$, and this amount of errors is allowed in our analysis (cf. line 654).\n\n## Intuition on the proof\nWe will discuss more on the intuition of the proof in the revision. Here is the general idea. The key equation in \nour analysis is line 654-655. If there is no noise, then it is a linear system with coefficient matrix of the linear \nterm being $[[0, 1], [1, 0]]$ and drift $[1/T, 0]$. The drift term will provide the initial signal, guiding the \ndynamics toward the correct direction and then the first linear term will amplify the signal and fit the target. \nThe linear term is approximately $0$ around initialization, so we have to rely on the drift term first, and to \ndistinguish this $1/T$-order signal from the noise, we need $\\mathrm{poly}(T)$ samples. When we have accumulated enough signals, \nthe linear term will have order $1/\\mathrm{poly}(Q, N)$, and we can rely on it to learn the target, which costs only $\\mathrm{poly}(N, Q)$\nsamples. In transfer learning, we use the pretrained model, instead of the small drift term, to get the initial \ncorrelation, and this allows us to skip the first sample-intensive stage.", "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": "cidbRlKkpm", "reviewer": "Reviewer_9wme", "summary": "The paper studies the training dynamics of a one-layer linear transformer on the sparse contextual bigram task. Using $\\ell_1-$regularization and proximal gradient descent, they show that the training process goes through a sample-intensive phase and then a sample-efficient phase. They also extend the results to transfer learning.", "questions": "1. Could the authors comment more about the choice of the setting? How does sparsity affect the results? I do not find the parameter Q to have a special effect on the results besides serving as an extra parameter. \n2. Can authors give some implications of their results on the training dynamics with the complete transformer structure?", "limitations": "Although the theory is solid, and some findings are interesting, I do not find any important intuitions that can be gained from the author's discussion. In particular,\n1. The current setting makes too many modifications to the original transformer structure and training practice. Such simplification is OK. However, the results do not facilitate understanding transformers' training under more realistic settings.\n2. The paper lacks experiments under a more realistic setting with the standard transformer structure and training method.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written and organized.\n2. The theory is solid. The discussion of the two phases in training dynamics is interesting.", "weaknesses": "1. Could the authors comment more about the choice of the initialization? Could you show, at least in experiments, the consequence of another choice of initializations?\n2. The loss curve does not look like converging in Figure 1. It would be important to show the loss curve with more training rounds.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your reviews. We are glad to answer your questions.\n\n### Initialization\n\nWe initialize the first layer weights with $a^{(k)}_t = 1/T$, which corresponds to the uniform attention. In practice,\nweights are initialized to be small and softmax is used. As a result, the attention pattern is also close to uniform \nat initialization. In addition, our result is stable under small perturbation at initialization. After all, since we are using \nSGD, some errors will be introduced at every step. \n\nFor the second layer, we initialize $\\mathbf{V} = \\mathbf{\\mu}\\mathbf{1}^\\top$. This is the trivial transition matrix with stationary \ndistribution $\\mathbf{\\mu}$. It corresponds to the transition where the next token follows the distribution $\\mathbf{\\mu}$ no matter \nwhat the current token is. We choose this initialization to simplify the analysis. Another reasonable candidate is \n$\\mathbf{V} = \\mathbf{1}\\mathbf{1}^\\top/N$, which corresponds to choosing the next token uniformly at random. As in the situation of the \nfirst layer, this is what small initialization + softmax approximately implements. \n\nSince $\\mathbf{\\mu}$ is assumed to be well-conditioned, the difference between these two initialization schemes scales with \n$1/N$. Note that in our transfer learning result, the error also has this scale at initialization. Hence, with a similar\nargument, one can show that this uniform initialization scheme leads to the same result. The main reason behind this \nflexibility is that the population dynamics are rather stable under these type of errors (cf. line 654). \n\n### Training curves\n\nThanks for pointing this out. We have added the updated plots into PDF file associated with the global response. \nWe can observe that with $l_1$-regularized SGD can reach a much smaller loss/distance to the ground-truth faster \nthan vanilla SGD. (The distance to the ground-truth will eventually increase as the ground-truth is not the unique \nminimizer and in both cases, SGD will overshoot due to $l_1$-regularization or noise.)\n\n### Role of the sparsity parameter $Q$\nOne of the goals of the SCB model is to model the linguistic phenomenon where the next token can depend on **a few** pieces \nof the text that depend on the recent context but can potentially be far away from the current position. The sparsity\nparameter $Q$ describes the amount of long-term context on which the next token depends. We believe it is reasonable to \nassume $Q$ to be much smaller than the text length $T$, as it is hard to imagine all previous tokens can directly \ninfluence the next token. \n\nThen, it is natural to ask if we can obtain sample complexity that is $\\mathrm{poly}(Q)$ instead of $\\mathrm{poly}(T)$. Our result says \nthat while it might be unlikely to obtain $\\mathrm{poly}(Q)$ sample complexity when training from scratch, doing transfer \nlearning can indeed reduce the sample complexity to $\\mathrm{poly}(Q)$, as it allows us to skip the sample-intensive initial \nsignal-boosting stage. \n\n### Comparison with and implications on more practical architectures\nWe have discussed the relationship between the training dynamics of softmax transformers and our linear model in the \nglobal response. In short, we expect them to be approximately equivalent, up to a change of learning rate, toward the \nbeginning of training, and softmax transformers will automatically implement something similar to $\\ell_1$-regularization\nin Stage~2. \n\nIt might be hard to extrapolate our results to transformers with more layers rigorously. However, we want to stress that \nanalyzing the simpler shallow models can be served as the first step toward understanding the more practical models, \nand this approach is widely adopted in deep learning theory and also the in the study of transformers (cf. [1-3]). \n\nStill, at a heuristic level, our transfer learning result can potentially lead to new explanation on the effectiveness \nof transfer learning and the pretraining-finetuning pattern. The usual story is during the pretraining stage, the model \ncan learn representations that can be reused in downstream tasks. Our results suggest that at least in certain cases, \nthis is not necessary and having some correlation between the learned and downstream features (which is much easier than\napproximately recovering the downstream features) is sufficient to greatly improve the sample complexity. \n\n[1] Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, Leon Bottou. Birth of a Transformer: A Memory Viewpoint. (2023)\n\n[2] Yuchen Li, Yuanzhi Li, Andrej Risteski. How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding. (2023)\n\n[3] Ruiqi Zhang, Spencer Frei, Peter L. Bartlett. Trained Transformers Learn Linear Models In-Context. (2023)", "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:
## Relationship between our model and softmax transformers

As several reviewers 
asked about the difference between the softmax transformers and our linear model, we'd like to further discuss their relationship. In short, they have qualitatively similar behaviors: there will be sample-intensive
initial stage, and when there is nontrivial correlation between the model and the target, SGD will become much more 
sample efficient. In the following, for simplicity, we will assume $N = 1$, write $a := a^{(1)}$, and the ground-truth 
$q$ is $e_1 = (1, 0, \dots, 0)$. Most of our argument bellow can be generalized to the general setting at least at a 
heuristic level. In our linear model $f(X) = VXa$, we directly optimize $a$ while keeping the constraint $\sum_t a_t = 1$. 
By the softmax transformer, we mean the model $f(X) = VX\sigma(w) =: VXa$ where $\sigma$ is the softmax function 
and $w \in \mathbb{R}^{T}$ is the trainable weights. 

Let $l$ denote the (per-sample) loss. Since the Jacobian of the softmax function is $\mathrm{diag}(a) - aa^\top$, 
we have 
$$
  \nabla_{a} l(f(X)) 
  = (VX)^\top \nabla l(f(X)),
  \quad 
  \nabla_{w} l(f(X)) 
  = \left( \mathrm{diag}(a) - aa^\top \right) (VX)^\top\nabla l(f(X)). 
$$.
As a result, the dynamics of the attention weights $a$ are controlled by 
$$ 
  a(\tau+1)
  \approx \begin{cases}
    a(\tau) - \eta \left(I - \frac{\mathbf{1}\mathbf{1}^\top}{T}\right) (VX)^\top \nabla l(f(X)), & \text{in our linear model}, \\\\
    %
    a(\tau) - \eta \left( \mathrm{diag}(a) - aa^\top \right)^2 (VX)^\top\nabla l(f(X)), 
    & \text{in softmax transformers}. 
  \end{cases}
$$ 
In other words, the main difference is that there will be a preconditioning matrix $\left( \mathrm{diag}(a) - aa^\top \right)^2$ in the dynamics of softmax transformers. In particular, this implies with adapative second-order methods,
the training dynamics of these two models can potentially be very similar. Moreover, even with SGD, there is a close 
relationship between their training dynamics, which we will now explain. 

Near initialization, i.e., when the attention pattern is still close to the uniform attention, we have 
$\left( \mathrm{diag}(a) - aa^\top \right)^2 \approx \frac{1}{T^2} \left(I - \frac{\mathbf{1}\mathbf{1}^\top}{T}\right)$. 
In other words, our linear model and softmax transformers are approximately equivalent, up to a change of learning rates.

Now, suppose that there is a nontrivial correlation between $a$ and $q = e_1$, say, $a_1$ is a small constant while 
all other entries are $O(1/T)$. In this case, we have $\left( \mathrm{diag}(a) - aa^\top \right)^2 \approx 
a_1(1 - a_1) e_1 e_1^\top + O(1/T)$. Effectively, softmax transformers automatically adjust the learning rate 
according to $a_t$ and roughly ignore those positions with a small attention weight to stabilize the gradients. 
Note that this is also what $\ell_1$-regularization does in our algorithm. In fact, mimicking this behavior is one of 
the motivations of using $\ell_1$-regularization. 

## Empirical results
Empirically, we run further experiments to stress the resemblance between softmax attention and our linear attention model (with full attention to both the token and positional encoding parts, see the attached PDF). For the effective token-position attention block, both softmax and linear attention has the same pattern, succeeded in recovering the ground-truth.

We thank Reviewer fyHi and 9wme for providing constructive suggestions on the drawbacks of our experiments, and follow-up experiments are added in the attached PDF. We tried different orders of state number $N$ and show that $\ell_1$ regularization is necessary and outperforms SGD when batch size is small (gradient noise is large).  Due to computation limitation, we experiment with real batched gradient on $N\leq 20, T\leq 5000$, and do SGD simulation by combining our population gradient + gaussian noise (to mimick the batch gradient noise) for $N\leq 500, T\leq 100000$. 

We also corroborate our previous experiments with the new test loss plot to show convergence of training. Note that since there are multiple global minima for the linear attention, $\ell_1$ regularized dynamics eventually will make $A$ and $Q$ deviate from the ground-truth while representing the same function. That is why the loss converges but the distance to the ground-truth increases after some point. That's why we need the final normalization step to recover the ground-truth. Another point is that according to our theory, the regularization will eventually distort the learned pattern when trained for too many iterations. Empirically, the loss also increases a little after it converges. Therefore, we need to early stop and do the normalization before the distortion happens.

Meta-review: This paper tries to improve our understanding of training dynamics of transformer-based language models. In this regard, the authors propose a novel Sparse Contextual Bigram (SCB) model and use it to analyzes the dynamics. In this model, it is shown that training can be split into two stages: an initial sample-intensive stage and a more sample-efficient stage. This existence of two stages has many consequences, for example in transfer learning scenarios, the first stage can be skipped if there's significant correlation between the pretraining and downstream tasks. The reviewers generally found the paper to be technically sound and well-presented, with novel theoretical contributions. However, a few limitations and concerns were also raised: practical relevance of the SCB data generation model, simplification of the transformer architecture, and limited experimental validation. We thank the authors and reviewers for engaging during the discussion period towards improving the paper. In their rebuttal, the authors addressed many of these concerns, providing additional experimental results, discussing connections to softmax transformers, and offering more intuition behind the proofs. Overall, the paper presents a novel and rigorous theoretical analysis that contributes to our understanding of transformer training dynamics. While the simplifications made limit direct practical applications, the work provides valuable insights and a foundation for further research on more complex transformer architectures.
Decision: Accept (poster)