The Closeness of In-Context Learning and Weight Shifting for Softmax Regression
Authors: Shuai Li
Zhao Song
Yu Xia
Tong Yu
Tianyi Zhou
Venue: NeurIPS 2024 poster (2024)
Keywords: in-context learning
Primary area: optimization_for_deep_networks

Abstract: Large language models (LLMs) are known for their exceptional performance in natural language processing, making them highly effective in many human life-related tasks. The attention mechanism in the Transformer architecture is a critical component of LLMs, as it allows the model to selectively focus on specific input parts. The softmax unit, which is a key part of the attention mechanism, normalizes the attention scores. Hence, the performance of LLMs in various NLP tasks depends significantly on the crucial role played by the attention mechanism with the softmax unit.

In-context learning is one of the celebrated abilities of recent LLMs. 
Without further parameter updates, Transformers can learn to predict based on few in-context examples. 
However, the reason why Transformers becomes in-context learners is not well understood.
Recently, in-context learning has been studied from a mathematical perspective with simplified linear self-attention without softmax unit. 
Based on a linear regression formulation $\min_x\| Ax  - b \|_2$, existing works show linear Transformers' capability of learning linear functions in context. The capability of Transformers with softmax unit approaching full Transformers, however, remains unexplored.

In this work, we study the in-context learning based on a softmax regression formulation $\min_{x} \| \langle \exp(Ax), {\bf 1}_n \rangle^{-1} \exp(Ax) - b \|_2$. We show the upper bounds of the data transformations induced by a single self-attention layer with softmax unit and by gradient-descent on a $\ell_2$ regression loss for softmax prediction function.
Our theoretical results imply that when training self-attention-only Transformers for fundamental regression tasks, the models learned by gradient-descent and Transformers show great similarity.

The Closeness of In-Context Learning and Weight Shifting for Softmax Regression
Shuai Li
Zhao Song
Shanghai Jiao Tong University shuai1i8@sjtu. edu.
Simons Institute for the Theory of Computing, UC Berkeley magic. linuxkde@gmail...........mmmm..m com
Yu Xia
Tong Yu Adobe Research tyuSadobe com
Tianyi Zhou
University of California, San Diego yux0780ucsd.em edu
University of Southern California tzhou0296usc.edd
Abstract
Large language models (LLMs) are known for their exceptional performance in natural language processing, making them highly effective many human life- related tasks. The attention mechanism in the Transformer architecture is critical component of LLMs, as it allows the model to selectively focus on specific input parts. The softmax unit, which is a key part of the attention mechanism, normalizes the attention scores. Hence, the performance of LLMs various NLP tasks depends significantly on the crucial role played by the attention mechanism with the softmax unit.
In-context learning is one of the celebrated abilities of recent LLMs. Without fur- ther parameter updates, Transformers can learn to predict based on few io-context examples. However, the reason why Transformers becomes in-context learners is not well understood. Recently, in-context learning has been studied from a mathematical perspective with simplified linear self-attention without softmax unit. Based on a linear regression formulation ming |r - bll2, existing works show linear Transformers' capability of learning linear functions in context. The capa- bility of Transformers with softmax unit approaching full Transformers, however, remains unexplored.
In this work, we study the in-context learning based on a softmax regression formulation min, ((exp(Ax)) 1,,) -x))==exp(AA (xr(Ax)== bll2- We show the upper bounds of the data transformations induced by a single self-attention layer with softmax unit and by gradient-descen on a la regression loss for softmax prediction function. Our theoretical results imply that when training self-attention-only Transformers for fundamental regression tasks, the models learned by gradient-descent and Transformers show great similarity.
Introduction
In recent years, there has been a significant increase in research and development the field of effective way tackle
Artificial Intelligence, with large language models (LLMs) emerging complex tasks. Transformers have achieved state-of-the-art results various NLP tasks, such as machine translation [PCR19, GHG 20] and text generation [LSX 22). As result, they have become the preferred architecture for NL.P, where BERT [DCLT18), GPT-3 [BMR+ 20], PaLM [CND+22] proposed. They have demonstrated remarkable learning and reasoning capabilities and have proven to be more efficient than traditional models when processing natural language.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Additionally, LLMs can be fine-tuned for multiple purposes without requiring a new build from scratch, making them versatile tool for Al applications. Moreover, recent studies on the in-context learning abilities of L.LSs have demonstrated that even without further fine-tuning, LLMs can generalize to new tasks with only a few demonstration examples in the prompt. To understand how LLMs become in-context learners, recent works have studied the problem and provided mathematical explanations from the Transformer architecture perspective, showing a simplified linear self-attention layer of Transformer can learn linear functions similarly step of gradient descent [ONR 22, ASA+22, GTLV22, 24J. While such linear approximation of full Transformers is overly simplistie, studies more complex Transformer architecture needed to further explain the in-context learning phenomenon.
Transformers have specific type of sequence- coaennceettt-seeeee -equenne neural network architecture. They utilize the attention mechanism [VSP+17, RNS+ 18, DCLT18, BMR+20) that allows them to capture long- range dependencies and context from input data effectively. The core of the attention mechanism the attention matrix which comprised of rows and columns, corresponding to individual words or "tokens". The attention matrix represents the relationships within the given text. It measures the importance of each token in a sequence as it relates to the desired output. During the training process, the attention matrix is learned and optimized to improve the accuracy of the model's predictions. Through the attention mechanism, each input token is evaluated based on its relevance to the desired output by assigning token score. This score is determined by similarity function that compares the current output state with input states.
Theoretically, the attention matrix is comprised of the query matrix Q E IRadd the key matrix K E Ra and the value matrix V E Igaxd Following [ZHDK23, AS23, BSZ24, AS24b, A$24c, AS24a], the computation of the normalized attention function is defined as D-1 exp(QK VV. Following the transformer literature, we apply exp to matrix entry-wise way. Here E gnx= diagonal matrix that defined as D = diag(exp((()) (exp(QK' ).,, Intuitively, D denotes the softmax normalization matrix. more general computation formulation can be written as
D
X
diagonal matrix
nXn
axd dxd
where
== =
In the above setting, treat Q K.V R weights and X is the input sentence data that has length n and each word embedding size is d. In the remaining of the part, we will switch to notation A and use A to denote sentence. Mathematically, the attention computation problem can be formulated as regression problem in the following sense Definition 1.1. We consider the following problem
||D_ exp(AXAT) BlI X XER
min
where A E Raxd can be treated length-n document and each word has length-d embedding size. Here D = diag(AX 1 For any given A E Raxd and B E Rax", the goal is to find some weight X to optimize the above objective function.
In contrast to the formulation in [ZHDK23, AS23, BSZ24], the parameter X in Definition 1.1 equivalent the QKT R in the generalized version of [ZHDK23, AS23, BSZ24] (e.g.
replacing Q E Baxd by XQ where X E Raxd and Q E Rdxd, Similarly for K and V. In such scenario, X can be viewed as matrix representation of length-n sentence.).
number of work JASA+ + 22, GTLV22, ONR + 22] study the in-context learning from mathematical perspective in much simplified setting than Definition 1.1, which is linear regression formulation as in Definition 1.2. They show linear Transformer without softmax unit in its attention layer can mimic the ability of gradient descent in learning linear functions in context. While the softmax unit plays an important role attention computations of full Transformers, their simplification of the softmax unit leaves gap in explaining LLMs' in-context learning abilities.
Definition 1.2. Given matrix A E R and b E R the goal is to solve
min Ar bll2

---

Several theoretical transformer work have studied either exponential regression [GMS23, LS223] or softmax regression problem [DLS23, LLSS24a]. In this work, to take a step forward to understand the softmax unit in the attention scheme in LLMs. We consider the following softmax regression and study the in-context learning phenomena and its closeness to gradient descent from the data transformation perspective.
Definition 1.3 (Softmax Regression), Given A E Baxd and a vector b E R,, the goal is to solve
min
|| (exp(Ax), 1.)) exp(Ar) bb/2
aERd
We remark that the Definition 1.3 of Softmax Regression a formulation in between Definition 1.2 and Definition 1.1.
We state our major result as follows:
Theorem 1.4 (Bounded shift for in-context learning. informal version of the combination of The-
orem 4.2 and Theorem 4.3). If the following conditions hold: Let A E Rxxd Let b E IR" ||||| & R. Let || R. [|A(x++1 r1)ll < 0.01. I1(A++1 At)arl 0.01. Let R 2 Let
M == n exp(10R"), We consider the softmax regression (Definition 1.3) problem
min ((exp((x)) 1n) exp(Ax) bl[2-
Part 1. If we move the to 2(+1, then we 're solving new softmax regression with ming (eexp(Ax)) 1,) exp((xt (xx) bll2 where 16 b1|2 S M || xell2
1,)
Part 2. If we move the A, to Act then we'ne solving new softmax regression with min, (eexp(Ax) 1,, 1,, exp(Ax) - bll2 where bll2 S M ||Ar+1 Arll
Recall that A Rexd denotes length-n document and each word has the length-d embedding size and x denotes the simplified version of QK One-step gradient descent can be treated as an update to the model's weight x. Thus, part of our result (Theorem 1.4) implies that the data transformation of induced by gradient-descent on the &2 regression loss is bounded by M || - . seel2-
According to [ONR 22], to do in-context learning, self-attention layer update can be treated as an update to the tokenized document A. For detailed derivation, please refer to [ONR* + 22]. Thus, part of our result (Theorem 1.4) implics that the data transformation of b induced by single self-attention layer is bounded by M ||| Aell-
We remark that the data transformation of b induced by 1) single self-attention layer and by 2) crudient-descent on the & regression loss are both bounded. The bounded transformation of implies that when training self-attention-only Transformers for fundamental regression tasks, the models learned by gradient-descent and Transformers show great similarity.
Roadmap. In Section 2, we give some preliminaries. In Section 3, we compute the gradient of the loss function with softmax function with respect to x. Those functions include (xx) -1 (#x and (z). In Section 4, we give our formal theoretical results, validated by numerical experiments presented in Section In Section 6, we conclude our paper.
Preliminary
In Section 2.1, we introduce the notations used in this paper. In Section 2.2, we give some facts about the basic algebra. In Section 2.3, we propose the lower bound on (exp(Ar), 1,))
2.1 Notations
,,). For positive integer we use [rx] to denote m)), for any positive integer n. We use E[[] denote expectation. We use Pr[.) to denote probability. We use 1, denote the vector where all entries are one. We use On to denote the vector where all entries are zero. The identity matrix of size n X n is represented by I, for positive integer n. The symbol refers to real numbers and Rzo represents non-negative real numbers. For any vector E R. exp(x) E R denotes vector where the i-th entry is exp(x;) and ||æ|| represents its l norm, that is, ||l|| = EE xf) We use

---

| to denote maXig[n] | For any vector z E IR" and vector R. we use (æ, y) to denote the inner product of vector x and y. The notation B, used indicate the i-th row of matrix B. If and are two column vectors in Re, then b denotes a column vector where (a o b); agb. For square and full rank matrix B, we use B-1 to denote the true inverse B.
2.2
Basic Algebras
Fact 2.1. For vectors z, y E R", have
|ar o y|ا2 ||0| |00 ||u||2
| S || < /
exp(x>)l>> S cxp([1<<122)
For any || - ull S 0.01, we have || exp(x) exp(v)ll2 S ||ex(())]l2 21|x vl Fact 2.2. For matrices X, we have
|| |||
|X 2 ||||| ||X YIl
|X +Y| v |X||| + ||
|X Y|| E ||||| ||
1/X < a Y, then |||| S a |||||
2.3 Lower bound on 8
Lemma 2.3. If the following conditions holds
||||| S R
| 12 S R
Let 3 be lower bound on (exp(Ax)) 1,,)
Then we have
B 2 exp(-R²)
Proof. We have
(exp(Ax), 1,) Σ exp((Az))
min exp (AA)))))
1E|m]
min exp( -1(A)))))))))))
iE|n
exp(
max (Ax);))
iE(n)
exp( ||
exp(
|
exp( R²))
the 1st step follows from simple algebra, the 2nd step comes from simple algebra, the 3rd step follows from the fact that exp(x) 2 exp(-Iz[), the 4th step follows from the fact that exp(-r) is monotonically decreasing, the 5th step comes from definition of fox norm, the 6th step follows from Fact 2. the 7th step follows from the assumption A and

---

Softmax Function with Respect to z
In Section 3.1, we give the definitions used the computation. In Section 3.2, we compute the gradient of the loss function with softmax function with respect to æ. Those functions includes a(x) -1 axx) and f()).
3.1 Definitions
We define function softmax follows
Definition 3.1 (Function f. Definition 5.1 in [DLS23]). Given a matrix A E [Raxd Let 1n denote length-n vector that all entries are ones. We define prediction function f R + R2 as follows
f(z) == (exp(Ar), 1x)=+ exp(Ax).
Definition 3.2 LLsss function Texps Definition 5.3 [DISS23]) Given a matrix Raxd and vector b R We define loss function Lexp R > as follows
[x) 0.5 (exp((xx) 1,) exp(Ax) - bll2-
For convenient, we define two helpful notations and Definition 3.3 (Normalized coefficients, Definition 5.4 in [DLS23]). We define a Rd - R as follows
(x) (exp(Ar)) In))
Then, we can rewrite (() (see Definition 3.1) and Lexp (2) (see Definition 3.2) as follows
f(x) = axx) cxp(Ax)
Txxp(() = 0.5 ||a(x) exp(Ar) - -112-
Lexp (x) = 0.5 [[(()) 6113.
Definition 3.4 (Definition 5.5 in [DLS223). We define function Rd E R follows c(x) = f(x) b.
Then we can rewrite (x) (see Definition 3.2) as follows
[Lexp(x)) 0.5 ||
3.2 Gradient Computations
We state lemma from previous work,
Lemma 3.5 (Gradient, Lemma 5.6 in [DLS23]). If the following conditions hold
Given matrix A E Rgad and vector b E R"
Let a(x)) be defined in Definition 3.3.
Let f(x) be defined in Definition 3.1.
Let c(x) be defined in Definition 3.4.
Let (x) be defined in Definition 3.2.
For each i E [d], we have
Purt
d exxp(Ax) exp(Ax) o A dæ;
5

---

Part 2.
d(exp(Ax) In) (exp(Ar)) A.,))
dx;
Part
da(xz) -a(x)- -a(x)=-(((t))AAA
dr;
Part 4.
df(x) de(x) = (f(x), f(x) + f(x) o A.S dæ;
dæ;
Part 5.
dLexp(z)
A.,
((z) (clx), diag (fig(
dr;
1xn
nx1
scalar
nxa
XX
4
Main Results
In Section 4.1, we show the lipschitz bound function f. In Section 4.2, we show our upper bound result of 5g with respect v. In Section 4.3, show upper bound result of ds with respect A.
4.1 Lipschitz Bound
To bound the shift of b, we first show the Lipschitz property for the basic functions:
(exp(Ax) - exp(Ay)le s 2/nnRexpl (?²) ||x - y/12
a(()) a < (xp Ar) - exp(Ay) 1| vn
la(s)++ 3-2 (a(x)))())))))))) av)
We can show that
Lemma 4.1. If the following conditions hold
Let E (0,1).
Let 86,1 TR. be defined as Definition B.3.
Let 8,.2 IR." be defined as Definition B.3.
Let 8g 86,1 86,2-
Let
We have
Part 1.
16, I 2 cxp((R))) || - 2x112
Part 2.
18,2 23- (RR -1 2el12
Part 3.
|| [((++1) f(x)) ||2 < 2 Rexp(2112) || xxl12

---

Proof. Proof of Part 1. We have
16, < -
[a(0441) o(()) exp (t2)
S 8- (0(2)41) a(z)) vni vippp
$ 8 2 vn exp((A))1) exp(Ax:) ll2 vn cxpp(R2)
- 8 vn 2 (t |++1 =e|l2 vn exp(1e))
23 2n 1.5 R Rexp(2R*)) || - 2xll2
where the first step follows from definition, the second step follows from assumption A and the third step follows Lemma B.7, the forth step follows from Lemma B.6, the fifth step follows from Lemma B.5.
Proof of Part
We have
6 la S (a(x exp((Ax41) exp(Axx))
S 8 exp( (xx++) - exp( (xx;))
< 2//Rexx((2R)) ||
2z112
where the first step follows from definition, the 2nd step comes from Lemma B.5. Proof of Part 3.
We have
|6ll2
|106,1 + 20.2112
| 112 + ||00,2112
2 2 1.5 exp (2R") || rella + 2. 'n R eexp(2RR)
|| zell2
2 n (2x)(2Re)
|| - zxll2 2.3 2n' 1.5 Rexp(21l2)
||
xell2
< 2 ReesSeepp Rexp(2R") (2 1aai--rell xella
where the 1st step follows from the definition of dg, the 2nd step follows from triangle inequality, the 3rd step follows from the results in Part 1 and Part 2, the 4th step follows from the fact that 2 1 and B 2 1, the 5th step follows from simple algebra.
Similarly, we can show the Lipschitz property of function f with respect to A as the following
||Ar+1 Arll2
Due to space limitation, we defer formal lemma and proof to D.2.
4.2 Shifting Weight Parameter
Theorem 4.2 (Bounded shift for shifting the weight parameter, formal of Theorem 1.4), If the following conditions hold
Let A Raxd
|All
|A(x141 - 12)0 < 0.01
Let
Let : 1.5 exp())

---

We consider the softmax regression problem
min (c(p))(Ax), 1.,) exp(Az) bll2
If we move the 2e to Ft+1 then we' re solving a new softmax regression problem with
min ([(ex((Ax)) 1n)) exp(Az) 1x)_A))((Ax) bll2
where
1 b112 S M. 11,1111111111 - xxll2
Proof. We have
| b/12 1 43 1.5
" Rexp(21(2) || xell2
I 4n 1.5 R2xp(2R exp(2R%) || xelle
< n 1.5 (4R) exp (4 || xxel2
x n 15 cxp( exp(4R²) || =xel2
n 1.5 exp(10 || xx.12
2 || - xell2
where the 1st step follows from Lemma 4.1, the 2nd step comes from Lemma 2.3, the 3rd step comes from simple algebra, the 4th step follows from simple algebra, the 5th step follows from simple algebra and the 6th step follows from the definition of M.
4.3 Shifting Sentence Data A
Theorem 4.3 (Bounded shift for in-context learning, formal of Theorem 1.4). If the following conditions hold
Let Raxd
R
[1(A+1 - lo) < 0.01
Let R24
Let = 1.5 exp(1)
We consider the softmax regression problem
min [exep(Ax)) 1.)) exp(Ax) bll2
If we move the A, A441 then we re solving new softmax regression problem with min (((ex((Ar)) 1,) exp(Ax)) bll2
where
|b bll2 S M || Arll-
Proof. We have
[ b112 _ 2n 1.5 Rexp(2R2) || Aell
< 4n 5..5pp R exeepp22 (2R²) exp(2 ||Att1 A.ll
< n 1.5 (4R) expl AR²) || - Aell
x n 1.5 exp( exp ||441 - Aell
< n 1.5 exp(10R²)) | A
Asll
< M. M. || A.ll
where the 1st step follows from Lemma D.5, the 2nd step follows from Lemma 2.3, the 3rd step follows from simple algebra, the 4th step comes from simple algebra, the 5th step comes from simple algebra and the 6th step follows from the definition of M.

---

Numerical Experiments
In this section, we present our numerical experiments to validate our theoretical results that when training self-attention-onll Transformers for softmax regression tasks, the models learned by gradient- descent and Transformers show great similarity.
5.1 Experiments Setup
According Definition 1.3, we construct the synthetic softmax regression tasks consists of randomly sampled length-n documents A E Rgxxd where each word has the d-dimensional embedding and targets b R". Each document is generated from unique random seed. In our experiments, choose a set of different document length n and set of different embedding size d.
Following [ONR 22], compare the following two models in our experiment
* trained single self-attention (SA) layer with softmax unit approximating full Transformers. softmax regression model trained with one-step gradient descent (GD).
The training objective for both models is defined as in Definition 1.3. the single self-attention layer with softmax unit, choose the learning rate 7/ 0.005. For the softmax regression model, we determine the optimal learning rate 7JDD by minimizing the l regression loss over a training set of 10 tasks through line search.
To compare the trained single self-attention layer with a softmax unit and the softmax regression model trained with one-step gradient descent, we sample 103 tasks and record the losses of two models. In addition, we follow [ONR* #22] to record
Pred Diff: the predictions difference measured with the & norm:
UsA (A) (GD(2>l22
where ğSA (A) corresponds to the b in Theorem 4.2, and (CD(2) corresponds the 6 in Theorem 4.3.
Model Cos: the cosine similarity between the sensitivities of two models:
(T) (A))
CosSim
dr 0A
Model Diff: the model sensitivity difference measured with the & norm:
(x)
||2
Ör aA
All experiments run single NVIDIA RTX2080TI GPU with 10 independent repetitions.
5.2 Different Document Lengths
For synthetic softmax regression tasks of document length E (200, 1000) and word embedding size d = 20, the comparison results between a trained single self-attention layer and one-step gradient descent are shown in Figure and Figure 2. Due to space limitation, we present more results with different document length n {25, 50, 100, 200, 400, 1000} in Appendix E.
We compare two models' losses over training steps of Transformers in Figure la and Figure 2a. In Figure lb and Figure 2b, we show the differences and similarities of two models over the training steps. From the results, we find identical performances of the two models measured in losses. We also observe considerable alignment of the two models across tasks of different document lengths, indicated by decreasing prediction and model difference and increasing cosine similarity between models. Besides, comparing results with different nı, we observe that with larger document length, which is common in practical NLP tasks, more training steps are required for Transformers to exhibit such similarities. This shows the crucial role of pretraining stage of Transformers for their in-context learning ability.

---

Cred Doceet Thaine
35
Fooes Perstorer
KeeeOn
0.3
300
28.
800
15
12
0.76
13-
100
0.4
276
455
26
310
3850
206
430
000
500
1000
0 380
430 030
687
5000
00-
200
420
000
560
3100
(a) Losses over training (b) Difference and similarity
(a) Losses over training (b) Difference and similarity
sleps of Transformer
training steps
steps of Transformer
training steps
Figure 1: Comparison on softmax regression tasks Figure 2: Comparison on softmax regression tasks of document length n 200.
of document length 1000.
5.3 Different Word Embedding Sizes
We also compare trained single self-attention layer and one-step gradient descent on synthetic softmax regression tasks of different word embedding sizes and document length n 200. Similarly, we measure two models' losses and similarities over training steps on each set of tasks. Due to space limitation, we follow [ONR* 22] to show in Figure 3 the loss comparisons at the end of training over different embedding size d E (5, 10, 20, 35, 50). The complete loss curves and measurements of model difference and similarity presented in Appendix E.
From the results, we again observe similar performances and close aligmment of the two models with different word embed- ding sizes.
0.8-
Gradient Descent
Trained Transfomer
0.7
0.6
To summarize, our numerical results validate our theoretical results in Section showing that when training self-attention- only Transformers for softmax regression tasks, the models learned by gradient-descent and Transformers show great sim- ilarity. Note that due to the non- non-linearity of softmax regression, it is not expected for models match exactly as implied in our theoretical results in Section 4, which is also observed in our numerical findings.
0.5
20
30 Input Dim
40
50
Figure Loss comparisons with different word embedding sizes d.
6 Conclusion
The attention mechanism that incorporates the softmax unit crucial aspect of Large Language Models (LLMs) and significantly contributes to their extraordinary performance in various Natural Language Processing (NLP) tasks. The ability to learn in-context is highly valued in recent LLMs. and comprehending this concept is vital when querying LLMs. In this study, taking a step further from prior works' studies on linear Transformer's ability of learning linear functions, we examined the in-context learning process from a softmax regression perspective of Transformer's attention mechanism. We established the bound on the data transformations brought about by a single self- attention layer with softmax unit and gradient descent on an L2 regression loss. Our findings suggest that the update acquired through gradient descent and in-context learning are highly similar when training seff-attention-nnly Transformers for softmax regression tasks, which is also validated through our preliminary experimental results. These results offer insights into the theoretical underpinnings of in in-context learning in Transformers and can aid in improving the understanding and performance of LLMs in various NLP tasks.
Acknowledgments
The authors would like to thank Jerry Yao-Chieh Hu, Zhenmei Shi, Lichen Zhang and Yufa Zhou for helping preparing for camera-ready version of this paper. For more information related to the paper and adjacent topics, see /tattss://wwwt wttpeuuube.cc////aaaaaae homm///aaasoog203111 and https:// space. bilibili. /00.///5//555777766606666
10

---

References
[Ano24a] Anonymous. Fundamental limits of prompt tuning transformers: Universality, capacity and efficiency. In Submitted to The Thirteenth International Conference on Learning Representations, 2024. under review.
[Ano24b] Anonymous. On statistical rates of conditional diffusion transformer: Approximation and estimation. In Submitted to The Thirteenth International Conference on Learning Representations, 2024. under review.
[AS23] Josh Alman and Zhao Song. Fast attention requires bounded entries. In NeurlPS, 2023.
[AS24a] Josh Alman and Zhao Song. Fast rope attention: Combining the polynomial method and fast fourier transform. In manuscript, 2024.
[AS24b] Josh Alman and Zhao Song. The fine-grained complexity gradient computation for training large language models. In NeurIPS, arXiv preprint arXiv:2402.04497, 2024.
[AS24c] Josh Alman and Zhao Song. How to capture higher-order correlations? generalizing matrix softmax attention to kronecker computation. In ICLR, 2024.
[ASA + 22] 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 arXivv2211.15661, 2022.
[BAG20] Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the Ability and Limitations of Transformers to Recognize Formal Languages. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 7096-7116. Online, November 2020. Association for Computational Linguistics.
[Be122] Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances. Compu- tational Linguistics, 48(1)2207-219 March 2022.
[BRR + 20] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances neural information processing systems, 33:1 13:117711901, 2020.
[BP66] Leonard Baum and Ted Petrie. Statistical inference for probabilistic functions of finite state markov chains. The annals of mathematical statisties, 376):1554-1563, 1966.
[BPG20] Satwik Bhattamishra, Arkil Patel, and Navin Goyal. On the computational power of transformers and its implications in sequence modeling. In Proceedings of the 24th Conference on Computational Natural Language Learning, pages 455-475, Online, November 2020. Association for Computational Linguistics.
[BSZ24] Jan den Brand, Zhao Song, and Tianyi Zhou. Algorithm and hardness for dy- namic attention maintenance in large language models. In ICML. arXiv preprint arXiv:2304.0224 2024.
[CDL+22] Beidi Chen, Tri Dao, Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. In International Conference on Learning Representations (ICLR), 2022.
[CDW+21] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scat- terbrain: Unifying sparse and low-rank attention. Advances Neural Information Processing Systems (NeurIPS), 34:17413-17424 2021.
[CGRS19] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 2019.

---

[CL Bo Chen, Xiaoyu Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. Bypassing the exponential dependency: Looped transformers efficiently learn in-context by multi-step gradient descent. arXiv preprint arXiv:2410. 1266, 2024.
[CLS* Bo Chen, Yingyu Liang, Zhizhou Sha, Zhenmei Shi, and Zhao Song. Hsr-enhanced sparse attention acceleration. arXiv preprini arXiv:2410. 10165, 2024.
[CND+22] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebas- tian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv2224.02311,2000 2022.
[DCL*21] Tri Dao, Beidi Chen, Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. arXiv preprint 1mXiv2212.0000 00.00299 2021.
[DCLT18] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre- training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04880 1810:0800, 2018.
[DGS23] Yichuan Deng, Yeqi Gao, and Zhao Song. Solving tensor low cycle rank approximation. arXiv preprini 23Xx::23004065594, 2023.
[DGV+18] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint 1807.038 2018.
[DKOD20] Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrf- efficient attention using asymmetric clustering. Advances in Neural Information Processing Systems (NeurIPS), 33:6476-6449 2020.
[DLS23] Yichuan Deng, Zhihang Li, and Zhao Song. Attention scheme inspired softmax regression. arXiv preprint arXiv:2304. 2001::2334.0411,,,, 2023.
[DMS23) Yichuan Deng. Sridhar Mahadevan, and Zhao Song. Randomized and deterministic attention sparsification algorithms for over-parameterized feature dimension. arxiv preprint: arxiv 03426, 2023.
[EGK22]] Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. arXiv preprint arXiv:2110. 211.11000000 2021.
[EGZ20] Javid Ebrahimi, Dhruv Gelda, and Wei Zhang. How can self-attention networks recognize Dyck-n languages? In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 4301-4306, Online, November 2020. Association for Computational Linguistics.
[GHG+20] Peng Gao, Chiori Hori, Shijie Geng, Takaaki Hori, and Jonathan Le Roux. Multi-pass transformer for machine translation. arXiv preprint arXiv:2009.1333 2020.
[GMS23] Yeqi Gao, Sridhar Mahadevan, and Zhao Song. An over-parameterizze exponential regression. arXiv preprint arXiv:2303.16500 2023.
[GSYZ24] Yeqi Gao, Zhao Song, Xin Yang, and Yufa Zhou. Differentially private attention computation. In Neurips Safe Generative Al Workshop 2024, 2024.
[GTLV22] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can trans- formers learn in-context? a case study of simple function classes. arXiv preprint arXiv:2208.0 2022.
[HCL+24] Jerry Yao-Chieh Hu, Pei-Hsuan Chang, Haozheng Luo, Hong- Yu Chen, Weijian Li, Wei-Po Wang, and Han Liu. Outlier-efficient hopfield layers for large transformer- based models. Forty-first International Conference on Machine Learning (ICMI., 2024.
12

---

[HCW*24] Jerry Yao-Chieh Hu, Bo- Yu Chen, Dennis Wu, Feng Ruan, and Han Liu. Nonparametric modern hopfield models. arXiv preprint 2at::2004033000. 2024.
[HL19] John Hewitt and Percy Liang. Designing and interpreting probes with control tasks. In Proceedings of the 2019 Conference Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNIPP) pages 2733-2743, Hong Kong, China, November 2019. Association for Computational Linguisties.
(HLSL24) Jerry Yao-Chieh Hu, Thomas Lin, Zhao Song, and Han Liu. On computational limits of modern hopfield models: A fine-grained complexity analysis. In Forty-first International Conference on Machine Learning (ICMI), 2024.
[HWL244] Jerry Yao-Chieh Hu, Dennis Wu, and Han Liu. Provably optimal memory capacity for modern hopfield models: Transformer-compatbbt dense associative memories as spherical codes. In Thirty-eighth Conference Neural Information Processing Systems (NeurIPS), 2024.
[HWL+24b] Jerry Yao-Chieh Hu, Weimin Wu, Zhuoru Li, Sophia Pi Zhao Song, and Han Liu. On statistical rates and provably efficient criteria of latent diffusion transformers (dits). In Thirty-eighth Conference on Neural Information Processing Systems (NeurIPS), 2024.
[3YW+23] Jerry Yao-Chieh Hu, Donglin Yang, Dennis Wu, Chenwei Xu, Bo- Yu Chen, and Han Liu. On sparse modern hopfield model. In Thirty-sevntth Conference on Neural Information Processing Systems (NeurIPS), 2023.
[KKL20] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient trans- former. arXiv preprint arXiv:2001.0455) 2020.
[KS23] Tokio Kajitsuka and Issei Sato. Are transformers with one layer self-attention using low-rank weight matrices universal approximators? arXiv preprint rTXiv:2307.14023 2023.
[KVPF20] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Trans- formers are mns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156-5165. PMLR, 2020.
[LLS+24a) Chenyang Li, Yingyu Liang, Zhenmei Shi, Zhao Song, and Tianyi Zhou. Fourier circuits in neural networks and transformers: A case study of modular arithmetic with multiple inputs. arXiv preprint (et::2202.094699 2024.
[LLS+24b) Xiaoyu Li, Yingyu Liang. Zhenmei Shi, Zhao Song, and Yufa Zhou. Fine-grained attention i/o complexity: Comprehensive analysis for backward passes. arXiv preprint arXiv:2410.09397, 2024.
[LLS+24c] Yingyu Liang, Heshan Liu, Zhenmei Shi, Zhao Song, Zhuoyan Xu, and Junze Yin. Conv-basis: A new paradigm for efficient attention inference and gradient computation in transformers. arXiv preprint arXiv:2405.00 052199 2024.
[LLS+24d] Yingyu Liang, Jiangxuan Long, Zhenmei Shi, Zhao Song, and Yufa Zhou. Beyond linear approximations: A novel pruning approach for attention matrix. arXiv preprint arXiv:2410. /261, 2024.
[LLSS24a] Chenyang Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. Exploring the frontiers of softmax: Provable optimization, applications diffusion model, and beyond. arXiv preprint arXiv:2405.03255. 2024.
(LLSS244b Xiaoyu Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. A tighter complexity analysis of sparsegpt. arXiv preprint arXiv:2408. 12 2511, 2024.
[LSS$244] Yingyu Liang, Zhizhou Sha, Zhenmei Shi, Zhao Song, and Yufa Zhou. Looped relu mips may be all you need as practical programmable computers. arXiv preprint arXiv:2410.09355.000 2024.
13

---

[LSS*244] Yingyu Liang, Zhizhou Sha, Zhenmei Shi, Zhao Song, and Yufa Zhou. Multi-layer transformers gradient can be approximated in almost linear time. arXiv preprint arXiv:2408. 13233, 2024.
[LSSY24] Yingyu Liang, Zhenmei Shi, Zhao Song, and Chiwun Yang. Toward infinite-long prefix in transformer. arXiv preprini urXiv:2406. 14036, 2024.
[LSSZ24a] Yingyu Liang, Zhenmei Shi, Zhao Song, and Yufa Zhou. Differential privacy of cross-attention with provable guarantee. arXiv preprint arXiv:2407. 14717, 2024.
[LSSZ244] Yingyu Liang, Zhenmei Shi, Zhao Song, and Yufa Zhou, Tensor attention train- ing: Provably efficient learning of higher-order transformers. arXiv preprint arXiv:2405. 16411, 2024.
22] Renqian Luo, Liai Sun, Yingce Xia, Tao Qin, Sheng Zhang, Hoifung Poon, and Tie- Yan Liu. Biogpt: generative pre-trained transformer for biomedical text generation and mining. Briefings in Bioinformatics, 23(6), 2022.
[LSZ23] Zhihang Li, Zhao Song, and Tianyi Zhou. Solving regularized exp, cosh and sinh regression problems. arXiv preprint, 2303.15725, 2023.
[LWD+23] 23) Zichang Liu, Jue Wang. Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, and Beidi Chen. Deja vu: Contextual sparsity for efficient Ilms inference time. In Manuscript, 2023.
[2] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexan- der Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07777 2022.
[PCR19] Gabriele Prato, Ella Charlaix, and Mehdi Rezagholizadeh. Fully quantized transformer for machine translation. arXiv preprint a9Xiv: 10485, 2019.
[PMB19] Jorge Pérez, Javier Marinković, and Pablo Barceló. On the turing completeness of modern neural network architectures. arXiv preprint /00/::/99000000 2019.
18] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et Improving language understanding by generative pre-training. 2018.
24] Zhenmei Shi, Yifei Ming, Xuan-Phi Nguyen, Yingyu Liang, and Shafiq Joty. Discover- ing the gems in early layers: Accelerating context Ilms with 1000x input token reduction. arXiv preprint arXiv:2409. 17422, 2024.
[SSZ*24al 24a] Xuan Shen, Zhao Song. Yufa Zhou, Bo Chen, Yanyu Li, Yifan Gong, Kai Zhang, Hao Tan, Jason Kuen, Henghui Ding. Zhihao Shu, Wei Niu, Pu Zhao, Yanzhi Wang, and Jiuxiang Gu. Lazydit: Lazy learning for the acceleration of diffusion transformers, 2024.
[SSZ+24b) Xuan Shen, Zhao Song, Yufa Zhou, Bo Chen, Jing Liu, Ruiyi Zhang, Ryan A. Rossi, Hao Tan, Tong Yu, Xiang Chen, Yufan Zhou, Tong Sun, Pu Zhao, Yanzhi Wang, and Jiuxiang Gu. Numerical pruning for efficient autoregressive models, 2024.
[SZKS21] Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint :XXi:21033007011 2021.
[TDP19] Ian Tenney, Dipanjan Das, and Ellie Pavlick. BERT rediscovers the classical NLP pipeline. In Proceedings of the 57th Annual Meeting of the Association for Com- putational Linguistics, pages 4693,4601, Florence, Italy, July 2019. Association for Computational Linguistics.
[VB19] Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in a transformer language model. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP,, pages 63-76, Florence, Italy, August 2019. Association for Computational Linguistics.
14

---

[VBC20] James Vuckovic, Aristide Baratin, and Remi Tachet des Combes, A mathematical theory of attention. arXiv preprint arXiv:2007.02876 2020.
[VSP*17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[WCM21] Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: case study approximating turing machines with transformers. arXiv preprint vrXiv:2107.3316. 2021.
[WHHL24] Dennis Wu, Jerry Yao-Chich Hu, Teng-Yup Yun Hsiao, and Han Liu. Uniform memory retrieval with larger capacity for modern hopfield models. In Forty-first International Conference on Machine Learning (ICML), 2024.
[WHL+24] Dennis Wu, Jerry Yao-Chieh Hu, Weijian Li, Bo- Yu Chen, and Han Liu, STanhop: Sparse tandem hopfield model for memory-enhancee time series prediction. In The Twelfth International Conference Learning Representations (ICLR), 2024.
IWLK+20 Sinong Wang. Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006. 204,0,000000768,, 2020.
[4HH*24] Chenwei Xu, Yu-Chao Huang, Jerry Yao-Chieh Hu, Weijian Li, Ammar Gilani, Hsi- Sheng Goan, and Han Liu. Bishop: Bi-directional cellular learning for tabular data with generalized sparse modern hopfield model. In Forty-first International Conference on Machine Learning (ICML), 2024.
[XRLM21] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:21 11.00000, 2021.
[XSL24] Zhuoyan Xu, Zhenmei Shi, and Yingyu Liang. Do large language models have compositional ability? an investigation into limitations and scalability. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024.
[YBR+20] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiy Kumar. Are transformers universal approximators of geguence-to-sequenee functions? In International Conference on Learning Representations, 2020.
[YPPN21] Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self- attention networks can process bounded hierarchical languages. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguisties and the 11th International Joint Conference on Natural Language Processing (Volume Long Papers), pages 3770-3785, Online, August 2021. Association for Computational Linguistics.
[ZBB+22] Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: synthetic reasoning task, 2022.
[ZHDK233 Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating transformers via kernel density estimation. arXiv preprint arXiv:2302.02455/ 2023.
[ZKV+20] Jingzhao Zhang, Sai Pranceth Karimireddy, Andreas Veit, Seungyeon Kim. Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33: 33:15383-15393, 2020.
[ZPGA23] Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? arXiv preprint arXiv:2303.08/ 2023.
15

---

Appendix
Roadmap. In Section A, we introduce some related works. In Section B, compute the Lipschitz with respect to In Section C, we give some definitions related to the softmax function of A. In Section D, we compute the Lipschitz with respect to A. In Section E, we show our complete numerical experiments that support our theoretical results.
A Related Work
A.1 In-Context Learning
[ASA+22] indicate that Transformer-baset in-context learners are able to perform traditional learning algorithms implicitly. This is achieved by encoding smaller models within their internal activations. These smaller models are updated by the given context. They theoretically investigate the learning algorithms that Transformer decoders can implement. They demonstrate that Transformers need only limited number of layers and hidden units implement various linear regression algorithms. For d-dimensional regression problems, a ((d))hidden--iz) Transformer can perform a single step of gradient descent. They also demonstrate its ability update ridge regression problem. The study reveals that Transformers theoretically have the ability to perform multiple linear regression algorithms.
[GTLV22] concentrate on training Transformer to learn certain functions, under in-context conditions. The goal is to have more comprehensive understanding of in-context learning and determine if Transformers can learn the majority functions within a given class after training. They found that in-context learning is possible even when there is distribution shift between training and inference data or between in-context examples and query inputs. In addition, they find out that Transformers can learn more complex function classes such sparse linear functions, two-layer neural networks, and decision trees. These trained Transformers have comparable performance to task-specific learning algorithms.
1ONR+222 demonstrate the similarity between the training process of Transformers in in-context tasks and some meta-learning formulations based on gradient descent. During training Transformers for auto-regressive tasks, the implementation of in-context learning the Transformer forward pass is carried out through gradient-base -basee optimization of implicit auto-regressive inner loss that is constructed from the in-context data.
Formally speaking, they consider the following problem min, || bll2 defined in Definition 1.2. They show that one step of gradient descent carries out data transformation as follows:
A(x ( + 5z) bll2 || (b 66) l2
= -||Ar bl12
where 6r denotes the one-step gradient descent on and ôy denotes the corresponding data transfor- mation on b. They also show that self-attention layer is in principle capable of exploiting statistics in the current training data samples. Concretely, let Q, K, V IRddd denotes the weights for the query matrix, key matrix, and value matrix respectively. The linear self-attention layer updates an input sample by doing the following data transformation:
6, = + PVP Qj
where denotes the updated and P denotes the projection matrix such that a Transformer step i, on every j is identical to the gradient-induced dynamics by. This equivalence implies that when training serera-ssl--tttettoononjjj Transformers for fundamental regression tasks, the models learned by GD and Transformers show great similarity.
[XRLM21] explores the occurrence of in-context learning during pre-training when documents exhibit long-range coherence. The Language Model (LLM) develops the ability to generate coherent next tokens by deducing a latent document-level concept. During testing, in-context learning is observed when the LLM deduces a shared latent concept between examples in a prompt. They demonstrate that in-context learning happens even when there is distribution mismatch between prompts and pretraining data, especially when the pretraining distribution is mixture of Hidden
16

---

Markov Models [BP66]. Theoretically, they show that the enror of the in-context predictor is optimal when a distinguishability condition holds. In cases where this condition does not hold, the expected emoor still reduces as the length of each example increases. This finding highlights the importance of both input and input-output mapping for in-context learning.
A.2 Transformer Theory
The advancements of Transformers have been noteworthy, however, their learning mechanisms are not completely comprehensible yet. Although these models have performed remarkably well in structured and reasoning activities, our comprehension
their mathematical foundations lags significantly behind. Past research has indicated that the outstanding performance of Transformer-basee models can be attributed to the information within their components, such as multi-head attention. Various studies (TDP19, VB19, HL19, Bel22, LLS + 24a. XSL24] have presented empirical proof that these components carry a substantial amount of information, which can help resolve different probing tasks.
Recent research has investigated the potential of Transformers through both theoretical and experimen- tal methods, including Turing completeness [BPG20], function approximation [YBR* 20, CDW+ 21], formal language representation [BAG20, EGZ20, YPPN21J, and abstract algebraic operation learning [ZBB* C2B. Some of these studies have indicated that Transformers may act universal approxima- tors for 1PMB19, BPG20]. [LWD*23) 23] demonstrate the existence of contextual sparsity in LLM, which can be accurately predicted. They exploit the sparsity to speed up LLM inference without degrad- ing the performance from both theoretical perspective and an empirical perspective. [DCII+21] proposed the Pixelated Butterfly model that uses a simple fixed sparsity pattern to speed up the training of Transformer. Other studies have focused the expressiveness of attention within Transformers 18, VBC20, ZKV+20, EGKZ21, SZKS21, WCM21, LSSY24, LSS 24a] and differentially private attention mechanisms [GSYZ24, LSSZ24a]. Recently, modern Hopfield models [HYW 23, HLSL24, WHHL24, HCL+24, 24, HWL24a. HCW 24] have introduced Hopfield layers as powerful alternatives for transformer attention, offering solid theoretical guarantees and strong empirical performance [XHH 24, WHL++44]. Additionally, the statistical and computational theory of transformer-base diffusion models, specifically Diffusion Transformers (DiTs), has been studied in depth [HWL 224,, Ano24b].
operations [YBR*20 KS23, Ano24a] and emulate Turing machines
Furthermore, [ZPGA23] has demonstrated that moderately sized masked language models may effectively parse and recognize syntactic information that helps in the partial reconstruction of a parse tree. Inspired by the language grammar model studied by [ZPGA23], [DGS23] consider the tensor cycle rank approximation problem. [GMS23] consider the exponential regression in neural tangent kernel over-parameeterization setting. [LSZ23] studied the computation of regularized version of the exponential regression problem but they ignore the normalization factor. [DLS23] consider the softmax regression which considers the normalization factor compared to exponential regression problems (GMS23, LSZ23J. The majority of LLMs can perform attention computations in an approximate manner during inference, as long as there are sufficient guarantees of precision. This perspective has been studied by various research, including [CGRS19, KKI.20, WI.K 20, DKOD20, KVPF20, C1 CDL 22, LLSS24b, L LSS+ SM LLS, SSZ**44b 24b, 24a] With this in mind, [ZHDK23, AS23, BSZ24, DMS23, HYW-23, LLS 224, CLS+ LSSZ22b, HLSL24,, +WL+24b, HWL224]] have studied the attention matrix computation from the hardness perspective and developed faster algorithms.
B Lipschitz with respect to x
In Section I. we give the preliminary to compute the Lipschitz. Section B.2, we compute the I Lipschitz of function exp( (Ar) with respect to æ. In Section B.3, we compute the Lipschitz of the function a with respect to z. In Section B.4, we compute the Lipschitz of function a= with respect to x.
17

---

B.1 Preliminary
We can compute
dL
dr g(z)
Let rj > denote the learning rate.
We update
R(41 2t g(1e)
Definition B.1. We define 6, E IR.h to be the vector that satisfies the following conditions
()) 1a) exp((2++1) bll [[(ex((Ax;) 1,,) 11 exx((Ar) b 6
Let 1-1,-11 denote vector that each entry can be either -1 +1. In the worst case, there 2" possible solutions, c.g.,
(/cxp(Ax++1), 1,,) exp(Ax++1) - (oxp(Ax²), 1,,) - exp(Ar:)) о [-1,+1}"
The norm of all the choices are the same. Thus, it is sufficient to only consider one solution as follows.
Claim B.2. We can write 66 as follows
8, (exp(Axt+1) (n)=-exp(AAtr)) (exp(Ax)), 1n) exp(AAt)
f(x++1)
f(x.)
Proof. The proof directly follows from Definition B.I.
For convenience, we split 6, into two terms, and provide the following definitions Definition B.3. We define
σ,, ((exp(Axt+1), 1,) (cxp(Axx),,1. 1,) 1n)- cxp(Axt+1)
86,2 = (exp(Ax;) 1m, 11 (exp(Axx+1) exp(Ax.))
Thus, have
Lemma B.4. We have
8, 56,1 86,2
We can rewrite du, follows
5 = - exp(Arcti),
We can rewrite 6, as follows
8,,2 = a (
Proof. We have
86 = 66,1 + 5b.2
= a((141) exp(Axxx)))) - exp((Arp)) -xpp))
alxx)
a(x1) exp(Ar(+) - - exp(Ar)
=a exp Af exp( a(xx) -xx(Ax)))xxx(
(exp(Az(+1) 1, exp(AAt41) (exp(Ar;); 1n) exp(Ax;),
where the step follows from the definitions of 5g. the 2nd step follows from the definitions of 86, and 8,,2, the 3rd step follows from simple algebra, the 4th step comes from the definition of 0
18

---

B.2
Lipschitz for function exp(AAz) with respect to a
Lemma B.5. the following conditions holds
Les A E Raxd
Ler - =2) < 0.01
Lel R
Let v; y satisfy that ||0|| v R and llvll S R
Then we have
(exp(Ax) exp(Au)llz S
2VTReep(((t) -exp(A))-expAA)llll ||---~-0 y/12-
Proof. We have
(exp(Ax) exp(Au)]le S expp(x))l2 -
S vn exp( A| y)l
I - 2 y)1l2
I Vnexp(R²) ||| | - yll2
I 2/7Rexp( || - Eاا
where the step follows from ||A(y lx < 0.01 and Fact 2.1, the 2nd step comes from Fact 2.1, the 3rd step follows from Fact 2.2, the 4th step follows from Fact 2.2, the last step follows from || S R.
B.3 Lipschitz for function (zz) with respect to x
We state tool from previous work [DLS23].
Lemma B.6 (Lemma 7.2 in [DLS23]). If the following conditions hold
Let (xx be defined Definition 3.3
Then we have
(a(x)) - a(y) exp(Ax)) exp(Au)lla Vn.
B.4 Lipschitz for function a(z)-1 with respect to æ
We state tool from previous work [DIS233..
Lemma B.7 (Lemmaa 7.2 in [DLS23]). If the following conditions hold
Let (exp(Ax), 1x) 2
Let (exp(Ay)) 1m) 2 3
Then, we have
-(()) 8-22 [a(x) -ay))-
C Softmax Function with respect to A
this section, we consider the function with respect to A. define function softmax as follows Definition C.1 (Function f. Reparameterized by A in Definition 3.1). Given matrix A Raaxd Let 1 denote a length-n vector that all entries are ones. We define prediction function f Buxd I" as follows
f(A) (exp(Ar), 1m) exp(Ar).
19

---

Similarly, we reparameterized z by A for our loss function L. We define loss function L as follows Definition C.2 (Loss function Lexp, Reparameterized x by A in Definition 3.2). Given a matrix E Rxx and vector E RRaxa We define loss function Texp Raxd + R as follows
Lexp (A) = 0.5 [[(ex((Ax), ln) exp(Ax) bll2-
For convenience, we define two helpful notations a and e with respect to A as follows:
Definition C.3 (Normalized coefficients, Reparameterized x by in Definition 3.3). We define
1axd + R as follows
(AA) == (exp(Ar), 1n).
Then, we can rewrite f(A) (see Definition C.I) and Lexp (A) (see Definition C.2) as follows
exp(Ax).
(Ax = 0.5 ||a(A) -1 exp(Ax) - 61ß..
(A) 0.5 (I/(A) bll2-
Definition C.4 (Reparameterizeed by A in Definition 3.4). We define function Raxd IR.
follows
d(A) = f(A)
Then we can rewrite (A) (see Definition C.2) as follows
Lexp (A) 0.5 ||c(A)]
D Lipschitz with respect to A
In Section D.I, give the preliminary compute the Lipschitz. In Section D.2, show the upper bound of 6g with respect to A. In Section D.3, we compute the Lipschitz function (xp(Ar) with respect A. In Section D.4, we compute the Lipschitz of the function with respect A. In Section D.5, we compute the Lipschitz of function a with respect to
D.1 Preliminary
We define 6, as follows
Definition D.1 (Reparameterizeed by A in Definition B.1). We define 6 E TR to be the vector that satisfies the following conditions
(exp(AA+12)) 1,,) exp(Ae414) bll2 |I(exp(Az)). 1.)) exp(A(z) b 5ol12
Claim D.2 (Reparameterized by A in Definition B.2). We write 86 follows
6,
1,) exp(Ar+1æ) (exp(Axt), 1n) exp(xxx(()
(lA++))
((A))
Proof. The proof directly follows from Definition D.I.
For convenient, we split & into two terms, and provide the following definitions Definition D.3 (Reparameterizee x by A in Definition B.3). We define
8,,1 = (eexp(A.+12) 1,) ) --1 (exp(A,xx), 1n) -1) exp(Ac+12) 8,22 == (exp(A,r), 1,, 1 (exp(A(411) - exp(A(z))
Thus, have
Lemma D.4 (Reparameterized by A in Lemma B.4). We have
20

---

We can rewrite 6, E R= as follows
ôb 50,1 + 0b,2
We can rewrite ób,1 E RR as follows
δ,,1 =
exp(A1413),
We can rewrite 6b,2 R follows
6,,2 = a(A,)) (exp(A++1x) -exp(AA)))))(()))
Proof. We have
8n = 6n,1 + 6b.2
= = exp - a(()) exp(A4413)+ 1x)(A)))()))))
a(Ai) exp (2) -xAA)-eep(A
-a exp(Axz)
=a(A(+1) exp ( ( - exp(Ata)
(exp(A++10). 1,,) exp(A(+11)) (exp( Azx); 1n) exp(A.a);
where the 1st step follows from the definitions of og, the 2nd step follows from the definitions of 8',1 and 2,,2, the 3rd step comes from simple algebra, the 4th step comes from the definition of a.
D.2 Upper Bounding with respect A
We can show that Lemma D.5 (Reparameterized x by A in Lemma 4.1). If the following conditions hold
Let E (0, 1).
Let 86,1 E IR be defined as Definition D.3.
Let 6,,2 E " be defined as Definition D.3.
Let 86 = 86,1 + 66,2.
Let R 2 4
We have
Part
16.,1122 5 2,3 2/322 exxp((22 ||A Aelle
Part 2.
110,2112 s 23-1 (TRexx)(P))) -1Ac,1 - Arll2
Part 3.
(((+1)) - ((()) s 2exp2)n ||Ar+1 - Arl12
Proof. Proof of Part 1. We have
[ 2|2 S (a(A+11) a(A:) /
< -1- a(A)) vn exp(R²)
$ B la(A(41) a(A,) | exp( R?2)
$ B- Vn. |eep((A exp((Ax))1 122 vn ex)(
3- Vn. 2 exKRexppPPP A Aell Vn exp(R?)
21

---

2 exp(22 ||A++1 Aell
where the first step follows from definition, the second step follows from assumption on and æ, the third step follows Lemma D.8, the forth step follows from Lemma D.7, the fifth step follows from Lemma D.6.
Proof of Part 2.
We have
| 12 S (a(A+t1) (x - exp(A
$ 3 |eexp(A(110) - exp(Axx)]1
s 5 2vTRee22)))))) || Arll
Proof of Part 3.
We have
16%112 = |166,1 + 60
< 1166,1112 2 16,,2122
< 2 (eex((2R)) ||A441 Arl 2 Re..seexi (exx(21()) ||A441 AAll
< 2 Rexp(2R²) A441 All++ 23 2-2-2n Rcxp(2R2) ||A441 AAll
3 43 2 Rexp(2 ||Ac+1
where the 1st step follows from the definition of 6s. the 2nd step comes from triangle inequality, the 3rd step comes from the results in Part and Part 2, the 4th step follows from the fact that n 2 and 8-1 2 1, the 5th step follows from simple algebra.
D.3
Lipschitz for function exp(Ar) with respect to A
Lemma D.6 (Reparameterized x by A in Lemma B.5). If the following conditions holds
Let A, B E Rxd
Let [er] - B)x los 0.01
Let ||AI| S R
Let x. satisfy that ||~r|| S
Then we have
|| exx((Br) I 2/TRexp((H)) I|A Bl.
Proof. We have
[exp(Ar) exp(Br)ll2 exp(Ax)]le -
Vn.expi ex.(eexplllllllp -A B)x
s - 22] B)z||2
$ /nexp( 2 ||a||
< 2 ||A- -BI
where the 1st step follows from ||A 1o 0.01 and Fact 2.1, the 2nd step follows from Fact 2.1, the 3rd step follows from Fact 2.2, the 4th step comes from Fact 2.2, the last step follows from || < R.
D.4 Lipschitz for function (AA) with respect A
Lemma D.7 (Reparameterizee x by A in Lemma B.6). the following conditions hold
Let (AA) be defined as Definition C.3
22

---

Then we have
(a(A) aa()] $ - exp(Bx)l2 vn.
Proof. We have
Ia(A) aBB) = (eexp(Aa] exp(Bx), 1n)
I cxp(AAx exp( Bx) vn.
where the 1st step comes from the definition of a(x), the 2nd step follows from Cauchy-Schwer inequality (Fact 2.1).
D.5
Lipschitz for function (AA) with respect to
Lemma D.8 (Reparameterizeed x by A in Lemma B.7). If the following conditions hold
Let (exp(Ar), 1n) 2
Let (exp( Bx), 1m) 3
Then, we have
(a(A) - a(B))-
Proof. We can show that
- - a(B)
33 3 la(A) - a(B)
where the 1st step follows from simple algebra, the 2nd step follows from a(A) 2 BBa(B) 2 8. O
E
Experiments
In this section, show the complete numerical experimental results supporting theoretical results that when training self-attention-only Transformers for softmax regression tasks, the models learned by gradient-descent and Transformers show great similarity.
Experiments setup. According to Definition 1.3, we construct the synthetic softmax regression tasks consists of randomly sampled length-n documents A E Raxd where each word has the d-dimensional embedding and targets
R In our experiments we choose a set of different (25, 50, 100, 200, 400} and set of different embedding size d (5, 10, 20, 35, 50). Following [ONR* 22], we compare the two models in our experiment: a trained single self-attention (SA) layer with softmax unit approximating the full Transformers, and softmax regression model trained with one-step gradient descent. The training objective for both models is defined as Definition 1.3. For the single self-attention layer with softmax unit, we choose the learning rate 7/SA = 0.005. For the softmax regression model, we determine the optimal learning rate 1JDD by minimizing the &a regression loss over training set of 103 tasks through line search.
value of document length
To compare the trained single self-attention layer with softmax unit and the softmax regression model trained with one-step gradient descent, we sample 10 tasks and record the losses of two models. In addition, we follow [ONR* + +22 record
Pred Diff: the predictions difference measured with the & norm:
[SA (GD(2))
where USA( (A) is corresponding the b in Theorem 4.2, and (GD(I) is corresponding to the in Theorem 4.3.
Model Cos: the cosine similarity between the sensitivities of two models:
CosSim doa (x)
Oīs (A)
0A
23

---

Model Diff: the model sensitivity difference measured with the &g norm:
DoGD x))
(ASA(A)
2|
Or
aA
All experiments run on single NVIDIA RTX2080TI GPU with 10 independent repetitions.
Results on tasks of different document lengths. The results of the comparisons between a trained single self-attention layer and one-step gradient descent on synthetic softmax regression tasks of document length n E {25, 50, 100, 200, 400, 1000} and word embedding size d 20 are shown in Figure 4-9, We measure two models' losses and similarities over the training steps of the SA layer for each set of tasks. From the results, we observe identical performances of the two models measured in losses. We also observe considerable alignment of the two models across tasks of different document lengths, indicated by decreasing prediction and model difference and increasing cosine similarity between models.
Results on tasks of different word embedding sizes. The results of the comparisons between a trained single self-attentin layer and one-step gradient descent on synthetic softmax regression tasks of document length n = 200 and word embedding size d E (5. 10, 20, 35, 50} are shown in Figure and 10-13. Similarly, measure two models' losses and similarities over training steps of the SA layer for each set of tasks. We again observe similar performances and close aligmment the two models.
In conclusion, our experimental results empirically validate our theoretical results in Section showing that when training seff-attention-only Transformers for softmax regression tasks, the models learned by gradient-descent and Transformers show great similarity. Due to the non-linearity of softmax regression, it is not expected for models to match exactly as implied in our theoretical results in Section 4, which is also observed in our experimental findings.
Gradient Transfomer
Preds
Model
1.5
1.2
1.0
0.80
Norm
9 1.2
0.8
0.75
0.5-
0.70
0.9
0.65
0.6
200
400
600
800
1000
200
400
600
800
1000
Training Step
Training Step
(a) Losses over training steps of Trans- former
(b) Difference and similarity over training steps
Figure 4: Comparison between trained single-SA SA-layer Transformer and one-step GD on softmax regression tasks of document length 25 and embedding size d 20.
F
Limitations
Our findings are restricted to small Transformer and simple regression problems. One interesting direction for further investigation is to acquire a comprehensive perception of in-context learning larger models. To our best knowledge, we believe this work does not have any negative societal imppct.
G Impact Statements
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of work, none which we feel must be specifically highlighted here.
24

---

Gradient Descent Trained Transfomer
1.8
Model DVf
1.00
2.0-
0.95
1.5
Norm
1.5
0.90
1.2
0.85
1.0
0.80
0.75
0.0
-0.70
200
400
600
800
1000
200
400
600
800
1000
Training Stop
Training Step
(a) Losses over training steps of Trans- former
Difference and similarity over training steps
Figure 5: Comparison between trained Single-SA-layee Transformer and one-step GD softmax regression tasks of document length n = 50 and embedding size d = 20.
Gradient Descent
Preds Diff
Model Das
Taanec
4,0
3.0
2.5
3.0
0.9
Norm
Cosine Sm
2.0
2.0
1.5
1.0
0.7
200
400
600
800
1000
200
400
600
800
1000
0.6
Training
Training Step
(a) Losses over training steps Trans-
(b) Difference and similarity over training steps
former
Figure Comparison between trained one-SA-layer Transformer and one-step GD on softmax regression tasks of document length 100 and embedding size d = 20.
Gradient Descent
Preds DIfT Model
Model Cos
3.6
1.000
0.95
0.90
L2 Norm
0.85
2.0
0.80
1.0-
0.75
0.70
0.4
0.0
400
Training Step
Training Step
(a) Losses over training steps of Trans- former
(b) Difference and similarity over training steps
Figure 7: Comparison between trained one-SA-layer Transformer and one-step GD on softmax regression tasks of document length 200 and embedding size d 20.
25

---

Gradiant Descerk Traine mransfomer
Preds Dif
Model Cos
3.0
8.0
Model Dif
1.00
0.95
2.5
Norm
Lo55
4.0-
0.80
0.75
200
400
600
800
1000
200
400
600
800
1000
Training Step
Training Step
(a) Losses over training steps Trans- former
(b) Difference and similarity over training Steps
Figure 8: Comparison between trained me-SS-llayer Transformer and one-step GD on softmax regression tasks of document length n 400 and embedding size d = 20.
Gradient Descent
Preds
Model
Model
6.6
30.0
5.4
L2 Norm
0.0
50
10.0
0.6
0.6
200
400
600
800
1000
200
400
600
800
1000
0.5
Training Step
Training Stop
(a) Losses training steps of Trans- former
(b) Difference and similarity training steps
Figure 9: Comparison between trained one-SA-layer Transformer and one-step GD on softmax regression tasks of document length n 1000 and embedding size 20.
Gradient Descert
Preds DIf
Model Cos
6.0
6.0
Model Dit
1.00
5.0
0.95
Norm
4.0-
0.90
2 3.0
3.0
0.85
2.0
2.0
1.0
0.0
0.75
200
600
800
1000
800
1000
Training Step
Training Step
(a) Losses over training steps of Trans- former
(b) Difference and similarity over training steps
Figure 10: Comparison between trained one-SA-layer Transformer and one-step GD on softmax regression tasks of document length 200 and embedding size d 5.
26

---

Gradient Descent
Trained Transfomer
Model DVf
1.00
4.0
0.95
12 Norm
3.0
0.90
0
0.85
2.0
1.-
0.2
0.0
0.70
200
400
600
800
1000
200
400
6.00
800
1000
Training Stop
Traininng Step
(a) Losses over training steps of Trans- former
Difference and similarity over training steps
Figure 11: Comparison between trained one-SA-layer Transformer and one-step GD softmax regression tasks of document length n = 200 and embedding size 10.
Gradient Descent
Freds Dif Model
Model Cos
Transtormer
1.0
3.6
Sim
Norm
2.4
1.8
2.0
1.2
1.0
0.6
0.0
1.5
0
200
400
600
800
1000
200
600
800
1000
Training Step
Training Step
(a) Losses over training steps of Trans- former
(b) Difference and similarity training steps
Figure 12: Comparison between trained one- -A-lyrr Transformer and one-step GD softmax regression tasks of document length n 200 and embedding size 35.
Gradient Descent
Freds Dfff
Model Cos
Twired Thansformer
Model DIf
1.0
3.0
0.9
4.0
Ncm
0.7
2.0
0.6
1.0
0.5
0.5
0.0
200
0.4
Training Step
Training Step
(a) Losses over training steps of Trans- former
(b) Difference and similarity over training steps
Figure 13: Comparison between trained one-SA-layee Transformer and one-step GD on softmax regression tasks of document length 200 and embedding size d 50.
27

---

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 propose our main results at the end of 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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long 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:
The limitations is discussed in Section F.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 complete (and correct) proof?
Answer: [Yes]
28

---

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

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [No]
Justification:
The data and code are planned to be released upon acceptance and approval.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (uttps://nips..
wrreissssssooolooooo onPolicy) for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: //nips. wwtppblic/guttecccccccccccccccct gewei/iiggdde///eeeeeeeeeseessee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
The experimental setup described in Section 5.1 and Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification:
The results plotted in Section 5 and Appendix E include error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, statistical significance tests, 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 with given experimental conditions).
30

---

The method for calculating the error bars should be explained (closed form formula, call 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 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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The compute resource information is provided in Section 5.1 and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't I make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https http:////./e..eeeeeeeeeee http://uuurip..ccpppboooo cttppccicccc/cccccccccccccccclct wwwublic/Ethcssuiideeeo ctwuieeeeeeeeeees
Answer: [Yes]
Justification: The authors have read the NeurIPS Code Ethies and made sure the paper followsthe NeurIPS Code of Ethics in every aspect.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The potential societal impact is discussed in Section G.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
31

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that 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 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 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe 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)?
Question:
Answer: [NA]
Justification:
The paper does not release new dataset or model.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
The assets are properly mentioned and cited in Section and Appendix E.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
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 use in package should be provided. For popular datasets, paperswi www.wwtttt......eee com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
32

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 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 the dataset/code/moodl as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According 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 equivalent approval/review based 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 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.
33


Reviews:
[{"review_id": "2c63CcfR8b", "reviewer": "Reviewer_LX6R", "summary": "This research delves into enhancing the comprehension of in-context learning through theoretical analysis. It builds on prior studies showcasing how a single self-attention layer can learn gradient steps in linear regression contexts. The authors extend this concept to softmax regression and give the upper bounds of the data transformations driven by gradient descent for a single self-attention layer. Empirical results in the appendix validate these theoretical advancements.", "questions": "See above.", "limitations": "The authors adequately addressed the limitations", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- A comprehensive preliminaries and mathematical notations are defined properly\n- Studying the relationship between in-context learning and gradient-decent is helpful for understanding current LLMs", "weaknesses": "- The presentation is hard to follow. It's better to organize the formulation in a question-driven format, and it's unclear why bounding the single step of data transformation relates to building a connection between in-context learning and softmax weight shift. Consider optimizing the presentation of certain mathematical proofs by relocating them to the appendix for conciseness, and leaving more place in the main paper for intuition illustration and experiment.\n- As the theoretical analysis and experiment only consider a single attention layer, it's unclear whether other components in Transformers, such as MLP and Layer Norm, will affect the conclusion in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "***Q1: The presentation is hard to follow. It's better to organize the formulation in a question-driven format, and it's unclear why bounding the single step of data transformation relates to building a connection between in-context learning and softmax weight shift. Consider optimizing the presentation of certain mathematical proofs by relocating them to the appendix for conciseness, and leaving more place in the main paper for intuition illustration and experiment.***\n\nThank you for your feedback. We agree that the presentation can be enhanced for better clarity. We have restructured the paper by keeping necessary theoretical formulation and key findings in the main paper while relocating unused facts for proofs to the appendix. Our studied connection between in-context learning and softmax weight shift builds on the previous study [1] showing matching weights between in-context learning of a single linear transformer layer and gradient descent. We will include more intuitive explanations and illustrations on this part to aid comprehension.\n\n***Q2: As the theoretical analysis and experiment only consider a single attention layer, it's unclear whether other components in Transformers, such as MLP and Layer Norm, will affect the conclusion in the paper.***\n\nThank you for your comments. Our current analysis focuses on a single attention layer with a softmax unit introducing non-linearity in full Transformers, which extends previous works [1, 2, 3] considering only a single linear attention layer. Similar to [1, 2, 3], we believe our current work establishes the basic fundamentals necessary for future studies introducing more complex Transformer components.\n\n[1] Von Oswald, Johannes, et al. \"Transformers Learn In-Context by Gradient Descent. ICML 2023\".\n\n[2] Zhang, Ruiqi, et al. \"Trained Transformers Learn Linear Models In-Context\". Journal of Machine Learning Research 25.49 (2024): 1-55.\n\n[3] Mahankali, Arvind, 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\". ICLR 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nwnsSB2DLJ", "reviewer": "Reviewer_EkHX", "summary": "This paper explores the relationship between in-context learning and weight shifting in the context of softmax regression for large language models. The authors delve into the mathematical aspects and study the in-context learning based on a softmax regression formulation. They present theoretical results that imply when training self-attention-only Transformers for fundamental regression tasks, the models learned by gradient descent and Transformers show great similarity. The paper shows the upper bounds of the data transformations induced by a single self-attention layer with a softmax unit and by gradient descent on a ℓ2 regression loss for softmax prediction function. The authors provide a formal theoretical result and numerical experiments to validate the theoretical results.", "questions": "Please refer to Weaknesses.", "limitations": "Please refer to Weaknesses.", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents a novel perspective on the in-context learning abilities of large language models (LLMs) by examining the softmax regression formulation. This approach is innovative as it bridges the gap between the attention mechanism's role in LLMs and the mathematical understanding of in-context learning. The authors' decision to study the bounded shifts in data transformations induced by a single self-attention layer with a softmax unit is a significant contribution to the field, offering a fresh angle on the problem that has not been extensively explored in prior work.\n\n2. The paper is well-structured and clearly written.\n\n3. The quality of the paper is high, as evidenced by the rigorous theoretical analysis and the comprehensive numerical experiments conducted to validate the theoretical findings.", "weaknesses": "1. While the paper compares the performance of a single self-attention layer with a softmax unit and a softmax regression model trained with one-step gradient descent, it would be informative to include comparisons with state-of-the-art models and methods.\n\n2. The paper could benefit from a deeper exploration of the explainability and interpretability of the models in the context of softmax regression.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "***Q1: While the paper compares the performance of a single self-attention layer with a softmax unit and a softmax regression model trained with one-step gradient descent, it would be informative to include comparisons with state-of-the-art models and methods.***\n\nThank you for your suggestion. Our current analysis focuses on establishing a theoretical foundation for understanding in-context learning of Transformers through softmax regression tasks. Therefore, similar to previous works [1, 2, 3] studying a single linear attention layer, we provide our analysis and results on a single self-attention layer with a softmax unit. Empirical comparisons with state-of-the-art models, as you suggest, would indeed provide more informative and interesting findings, which we leave as a future work.\n\n***Q2: The paper could benefit from a deeper exploration of the explainability and interpretability of the models in the context of softmax regression.***\n\nThank you for your suggestion. We will include additional discussion on the explainability and interpretability of models in our studied softmax regression setting.\n\n[1] Von Oswald, Johannes, et al. \"Transformers Learn In-Context by Gradient Descent. ICML 2023\".\n\n[2] Zhang, Ruiqi, et al. \"Trained Transformers Learn Linear Models In-Context\". Journal of Machine Learning Research 25.49 (2024): 1-55.\n\n[3] Mahankali, Arvind, 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\". ICLR 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "17injgF0nR", "reviewer": "Reviewer_K6fK", "summary": "This paper aims to investigate why transformers possess the capability of in-context learning from a theoretical perspective. Previous works have shown simplified self-attention layer’s capability of learning linear functions in context. This work conducts further research based on softmax regression, as softmax-based algorithms are more complex and closer to algorithms used in actual LLMs. Through mathematical analysis and experimental validation, the authors conclude that the update acquired through gradient descent and in-context learning are similar when training simplified transformers for softmax regression tasks.\n\nI am not fully follow this paper and I wish to learn more intuitive understanding of this paper during rebuttal (from the authors and other reviewers). I may adjust my rating score after rebuttal period.", "questions": "None", "limitations": "The limitations are fine with me.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1.Explaining the reasons why LLMs can learn from context at a theoretical level is of significant importance, and this paper has made a valuable exploration into this issue.\n\n2.There is a thorough and comprehensive mathematical analysis to demonstrate the similarity between the models learned by transformers and gradient-descent on softmax regression.\n\n3.Theoretical results and experimental results corroborate each other in this paper.", "weaknesses": "1.This paper appears to build upon previous research that explored in-context learning capabilities of transformer using linear regression, but instead opts for a softmax regression approach, lacking significant innovation and contribution.\n\n2.Although this paper represents a further advancement over previous studies based on linear regression, it still employs a highly simplified Transformer model, which is insufficient to fully explain the principles of LLMs’ in-context learning ability.\n\n3.I suggest that the structure of the paper could be organized more clearly, allowing readers to grasp the overall framework first. The details of the two models compared in the experiments should be explained more thoroughly. (1) The section 2 and 3 are not so well-organized, where some formula lacks necessary comments and explanations. (2) The introduction section is not so intuitive to follow.  (3) There is almost no textual descriptions about the proposed method section (section 3). I believe most NeurIPS reviewers are hard to follow the section 3 without enough descriptions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "***Q1: This paper appears to build upon previous research that explored in-context learning capabilities of transformer using linear regression, but instead opts for a softmax regression approach, lacking significant innovation and contribution.***\n\n\nThank you for your comments. Different from previous works studying linear models with a linear self-attention layer, we take a step further to consider the important softmax unit for Transformers and provide analysis through softmax regression. Such extension approximating the non-linearity of Transformers is non-trivial and we believe our theoretical results provide novel insights for understanding the in-context learning capability of Transformers​.\n\n\n***Q2: Although this paper represents a further advancement over previous studies based on linear regression, it still employs a highly simplified Transformer model, which is insufficient to fully explain the principles of LLMs’ in-context learning ability.***\n\n\nThanks for your comments. We kindly argue that it is common to study shallow networks for theoretical analysis [1,2,3, 4]. While our analysis uses a simplified Transformer model, it serves as a foundational study, which can be extended to more complex Transformer architectures in future works and validated in more complex settings. We believe such study is crucial for understanding the fundamental mechanisms before expanding to more complex scenarios.\n\n\n[1] Nanda, Neel, et al. \"Progress measures for grokking via mechanistic interpretability.\" ICLR2023.\n\n\n[2] Morwani, Depen, et al. \"Feature emergence via margin maximization: case studies in algebraic tasks.\" ICLR2024.\n\n\n[3]Alman, Josh, and Zhao Song. \"Fast attention requires bounded entries.\" NeurIPS 2024.\n\n\n[4] Mahankali, Arvind, 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\". ICLR 2024.\n\n\n***Q3: I suggest that the structure of the paper could be organized more clearly, allowing readers to grasp the overall framework first. The details of the two models compared in the experiments should be explained more thoroughly. (1) The section 2 and 3 are not so well-organized, where some formula lacks necessary comments and explanations. (2) The introduction section is not so intuitive to follow. (3) There is almost no textual descriptions about the proposed method section (section 3). I believe most NeurIPS reviewers are hard to follow the section 3 without enough descriptions.***\n\nThank you for your feedback. In Section 2, we provide the notations and some basic algebra used in the proof of our main results. In Section 3, we formally define the network and the loss function, which are necessary to obtain our theoretical results. To help readers better understand our paper, we have added in the revised version more intuitive explanations in Sections 2 and 3 between the definitions, and moved some unused facts to the Appendix to make the paper more concise and easier to understand.\nRegarding the introduction, since our main findings are from the theoretical results, it is necessary to define the notions and concepts used in Theorem 1.4 and such necessity may make the introduction less intuitive to follow. To help readers better follow this section, we have included additional discussions to present our theoretical findings of similarity between the models learned by gradient descent and Transformers as well as the experimental results.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper provides a theoretical analysis of in-context learning in Transformers, focusing on the role of the softmax unit within self-attention layers. By examining the upper bounds of data transformations and comparing these with gradient descent on a regression loss, the authors offer insights into the similarities between models learned by Transformers and traditional methods. While the study is primarily theoretical, it fills a critical gap in understanding how softmax contributes to in-context learning. Based on the reviews, I recommend to accept the paper.
Decision: Accept (poster)