Title: Learning and Transferring Sparse Contextual Bigrams with Linear Transformers
Venue: NeurIPS 2024
Accepted: yes
Keywords: transformer

sample complexity

optimization
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

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 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 this setting and discuss its relationship with the usual softmax-based transformers.

1 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 et al., 2021). The core of transformers is the self-attention layer (Vaswani et al., 2017), which can attend to any subset of the input sequence to output a 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 in 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 as global knowledge (OpenAI, 2023) acquired during training.

To theoretically understand how transformers learn both capabilities, we propose a minimalist datagenerating 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 𝑷 and a sparse set of prior tokens that is determined by the last token. In particular, SCB can be represented by a one-layer linear transformer -asimplified architecture that can serve as an abstraction for studying transformer optimization (Ahn et al., 2023), which makes it suitable for theoretical analysis.

In this paper, we investigate the training dynamics and sample complexity of training a linear transformer to learn the SCB task using a stochastic gradient-based 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 ℓ 1-regularized 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 hyperparameters, we prove that our algorithm can recover the ground-truth with polynomial dependence on the sequence length 𝑇 , number of states 𝑁 , and the sparsity parameter 𝑄 ≪ 𝑇 . We show that the training first goes through an initial sample-intensive stage which boosts the signal with poly ( 𝑇 ) samples, followed by a more sample-efficient stage to achieve final convergence with poly ( 𝑁, 𝑄 ) samples. We empirically verify that our gradient-based methods converge to the ground truth with a small batch size, while unregularized stochastic gradient descent fails due to the large variance.

Transfer Learning: We prove that, when there is a nontrivial correlation between the pretraining and downstream tasks, we can transfer a pre-trained model to bypass the first sample intensive stage, so that our algorithm converges to the ground truth of the downstream task with only poly ( 𝑁, 𝑄 ) 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 a Vision Transformer (ViT) (Dosovitskiy et al., 2020) trained through GD, augmented with positional-embedding attention matrices, can effectively capture spatial structures. Li et al. (2023) investigated the sample complexity necessary to achieve good generalization performance on a similar ViT model. Tarzanagh et al. (2023) established a connection between the optimization landscape of self-attention and the formulation of a 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-attention 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 on this task. Huang et al. (2023) provided a convergence guarantee for a one-layer transformer with softmax attention on a similar task where the in-context tokens are drawn from a 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 a 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 𝑛 -gram tasks, and observed multi-phase dynamics. Makkuva et al. (2024) studied the loss landscape of transformers trained on sequences sampled from a single Markov Chain. Our SCB model extends the classical bigram models to allow context-dependent sparse attention on previous tokens.

Several works, including Tian et 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 a 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, our 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 et 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) 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 a 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 2 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-generating model, the one-layer linear transformer architecture, and the training algorithm.

Notations. We use [ 𝑇 ] to denote the set { 1 , 2 , ..., 𝑇 } . Matrices and vectors are denoted in upper-case bold letters ( 𝑨 , 𝑽 , 𝚫 , etc.) and lower-case bold letters ( 𝒂 , 𝒒 , etc.), respectively. For norm, ∥·∥ denotes ℓ 2 norm and ∥ · ∥ 𝐹 denotes the Frobenius norm. Additionally, for 𝝁 ∈ R 𝑁 , ∥ 𝑨 ∥ 𝝁 denotes 𝜇 -norm for matrix 𝑨 ∈ R 𝑑 × 𝑁 for arbitrary 𝑑 , which is defined as ∥ 𝑨 ∥ 2 𝝁 : = Tr ( 𝑨 diag ( 𝝁 ) 𝑨 ⊤ ) . We use ✶ {·} to denote the indicator function. We use ˜ 𝑂 (·) to 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 𝑷 ∈ R 𝑁 × 𝑁 where 𝑃 𝑛,𝑚 = P [ 𝑋 𝑡 + 1 = 𝑛 | 𝑋 𝑡 = 𝑚 ] , which is achievable through a linear model (0-layer transformer).

A 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 can model situations such as generating the words after the phrase 'by Theorem 3.2', which requires us to retrieve the statement of 'Theorem 3.2'. Here, we propose a simple extension of this type, which we call the Sparse Contextual Bigram (SCB). The contextual information is encoded by a sparse probability vector 𝒒 determined by the last token. To generate the next token, the model retrieves the tokens referenced by 𝒒 and applies the transition matrix 𝑷 (global knowledge) to one of them according to the distribution 𝒒 .

Formally, our data-generating model SCB can be described as follows. Let 𝑇 be the sequence length and [ 𝑁 ] the vocabulary. Let 𝑷 ∈ R 𝑁 × 𝑁 be a transition matrix, with column 1 𝑷 𝑘 being the transition probability vector of token 𝑘 . Suppose that 𝝁 ∈ R 𝑁 ≥ 0 is the stationary distribution of 𝑷 .

Each input sequence consists of 𝑇 + 1 tokens ( 𝑥 1 , . . . , 𝑥 𝑇 + 1 ) , i.i.d. sampled from distribution 𝝁 . The output token (label) is generated as follows. For each 𝑘 ∈ [ 𝑁 ] , there is a probability vector 𝒒 ( 𝑘 ) ∈ 𝑅 𝑇 ≥ 0 that represents the tokens the model needs to attend to when the last token 𝑥 𝑇 + 1 is 𝑘 . For notational

simplicity, we will write 𝒙 = ( 𝑥 1 , . . . , 𝑥 𝑇 ) , 𝑿 = ( 𝒆 𝑥 1 , . . . , 𝒆 𝑥 𝑇 ) and 𝑸 = ( 𝒒 ( 1 ) , . . . , 𝒒 ( 𝑁 ) ) ∈ R 𝑇 × 𝑁 . When 𝑥 𝑇 + 1 = 𝑘 , the output token 𝑥 𝑜 is sampled from the distribution

In words, we first sample a position 𝑠 ∈ [ 𝑇 ] according to 𝒒 ( 𝑘 ) and then run one step of the Markov Chain 𝑷 from 𝑥 𝑠 to generate 𝑥 𝑜 . Note that this model can be represented by a one-layer linear transformer (see the next subsection for details). We make the following assumptions on SCB task.

Assumption 2.1 (Sparse Contextual Bigram, SCB) . In the SCB task, we assume the following:

(a) ( 𝑄 -sparse) For some 𝑄 ≪ 𝑇 and each of 𝒒 ( 𝑘 ) , at most 𝑄 entries are nonzero.

(b) (Well-conditioned) There exists some constant 𝐶 ≥ 1 such that for every 𝑘 ∈ [ 𝑁 ] and 𝑡 ∈ [ 𝑇 ] , 𝑞 ( 𝑘 ) 𝑡 ∈ [ 1 /( 𝐶𝑄 ) , 𝐶 / 𝑄 ] if it is nonzero, and 𝜇 𝑘 ∈ [ 1 /( 𝐶𝑁 ) , 𝐶 / 𝑁 ] .

(c) (Nontrivial transition) ∥ 𝑷 ∥ 2 𝜇 - ∥ 𝝁 ∥ 2 ≥ ∥ 𝝁 ∥ 2 .

(d) (Long sequence) 𝑇 ≥ ( 𝑁𝑄 ) 10 .

Remark on condition (c) . We say the transition 𝑷 is trivial if the transition probability vectors are all the same, i.e., 𝑷 = 𝝁 1 ⊤ . In this case, we have ∥ 𝑷 ∥ 2 𝜇 = ⟨ 𝝁 1 ⊤ , 𝝁𝝁 ⊤ ⟩ = ∥ 𝝁 ∥ 2 . Requiring ∥ 𝑷 ∥ 2 𝜇 - ∥ 𝝁 ∥ 2 ≥ ∥ 𝝁 ∥ 2 rules out situations where 𝑷 is too close to the trivial one. Also, note that for any well-conditioned 𝝁 , we have ∥ 𝝁 ∥ 2 ≥ Ω ( 1 / 𝑁 ) . ♣

In this work, we focus on the case where ( 𝒙 , 𝑥 𝑇 + 1 , 𝑥 𝑜 ) are given as (one data point of) the training data with ( 𝑥 1 , . . . , 𝑥 𝑇 + 1 ) i.i.d. sampled from 𝜇 . The SCB task can be extended to a sequence-to-sequence model: we drop 𝑥 1 and append 𝑥 𝑜 to get a new input sequence ( 𝑥 2 , . . . , 𝑥 𝑇 + 1 , 𝑥 𝑜 ) , and then repeat the same sampling procedure to generate another token. This generates a sequence ( 𝑥 𝑡 ) ∞ 𝑡 = 1 where ( 𝑥 𝑇 + 2 , 𝑥 𝑇 + 1 , . . . ) are not independent, and this makes our model a true language model. We leave the study of the more complicated learning-from-( 𝑥 𝑡 ) ∞ 𝑡 = 1 task to future works.

2.2 Transformer architecture

Our learner model is a one-layer single-head linear transformer (Akyürek et al., 2022; Zhang et al., 2023b; Ahn et al., 2023). A general linear transformer can be expressed as: 𝑭 ( 𝒙 , 𝑥 𝑇 + 1; 𝑽 , 𝑨 ) = 𝑽𝑬 ( 𝑬 ⊤ 𝑨𝑬 ) , where 𝑬 is the embedding of the input tokens and positions, and 𝑨 , 𝑽 are the parameters of the attention and output layers, respectively. In our setting, we only need a simpler model:

where 𝑽 ∈ R 𝑁 × 𝑁 and 𝑨 ∈ R 𝑇 × 𝑁 are the trainable parameters, and 𝒂 ( 𝑘 ) denotes the 𝑘 -th column of 𝑨 . This model uses cross-attention (replacing the last 𝑬 with 𝒆 𝑥 𝑇 + 1 ), uses only the positional embeddings together with the last token to compute the attention weights (replacing the second 𝑬 with 𝑰 𝑇 ), and discards the positional embeddings in the output layer (replacing the first 𝑬 with 𝑿 ). This is equivalent to manually set certain blocks in the weight matrices to 0, which is a common practice in the theoretical literature to simplify the analysis (Nichani et al., 2024; Huang et al., 2023; Zhang et al., 2023b).

Note that our data-generating model (1) can be represented using (2) by setting 𝑨 = 𝑸 and 𝑽 = 𝑷 . 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 𝑷 and 𝑸 are known when choosing the initialization and learning rate. The goal here is to recover 𝑷 and 𝑸 . Our loss function is the ℓ 1-regularized MSE loss. The standard way to optimize an ℓ 1-regularized loss is to use the proximal gradient descent. We adopt this algorithm with several additional pre-conditioning and a projection step to ensure some basic properties.

Formally, let the per-sample loss be defined as

We initialize 𝑨 = 1 𝑇 1 ⊤ 𝑁 / 𝑇 to have uniform attention and 𝑽 = 𝝁 1 ⊤ to be the trivial transition. At each step 𝜏 ≥ 0, we sample 𝐵 𝜏 fresh samples { 𝒙 ( 𝑖 ) , 𝑥 ( 𝑖 ) 𝑇 + 1 , 𝑥 ( 𝑖 ) 𝑜 } 𝐵 𝜏 𝑖 = 1 to form a mini-batch. The ℓ 1-regularized mini-batch loss is defined as

where 𝜆 > 0 is a parameter that controls the strength of regularization. Let ∇ ( 𝐵 𝜏 ) 𝑽 𝑙 and ∇ ( 𝐵 𝜏 ) 𝑨 𝑙 denote the mini-batch gradients of the original 𝑙 w.r.t. 𝑽 and 𝑨 , respectively. We then define the preconditioned gradients as

Here, the 1 / 𝝁 rescaling plays a role similar to importance sampling. We multiply ∇ ( 𝐵 𝜏 ) 𝑽 𝑙 with 𝑰 -11 ⊤ / 𝑁 and 𝑰 -𝝁𝝁 ⊤ /∥ 𝝁 ∥ 2 to ensure at least 1 ⊤ 𝑁 𝑽 = 1 ⊤ 𝑁 , 1 ⊤ 𝑇 𝑨 = 1 ⊤ 𝑁 , and 𝑽𝝁 = 𝝁 always hold throughout training. Note that we project each column of 𝑽 to the affine space { 𝒗 ∈ 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

where 𝜂 𝑉 > 0 is the step size. Now, consider the attention layer. Due to the existence of the ℓ 1-regularization, the update rule becomes a simple variant of the standard proximal gradient descent. Formally, for step size 𝜂 𝐴 > 0, each 𝑘 ∈ [ 𝑁 ] and 𝑡 ∈ [ 𝑇 ] , we have

For the proximal step, we will later show that no 𝑎 ( 𝑘 ) 𝑡 can ever become smaller than -𝜆 , so it suffices to consider those two cases. During the proximal step, all small 𝑎 ( 𝑘 ) 𝑡 are set to 0, and 𝜆 is subtracted from all large coordinates. For notational simplicity, we define 𝒈 ( 𝑘 ) 𝜆,𝜏 : = -𝜂 -1 𝐴 ( 𝒂 ( 𝑘 ) 𝜏 + 1 -𝒂 ( 𝑘 ) 𝜏 ) so that we can write the update as 𝒂 ( 𝑘 ) 𝜏 + 1 = 𝒂 ( 𝑘 ) 𝜏 -𝜂 𝐴 𝒈 ( 𝑘 ) 𝜏 . We will choose 𝜆 = 0 in certain stages of training. In this case, (6) becomes the usual projected preconditioned gradient descent and we have

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 1 and more details on the projection/normalization steps are provided in Appendix E and F. When we train the model from scratch, all three stages are used and the initialization is 𝑽 0 = 𝝁 1 ⊤ 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, we will run Algorithm 1 from Stage 2.

Algorithm 1 Projected preconditioned ℓ 1-proximal gradient descent

Input: Stationary distribution 𝝁 ; initialization 𝑽 0 , 𝑨 0; learning rates 𝜂 ( 𝑖 ) 𝐴 , 𝜂 ( 𝑖 ) 𝑉 , 𝑖 ∈ [ 3 ] ; threshold 𝜆 0; regularization strength ˆ 𝜆 ; times T 1 , T 2 , T 3

Stage 1: Run (5) and (6) with 𝜂 𝐴 = 𝜂 ( 1 ) 𝐴 , 𝜂 𝑉 = 𝜂 ( 1 ) 𝑉 , 𝜆 = 0 for T 1 steps;

∀

∈ [

]

,

ˆ

𝒂

=

[

(

𝑡

{

(

𝑡

0

}]

𝑡

,

𝒂

) ←(

𝑰

𝑇

-

1

𝑇

1

⊤

𝑇

/

)

ˆ

𝒂

Thresholding-normalization: ∀ 𝑘 ∈ [ 𝑛 ] , ˆ 𝒂 ( 𝑘 ) = [ 𝑎 ( 𝑘 ) 𝑡 ✶ { 𝑎 ( 𝑘 ) 𝑡 ≥ Ω ( 1 / 𝑄 )}] 𝑡 . 𝒂 ( 𝑘 ) ← ˆ 𝒂 ( 𝑘 ) / 1 ⊤ ˆ 𝒂 ( 𝑘 ) Stage 3: Run (5) and (6) with 𝜂 𝐴 = 0, 𝜂 𝑉 = 𝜂 ( 3 ) , 𝜆 = 0 for T - T steps;

Thresholding-projection: ✶ Stage 2: Run (5) and (6) with 𝜂 𝐴 = 𝜂 ( 2 ) 𝐴 , 𝜂 𝑉 = 𝜂 ( 2 ) 𝑉 , 𝜆 = ˆ 𝜆 for T 2 - T 1 steps;

𝑘

𝑛

𝑎

𝑎

𝜆

𝑉 3 2 Output: 𝑨 T 3 , 𝑽 T 3 .

3 Results for training from scratch

In this section, we consider the situations where we train the model from scratch, i.e., the initialization is 𝑽 0 = 𝝁 1 ⊤ and 𝑨 0 = 11 ⊤ / 𝑇 and discuss the ideas of the proof of the following theorem.

Theorem 3.1 (Theorem G.1) . Let 𝜀 > 0 be our target accuracy and T 1 = min { 𝜏 ≥ 0 : max { 𝛼 𝑉,𝜏 , 𝛼 𝐴,𝜏 } ≥ Θ ( 1 /( 𝑄𝑁 ))} . We can choose the hyperparameters in Algorithm 1 such that within poly ( 𝑁, 𝑄, 1 / 𝜀, log 𝑇 ) steps, we have ∥ 𝑨 -𝑸 ∥ 2 𝜇 ≤ 𝜀 and ∥ 𝑽 -𝑷 ∥ 2 𝜇 ≤ 𝜀 with probability at least 1 -𝛿 and the numbers of samples used before and after T 1 are poly ( 𝑇, 𝛿 ) and poly ( 𝑁, 𝑄, 1 / 𝜀, log 𝑇, 𝛿 ) , respectively.

The overall strategy is analyzing the population process and then controlling the distance between the mini-batch trajectory and the population process 2 . In Section 3.1, we discuss the key properties of the population process that simplify the analysis. After that, we describe the dynamics of the algorithm and the signal-noise-ratio (SNR) in each of the three stages of Algorithm 1 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 ease of presentation, we assume 𝜆 = 0 and access to the population loss L : = E 𝑙 . In other words, we consider the projected preconditioned gradient descent. By Lemma B.8, the dynamics of the population process is controlled by

One can prove via induction on 𝜏 that 𝑽 (resp. 𝑨 ) always stays on the straight line crossing 𝝁 1 ⊤ and 𝑷 (resp. 11 ⊤ / 𝑇 and 𝑸 ). In other words, there exists some time-dependent real numbers 𝛼 𝑉,𝜏 , 𝛼 𝐴,𝜏 , 𝛽 𝑉,𝜏 : = 1 -𝛼 𝑉,𝜏 , 𝛽 𝐴,𝜏 : = 1 -𝛼 𝐴,𝜏 such that 𝑽 𝜏 = 𝛼 𝑉,𝜏 𝑷 + 𝛽 𝑉,𝜏 𝝁 1 ⊤ and 𝑨 𝜏 = 𝛼 𝐴,𝜏 𝑸 + 𝛽 𝐴,𝜏 11 ⊤ / 𝑇 . The same calculation yields the following equations that govern the dynamics of 𝛼 𝑉 and 𝛼 𝐴 :

where 𝛼 𝑉, 0 = 𝛼 𝐴, 0 = 0, 𝐾 𝑃 : = ∥ 𝑷 ∥ 2 𝜇 - ∥ 𝝁 ∥ 2 ≳ 1 / 𝑁 and 𝐾 𝑄 : = ∥ 𝑸 ∥ 2 𝜇 -1 / 𝑇 ≳ 1 / 𝑄 . Choose 𝜂 𝑉 = 𝜂 / 𝐾 𝑄 and 𝜂 𝐴 = 𝜂 / 𝐾 𝑃 , and we can write the above in matrix form as

(

𝑘

)

𝑘

)

𝑘

)

(

𝑘

𝑇

(

𝑘

)

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 𝛼 '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 𝛼 𝑉 and/or 𝛼 𝐴 3 .

Note that when both 𝛼 𝑉 , 𝛼 𝐴 are still small, the population dynamics of 𝛼 𝑉 , 𝛼 𝐴 are a linear system with coefficient matrix 𝜂 GLYPH<2> 0 1 1 0 GLYPH<3> and drift h 𝜂𝑇 / 𝐾 𝑄 0 i . 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 0 at initial and the initial signal provided by the drift term has order 1 / 𝑇 , we should expect that poly 𝑇 samples are necessary to distinguish it from noises (Stage 1). After the signal becomes reasonably large, the first term will have order 1 / poly ( 𝑁, 𝑄 ) , and we can then rely on it (combined with the 𝑙 1-regularization) instead of the drift term to learn the model (Stage 2).

3.2 Stage 1: boosting the signal

At initialization, we have 𝛼 𝑉 = 𝛼 𝐴 = 0. We define Stage 1 to be the phase until at least one of them has grown from 0 to some small 𝜎 1 = 1 / poly ( 𝑁, 𝑄 ) . Note that in this stage, the mini-batch version of (8) is approximately equivalent to

where 𝜺 noise and 𝜺 approx 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 𝜺 noise and 𝜺 approx were 0, then the second term on the RHS of (9) could provide a small positive signal to 𝛼 and the first term would quickly boost it to 𝜎 1 within T 1 = log ( 𝑇 )/ 𝜂 iterations. In order for the above analysis to work, we need both 𝜺 noise and 𝜺 approx to be at least 𝑂 ( 1 )/ 𝑇 small. Since, unfortunately, 𝜺 noise does not scale with 1 / 𝑇 , we need poly ( 𝑇 ) samples to ensure these conditions.

We conjecture that this poly ( 𝑇 ) 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 to reveal the signal or poly ( 𝑇 ) steps (whence also poly ( 𝑇 ) 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 𝛼 𝑉 and 𝛼 𝐴 is 𝜎 1 = 1 / poly ( 𝑄, 𝑁 ) large. Hence, one may expect that the signal is 1 / poly ( 𝑄, 𝑁 ) large now so that we no longer need to make the noises 1 / 𝑇 small and therefore, only poly ( 𝑁, 𝑄 ) samples are needed. Unfortunately, this argument will not work directly, since the variance of the mini-batch gradients scales with 𝑇 . 4 Therefore, we still need poly ( 𝑇 ) samples to reduce the squared ℓ 2-norm of 𝜺 noise from Ω ( 𝑇 ) to 1 / poly ( 𝑄, 𝑁 ) . To address this issue, we introduce the ℓ 1-regularizer and use a variant of proximal gradient descent.

The idea is, while the concentration in the ℓ 2 sense is difficult, controlling the ℓ ∞ -error is easy as every entry of ∇ 𝒂 ( 𝑘 ) 𝑙 is bounded whence subgaussian. As a result, we can make the coordinatewise difference between the population and mini-batch gradients 1 / poly ( 𝑁, 𝑄 ) small using only poly ( 𝑁, 𝑄 ) log 𝑇 samples by a standard concentration argument. Moreover, we have (cf. the proof of Lemma E.3)

Thus, as long as 𝛼 𝑉 ≥ 1 / poly ( 𝑁, 𝑄 ) , the ℓ ∞ -norm of the gradient noise being small is enough to create a separation between those useful entries ( 𝑞 ( 𝑘 ) 𝑡 ≠ 0) and useless entries ( 𝑞 ( 𝑘 ) 𝑡 = 0) and ensure the ℓ 2-error of those 𝑄 useful entries is small.

The above analysis suggests removing all small entries from the gradient will work. Now, we claim that ℓ 1-regularization and proximal gradient descent naturally implement this strategy, at least approximately. We believe softmax-based 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 1 and after the thresholding-projection step - which is approximately equivalent to running one proximal step first - we know that all useful 𝑎 ( 𝑘 ) 𝑡 are at least Ω ( 𝛼 𝑉 / 𝑄 ) = 1 / poly ( 𝑁, 𝑄 ) , while all useless entries are of size 𝑂 ( 1 )/ 𝑇 . By our previous discussion, we know that if 𝜆 is chosen appropriately, with poly ( 𝑁, 𝑄, log 𝑇 ) 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 𝜆 . Thus, after one gradient step, the absolute value of each of those useless 𝑎 ( 𝑘 ) 𝑡 is still much smaller than 𝜆 . As a result, they will be set to 0 in the proximal step (and to 𝑂 ( 1 / 𝑇 ) after the projection step), which is equivalent to filtering out all those entries, up to a small bias. Therefore, the proximal gradient updates stay close to the population trajectory, and the growth of the signals 𝛼 𝐴 , 𝛼 𝑉 can be analyzed using the population dynamics.

We end Stage 2 when ( 𝛼 𝑉 + 𝛼 𝐴 )/ 2 ≈ 1. Similar to Stage 1, this also only takes T 2 = ˜ 𝑂 ( 1 / 𝜂 ) steps. We also show that the difference between the mini-batch trajectory and the 'population trajectory' can decrease to a small value (cf. Lemma E.10). 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 a minor issue regarding | 𝛼 𝑉 -𝛼 𝐴 | . Taylor expand (8) around ( 𝛼 𝐴 , 𝛼 𝑉 ) = ( 1 , 1 ) and one will notice that although ( 𝛼 𝑉 + 𝛼 𝐴 )/ 2 can converge to 1 at a linear rate (and the approximation error also decreases exponentially fast), the convergence rate of 𝛼 𝐴 -𝛼 𝑉 is much slower, and the process will get stuck around ( 1 + 𝛿, 1 -𝛿 ) for some small nonzero 𝛿 , instead of converging to ( 1 , 1 ) . To accelerate this process, we directly round 𝑨 via normalization, which is possible only after the approximation error becomes small in Stage 2. Then we freeze 𝑨 and train 𝑽 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. It is often assumed that the downstream and pretraining tasks share a common structure or representations/features, 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-intensive signal-boosting stage.

Formally, we consider the following setting. Let ˆ 𝑷 be the transition matrix of the pretraining task and ( 𝑷 , 𝑸 ) the transition matrix and 𝑸 -matrix of the downstream task. We still assume Assumption 2.1. In addition, we assume ˆ 𝑷 and 𝑷 share the same stationary distribution 𝝁 , GLYPH<13> GLYPH<13> ˆ 𝑷 GLYPH<13> GLYPH<13> 2 𝜇 = Θ ( 1 ) ∥ 𝑷 ∥ 2 𝜇 and GLYPH<10> ˆ 𝑷 , 𝑷 GLYPH<11> 𝜇 -∥ 𝝁 ∥ 2 ≥ ∥ 𝝁 ∥ 2 . The last condition can be viewed as the transfer learning version of condition (c) of Assumption 2.1. Note that we allow the correlation between ˆ 𝑷 and 𝑷 to be as small as 𝑜 ( 1 ) . Theorem 4.1 (informal version of Theorem H.3) . Consider the above setting. Initialize 𝑨 = 11 ⊤ / 𝑇 , 𝑽 = 𝜃 ˆ 𝑷 + ( 1 -𝜃 ) 𝝁 1 ⊤ for some small 𝜃 > 0 , and run one step of gradient update on 𝑨 . Then, running Algorithm 1 from Stage 2 allows us to recover ( 𝑷 , 𝑸 ) to 𝜀 -accuracy with high probability with poly ( 𝑁, 𝑄, 1 / 𝜀, log 𝑇 ) samples.

To intuitively see why using poly ( 𝑁, 𝑄, 1 / 𝜀, log 𝑇 ) samples is possible, recall from (9) that the reason we need poly ( 𝑇 ) samples in Stage 1 is the signal is additive and has order 𝑂 ( 1 / 𝑇 ) , so we need the size of the noise to be at most 𝑂 ( 1 / 𝑇 ) . On the other hand, when we initialize 𝑽 = 𝜃 ˆ 𝑷 + ( 1 -𝜃 ) 𝝁 1 ⊤ , we have 𝛼 𝑉 ≥ Θ ( 𝜃 /( 𝑁𝐾 𝑃 )) ≫ 1 / 𝑇 . Then we can rely on 𝛼 𝑉 , instead of the 1 / 𝑇 -sized additive signal, to boost the signal of 𝛼 𝐴 to 𝜔 ( 1 / 𝑇 ) in one step, which leads to a sample complexity that depends on 𝛼 𝑉 instead of the 1 / 𝑇 -sized additive signal. Then, we can reuse the analysis of Stage 2 and 3 to show that the downstream ( 𝑷 , 𝑸 ) can be approximately recovered using poly ( 𝑁, 𝑄, 1 / 𝜀, log 𝑇 ) .

Note that unlike the case of training from scratch, when performing transfer learning, the initial approximation error ∥ 𝚫 𝑉 ∥ 𝜇 , i.e., the distance between 𝑽 and its population projection, can be much larger than the signal 𝛼 𝑉 , 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 ∥ 𝚫 𝑉 ∥ 2 𝜇 on the dynamics scales with 𝛼 𝐴 (≈ 𝛼 𝑉 ) , which is small. In addition, we also show that as long as 𝛼 𝐴 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 in the signal-boosting stage until we enter the regime of the Stage 2 analysis.

5 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 a randomly sampled transition matrix 𝑷 together with its stationary 𝜇 , and the ground truth attention pattern 𝑸 . We choose the number of states 𝑁 = 3, sparsity 𝑄 = 2, and the sequence length 𝑇 = 5000 ≫ 𝑁, 𝑄 . We use a batch size 𝐵 = 64 to run the online projected proximal gradient descent with 𝜆 = 1e-5 and the vanilla SGD for T = 1000 iterations. Through the signal boosting stage 𝜏 ∈ [ 0 , 400 ] , we use 𝜂 1 = 0 . 01 to accelerate the process. After 𝜏 > 400, we use 𝜂 2 = 0 . 005 for further improvement. For SGD, we add another set of experiments with 𝜂 ′ 2 = 0 . 001 to prevent potential instability. For more details, see Appendix I.

5.1 Convergence

Our experiments (cf. Fig. 1) show that after switching to proximal gradient descent after Stage 1 (the signal-boosting stage), both ∥ 𝑷 -𝑽 ∥ 𝜇 and ∥ 𝑨 -𝑸 ∥ 𝜇 decrease faster than SGD. The final distance to the ground-truth after normalization gets close to 0, 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 is too slow when a smaller learning rate is used. This phenomenon verifies our theory that the variance of the original stochastic gradient will be too large for SGD to converge when 𝑇 ≫ 𝑄, 𝑁 , while proximal gradient descent with an ℓ 1 regularizer can resolve this issue.

0

200

400

600

800

1000

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

Distance to ground truth

‖

P

-

V

‖

µ

, ours

‖

P

-

V

‖

µ

, SGD

‖

P

-

V

‖

µ

, SGD small

Final error, ours

0

200

400

600

800

1000

0

.

2

0

.

4

0

.

6

0

.

8

Distance to ground truth

‖

Q

-

A

‖

µ

, ours

‖

Q

-

A

‖

µ

, SGD

‖

Q

-

A

‖

µ

, SGD small

Final error, ours

iteration iteration

5.2 Relationship between our model and softmax transformers

Weclaim that they have our linear model and softmax transformers have qualitatively similar behaviors: there will be a sample-intensive initial stage, and after the model and the target have a nontrivial correlation, proximal gradient descent/SGD will become much more sample efficient.

For ease of presentation, in the following, we will assume 𝑁 = 1, write 𝒂 : = 𝒂 ( 1 ) , and assume the ground-truth 𝒒 : = 𝒒 ( 1 ) is 𝒆 1 = ( 1 , 0 , . . . , 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 𝑓 ( 𝑿 ; 𝑽 , 𝒂 ) = 𝑽𝑿𝒂 . By a softmax transformer, we mean the model 𝑓 𝜎 ( 𝑿 ; 𝑽 , 𝒘 ) = 𝑽𝑿 𝜎 ( 𝒘 ) = : 𝑽𝑿𝒂 𝜎 where 𝜎 is the softmax function and 𝒘 ∈ R 𝑇 is the trainable first-layer weights.

0

20

40

60

0

10

20

30

40

50

60

0

20

40

60

0

10

20

30

40

50

60

Linear Transformer

0

20

40

60

0

5

0

20

40

60

0

5

0

1000

2000

3000

4000

0.44

0.45

0.46

0.47

0.48

0.49

0.50

linear

softmax

Let 𝑙 denote the (per-sample) loss. We have ∇ 𝒘 𝑙 ( 𝑓 𝜎 ( 𝑿 )) = GLYPH<0> diag ( 𝒂 𝜎 ) -𝒂 𝜎 𝒂 ⊤ 𝜎 GLYPH<1> ( 𝑽𝑿 ) ⊤ ∇ 𝑙 ( 𝑓 𝜎 ( 𝑿 )) . As a result, the dynamics of the attention weights are controlled by

In other words, the main difference is that there will be a preconditioning matrix ( diag ( 𝒂 ) -𝒂𝒂 ⊤ ) 2 in the dynamics of softmax transformers.

Near initialization, i.e., when the attention pattern is still close to the uniform attention, we have GLYPH<0> diag ( 𝒂 𝜎 ) -𝒂 𝜎 𝒂 ⊤ 𝜎 GLYPH<1> 2 ≈ 1 𝑇 2 GLYPH<16> 𝑰 -11 ⊤ 𝑇 GLYPH<17> . In other words, our linear model and softmax transformers are approximately equivalent up to a change in learning rates.

Now, suppose that there is a nontrivial correlation between 𝒂 𝜎 and 𝒒 = 𝒆 1, say, 𝑎 𝜎, 1 is a small constant while all other entries are 𝑂 ( 1 / 𝑇 ) . In this case, we have GLYPH<0> diag ( 𝒂 𝜎 ) -𝒂 𝜎 𝒂 ⊤ 𝜎 GLYPH<1> 2 ≈ 𝑎 𝜎, 1 ( 1 -𝑎 𝜎, 1 ) 𝒆 1 𝒆 ⊤ 1 + 𝑂 ( 1 / 𝑇 ) . Effectively, softmax transformers automatically adjust the learning rate according to 𝑎 𝜎,𝑡 and roughly ignore those positions with a small attention weight to stabilize the gradients. Note that this is also what ℓ 1-regularization does in our algorithm. In fact, mimicking this behavior is one of the motivations of using ℓ 1-regularization in our linear setting. We run further experiments to highlight the resemblance between softmax attention and our linear attention model (Figure 2).

6 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. In particular, we show when trained from scratch, the training process can be split into two stages, where the first stage uses a lot of samples to boost the signal from zero to a 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 is a nontrivial correlation between the pretraining and downstream tasks, the first sample intensive stage can be bypassed.

Our data-generating model and results also lead to 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?

References

OpenAI. Gpt-4 technical report, 2023.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 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 Romera-Paredes, 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://api.semanticscholar.org/CorpusID:235959867 .

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning , pages 5793-5831. PMLR, 2022.

Bingbin Liu, Jordan T 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.05217 , 2023.

Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. arXiv preprint arXiv:2105.11115 , 2021.

Fabio Petroni, Tim Rocktäschel, Patrick Lewis, Anton Bakhtin, Yuxiang Wu, Alexander H Miller, and Sebastian Riedel. Language models as knowledge bases? arXiv preprint arXiv:1909.01066 , 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-2650. USENIX Association, August 2021. ISBN 978-1939133-24-3. URL https://www.usenix.org/conference/usenixsecurity21/ presentation/carlini-extracting .

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 in Neural Information Processing Systems , 35:37822-37836, 2022.

Hongkang Li, Meng Wang, Sĳia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015 , 2023.

Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv:2308.16898 , 2023.

Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. arXiv preprint arXiv:2305.16380 , 2023a.

Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. 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.03576 , 2023.

Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927 , 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 arXiv:2402.19442 , 2024.

Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: 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:2402.11004 , 2024.

Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: A framework for principled analysis of transformers via markov chains. arXiv preprint arXiv:2402.04161 , 2024.

Juno Kim and Taĳi Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. arXiv preprint arXiv:2402.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:2406.06893 , 2024.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

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-7862, 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 arXiv:2002.09434 , 2020.

Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A theoretical analysis of contrastive unsupervised representation learning. arXiv preprint arXiv:1902.09229 , 2019.

Steve Hanneke, Samory Kpotufe, and Yasaman Mahdaviyeh. Limits of model selection under transfer learning. In 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.

Ye Tian and Yang Feng. Transfer learning under high-dimensional generalized linear models. Journal of the American Statistical Association , 118(544):2684-2697, 2023.

Zhe Fei and Yi Li. Estimation and inference for high dimensional generalized linear models: A splitting and smoothing approach. Journal of Machine Learning Research , 22(58):1-32, 2021.

Xuhui Zhang, Jose Blanchet, Soumyadip Ghosh, and Mark S Squillante. A class of geometric structures in transfer learning: Minimax bounds and optimality. In International Conference on Artificial Intelligence and Statistics , pages 3794-3820. PMLR, 2022.

Peizhong Ju, Sen Lin, Mark S Squillante, Yingbin Liang, and Ness B Shroff. Generalization performance of transfer learning: Overparameterized and underparameterized regimes. arXiv preprint arXiv:2306.04901 , 2023.

Yehuda Dar and Richard G Baraniuk. Double double descent: On generalization errors in transfer learning between linear regression tasks. SIAM Journal on Mathematics of Data Science , 4(4): 1447-1472, 2022.

Andrew K Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. arXiv preprint arXiv:1809.10374 , 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. In Conference on Learning Theory , pages 5413-5452. PMLR, 2022.

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661 , 2022.

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 or 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 is 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. It would be interesting (and more challenging) to consider more natural/practical training algorithms.

Finally, for our data-generating model, we only use it to generate one next token, instead of repeatedly apply SCB on the previous generated results to obtain a long sequence. In our setting, the contextual tokens are independent. While this simplifies the analysis, it 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 to the end of this section.

B.1 Probabilities

Lemma B.1. For any 𝑡 ∈ [ 𝑇 ] and 𝑘, 𝑛, 𝑚 ∈ [ 𝑁 ] , we have

Lemma B.2. For any 𝑠 ≠ 𝑡 ∈ [ 𝑇 ] , 𝑘, 𝑛, 𝑚, 𝑙 ∈ [ 𝑁 ] , we have

B.2 Gradient and Expectations

Lemma B.3. Suppose the last input token 𝒙 𝑇 + 1 and 𝒂 : = 𝑨𝒙 𝑇 + 1 . The gradients of the objective are

Lemma B.4. For any 𝑘 ∈ [ 𝑁 ] and 𝑠, 𝑡 ∈ [ 𝑇 ] , we have

Lemma B.5. For any 𝑽 ∈ R 𝑁 × 𝑁 with 𝑽𝝁 = 𝝁 and 𝑠, 𝑡 ∈ [ 𝑇 ] , we have

Lemma B.6. For any 𝑽 ∈ R 𝑁 × 𝑁 with 𝑽𝝁 = 𝝁 and 𝑡 ∈ [ 𝑇 ] , we have

Lemma B.7 (Expected gradients) .

Lemma B.8 (Expected preconditioned gradients) .

B.3 Deferred proofs of this section

B.3.1 Probabilities

Proof of Lemma B.1. For notational simplicity, define 𝒙 -𝑡 = ( 𝑥 1 , . . . , 𝑥 𝑡 -1 , 𝑥 𝑡 + 1 , . . . , 𝑥 𝑇 ) . We compute

By the independence assumption, we have P ( 𝑥 𝑡 = 𝑚, 𝒙 -𝑡 = ˆ 𝒎 | 𝑥 𝑇 + 1 = 𝑘 ) = P ( 𝑥 𝑡 = 𝑚, 𝒙 -𝑡 = ˆ 𝒎 ) = P ( 𝑥 𝑡 = 𝑚 ) P ( 𝒙 -𝑡 = ˆ 𝒎 ) . For the first factor, we have P ( 𝑥 𝑜 = 𝑛 | 𝑥 𝑡 = 𝑚, 𝒙 -𝑡 = ˆ 𝒎 , 𝑥 𝑇 + 1 = 𝑘 ) = 𝑄 ( 𝑘 ) 𝑡 𝑃 𝑛,𝑚 + ˝ 𝑠 ≠ 𝑡 𝑄 ( 𝑘 ) 𝑠 𝑃 𝑛, ˆ 𝑚 𝑠 . Therefore,

Note that for any 𝑠 ≠ 𝑡

Thus,

Proof of Lemma B.2. For notational simplicity, let 𝒙 -𝑠, -𝑡 denote the vector obtained by removing the 𝑠, 𝑡 coordinates from 𝒙 . Then, we compute

Therefore,

B.3.2 Expectations

Proof of Lemma B.3. For each sample 𝑿 , we have

and 𝒂 Then we have the matrix differential:

d

𝑙

=

𝑽𝑿𝒂

-

𝒆

𝑥

𝑜

d

𝑽𝑿𝒂

+

𝑽𝑿𝒂

-

𝒆

𝑥

𝑜

𝑽𝑿

Proof of Lemma B.4. When 𝑠 ≠ 𝑡 , we have E ( 𝑘 ) [ 𝒆 𝑥 𝑠 𝒆 ⊤ 𝑥 𝑡 ] = E ( 𝑘 ) [ 𝒆 𝑥 𝑠 ] E ( 𝑘 ) [ 𝒆 ⊤ 𝑥 𝑡 ] = 𝝁𝝁 ⊤ . When 𝑠 = 𝑡 , we have E ( 𝑘 ) [ 𝒆 𝑥 𝑡 𝒆 ⊤ 𝑥 𝑡 ] = ˝ 𝑁 𝑘 = 1 𝜇 𝑘 𝒆 𝑘 𝒆 ⊤ 𝑘 = diag ( 𝝁 ) .

Proof of Lemma B.5. When 𝑠 ≠ 𝑡 , we have

When 𝑠 = 𝑡 , we have ˝ 𝑁 𝑛 = 1 E 𝑉 2 𝑛,𝑥 𝑡 = ˝ 𝑁 𝑛 = 1 ˝ 𝑁 𝑘 = 1 𝜇 𝑘 𝑉 2 𝑛,𝑘 = ∥ 𝑽 ∥ 2 𝜇

.

Proof of Lemma B.6. Recall Lemma B.1. Then, we compute

Proof of Lemma B.7. First, we consider ∇ 𝑽 𝑙 = GLYPH<0> 𝑽𝑿𝒂 -𝒆 𝑥 𝑜 GLYPH<1> ( 𝑿𝒂 ) ⊤ , and compute E ( 𝑿𝒂 )( 𝑿𝒂 ) ⊤ and E 𝒆 𝑥 𝑜 ( 𝑿𝒂 ) ⊤ . Write 𝑿𝒂 = ˝ 𝑇 𝑡 = 1 𝑎 𝑡 𝒆 𝑥 𝑡 . Then, we have

GLYPH<0>

GLYPH<1>

⊤

GLYPH<0>

GLYPH<1>

⊤

d

𝒂

where the last line comes from Lemma B.4. Then, we compute

where the second line comes from Lemma B.1. Thus, for ∇ 𝑽 𝑙 , we have

Now, consider ∇ 𝒂 ( 𝑘 ) 𝑙 = ✶ { 𝑥 𝑇 + 1 = 𝑘 }( 𝑽𝑿 ) ⊤ GLYPH<0> 𝑽𝑿𝒂 -𝒆 𝑥 𝑜 GLYPH<1> and compute E ( 𝑘 ) ( 𝑽𝑿 ) ⊤ ( 𝑽𝑿 ) and E ( 𝑘 ) ( 𝑽𝑿 ) ⊤ 𝒆 𝑥 𝑖 . By Lemma B.5, for each 𝑠, 𝑡 ∈ [ 𝑇 ] , we have

In matrix form, this is

Then, by Lemma B.6, for each 𝑡 ∈ [ 𝑇 ] , we have

In matrix for, this is

Combine these together, and we obtain

Proof of Lemma B.8. Recall from Lemma B.7 that

For ∇ 𝑽 , we have

In particular, note that the 𝝁𝝁 ⊤ /∥ 𝝁 ∥ 2 terms will cancel with each other. Thus, we have

For ˆ ∇ 𝒂 ( 𝑘 ) 𝑙 , we have

B.4 Concentration

In this section, we provide concentration inequalities for the gradients of the loss function. The concentration is applied on the gradient noise term

𝒉

GLYPH<16>

𝐵

)

𝑨

, 𝜏

:

=

ˆ

∇ (

𝑨

𝑙

-

E

ˆ

∇

𝑨

𝑙

where ˆ ∇ ( 𝐵 ) 𝑽 𝑙 and ˆ ∇ ( 𝐵 ) 𝑨 𝑙 are the preconditioned empirical gradients computed from a batch of size 𝐵 . Here, we first consider the concentration of the original gradients:

✶ and then consider the concentration of the preconditioned gradients. In this paper, we focus on ∥·∥ 𝜇 as the mostly used metric for the gradient matrices.

First we prove a naive concentration w.r.t. any random vector 𝒚 with bounded second moment with any ∥·∥ .

Lemma B.9. Fix 𝛿, 𝜀 > 0 . Let 𝒚 be a 𝐷 -dimensional random vector with E ∥ 𝒚 ∥ 2 ≤ 𝐺 . Define 𝒚 ( 𝐵 ) : = 𝐵 -1 ˝ 𝐵 𝑖 = 1 𝒚 𝑖 where ( 𝒚 𝑖 ) 𝑖 are i.i.d. versions of 𝒚 . If

then with probability at least 1 -𝛿 , we have ∥ 𝒚 -E 𝒚 ∥ ≤ 𝜀 .

Proof of Lemma B.9. Assume w.l.o.g. that E 𝒚 = 0. First, note that

GLYPH<17>

Hence, by the Markov inequality, we have

Thus, for fixed 𝜀, 𝛿 ∈ ( 0 , 1 ) , if we choose 𝐵 = 𝐺 /( 𝛿𝜀 2 ) , then we have with probability at least 1 -𝛿 , GLYPH<13> GLYPH<13> 𝒚 ( 𝐵 ) GLYPH<13> GLYPH<13> ≤ 𝜀 .

Now we upper bound the infinity norm of the preconditioned gradients to apply concentration.

Lemma B.10. Suppose that ∥ Vec 𝑽 ∥∞ = 𝑂 ( 1 ) , 𝑽 = ˜ 𝑽 + 𝚫 𝑽 , 𝑨 = ˜ 𝑨 + 𝚫 𝑨 , where ˜ 𝑽 ∈ R 𝑁 × 𝑁 is a transition probability matrix, and in the attention matrix ˜ 𝑨 each column ˜ 𝒂 ( 𝑘 ) is a probability vector. Moreover, ∥ 𝚫 𝑨 ∥ 2 𝐹 , ∥ 𝚫 𝑽 ∥ 2 𝐹 ≤ 𝑂 ( 1 / 𝑇 ) . Then, we have

Proof. We first consider the infinity norm of the original gradient. Recall that the gradient for 𝑽 and 𝑨 are

and the preconditioned gradient is:

We first consider the maximum absolute value in the original gradients. For ∇ 𝒂 ( 𝑘 ) 𝑙 , we have

The first term can be upper-bounded in the following way:

Since 𝑽 = ˜ 𝑽 + 𝚫 𝑽 and ∥ 𝚫 𝑽 ∥ 2 𝐹 , ∥ 𝚫 𝑨 ∥ 2 𝐹 ≤ 𝑂 ( 1 / 𝑇 ) , we have max 𝑠,𝑡 GLYPH<12> GLYPH<12> ( 𝑽𝒆 𝑥 𝑠 ) ⊤ 𝑽𝒆 𝑥 𝑡 GLYPH<12> GLYPH<12> upper bounded by 𝑂 ( 1 ) . Therefore

And similarly, the second term max 𝑠 ∈[ 𝑇 ] GLYPH<12> GLYPH<12> ( 𝑽𝒆 𝑥 𝑠 ) ⊤ 𝒆 𝑥 𝑜 GLYPH<12> GLYPH<12> can be bounded by 𝑂 ( 1 ) because the infinity norm of 𝑽 is also upper bounded by 𝑂 ( 1 ) . Therefore, we know ∥∇ 𝒂 ( 𝑘 ) 𝑙 ∥∞ ≤ 𝑂 ( 1 ) .

Now we consider the preconditioned gradient ˆ ∇ ( 𝐵 𝜏 ) 𝒂 ( 𝑘 ) 𝑙 :

since 𝜇 𝑘 ≥ 𝑐 𝑁 for all 𝑘 ∈ [ 𝑁 ] .

We use similar technique on ˆ ∇ ( 𝐵 𝜏 ) 𝑽 𝑙 . First, we prove the infinity norm upper bound on the original gradient.

And therefore, the preconditioned gradient can also be bounded.

since 𝜇 𝑘 ≥ 𝑐 𝑁 for all 𝑘 ∈ [ 𝑁 ] . 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 𝑨 and 𝑽 .

Corollary B.11. With the same setting in Lemma B.10 and GLYPH<13> GLYPH<13> ˆ ∇ 𝒂 ( 𝑘 ) 𝑙 GLYPH<13> GLYPH<13> ∞ ≤ 𝑂 ( 𝑁 ) , GLYPH<13> GLYPH<13> Vec ˆ ∇ 𝑽 𝑙 GLYPH<13> GLYPH<13> ∞ ≤ 𝑂 ( 𝑁 ) . Moreover, ∥ 𝚫 𝑨 ∥ 2 𝐹 , ∥ 𝚫 𝑽 ∥ 2 𝐹 ≤ 𝑂 ( 1 / 𝑇 ) . Then, we have

Proof. We directly upper bound GLYPH<13> GLYPH<13> ˆ ∇ 𝑨 𝑙 GLYPH<13> GLYPH<13> 2 𝜇 and GLYPH<13> GLYPH<13> ˆ ∇ 𝑽 𝑙 GLYPH<13> GLYPH<13> 2 𝜇 using the upper bound on infinity norm. Since ∥ 𝚫 𝑨 ∥ 2 𝐹 , ∥ 𝚫 𝑽 ∥ 2 𝐹 ≤ 𝑂 ( 1 / 𝑇 ) , we have the infinity norm be upper bounded by

Then, we can first bound the Frobenius norm ∥ ˆ ∇ 𝑨 𝑙 ∥ 2 𝐹 , ∥ ˆ ∇ 𝑽 𝑙 ∥ 2 𝐹 . We have 𝑽 ∈ R 𝑁 × 𝑁 , 𝑨 ∈ R 𝑇 × 𝑁 , so

That leads to:

where the second inequality comes from the assumption that 𝜇 ∼ Θ ( 1 / 𝑁 ) .

Nowwith 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 to be bounded in the signal dynamics:

Lemma B.12. Fix 𝜀, 𝛿 > 0 . Under Assumption 2.1, suppose GLYPH<13> GLYPH<13> ˆ ∇ 𝒂 ( 𝑘 ) 𝑙 GLYPH<13> GLYPH<13> ∞ ≤ Θ ( 𝑁 ) , GLYPH<13> GLYPH<13> ˆ ∇ 𝑽 𝑙 GLYPH<13> GLYPH<13> ∞ ≤ Θ ( 𝑁 ) . If 𝐵 ≥ Θ ( 1 ) max GLYPH<0> 𝑁 4 , 𝑄 2 𝑁 2 GLYPH<1> 𝑁 2 log 4 𝛿 𝜖 2 𝐾 2 𝑷 𝐾 2 𝑸 , then with probability at least 1 -𝛿 , we have:

Proof. Note 𝒉 𝑨 = 1 𝐵 ˝ 𝑖 ˆ ∇ 𝒂 ( 𝑘 𝑖 ) 𝑙 ( 𝑖 ) -E ˆ ∇ 𝒂 ( 𝑘 ) 𝑙 , thus we have the upper bound for each coordinate of the gradient error bounded by Θ ( 𝑁 ) . Similarly, we have the upper bound for each coordinate of the gradient error of 𝒉 𝑽 bounded by Θ ( 𝑁 ) .

Then, we can bound the infinity norm of ⟨ ˆ ∇ 𝑽 𝑙 ( 𝑖 ) , 𝑷 ⟩ 𝜇 and ⟨ ˆ ∇ 𝑨 𝑙 ( 𝑖 ) , 𝑷 ⟩ 𝜇 :

Note that E GLYPH<2> ∇ 𝒂 ( 𝑘 𝑖 ) 𝑙 ( 𝑖 ) -E ∇ 𝒂 ( 𝑘 ) 𝑙 GLYPH<3> = 0 , E [∇ 𝑽 𝑙 ( 𝑖 ) -E ∇ 𝑽 𝑙 ] = 0, which means the two terms above have expectation 0. Since 𝒉 𝑨 , 𝒉 𝑽 are both averages of 𝐵 gradients of a single sample, we use Hoeffding Inequality:

By union bound, if 𝐵 ≥ max GLYPH<0> 𝑁 2 , 𝑄 2 GLYPH<1> 𝑁 4 log 4 𝛿 𝜖 2 𝐾 2 𝑷 𝐾 2 𝑸 it has at least 1 -𝛿 probability, s.t.

Then we finish this section with the concentration of the second order terms ∥ 𝒉 𝑨 ∥ 2 and ∥ 𝒉 𝑽 ∥ 2 , which need to be bounded in the error evolution.

Lemma B.13. Fix 𝜀, 𝛿 > 0 . Under Assumption 2.1, if E ∥ ˆ ∇ 𝑽 𝑙 ∥ 2 𝜇 = Tmp 3 = 𝑂 ( 𝑁 3 ) , E ∥ ˆ ∇ 𝑨 𝑙 ∥ 2 𝜇 = Tmp 4 = 𝑂 ( 𝑇𝑁 2 ) , 𝐵 ≥ Θ ( 𝑇𝑁 2 ) 𝛿𝜖 2 , then with probability at least 1 -𝛿 , we have:

Proof. Similar to Lemma B.12, E 𝒉 𝑽 = 0 , E 𝒉 𝑨 = 0. By Lemma B.9, when we pick 𝐵 ≥ Θ ( 𝑇𝑁 2 ) 𝛿𝜖 2 we have ∥ 𝒉 𝑨 ∥ 𝜇 ≤ 𝜀, ∥ 𝒉 𝑽 ∥ 𝜇 ≤ 𝜀 with probability at least 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 dynamics of projected process and the distance of the true SGD process to the space of population trajectories. All proofs - except for those short ones are deferred to the end of this section.

C.1 Definition of the population projection

The main reason we analyze the population process first is that on the population trajectory, both layers possess special structures. Recall that

and we initialize 𝑽 0 -𝝁 1 ⊤ = 0 and 𝑨 -11 ⊤ / 𝑇 = 0. Note that for any ( 𝒛 𝜏 ) 𝜏 , if 𝒛 0 = 0 and / 𝒛 = 𝐴 𝒛 + 𝐵 𝒛 ∗ , then 𝒛 𝜏 ∝ 𝒛 ∗ for all 𝜏 ≥ 0. In other words, 𝒛 𝜏 moves only along the direction 𝒛 ∗ and therefore, can be characterized by a single real number. This is exactly the same case of 𝑽 -𝝁 1 ⊤ and 𝑨 -11 ⊤ / 𝑇 in the population case. Hence, in the population case, 𝑽 stays on the line crossing 𝝁 1 ⊤ and 𝑷 , and 𝑨 stays on the line crossing 11 ⊤ / 𝑇 and 𝑸 .

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 𝑽 , 𝑨 satisfying 1 ⊤ 𝑁 𝑽 = 1 ⊤ 𝑁 , 1 ⊤ 𝑇 𝑨 = 1 ⊤ 𝑁 , and 𝑽𝝁 = 𝝁 , we define

By setting the derivative to be 0, we can obtain the following closed-form formulas for 𝛼 𝑉 and 𝛼 𝐴 .

Note: Without specification, we drop the the time subscript 𝜏 and consider 𝛼 𝑉 : = 𝛼 𝑉,𝜏 for similicity.

For notational simplicity, we define 𝛽 𝑉 = 1 -𝛼 𝑉 , 𝛽 𝐴 = 1 -𝛼 𝐴 ,

Then, define 𝚫 𝑉 = 𝑽 -˜ 𝑽 and 𝚫 𝐴 = 𝑨 -˜ 𝑨 so that we can decompose 𝑽 = ˜ 𝑽 + 𝚫 𝑉 and 𝑨 = ˜ 𝑨 + 𝚫 𝐴 . By our construction, we have ˜ 𝑽 -𝝁 1 ⊤ ⊥ 𝚫 𝑉 and similarly for 𝚫 𝐴 . We will now show that we can in fact drop 𝝁 1 ⊤ .

Lemma C.2. For any 𝑽 ′ = 𝜃 𝑷 + ( 1 -𝜃 ) 𝝁 1 ⊤ and 𝑨 ′ = 𝜃 𝑸 + ( 1 -𝜃 ) 11 ⊤ / 𝑇 with 𝜃 ∈ R , we have ⟨ 𝚫 𝑉 , 𝑽 ′ ⟩ 𝜇 = 0 and ⟨ 𝚫 𝐴 , 𝑨 ′ ⟩ 𝜇 = 0 . In particular, we have ⟨ 𝚫 𝑉 , 𝑷 ⟩ 𝜇 = GLYPH<10> 𝚫 𝑉 , ˜ 𝑽 GLYPH<11> 𝜇 = 0 and

The following lemma the basic definitions and results about the population projection.

Lemma C.3 (Definitions and basic results on the population projection) . Suppose that 𝑽 , 𝑨 satisfy 1 ⊤ 𝑁 𝑽 = 1 ⊤ 𝑁 , 1 ⊤ 𝑇 𝑨 = 1 ⊤ 𝑁 , and 𝑽𝝁 = 𝝁 . We define the following:

Moreover, by Lemma C.2, the following hold.

C.2 Dynamics of the population projected process and the approximation error

We write

where the expectations are taken over the fresh samples at step 𝜏 .

First, we expand the expected preconditioned gradients around the population projection.

Lemma C.4 (Expanding the gradients) .

Then, we compute the dynamics of the projected process.

Note that ˜ 𝑽 = 𝝁 1 ⊤ + 𝛼 𝑉 ( 𝑷 -𝝁 1 ⊤ ) and ˜ 𝑨 = 1 𝑇 1 ⊤ 𝑁 / 𝑇 + 𝛼 𝐴 ( 𝑸 -1 𝑇 1 ⊤ 𝑁 / 𝑇 ) . Hence, this also gives formulas for ˜ 𝑽 𝜏 + 1 and ˜ 𝑨 𝜏 + 1 .

Now, we consider the dynamics of the errors.

Lemma C.6 (Dynamics of the errors) .

C.3 Omitted proofs in this section

Proof of Lemma C.1. We compute

Set the derivative to be 0, and we get 𝛼 𝑉 = (⟨ 𝑽 , 𝑷 ⟩ 𝜇 - ∥ 𝝁 ∥ 2 )/ 𝐾 𝑃 . Similarly, we compute

Again, set the derivative to be 0, and we get 𝛼 𝐴 = GLYPH<16> ⟨ 𝑨 , 𝑸 ⟩ 𝜇 -1 / 𝑇 GLYPH<17> /(∥ 𝑸 ∥ 2 𝜇 -1 / 𝑇 ) .

Proof of Lemma C.4. Recall from Lemma B.8 that

First, consider the dynamics of 𝑨 . By Lemma C.3, we can further decompose it as

Similarly, we can rewrite the expected preconditioned gradient of 𝑽 as

Proof of Lemma C.5. Recall that 𝛼 𝑉 = ⟨ 𝑽 , 𝑷 ⟩ 𝜇 / 𝐾 𝑃 and 𝛼 𝐴 = ⟨ 𝑨 , 𝑸 ⟩ 𝜇 / 𝐾 𝑄 . First, consider the dynamics of 𝑽 . By Lemma C.4 and Lemma C.3, we have

Similarly, for 𝑽 , we have

Proof of Lemma C.6. First, consider the dynamics of 𝚫 𝑨 , which is given by

Decompose 𝑨 into ˜ 𝑨 + 𝚫 𝑨 , rearrange terms, and we obtain

Now, consider 𝚫 𝑽 . Similar to the previous calculation, we have

Again, note that 𝚫 𝑽 ⊥ 𝜇 𝑷 -𝝁 1 ⊤ , ∥ 𝑷 -𝝁 1 ⊤ ∥ 2 𝜇 = 𝐾 𝑃 , and ⟨ 𝝁 1 ⊤ , 𝒉 𝑽 ⟩ 𝜇 = 0. Hence, we have

D Stage 1: signal boosting

In this section, we assume both 𝛼 𝑉 and 𝛼 𝐴 are close to 0. In this case, we can approximate Lemma C.5 with

We can also write this matrix form as

Suppose that ∥ 𝚫 𝐴 ∥ 2 𝜇 / 𝐾 𝑄 and ∥ 𝚫 𝑉 ∥ 2 𝜇 / 𝐾 𝑃 are both bounded by 𝛿 2 . Then, we have

As long as 𝛿 ≪ 1 and we choose a sufficiently large batch size so that the second line is bounded by 𝜂 /( 2 𝑇𝐾 𝑄 ) , 𝛼 𝑉 + 𝛼 𝐴 grows exponentially fast. Similarly, one can also bound the difference between 𝛼 𝑉 and 𝛼 𝐴 . Formally, we have the following lemma.

Lemma D.1 (Main result of Stage 1) . Define the end of Stage 1 as

Suppose 𝜂 ≤ 1 / 10 , 𝜂 𝑉,𝜏 = 𝜂 / 𝐾 𝑄 , and 𝜂 𝐴,𝜏 = 𝜂 / 𝐾 𝑃 for some 𝜂 ≤ min { 𝐾 𝑃 , 𝐾 𝑄 } . Let 𝐵 𝜏 > 0 be the number fresh samples we use at step 𝜏 . Suppose that 𝐵 𝜏 ≥ ˜ 𝑂 GLYPH<18> 𝑇 2 𝑄 4 𝑁 5 min { 𝐾 3 𝑃 ,𝐾 3 𝑄 } GLYPH<19> are chosen s.t. with probability 1 -𝛿 𝜏

Then, the following hold with probability at least 1 -𝛿 𝑃 :

(a) T 1 ≤ Θ ( 1 ) log GLYPH<0> 𝑇𝐾 𝑄 GLYPH<1> / 𝜂 .

(b) Throughout Stage 1, ∥ 𝚫 𝐴 ∥ 2 𝜇 ≤ Δ / 𝑇 and ∥ 𝚫 𝑉 ∥ 2 𝜇 ≤ Δ / 𝑇 , where

(c) At T 1 , we have 𝛼 𝑉 + 𝛼 𝐴 = Θ ( 1 𝑄𝑁 ) and | 𝛼 𝑉 -𝛼 𝐴 | ≤ Θ ( 1 ) 𝑇𝐾 𝑄 .

(d) For all 𝜏 ≤ T 1 , we have ∥ 𝑽 ∥∞ ≤ Θ ( 1 ) , ∥ 𝑨 ∥∞ ≤ Θ ( 1 / 𝑄 ) .

The proof of this lemma is a large induction argument. We will first assume the bounds on ∥ 𝚫 𝐴 ∥ 𝜇 and ∥ 𝚫 𝑉 ∥ 𝜇 are true, so that the approximation (11) is valid. This will give us an upper bound on the length of Stage 1. 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 𝛼 𝑉 𝛼 𝐴 ≤ 1 4 by definition of Stage I, we can rewrite Lemma C.5 as

where the second line comes from induction hypothesis (b) and (12). Recursively expand the RHS, and we obtain

Since the RHS is upper bounded by 1 by definition of stage I, and we obtain

Part II of the proof of Lemma D.1: Upper bounds on ∥ 𝚫 ∥ 𝜇

Proof. Recall from Lemma C.6 that

By part I of the proof, Stage 1 takes at most Θ ( 1 ) log GLYPH<0> 𝑇𝐾 𝑄 GLYPH<1> / 𝜂 steps. Hence, it suffices to bound the increase of these ∥ 𝚫 ∥ 2 𝜇 in this many steps. Recall the notation Δ = min n 1 4 𝐾 𝑄 , 1 4 𝐾 𝑃 , Θ ( 1 ) 𝑄𝑁 1 / 2 o . By induction hypothesis (b), we have for all 𝜏 ,

Since Stage I at most takes T 1 = Θ ( 1 ) log GLYPH<0> 𝑇𝐾 𝑄 GLYPH<1> / 𝜂 steps, the increase of GLYPH<13> GLYPH<13> 𝚫 𝑨 , 𝜏 GLYPH<13> GLYPH<13> 2 𝜇 and GLYPH<13> GLYPH<13> 𝚫 𝑨 , 𝜏 GLYPH<13> GLYPH<13> 2 𝜇 in 𝜏 ≤ T 1 are at most (since 𝚫 𝑨 , 0 = 0 , 𝚫 𝑽 , 0 = 0):

Therefore, we completed the induction for the error terms.

Part III of the proof of Lemma D.1: Ending state

Proof. First, consider the distance between 𝛼 𝑉 and 𝛼 𝐴 . Similar to the part I of the proof, we rewrite Lemma C.5 as

Thus, whenever 𝛼 𝑉,𝜏 -𝛼 𝐴,𝜏 ≥ Ω ( 1 /( 𝑇𝐾 𝑄 )) , it will start to decrease. Since the amount of increase at each step is also upper bounded by 𝑂 ( 1 /( 𝑇𝐾 𝑄 )) , this implies | 𝛼 𝑉,𝜏 -𝛼 𝐴,𝜏 | ≤ 𝑂 ( 1 /( 𝑇𝐾 𝑄 )) . The other direction can be proved in the same way.

Finally, we bound the possible amount of overshot. By the part I of the proof, we can also upper bound the signal term growth

Since 𝛼 𝑉,𝜏 + 𝛼 𝐴,𝜏 ≤ 1 for all 𝜏 ≤ T 1, we have 𝛼 𝑉 + 𝛼 𝐴 = Θ ( 𝑄𝑁 ) at time T

Part IV of the proof of Lemma D.1: Upper bound on Infinity norm of 𝑽 and 𝑨

Here we consider the upper bound of the weights 𝑽 and 𝑨 , which can be used in the concentration section below.

Proof. First, we upper bound the infinity norm of 𝑽 𝜏 .

and we can upper bound ∥ 𝚫 𝑽 ∥ 𝐹 by its 𝜇 -norm:

Thus we have ∥ 𝑽 𝜏 ∥∞ ≤ GLYPH<13> GLYPH<13> ˜ 𝑽 𝜏 GLYPH<13> GLYPH<13> ∞ + Θ ( 1 ) √ 𝑁 ∥ 𝚫 𝑽 ∥ 𝜇 ≤ GLYPH<13> GLYPH<13> ˜ 𝑽 𝜏 GLYPH<13> GLYPH<13> ∞ + Θ ( 1 𝑄 2 𝑁 ) by Induction hypothesis (b). And we can further bound GLYPH<13> GLYPH<13> ˜ 𝑽 𝜏 GLYPH<13> GLYPH<13> :

∞

Therefore, we have ∥ 𝑽 𝜏 ∥∞ ≤ Θ ( 1 ) .

Similarly, for 𝑨 we have (since GLYPH<13> GLYPH<13> ˜ 𝑨 𝜏 GLYPH<13> GLYPH<13> ∞ = GLYPH<13> GLYPH<13> 1 𝑇 1 ⊤ 𝑁 / 𝑇 + 𝛼 𝐴,𝜏 ( 𝑸 -1 𝑇 1 ⊤ 𝑁 / 𝑇 ) GLYPH<13> GLYPH<13> ∞ ≤ 𝐶 / 𝑄 .)

Part V of the proof of Lemma D.1: Concentration

Finally, we need to ensure that with high probability, all the error terms 𝒉 cannot exceed the given bounds throughout 𝑡 ≤ T 1. We use Lemma B.12 and Lemma B.13 to bound the concentration of the error terms.

By Lemma B.13 and union bound, we have that if 𝐵 𝜏 ≥ Θ ( 1 ) T 1 · 𝑇 2 𝑄 4 𝑁 5 T 1 log 2 ( 𝑇𝐾 𝑄 ) 𝛿 𝜏 min { 𝐾 3 𝑃 ,𝐾 3 𝑄 , 1 } , then with probability at least 1 -𝛿 𝜏 / 2, the following holds for all 𝑡 ≤ T 1:

By Lemma B.12 and union bound, we have that if 𝐵 𝜏 ≥ Θ ( 1 ) T 1 · 𝑇 2 max { 𝑁 2 ,𝑄 2 } 𝑁 4 log GLYPH<18> 16 log ( 𝑇𝐾 𝑄 ) 𝛿𝜏 𝜂 GLYPH<19> 𝐾 2 𝑃 , then with probability at least 1 -𝛿 𝜏 / 2, the following holds for all 𝑡 ≤ T 1:

And by union bound, we have that with probability at least 1 -𝛿 𝜏 , all the bounds above hold for all 𝑡 ≤ T 1. Therefore, we conclude the proof of Lemma D.1.

E Stage 2: learning the model

In this Stage 2, we use a positive 𝜆 for the ℓ 1-regularization. In this case, we can write the update rule of each 𝒂 ( 𝑘 ) as

For notational simplicity, we define

Then, note that

We further define 𝑮 𝜆,𝜏 = -𝜂 -1 𝐴 ( 𝑨 𝜏 + 1 -𝑨 𝜏 ) as the full gradient of the matrix 𝑨 .

First, we will show that with appropriate rounding at the beginning of Stage 2, we can ensure 𝒈 ( 𝑘 ) 𝜏 ≈ 𝜂 E ˆ ∇ 𝒂 ( 𝑘 ) 𝑙 using only 𝐵 𝜏 ∝ log ( 𝑇 ) fresh samples at each step (Section E.1).

E.1 Rounding and gradient denoising

Recall that the first step of the Stage 2 is rounding each 𝒂 ( 𝑘 ) by setting all small coordinates to 0 and then projecting it back to the affine space the probability simplex lies. Since after Stage 1, there will be a separation between 𝑎 ( 𝑘 ) 𝑡 with 𝑡 ∈ 𝒒 ( 𝑘 ) and 𝑡 ∉ 𝒒 ( 𝑘 ) , this rounding step makes all 𝑎 ( 𝑘 ) 𝑡 with 𝑡 ∉ 𝒒 ( 𝑘 ) have the same small value.

Also recall that we use ℓ 1-regularization and proximal gradients in the Stage 2. Effectively, the ℓ 1-regularization ensures those useless 𝑎 ( 𝑘 ) 𝑡 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 𝒒 ( 𝑘 ) , analyze its influence on 𝛼 𝐴 and the distance to the population subspace. Then, we analyze the effect of use of the proximal gradients and show that with poly ( 𝑁, 𝑄 ) log ( 𝑇 ) 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

Then, we can choose a threshold 𝜃 Tmp = Θ ( 𝛼 𝑉 / 𝑄 ) s.t. 𝑎 ( 𝑘 ) 𝑡 ≤ 𝜃 Tmp iff 𝑞 ( 𝑘 ) 𝑡 = 0 .

Proof. Note that there exists some universal constant 𝑐 > 0 such that 𝒒 ( 𝑘 ) 𝑡 ≥ 𝑐 / 𝑄 for all nonzero 𝑞 ( 𝑘 ) 𝑡 and 𝜇 𝑘 ≥ 𝑐 / 𝑁 for all 𝑘 ∈ [ 𝑁 ] . Recall that the population process ˜ 𝑨 = 𝛼 𝐴 𝑸 + ( 1 -𝛼 𝐴 ) 1 𝑇 1 ⊤ 𝑁 / 𝑇 . Hence, for all 𝑘 ∈ [ 𝑁 ] and 𝑡 ∈ [ 𝑇 ] ,

Hence, for any 𝑘 ∈ [ 𝑁 ] and 𝑡 ∈ [ 𝑇 ] , we have | 𝑎 ( 𝑘 ) 𝑡 -˜ 𝑎 ( 𝑘 ) 𝑡 | ≤ ∥ 𝚫 𝑨 ∥ 𝜇 √︁ 𝑁 / 𝑐 . Combine these together, and we obtain

Hence, in order to get a separation,

Lemma E.2 (Effect of rounding) . Under the conditions of Lemma E.1, choose 𝜆 as in Lemma E.1, and set

We have 𝛼 𝐴 ← 𝛼 𝐴 + 𝑂 ( 1 )/ 𝑇 and ∥ 𝚫 𝐴 ∥ 2 𝜇 ←∥ 𝚫 𝐴 ∥ 2 𝜇 + 𝑂 ( 1 )/ 𝑇 .

Proof. For notational simplicity, put 𝒃 ( 𝑘 ) = 𝒂 ( 𝑘 ) ⊙ ( ✶ { 𝑎 ( 𝑘 ) 𝑡 ≥ 𝜆 }) 𝑇 𝑡 = 1 . By Lemma E.1, we know 𝒃 ( 𝑘 ) is supported within 𝒒 ( 𝑘 ) . Set 𝑏 ( 𝑘 ) = ˝ 𝑇 𝑡 = 1 𝑏 ( 𝑘 ) 𝑡 . Then, we can write

Recall that 𝛼 𝐴 = 𝐾 𝐴𝑄 / 𝐾 𝑄 where 𝐾 𝐴𝑄 = ⟨ 𝑨 , 𝑸 ⟩ 𝜇 = ˝ 𝑁 𝑘 = 1 𝜇 𝑘 GLYPH<10> 𝒂 ( 𝑘 ) , 𝒒 ( 𝑘 ) GLYPH<11> . Hence, for each 𝑘 ∈ [ 𝑁 ] , we have

Now, consider the effect of projection on the distance to the population subspace. We will use subscript new to indicate values after rounding and use notations such as 𝒂 ( 𝑘 ) to denote the values before rounding. Recall that ∥ 𝚫 𝑨 ∥ 2 𝜇 = ˝ 𝑁 𝑘 = 1 𝜇 𝑘 GLYPH<13> GLYPH<13> 𝒂 ( 𝑘 ) -˜ 𝒂 ( 𝑘 ) GLYPH<13> GLYPH<13> 2 . We have

Note that

As a result,

Hence,

As we have seen in Stage 1, the norm of ∇ 𝒂 ( 𝑘 ) 𝑙 can scale linearly with 𝑇 . Hence, in order to make 𝒉 𝒂 ( 𝑘 ) , 𝜏 has size 𝑂 ( 1 ) in terms of ∥·∥ 2 , it is necessary to use poly ( 𝑇 ) samples, which is undesirable. However, note that all entries of ∇ 𝒂 ( 𝑘 ) 𝑙 are bounded, and therefore are subgaussian. Hence, for each entry of ∇ 𝒂 ( 𝑘 ) 𝑙 , with 𝑂 ( log 𝑇 ) samples, we can make sure the relative error is small with probability at least 1 -1 / poly ( 𝑇 ) . By union bound, this means the ∥·∥∞ error can be made small using only log ( 𝑇 ) samples. Note that there is a separation between the signal and noise parts of E ∇ 𝒂 ( 𝑘 ) 𝑙 . This implies that we can distinguish them using log ( 𝑇 ) samples and directly remove the noise part. Formally, we have the following lemma.

Lemma E.3 (Gradient denoising) . Given 𝜀 ∈ GLYPH<16> Θ GLYPH<16> 𝑄 2 𝑁 3 √ 𝑇 GLYPH<17> , 0 . 1 GLYPH<17> . Suppose that for any 𝑘, 𝑚, 𝑛 ∈ [ 𝑁 ] , | 𝒂 ( 𝑘 ) 𝑡 | ≤ 1 / 𝑇 if 𝑡 ∉ 𝒒 ( 𝑘 ) , | 𝑉 𝑛,𝑚 | ≤ 𝑂 ( 1 ) , and ∥ 𝚫 𝐴 ∥ 𝜇 ≤ 𝑐 ( 1 -𝛼 𝑉 )/( √ 𝑁𝑄 ) and ∥ 𝚫 𝑉 ∥ 2 𝜇 ≤ 𝑐 ( 1 -𝛼 𝑉 ) 𝐾 𝑃 for some small constant 𝑐 > 0 . For the target accuracy 𝜀 Tmp = Θ ( 𝜀 𝑄𝑁 2 ) , we have with probability 1 -𝛿 𝜏 . If we choose

for some large constant 𝐶 > 0 . Then, for each 𝑘 ∈ [ 𝑁 ] , with probability at least 1 -𝛿 Tmp , we have

Proof. For 𝑠 ∈ [ 𝑇 ] , recall that

First, consider the expectations. If 𝑠 ∉ 𝒒 ( 𝑘 ) , then by our assumption, we have | E 𝜕 𝑎 ( 𝑘 ) 𝑠 𝑙 | = 𝜇 𝑘 𝐾 𝑉 𝑎 𝑠 = 𝑂 ( 𝐾 𝑉 /( 𝑁𝑇 )) . Meanwhile, for 𝑠 ∈ 𝒒 ( 𝑘 ) , we have

As a result, we have for any 𝑘 ∈ [ 𝑁 ] and 𝑠 ∈ 𝒒 ( 𝑘 ) ,



for some small constant 𝑐 > 0. Then, for the size of each entry, we have

Therefore, 𝜕 𝑎 ( 𝑘 ) 𝑠 𝑙 is 𝑁 2 𝐾 4 Tmp -subgaussian, whence 𝜕 ( 𝐵 𝜏 ) 𝑎 ( 𝑘 ) 𝑠 𝑙 is 𝑁 2 𝐾 4 Tmp / 𝐵 𝜏 -subgaussian. As a result, for each 𝜉 > 0, we have

Apply union bound and we get

Recall that the separation between the expectations is 𝑐𝛼 𝑉 𝐾 𝑃 /( 𝑁𝑄 ) . Hence, it suffices to choose 𝜉 = 𝑐𝛼 𝑉 𝐾 𝑃 /( 2 𝑁𝑄 ) . Then, to make the failure probability at most 𝛿 Tmp , we can choose 𝐵 𝜏 as follows:

for some large constant 𝐶 > 0. Then, to boost the accuracy from 1 / 2 to 𝜀 Tmp , it suffices to increase the batch size to 𝐶 𝑁 4 𝑄 2 𝜀 2 Tmp 𝛼 2 𝑉 𝐾 2 𝑃 log GLYPH<16> 𝐶𝑁𝑇 𝛿 𝜏 GLYPH<17> .

Using this lemma, we can pick 𝜆 = Θ ( 𝜀𝐾 𝑃 𝑄 2 𝑁 3 ) to sparsify our proximal gradient. Notice that our proximal gradient is a 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 ˆ 𝒉 𝑨 , 𝜏 = 𝑮 ( 𝑘 ) 𝜆,𝜏 -E ˆ ∇ 𝑨 𝑙 instead of 𝒉 because of the bias introduced by the proximal gradient.

Lemma E.4. Under the same setting of Lemma E.3, if the batch size

the noise attention score | 𝑎 ( 𝑘 ) 𝑡 | ≤ 𝑂 ( 1 / 𝑇 ) for all 𝑡 ∉ 𝒒 ( 𝑘 ) and all 𝑘 ∈ [ 𝑁 ] , and 𝜆 = Θ GLYPH<16> 𝜀𝐾 𝑃 𝑄 2 𝑁 3 GLYPH<17> , then with probability 1 -𝛿 𝜏 , the gradient error at iteration 𝜏

Proof. For notational simplicity, we drop the superscript 𝑘 . The goal is to estimate the difference between E ˆ ∇ 𝒂 𝑙 and 𝒈 𝜆,𝜏 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 { 𝑎 𝑡,𝜏 } 𝑡 ∉ 𝒒 are all the same in Stage 2, we can write

and recall the projected preconditioned gradient for 𝑨 .

Therefore, we have

Now, we consider 𝒈 by calculating the expression of 𝒂 𝜏 + 1. First, note that by Lemma E.3, we have 𝑎 ′′ 𝑡,𝜏 + 1 = GLYPH<16> 𝑎 ′ 𝑡,𝜏 + 1 -𝜆 GLYPH<17> ✶ { 𝑡 ∈ 𝒒 } . This implies that { 𝑎 𝑡,𝜏 + 1 } 𝑡 ∉ 𝒒 are all the same and the value is at most 1 / 𝑇 . Moreover, it also implies that it suffices to focus on [ 𝒂 ′′ 𝜏 + 1 ] 𝒒 , for which we have

Therefore, ˝ 𝑡 ∈ 𝒒 𝑎 ′ 𝑡,𝜏 + 1 = 𝑏 𝜏 -𝜂 𝜇 𝑘 ˝ 𝑡 ∈ 𝒒 𝜕 ( 𝐵 𝜏 ) 𝑎 𝑡 𝑙 -𝑄𝜆 and

Thus, we can write an explicit update

Then the gradient error at step 𝜏 can be decomposed into:

Here the gradient bias error can be further simplified to

First, we estimate the concentration error. Similar to Lemma B.10, we first upper bound the infinity norm of the gradient. Consider the maximum absolute value in the original gradients. For ∇ 𝒂 ( 𝑘 ) 𝑙 , we have

The first term can be upper-bounded in the following way:

The second inequality is due to | 𝑎 ( 𝑘 ) 𝑡 | ≤ 𝑂 ( 1 / 𝑇 ) for 𝑡 ∉ 𝒒 . Since 𝑽 𝑛,𝑚 ≤ 𝑂 ( 1 ) , we have max 𝑠,𝑡 GLYPH<12> GLYPH<12> ( 𝑽𝒆 𝑥 𝑠 ) ⊤ 𝑽𝒆 𝑥 𝑡 GLYPH<12> GLYPH<12> upper bounded by 𝑂 ( 1 ) . Therefore

And similarly, the second term max 𝑠 ∈[ 𝑇 ] GLYPH<12> GLYPH<12> ( 𝑽𝒆 𝑥 𝑠 ) ⊤ 𝒆 𝑥 𝑜 GLYPH<12> GLYPH<12> can be bounded by 𝑂 ( 1 ) because the infinity norm of 𝑽 is also upper bounded by 𝑂 ( 1 ) . Therefore, we know ∥∇ 𝒂 ( 𝑘 ) 𝑙 ∥∞ ≤ 𝑂 ( 1 ) .

Now we consider the preconditioned gradient ˆ ∇ ( 𝐵 𝜏 ) 𝒂 𝑙 :

since 𝜇 𝑘 ≥ 𝑐 𝑁 for all 𝑘 ∈ [ 𝑁 ] . Now since h ˆ ∇ ( 𝐵 𝜏 ) 𝒂 𝑙 i 𝒒 is 𝑄 -sparse, we have E GLYPH<13> GLYPH<13> GLYPH<13> GLYPH<2> ˆ ∇ 𝑨 𝑙 GLYPH<3> 𝒒 GLYPH<13> GLYPH<13> GLYPH<13> 2 𝜇 ≤ 𝑂 ( 𝑄𝑁 2 ) . By Lemma B.9, when 𝐵 𝜏 ≥ Θ ( 𝑁 6 𝑄 3 ) 2 , with probability 1 -𝛿 𝜏 ,

𝛿𝜀 2

Then, consider the gradient bias term. With the selected 𝜆 = Θ GLYPH<16> 𝜀𝐾 𝑃 𝑄 2 𝑁 3 GLYPH<17> and GLYPH<13> GLYPH<13> GLYPH<13> 𝜕 ( 𝐵 𝜏 ) 𝑎 𝑡 𝑙 GLYPH<13> GLYPH<13> GLYPH<13> ≤ 𝑂 GLYPH<16> 𝜀𝐾 𝑃 𝑁 3 𝑄 2 GLYPH<17> for 𝑡 ∉ 𝒒 , with probability 1 -𝛿 𝜏 / 2 we have the 𝜇 -norm of first term (since 𝐾 𝑃 ≤ 𝑂 ( 𝑁 ) ):

and the second term can be upper-bounded by

since 1 √ 𝑇 ≤ 𝑂 GLYPH<16> 𝜀 𝑄 2 𝑁 3 GLYPH<17> , 𝜕 ( 𝐵 𝜏 ) 𝑎 𝑡 𝑙 = 𝑂 ( 1 ) , and ∥ 𝑽 ∥ 2 𝜇 ≤ 𝑁 .

Combine all three terms and by union bound, we have with probability 1 -𝛿 𝜏

since 𝜇 𝑘 = Θ ( 1 / 𝑁 ) .

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 1 and the final error. In particular, we show that eventually we will have 𝛼 𝑉 + 𝛼 𝐴 ≈ 2, ∥ 𝚫 𝐴 ∥ 2 𝜇 ≈ 0 and ∥ 𝚫 𝑉 ∥ 2 𝜇 ≈ 0 5 .

For notational simplicity, define 𝛿 2 𝐴 = ∥ 𝚫 𝐴 ∥ 2 𝜇 / 𝐾 𝑄 and 𝛿 2 𝑉 = ∥ 𝚫 𝑉 ∥ 2 𝜇 / 𝐾 𝑃 . Recall Lemma C.5 and Lemma C.6 and that we choose 𝜂 𝑉 = 𝜂 / 𝐾 𝑄 , 𝜂 𝐴 = 𝜂 / 𝐾 𝑃 . The dynamics of the signals and the errors can be described using 6

and

E.2.1 Lemmas for the dynamics

Before we come to the final convergence analysis, we first simplify the dynamics with some basic lemmas.

Lemma E.5 (Dynamics of the errors) . For the errors, we have

Proof. First, we write

«

‹

For the first term, we have

Thus, for 𝛿 𝐴 , we have

«

‹

Similarly, for 𝛿 𝑉 , we have

Lemma E.6 (Dynamics of 𝛼 𝐴 -𝛼 𝑉 ) . The difference between the signals evolves as follows

Proof. First, we write

Therefore, we have

Then, for Tmp , we compute

In particular, this implies | Tmp | ≤ 1. Thus, we have

Lemma E.7. Suppose that both 𝛼 𝑉 , 𝛼 𝐴 are at most 1 . Then, we have

Proof. First, we write

For the signal growth, note that 𝛼 𝐴 𝛼 𝑉 ≤ ( 𝛼 2 𝑉 + 𝛼 2 𝐴 )/ 2 ≤ ( 𝛼 𝑉 + 𝛼 𝐴 )/ 2. Hence,

For the error terms, we have

Thus, we have

and, therefore,

Lemma E.8. Put 𝜀 𝐴 = 1 -𝛼 𝐴 and 𝜀 𝑉 = 1 -𝛼 𝑉 . When | 𝜀 𝐴 | , | 𝜀 𝑉 | ≤ 1 / 2 , we have

«

Proof. Similar to the proof of the previous lemma, we write

= : 𝛼 𝑉,𝜏 + 𝛼 𝐴,𝜏 + 𝜂 Tmp sig + 𝜂 Tmp err .

For the signal term, we have

For the error term, we have

‹

Combine these together, and we obtain

Thus,

Note that

Recall | 𝜀 𝐴 | , | 𝜀 𝑉 | ≤ 1 / 2. Thus,

Lemma E.9. Suppose that ( 𝑋 𝜏 ) 𝜏 satisfies 𝑋 𝜏 + 1 ≤ 𝑒 -𝐴 𝑋 𝜏 + 𝐵 for some 𝐴 ∈ ( 0 , 1 ] , 𝐵 ≥ 0 . If 𝑋 0 ≤ 2 𝐵 / 𝐴 , then we have 𝑋 𝜏 ≤ 3 𝐵 / 𝐴 for all 𝜏 ≥ 0 . If 𝑋 0 ≥ 2 𝐵 / 𝐴 , then we have 𝑋 𝜏 ≤ 3 𝐵 / 𝐴 for all 𝜏 ≥ 6 𝐴 log GLYPH<16> 𝑋 0 𝐴 3 𝐵 GLYPH<17> .

Proof. Since 𝑒 -𝐴 ≤ 1 -𝐴 / 2 for 𝐴 ∈ ( 0 , 1 ] , we have 𝑋 𝜏 + 1 ≤ 𝑋 𝜏 -𝐴𝑋 𝜏 / 2 + 𝐵 . Hence, whenever 𝑋 𝜏 ≥ 2 𝐵 / 𝐴 , we will have 𝑋 𝜏 + 1 ≤ 𝑋 𝜏 . Moreover, if 𝑋 𝜏 < 2 𝐵 / 𝐴 , we have 𝑋 𝜏 + 1 ≤ 2 𝐵 / 𝐴 + 𝐵 ≤ 3 𝐵 / 𝐴 . This proves the first part of the lemma.

Now, suppose that 𝑋 0 ≥ 2 𝐵 / 𝐴 . When 𝑋 𝜏 ≥ 3 𝐵 / 𝐴 , we have

Thus, it takes at most 6 𝐴 log GLYPH<16> 𝑋 0 𝐴 3 𝐵 GLYPH<17> steps to reduce 𝑋 𝜏 from 𝑋 0 to 3 𝐵 / 𝐴 . After that, the previous analysis applies.

E.2.2 Main lemma of Stage 2

We split the analysis of Stage 2 into two substage. Let 𝑐 𝛼 ∈ ( 0 , 0 . 05 ) be a small constant. Define

We call { 𝜏 : 𝜏 ≤ T 2 . 1 } stage 2.1 and { 𝜏 : 𝜏 ≥ T 2 . 1 } stage 2.2. For notational simplicity, we define

  Note that this can be made essentially arbitrarily small by choosing a large enough batch size.

Lemma E.10. Suppose that the following hold at the beginning of Stage 2 (after thresholding and projection):

Let 𝛿 ∗ be our target value for 𝛿 𝐴 and 𝛿 𝑉 . Choose 𝜆 as in Lemma E.4. Suppose that

Then, within 𝑂 ( 1 /( 𝜂𝛼 ( 2 ) ) + log ( 1 / 𝛿 ∗)/ 𝜂 ) steps, we will have 𝛿 2 𝐴 , 𝛿 2 𝑉 ≤ 𝛿 2 ∗ and 𝛼 𝑉 , 𝛼 𝐴 ∈ ( 0 . 9 , 1 . 1 ) .

Remark . Note that our conditions on 𝛼 𝑉 , 𝛼 𝐴 and 𝛿 ( 2 ) are much weaker that what one can obtain from Stage 1. This allows us to apply the analysis here to transfer learning. ♣

The following proof should be treated as a large induction argument though we do not explicitly write down the induction as in the proof of Stage 1. In particular, we will show (by induction) that the approximation errors 𝛿 𝐴 and 𝛿 𝑉 are small, so that most of the naïve bounds on the entries of ˜ 𝑨 and ˜ 𝑽 can be transferred to 𝑨 and 𝑽 . In particular, | 𝑉 𝑛,𝑚 | = 𝑂 ( 1 ) for all 𝑛, 𝑚 ∈ [ 𝑁 ] so that our bounds in Section E.1 are valid, and | 𝑎 ( 𝑘 ) 𝑡 | = 𝑂 ( 1 ) for all 𝑘 ∈ [ 𝑁 ] , 𝑡 ∈ [ 𝑁 ] , which implies that after the projection step, 𝑎 ( 𝑘 ) 𝑡 = 𝑂 ( 1 / 𝑇 ) for all 𝑡 ∉ 𝒒 ( 𝑘 ) .

Proof of the Lemma E.10

Commonresults 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 𝛼 𝑉 , 𝛼 𝐴 ≥ 𝛼 ( 2 ) and 𝛿 2 𝑉 + 𝛿 2 𝐴 ≤ 𝛿 ( 2 ) hold throughout Stage 2. Recall from Lemma E.5 that

Hence, as long as Errsmall ≤ 𝑂 ( 1 ) 𝛿 ( 2 ) /( 𝛼 ( 2 ) ) 2 , we can ensure 𝛿 2 𝐴,𝜏 + 1 + 𝛿 2 𝑉,𝜏 + 1 ≤ ( 𝛿 ( 2 ) ) 2 always hold.

Stage 2.1: signal growth By Lemma E.7, we have

When ( 𝛼 𝑉 + 𝛼 𝐴 )/ 2 ≤ 1 -𝑐 𝛼 and 𝛼 𝑉 , 𝛼 𝐴 ≥ 𝛼 ( 2 ) , we have

Hence, as long as Errsmall ≤ 𝛼 ( 2 ) 𝑐 𝛼 / 2 and 𝛿 ( 2 ) ≤ 𝑐 𝛼 / 2, we have

Thus, stage 2.1 takes at most 𝑂 ( 1 /( 𝜂𝛼 ( 2 ) )) steps.

Stage 2.1: difference between the 𝛼 's By Lemma E.6, we have

By the AM-GM inequality, we have 𝛼 𝐴 𝛼 𝑉 ≤ (( 𝛼 𝐴 + 𝛼 𝑉 )/ 2 ) 2 ≤ ( 1 -𝑐 𝛼 ) 2 . Therefore, 1 -𝛼 𝐴 𝛼 𝑉 ≥ 𝑐 𝛼 and the above inequality can be further rewritten as

Thus, by (the proof of) Lemma E.9, we have

Stage 2.2: error decrease By Lemma E.5, we have

In other words, Stage 2.2 takes at most 𝑂 GLYPH<16> log ( 1 / 𝛿 ∗ ) 𝜂 GLYPH<17> steps.

Stage 2.2: stability of 𝛼 𝑉 ± 𝛼 𝐴 Recall from Lemma E.8 and Lemma E.6 that

We wish to maintain the induction hypotheses | 𝛼 𝑉 -𝛼 𝐴 | = | 𝜀 𝑉 -𝜀 𝐴 | = 𝜃 -for some 𝜃 -= 𝑜 ( 1 ) and ( 𝛼 𝐴 + 𝛼 𝑉 )/ 2 ≤ 1 + 𝑐 / log ( 1 / 𝛿 ∗ ) = : 1 + 𝜃 + .

First, assume these conditions are true. Then, we have

Since Stage 2.2 takes at most 𝑂 ( log ( 1 / 𝛿 ∗)/ 𝜂 ) steps, when the constant 𝑐 is small, we have

We can choose the parameters appropriately so that 𝜃 -< 𝑜 ( 𝜃 +) . Then, when 𝜀 𝐴 + 𝜀 𝑉 ∈ (-2 𝜃 + , -𝜃 +) , we have, for some large universal constant 𝐶 > 0,

This establishes the induction hypotheses on 𝛼 𝑉 ± 𝛼 𝐴 .

F Stage 3: Final Convergence

As mentioned last section, we cannot ensure 𝛼 𝐴 ≈ 𝛼 𝑉 since 𝛼 𝐴 -𝛼 𝑉 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 𝜀 -error.

First, we formally define our rounding procedure. Let 𝑐 > 0 be a small constant (cf. the proof of Lemma F.1). For each 𝑘 ∈ [ 𝑁 ] , define

✶ This 𝑨 ∗ : = [ 𝒂 (∗ ,𝑘 ) ] 𝑘 ∈[ 𝑁 ] is our rounded version of 𝑨 . For the error between 𝑨 ∗ and 𝑸 , we have the following lemma.

LemmaF.1 (Rounding 𝑨 ) . Let 𝜀 ∈ GLYPH<0> Θ ( 𝑄 2 / 𝑇 2 ) , 0 . 1 GLYPH<1> be our target accuracy. Suppose that 𝛼 𝐴 = 1 -𝜀 𝐴 for some | 𝜀 𝐴 | ≤ 0 . 1 and ∥ 𝚫 𝐴 ∥ 2 𝜇 ≤ 𝛿 2 𝐴 for some 0 < 𝛿 𝐴 ≪ 1 /( 𝑄 √ 𝑁 ) . Then, after rounding, we have

In particular, to achieve 𝜀 accuracy in terms of ∥·∥ 𝜇 , we only need | 𝜀 𝐴 | ≤ 0 . 1 and 𝛿 2 𝐴 ≤ 𝑂 ( 𝜀 / √ 𝑁𝑄 ) .

Remark . In particular, this lemma implies that as long as 𝜀 𝐴 is not too large, after rounding, the error depends solely on ∥ 𝚫 𝐴 ∥ 2 𝜇 . ♣

Proof. For notational simplicity, we omit the superscript 𝑘 for now. Write

Note that ∥ 𝚫 𝒂 ∥ ∞ ≤ ∥ 𝚫 𝒂 ∥ 2 ≤ 𝑂 ( 1 ) √ 𝑁 ∥ 𝚫 𝑨 ∥ 𝜇 ≤ 𝑂 ( √ 𝑁𝛿 𝐴 ) . Since all nonzero 𝑞 𝑠 are lower bounded by Ω ( 1 / 𝑄 ) , we have, for any 𝑠 ∈ 𝒒 and 𝑡 ∉ 𝒒 ,

Hence, we can choose a small constant 𝑐 > 0, so that

✶ where for 𝒗 ∈ R 𝑇 , 𝒗 𝒒 ∈ R 𝑇 is defined as [ 𝑣 𝑡 ✶ { 𝑡 ∈ 𝒒 }] 𝑡 ∈[ 𝑇 ] here. Now, consider the difference between 𝒒 and ˆ 𝒂 / 1 ⊤ ˆ 𝒂 . We have and therefore,

Thus,

Lemma F.2. Let 𝜀 ∈ GLYPH<0> Θ ( 1 )/( 𝐾 𝑄 𝑇 ) , 0 . 1 GLYPH<1> be our target accuracy. Suppose that ∥ 𝑨 -𝑸 ∥ 𝜇 ≤ 𝛿 𝐴, ∗ ≤ 0 . 01 𝜀 and ∥ 𝚫 𝑉 ∥ 𝜇 ≤ 𝛿 𝑉 ≤ 0 . 01 at the beginning of Stage 3, and ∥ 𝒉 𝑽 ∥ 𝜇 ≤ 𝑐𝜀 √ 𝐾 𝑃 𝐾 𝑄 for all 𝜏 and a sufficiently small constant 𝑐 . Then, we have ∥ 𝑽 -𝑷 ∥ 2 𝜇 ≤ 𝜀 for all 𝜏 ≥ T 2 + Θ ( log ( 1 / 𝜀 )/ 𝜂 ) .

Proof. Under the condition ∥ 𝑨 -𝑸 ∥ 𝜇 ≤ 𝛿 𝐴, ∗ , we have ∥ 𝚫 𝐴 ∥ 𝜇 ≤ ∥ 𝑨 -𝑸 ∥ 𝜇 ≤ 𝛿 𝐴, ∗ and

Recall that we only train 𝑽 in Stage 3. Note that by Lemma C.3,

Hence, to get 𝜀 accuracy, it suffices to have ( 1 -𝛼 𝑉 ) 2 ≤ 𝜀 / 2 and ∥ 𝚫 𝑉 ∥ 2 𝜇 ≤ 𝜀 / 2.

First, for 𝛼 𝑉 , by Lemma C.5 and 𝜂 𝑉 = 𝜂 / 𝐾 𝑄 , we have

Hence,

For ∥ 𝚫 𝐴 ∥ 2 𝜇 , by (the proof of) Lemma E.5, we have

Thus, by Lemma E.9, we have ( 1 -𝛼 𝑉 ) 2 ≤ 𝜀 / 2and ∥ 𝚫 𝑉 ∥ 2 𝜇 ≤ 𝜀 / 2for all 𝜏 ≥ T 2 + Θ ( log ( 1 / 𝜀 )/ 𝜂 ) .

Corollary F.3. Let 𝜀 ∈ ( Θ ( 1 )/( 𝐾 𝑄 𝑇 ) , 0 . 1 ) be our target accuracy. Suppose that 𝛼 𝐴 ∈ ( 0 . 9 , 1 . 1 ) , ∥ 𝚫 𝐴 ∥ 2 𝜇 ≤ 𝑂 ( 𝜀 /( 𝑄 √ 𝑁 )) , and ∥ 𝚫 𝑉 ∥ 2 𝜇 ≤ 0 . 01 . Then with poly ( 𝑁, 𝑄, 1 / 𝜀 ) samples, we have with high probability that ∥ 𝑨 -𝑸 ∥ 2 𝜇 ≤ 𝜀 and ∥ 𝑽 -𝑷 ∥ 2 𝜇 ≤ 𝜀 after Stage 3, which takes 𝑂 ( log ( 1 / 𝜀 )/ 𝜂 ) 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 our target accuracy and T 1 = min { 𝜏 ≥ 0 : max { 𝛼 𝑉,𝜏 , 𝛼 𝐴,𝜏 } ≥ Θ ( 1 /( 𝑄𝑁 ))} . We can choose the hyperparameters in Algorithm 1 such that within 𝑂 GLYPH<0> log ( 𝑇 )/ 𝜂 1 + 1 /( 𝜂𝛼 ( 2 ) ) + log ( 1 / 𝜀 )/ 𝜂 GLYPH<1> steps, we have ∥ 𝑨 -𝑸 ∥ 2 𝜇 ≤ 𝜀 and ∥ 𝑽 -𝑷 ∥ 2 𝜇 ≤ 𝜀 with probability at least 1 -𝛿 and the number of samples used before and after T 1 are poly ( 𝑇, 𝛿 ) and poly ( 𝑁, 𝑄, 1 / 𝜀, log 𝑇, 𝛿 ) , respectively.

Proof. The results for Stage 1 follow directly from Lemma D.1. Now, consider the results for Stage 2 and 3. First, by Corollary F.3, it suffices to make sure at time T 2, we have 𝛼 𝐴 ∈ ( 0 . 9 , 1 . 1 ) , ∥ Δ 𝐴 ∥ 2 𝜇 ≤ 𝑂 ( 𝜀 /( 𝑄 √ 𝑁 )) and ∥ 𝚫 𝑉 ∥ 2 𝜇 ≤ 0 . 01 (with high probability). By Lemma D.1, we know the following hold at time T 1 w.h.p:

Therefore, by Lemma E.1, if we choose the threshold to be Θ ( 1 /( 𝑄 2 𝑁 )) , then after thresholding and projection, we have 𝑎 ( 𝑘 ) 𝑡 = 𝑂 ( 1 / 𝑇 ) for all 𝑡 ∉ 𝒒 ( 𝑘 ) . Thus, by Lemma E.3 and Lemma E.4, with poly ( 𝑄, 𝑁, 1 / 𝜀, log ( 𝑇 )) samples, we have with high probability that 𝑎 ( 𝑘 ) 𝑡 = 𝑂 ( 1 / 𝑇 ) for all 𝑡 ∉ 𝒒 ( 𝑘 ) holds and GLYPH<13> GLYPH<13> GLYPH<13> ˆ 𝒉 𝐴 GLYPH<13> GLYPH<13> GLYPH<13> 𝜇 satisfies the requirements in Lemma E.10 throughout Stage 2. Thus, by Lemma E.10 with 𝛿 2 ∗ = 𝑂 ( 𝜀 / 𝑄 √ 𝑁 ) , at the end of Stage 2, we have with high probability that 𝛼 𝐴 ∈ ( 0 . 9 , 1 . 1 ) , ∥ Δ 𝐴 ∥ 2 𝜇 ≤ 𝑂 ( 𝜀 /( 𝑄 √ 𝑁 )) and ∥ 𝚫 𝑉 ∥ 2 𝜇 ≤ 0 . 01. When combined with Corollary F.3, this completes the proof.

H Transfer learning

Lemma H.1 (Initialization) . Suppose that we have learned ˆ 𝑷 and GLYPH<10> ˆ 𝑷 , 𝑷 GLYPH<11> 𝜇 ≥ 2 ∥ 𝝁 ∥ 2 , GLYPH<13> GLYPH<13> ˆ 𝑷 GLYPH<13> GLYPH<13> 2 𝜇 = Θ ( 1 ) ∥ 𝑷 ∥ 2 𝜇 . Let 𝑽 = 𝜃 ˆ 𝑷 + ( 1 -𝜃 ) 𝝁 1 ⊤ for some 𝜃 ∈ ( 0 , 1 ) . We have

Proof. First, consider 𝚫 𝑉 . Since ∥ 𝚫 𝑉 ∥ 𝜇 is the distance to a projection, we have

For 𝛼 𝑉 , we compute

In particular, this implies 𝛼 𝑉 ≥ Θ GLYPH<16> 𝜃 𝑁𝐾 𝑃 GLYPH<17> .

Lemma H.2 (First gradient step) . Let 𝑨 = 11 ⊤ / 𝑇 and 𝑽 given by Lemma H.1. Let 𝜀 Tmp ≤ 𝑂 ( 1 /( 𝑁𝐾 𝑃 )) . Run one gradient step with poly ( 𝑁, 𝑄, log 𝑇, 1 / 𝜀 Tmp ) samples and 𝜂 = 1 , remove all entries of 𝒂 ( 𝑘 ) with | 𝑎 ( 𝑘 ) 𝑡 | ≤ Θ ( 𝜃 /( 𝑁𝐾 𝑃 )) and the replace 𝒂 ( 𝑘 ) with the projection ( 𝑰 -11 ⊤ / 𝑇 ) 𝒂 ( 𝑘 ) .

With high probability, we have 𝛼 𝐴 = ( 1 ± 𝑂 ( 𝜀 Tmp )) 𝛼 𝑉 and ∥ 𝚫 𝐴 ∥ 2 𝜇 ≤ 𝑂 GLYPH<18> 𝜀 2 Tmp 𝛼 2 𝑉 𝑄 + 1 𝑇 GLYPH<19> .

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 𝑠 ∈ [ 𝑇 ] , we have

Also recall from Lemma C.3 that 𝐾 𝑉𝑃 = 𝛼 𝑉 𝐾 𝑃 and 𝐾 𝑉 = 𝛼 2 𝑉 𝐾 𝑃 + ∥ 𝚫 𝑉 ∥ 2 𝜇 . For any 𝑠 ∈ 𝒒 ( 𝑘 ) , we have

where the inequality comes from 𝜇 𝑘 𝛼 𝑉 𝐾 𝑃 𝑞 𝑠 ≥ Ω ( 𝛼 𝑉 / 𝑁 2 / 𝑄 ) ≫ 1 / 𝑇 . Meanwhile, for any 𝑠 ∉ [ 𝑇 ] , we have E 𝜕 𝑎 ( 𝑘 ) 𝑠 𝑙 = 𝑂 ( 1 / 𝑇 ) .

Meanwhile, for any 𝑠 ∈ [ 𝑇 ] we have

Thus, by some standard concentration argument similar to the one in Lemma E.3, we can show that with poly ( 𝑁, 𝑄, 1 / 𝛼 𝑉 , 1 / 𝜀 Tmp , log 𝑇 ) samples, we can make sure with high probability,

Thus, after one gradient step with 𝜂 = 1, we have

Recall from Lemma H.1 that 𝛼 𝑉 ∈ [ Θ ( 𝜃 /( 𝑁𝐾 𝑃 )) , Θ ( 𝜃 )] . Hence, we can choose the threshold 𝜆 0 to be Θ ( 𝜃 /( 𝑁𝐾 𝑃 )) and 𝜀 Tmp ≤ 1 /( 𝑁𝐾 𝑃 ) so that

Now, set

Note that this implies that after the first step, we have

and

Theorem H.3 (Main theorem for transfer learning) . Let 𝜀 > 0 be our target accuracy. Consider the same setting of Lemma H.1 and Lemma H.2. Choose 𝜃 = √︁ 𝑂 ( 1 / log ( 1 / 𝜀 )) and 𝜀 Tmp = 𝑂 GLYPH<16> 1 log ( 1 / 𝜀 ) GLYPH<17> . Then, after one step of update on 𝑨 as in Lemma H.2, 𝑨 and 𝑽 satisfies the conditions of Lemma E.10, and therefore we can learn ( 𝑷 , 𝑸 ) to 𝜀 -accuracy using poly ( 𝑁, 𝑄, 1 / 𝜀, 1 / 𝛿 ) samples with probability at least 1 -𝛿 within poly ( 𝑁, 𝑄, 1 / 𝜀, 1 / 𝛿 ) 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 𝑇 = 5000 , 𝑄 = 2 , 𝑁 = 3. The architecture is

and the data model is the SCB (1) data-generating model. The batch size is 𝐵 = 64 and the regularization hyperparameter is 𝜆 = 1e-5. The total time is T = 1000 iterations where stage 1 takes 𝜏 ∈ [ 0 , 400 ] with learning rate 𝜂 1 = 0 . 01. After 𝜏 > 400, we use 𝜂 2 = 0 . 005 for further improvement (stages 2 and 3).

Hyperparameter selection. Due to the limitation of computational resources, we do experiments with 𝑁 ∈ [ 3 , 20 ] for real-world batched gradient experiments, and 𝑁 ∈ { 100 , 500 } , 𝑇 = 100000 experiments by using Gaussian noise SGD simulations based on the dynamics of Lemma C.5 and C.6. As 𝑇 needs to scale with 𝑁 polynomially, it would be beyond our computation capability to

experiment with larger 𝑁 . As for other hyperparameters, 𝜆 is chosen based on our theoretical results (Theorem 3.1 and G.1): 𝜆 ∼ Θ ( 𝜖𝐾 𝑃 / 𝑄 2 𝑁 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. 𝜂 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

✶ and we do further gradient descent to recover 𝑽 . In this case, we will directly use the linear regression solution on population loss for 𝑽 .

Here we report in addition: (1) original signal projection on the population process trajectory 𝛼 𝐴 , 𝛼 𝑉 and the distance to the trajectory 𝚫 𝑨 , 𝚫 𝑽 . (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

200

400

600

800

1000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

Signal Projection

α

A

, ours

α

V

, ours

α

A

, SGD

α

V

, SGD

0

200

400

600

800

1000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

Distance to projection

∆

A

, ours

∆

V

, ours

∆

A

, SGD

∆

V

, SGD

iteration iteration

0

200

400

600

800

1000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Similarity (Proximal)

〈

¯

P,

¯

V

〉

, original

〈

¯

A,

¯

Q

〉

, original

〈

¯

P,

¯

V

〉

, normalized

〈

¯

A,

¯

Q

〉

, normalized

0

200

400

600

800

1000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Similarity (SGD)

〈

¯

P,

¯

V

〉

, original

〈

¯

A,

¯

Q

〉

, original

〈

¯

P,

¯

V

〉

, normalized

〈

¯

A,

¯

Q

〉

, normalized

iteration iteration

Wetried different orders of state number 𝑁 and show that ℓ 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 𝑁 ≤ 20 , 𝑇 ≤ 5000, and do SGD simulation by combining our population gradient + Gaussian noise (to mimic the batch gradient noise) for 𝑁 ≤ 500 , 𝑇 ≤ 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 linear attention, ℓ 1 regularized dynamics eventually will make 𝐴 and 𝑄 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, 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 a little after it converges. Therefore, we must stop early and normalize before the distortion happens.

0

250

500

750

1000

1250

1500

iteration

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

Distance to ground truth

‖

P

-

V

‖

µ

, ours

‖

P

-

V

‖

µ

, SGD

‖

P

-

V

‖

µ

, SGD small

Final error, ours

0

250

500

750

1000

1250

1500

iteration

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

Distance to ground truth

‖

Q

-

A

‖

µ

, ours

‖

Q

-

A

‖

µ

, SGD

‖

Q

-

A

‖

µ

, SGD small

Final error, ours

0

250

500

750

1000

1250

1500

iteration

0

.

300

0

.

305

0

.

310

0

.

315

0

.

320

0

.

325

Test loss

ours

SGD

SGD small

Final error, ours

0

250

500

750

1000

1250

1500

iteration

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

Distance to ground truth

‖

P

-

V

‖

µ

, ours

‖

P

-

V

‖

µ

, SGD

‖

P

-

V

‖

µ

, SGD small

Final error, ours

0

250

500

750

1000

1250

1500

iteration

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Distance to ground truth

‖

Q

-

A

‖

µ

, ours

‖

Q

-

A

‖

µ

, SGD

‖

Q

-

A

‖

µ

, SGD small

Final error, ours

0

250

500

750

1000

1250

1500

iteration

0

.

435

0

.

440

0

.

445

Test loss

ours

SGD

SGD small

Final error, ours

0

500

1000

1500

2000

2500

3000

iteration

0

.

00

0

.

05

0

.

10

0

.

15

Distance to ground truth

‖

P

-

V

‖

µ

, ours

‖

P

-

V

‖

µ

, SGD

‖

P

-

V

‖

µ

, SGD small

Final error, ours

0

500

1000

1500

2000

2500

3000

iteration

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

Distance to ground truth

‖

Q

-

A

‖

µ

, ours

‖

Q

-

A

‖

µ

, SGD

‖

Q

-

A

‖

µ

, SGD small

Final error, ours

0

500

1000

1500

2000

iteration

0

.

468

0

.

470

0

.

472

0

.

474

Test loss

ours

SGD

SGD small

Final error, ours

0

2000

4000

6000

8000

10000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

with

/lscript

1

reg (batch size = 1.0E+06)

α

V

α

A

∆

2

A

/K

Q

∆

2

A

/K

P

Stage 2 starts

0

2000

4000

6000

8000

10000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

without

/lscript

1

reg (batch size = 1.0E+06)

α

V

α

A

∆

2

A

/K

Q

∆

2

A

/K

P

Stage 2 starts

0

2000

4000

6000

8000

10000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

without

/lscript

1

reg (batch size = 1.0E+07)

α

V

α

A

∆

2

A

/K

Q

∆

2

A

/K

P

Stage 2 starts

0

2000

4000

6000

8000

10000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

without

/lscript

1

reg (batch size = 1.0E+08)

α

V

α

A

∆

2

A

/K

Q

∆

2

A

/K

P

Stage 2 starts

0

2000

4000

6000

8000

10000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

1

.

4

with

/lscript

1

reg (batch size = 1.0E+07)

α

V

α

A

∆

2

A

/K

Q

∆

2

A

/K

P

Stage 2 starts

0

2000

4000

6000

8000

10000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

without

/lscript

1

reg (batch size = 1.0E+07)

α

V

α

A

∆

2

A

/K

Q

∆

2

A

/K

P

Stage 2 starts

0

2000

4000

6000

8000

10000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

without

/lscript

1

reg (batch size = 5.0E+07)

α

V

α

A

∆

2

A

/K

Q

∆

2

A

/K

P

Stage 2 starts

0

2000

4000

6000

8000

10000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

without

/lscript

1

reg (batch size = 1.0E+08)

α

V

α

A

∆

2

A

/K

Q

∆

2

A

/K

P

Stage 2 starts

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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: It 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 in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren'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 in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to 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 a complete (and correct) proof?

Answer: [Yes]

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix 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 or 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.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is 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 might suffice, or if the contribution is a specific model and empirical evaluation, it may 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all 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 a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) Werecognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https://nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

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 experimental setting should be presented in the core of the paper to a 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 performance/accuracy results.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, 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 are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

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 state that they have a 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 or 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 is included in Section 5 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 in 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://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This work is a pure theory paper without potential harmful effect.

Guidelines:

The answer NA means that the authors have not 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., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: 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 societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper is a theoretical work without such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: The paper does not use existing assets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should state which version of the asset is used and, if possible, include a URL.

The authors should cite the original paper that produced the code package or dataset.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: 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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset 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 as 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 not involve crowdsourcing 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 to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: 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.

STRENGTHS:
1. The paper is well-written and organized.
2. 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?
2. The loss curve does not look like converging in Figure 1. It would be important to show the loss curve with more training rounds.

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. 
2. 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,
1. 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.
2. The paper lacks experiments under a more realistic setting with the standard transformer structure and training method.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\n- See weakness.\n- What's $I_T$ in equation 2?\n\nLIMITATIONS:\nThey discuss the limitations in appendix A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nNaively, 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.\n\nQUESTIONS:\nNA\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-21"}