JoMA: Demystifying Multilayer Transformers via Joint Dynamics of MLP and Attention
Authors: Yuandong Tian
Yiping Wang
Zhenyu Zhang
Beidi Chen
Simon Shaolei Du
Venue: ICLR 2024 poster (2024)
Keywords: multilayer transformer
training dynamics
theoretical analysis
self-attention
interpretability
neural network understanding
Primary area: optimization
TL;DR: We analyze the training dynamics of multilayer transformer, characterizing the role of self-attention, MLP nonlinearity, and the learning procedure of hierarchical structure, if the data follow hierarchical generative models.

Abstract: We propose Joint MLP/Attention (JoMA) dynamics, a novel mathematical framework to understand the training procedure of multilayer Transformer architectures. This is achieved by integrating out the self-attention layer in Transformers, producing a modified dynamics of MLP layers only. JoMA removes unrealistic assumptions in previous analysis (e.g., lack of residual connection), and predicts that the attention first becomes sparse (to learn salient tokens), then dense (to learn less salient tokens) in the presence of nonlinear activations, while in the linear case, it is consistent with existing works. We leverage JoMA to qualitatively explains how tokens are combined to form hierarchies in multilayer Transformers, when the input tokens are generated by a latent hierarchical generative model. Experiments on models trained from real-world dataset (Wikitext2/Wikitext103) and various pre- trained models (OPT, Pythia) verify our theoretical findings. The code is at https://github.com/facebookresearch/luckmatters/tree/yuandong3.

Published as a conference paper at ICLR 2024
JoMA: DEMYSTIFYING MULTILAYER TRANSFORMERS VIA JOINT DYNAMICS OF MLP AND ATTENTION
Yuandong Tian
Yiping Wang
Zhenyu Zhang
AIGMeta (FAIR)
University of Washington
University of Texas Austin zhenyu. zhenyu.zhangduteer www.wauteaaaaa....oooo edu
yuandong@meta.com com
ypwang616cs washington.err edu
Beidi Chen
Simon Du
Carnegie Mellon University, Al@Meta (FAIR) beidic@meta. com, beidickandrev. cmu.edu edu
University of Washington ssduêcs wasssnatshhnnnn edu
ABSTRACT
We propose Joint MLP/Attenin Attention (JOMA) dynamics, a novel mathematical framework to understand the training procedure of multilayer Transformer architectures. This is achieved by integrating out the self-attention layer in Transformers, producing a modified dynamics MLP layers only. JoMA removes unrcalistic assumptions from previous analysis (e.g., lack of resid- ual connection) and predicts that the attention first becomes sparse (to learn salient tokens), then dense (lo learn less salient tokens) in the pres- ence of nonlinear activations, while in the linear case, it is consistent with existing works that show allention becomes sparse over time. We leverage JoMA to qualitatively explains how tokens are combined to form hierarchies in multilayer Transformers, when the input tokens are generated by a la- tent hierarchical generative model. Experiments on models trained from real-world dataset (Wikitext2/ 1eatteext//iiitextt//ee WWikiexxt 103) and various pre-trained mod- els (OPT, Pythia) verify our theoretical findings. The code is at
INTRODUCTION
Since its debut. Transformers (Vaswani et al. 2017) have been extensively used in many applications and demonstrates impressive performance (Dosovitskiy et al. 2020 OpenAl 2023) compared to domain-specife models (e-g., CNN in computer vision, GNN graph modeling, RNN /LN/MSTM in language modeling, etc). In all these scenarios, the basic Trans- former block, which consists of one self-attention plus two-layer nonlinear MLP, plays a critical role. natural question arises:
How the basic Transformer block leads to effective learning?
Due to the complexity and nonlincarity of Transformer architectures, it remains a highly nontrivial open problem to find a unified mathematical framework that characterizes the learning mechanism of multi-layer transformers. Existing works mostly focus on 1-layer Transformer (Li et al.] 2023a Tarzanagh et al.) 2023b) with fixed MLP Tarzanagh et al.) 2023a layer, linear activation functions (Tian et al. 2023 and local gradient steps at iitiaaization (Bictti et al. 2023 Oymak al.] 2023
etc.
In this paper, we propose a novel joint dynamics of sell-atlention plus MLP, based on Joint MLP/ Attention Integral (JoMA), a first integral that combines the lower layer of the MLP and self attention layers. Leveraging this joint dynamics, the self-attention shown to have more fine-grained and delicate behavior: it first becomes sporse as in the linear case (Tian et al.) 2023), only attends tokens that frequently co-occur with the query, and then becomes denser and gradually includes tokens with less frequent co-occurrence, in the case of nonlinear activation. This shows changing inductive bias in the Transformer training: first the model focuses on most salient features, then extends to less salient ones.
Another natural question arises: why such a learning pattern is preferred? While for 1-layer this does not give any benefits, in multilayer Transformer setting, we show qualitatively that such a dynamics plays an important role. To demonstrato that this is the case, we assume hierarchical tree generative model for the input tokens. In this model, starting from the

---

Published as a conference paper at ICLR 2024
upper level latent variables (in which the top-most is the class label of the imput sequence), abbreviated as sVs, generales the latents LV,-1 in the lower layer, until reaching the loken level (s = 0). With this model, we show that the tokens generated by the lowest latents LV, co-occur lot and thus can be picked up first. by the at aetntion dynamics as "salient features". This leads to learning of such token combinations in hidden MLP nodes, which triggers self-attention grouping at = 1, etc. In this way, the non-salient co-occurrences are naturally explained by the top level hierarchy, rather than incorrectly learned by the lower layer as spurious correlation, which is fortunately delayed by the attention mechanism. Our theoretical finding is consistent. with both the pre-trained models such as OPT/Pythi and models trained from scratch using real-world dataset (Wikitext2 and Wikitext
We show that JoMA overcomes several main limitations from Scan&Snapp (Tian et al. 2023 JoMA incorporates residual connections and MLP nonlinearity as key ingredients, analyzes joint training of MLP and self-attention layer, and qualitatively explains dynamics of multi- layer Transformers. For linear activation, JoMA coincides with Scan&Snap, ie.., the attention becomes sparse during training.
1.1 RELATED WORK
Training Dynamics of Neural Networks. Earlier research
has delved into training
dynamics within multi-layer linear neural networks (Arora al.
2018
Bartlett al. 2018
the teacher-student setting (Brutzkus & Globerson 2017 Tian
2017
Soltanolkotabi
2017
Du et. al. 2017:
2018a; Xu & Du 2023), and unlinite-width limits
Jacot al. 2018 Du et al.
2018b
Allen Zhu et al. 2019
Oymak & Soltanolkotabi 2020 Li & Liang. 2018 Nguyen &
Pham
2020
Fang et al.
2021
b. This includes extensions to cawpnttton--add-moooooe
(Hron
et al.
2020
Yang et al.
2022
In self-supervised learning, analysis exists for dynamics In
deep linear networks TTia
2022 and the impact of nonlinearity Tian 2023).
Dynamics for Attention-based models. Zhang et al. (2020) delves into adaptive opti- mization techniques. Jelassi et al. (2022) demonstrates that the vision transformer (Dosovit- skiy ct al. 2020) trained via gradient descent can discern spatial structures. Li ct al. [022cc] illustrates that single-layer Transformer can learn a constrained topic model, where each word is tied to a single topic, using l loss BERT-like framework (Devlin et al.) 2018), and certain assumptions
attention patterns. /Snell et al. (2021) investigate the training dynam- ics single- head attention in mimicking Seq2Seq learning. Tian et al. (2023) characterizes the SGD training dynamics of a 1-layer Transformer and shows that with cross-entropy loss, the model will pay more attention to the key tokens that frequently co-occur with the query token. Oymak et al. (2023) constructs the attention-
contextual mixture model and demonstrates how the prompt can attend to the sparse context-relevant tokens via gradient descent. [Tazanaahh el al. (002b)) also finds that running gradient descent will converge in direction to the max-margu solution that separates the locally optimal tokens from others, and (Trzanaah el al. (2023a) further disclose the connection between the optimization geom- etry of sell-allention and hard-margin SVM problem. For the in-context learning scenario, several recent works analyze linear transformers trained on random instances for linear gression tasks from the perspective of loss landscape (Boix-Adsera Adsera el al. 2023 Zhang el al., 2023). While these studies also study the optimization dynamics of attention based models, they do not reveal the phenomena that we discuss.
Expressiveness of Attention-based Models. The universal approximation abilities of attention-based models have been studied extensively (Yun et al.
Bhattamishral jotaal, al. (2200a b Dehghani et al.) 2018 Pérez et al.) 2021). More recent studies ofter detailed
2019
insights into their expressiveness for specific functions across various scenarios, sometimes
2021 Likhosherstov incorporating statistical evaluations (Edelman et al. 2022 Elhage et al.
(et al. 2021 Akyürek et al.) 2022; Zhao et al. 2023 Yao et al. 2021 Anil et al. 2022
Barak et al.) 2022). A fruitful line of work studied the in-context learning capabilities of the Transformer (Dong et al.) 2022), linking gradient descent in cuassffiaatio/rogrsssii
learning to the feedforward actions in Transformmer layers (Garg et al. 2022; Von Oswald ct al. 2022 Bai et al. 2023, Olsson et al.) 2022; Akyürek et al. 2022 Li et al. 2023b) However, unlike our study, these work do not characterize the training dynamics.
2 PROBLEM SETTING
Let the total vocabulary size be M, in which Mc is the number of contextual tokens and Me is the number of query tokens. Consider one layer in multilayer transformer (Fig. 1 (b):
he d(wp f), f = Ucb + uhy b = o(2)) x//A A
(1)
2

---

Published as a conference paper at ICLR 2024
(ww)
MLP loyer)
MLP Nortinssrit
MUP layeri
MLP ioyeri
Noriinearity gf
w
Nonlinesrity
Nonlincrrtyy
Nonlinsarity
MIP (lower layeri
JoMA
MEP llower
MLP
f = Ugb
llower
Modifice MLP llawer lsyer)
Modified llower layer)
(x,) xx//
altendon
attennion
Xg Sef-attention
Layer 0
Layer
Layer
Layer
"Tii apple
Figure 1: (a) Overview of JoMA framework. Using the invariant of training dynamics,, the self- attention layer and the lower layer MLP can be merged together to yield a MLP laver with modified dynamics (Theorem 1, which explains the behaviors attention linear (Sec. 3.1 and
nonlinear (Sec. 1 ML.P activation Φ. as well as hierarchical concept learning in multilayer cases (Sec. 5) (b) Problem setting. JoMA frameworks support different kind of attentions, including
linear altention by == Zz exp atlention by == zqe /A and softmax b == mee / E,1
Input/outputs. x == [ma]] E &e is the input frequency vector for contextual token 1 < S Mc, I q S Mq is the query token index, K is the number of nodes in the hidden MLP layer, whose outputs are hg- All the quantities above vary across different sample index (i.c,, 21 = 11
qfi)). In addition, & is the nonlincarity (c.g., RcLU).
Model weights. bn = [=w]] E Me is the (unnormalized) attention logits given query 4 and w/ E Rdd are the weights for the lower MLP layer. These will be analyzed in the paper.
The Attention Mcchanism. In this paper, we mainly study three kinds of attention:
Linear Allention Von Oswald et al. 2022): (xx) = x and A 1;
Exp Allention: (x) = exp(x) and
const;
Softmaz Attention Vaswani et al.
2017): (xx) exp(x) and A: == 1 (o(zz) x)).
Here o is the Hadamard (element-wise) product. b RM are the attention scores for con- textual tokens, given by point-wise attention function a. is the normalization constant.
Embedding vectors. ui is the embedding vector for token 1. We assume that the embed- ding dimension d is sufficiently large and thus u, uy = II([ = 1'), i.e., {u} are orthonormal umel E Rd Me be the matrix that encodes all embedding vectors of contextual tokens. Then U. = I. Appendix B.1 verifies the orthogonality assumption in multiple pre-trained models (Pythia, LLaMA, ctc).
bases. Let Uc = [1, .2...., usel
Residual connections
introduedd as an additional term un in Eqn. [ which captures the critical component in Transformer architecture. Note that we do not model value matrix Wy since it can be merged into the embedding vectors (e.g., by un Wyw),, while Wx and WQ are already implicitly modeled by the self-attention logits 2gt- ux, W Wx
Gradient backpropagation in multilayers. In multilayer setting, the gradient gets backpropagated from top layer. Specifically, let 9h, [i] be the backpropagated gradient sent to node at sample i. For 1-layer Transformer with softmax loss directly applied to the hidden nodes of MLP, we have gh. [i] [[[o[[]] = k), where yoli] is the label to be predicted for sample i. For brevity, we often omit sample index if there is no ambiguity.
Assumption (Stationary backpropagated gradient gh2). Expectation terms involving ghz (e.g., E [9h, x)) remains conslant during bruining.
Note that this is true for layer-wise training: optimizing the weights for a specific Trans- former layer, while fixing the weights of others and thus the statistics of backpropagaled are stationary. For joint training, this condition also holds approximately since the weights
change gradually during the training process. Under Assumption I Appendix A.1 gives an equivalent formulation terms of per-hidden node loss.
3

---

Published as a conference paper at ICLR 2024
X
X
0.75
Comportual tokknn indoxt
Normallzod Conclation
NCS2vftl. Ze10)
0.25
NC.2c10.22d19
Comtuttul
0.00
NOKrr10..2edt)
0.5
NC
0.00
NC talti
-0.50
-025
Class labely
Class labal
1500
Number Batches
Number Batches
Figure Test training dynamics with linear MLP activation (Ø(x) 2) under softmax attention. Left Two: The distribution æ smoothly transits over different class labels. Right Two: The distribution
over different classes are randomly generated. both cases, estimated an (t) by the first integral (Theorem D] despite assumptions on bm, shows high correlation with the ground truth self-attention logits zm()), while its components &m1 (1) Σx vz(() and
(t) = - Ex (t)lib), do not.
Training Dynamics. Define the conditional expectation Bg-m []] == E [-lg = m]. Now
us consider the dynamiis of wr and Rm, if we train the model with a batch of inpuls that always end up with query qii m, then:
ih = Zggm [gh Mgf,,
=mym (Ob/0zm)) Utgr
(2)
Here 1/4 = d'w f) is the derivative of current activation and gs == Ex Ihu Shrnh/e kg
JOMA: EXISTENCE OF JOINT DYNAMICS OF ATTENTION AND MLP
While the learning dynamics of wa and 2m can be complicated, surprisingly, training dy- namics suggests that the attention logits (t) have close-form relationship with respect to the MLP weights (t),), which lays the foundation of our JoMA framework:
Theorem 1 (JoMA). Lel Ua = w]w,, then the dynamics of Eqn. salisfies the invariants:
Linear attention. The dynamics satisfies (t) = E3 (t +c.
Exp allention. The dynamics salisfies 2m (1) E. (1) tc.
Softmax attention.
If bm = Eg=m [6] is a constant over time and
Ey=m EE 9h. N1 bb 6 Eyzm |Ex Ihn b,, then the dynamics satisfies zm sm(t) =
1 Ex vZ(() - (t) (1.m +c.
Under zero initialization (wx (0) = 0, 2m (0) = 0), then the time-independent constant 0.
Therefore, we don't need to explicitly update self-attention, since is already implicitly incorporated in the lower layer of MLP weight! For softmax attention, we verify that even with the assumption, the invariance proposed by Theorem still predicts 2t(() fairly well.
3.1 LINEAR ACTIVATIONS: WINNER-TANE TAKE-ALL ALL
Now we can solve the dynamics of wx(() (Eqn. 2), by plugging in the close-form solution of self-attention. For simplicity, we consider exp attention with K = 1 (i.e., single hidden MLP node). Let. Am. Emn= [9,, Mar. then vs dynamics is (2% written v):
Am exp(zm) Am exp(x? exp(o2/22+)) c)
(3)
In the case of linear activations ((x) x, h/ = According to Assumption 1 not depend on v and we arrive at the following thcorem:
Am does
Theorem 2 (Linear Dynamics with Self Self-attention). With linear MLP activation and zero initialization, for erp attention any two tokens + satisfy the following invariants:
orf (vx(t)/2) ct)(2x(t)/2) Am
(4)
Apm
where Alm Eg=m 2i) and erf(x) 3s See -t dt is Gauss error function.
4

---

Published as a conference paper at ICLR 2024
Initialization
after convergence
Number of contesttull tokons (N.)
0.02
componento
1.5
component2
0.01
compurent2
companent:
1.0
0.00
0.0
0.01
-0.5
0.5
10
-0.03
-15
500
1000
1500
Number MiniBatches
Figure 3: Growth different components vo(l) (First few components the first column of V(t) in linear MLP activation and softmax attention. As predicted by Sec. 3.1 after convergence, only some components of du grows while the remaining components is saatrated after initial growing, consistent with Theorem Leven if is derived from JoMA's ' approximation in Theorem [ Each node l (and thus wa) receives back- propagated gradient from k-th class via cross-entropy loss.
Remarks. The dynamics suggests that the weights become one-hot hot over training, Specif- x oo and other (t) converges to
ically, let 1 = arg max, |m then (t) - finite numbers, because of the constraint imposed by Equ. (see Fig. 3] For softmax atten- tion, there is an additional sample-dependent normalization constant A[i,, if A[i remains constant across samples and all elements of 6m are the same, then Theorem also applies.
Beyond distinct/ /oott//ommoonn tokens. Alm == Er,q=m [9h, [ m is product of token discriminancy (i.e., 1g..gmm [9h.] > O means Loken 1 positively correlated to backpropagaled gradient 9h, or label in the 1-layer case) and loken frequency (i.e.. P(4/m)) how frequent / appears given m). This covers broader spectrum of tokens than Tian et al. (2023). which only discusses distinct (i.e., large [Atml) and common tokens (i.e.; when A/m R3 R).
4 TRAINING DYNAMICS UNDER NONLINEAR ACTIVATIONS
In nonlinear case, the dynamics turns out to be very different. In this case, Am is no longer a constant, but will change. As result, the dynamics also changes substantially.
Theorem 3 (Dynamies of nonlinear activation with uniform attention). If æ is sampled from a misture of isotropic distributions centered at :R,,,, where each a. E Rd and gradient ghe constant within each mirture, then:
b Am
Σ ag01 (re)
Σ aal2(rr)
(5)
|0012
here dc == Ey [aha P(c], To == TTa. te + is the affinity to &e and the "bias" term 5(1)== So Bqum [gh, My] dl, 0, and O2 depend derivalive of nonlinearily sp o' and data distribution but not v. If 2b is monotonous with (--20) = and (+++)) = 1, so does 0j.
Appendix A.3.2 presents eritical point analysis. Here focus simplified one when is constrained to be unit vector, which leads to the following modified dynamics (P, v = 0):
= P. = ar.
Σ (r)) zellime - (v (c)0)
(6)
where == Ez/l|Eell We consider when v is aligned with one cluster De but far away from
others, then Ve >> rel for d' 1 c and 01 >> (rer) since is monotonously increasing. Hence Me dominates and let / = Me for brevity. Similar to Equ. 3, we use close-form simplification of JoMA to incorporate self-attention, which leads to (we use exp attention):
x (μ - v) o exp(v?/2)
(7)
Here we omit. the scalar terms and study when v is close to μ. in which v T -===11 O(I w113) ze 1. It is clear that the critical point v, = does not change after adding the term exp(v² 12). However, the convergence speed changes drastically. As shown in the following lemma, the convergence speed towards salient component of μ (i.e., component with large magnitude) is much faster than non-salient. ones:
"Since zi is the empirical frequency of token I in sample have Alm B.-m [gh. Σ, Sbg P(]g m, () m) Σ, 9h. P)P(( m. P)P(ig m) BBg.=m [gkg.] P(]m).
5

---

Published as a conference paper at ICLR 2024
Colored line: dynamics of v(t). Dashed line: target μ
Entropy changes over time
5000
2.25
4000
Intropyvvva
2.00
3000
1.75
2-
2000
1.50
1.25
1000
1.00
0
0
1000
2000
3000
4000
5000
Sorted index v components
eiteration
Figure 4: Dynamics of nonlinear MLP with self-attention components included (Equ. 7 Left: Training dynamics (color indicating training steps). The salient components (i.e., components with large magnitude (4) of e(t) are learned first, followed by non-salient ones. Right: Entropy of the attention (i.e., (softaaxs drops when salient components are learned first, and then rebounces when other components catch up.
Theorem 4 (Convergence speed of salient vs. non-salient components). Let (t) := ,,(t)/p; be the convergence metric for component j (t) = means that the component converges). For nonlinear dynamics with attention (Eqn. 7. then
In (0))__()))
et²//2
(1 Asx(t))
(8)
In (0)/6(())
ext/2
Here (t)(t) - Aym(t)- (a0tt)-ea//1) In 1n (0)))_).)) where [)x((t)] C5k and Cşs only depends on (0) and (0). So when (t) < (0)( expl Cys exp(ux)), we have [A(t) < 1.
Remarks. For linear attention, the ratio is different but the derivation similar and simpler. Note that the convergence speed heavily depends on the magnitude of Hÿ- If Hj > μi,, then (t) < (t) and (t) converges much faster than (t). Therefore, the salient (i.e., large) components is learned first, and the non-salient (i.e., small) component learned later, due to the modulation of the extra term exp(we /2) thanks to self-attention, as demonstrated in Fig.
A follow-up question arises: What is the intuition behind salient and non-salient components in μu? Note that μ is an fz-nommllized version of the conditional token frequency æ, given the query q m. In this case, similar to Theorem 2 (and TTan et al. (2023)), ), again see that if a contextual token co-occurs a lot with the query m, then the corresponding component becomes larger and the growth speed of i towards i is much faster.
Relationship with rank of MLP lower layer. Since MLP and attention layer has joint dynamics (Theorem [], this also suggests that in the MLP layer, the rank of lower layer matrix W (which projects into the hidden nodes) will first drop since the weight components that correspond to high target value H4 grow first, and then bounce back to higher rank when the components that correspond to low target value Hí catch up later.
How SELF-ATTENTION LEARNS HIERARCHICAL DATA DISTRIBUTION?
A critical difference between the training dynamics of lincar and nonlinear MLP is that in the nonlinear case, although slowly, the non-salient components will still grow, and the entropy of the attention bounces back later. While for 1-layer Transformer, this may only slow the training with no clear benefits, the importance of such a behavior is manifested we think about the dynamics of multiple Transformer layers trained on data distribution generated in a hierarchical manner.
Consider a simple generative hierarchical binary latent tree model (HBLT) Tian et al. 2020) (Fig. 7 (a))) in which we have latent (unobservable) binary variables at layer that generate latents at layer
1, until the obser observable tokens are generated at the lowest level (s 0). The topmost layer is the class label yo, which can take D discrete values. In HBLT, the generation process of ye at layer
given Ya layer s be characterized by their conditional probability IP(y, 1ly, = = Plya 0/ya 0] = {1 + p). The uncertainty hyperparameter p E [-1,]] determines how much the top level latents can determine the values of the low level ones. Please check Appendix A.. for its formal definition.
With HBLT, we can compute the co-occurrence frequency of two tokens | and m, as function of the depth of their common latent. ancestor (CLA)):

---

Published as a conference paper at ICLR 2024
label (observed)
(a) Setting
Vo
(b) Shallow Latent Distributon
yo
(c) Deep Lasens
ostribution
Yo
Ya
CLA m))
Ya
Latent
variables
obserced)
V
CLA me)
y
V
CLA m))
Tokens (observed)
772
m
772
772
772
Figure 5: (a) Hierarchical binary tree generative models. Except for yo that is the observable label of a sequence and can take D discrete labels, all latent variables follow binomial distribution. A binary leaf variable yi = indicates that token appears the sequence. (b) Attention dynamics in multi-layer setting. There is a strong co-occurrence between the query m and the token 1, but a weak co-oceurrence between m l' As result, m associates with ( first, and eventually associates with 1', even if they co-occur weakly, according to Theorem 4 (c) Tf there exists an additional layer and Yar in the latent hierarchy, the association m-l and m' -I' will be learned first due to their high co-occurrence. Once the lower hierarchy gets learned and some hidden nodes in MLP represents ya and yor (see Sec. for experimental validation), on the next level, yo and yer shows strong co-occurrence and gets picked up by the self-attention mechanism to form even higher level features. In contrast, the association of of l'-m is much slower and does not affect latent hierarchy learning, showing that self-attention mechanism is adaptive to the structure of data distribution.
Theorem 5 (Token Co-occurrence in HBLT(p))- If loken I and m. have common latent
ancestor (CLA) of depth (Fig. 3(c)), then PP[y = 1/ym = 1] [1p2" -20"-'p0 pa
where
1- pl too
L is the lotal deplh of the hierurchy and po == P. Pl in which po = [P[yo k]] E RD and P.oo = Ply 0/8o- ER,, where {ur are the immediate children of the root node Vo-
Remarks. If yo takes multiple values (many classes) and each class only trigger one specific latent binary variables, then most of the top layer latents are very sparsely triggered and thus Po is very close to 1. If is also close to 1, then for deep hicrarchy and shallow common ancestor, Ply: 1 ll 1] 1. To see this, assume P = Po e, then we have:
1
1 +++1 2Hc 2(1 Lc) + + H +0(2)
Ply- ym 1]
(9)
2
(1 Le)
This means that two tokens and m co-occur a lot, if they have a shallow CLA (H small) that is close to both tokens. If their CLA is high in the hierarchy (e.g., l' and m). then the token l' and have much weaker co-occurrence and P((()m) (and thus zy and pr) is small.
With this generative model, we can analyze qualitatively the learning dynamics of JoMA: first it focuses on associating the tokens in the same lowest hierarchy as the query (and these tokens co-occur a lot with m), then gradually reaches out to other tokens I' that co-occur less with m, if they have not been picked up by other tokens (Fig. b]]) if co-occurs a lot with some other m', then m-l and m'-l' form their own lower hierarchy, respectively. This leads to learning of high-level features ys and yer which has high correlation are associated in the higher level. Therefore, the latent hierarchy is implicitly learned.
EXPERIMENTS
Dynamics of Attention Sparsity. Fig. shows how attention sparsity changes over time when training from scratch. We use 10 learning rate and test our hypothesis on Wikitext2/ Wikktext2/Wikittxt.tk Wikitext103 (Merity et al. 2016 (top/bottom row). Fig. 8 further shows that difforent lcarning rate leads to different attention sparsity patterns. With large learning rate, attention becomes extremely sparse as in (Tian et al.) 2023 Interestingly, the attention patterns, which coincide with our theoretical analysis, yield the best validation score.
We also tested our hypothesis in OPT iZhang et al. derman el. al., 2023 (Pythia-70M/144B// (691)) tre-trained models, both of which has public intermediate checkpoints. While the attention patterns show less salient frop-and-boumt patterns, the dynamics of stable ranks of the MLP lower layer (projection into hidden neu- rons) show much salient such structures for top layers, and dropping curves for bottom layers since they are suppressed by top-level learning (Sec. 5). Note that stable ranks only depend on the model parameters and thus may be more reliable than attention sparsity.
2022) (OPT-2.7B) and Pythia Bi-

---

Published as a conference paper at ICLR 2024
Layers: 1, val Joss: 5.357
vol_oos:
Layers: 2. val_less:
Jess: 5.255
Layers: val loss: 5.169
Layers: 10, val loss: 5.110
2:50
layaro
layero layor2
Aotmmon Entropy
2.25
2.00
1.50
1.25
1.00
3'
20 Layers:
40
Layers:
5.047
Layers: 2.
Layers: 4.679
2.5
layero
Aetortion Entrayy
2.0
layerO
50
100 150
200
5
50
100
150
200
0
100 200
100
200
Minibatch
Minibatch 00
Minibatch (kl
Miribatch 00
Figure Dynamics of attention sparsity. Tn 1-layer setting, The curves bear strong resemblance to our theoretical prediction (Fig. 4 in multi-layer settings, the attention entropy in top Transformer layers has a similar shape, while the entropy in bottom layers are suppressed due to layer interactions (Sec. Top row: Wikitext2, Bottom row: Wikitext103.
Tythe TON. Lepur: val 6:221
Pythe 70%, Layans val loa 4,22
gem-278 Layers 30, vil Vae 2,224
071.2.1 Laysn: 32, veliess 3.26
3643
012
0035
3.55
3122
Prestn
Pyerb-3.40 Layart 24, va We 2,425
1.1%
2/711.1600 10, Sypre 26 vsl ere 2,025
Py01 Layer: 32, vll lec: 2.350
0r1
22733.6 Layare 30, vs lone 2.358
0013
3.54
12
A623
3.62
3005
Figure Dynamics of attention sparsity and stable rank OPT-2.7B and 1yythia-70M/1.411 Results are evaluated on Wikitext103 (eeitty et al. 2016).
1e-06, val: 6.43 Se-06, val: S. 5.90 Ir: 1e-05, val: 5.69 In: 50-05, val: 5.27ir: 0.0001, val: 5.21 0,001, val: 5.32 Ir:0.01, val: 6.76
2.0
1.5
3.0
0.5
1agw2
3.0
20
20
20
40
20
40
Minibatchh
Mintbatch 16
Mingbatch
Mivibatch
Minibatch
Minlbatch
Minibatch
Ir: 1e-06, val:: 6.41 In: 5e-06, val: U...5.. 5.89
1e-05, val: 5.67
Se-05, val:: 5.23
0,0001, val: 5.17
0.001, val: 5.30
0.01. val: 6.76
2.5
2.0
1.5
Inyero
2.0
0.5
byyer
13y274
0.0
20
40
40
20
20 40
Minbatcs
Minibatch
Hirbbath
Minbstch
Min
Mintzatch
Minibanch
Figure 8: Effect of different learning rates attention sparsity. Different learning rates lead to different dynamics of attention sparsity, and the altention patterus consistent with theoretical analysis (Fig. 4 give the lowest validation losses.

---

Published as a conference paper at ICLR 2024
C= 20, No.
C 20,
Noh (20, 30)
C= 30, Noh
(No, N1)
(10, 20)
(20, 30)
(10, 20)
(10, 20)
(20, 30)
NCorr (s =0)
0.99 = 0.01
0.97 0.02
1.00 = 0.00
0.96 = 0.02
0.99 = 0.01
0.94 0.04
NCorr (s = 1)
0.81 = 0.05
0.80 = 0.05
0.69 = 0.05
0.68 0.04
0.73 = 0.08
+.44:003
C = 30 Nob =
C = 50,
Noa
C = 50,
Noh. =3
(No, N1)
(10, 20)
(20, 30)
(10, 20)
(20, 30)
(10, 20)
(20, 30)
NCorr (s =0= 0)
0.99 + 0.01
0.95 + 0.03
0.99 + 0.01
0.95 - 0.03
0.99 + 0.01
0.95 + 0.03
NCorT (s 1) 0.72 + 0.04
0.66 + 0.02
0.58 + 0.02
0.55 0.01
0.64 + 0.02
0.61 + 0.04
Table 1: Normalized correlation between the latents and their best matched hidden node MLP of the same layer. All experiments are run with random seeds.
Validation of Alignment between latents and hidden nodes in MLP. Sec. 5 is based on an assumption that the hidden nodes in MLP layer will learn the latent variables. We verify this assumption in synthetic data sampled by HBLT, which generate latent variables in a top-down manner, until the final tokens are generated. The latent hierarchy has hyperparameters: number of latents per layer (N,) and number of children per latent (Nen..
C' is the number of classes. Adam optimizer is used with learning rate 10-5. Vocabulary size M = 100, sequence length T = 30 and embedding dimension d = 1024.
We use 3-layer generative model as well as 3-layer Transformer models. We indeod per- ceive high correlations between the latents and the hidden neurons between corresponding layers. Note that latents are known during imput generation procedure but are not known to the transformer being trained. We take the maximal activation of each neuron across the sequence length, and compute normalized correlation between maximal activation of each neuron and latents, after centeralizing across the sample dimension. Tbl. [ shows that indeed in the learned models, for each latent, there exists al least one hidden node in MLP that has high normalized correlation with it, in particular in the lowest layer. When the generative models becomes more complicated (i.e., both Neh and N₁ become larger), the correlation goes down a bit.
7 DISCUSSION
Deal with almost orthogonal embeddings. In this paper, we focus on fired ort orthonormal embeddings vectors. However, in real-world Transformer training, the assumption may not be valid, since often the embedding dimension d is smaller than the number of vocabulary M so the embedding vectors cannot be orthogonal to each other. In this setting, one reasonable assumption is that the embedding vectors are almost orthogonal. Thanks to lemma, one interesting property of high-dimensional space is that for M embedding vectors achieve almost orthogonality |u, ur S e, only d 2 Se- - M is needed. As a result, our JoMA framework (Theorem D will have additional e-related terms and we leave the detailed analysis as one of our future work.
Training embedding vectors. Another factor that not considered in JoMA is that the embedding vectors are also trained simultancously. This could further boost the effi- ciency of Transformer architecture, since concepts with similar semantics will learn similar embeddings. This ossentially reduccs the vocabulary size at cach layer for learning to be more effective, and leads to better generalization. For example, in each hidden layer 4d hidden neurons are computed, which does not mean there are 4d independent intermediate "tokens", because many of their emboddings are highly correlated.
Self-attention computed from embedding. JoMA arrives at the joint dynamics of MLP and attention by assuming that the pairwise attention score Z is an independent parameters optimized under SGD dynamics. In practice, 2 = WW x is also parameterized by the embedding matrix, which allow generalization to tokens with similar embeddings, and may accelerate the training dynamics of Z. We leave it in the future works.
CONCLUSION
We propose JoMA, a framework that characterizes the joint training dynamics of nonlinear MLP and attention layer, by integrating out the sell-attention logits. The resulting dynam- ics connects the dynamics of nonlinear MLP lower layer weights (projection into hidden neurons) and self-allention, and shows that the attention first becomes sparse (or weights becomes low rank) and then becomes dense (or weights becomes high rank). Furthermore, we qualitatively give a learning mechanism of multilayer Transformer that reveals how self- attentions at different. layers interact. with each other to learn the latent feature hierarchy.
9

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
Simon S. Du is supported by supported by NSF IIS 2110170, NSF DMS 2134106, NSF CCF 2212261, NSF Us 2143493, NSF CCF 2019844, NSF I1S 2229881.
REFERENCES
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learn- ing algorithm is in-context learning? investigations with lincar models. arXiv preprint arXiv:2211.1560) 2022.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via paaamttrririiioonnn In International Conference on Machine Learning, pp. 242 252. PMLR, 2019.
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowyez, Vedant Misra, Vinay Ra- masesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. arxiv preprint arXiv:2207.04900 2022.
Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. arXiv preprint orXiv:1810.00 022811, 2018.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisti- cians: Provable in-context learning with in-context algorithm selection. orXw preprint ar 0ix:023006.0 04637, 2023.
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Ad- vances in Newral Information Processing Systems, 35:21750 217644 2022.
Peter Bartlett, Dave Helmbold, and Philip Long. Gradient descent with identity initializa- tion efficienlly learns positive definite linear transformations by deep residual networks. In International conference on machine learniing, pp. 521 530. PMLR, 2018.
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint Xr::::2009. 1 1124,, 2020a.
Satwik Bhattamishra, Arkil Patel, and Navin Goyal. On the computational power of trans- formers and its implications in sequence modeling. arXiv preprint arXiv:2006. 092866 2020b.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Authony, Herbic Bradlcy, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397- 2430. PMLR, 2023.
Alberto Bictti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Botton. Birth of transformer: A memory viewpoint. arXi preprint arXi::23060088/ 2023.
Enric Boix-Adsera Etai Littwin, Emmanuel Abbe, Samy Bengio, and Joshua Susskind. Transformers learn through gradual rank increase. arXiv preprint arXm::230.. 07042, 2023.
Alon Brutzkus and Amir Globerson. Globally optimal gradient descent for a convnet with gaussian imputs. In International conference on machine learning, pp. 605 614. PMLR, 2017.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXie preprint 11x7: 2018.
Jacob Devlin, Ming Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre- training of deep bidirectional transformers for language understanding. arXiv preprint Xrx:o:1810 805, 2018.
10

---

Published as a conference paper at ICLR 2024
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhilang Sui. A survey for in-conlexi learning. arXiv preprinl ar Xiv:23301 000344, 2022.
Alexcy Dosovitskiy, Lucas Beyer, Alexander Kolosnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXw::2000.1100 2020.
Simon Du, Jason Lee, Yuandong Tian, Aarti Singh, and Barnabas Poczos. Gradient descont learus one- nue-hidden-layer enn: Don't be afraid of spurious local minima. In International Conference on Machine Learning, pp. 1339-1348. PMLR, 2018a.
Simon S Du, Jason D Lee, and Yuandong Tian. When is convolutional filter easy to learn? arXi preprint 1r 709.06619, 2017.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes wer-paaraetterized neural networks, 2018b. URL https://arxiv.o org/abs/ abs/ 1810. 02054.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang, Inductive biases and variable creation in self-atlention mechanisms. In International Conference on Machine Learning, pp. 4793-5831. PMLR, 2022.
N Elhage, N Nanda, C Olsson, T Henighan, N Joseph, B Mann, A Askell, Y Bai, A Chen, T Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuils Thread, 2021.
Cong Fang, Jason Lee, Pengkun Yang, and Tong Zhang. Modeling from features: a mean- field framework for weer-paramoterized deep neural networks. In Conference learning theory, pp. 1887 1936. PMLR, 2021.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes, Advances in Neural Information Processing Systems, 35:30583-30595 2022.
Jiri Hron, Yasaman Bahri, Jascha Sohl-Diekstein, and Roman Novak. Infinite attention: Nugp and ntk for deep attention networks. In International Conference on Machine Learning, pp. 4376-4386. PMLR, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Aduances in neural information processing systems, 31, 2018.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Newral Information Processing Systems, 35:37822 37836, 2022.
Hongkang Li. Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of
shallow vision transformers: Learning, generalization, and sample complexity. In The Eleventh International Conference on Learning Representations, 2023a. URL https: wtwp:eewwwwewweefffrreeeeee //openreview net wwt..e//oorummddddcillv....a.ii
Shuai Li, Zhao Song, Yu Xia, Tong Yu, and Tianyi Zhou. The closeness of in-context learning and weight shifting for softmax regression. Xive preprint XXXi:2204411276 2023b.
Yuanzhi Li and Yingyu Liang, Learning overparameterized neural networks via stochastic gradient descent on structured data. Advances in neural information processing systems, 31, 2018.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. arXiv preprint arXiv:2303.04245 2023c.
Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self-attention matrices. arXiv preprint NXi:22006.0976, 2021.
11

---

Published as a conference paper at ICLR 2024
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. mixture models. ar preprinl or Xiv: 07843, 2016.
Pointer sentinel
Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multilayer neural networks. arXiv preprint arXiv:2001.11443 2020.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXie preprint XrXiv:2209.110 10009..11899 11895, 2022.
OpenAL. Gpt-4 technical report, 2023.
Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparamettrization:: Global convergence guarantecs for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory, 1(1):84 105, 2020.
Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. ICML, 2023.
Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research, 22(1):3463-3444 -3033 2021.
Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint (kk::2103.07600) 2021.
Mahdi Soltanolkotabi. Learning relus via gradient descent. Advances in neural informalion processing systems, 30, 2017.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Trans- formers as support vector machines. arXiv preprint Xrxvv22308.168 16898, 2023a.
Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. arXiv preprint rrXiv:2306.13550 3(7):47, 2023b.
Yuandong Tian. An analytical formula of population gradient for two-layered relu network and its applications in convergence and critical point analysis. In International conference on machine learning, pp.
PMLR, 2017.
Yuandong Tian. Understanding the role of nonlinearity in training dynamics of contrastive learning. arXiv preprint arXi::2206. 01342, 2022.
Yuandong Tian. Understanding the role of nonlincarity in training dynamics of contrastive learning. ICLR, 2023.
Yuandong Tian, Lantao Yu, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning with dual deep networks. arXiv preprint arXi::20100000 00578, 2020.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer, 2023.
Ashish Vaswani, Noam Shazoer, Niki Parmar, Jakob Uszkoreit, Llion Joncs, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. 2017. URL
https://arrae http///xxwvoog//dff./// / 03762.pdr
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. ar preprint orXi:22212.070 07677, 2022.
Weihang Xu and Simon S Du. Over-parameterization exponentially slows down gradient descent for learning a single neuron. arXiv preprint 2rXiv:2302.1003, 2023.
Greg Yang, Edward J Hu, Igor Babuschkin, Sxymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs w: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint or orXi:2203.008 03466, 2022.
12

---

Published as a conference paper at ICLR 2024
Shunyu Yao, Binghui Peng, Christos Papadimitrion, and Karthik Narasimhan. Self- attention networks can process bounded hierarchical languages. ar prepriul arXiv:2105.1111: 2021.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequennce-t--sequeeeee functions? arXiv preprint arXi::1112.1007 2019.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungveon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models?
Advances in Neural Information Processing Systems, 13:15333-15333, 2020.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. arXiv preprint orXiv:2306. 09927, 2023.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, al. Opt: Open pre-trained transformer language models. aie preprint orXiv:2205. 0at.:20550006,,, 2022.
Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? or Xrxi preprint ar arXiv:2303. St::2303.8117, 2023.
13

---

Published as a conference paper at ICLR 2024
A PROOrS
A.1 PER-HIDDEN LOSS PORMULATION
Our Assumption has an equivalent per-hidden node loss:
max ED f1o (2m))
Σ gh heg
max {ar (2m)
BiwD
Σ ghe ( [0]
(10)
where ghg [t] the backpropagated gradient sent to node he sample
A.2
JOMA FRAMEWORK (SECTION
3
Theorem 1 (JoMA). Let 1s == Ug wk, then the dynamics of Equ. satisfies the incariants:
Linear attention. The dynamics satisfies (t) = Ex (t +c.
Exp attention. The dynamics satisfies 2m (1) 1/2 Ex 1 (t) + cc
If bm E E==m EE 9h. bb 6m Byzm
Softmax ottention.
Bg=m [b] is constant time and 1Ex Ihe Ihe b,b,, then the dynamics solisfies zm el())
2 Σx vt(() 1y (t) 13) +c.
Under zero initialization (wx (0) = 0, Zm zm( = 0), then the time-independent constant c c= 0.
Proof- Let I= 0b/0zm. Plugging the dynamics of wy into the dynamics of self-attention logits Zm, we have:
2m Eggm LTUC Σ ghe ki
Egm h
T vx)
(11)
Before we start, we first define (1) = So Fm lon, (L)))( di' Therefore, É Eyzm (9h, My] Intuitively, & is the bias of node k, regardless of whether there exists an actual bias parameter to optimize.
Notice that U + Ue uy with orthonormal condition between contextual and query tokens: U/ um 0, and thus U b, which leads to
d = Ue tb Up Eg=m [oh, M&fl Eg=m (gin hi
(12)
Unnormalized attention (A == const). In this case, we have b = o(xm) a x/A and (zm() o x))A A= diag [d((=m) diag(b) and thus ((2m)
diag(o'((
2m Σ Eymm (ghe higL Un] diag (2(2m)) (zm)
Σ
By=m (ghe hibb] o un
(13)
o' (Zm)
diag
Σ is o ukk
(14)
σ 2m
which leads to
[(mm) 0' ('(
diag
2m Σ Uz Uk
(15)
Therefore, for linear attention, a (2m)) (2m) = Zm: by integrating both sides, we have 2 (t) = E. vz(() c. For exp attention, (0zm)//((zz) = 1, then by integrating both sides, we have (t) Ee (t( +c.
Softmax attention. In this case, we have L diag(b) Tb Therefore,
Eyzm (gh, didaa(()) Ue wik Byg=m Jgh, M[b vs Ug
(16)
14

---

Published as a conference paper at ICLR 2024
where is the Hadamard (element-wise) product. Now Therefore, have: E Bym (gh, Rig 1 Ue us = 0% t/
(17)
Given the assumption that b is uncorrelated with Σx Ihh h (e.g., due to top-down gradient information), and let 6m = Eg=m [b), we have:
Zm Σ
Us U 0g
(18)
If we further assume that pm is constant over time, then we can integrate both side to get close-form solution between (1) and {1x(
Zm (m(() (t) = 2 Σ (a) | (1xm)2mm)) +
(19)
Theorem 2 (Linear Dynamics with Self-attention). With linear MLP octivalion and zero initialization, for exp oftention any two tokens 1 + satisfy the following invariants:
erf (ur(t)/2) erf(ev(t)/2) Alm
(4)
Avm
where Atm Eg=m mi] and erf(x) / So e e-t² dt is Gauss error function.
Proof. Duc to the assumption, we have:
di Ey=m [i] xil exp((mm)/44 Aim exp(2m))/A
(20)
where Aim = Rg=m [g]. m = m, y(i),, then Aim Ry.g=m [ghe] P(]m). Note that for linear model, Atm is a constant over time.
Plugging in the close-form solution for exp attention, the dynamics becomes
in Aim 2
/2
(21)
Assuming q = O, then for any two tokens + 11, we gel.
Am exp(Zmi)
Alm exp (07 (27/2)
(22)
dy Drm exp(2xnt)
Arm (xppu//22
which can be integrated using erl(-) function (i.e., Gaussian CRF: erf(x) = V7 S0
2
-t dt):
erf (vn(2)
erf(er (1)/2)
+ cu
(23)
! Am
Aum
if v0) 0, then ew O.
A.3 DYNAMICS OF NONLINEAR ACTIVATIONS (SEC.
A.3.1 WITHOUT ATTENTION (OR EQUIVALENTLY, WITH UNIFORM ATTENTION)
Lemma 1 (Expoctation of Hyperplane function under Isotropic distribution). For any isotropic distribution p(æ 2) with mean t in subspace spanned by orthonormal bases R, if + O, we have:
Ep [ uT x + E)] = (r()) x |
Oz(re)
01 (1(rv) |
RR v. |0
Rp [v(u x +5]]
(24)
where To = - + is the (signed) distance between the distribution meon & and the affine hyperplane (v. E). (r(r) and 02(r) only depends on & and the underlying distribution but not v. Additionally,
(r) is monotonously increasing, then 01 (r) is also monotonous increasing;
15

---

Published as a conference paper at ICLR 2024
(f)) 2 0, then (1(r) 2 0;
If (-(o)) = 0, (+000 1, then 01(- = 0 and 01 (+00)
If (-<o)) = O, then (2(- 0.
Proof. Note that x' isotropic in span R) and thus p(x') just depends on ||2||], we let Po Rt + Rt satisfies Pa (o( = p(x'). Our goal is to calculate
/ Jppan(()
Ep (w x E)]
x x E)p(x uddx
(25)
R)
(x' + u))(w" aaa a' rw (lu)p(())d
(26)
Japan(R) RR)
where x' := x 1--- is isotropic. Since R... w is the projection w onto space span( R), we denote v == RR w and y' == TTx x = Ta since w' lies in span Then let any hyper-plane through v, which divide span(R) R)) into two symmetric part V- and V- V-oundarry zero measurement set and can be ignored), we have,
Sgrnnce a'vl(w" x' Ja' "ww)(("'dd
P1
(27)
J/r + [/e )a''b("T TT x' + tw)p(p')dd
(28)
x [ TTxd
|ull [ell| (00 Taa an + rw rwm) )aa
(29)
1l JepantR) /vlll + (r)
(30)
|aa
Eqn. 29
holds since for every x' E V+, we can always find unique x" eV_ defined as
x" =- vT
TT
2y
11 v) + |/ v
x
(31)
|0
where x" and w' satisfy ||~"|| || v. xx" w w', and have equal reverse component
+(x' TT [|a| v) perpendicular to u. Thus for the x' in Equ. 28 only the component parallel to remains. Furthermore, let {uI..... Un-1, w/[|o|] to be orthonormal bases of span [R) and denote x u. x Vi E in - 1], then we have
of Voll + +rw)d 1
/1/
fo p(a')da" ...dda...ddr,, ..r''
P
(32)
|
V/v(2l/ + rw ( (un))
(33)
1
Here Pn Pn(y) is the probability density function of y' obtained from x'. For the trivial case where = 1, clearly pr())') po(lv/) = =(y'). If n 2 2, it can be further calculated as:
f ... 3. poly (z1)2
Pn()')
(2 (2/)2) dri dz
(34)
Pol vy2 + +22). -
(35)
(n () 11/2
to0
po 1/22 + 2dl
(36)
I(
n/2-11
Pol Pol 1/2 +13) 2d,
even
(n - -1)!!
Jo
=
27(1-1)/2 1)/2
+00
(37)
Po PolV +1)) rr 1m-2
odd
("=3)!
16

---

Published as a conference paper at ICLR 2024
where (R) l(n/2+1) R represents the surface area of an n-dimensional hyper-sphere
nx²
of radius r denotes the gamma function and we use the property that I + 1) = n! and T(n + 2 = (2n - 1) a for any N+.
Similarly, for another term we have
P2 = SSPantR)) μ ".w(ww" T Ta' + trw)p(x')da
(38)
[//00 w/y +rw)pa(/ddd' /
t00
(39)
|0||
(40)
Finally, let
S s/ byy + (
(rw))
(41)
(2(ww)
(42)
Then we arrive at the conclusion.
Theorem 3 (Dynamics of nonlinear activation with uniform attention). If æ is sampled
from a minture of C isotropic distributions centered at (#1,..., where each Ta E RRd and gradient Ih are constant within each misture, then:
ò Am
|1012 Σ (r + 10 Σ a.02(
(5)
here Og E. [ P(c], re vT o. E the affinity to Bh and the "bias" term
E(t) = So B-=m [gh, Mj] dt, 01 and Oz depend on derivative of nonlinearity up = ø' and data distribution but not v. If sb is monotonous with (00-0)) = and (+000) = 1, so does 0j.
Proof. Since backpropagated gradient ghs is constant within each of its mixed components, we have:
E q 19he b₂b] = Σ E Ey-mm (g bic = jl
(43)
By=m,c=j [ohnl |Pc = BB [b
(44)
Σ Ca m;) [bo'(w) [T]]
(45)
Let vb 0'. Note that w f w TU.b ug) v b and with uniform attention b x,, we have:
Am = Eo [ov(v] x +
(46)
Using Lemma leads to the conclusion.
Remarks. Note that if & is linear, then /b = 0, = and 02 0. In this case, 0, constant, which marks a key difference between linear and nonlinear dynamics.
A.3.2 (TENTATIVE) CRITICAL POINT ANALYSIS OF DYNAMICS IN THEOREM
Lemma (Property of 0,, O2 with homogeneous activation). If ((x)) = 2 (x)) a homo- geneous activation function and + = ő', then we have:
dr +r)()) = (r))
(47)
17

---

Published as a conference paper at ICLR 2024
Integrating both sides and get:
02(r) + r0((r) (0((()==((r) =F(r)= = [ ['aam)ar++
(48)
Let == and it is clear that C 0z(()) Thus
(2(r) + r0 (r) F(r) ( (((r))rr
f + +2((0)
(49)
of 2 then F(r) is a monotonous increasing function with F(+00) to0. Furthermore, if --o (r) = 0 and (-))) = 0, then 02(-x0) 0 and F(-00) = 0 and thus
F(r) 2. 0.
Proof. Simply verify Eqn. 47 is true.
Overall, the dynamics can be quite complicated. We consider a special C C= 2 with one positive (a+, r+ and #+) and one negative (a_, r- and Ø-) distribution.
Lemma 3 (Existence of critical point of dynamics with ReLU activation). For any homoge-
neous activation ((x) r (x), any stationary point of Eqn. 5 must satisfy E, 0, (r,,) 0, where F(r) = (0) + 5 0, (f") is a monotonous increasing function.
Proof. We rewrite the dynamics equations for the nonlinear activation without attention case:
1u11 Σ Ea,0((1)) + b En (002(1s)),
1
é= |0ll
Σ Eg0((1s)
(50)
Notice that u- = ry - E, this gives that:
| - - +)++ +)<=2,00007 (,02(1,)
(51)
Σ4(50(1))) +02(r)) - Σ (rs)
(52)
= -afF(r) - 110
(53)
in which the last equality is because the dynamics of E and due Lemma 2 Now leverage the condition of stationary points (ò 0 and & = 0), we arrive at the necessary conditions at the stationary points:
- =0
(54)
Note that in general, the scalar condition above is only necessary but not sufficient. Equ. has Mc + equations but we only have two scalar equations (Eqn. 50 and
||ol|25 E, (r,)) = 0). However, we can get a better characterization of the stationary points if there are only two components at and a_:
A special case: one positive and one negative samples In this case, we have (here r+ = T x+ + and r_ == T x_ + +)::
a F(44) - a =0
(55)
So the sufficient and necessary condition for (v,E) be the critical point is that
F(r+) (1+)) a F(r_) (r(r) a+
(56)
Without loss of generality, we consider the case where & is ReLU and (b(r = I/r > 0).
Note that 0, is a monotonously inereasing function, we have 8-1 (0, - such that (0, (r)) = r for any r E R. And we demote G: (0, 1) + B which satisfies: (v))
G(y) F(O (v))
(57)
18

---

Published as a conference paper at ICLR 2024
G(y)
G(y+)
G
0
y-
y+ 1
Figure 9: The plot of function G(y).
and := (r+), , y- == (r-). Then if we can find some line 4. y kr for some
1 E R such that la has at least two points of intersection (Vt, ky,), 1, 2 with curve G and a-/a+= = 01/32 or a-/a+ = 92/U1, then we can always find some v and & such that Equ. 56 holds.
On the other hand, it's easy find that (Fig.
9
dG(v)
(x) = (1)
> 0
dy
Pn(x)
lim Gly) = lim F(r) +00
s+1
r-++00
lim G(y) = lim F(r) lim r0](r) x=0
---
----0
Note that since G G(y=)/G(y = y+/y-, we have G(y+)/v+ G(y-)/v- and thus (U+, G(y+)) and G(y)) are lying the same straight line.
For finding the sufficient condition, we focus on the range 2 0 and (zr) > 1.2 Then in order that line 4 y kx for some has at least two points of intersection with curve G, we just need to let
(() dG(y)
1== = 02( (0) = pn(()) ( s'pals/)dy' 2
(58)
(0) dy
For convenience, let S = ((,,v)ly kx) and So {(x v) G(x)} to be the image of the needed functions. Denote R² - R Fi x1((x,)) y)) x for any x, E R, Ti(S)= (s)(s) E S}. Therefore, if Eqn. 58 holds, then the following set S will not be empty.
S =
Va + x2 E E (Sr. (Sa))
(59)
KER
And Eqn. has critical points a+ a4/a_ E S. And it's easy find that Vs E S, 8 (H, U (1, Similar results also hold for other homogeneous activations.
Remarks. It is often the case that y- < 1/2 and y+ 1/2, since G(y) when 1/2 is convex and there will be at most two intersection between a convex function and a straight line. This means that rf > 0 and r- = = < 0.
19

---

Published as a conference paper at ICLR 2024
Pythia-160M
Every morning. the slowly seokens with the distant hurn of maficc and choppin of offprresss John takos monaii savur the poacoful ambrance before he walks blck, prestng familiar and the fesh
realm physies, votur subjasial iemper tuat long
aimsophere
undangees phase
gie, prodtaiig
technelogical cuttaary applications.
The Sahara Desert, North Bind largst the and
dancs and sconching tomperatons. hans enviroeents
snouee speckss adipted
Novels, beyood deeir estestaizmen value, senv minses to seciey, ofter reflecting culural, social, and politiall mances thii Authon lke George Oewell and June Ausion recal their works
provide traights into the world they lival
Pythia-70M
are their independent nature. Many people appreciate then their tow-maittetanee lifesle, uszaional platime.
with just condoraabl spoc zap
Rainfieces are vital for the Earth's ecssystem. They provide labitai for comicee specios of which not found any where clsc. Additionally they play cracial
regulating clinre producing oxyge.
The Eifld Lower, Iconie landmark in Puss, originally cnssructed temporay cbibi for the 1959 Words Fai: Over the year, i: las becoene systhol of nomance and inchtiouural strreting millions tourion anmullyy.
The buman digestive systems comples actwork of organs wocking togeterr interine, each per play
csaental musicsts, Beginning mouth coaring our bodice nosire eaegy vicemine nocded daily
Examples pallern superposition: the same neuron in MLP hidden layers can be activated by multiple irrelevant combinations of tokens (A and B each group, e.g., same neuron activated by both "Every morning" and "In the realm physics"), in Pythia-70M and Pythia-160M, models. Bold tokens are what the query token attends to.
Figure 10:
A.4 SEVERAL REMARKS
The intuition behind &: Note that while node in MLP layer does not have an explicit bias lerm, our analysis above demonstrates that there exists an "implicit bias" term (t) embedded in the weight vector wy:
wll) w0) + Uc(v(1) - v(0) + umE(())
(60)
This bias term allows encoding of the query embedding um into the weight, and the negative bias 5 0 ensures that given the query g = m, there needs to be positive inner product between v. (i.c., the "pattern template") and the input contextual tokens, in order to activate the node k.
Pattern superposition. Note that due such mechanism, one single weight may contain multiple query vectors (e.g., Um. and Um,)) and their associated pattern templates (e.g., Umy and Umg), as long as they orthogonal to each other. Specifically, if w Um: - Em, Um, + Um. Emz Umzy then can match both pattern and pattern 2. We called this "pattern superposition", as demonstrated in Fig. 10
Lemma 4. If ((x) is homogeneous, i.e., δ(x) = (')x), then there erist constant c_ _. E depend on d such that ((x) = c- 1s 0) + c [ 0], and thus
d0 = (c_ + + (r)) dr
d0²
- (c_ + +c+rr Pr(
(61)
dr
Proof. For any x O, we have
6(x + dx) - ((x)
d'(z) = lim
(62)
6
dz
8'(x 6x) (x)
lim
x
lim o' d' + 8x)
(63)
60-0+
5x
o' (x + 8x) - 6'(x) Φ(xx)
lim
(64)
6x-
dx
(65)
So for any > 0, Φ'(x) must be constant, and similar results hold for < O. Then by direct calculation, get the results.
20

---

Published as a conference paper at ICLR 2024
A.4.1 WITH SELF-ATTENTION
Lemma 5. Let g(y) 1-e
Then max,209(y) 3/2
V2.
Proof. Any of its stationary point ys must satisfics gg(v)) = 0, which gives:
e v.
(66)
2y² +
Therefore, at any stationary points, have:
2y.
2
g(u-)
S
(67)
2y +1
2y. + Vo
V2
since g(0) g(+00) 0, the conclusion follows.
Lemma (Bound of Gaussian integral). Let G(y) = 00 S/2 @2 then $ G(y) S for 2 0.
Proof. G(y)) 2 is obvious. Note that
2 G(y) = e e- 1 (2 SE0D col2ds " 2 = v2g(u/V2)
Applying Lemma gives the conclusion.
Theorem 4 (Convergence speed of salient
vs. non-salient components). Let (1) = 1
vy (1(t)pz be the convergence metric for component j (5, (t) 0 means that the component.
converges). For nonlinear dynamics with attention (Eqn. 7. then
In (0)/__)))
c/,//2
ex2/2 (1 + Azx(t))
(8)
In (0)/bx()
Here (1) = (t) 22 In - (6x(0)6_())) where (1)x(()) S Cjk and Cjk only depends on 60) and (0). So when (6x(t) < (0)( expl- exp(xx)), we have [A(t)) < 1.
Proof. We first consider when μ 0. We can write down the dynamics in a component wise manner, since all components share the same scalar constant:
0j (e3 0/2
(68)
8 (μu 2/ 822/2
which gives the following separable form:
bye-v9 1//2
U&E v//2
(69)
Mi Uj
11 Uk
Let
1TH TH 0²
e ==²22112
Top /1 du
dx
(a= v/n)
(70)
Integrating both sides of Eqn. equation at time t:
69 from t = 0 to t, the dynamics must satisfy the following
F(r;(1),r;()) (z) F(rx((), (0),
(71)
where (t) = (t)/ According to the dynamics, ry(t) - 1 and the question is how fast the convergence is. Depending on the initialization, r1( or r1())= 1.
Eqn. 71 implicitly gives the relationship between rt()) and (t) (and thus 8y(t) and 8u (t)). Now the question is how to bound F(r, ro, 11,, which does not. have close-form solutions.
21

---

Published as a conference paper at ICLR 2024
Note that we have:
OF
x²e 1 dx
11
(72)
dp
1 -t
To
( c-12o212
x
11
Jro
I
e
1-x dx
(73)
x2 x222212 /2 (r;
11
(1 2)
(74)
π
ry
Toju
1
-ōu - - - 1u ro. (75)
V
erf
V2
erf
122
+ -
f
ro.
Let
Tojl
- ( (e- -2-0121
1
<(r, To. (1)
erf V2
erf
+
(76)
v2
11
Applying Lemma and notice that μ > 0, we have
IS, n)1 S V2T v2((rol + 1rl)//2 < V2T + m2x(2 |rol + 1) =: M (ro)
(77)
which means that [S(r, To, μn is uniformly bounded, regardless of p. and r(t) (note that r is bounded and will converge 1 from the dynamics). Integrating both side and have:
δ Op
(2
<(r, uoe
(78)
12 ((r. - F(r,ro,0) = I."
(79)
For, = / ro, 0) + 1 I C(r, ro, 12 (80)
Note that Fr,,ro,0) 0) has close form:
For, ST 1 dr In 1-
1
(81)
1-r
has a close-form solution that works for both ro < < and ro > (the situations that 1 is between To and r won't happen). Using mean-valis value theorem, we have:
F(r, = 22
1-ro I-r
M)e- + Mo, 1212 I" dx 12d
(82)
Applying Lemma 6 have the following bound for F(r.)):
-M(ro) S F(r,p) = In --r
1 1- To
< M(ro)
(83)
When is close 1 (near convergence), the term e- -2² In 1-ra (with fixed μ and fixed ro)
T-r
is huge compared to the constant M (ro),, which is /2m + 1.5 xs 4.0066 for e.g., Irol 1/2, and thus F(r,u) + -2 22/2
I-r"
To be more concrete, note that 8(t) v(t)/p((pppppp) r(t), we let 8(0)
p()(1),)) = F(1 5(1,, ())) - -1/2
(-M(ro), 8(1)
(84)
And using Equ. 71 we have:
F(1 -,()),,-6,(0), = F(1 (x(t), (k()).
(85)
Then
)tx(t) p(5x(t), - -1)
(86)
= 13// In
6,(
/2/2 In
(0)
(87)
(t
dx(t)
and (Ase(1)) S (())(0)) + ((r(()).. Then arrive at the conclusion.
22

---

Published as a conference paper at ICLR 2024
A.5 HIERARCHICAL LATENT TREE MODELS (SECTION
5
We formally introduce the definition of HBLT here. Let Me be a binary variable at layer (upper layer and be binary variable at layer (lower layer). We use a 2x2 matrix a to represent their conditional probability:
Plys 0a = 0)
Plya = 0 yo = Plys = 1a =
Pa = [Plyslya]]
(88)
Plus 1/ 0]
Definition 1. Define 2 matrix (p) = 1
1 +
1 1- P and 2-dimensional vector
1 1-p p
1+p
p(p) 1/11 + p, T]T for P E [-1,1]. Lemma (Property of M())- M()) has the following properties:
M(p) is symmetric matria.
M = 12.
M(P1) = M(P1)2)- So matric molliiplication in (M(D)) is comma- nicatine and isomorphic to scalar multiplication.
M(pi)p(p2) p(p122).
Proof. The first two are trivial properties. For the third one, notice that M(p) = (11""
pee ) in which e === [1, [1,- Therefore, TT 2 and 1 e 0 and thus:
(p2) = 1 (112 + + pieee)((11 + Pree = (11 #iPzee = M (pip2)
1
(89)
For the last one, note that p(p) = {1 pe) and the conclusion follows.
Definition (Definition of HBLT). In HBLT(p), Pala = (P3ta)) where Pala E [-1,11] is the uncertainty parameter. In particular, Pāja p, then just write the entire HBLT model as HBLT(p).
Lemma 8. For latent ya and its descendent Yy, we have:
Paja = P718, Pp.13. ... Pazla = M (Pyla)
(90)
where Prla == P918, ... Payla and a x B1 > B2 > > Bb. is the descendent chain from Yo to y---
Proof. Due to the tree structure of HBLT, we have:
Pluy|ua]= = Σ Plyyluo, J(yo, [e,. ... Plya, [a]
(91)
which is precisely how the entries of Pris, PA PBolo get computed. By leveraging the property of M (p), we arrive at the conclusion.
Theorem 5 (Token Co-occurrence in HBLT(p)). If token I and m have common latent
ancestor (CLA) of depth (Fig. (()),, then Ply: = 1yym = 1] = 1+ -2- 200--to) where 1-p' 1-ppl1pp
L is the total depth of the hierarchy and Pu == p. P in which po = [P/v k]] RD and p.o = P(y = 0/ = k]] RDD, where {ur} are the immediate children of the root node yo-
Proof. Let the common latent ancestor (CL.A) of yo, and yos be ye: then we have:
Plys, Vo.l Σ Plys, lvloolvv/PPsolluuele
(92)
Let P3,0, [P[ya, y.. then have:
PA (() (Ppal))
(93)
23

---

Published as a conference paper at ICLR 2024
Pe
where D(c) == (Fla(((()))) =
is a diagonal matrix, and Pa = 21p/
1 P- Pe
1. Note that
= e D(c)e==)
(ole)e=e-D((0 Pc
(94)
And M()) = 2111" + ppee))
therefore we have:
p3.g2 M(AI (Po21))
(95)
= (11" + Ppllee )D(c) + ppoleee))
(96)
4
= (11" + PB PoP + PAIPPell
+ P8glopele
(97)
4
Now we compute Pe. Note that
Plyel Σ Pl [3o]
(98)
Let P. == := [Ply]]] be a 2-dimensional vector. Then we have P. = Pell Po p ((ecoPo) where Pu is the probability distribution of class label you which can be categorical of size C:
Pe
Po PelnoP
Σ
Delvi DIIDOP
(99)
91
1 + Pıjo
1+ P210
... 1 PoIa
M(pq)) 2
Pijo
1 P2jo
1 Pata
Po
(100)
+ P.PPP
M(p.1) 2 - P. oPo
(101)
M(PQLPTTOPO)
(102)
in which y1 the last binary variable right below the root node class label yo-
Therefore, Pe = PeliPos where Po := p Po is the uncertainty parameter of the root node yo-
all Paia = p for immediate parent Yo and child ya, ys, is for token I and ya, for token m, then Pe = PBale = =p", and Pajt p 1- 1-1-- and thus we have:
Ply 1, 3m
1 + 2H - n
2p" Pe
Plyt 1/ym =
(103)
P Plym =1]
2
1 + p 2p--110 Pa
(104)
2
ph- Po
and the conclusion follows.
B MORE EXPERIMENT RESULTS
B.1 ORTHOCONALITY OF EMBEDDING VECTORS
We verify the orthogonality assumption mentioned in our problem setting (Sec. 2 The orthogonality is measured by absolute cosine similarity cossim(a1, E [0, 1] of two vectors 2 and Ø2:
x
x2
cossim( 22)
(105)
|
Here the two vectors 2i and #2 are column vectors of the out-projection (or upper) matrix of MLPs different layers, each corresponding one hidden neuron. For a MLP layer with model dimension d and hidden dimension 4d, there will be 4d such column vectors. We measure the average cosine similarity across all 2d(4d 1) pairs and report in the figure.
24

---

Published as a conference paper at ICLR 2024
While 4d d-dimensional vectors have to be linearly dependent, they indeed almost orthog- onal (i.e., cossim (#1, (2) < 1) throughout the training process, as shown below. Fig. 11 we show cosine similiarity over the entire training process of Pythia models of different sizes.
Fig. 12 further checks the training curve early training stages, since Pythia checkpoints are more densely sampled around early training stages, i.c., "steps O (initialization). 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1000, and then every 1,000 subsequent steps" (Biderman et al. 2023). Finally, for models whose intermediate checkpoints are not available, we show the cosine similarity in the publicly released pre- trained models (Fig. 13
Pythie-70M: Laye: Dim: 51.2
Pydra-160% 1ytha-160M: Layer: 12, Drn:
Pythis-l 24, Dim: 1034
0.035
0,036
Chatre 230 Smieet
64038
Crunee S1081
15
5:039
0.022
0.034
6:632
0.030
0.0
0.033
0,024
Pythi- lnyer:
Pythi- Layen
Pythis-2 Laver Dimm
0.00
14
0.0375
0.0330
0.6390
0.0325
0.030
0.28
0.044
0.0279
0.0250
0.0229
0.0200
0.0135
6:8375
0.015
Mi-bbrch
Figure 11: Orthogonality embeddings of MLP in LLMs during the whole training process.
1ytha-77MM; Laper: B.
Pythic 160M:
1ydta-410M Layer: 1024
0 0.69200
15
2.3580
000014
Colvee
3.00112
003335
3.3576
0.00911
030337
2,03344
355
140
125
100
125
Minbatcn
Pi BiL Du 2018
Pythi- ABI Layer 24: Dlmm: 2048
Pythia Layer: Dir:: 2560
6.077744
0.017704
6.017782
0.017343
K017766
8017758
0.11755
3.017755
6013754
2.177554
2017753
0003095
0403824
Mintacch
Figure 12: Orthogonality of embeddings of MLP in LLMs during the early training stage.
B.2 ATTENTION ENTROPY FOR ENCODER-DECODRS MODELS
We also measure how allention entropy, as well as stable rank of the in-projection (or lower) matrix in MLP, changes over time for encoder-decoder models like BERT, as shown Fig. 14 The behavior very similar to the decoder-only case (Fig. 7 further verilying our theoretical findings.
25

---

Published as a conference paper at ICLR 2024
Final Checkpoint
0.5
BERT-Base
OPT-6.7B
0.4
LLAMA-2-78
Cosine Similarity
VIT-Huge
0.3
0.1
0.0
5
10
15
20
25
30
Layer Index
Figure Orthogonality measures in model architectures (BERT-Base OPT-6.7B, LLMMA-2-7B, ViT- VugeH with only final checkpoint available.
BERT, Layers: 12, loss: 1.517
BERT, Layers: 12, val_loss: 1.517
10
10
0.08
Attention Entropy
Stable Rank
0.04
3-
2
2-
0.00
500
1000
1500
2000
500
1000
1500
200D
Minibatch
Minibatch
Figure 14: (Left) Attention entropy of BERT; (Right) Stable rank in BERT.
26


Reviews:
[{"review_id": "KBO7SLYRwy", "reviewer": "Reviewer_P4nM", "summary": "The paper proposes a mathematical framework called JoMA to understand the training dynamics of multilayer Transformers. JoMA integrates out the self-attention layer and produces a modified dynamics of MLP layers only. JoMA reveals the training dynamics under linear or non-linear attention, as well as how the attention sparsity change over time, and how they depend on the input distribution.. The paper also verifies its theoretical findings with experiments on real-world data.", "questions": "1. This paper analyze the training dynamics of multi-layer transformer and reveals how the attention is learned through time (first learning attention on tokens with most co-occurrence, and then expanding attention on other tokens). One question is, how is this dynamics dependent on model size? As we see from empirical results in the literature, larger model tends to learn attention that involves more complicated reasoning instead of simple co-occurrence. Can this factor be reflected by the analysis in the paper? How does the training dynamics (and entropy of attention) change when model size changes?\n\n2. Is the proposed JoMA also applicable to vision transformers? Since the input to vision transformer is different from text transformer in that the tokens are continuous instead of discrete and several assumptions in the paper may not hold. For example, the paper assumes tokens are orthogonal to each other, but for vision input lots of input tokens are very similar.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "By integrating out the self-attention layer and analyzing only the modified dynamics of MLP layers, JoMA provides a fresh perspective on how these models learn and adapt during training. The paper demonstrates rigorous theoretical analysis on training dynamics, especially on how attention evolves over time. An interesting observation from the paper is that MLP and attention shares the same dynamics, which leads to several intriguing conclusions, for example, the rank of lower mlp layers first goes down and then bounces up. The mathematical derivations are presented clearly, making it accessible to researchers with different backgrounds.", "weaknesses": "1. One assumption in the paper is that the embedding vectors of all tokens are orthogonal to each other. Does this need to hold for token embedding in every layer? Have the authors verified this assumption is valid in real cases? How much does this assumption affect the following conclusions?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ULCa8Zdr0v", "reviewer": "Reviewer_h5HX", "summary": "This paper named JoMA, explores to explain the training procedure of multi-layer transformers. JoMA delivers a conclusion that transformers learn salient tokens at low layers while learning less salient tokens at high layers when the nonlinear activations involve training. To build a mathematical framework to show the learning mechanism of multi-layer transformer, JoMA proposes joint dynamics of self-attention and MLP.\n\nCompared with former works that focus on shallow transformer networks, linear activation or local gradient steps, JoMA builds a unified mathematical framework to characterize the learning mechanism of multi-layer transformers. Analyzing deep transformer networks is more challenging.", "questions": "-", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "JoMA is able to analyze deep and sophisticated transformer networks.\n\nIt's interesting to use dynamics to unify the MLP, self-attention, and non-linear activation.\n\nThe derived conclusion of how self-attention learns hierarchical data distribution is meaningful and can guide the design of follow-up transformer-based network.", "weaknesses": "It would be better to discuss the limitation of JoMA.\n\nIt's not clear why chose OPT and Pythis for verification. Can the findings in this work coincide with BERT-style models?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZYPLhkltry", "reviewer": "Reviewer_cMvf", "summary": "This paper proposes a novel mathematical framework called JoMA to understand the training procedure of multilayer Transformer architectures. By integrating out the self-attention layer in Transformers, the authors focus on the dynamics of MLP layers. The framework removes unrealistic assumptions and provides insights into the behavior of attention mechanisms. The authors leverage JoMA to explain how tokens form hierarchies in multilayer Transformers using a latent hierarchical generative model. Experimental results on real-world datasets and pre-trained models validate the theoretical findings.", "questions": "Please see Weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed JoMA framework offers a new perspective on understanding the training dynamics of multilayer Transformers. By integrating out the self-attention layer, the authors provide valuable insights into the behavior of MLP layers.\n\n2. This paper presents theoretical findings that explain how tokens are combined to form hierarchies in multilayer Transformers. These findings are supported by experiments on real-world datasets and pre-trained models, enhancing the credibility of the results.\n\n3. This paper is well-written and structured, making it easy to follow the proposed framework and understand the theoretical explanations. The inclusion of references to related works in the field strengthens the paper's contribution.", "weaknesses": "1. Limited practical implications. This paper focuses on theoretical analysis and understanding the training dynamics of multilayer Transformers. It would be valuable to discuss potential practical applications or implications of the proposed framework in real-world scenarios.  \n\n2 Several lines of work around the learning preference of DNNs need to discussed [cite1-3].\n\n[cite1] Devansh Arpit, Stanisław Jastrz˛ebski, Nicolas Ballas, David Krueger, Emmanuel Bengio,\nMaxinder S Kanwal, Tegan Maharaj, Asja Fischer, Aaron Courville, Yoshua Bengio, et al. A\ncloser look at memorization in deep networks. In International conference on machine learning,\npages 233–242. PMLR, 2017  \n[cite2] Karttikeya Mangalam and Vinay Uday Prabhu. Do deep neural networks learn shallow learnable\nexamples first? 2019.  \n[cite3] Huh M, Mobahi H, Zhang R, et al. The low-rank simplicity bias in deep networks[J]. arXiv:2103.10427, 2021.\n\n3  I am not sure how the insights from this paper could help us train/design better DNNs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rcAGDvJMNL", "reviewer": "Reviewer_5uGr", "summary": "The paper introduces Joint MLP/Attention (JoMA) dynamics, a mathematical framework designed to understand the training process of multilayer Transformer architectures. This is accomplished by integrating out the self-attention layer in Transformers, resulting in a modified dynamics that focuses solely on the MLP layers. JoMA addresses limitations present in previous analyses by removing unrealistic assumptions, such as the lack of residual connections.\n\nThe authors predict that attention in JoMA initially becomes sparse to learn salient tokens and then transitions to dense attention to capture less salient tokens when nonlinear activations are incorporated. In the case of linear activations, the predictions align with existing literature. Additionally, JoMA is leveraged to provide qualitative explanations of how tokens combine to form hierarchies in multilayer Transformers, specifically in scenarios where input tokens are generated by a latent hierarchical generative model.\n\nTo validate their theoretical findings, the authors conduct experiments using real-world datasets (Wikitext2/Wikitext103) and various pre-trained models (OPT, Pythia). The experimental results support their claims and contribute to the credibility of their research.", "questions": "1. Why is Figure 1(b) labeled as the \"problem setting\"?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ The paper introduces the JoMA dynamics, which is a new mathematical framework specifically designed to understand the training procedure of multilayer Transformer architectures.\n\n+ The authors highlight that previous analyses of similar models often make unrealistic assumptions, such as the lack of residual connections. In contrast, the JoMA framework removes these assumptions, making it more accurate and realistic in capturing the training dynamics of multilayer Transformers.", "weaknesses": "- Although the paper mentions experimental validation using real-world datasets and pre-trained models, it does not provide detailed quantitative evaluation metrics or comparisons against existing approaches. Including quantitative analysis would enhance the robustness and comprehensiveness of the findings.\n\n- The paper does not delve into specific implementation details of the JoMA framework, such as hyperparameter choices or training procedures. Providing more information on these aspects would enable researchers to replicate the experiments and further validate the proposed approach.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)