The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit
Authors: Lorenzo Noci
Chuning Li
Mufan Bill Li
Bobby He
Thomas Hofmann
Chris J. Maddison
Daniel M. Roy
Venue: NeurIPS 2023 poster (2023)
Keywords: Deep Learning Theory
Covariance SDE
Attention Mechanism
Infinite-Depth-and-Width
Scaling Limit
TL;DR: We study the proportional infinite-depth-and-width limit of Transformers at initialization in order to devise a modified attention mechanism that avoids known degeneracy issues.

Abstract: In deep learning theory, the covariance matrix of the representations serves as a
proxy to examine the network’s trainability. Motivated by the success of Transform-
ers, we study the covariance matrix of a modified Softmax-based attention model
with skip connections in the proportional limit of infinite-depth-and-width. We
show that at initialization the limiting distribution can be described by a stochastic
differential equation (SDE) indexed by the depth-to-width ratio. To achieve a
well-defined stochastic limit, the Transformer’s attention mechanism is modified
by centering the Softmax output at identity, and scaling the Softmax logits by a
width-dependent temperature parameter. We examine the stability of the network
through the corresponding SDE, showing how the scale of both the drift and diffu-
sion can be elegantly controlled with the aid of residual connections. The existence
of a stable SDE implies that the covariance structure is well-behaved, even for very
large depth and width, thus preventing the notorious issues of rank degeneracy
in deep attention models. Finally, we show, through simulations, that the SDE
provides a surprisingly good description of the corresponding finite-size model.
We coin the name shaped Transformer for these architectural modifications.

The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit
Lorenzo Noci*t
Chuning Li
Mufan (Bill) Li
Bobby Hel
Thomas Hofmann
Chris Maddison
Daniel M. Roy
Abstract
In deep learning theory, the covariance matrix of the representations serves as a proxy examine the network's 's trainability. Motivated by the success of Transform- ers, we study the covariance matrix of modified Softmax-base attention model with skip connections in the proportional limit of infinite-depth-ande width We show that at initialization the limiting distribution can be described by a stochastic differential equation (SDE) indexed by the depth-to-widih ratio. To achieve well-defined stochastic limit, the Transformer's attention mechanism is modified by centering the Softmax output identity, and scaling the Softmax logits by width-dependent temperature parameter. We examine the stability the network through the corresponding SDE, showing how the scale of both the drift and diffu- sion can be elegantly controlled with the aid of residual connections. The existence of stable SDE implies that the covariance structure is well-behaved, even for very large depth and width, thus preventing the notorious issues of rank degeneracy in deep attention models. Finally, we show, through simulations, that the SDE provides a surprisingly good description of the corresponding finite-size model. We coin the name shaped Transformer for these architectural modifications.
Introduction
Pre-trained large language models have experienced a remarkable increase in popularity due to their cerily human-like ability to puzzle through complex reasoning tasks, solve coding challenges, and produce pages of logically sound text [1]. Arguably, the Transformer is the foundation of these successes [2]. Recent research has found evidence for scaling laws, linking the performance of these architectures to their parameter counts and the quantity of training data, fueling the desire to train deeper and wider models on ever larger datasets in order to unlock new levels of performance [3-7].
Bundled with the increased expressivity of deep architectures, however, increased numerical instability, both in the forward pass and gradients, which hinders training. One of the clearest examples of instability is the so-called rank collapse phenomenon [8, 9]- the observation that, in Softmax-based attention models, the network's representation of different tokens tend to perfectly align at large depth. The resulting poorly conditioned covariance and correlation between tokens leads exploding and/or vanishing gradients at initialization, disrupting gradient updates of the affected parameters. This situation violates a well-known guiding principle from the literature of deep signal propagation: a stable covariance is necessary condition for stable training [10-15]. In fact,
"gqual contribution. Correspondence to: lorenzo. noci@inf. ettz. ch, chuning Li@miil.... utoronto. ca. muf iafa..... 1i0mall......mmmmmm utoronto.co
TETH Zurich
I University Toronto Vector Institute
$ University of Oxford
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

1
1-
Mean Correattion pya
0.8
0.75
Density
0.6
0.50
0.25
0.2
0-
50
100
150
1
-0.5
0.5
Layer
Correlation Pa
Shaped (Ours) Unshaped Unshaped+PPr-
-SDE (Ours) Shaped
Figure 1: Our shaped Transformer prevents token representations from becoming perfectly aligned, i.e. rank collapse. Left: mean correlation p3 of Transformers (Eq. 11) with and without shaped attention (Eq. 9) and Pre-LN [48]. Right: kernel density estimate and histogram of correlations from covariance SDE in Theorem 4.2 and shaped attention NN. Here we note correlation converging to 1 implies poorly conditioned covariance matrix. Simulated with 200, d = 150, 7 = 1/ 1/v8, Po 0.2, SDE step size 0.01, and 212 samples.
the instability of Transformers is evident when considering the critical role of hyperparameter tuning and the judicious use of normalization layers. In this work, we study Transformers in a novel infinite limit, rectify sources of instability with a novel modification, and derive the SDEs characterizing the covariance and output distribution.
Scaling limits have been used successfully to provide guidance on architecture [16-18] and tuning hyperparameters settings [19]. Our work represents a contribution in this direction. The ability to use such limits to diagnose instabilities depends on their tractability and faithfulness to real-world (finite) networks. In this regard, not all limits are created equal. In particular, the faithfulness of scaling limits depends critically how other parameters are scaled with width. One of the simplest (and thus most popular) limits to work with the "NTK" limit [20-24] treats the depth of the network as fixed. As result, at initialization, this limit does not accumulate sufficient random fluctuations over the depth of the network, leading to deterministic covariance matrices that do not agree with those of standard (finite) networks. Such networks have another defect: they are incapable of learning features in the limit [25]. Various other limits have been studied, towards identifying tractable yet faithful models of initialization and/or training. These include mean field limits [26-29] and the perturbative regime [30-36].
This work operates in a relatively new regime- the proportional infinite depth-and-wind limit where depth d and width n diverge as the ratio d/n tends to a positive constant. This limit, first analyzed by Hanin and Nica [37], has been the recent subject of study in the context of neural network [38-11, 18]. A related line of work also studied the Lyapunov exponent for products of random matrices [42-45]. This regime retains the network's stochasticity and, initialization, has been shown to closely resemble the behaviour of finite architectures, yet still yield relatively simple limiting description, expressible in terms stochastic differential equations [40, 18]. In this work, we fully characterize the initial output distributions of a network with skip connections and Softmax- attention mechanisms, in the proportional Shinite-depth-amg aaruhh-andwwar width limit.
Inspired by the idea of shaping activation functions [16-18, 46], our theoretical approach finds an adequately modified attention mechanism via its SDE limit. Our modification involves making the attention matrix closer to the identity, and appropriately choosing the temperature parameter which re-scales the logits of the Softmax. Similar to shaping activation functions, the temperature scaling we devise lincarizes and reduces the saturation of the Softmax, a known source of training instability in Transformers [47]. In order to model the feedforward layer a Transformer's block, we extend existing results [18] to derive an SDE for the proportional limit of shaped- feedforward layti-layer perceptrons (MIPs) with skip connections. Combined, we fully characterize the output distribution of Transformer with shaped non-lincaritics (Corollary 4.3).

---

Shaped
1.5
Shaped Pre-LN Post-LN
Gradient Norm Value Weight
Pre-LN
Gradent Nom of Query Weights
1.4
20
1.2
1.9
15
0.8
10.
0.6
0.4
5.
0.2
0.0
20
40
60
80
1.00
120
140
100
120
140
Layer Z
Layer
(a) Value Weights WY
(b) Query Weights W
function of depth, with and without shaped attention. The architecture is the same as in Figure but with autoregressive causal masking, and the task next-token prediction code data. Left: Value weights WY for shaped attention, standard Pre-LN, and the original Post-LN block [2]. Right: the same gradient norm plot but for Query weights W.O. We find that shaping the attention mechanism successfully prevents gradients from vanishing, while unshaped Transformers suffer from rapidly vanishing gradients. Interestingly, only the Post-LN query gradients vanish, but value gradients are stable across depths, which is consistent with the findings of Noci et al. [9]. On the other hand, shaped attention has stable gradients for both parameters inside and outside the Softmax nonlinearity.
Figure 2: Comparing gradients norms initialization for different parameters
Notably, our modification successfully prevents a poorly conditioned covariance matrix, whereas the vanilla Softmax-bai based attention model without LayerNorm [49] fails in this regard, and the corresponding Pre-LN architecture provides only marginal improvements (see Figure 1). Given that our modification is inspired by previous work on shaping activation functions, we coin the terms shaped attention for the proposed attention mechanism and shaped Transformer for the overall architecture that includes the MLP block and residual connections. Through simulations (e.g.. Figure 1), we show that the limiting neural covariance SDE approximates the distribution of finite- size Transformers with shaped attention mechanism surprisingly well. We also provide preliminary training experiments for our proposed shaped attention architecture on standard language modeling tasks, demonstrating the feasibility of the architecture in practice (see Section 5 and Appendix D).
In summary, our contributions are as follows:
1. We study the effect of skip connections in the proportional limit, showing that under precise relation between the scaling parameters of the shortcut and residual branches, the feature covariance converges to the solution of weighted version of the neural covariance SDE for MLPs (Theorem 3.2). The dependence on the depthhto-wiitt ratio implies the existence of stable non-commutative limit for residual networks, complementing the commutative limit studied in Hayou and Yang [50].
2.
We propose shaped attention, where we modify the Softmax- attention mechanism to be perturbation the identity. We demonstrate that shaped attention successfully prevents the degeneracy of correlation in contrast to existing Transformer architectures (Figure l). The enhanced stability in the forward pass is reflected in the gradients, which are also stable with depth, empirically show in Figure 2.
3. For the proposed shaped attention architecture, we derive the neural covariance SDE character- izing the initial distribution in the proportional limit (Thcorem 4.2). Consequently, we provide the first characterization of Transformer-type architectures, i.c. the shaped Transformer, in the large depth-am and -ad--anddd---e width regime (Corollary 4.3).
4. We provide simulations to validate the theory and interpret the effects of network hyperpara- maters on the covariance matrix of the shaped Transformer. Specifically, we study finite time stability of the SDE and provide explicit guidance on hyperparameters prevent numerical instability.

---

The paper is organized follows: In Section 2, we provide the basic setup and some background on existing results. In Section 3, we generalize the SDE results of M. Li et al. [18] to include skip connections. This serves as model to understand the effect of skip connections in isolation from the attention model. In Section 4, we present our main result, first pinpointing the origins of instability in the Softmax, then showing how the modifications underlying shaped attention allow us derive a non-trivial SDE limit. Finally, in Section 5, we discuss the implications of our results and some future directions. Proofs of all theorems and additional experiments are deferred to the Appendix.
2 Background
Setup. Let X₂ E R be the data matrix representing a sequence of m tokens embedded in dimensions layer E [d], where d is the depth of the network. We elide the explicit dependence on é when it is clear from the context, and use superscript Greek letters to indicate specific tokens' representations, for instance x² E R is the a-th row of Xe. We consider the following attention model with residual connections:
Vn WY
where y, E [0. 1] are parameters that control the strength of the shorteut and residual branch, respectively, W.Y E [Rnx" is the weight matrix of the values, and A, E mm is the attention matrix. We consider Softmax-based scaled dot-product attention, where At has the form:
Ag = Softmax (+x²
1 We K.T.T X
Vn
Vn
where the Softmax is applied row-wise, W W.K Rux, are additional random weights, and is a temperature parameter, which controls the entropy of the distribution. Here we let all the weight matrices W.O, W.K. W. have (0, 1)-iid entries. In the case where 7, 1, with the application of LayerNorm on the residual branch [48], and with T = VTg, we recover the attention block of the vanilla "Pre-LN" Transformer architecture [2]. Here we note that we pull the conventional n n factor outside of the weight matrices, which preserves the forward pass, and yields equivalent training dynamics up to a reparameterization of the learning rate [25]. In this work, we consider unnormalized architectures, and control the variance propagation with the condition X² + +2² = 1 [40]. We are interested in studying the so-called neural covariance for the attention model (Eq. 1) in the proportional limit.
Neural Covariance. In deep learning theory, researchers have long sought to understand how networks internally represent different inputs and how different architectural choices affect these representations. The approach followed by work on signal propagation has been to study how the relative aligmment of different inputs evolves across the network, as measured by the neural covariance v (uf zgg) (or قمم == (Vaay - if interested only in the correlation). At initialization, characterizations of this covariance structure have been exploited to infer important properties of neural networks [10, 11]. As an example, in the sequential infinite-width--ie width-then--eye
limit, the correlation Pd of MLPs is known to converge to fixed point independent of the input [6]. In this regime, the model is not able to discriminate different data points, which severely hinders training, as the gradient step for the deep layers is taken in the same direction regardless of the input. In the context of Softmax- matrix X₂ loses rank doubly exponentially fast with depth, and Noci al. [9] showed how this leads to vanishing gradients of the queries and keys parameters, thus further highlighting how the stability of forward and backward passes are deeply entangled (see also Figure 2).
attention models, Dong et al. [8] proved that the feature
Stabilizing the Effect of Non-Linear Layers. Central to the issue of degeneracy of the neural covariance are commonly used non-linear activation functions that severely deviate from the identity, The recent line of work of Deep Kernel Shaping (DKS) [16-18] addresses the issue by considering the cumulative amount of non-linearity throughout layers, and shaping the activation function by making it closer to the identity map. Inspired by this line of work, B. He et al. [46] devise an initialization for Transformers that avoid the rank collapse problem without the aid of skip connections or LayerNorm.
In an alternative approach, the line of work behind Stable ResNets [50-53] considers scaling the residual branches by y /cptt, and postulates this scaling is sufficient to stabilize the neural

---

covariance with minimal assumptions the activation function. Noci al. [9] adopts this scaling to give precise formulas on the expected covariance a Transformer at initialization. In this work, we consider 7 constant width and depth, and derive a complementary limiting result.
Width Limit. In the context of feed-forward MLPs, the output distribution with respect to single input was studied in [37, 41], where it was shown that for the ReLU nonlinearity, the norm of the activations Voo converges to a log-normal random variable. To resolve the degeneracy of covariance and provide characterization of output distributions for multiple inputs, M. Li et al. [18] shapes the ReLU by setting its slope 1/v //wwtth-wwwwwy from lincarity. In the proportional limit, the effect of the non-linearity accumulates over the layers, and the covariance matrix V = [V las converges weakly to the solution of the SDE
The Proportional Safinite-eepth-ae
JaB
12/2 dV = BRALL (V)) dt + 2im (V)) dB,,
(3)
where the formulae for coefficients UROLU, Eiin can be found in Theorem 3.2.
We note that the output neuron distributions are directly recovered as a conditional Gaussian with covariance Vr for T n²' [20-22]. For example, the i-th output X conditioned on Va are asymptotically iid. W VT) as d, - 30. The reader referred to Appendix for more technical background on the covariance SDE and the convergence result.
in a similar spirit as the neural network Gaussian process (NNGP) results
While the existing results are limited initialization, we remind the reader that this necessary step before we can study training dynamics. In particular, the NNGP techniques developed for widhh networks at initialization were directly used to study the training dynamics in the same limit [24, 54]. We will provide further discussions on this topic in Section 5.
3 Warm-Up: a Neural Covariance SDE for ResNets
To understand the effect of skip connections, it is helpful to look at simplified model composed of shaped ReLU-activated layer and skip connections:
Xe+1 = AX, + to, (xe
Wio vn
V
(4)
n
where (x(x)= == s_ max (x, 0) + s_ min(z, 0) is the shaped ReLU with slopes 82 := 1 + c 1/2
for some constants c_, c_ E R We assume i.idd weights (WP") ).y (WPoa) )i id N 1), and
c- -.=Ea,a (g)² for g NO0 1) a constant that ensures that the activations are normalized [55]. Notice that this is the form of the feedforward layer in a Transformer [2].
We will next define the notion of convergence for covariance matrices and state our first main result. We refer the reader to Appendix A for more precise details on the Skorohod topology.
Definition 3.1 (Convergence of Covariance). Let X E Rmmxn be the L-th layer matrix of representa- tions, and define the feature covariance as Vz = n X Let v.(m) V(tml E Rm((++1/22 be the the continuous time interpolation of the upper triangular entries vector, We say the covariance (n) converges to V, if in the limit as n, d + cc, {VeJeeuu weakly in the Skorohod topology.
d + T, the process [V((") Jeeoorr converges to
Theorem 3.2. Let X be the hidden layers of ResNet defined in Eq. 4 with X2 + +² 1, where both X and y do not depend on d, n. Then the feature covariance Ve converges to the solution of the following SDE (in the sense of Definition 3.1)
1 dV:; = (Vi) dt + Enes (V.) dB,, Vo = -XoX
(5)
((V)=>?buel))) where (V) = bbb (V) = (2(x(p=0) V88 asa with pri (V 1/2 and
v(p) (c+ cc-²
( 11-p² parccos
(6)
2x
furthermore, Ene (V) = =2: 2in =
yaw 18 Ja59.55w
5

---

Residual 1/v depth
SDF
NN
Correlation 95th percentlle
MLP
0.8
0.5
y=
Density
0.6
NN
0.
1/ d d-scaling
0.0
-05
0.6
Correlation P8
Residual
Figure 3: Left: Kernel density estimates correlation Pa for various values the residual strength parameter 7, In particular, 7 I recovers shaped-ReLU MLP without skip connections, and = 1/V/d is the setting studied in Noci et al. [9] and Hayou and Yang [50]. Solid lines represent finite networks, while our SDE simulations are presented dashed lines. Right: 95th percentile of the absolute value of the correlation distribution
.9
function of Note reducing 7 reduces the concentration around pad = 1, and our SDE reliably approximates finite size networks. Simulated with n = 300, d = 100, Po
.² = 0.2, c+ = 0, c- = -1, and 213 samples.
Notice how the limiting SDE closely resembles the MIP case (Eq. 3), which is recovered exactly when ry = 1. The only difference is the extra 2 factor, which comes from the fact that in our definition each layer has effectively two times the number of weight matrices than the standard formulation for MLPs. As the drift depends solely on the nonlinearity, and the diffusion depends soley on the random weights, only the diffusion variance is doubled. The residual branch parameter y < 1 dampens both the drift and the variance of the Brownian motion by > thus it can be interpreted as a time change. In other words, the effect of y at initialization equivalent to reducing depth-to depth-to-wine to-width ratio, inline with existing intuitions that ResNets have a lower dffectivv-depth [56]. To visualize the stabilizing effect of y on the distribution, Figure (right) we plot the 95th percentile correlation as function of y. The increasing trend indicates larger probability of perfect alignment between two tokens. In Figure: (left) we plot the densities of both the residual SDE and the corresponding residual network for various values of y. Notice how the samples from the SDE well-approximates the histogram of finite network.
A Stable Non-Commutative Limit. Our results complement those of Hayou and Yang [50], where the authors have shown that for a similar ResNet under the parameter scaling X = 1, = 1/ d. the depth and width limits commute. More precisely, the covariance yo8 converges to the same limit regardless the order with respect which the limit is taken or the depth-to-width ratio. Furthermore, the limit is deterministic, and can be described by an ordinary differential equation (ODE). Intuitively, the convergence to a deterministic quantity occurs because ? = 1/vd suppresses the random fluctuations enough to vanish the limit. On the other hand, our results show that for y constant in d, the random fluctuations (Eq. 3), hence they do not vanish in the simultaneous limit. The most notable difference that our limiting regime is non-commuuttiie as it depends on the depth-to-wide ratio of the network. We remark that both regimes prevents degeneracy of covariance for residual architectures, forming two theories that complement each other.
the right order of O(n -1/2)
in the MLP case
4 Neural Covariance SDE for Softmax-Based Attention
4.1 Unshaped Attention and Its Taylor Expansion
A central piece to the neural covariance SDE theory for MLPs [18] is identifying the exact scaling of shaped activation functions. particular, the effect of the activations on the covariance Markov chain V must be on the same order as the random weights in an MLP, thus forming an approximate

---

Fuler-dissretizzaton
2 Ve+1 V22 6(v)) +
+ On Onn -/2),
Vn
where b, Σ. are deterministic coefficients, and & are random vectors with zero mean and identity covariance. From here onwards, moments bounded by universal constants (i.c. independent of rt). Since the update can be interpreted as discretization with step size n n-1 naturally the Markov chain converges to an SDE. We again note that a stable SDE implies stable covariance structure for finite size networks.
O(n-P) to denote random variable Z such that n2 Z has all
To achieve the goal for modified attention mechanisms, consider similar approach as M. Li et al. [18] for smooth activation functions, and Taylor expand the Softmax function in terms of a large temperature parameter To this end, let Yg to be the matrix of dot-products between queries, and keys, i.c. Y Xt 1 We 1 W.K.T X² Vn Vn
More specifically, given a row y*t E Rlxm of the logits Yg E m.mm we can Taylor expand the row-wise Softmax in terms of
Coftmax(T-1yg)) = 1T (0-
(Tm (yno3) +0(( (8) (10- - ya)z
m Tm
2
where ndah ;= the Eag0 and (1 the vector with squared entries of yt, and 1 E R the (columm) vector of ones. We note in practice T is often set to v Tik, which is often quite large and allows for asymptotic analysis [9].
We observe that the zero-th order term mn T is independent of 7. Considering the form of the attention block as AgX, / W/. this yields an update that is no longer a small perturbation of Ve. regardless of how chosen. Therefore, form Markov chain like Eq. 7, actually require A to be approximately the identity matrix.
4.2 Shaped Attention
To shape the Doftmax-attention mechanism as perturbation of the identity matrix, we propose the following modifications which we call the shaped attention $
A = I Softmax(r- 111 T= TOy
nink... T
(9)
The shaped attention presents three modifications to the Softmax attention in Eq. 2. Firstly, the zero-order term m m-111 of the Taylor expansion (Eq. 8) is removed as it causes a non-infinitesimal drift in the Markov Chain that ultimately leads instability in the covariance (see Section 4.1). Secondly, we also observe that when 7 is very large, the centered Softmax is a perturbation around zero. To recover an approximate Euler-update like Eq. 7, we simply add back the identity matrix. By biasing the attention matrix towards the identity, we encourage each token to self-attend. This type of modification was also previously considered by He et al. [46]. Finally, the Softmax's temperature is chosen to scale as /om/ for some constant
O, which guarantees non-degenerate limit as (d, rt) + co (Theorem 4.2). Note that the extra vn term is a departure from the standard parameterization.
In Figure 4, we show how removing any of the proposed changes individually alters the neural covariance structure, which becomes degenerate for large depths, while the proposed modifications remain stable. We stress that here for simplicity we focus on attention without masking. Shaped attention be extended to include masking (e.g. casual masking) by centering each i-th row of the Softmax matrix by a different factor 1 1/m,, where m. is the number of un-masked tokens in the i-th row.
'In principle, it could be possible to have ddse-tto- Softmax matrix when the logits are large. However, this regime also corresponds to very saturated Softmax, thus making training unstable [57]. As result, we will avoid this direction in this work.
7

---

10
6.6
shaped
Mean covariance 1/20
107
mi venter
Mean correlation
6.7
107
id, center
Id,
0.6
3017
only o'
only center
0.4
1011
only id
0.3
10
D0 122 13-4
Figure 4: Mean correlation (left) and covariance (right) (in absolute value) under various interventions on the proposed shaped attention. In particular, we remove either one or two of the three modifications from the shaped attention in Eq. 9. For instance "7² Ti center" indicates that proposed temperature, and we center by m but we do not add the identity matrix, while in "only id" we add the identity matrix but the covariance remains unstable due to incorrect scaling. Due to exploding covariance, we choose not include the cases "id, 7² nn, and "only id" in the correlation plot (but only in the covariance plot). Simulated with n 300, d 150, Po 0.2, 7 = 1/v2 and 23 samples.
the
/ne and do not center. We note in this "only id" case,
4.3
Main Result Neural Covariance SDEs for Shaped Attention Models and Shaped Transformers
Before we state our main results, we will first define a weakened notion of convergence, which is required whenever the drift and covariance coefficients not Lipschitz. This was also required for the case of shaped MLPs with smooth activations [18].
Definition 4.1 (Local Convergence). We say the covariance v(x) converges locally to V if the stopped process { (V 1:2 converges (VIT,,120 in the sense of Definition 3.1 for all stopping times of the form T, inf(t 0:
||V|| 2 r) with
Let the covariance with respect to the average token be defined as Vaa == m Er= Vay and the average trace be V = m- >x=1 Voe, We will need to compute a couple of important moments from the Taylor expansion terms of the Softmax (Lemma C.2)
S0b.Bw 1 1B('' -aann-jj) (Y - -y vo8 (Vów V82 VWB + 122).
-V) (10)
Vaa (v - 2V + +2 Sg nn 'E (yrd - アーソ (Ya03 )m²)
We are now ready to state our main result.
Theorem 4.2. Let Xe be the hidden layers of a residual attention network defined in Eq. I with shaped attention in Eq. parameters X² + 72 1 and Tuy/nnk, where X, % To all do not depend on d, n. Then the feature covariance V, converges locally to the solution of the following SDE (in the sense of Definition 4.1)
dV = b(V.dtt LVV
=o= -Xo00
where the drift has the following form
b(V) = 22
VUR Seng gov,Bn
1
vau S
gav
Vau cit
2 1,6==
2m ==1
550
the diffusion coefficient is defined by E(V)= = 72 22) T 2 18.655 and
A88,6w
1
Voa grBmww yan ywr Gan.vv vr gomwnn + van gqundn
m M,K=1
8

---

The drift depends the shaped attention mechanism through Si ST and S the moments of the first and second order terms of the Softmax's s Taylor expansion. On the other hand, the diffusion term depends on the attention solely through St, present the additional term A The presence of Ana8,ow is an intriguing difference compared to shaped ReLU networks, where the diffusion is not affected by the activation function. Both components of the SDE depend on averages over the tokens, reflecting the mixing property of the self-attention mechanism, in which every pair of tokens compared through dot products to form the attention weights. Finally, notice how the residual branch parameter y2 has a dampening effect on the scale of both the drift and the diffusion in similar way as in fully-connected residual network.
We are now ready to introduce the full shaped Transformer architecture, where we combine the attention and residual layers:
= =X, + 7AxX / WY X(+1 = AZ + to,
Ze - W Vn
WTT
(11)
where Ag the shaped attention defined by Eq. 9. We note that covariance SDE handle stacking of different layer types very conveniently by simply adding the drift and covariance of the diffusion coefficients, which we summarize in the Corollary below.
Corollary 4.3 (Shaped Transformer Covariance SDE). Let Xe be the hidden layers of a shaped transformer defined in Eq. 11 with parameters X² + 72 = and T Toy/nng, where X. all do not depend on d, n. Then the feature covariance Ve converges locally to the solution of the following SDE (in the sense of Definition 4.1)
dV2 [b(V;) + (ver(V))) +t [E(V)) EEm, dB,,
(12)
where the coefficients are defined in Theorem 3.2 and Theorem 4.2.
4.4 On Finite Time Stability of the SDE and Shaped Attention Networks
Although we did not observe numerical instability in majority of our simulations of the shaped attention networks and the corresponding SDE, we did observe that the drift component b(V)) in Theorem 4.2 is cubic in the entries of Vi. Whenever the drift is not Lipschitz in this case, we do not have general guarantees for the existence of a solution for all time (see the Feller test for explosions [58, Theorem 5.5.291). In fact, MLPs with smooth activations also yield non-Lipschitz drift coefficients as seen M. et al. [18].
However, locally Lipschitz coefficients are sufficient to guarantee the existence of local solutions, in the sense of up to stopping time [59, Proposition 6.9]. Not only does this fact help establish precise notion of convergence (Definition 4.1), we can also study the practical implications of this for finite sized attention networks. More specifically, we can inspect the effect of architectural changes to stopping time.
To demonstrate the potential numerical instabilities, we had to choose adversarial set of parameters: in particular, an unrealistically large norm (approx. 10yn) for the initial tokens Xp. which enlarges the eigenvalues of Vo to the order of 100. Given these initial conditions and large residual connection weight 7, we were able to consistently generate numerically unstable behaviour in shaped attention networks (see Figure 5 (left)).
That being said, it is very straight forward to stabilize the network by tweaking parameters such as 7, To and the depth-to-width ratio of the network. We demonstrate the effect of tuning 7 on both sample trajectories of the maximum eigenvalue intuitively expect, tuning smaller will delay the time scale of numerical instabilities, hence allowing for larger depth networks to remain stable.
V2 and the stopping time in Figure 5, As we may
5 Discussion
Architecture Design and Hyperparameter Tuning. Previous work have demonstrated the practical impact scaling limits can have on designing activation functions [16, 17] and tuning hyperparameters [19]. We follow this line of motivations and proposed novel attention mechanism, which successfully stabilizes the covariance structure in arbitrarily deep Transformers (e.g. Figure 1). The natural next

---

Y= 0.735
mecian 10th percentile
Stopping Time (Capped at 11
Y= 0.214
Maximum Eigervalue
10th
2011-
2g1r
zg1t
107
10"
10"
105-
107
107
107
107
25
200
layer
V
Figure 5: Left: Trajectories of the maximum eigenvalue of the covariance matrix in a shaped attention network, with adversarially large initial condition. Right: Stopping time of the shaped attention neural network, capped at I. Stopping time is defined as t* = d' /n with d' the maximum depth beyond which of the eigenvalues of the covariance matrix exceeds 104 or drops below 10-4. Simulated with n = d = 200, to 1, and 100 samples used for median and 10th percentile.
step is to investigate the scaling of gradients in the al. [19] illustrated, the existence of an infinite-width limit for the gradient implies the optimal hyperparameters for the training algorithm will also converge. This type of results allows for tuning of hyperparameters on networks with a much smaller width, yet extends easily to arbitrarily large networks that approximates the same limit, saving massive amounts of computing Given the existence an infinite-depth-nn width limit for the forward pass, we believe it's possible to extract optimal hyperparameters from networks with not only a much smaller width, but smaller depth as well.
limit. As Yang et
in the process.
Preliminary Experiments. Although this work is primarily theoretical, it is important to consider whether or not the proposed architecture useful practice. Given limited computing resources, we chose to only briefly test the feasibility of training the shaped Transformer. Nevertheless, preliminary experiments show promising results when it comes to training stability. In particular, the shaped Transformer (without LayerNorm) does indeed train approximately the same speed as well tuned Transformer architectures. Full details of the experiment and results can be found in Appendix D. A more comprehensive set of experiments with different tasks, datasets, and larger networks will be required to confidently determine the practical feasibility of the shaped Transformer, which we defer to future work.
Training Dynamics and Generalization. infinite-width NTK theories motivates our study of the proportional andnitt-depth-and--iddee limit. In particular, to address many of the open problems deep learning theory, we need a faithful and tractable description of training dynamics. Given the results at initialization, the proportional limit holds the potential for such a theory training as well. Another promising indicator that deep networks learn features in the proportional regime [38], which has been identified as a key advantage of neural networks over kernel methods [25, 60-66]. A precise theory of training will help understand other types of instabilities during training and improve existing optimization methods. Furthermore, determining the network which training converges to is a necessary step towards theory of generalization, demonstrated by the infinite-width approach [67]. In light of our results, we helieve that our theory sets the stage for future work on training and generalization in deep learning.
mentioned in the introduction, the limitations of
Acknowledgement
CL and ML would like to thank Keiran Paster for insightful discussions. LN would like to thank Sotiris Anagnostidis for support in pre-processing the dataset used for the training experiments of this manuscript. ML is supported by the Ontario Graduate Scholarship and Vector Institute. DMR is supported in part by Canada CIFAR AI Chair funding through the Vector Institute, an NSERC Discovery Grant, Ontario Early Researcher Award, a stipend provided by the Charles Simonyi Endowment, and New Frontiers in Research Exploration Grant.
10

---

References
[1] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. "Language models are few-shot learners". In; Advances in Neural Information Processing Systems 33 (2020), pp. 1877-1901.
[2] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and
I. Polosukhin. Attention is all you need. 2017. aiv: 1706. 03762.
[3] J. Hestness, S. Narang, N. Ardalani, G. Diamos, H. Jun, H. Kianinejad, M. Patwary, M. Ali, Y. Yang. and Zhou, Deep learning scaling is predictable, empirically. 2017, arXiv:
1712. 00409.
[4] J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford,
J. Wu, and D. Amodei. Scaling laws for neural language models. 2020, arXiv: 2001. 08361.
[5] J. Hoffmann, S. Borgeaud, A. Mensch, E, Buchatskaya, T. Cai, E. Rutherford, D. d. L.. Casas,
I. A. Hendricks, J. Welbi, A. Clark, et al. Training compute-optinal large language models.
2022. arXiv: 2203.15556.
[6] Y. Tay, M. Dehghani, S. Abnar, H. W. Chung, W. Fedus, J. Rao, S. Narang, V. Q. Tran, D. Yogatama, and D. Metzler. Scaling Laws vs Model Architectures How does Inductive Bias
Influence Scaling? 2022. arXiv: 2207 10551.
[7] A. Aghajanyan, L. Yu, A. Conneau, W.-N. Hsu, K. Hambardzumyan, S. Zhang, S. Roller,
N. Goyal, O. Levy, and L. Zettlemoyer. Scaling Lawsfor Generative Mixed-Modal Language
Models. 2023, arXiv:
[8] Y. Dong. J.-B. Cordonnier, and A. Loukas. "Attention is not all you need: Pure attention loses rank doubly exponentially with depth".. International Conference on Machine Learning.
PMLR. 2021, pp. 2793 2803.
[9] L. Noci, S. Anagnosiidis, L. Biggio, A. Orvieto, S. P. Singh, and A. Lucchi. Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse. 2022. arXiv: 2206. 03126.
[10] B. Poole, S. Lahiri, M. Raghu, J. Sohl-Dickstein, and S. Ganguli. "Exponential expressivity in deep neural networks through transient chaos". In: Advances Neural Information Processing Systems 29 (2016).
[11] S. S. Schoenholz, J. Gilmer, S. Ganguli, and J. Sohl-Dickstein. "Deep Information Propaga- tion". ICLR. 2017.
[12] G. Yang and S. S. Schoenholz. "Mean field residual networks: on the edge of c Advances in Neural Information Processing Systems. 2017, pp. 2865-2873.
[13] Hayou, A. Doucet, and Rousseau. "On the impact of the activation function on deep neural networks training". International Conference on Machine Learning. PMLR. 2019, pp. 2672-2680.
[14] M. Murray, V. Abrol, and J. Tanner. "Activation function design for deep networks: linearity and effective initialisation". In: Applied and Computational Harmonic Analysis 59 (2022), pp. 117-154.
[15] L. Xiao, J. Pennington, and S. Schoenholz. "Disentangling trainability and generalization in deep neural networks". International Conference on Machine Learning. PMLR. 2020,
pp. 10462-10472.
[16] J. Martens, A. Ballard, G. Desjardins, G. Swirszez, V. Dalibard, J. Sohl-Dickstein, and S. S. Schoenholz. Rapid training of deep neural networks without skip connections or normalization layers using deep kernel shaping. 2021. arXiv: 2110. 01765.
[17] G. Zhang, A. Botev, and Martens. Deep Learning without Shortcuts: Shaping the Kernel with Tailored Rectifiers. 2022. arvi: 2203. 08120..
[18] M. Li, M. Nica, and D. Roy. 'Thh neural covariance SDE: Shaped infinite depth-and-
networks at initialization"" In: Advances Neural Information Processing Systems 35 (2022), pp. 10795-10808
[19] G. Yang, E. J. Hu, I. Babuschkin, S. Sidor, X. Liu, D. Farhi, N. Ryder, J. Pachocki, W. Chen, and J. Gao. Tensor Programs V: Taning Large Neural Networks via Zero-Shot Hyperparameler Transfer. 2022. arXiv: 2203. 03466.
[20] R. M. Neal. "Bayesian learning for neural networks". PhD thesis. University of Toronto, 1995.
11

---

[21] J. Lee, Y. Bahri, R. Novak, S. S. Schoenholz, J. Pennington, and J. Sohl-Dickstein. "Deep Neural Networks as Gaussian Processes". ICLR. 2018.
[22] A. G. d. G. Matthews, M. Rowland, J. Hron, R. E. Turner, and Z. Ghahramani. Gaussian process behaviour in wide deep neural networks. 2018. arXiv: 1804. 11271.
[23] J. Hron, Y. Bahri, Sohl-Dickstein, and R. Novak. "Infinite attention: NNGP and NTK for deep attention networks". International Conference on Machine Learning. PMLR. 2020,
pp. 4376-4386.
[24] A. Jacot, Gabriel, and Hongler. "Neural tangent kerne:: Convergence and generalization
in neural networks". Advances in Information Processing Systems (NeurlPS). 2018. arXiv:
1806. 07572.
[25] G. Yang and E. Hu, Feature learning in infinite-width neural networks. 2020. arXiv: 2011. 14522.
[26] J. Sirignano and K. Spiliopoulos. "Mean field analysis of neural networks: A law of large numbers". In: SIAM Journal on Applied Mathematics 80.2 (2020). pp. 725-752.
[27] S. Mei, A. Montanari, and P.-M. Nguyen. "A mean field view of the landscape of two-layer neural networks". In: Proceedings of the National Academy of Sciences 115.33 (2018), E7665- E7671.
[28] I.. Chizat and F. Bach. "On the global convergence of gradient descent for over-parameterized models using optimal transport". In: Advances in Neural Information Processing Systems 31 (2018).
[29] G. M. Rotskoff and E. Sanden-Eijnder Trainability and Accuracy of Neural Networks: An Interacting Partiele System Approach. 2018. arXiv: 1805. 00915.
[30] Yaida.
processes and neural networks at finite widths". Mathematical and Scientific Machine Learning. PMLR. 2020, pp. 165-192.
[31] E. Dyer and G. Gur-Ari. "Asymptotics of wide networks from feynman diagrams". ", In: arXiv preprint arXiv: 1909. / 1304 (2019).
[32] D. A. Roberts, S. Yaida, and B. Hanin. The principles of deep learning theory. Cambridge
University Press, 2022.
[33] J. Zavatone- Veth and Pehlevan. "Exact marginal prior distributions of finite Bayesian neural networks". In: Advances in Neural Information Processing Systems 34 (2021).
[34] J. Zavatone-Vett A. Canatar, B. Ruben, and C. Pehlevan. " "Asymptotics of representation learning in finite Bayesian neural networks". In: Advances in Neural Information Processing Systems 34 (2021), pp. 24765-24777.
[35] B. Hanin. Correlation Functions in Random Fully Connected Neural Networks at Finite Width. 2022. arXiv: 2204. 01058.
[36] E. Dinan, S. Yaida, and Zhang. Effective Theory of Transformers at Initialization. 2023. arXiv: 2304. 02034.
[37] B. Hanin and M. Nica. "Products of many large random matrices and gradients in deep neural networks". In: Communications Mathematical Physics (2019), pp. 1-36.
[38] B. Hanin and M. Nica. Finite depth and width corrections to the neural tangent kernel. 2019.
arXiv: 1909. 05989.
[39] Z. Hu and H. Huang. "On the Random Conjugate Kernel and Neural Tangent Kernel". Inter- national Conference on Machine Learning. PMLR. 2021, pp. 4359-4368.
[40] M. B. Li, M. Nica, and D. Roy. "The future is log- Gag-saannn:: ResNets and their deeth- and- width limit at initialization". In: Advances in Neural Information Processing Systems 34 (2021), pp. 7852-7864.
[41] I. Noci, G. Bachmann, K. Roth, S. Nowozin, and T. Hofmann. "Precise characterization of the prior predictive distribution of deep ReLU networks". In: Advances in Neural Information Processing Systems 34 (2021).
[42] B. Hanin and G. Paouris. aooptssyp results for singular values of Gaussian matrix products". In: Geometric and Functional Analysis 31.2 (2021), pp. 268-324.
[43] G. Akemann, Z. Burda, and M. Kieburg. "From integrable chaotic systems: Universal local statistics of Lyapumov exponents". In: Europhysics Letters 126.4 (2019), 40001.
[44] J. Ipsen. "Lyapunov exponents for products of rectangular real, complex and quaternionic Ginibre matrices". In: Journal of Physics A: Mathematical and Theoretical 48. 15 (2015),
p. 155204.
12

---

[45] T. Jiang and Y. Qi. "Spectral radii of large non-Hermitian random matrices". In: Journal of Theoretical Probability 30.1 (2017), pp. 326-364.
[46] B. He, J. Martens, G. Zhang, A. Botev, A. Brock, S. L. Smith, and Y. W. Teh. Deep Trans- formers without Shortcuts: Modifying Self-attention for Faithful Signal Propagation. 2023.
arXiv: 2302. 10322.
[47] M. Dehghani, J. Djolonga, B. Mustafa, P. Padlewski, J. Heek, J. Gilmer, A. Steiner, M. Caron,
R. Geirhos, Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters.
2023. arXiv: 2302. 05442.
[48] R. Xiong. Y. Yang, D. He, K. Zheng, S. Zheng, C. Xing, H. Zhang, Y. Lan, L. Wang, and
T. Liu. "On layer normalization in the transformer architecture". International Conference on Machine Learning. PMLR. 2020, pp. 10524-10533.
[49]
J. L. Ba, R. Kiros, and G. E. Hinton. "Layer normalization". In: (2016), arXiv: .605000
[50]
S. Hayou and G. Yang. Width and Depth Limits Commute in Residual Networks. 2023. arXiv:
2302.00453.
[51] Hayou, Clerico, He, G. Deligiannidis, A. Doucet, and J. Rousseau. "Stable resnet". International Conference on Artificial Intelligence and Statistics. PMLR. 2021, pp. 1324- 1332.
[52] W. Tarnowski, P. Warchot, S. Jastrzebski, J. Tabor, and M. Nowak. "Dynamical isometry is achieved in residual networks in universal way for any activation function". The 22nd International Conference on Artificial Intelligence and Statistics. PMLR. 2019, pp. 2221- 2230.
[53] Hayou On the infinite-eptth limit of width neural networks. 2022. arXiv: 2210. 00688.
[54] G. Yang. "Tensor programs ii: Neural tangent kernel for any architecture". In: arXiv preprint MiXiv22006.14544 (2020).
[55] K. He, X. Zhang. S. Ren, and J. Sun. "Delving deep into rectifiers: Surpassing human-level performance on imagenet classification". Proceedings of the IEEE international conference on computer vision. 2015, pp. 1026-1034.
[56] A. Veit, J. Wilber, and S. Belongie. "Residual networks behave like ensembles of relatively shallow networks". In: Advances in neural information processing systems 29 (2016).
[57] S. Zhai. T. Likhomanenko, E. Littwin, D. Busbridge. J. Ramapuram, Y. Zhang, J. Gu, and
J. Susskind. Stabilizing Transformer Training by Preventing Attention Entropy Collapse. 2023.
arXiv: 2303. 06296.
[58] I. Karatzas and S. Shreve. Brownian motion and stochastic calculus. Vol. 113. Springer Science
& Business Media, 2012.
[59] J. Miller. Stochastic Calculus (Lecture Notes). http : / www statslab cam ac uk
/ jetpios/teachhiieee http:///teeeccinnenene////////// cetttturrenttsseeeeeteee pdf. 2015.
[60] B. Ghorbani, S. Mei. T. Misiakiewicz, and A. Montanari. "When do neural networks outper- form kernel methods?" In: Advances in Neural Information Processing Systems 33 (2020), pp. 14820-14830.
[61] E. Abbe, E. B. Adsera, and T. Misiakiewicz. "The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks". Conference on Learning Theory. PMLR. 2022, pp. 4782-4887.
[62] J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, and G. Yang. "Higg-dimennional asymp- totics of feature learning: How one gradient step improves the representation". Int. Conf. Learning Representations (ICLR). 2022. URL: https opwnrevieww net uet/for /fr//ormm? id= akddvRG6EGi.
[63] A. Damian, J. Lee, and M. Soltanolkotabi. "Neural networks can learn representations with gradient descent".. Conference on Learning Theory. PMI.R 2022, pp. 5413-5452.
[64] A. Mousavi-Hosseenn S. Park, M. Girotti, I. Mitliagkas, and M. A. Erdogdu. Neural Networks Efficiently Learn -imeeniinsi Representations with SGD. 2022. arXiv: 2209. 14863. and T. Misiakiewicz. SGD learning on neural networks: leap complexity and saddle-to-sadde dynamics. 2023. arXiv: 2302. 11055.
[65] E. Abbe, E.
[66] R. Berthier, A. Montanari, and K. Zhou, Learning time-scales in two-layers neural networks.
2023, Xiv: 2303. 00055.
13

---

[67] P. L. Bartlett, A. Montanari, and A. Rakhlin. "Deep learning: statistical viewpoint". In: Acfa numerica 30 (2021), pp. 87-201.
[68] O. Kallenberg. Foundations of Modern Probability. Probability theory and stochastic mod- elling. Springer, 2021. ISBN: 9783030618728.
[69]
HugeingFace. Wikipedia data set. https: http:///gnngggacccccoccttttttaaa
[70]
http/daaaaatts/booocoroopp HuggingFace. Bookcorpus data set. http::////nngggng...ceoo///tttt/
[71] D. P. Kingma and J. Ba. "Adam: A method for stochastic optimization". In: arXiv preprint (2014).
14
[72] A. Wang. A. Singh, J. Michael, Hill, O. Levy, and R. Bowman. "GLUE: A multi-task benchmark and analysis platform for natural language understanding". In: arXiv preprint 184:0:1804.0746/ (2018).
[73] HuggingFace. Hugging face Bert http:///tt..tttteeee/////oo//a/n
transf hha.a///02.0.aaaa////////e/aaaat
14

---

A Preliminaries: Covariance SDE Framework
In this section, we will review existing results on Markov chain convergence to SDE, as well as existing results for the covariance SDE. See also the Appendix of M. Li et al. [18] for more details.
Firstly, we will define the Skorohod Ji-topology, or just the Skorohod topology for short [68, Appendix 5]. The topology is used describe convergence of continuous time processes with discontinuities. in particular, Markov chains with a continuous time interpolation fits in this category. Let S be a complete separable metric space, and DR..s s be the space of cådlàg functions (right w! continuous with left limits) from R, + S. We use In - x to denote locally uniform convergence (uniform on compact subsets of R4), and consider the class of bijections À on 2 that À is strictly increasing with Àg = (can be interpreted as a time change).
Definition A.1. We define Skorohod convergence In E x DR.,S , if there exists sequence of bijections Xn satisfying above conditions and
Xn
u1 ld
In An _ x.
(13)
In particular, we call the topology that induces this convergence the Skorohod topology [68, Theorem A5.3].
On a heuristic level, if we have sequences of Markov chains Y" that satisfy the following type of Euler updates
6(y
(Y,") VT
Y = Y2
(14)
n
where & are iid random variables with zero mean and identity covariance, then we can interpolate this process in continuous time with X;" = [tn] and show that as n + co, we have that X converges to the solution of the following SDE (weakly with respect to the Skorohod topology)
dx,- b(x,) di (x,) dB,, Xo = lim Y
(15)
Our next theorem essentially weakens this result in several ways. Firstly, we don't need take the step size n-1 but instead replace it with 22 for all p > O. Next, we can allow the update to contain higher order terms that vanish, in particular, any terms of order O(n 2p-6) for all 6 > O. Here, we remind the readers that O(n P) denotes a random variable Z such that n2 Z has all moment bounded by universal constant (independent of rt). Thirdly, we can allow b(y) to be random. or more precisely replace it with b(y, w.) such that Eb(y. Wn) = b(y). Fourthly, we can also allow b, a to be a sequence bn, Jn such that they converge b, a on uniformly on compact sets. Finally, can also weaken the topology of convergence to locally, that is all processes are stopped by stopping time in Definition 4.
Now we will state the main technical result in this section.
Proposition A.2 (Convergence of Markov Chains to SDE, Proposition A.6, [18]). Let Y be a discrete time Markov chain on R defined by the following update for p, 8 >>> 0
(YF, on(Y"") n2
E+1 Y
cr O0n
(16)
np
where EE E N are iid random variables with zero mean, identity covariance, and moments uniformly bounded in n. Furthermore, wit are also iid random variables such that [Bbb,, (oXY",W)))p)) ,,)) = y] =
b,, (v) and b,, (y. w() has uniformly bounded moments in n. Finally, on is deterministic dunncinnistiicfuaciinnn and the remainder terms in O(n -22-5 have uniformly bounded moments in n.
Suppose ba, on are uniformly Lipschitz functions in and converges to b, a uniformly on compact sets, then in the limit as n - co, the process X = [imao converges in distribution to the solution
of the following SDE in the Skorohod topology of DR, RN
dx, b(x)) di + o (X,) dB,,
Xo = lim Y! n000
(17)
15

---

Suppose otherwise bu; on only locally Lipschitz (but still uniform in n.), then X converges locally to X in the same topology (see Definition 4. I. More precisely, for any fixed r > O, we consider the stopping times
T= r then the stopped process XiA+" converges in distribution to the stopped solution XtAr of the above SDE in the same topology.
: in
0: X 2
(18)
We will briefly recall the main result of M. et al. [18] next. As mentioned earlier in the text, the setting is for MLPs defined follows
Xe+1 = (Xe) We.
(19)
where (x(x) 8+ max (2, 0) + s min(x, 0) with 81 = =1+ vi = is the width of the network, and We,sy ~ NN 1) are iid random weights.
E (o, for g N/ 1),
Then it was shown that in the limit as n, d + oo with = = > 0, the covariance matrix Ve = X (or equivalent the post activation (X)) since the activation becomes infinitesimal) converges the solution of the following SDE in the Skorohod topology [18, Theorem 3.2]
1 dV; (V)) dt + Minn (V:) dBr, Vo XXx]
(20)
π
where the coefficients were given in Theorem 3.2.
We remark that the output is defined as Xout distribution as
W oan Xa E Roa, then we can recover the output Jn
Xout ~ (O, VT & Inae),
(21)
where we treated Xuu as a vector in mine.
B SDE for Residual Network
Recall that we adopt the following model:
X XXe + 7 1 os
(V wp
(22)
Vn
where (1(x) 84 max(x, 0) + min(2, 0) is the shaped ReLU with slopes 81 = + vn
for
some constants C.,, C_ R. We let the weights be (W?" ):y (W7Poo) iid (00 and Eo, (g) for g ~ (O, 1) is the He initialization constant [55].
From here onwards, we will define the filtration F3 a (1Xx Faslg). where a(-) denotes the sigma- algebra generated by the random variable. Furthermore, we will define the conditional expectation Eel-1= = - [Fe..
We are interested in studying the neural covariance, i.e. V.s : S (2 2 where e E [d] indexes the layers and d the depth. In particular, we are interested in understanding the simultancous limit (n, d) + co, where the ratio t = d/n remains constant.
From now on, we will remove the dependence on Defining the residual branch as f(X, W) Vn T. (VEX (EW)) W/post we can write Va as:
vas x2 + Xy(2", f(X, W) + X7k0 ((X, (F(,, W))
For the cross product terms, we get:
Ve (0t (ox (o, (XWP1) W
Vc
(x ((XX
popstt
W
Σ aiy wroe
in
(x "(x, WY
Ve
Σ
2 W
o Σ
23 WITe
x
16

---

For the term in 72:
(0. (XWPr)) W (o. (XWRR) Wrpoa) 8
(a (XWP)) Wroar) at (o. (XWPr)) W
nt2
π
Σ 2x
2e kj
2 Js Σ
K33''
n
m
R
k
We will define the following terms
True= eve nyn
W Wheat
Σ ag W3g
+ txgos
Σ
ji
2,
and
15
T
Σ o2 agwee
os Σ u W
j's os
n²
js
Hence we get the following update for Ve = e
(z) x²)):
a
E+1 = va
7d T
72
Tea0
vn
vn
It is easy to see that the cross product terms have (conditional) mean zero. For the term with y2:
Ee [T
23
where (1() (o) where gº, (N 1) and Bggg (x
Here
=||=
we recall [. = El [Fe] is the conditional expectation given the agma-allebba generated by
F. of (1xx (x))))))
M. Li et al. [18] showed that if the non linearity scaling exponent is p 1/2, then:
cK1(p) = -+ u(p) + Onn
where v(p) te-to_)?
(VI- p² + parccos(p)). ) Using this result, and summing and subtracting
2x
the mean of Tz:
71
2
(+1 = v2 72 Vaa V88 cK 1 (pas)
( +
Tra:
(T
Ex
(2
Vn
vn
= =x2 v2s + 2 /Vavvas (p
7t
22
T
+
(Te
Ee [T 1) + O(n -3/2).
Vn
Vni
Using X2 22 = 1:
v- F+1 = Vad + 2 /Va Vos (pa3)
7>
22
Tra
(T
E
(T + O( (n-3/))
vn
vn
Furthermore, we need second order moments of T and Tz. We derive them in the following Lemmas: Lemma B.1.
c + you c (0 Ex 8 pao
V8 /oaa wwo + y 1o (D 2 + 22 O(n
17

---

Proof. Recall the definition of T
+ +x,o, Σ az w337 ))
TTS
CvC nVn
W podl
Σ azg W
ji
We have that:
c
Ee 1. (,550
E
W
ju Ex x ||=00 [o(0))) + x |> (.. (2 (0(0)) + x. | 0=
n
a
x x Ee
||
ক ||00|| | o. (0)))) (.,00))))0
x |
+ |لاا||
vad VV33V (0 + wou 3333666 cKi (D=0)
1 VWaa (K ( + vow VVaa 1 (pad)
Lemma B.2.
Ee = V = vnval + Vyvay + +
Vn
Ee TT 2 T2u)
[T = 20 + 2 Ee [T2 3
On=
Proof. Recall the definition of T
T2 T
W
Whoal
f'a Os Σ of aggwe as
Σ Ff wore
n?
n
JJJ'=1
||
jt j'i as (0 (of Jos (g3))
n² 2
+.3'==1
For the mean, using the independence between Wrr and W/rust we have that:
c²
E
T
|
E
Wr post
E (0.1099)00199))
n² Vn
j's
| 10011|1031
E
n/n
j=1
||001
|
Ki (pa0)
c/m/vaay K1 ())),
which is the desired result. The final expression is the result of the aforementioned expansion for Ki [18]. For the covariance, we have that:
Eg [T 8 T
|
Σ
E
W pusd j'i
W
E
5
ki
1e's
2
VVaaV39y88va
Σ
(5ggrOks + dykdw + 65w (o..
n
18

---

Let's look each term of the separately. For the first term have that:
Σ
E
6..8 J.k
n²
Σ
J.k
n² Σ E (x(oo))ss
(0,00))011)) n²
Σ E (o. E (o.()))))0))))) sZk
= =n =nlKa + -n ( ( (P
where have defined the fourth moment of the shaped activation: K2(pns8u) E (o
for which it holds that [18, Lemma C.2]:
K2(pa8bu)= = =Egg gf ggg") + O(n O(n-1/3 p0 Dow pad now + p On
The other two summands can be solved similarly:
Σ
)o (o
= ni = K2(paBbu) + n22nn (p (p())
and
Σ 6 K
(o = K2(pa + n22nn (n- 1) 1 (p (100))
Hence, summing the three terms:
Eg (TT c² VVV K2(pasbu) + (n ( (pa0) (1
n (K2(p) + (n- 1) Ki (K
n2 (Kz(p (n - (p (p
12V K2 + nk n (1(p00) (( (
(K( K (p Bw) + (1 ) (08)) O(
Now, subtracting Eg T2
E2 [T Sw using the aforementioned expansions for K₁ and K2:
Ee [T T
Ee [T Ee hT = p (p no + p pow + paw g00)
valyów + pa + @5u + +(( [T
2V 21 255 Onn
where in the final step we have used the fact that 1 ((n This completes the proof.
Finally we also have that Ti and To are uncorrelated, i.e. Lemma B.3.
E2 [T TT 0.
Proof. It is easy to see that T involve odd standard Gaussian moments (in particular third order moments), which vanish due to the parity of the centered Gaussian measure.
Theorem 3.2. Let X₃ be the hidden layers of a ResNet defined in Eq. 4 with X2 + === 1, where both X and y do not depend d, Then the feature covariance Ve converges to the solution of the following SDE (in the sense of Definition 3.1)
1 = Vo
d.., = (V.) dt + Ene, (22) (V dd,.
(5)
where (ve = b (V) (V) = ?>>>>(((p>>) Vao 133 ass with pa = vas (Va and
upp) (c--
/1 11-p2 p² parccos
(6)
furthermore, Eneo (V) = -))=22220) 2i (nn(V = 2x2 2o yow 18) Ja59.55w |,00,6000
19

---

Proof.
We know that:
71
7²
T2
E+1
vn
+
vn
Using X² + = 1, and summing and subtracting the mean of Tz, we have that:
va 6+1 Vi - v2 Ee
7
722
Teas
T2 T2 )
+
nb
Vn
Using Lemma B.2 for the mean of T2, we have that:
Vas
1 4+ = Vi + (00
AA T
72
(Tae
- -Bg
T2us ))
+ Onn (/2))
which gives us an expression the right Markov Chain form with drift term VP For the diffusion term, we need to compute the covariance between two different neural covariances vag Et1 E+1- Using Lemmas B.1 to B.3, we have
Cov: (veti.
(NyTT E. [>>TT + 72 - - (2
+ +22T22 - Be (Tew1)) -
=X TT T TruiTa - Tras Eg [
+ Vau + 27 (Va80 + vaw )86) + 0 1)
=
= =272(Va66 + Va + +(n
where we use Cov, to denote the covariance conditioned on the sigma-algebra Jf o((Xx Finally, applying Proposition A.2, we get the desired result.
C SDE for Softmax-based Attention
C.1 Dot-Product Attention
Dot-product attention applies the Softmax row-wise to the following matrix:
Y
1 X
W W. Q.T xI,
71
W
where W.K. w RuxAN are Gaussian matrices with unit variance entries, and n& is the queries and keys' dimension. Here we study the first two moments of Yg. In particular, note that [Ee[Ye] = 0, where E, -- E Fi] denotes the conditional expectation given the sigma-algebra generated by F. ( [X (x)))))))
For the second moment have the following Lemma.
Lemma C.1. Ler
Y, Xx,W WOT X
be the dot-product attention parametrized by W.K. w e gux., Then:
Eg[ya8ydu = nu v Bww
(24)
Proof.
Ee (yaByów] 1 Σ n2
x XX W
kk'
20

---

For the expectation, have that:
B
B
33"
=
E
i
E [
E W
W W
Σ d
nx mx6000
where we recall W WQ..
Hence:
Ee Pa
72 n2
Σ
X X ng voVo
C.2
Shaping the Softmax
Recall that we have the following model:
1 WV
= XX
Vn
For the above model, Vj has the following form:
by (X&W X, AT + A.X.Wxx)) + A.X.W.W x AT. nyn
T2 n²
Ve+ = v
In order to have infinitesimal updates in Vz, we would intuitively like the attention matrix to be the form A03 8. Onn-1). In the case of the usual Softmax based attention, have:
A₂ = Softmax(r "Y;),
where - is temperature parameter that regulates the entropy of the resulting categorical distribu- tion; r is often chosen to be large (scale as a power of n or d), as low temperature results unstable training.
To transform the Softmax into the desired form of A = 8a, +0nn we first center the attention matrix around the identity matrix:
A = 1 + Softmax(--
and then examine the Taylor-expansion of the Softmax w.n.t
A = 1+ 1 I1T +0(T-)).
(25)
The first few terms of the expansion provides good approximation of A, when is large. Observe that for fixed m, the zero order term m -11" is not of O(n In particular, suppose that we use the attention model in Eq. 25. We can show that the expectation of the term in ? (responsible for the drift) has the following form:
Ee A Aauj = v20 7²
2
72
yow
688
yow + O1-
m
6,==1
where the terms scaled by leads to large incremental updates in expected value Ve w.r.t the previous layer, and precludes the possibility of convergence an SDE. Hence, we choose to re-center the Softmax by removing the term
1 11 as follows:
1 11
A = I Softmax (T-
m
21

---

which admits the following Taylor-expansion:
(ootmax(T-"Y)))--a -11 -1 | Ya 1 272m
( - ((Ya)3 -Ya3] 0,8
-0(T
where Ya := 1 Σn you, and (Ya)² 1 Ex (Yar)
Hence, up third order:
S 1a,3 + 2r²m [r 1yo0--Po -8 (Yaa)- + O(T
A
For sufficiently large т, the formulation above allows infinitesimal updates in Ve and will show rigorously in the rest of this section, permits convergence to an SDE.
C.3 Lemmas on Moments of Shaped Attention
Define:
FT yaB Ya
F22 = (y - -8022 ((Ya}3 ya²)
Hence, Ag can be written as:
A dag + Tm
F2 + 0
2r
We now compute the moments of Ag. We define the following quantities:
gab
1
Ey(Y
(Y) - Y
E,
Fo F&W.
7
74
Ee (yuu
(yya--Yay y012 ((Ya)z (Yaaz--a'))
E,
12k
7
where we recall Ee[
E1- Fe] is the conditional expectation given the igmma-algebe generated by
Fj = ((Xx)ket())
Lemma C.2 (Moments of Taylor Expansion).
S = vas (vow V8B VwB + V&2).
S2 = Voa (vō6 25 + 2V - v),
1 where V = 1 m En V and o m En x" is the average token.
Proof.
Using Lemma C.1 and linearity of expectation:
sq 1 7i (yas yag ya you va v yal (Vow - vov vwB + voz),
[ E, [Y Ee Sw ya Ee (Yay))
and:
82 1
Ee (yad - pay2] E. ((vajz vatt)
n
1
[Re((yadd))) 2R 2R + 2F Fo Ee (Y(()))
Thk
(Veay 22 + aa voa
yaa (v56 52 + 2V 28 - -V.
where V = Eg V83
22

---

Lemma C.3.
Ee[AAd A 8addpw
m2m S1
2 (5pw
ne
Sa6 5ab52 + (MxT
Proof.
Ee [2448 Adw) dusdRa
daw BBp yyad Ya
dad E, (rrW V3
1
B2(yas
72m²
5m
7m
+ dew
Ee
(rrd
pon - ((Ya)2 yaa)
27
+ 27
dad
Ee
(Y
yu ((Y Yo¹) 5|
+ O
dasd Bw
Spo.3w
+
Th
S2 Sa + Sab + O(
-2m1"
27
Lemma C.4.
Eg A A08" A86' Aww
dac 630
nik
03g. S/86 and + Sa dss Sp8r awl + d &wn SS8 S$B .68'
m2m
+ dgg 858 St gaa''ww www + d S saa' + Sw ST .99"
27 aaardgg0000.. S2n + San dag
11
Saa Sww S.0'' dew dggrd sgad O(NET
Proof.
Ee [Aaa' A88' 168' Aual Ee| ( Oaar
par'
1 Fgaa) (Bes) +
1 F&8
+
+
Tm
272m
Tm
2
66s + TM
1 F765
27 FV8' (Bw
1 Tm
Frow
Frow )
O(T-3)
272m
Eg
dag. dws Saw + 72m dag prow'
m2 72m dan dss Fs8.1000 poow 72m² 888 883 Faaa p30 FaaeswFF Faaa F66
+
1
?2m Saw F 1e00'
? 03g Fiad pluw
+
m2m² 5p3.8wwr
2 Saa Öggrr8s F
27 8 03g. 8ws F28
272m 8an 8 120
2-2m 8 dgg 0gg Dooo
+ 0
Using the linearity of expectation:
Ee Aaa' A88' A Aww"
dac 8333 888
Thk
(Sa' 83g S/8 aod
ST + dan Sp wwl
+ dan Swwr SpA .88
+ mm?
+ 53g 658 St ww'
+ ōgg Sww S .68'
S s .39
+ dow
272m (ga Stww + San Ogg S58
+ Sww 8B' + dss O(n
23

---

Note that the terms above be computed using Lemma C.2.
Lemma C.5.
Aaa' A 160' Anm² Ee [Aaa' 1189
Ee [188' AWw'
Eg
Tkk 72m2
058 Sq8 ww' + Oww SpB
88
+ 88g 688 ST
new
+ Oww _ + + 9).
Proof. The results is an immediate consequence of Lemma C.4 and Lemma C.3, where only the terms that do not cancel out are kept.
C.4 Neural Covariance SDE for Stable Attention
Recall that we have the following model:
X411 XX 7 1 WY
Vn
For the above model, Ve has the following form:
by
EXW,TXTA AT + AxXWxx]) + X2 n2
Ve+1 = vo + Byn
We define:
7 Ti
- -X-WT CXXWTXIAL + Axx-WoXT) a3
8 1 nVn
(A X AT)
Hence, the expression for Vi2 simplifies to:
xy T + 2² Vn
Vag = X2
Tea.
We need compute the moments for these two quantities: Lemma C.6 (Moments of Ti).
B [ 0,
and
Eg
T
28 + Ba + O(net 3).
Proof.
T To
1
(X,WTXX) XT + AxXWWXX)) a8 (XWTXXAA
Ax 6w
n²
1
(XxW X AT) 08 (XoWIXTAT)" (XXWIXTAAT"" OB (AcXoWexT) 8
+ (A.X.WxXT)) og (X.W, NAW,XTAN AT) + (AxX.WxX)) as (A.X)WoXT) &0
Let's look at the first summand:
1
(XeW, X, AT) aB (XWW, X 17)
Σ X
n
n 3j'
24

---

Hence, in expectation with respect to W:
E Σ
x
A
X
Au
1
E X Xx X
T
n²a
kj'
kk'jj'
X
n
DR
kk
Σ
AnX. A
An identical argument can be made for the remaining three summands. Hence, taking expectation with respect to the Softmax weights:
Eg
7 T
Σ
(vaágvme,
(A A + yawp B A + yor Ea AMM| + V g Es (00
UX
Now, using Lemma C.3:
Eg
Σ vun (vod AWM] + PoWEE 1 + E [A AA A4M] + V [A
Σ vun (V sgpddn + poudggrnn V Saw + v 8ap + O
yauyad + yow V8d yao vau + O(nNT
2(VodyR + g + O(nAT
Lemma C.7 (Moments of T2).
Ee
TT /n Σ V"ATE [A ASM]]
UK
Ee Ee Tobu Σ
Ee[AaM A AWn'l (nvvryow + Vor'Tan' + Voe vuu))
Proof.
T = m/n
1 (A.X.WWW X AT)
-E
A WWWWEWWse X5 A UK kik
nVn
Taking expectation with respect to W:
1 nyn (ArX.WoW? X AT) a8
n ΣΣ
A
A
PK kk's
A X A5
TI
UK ki
ΣΣ
A 38 X X
n
UK RE
ΣΣ
A
ATH A X
UK
Vuk
A0 Axx
Taking expectation w.r.t the Softmax weights, we get the desired result. For second moment, we can take the conditional expectation:
Ee (T Trow
E Σ a AN 55 Nnn X XXXXX XX X3
n
25

---

where recall - B[[[FF) is the conditional expectation given the sigma-algehra generated by Ft = (({Xx)<E(())
Using Isserlis Theorem, we have that:
W We We Wyr]= =dgōdes + 6 + ōRz d
Hence:
Σ
Aav A Aou' Awne
Σ
XzXxXX X3 (Dujōry + dZidRedesr
n
UND'R!
kk'jaey'
+ Óu
Aav A Aōv' AWn'
XXXXXX x + Xxx xy
+ XXX XX
kkjj
Σ
Aay A Aóu' Awn' nnvuuryur + navve VAK' + n veve polno)
n
Σ
A
9 A5 AWn (nVoRyUN you yiv + ve'n)
By taking expectation w.r.t the Softmax parameters, we get the desired result.
Lemma C.8 (Covariance of T2))
Ee T2
8 Tz/u Bx[T2 8
E TT
nikk
Anilbw vadysu + you +0
Pek
72
73
72
where:
Aasbww
Σ
S + saa - Buu gaurnn + V sav. 6n
m
MR
Proof. Using Lemma C.7, we have that:
Ee (T 8 Tu 11 T2 [E
TT
Σ
Ee AAv AlAawAA A A Aun't InvuNuuw you pan' + &vuk puln)
yoaroen E, nwv A E. M Awn
n
VnV'K'
Σ Vunqu'N'
Abd Awn' - Ee [ A Eg [Δoe' AWx
n
DND'K!
Σ
S /0 A you''yan + ven yerx)
Now we can use Lemma C.3 and Lemma C.4 to compute the moments of A. For the second summand, simply have:
Ep Aow A8k A8K Anv' Aunn dawdnndawddnn + O(nTT
hence:
Σ Ee[A0V ABn Aov A (pen' youe' + yun' padd + Daw O PNE'
For the first summand, recall from Lemma C.5 that:
Eg
o 388 5558 Awn"
Ee [ 200
Ee 58 A Aww
= 7kk
Saw S/8 wa' + S S33 .8
+ 5gg S wad
+ & S .66'
+ +
26

---

Hence:
Σ yuuqu'n [Bo[AarA gsm AWK Ee [Aar A Be [Ad A
VRV'K
Vuayurk St + S + dex Squnbrr Squnen' + dgn Snn Saaioo")
72m2
+ O(rnET
mn
Vaxgoosshe spaaa vanywu Sqn. vruggn Spusun + vou S squud
An
O ))).
We are now ready re-state and proof of Theorem 4.2.
Theorem 4.2. Let Ne be the hidden layers of a residual attention network defined in Eq. I with shaped attention in Eq. parameters X2 + q2 -1 and Toy/ung where X, 7, To all do not depend on d, n. Then the feature covariance Vj converges locally to the solution of the following SDE (in the sense of Definition 4.1)
dV. D(Vi)dt E(V))
Vo= - XoXO n
where the drift has the following form
MVV) = 2
VuR sav,Br
S1
(v gay -aa goBu
m
2m
M,R=1
1=1
as
the diffusion coefficient is defined by E(V)= = - 7222) Evin + -0 2 188855 and
m
A0
Σ
Snpnydsg@mmnn + Si vn squnun + pwu sav.6n)
m
Proof. Recall that:
V41 X²vas
xy T +
72 nn
Tra
Vn
Drift. Summing and subtracting Eg [T] (and using Lemma C.7), we have that:
vaB = 22 7² Σ VuNE [A [
+
Xy
T
72
(T
[ee(Ta))
nn
n
where recall Ee( R[- [Fe) is the conditional expectation given the sigma-algebra generated by F: = (1Xx)
Re-stating Lemma C.3, we have that:
EA[A"* A 8andgn + n&k 72m2
ST Sov,8x
+
Th
(6pu Sav + Saw + O(nTT 3).
2r
Plugging in the expression, and using 22 + 22
have that the drift is:
1 gav
1
X vas vva
+7²
yuu
-
(0a + daw + O(1TT
7mm
2
-va
2 +
Vue yougsa. Sa
1
Bu san + yau (o + O(neT
7
m2
2m
UK
27

---

From here, it is evident that in order to have the drift scaling as ((1/n) need to choose:
2 ng
(26)
where To > 0 a constant.
Covariance. Recall that:
vag = v2 + 72 Σ VuNBe [A [ok] + xy T 22 (T Er [T(
Pm
Furthermore, we have set X² + 72 and 7== Tinms to have the right scaling for the drift.
What's left to is to compute the conditional covariance for Ve+1 given Fg. Noting that T: TTh uncorrelated, i.e. Ee (T 0 (similarly to the case of Resnet with shaped ReLU Lemma B.3),
we have that:
Cove (ver voi)
Ee (AyTTER + ² Ee (NyTT ~
+72
X² 2 (ru 7 T T:u)
Taa
Eg [T
where Cov; to denote the conditional covariance given the sigma-algebra generated by Fg = of (((x)kee())
Using Lemma C.6 and Lemma C.8, we have that:
Cove (Vet1 V2t1)
>²72 (2(Va6 + vawpso)) + t7 Aasbu yadyow + vaw
1
= (2 (Vaóy you you
Aap6w
Now we can apply Proposition A.2 for locally Lipschitz drift and covariance coefficients, which gives us the desired result in local convergence in the Skorohod topology.
We will also restate and prove Corollary 4.3.
Corollary 4.3 (Shaped Transformer Covariance SDE). Let X₃ the hidden layers of a shaped transformer defined in Eq. 11 with parameters X2 - = and To To /nng, where X, 7 all do not depend on d, n. Then the feature covariance Vt converges locally to the solution of the following SDE (in the sense of Definition 4.1)
dV= = [b(V)) + bee((V)) dt + (VV.) + Evee dB,,
(12)
where the coefficients are defined in Theorem 3.2 and Theorem 4.2.
Proof. To combine the results of Theorem 3.2 and Theorem 4.2, it is sufficient to combine the following (simplified) iterated Markov updates into one Markov chain
U, = V2 b(V2) A/2.)
V4+1 = U2 (Ue) Ein Sin(U))//2g
(27)
Vn
where & & independent Since the limit, we have that either updates are infinitesimal, i.e.
and identity covariance random vectors.
U Vel 22000 0 almost surely,
(28)
then we can write (ROLU(U.) = Breeu. ( (Ve, we). where eventually we have that
lim Ee
= b(V2).
(29)
n n=000
28

---

so it will not contribute towards affecting the limit. Here we recall E. .1 R.[|F5) is the conditional expectation given the sigma-algebra generated by Fe a({Xx)ael) (ele)) We can treat Eii (Us( similarly to get the Markov chain update
V4+1 = V b(V)) UR..U. (Ve,wr)
E(vi) vn
Eii, (Vo)
(30)
n
vn
which converges to the following SDE with two Brownian motions using Proposition A.2
dV, = [b(V)) + brs (V))) dt + E(V)) dB, + (v.
(31)
Observe that since the two Brownian motions B,, By are independent, it's equivalent to write
E(V)) dB, Eres Eves(V (Vi) [V(V) + Ere (V)) dB.
(32)
We recover the desired results from considering a more general form the iterated Markov updates as in Proposition A.2, which do not hinder the above derivation.
D Preliminary Experiments
We perform preliminary experiments to understand the effect of shaped attention on training deep Transformer architectures. particular, we consider pre-training masked language modeling task, where we mask 15% of the tokens. We use a subset of the English Wikipedia 20220301.er and English bookcorpus datasets [69, 70]. As a baseline, we adopt Pre-LN Transformer encoder architecture with 18 or 24 blocks. For the residual feedforward layer, we shape the ReLU activation according to Eq. 4, by changing its negative slope 8 1 1-1/vl //vn instead of We then incorporate our shaped attention by replacing the attention mechanism as dictated in Eq. 9. We also add scalar multipliers Y1, 72 E IR both to the identity and centering terms of the shaped attention:
A = I Softmax(T -Y)) - 11T,
T T= Toy
(33)
and propose two ways to set 71, 72 during training. In both alternatives initialize "2 72 1, thus leveraging the stability properties of shaped attention at initialization. During training, we either:
1. Recover. Linearly decrease "/, 12 to zero with in the first 4000 steps, thus recovering the standard attention layer. Apply the same schedule for the Saaped-ReLLU slope s.,, recovering the usual ReLU activation. This approach recovers the vanilla Transformer architecture (without LayerNorm).
2. Learn. Learn all the shaping parameters 71, 1/2 and s..
The intuitive rationale behind these choices is that at initialization we want good signal propagation and a non-degenerate covariance (according to our theory. this requires the shaped attention and ReLU), On the other hand, we also allow the model to more dynamically make use of the nonlinearity during training to modify the correlation structure with Recover Learn. We report that without either adjustment, shaped atlention is still trainable but at much slower rates.
We also incorporate the t factor into the initialization of the queries and keys weights by decreasing the variance of their entries by factor of 1 This allows us to not re-tune the learning rate for the queries and keys. We stress that at at initialization, the two formulations (r = Vnn, and vnL with decreased weight's variance) are equivalent. In both alternatives we train all the skip connection parameters X, 7, and initialize y in the grid (0.05, 0,1,0.2) and set To = 1. We report training instabilities (loss divergence) for larger values of 7. All models inncluding the baselines use Adam [71] with learning rate warmup of 4000 steps, the learning rate is tuned in the grid (0.0001, 0.0005, 0.001, 0.005). We report the train/test loss after 100K optimization steps, averaged over 4 random seeds. All the other experimental details can be found in Appendix D.
The Shaped Transformer is Trainable. In Table 1, we compare the train/test loss of the two variant of shaped attention with the baseline Pre-LN model. Notice that our model (in both variants) achieves comparable performance to standard Transformers across all the reported values of y.
29

---

GLUE evaluation Furthermore, evaluate the trained models three datasets from the GLUE benchmark [72] and summarize the results in Table 2. These demonstrate that our shaped Transformers holds promise by outperforming our pre-In baseline.
d 18
=
24
Train Loss
Test Loss
Train Loss
Test Loss
0.05
2.0940.02
2.07 10.02
2.03
2.03
Learn
0.1
2. 00.02
10.01
2.02:10.01
2.02
0.2
2.07 2007
2.0640.02
2.09
2.09
0.05
2.05
2.04
2.02+1.00
2.00
Recover
0.1
2.05
2.04
2.06+0.03
2.00
0.2
2.01 10.03
2.03 20.m
2.054
2.00
Baseline
2.08
2.07
2.08+0.03
2.01
Table 1: Training and test loss of the proposed shaped attention, both in the "Recover" and "Learn" alternatives. Baseline refers to the vanilla Pre-INN Transformer. We report the best run under the learning rates (0.0001, 0.0005, 0.001, 0.005), averaged over 4 random seeds. We include the confidence interval of one standard deviation.
Model
COLA
MRPC
RTE
Baseline
0..00
0.504 10.000
d 18
Learn, 7 0.2
0.182 10.041
0.813 10.005
0.528 10.019
Recover, 7 = 0.2
0.221 10.042
0.80910.009
0.520
Baseline
0.022 10.055
0.785 1.0009
0.48 10.04
d 24
Learn, 7 = 0.05
0.150.0..02
0.812 10.004
0.513
Recover, 7 = 0.05
0.211 #0.039
0.803
0.529 10.019
BERT (Geiping et al.)
0.103
74.8
0.509
Table 2: Evaluation of the pretrained models on a subset of the GLUE benchmark. We take the checkpoints of the baselines and selected shaped Transformers (7 0.2 for the shallower model = 18, 7 = 0.05 for the deeper model d = 24). Then, we fine-tune on each of three datasets from GLUE (COLA, MRPC, RTE) for epochs and Adam optimizer Ir 5e with batch size 16, and report test evaluation metric for the corresponding task.
Entropy Collapse for Large Learning rates. To understand the sources of training instability, we keep track of the entropy of the probability distribution induced by the Softmax, it has been observed that the Transformer's 's training is unstable in the low-entropy regime [57]. The entropy calculated for each row of the Softmax matrix, and it averaged across rows and heads. The results are in Fig. 6. Notice how for the large learning rate regime observed in Fig. ??, the entropy collapses for the baseline model, but not for the proposed shaped Transformer. Entropy collapse indicates that the Softmax distribution degenerates to point mass, which is itself caused by large logits. Remarkably, this phenomenon does not affect the recover setting, despite recovering the Transformer architecture (without layer normalization) after the warm-up period.
D.1 Experimental details
Dataset. We use subset of the English Wikipedia 20220301. een and English bookcorpus datasets (69, 70]. The sentences are tokenized using by pre-training a tokenizer on the training set. We use vocabulary size of 32000, and maximum sequence length of 128 tokens.
Model parameters. We use an embedding size of 768 and multi-attention heads. The batch size is fixed to 32 sequences. All the initial weights sampled from (0, n n'), with the exception of the queries and keys' weights W Wº the shaped attention case, that are sampled from N/ n 33/2) (as explained in Appendix D). The feedforward layer maps the n 76S-dimensional embedding to the larger dimension 3072, in the Hugging face implementation of Bert [73].
30

---

10
124
gaggmma
Entrapy. Layer
0.1
10"1
0.2
Eriropy,
107
resgamma
0.001
0.2
10
10000
20000
10000
40000
50000
80000
10000
20000
30000
40000
50000
63000
step
181
10
10
10-
Entrap: Layer
10
10
1073
10"1
0.0002
10
10"
10000
20000
30000
50000
50000
00000
10000
20000
30030
40003
50000
0000C
stop
stop
Figure Dynamics of the mean entropy across heads for selected layers (first and seventh) and learning rates. (Above): shaped attention in recover setting. (Below): Baseline transformers. Notice that ir 0.001 the entropy collapses for the baseline model.
Optimization. We train using Adam [71] with betas parameters (0.9, 0.999) and learning rate chosen in the grid (0.0001, 0.0005, 0.001, 0.005). We do not use weight decay.
Computational Resources. The experiments are executed Nvidia DGX-1 GPU nodes equipped with 4 20-core Xeon E5-2698y processors, 512 GB of memory and 8 Nvidia V100 GPUs.
E
Additional Figures
1.2
1.0
Density
Covariance Vmne
SDE (Ours)
Shaped NN
Figure 7: Kernel density estimate and histogram of covariances from the covariance SDE in The-
orem 4.2 and shaped attention NN (Eq. 9). Simulated with 200, d 150, 7 1/v8, To
1.08 0.2, SDE step size 0.01, and 12 samples.
31

---

28.0
100
Simulated Median Stopping Time (Capped at 1) 1)
25.0
24.0
23.0
22.0
21.0
20.0
2-1.0
10-1
2-2.0-
2-3.0
2-4.0
2.5.0
2-6.0
2-7.0
2-8.0
10
2-9.0
2-10.0
72-Z
92-8
02-6
2
210
00-2
20-2
20-2
0002
Figure Median of the stopping time capped 1, of the shaped attention neural network with respect its parameters and To. Stopping time is defined t" d'/n with d' the maximum depth beyond which one of the eigenvalues of the covariance matrix exceeds 10 or drops below 10-4 Simulated with = d 200, and 100 samples used to estimate the median. To demonstrate the potential numerical instabilities, we had to choose an adversarial set of parameters: particular, an unrealistically large norm (approx. 10 TT) for the initial tokens Xo, which enlarges the eigenvalues of Vo to the order of 100.
32


Reviews:
[{"review_id": "0Qx2Vq6xf6", "reviewer": "Reviewer_J5pd", "summary": "**Please note that this is an emergency review, completed in a limited amount of time whilst I am traveling.**\n\nThis paper studies the properties of a class of attention-based networks, modified from the standard Transformer, in the limit in which depth and width tend proportionally to infinity. Technically, it extends recent work by Li, Nica, and Roy (NeurIPS 2022) on SDE-based descriptions of MLPs with random weights. The central result is a modification of the attention mechanism, referred to by the authors as shaped attention, that gives a non-trivial limit. ", "questions": "- Lines 46-47: In discussing the perturbative regime, work by Dyer and Gur Ari (2019) and by Zavatone-Veth, Canatar, Ruben, and Pehlevan (2021) should be cited alongside the work of Roberts, Yaida, and Hanin (2021). \n\n- Lines 48-53: For completeness, it might be useful to cite work on the proportional limit from random matrix theory, e.g. by G. Akemann and colleagues. \n\n- Line 243: \"Covariane\" -> \"Covariance\"\n\n- Lines 322-333: In relation to my comment about inference above, I think this paragraph would be substantially improved if the authors could offer a more concrete suggestion of how their approach could be leveraged to study training and generalization. \n\n- Lines 381-382: Why is the arXiv version of Li, Nica, and Roy cited in place of the NeurIPS 2022 publication? ", "limitations": "I think the authors adequately discuss the limitations of their work. ", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "In my abbreviated reading, I found the paper clear and enjoyable to read. I think it is a worthy contribution to the rapidly-growing literature on double-scaling limits of deep and wide networks, as it takes a substantial step closer to the architectures used in practice relative to prior works. The non-commutative scaling regime of ResNets studied by the authors nicely compliments recent work by Hayou and Yang on commutative infinite-width and -depth limits of residual networks. I found the observations regarding finite-time numerical instabilities to be particularly interesting. \n\nDue to time constraints, I have not had the chance to go through the proofs in detail. ", "weaknesses": "To my mind, the main weakness of this work is that it does not address inference. However, this reflects the larger issue that it is challenging to perform inference in conditionally-Gaussian processes (perhaps outside the nearly-Gaussian limit, where things can be studied perturbatively). Thus, this limitation does not significantly dampen my enthusiasm for the authors' work. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Summary \n\nThank you for your positive and constructive review.  We agree with the reviewer on the importance of studying the proportional limit for realistic architectures such as ResNets and attention models, as it paves the way for a faithful theory of neural networks beyond the infinite width limit. This is indeed one of the main motivations of our work. We also thank you for acknowledging the relevance of our results on resnets (in relation to Hayou and Yang), and for enjoying the analysis of the SDE’s stability. We will add or fix the citations you mentioned in the review. \n\nIn your review, you suggest a weakness of this work is that we do not develop inference. We would like to gently push back, and suggest that inference is ancillary to our goals, and that our work represents a key first step to building a more faithful model of training.\n\n## Inference missing?\n\nThe derivation of the NNGP in the infinite-width setting generated interest in the possibility of inference in this structure. In much the same way, as you've noticed, our work motivates inference in the conditional GP defined by the proportional limit. While we think inference in the NNGP was interesting, it is not our goal to study inference, so we do not view it as a limitation.\n\nIndeed, we believe that, looking back, the role of our work will be in defining boundary conditions for the yet-to-be-characterized processes describing training (by gradient methods), offering a model for training realistic networks that might offer us some theoretical insight, if it is tractable enough. This work is thus a first step towards understanding training.\n\nThat said, we did consider what inference might entail. Specifically, how might we sample an output at a specified input, conditional on a set of data points $\\{ x^\\alpha, y^\\alpha \\}_{\\alpha=1}^m$. We agree with you that this is not a trivial task. It is also orthogonal to our contributions. \n\nWhen we imagine the conditional diffusion dynamics, we expect this will create a new drift term in the SDE via Doob’s h-transform, which is not always easy to compute, since it can depend on the time marginals. In fact, this drift is closely related to score functions in diffusion models, which are known to be difficult to compute. After publication, we will be happy to discuss ideas around inference with any interested parties.\n\n## Training and Generalization \n\nFor future directions, we can consider the roadmap taken by infinite-width theories up to this point. For example, the calculations required for the Neural Tangent Kernel (NTK) are essentially identical to those of the Gaussian process results. While there are some technical difficulties, it is not unreasonable to foresee a similar SDE result for the NTK. This will immediately lead to calculations for the first step of training, which already allows to study which regimes allow for feature learning and learning rate transfer. \n\nGeneralization however, will require knowledge about the complete dynamics of training. This is currently not known for three-layer neural networks in the feature learning regime, and therefore we do not foresee infinite-depth limits to be any easier. The closest approach is the self-consistent dynamical mean field theory (DMFT) equations of Bordelon and Pehlevan [7], which we believe is possible to extend to infinite-depth with more careful calculations. Either way, it remains unclear how we can characterize the end of training, which is required for generalization results, but we suspect the high-level roadmap will serve as a helpful guide nonetheless. "}, {"review_id": "6QsyvESqC6", "reviewer": "Reviewer_4kLr", "summary": "This paper investigates the convariance matrix of a modified softmax-based attention model with skip connections in the proportional limit of infinite-depth-and-width. They develop a theory to describe  the limiting distribution of the convariance matrix using a SDE indexed by the depth-to-width ratio. Based on this result, they propose a shaped attention to achieve a well-defined limit and ensure the stablity of the network.", "questions": "Please see weaknesses.", "limitations": "The authors have adequately addressed the limitations, and there is no potential negative societal impact of their work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1 The authors derive neural convarince SDEs for ReNets and softmax-based attention, respectively. Then, they combine the attention and residual layers, and give the SDE for full transformer architeture. To the best of my knowledge, the theoretical contributions are novel. The analysis technique is standard and the proof is convincing.\n\n2 The theoretical contribution is also useful to guide the hyper-parameter tuning. The idea of shaped attention is very interesting, and the numerical simulations demonstrate the effectiveness of the proposed method.", "weaknesses": "1 All simulations are conducted at random initialization. It is better to conduct several experiments to show how the proposed initialization affect the training dynamics.\n\n2 There are several typos, e.g.,  lines 25, ''.One'' -> ''. One'', line 122 , ''importantproperties\"-> ''important properties\", line 242 \"remains -> remain\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive and constructive review. As the reviewer correctly points out, our research has a dual impact on the realms of deep learning. From a practical perspective, it guides architecture design (and hyperparameter tuning). From a theoretical perspective, it is a step towards a more precise theory of realistic neural networks, as our SDE is the first precise characterization of a Transformers covariance at initialization.\n\nTo begin, we will, of course, make our best effort to address all typos. Let us now address your main remark.\n\n## Simulations\n\nWith regards to simulations beyond initialization, we would like to first point to Appendix D, where we presented training results, beyond initialization, demonstrating that the modified (shaped Transformer) architecture (which can be scaled to infinite depth) trains to comparable accuracy as the standard (unshaped Transformer) architecture (which cannot be scaled). We stress that our architecture is freed of normalization layers.\n\nYour comment inspired us to offer more details on training dynamics, available in our \"one-pager\" (one-page PDF response) available to you. In the one-pager, we have simulated the gradient norm at initialization throughout different layers. For the queries and keys' parameters, we observe that while post-ln Transformers suffer from vanishing gradients, also the pre-ln Transformer's gradients severely diminish in magnitude with depth. This additional plot complements Figure 1 (left) of the paper in the sense that a high correlation between tokens in the forward pass is linked to vanishing gradients (see also [8]). On the other hand, the shaped Transformer has stable gradients, in the sense of depth-independent gradient magnitudes. \n\nNote that, in addition to our pre-training experiments in Appendix D mentioned above, we have added results for more seeds in our one pager. Just as in Appendix D, in these additional runs of our original training experiments, the shaped Transformer again matches the performance of well-tuned pre-ln Transformers. Note that we have not attempted to tune the training of the Shaped Transformer, so we might expect even better results given proper tuning, especially at greater depths, when we expect more instability for standard architectures. Regarding the training dynamics, we observe a higher trainability of our model with respect to the learning rate. This might be due to more resilience (in comparison to pre-ln Transformers) to entropy collapse, i.e. the regime where the softmax attends to only one token, in which case the network becomes untrainable (see Figures 5 and 6). \n\nIn summary, we believe that, combining our existing experiments in Appendix D and the new evidence presented in our one-pager (including the additional evaluation on GLUE), we have addressed the weakness you have raised. We hope you might reassess the potential impact of our paper accordingly."}, {"review_id": "lON1BVDPik", "reviewer": "Reviewer_5C73", "summary": "This paper studies a transformer architecture in the infinite width and depth limit where the ratio between the two remains constant. It is shown that, at initialization, the dynamics of the sample covariance matrix as it propagates through the layers converges to an SDE in the limit, where the SDE is indexed by the depth to width ratio. Moreover this limit is non-commutative and randomness preserving. The authors then propose the shaped attention - by constructing the attention matrix to resemble the identity function, and by properly scaling the residual connections and the softmax temperature, a more stable SDE is reached at the limit, theoretically allowing for more stable training.", "questions": "Is there any reason to prefer an SDE limit over an ODE one (assuming a feature learning limit rather than the NTK)? ", "limitations": "Limitations are adequately addressed", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents a rigorous analysis of the infinite depth-width limit, showing convergence of the dynamics (propagation of the data covariance at initialization) to a stochastic differential equation. \n2. The proposed shaped attention architecture is intriguing, meriting further investigation as to its potential benefits.\n", "weaknesses": "My main issue with the paper is the lack of any convincing experimental result on the potential benefit of the proposed architecture. While it seems reasonable that a more stable forward process would lead to a more stable training and better performance, this is not always the case in practice, and indeed the paper barely makes any effort to confirm this. The experiment that is presented in section D of the appendix is in my opinion not convincing for the following reasons: 1) only 2 random seeds are considered for generating the results which is not close to enough, and 2)  the relative improvement is inside the standard deviation anyway (for d=24).  Without further experimental confirmation, i'm afraid the entire motivation and analysis is shaken to the point of being pointless from a practical perspective (I am willing to increase my score if such evidence is presented).\nHaving said that, i do not think that an immediate practical application is mandatory for any paper, however given the abundance of literature on the subject of infinite width/depth limits i see the contribution of the paper at its current form as incremental, more fitting a workshop than the a main conference contribution. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your constructive feedback and, in particular, your recognition of our rigorous theoretical analysis.\n\nYour chief criticism regards \"the lack of any convincing experimental result on the potential benefit of the proposed architecture [to training in practice]\". Besides your concern that we don't use enough seeds (we've since run more seeds), it seems you were hoping for our new architecture to outperform the standard Transformer architecture and that, without this demonstration, the entire motivation and analysis is \"pointless from a practical motivation\".\n\nBelow we argue that you should adopt a different standard when judging our empirical work, and that our empirical findings are above and beyond what we might hope for when pursuing a proportional limit theory. \n\n##  Aside: no abundance of work on proportional limit\n\nFirst, we would like to quickly address the idea that there is an \"abundance of literature on the subject of infinite width/depth limits\". This is true of the infinite *width* limit, and we're only aware of one novel study of the Transformer in this limit. On the other hand, there are, to our knowledge, only 3 papers on the proportional infinite depth-and-width limit that can handle more than one data point. Moreover, the proportional limit holds promise based on empirical evidence [3,5] that it leads to a much more accurate model of finite networks. Note that no work has yet succeeded in giving a concise description of training dynamics in the proportional limit and this is a major open problem. In this context, our work is the first to identify a Transformer-like architecture that admits a nondegenerate proportional limit.\n\n## Role of empirical evaluation relates to motivation\n\nThere are a number of ways to motivate this work:\n\n1. A step towards a proportional limit model of Transformers and a theory for training dynamics. This is the principal motivation.\n\n2. An architectural modification to address known instability at initialization of Transformers at greater depth. This is a necessary step to achieve #1.\n\n3. Same as #2 but for improving training of Transformers at greater depth. This is a cherry on top.\n\nLimit theories have been immensely successful at deriving tractable models of neural networks. However, infinite-width theories have their limitations. This and related work on the proportional limit is pushing to derive more accurate models. \n\nThe standard Transformer architecture, up to our knowledge, does not admit a stable proportional limit. It is unstable due to \"degenerate signal propagation\" phenomena. Here is the key: *we must modify the Transformer architecture in order to use proportional limits to study Transformer-like architectures.*  \n\nGiven that it is *necessary* to find a stable-at-depth Transformer, we go looking for one. Our paper can be interpreted as a rigorous study of the degeneracy, leading us to a modification that solves the signal propagation problem in the forward pass and admits a proportional limit.\n\nIs any such modification equally good? Clearly no, because we might have solved the stability problem but created a useless architecture. Is our architecture useless? Clearly no. In fact, Appendix D offers an exciting realization: *the Shaped Transformer seems to train just as well*, at least at the scales we were able to evaluate. (Several better-funded labs have contacted us to ask questions and are launching internal, independent efforts to evaluate our ideas at GPT scale. This is the way open science can accelerate advances, in our opinion.)\n\n## Experiments\n\nIn light of the above, we never needed to beat unshaped Transformers: We wanted to find an architecture that *could* be studied in the proportional limit and was *worth* studying. In fact, we might expect better results with tuning. At greater scale and depth, we expect degeneracy to affect the standard architecture more and more, suggesting the potential for even more significant improvements.\n\nIn fact, thanks to your and the other reviewer’s suggestions, we are now in possession of more positive evidence. Results in our 1 page PDF suggest that shaping the attention model was *sufficient for a stable backward pass*, despite our theory addressing a stable forward pass. As you can see in Figure 1 of the one-pager, unshaped Transformers suffer significantly from vanishing gradients when depth is large in the query/key parameters. Meanwhile, *shaped Transformers do not have vanishing gradients at all*.\n\nFurther, at the request of reviewer h41x, we report performance on a few tasks from the GLUE benchmark. We find that, compared to the baseline of unshaped Transformers, the shaped Transformers perform slightly better than the pre-ln baselines (see answer to reviewer h41x). \n\nIn closing, the fact that we matched the performance of well-tuned Transformers, without LayerNorm and without re-tuning the optimization, is an exciting result. Our new results in the one-pager offer some even more exciting evidence that this architecture could be useful. Regardless, we think the architecture is clearly interesting enough to study in the proportional limit.\n\n## SDE vs ODE Limits\n\nThere is currently no significant evidence favoring either SDE or ODE limits. That being said, if we look at finite networks at scales that match the ODE 1/sqrt(depth) regime of [6], the SDE limit captures real randomness in the feature covariance matrix, which the ODE limit predicts is deterministic (see red curve in Figure 2). So we can conclude SDEs are better models here.\n\nWe believe the judgment will be determined by future results, such as which limit serves as a more faithful model of finite size networks during and after training, which are impossible to speculate on at the moment. At this time, we believe both directions are well-motivated, and should be developed concurrently and in a complementary manner.\n\nWe kindly hope you will be persuaded to reassess our work.\n"}, {"review_id": "Tj70NqmvUo", "reviewer": "Reviewer_h41x", "summary": "The paper analyzes the covariance matrix of transformer networks at initialization, which commonly serves to analyze the training stability of deep networks. Specifically, the limiting distribution at network initialization is characterized by an SDE in the infinite depth-and-width limit regime. To enable the analysis, the authors make minor modifications to the attention and MLP formulation, which do not seem to impact their functionality much. Finally, promising preliminary numerical results are presented, which indicate that the transformer network with modified attention and MLPs can train stably and produce non-trivial performance without normalization layers.", "questions": "* Could the proposed analysis be applied to decoder models with causally masked attention? Would masking the shaped attention lead to issues?\n* Could the dependence on the average token discussed in L264 be an artifact of the bias towards diagonal values in Eq. (9)? Are there potential issues with long-range dependencies when introducing this bias?\n* Assuming the authors used the embedding dim for the MLP hidden dim (both equal to 200), how would the results change when the depth, width and MLP hidden dim would be chosen with ratios like e.g. in BERT base?\n* Besides trainability without normalization layers, did the authors observe other signs of improved stability? For example fewer spikes in the loss curve or lower training loss variance?\n", "limitations": "There are no major limitations, please see “Weaknesses”. The authors are upfront that they only present an analysis for the network at initialization, and that more sophisticated analyses should be done in follow-up work.\n\nI’m not an expert in SDE-based analysis of neural networks and did not check the proofs in depth.\n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and relatively easy to follow. The analysis looks solid as far as I can tell. In particular, both the self-attention and FFN blocks, both with skip connections, are analyzed, and only relatively minor modifications are made to make the analysis tractable. The paper also presents promising experimental evidence that the proposed shaped attention can be trained stably without normalization layers.", "weaknesses": "I appreciate the analysis on BERT-style language modeling in the appendix. It would be great to train the model with a comparable setup to BERT-Base and evaluate it on GLUE. That would then be a first piece of evidence that the proposed modified could be deployed in practice.\n\nA curious thing is that the shaped attention in Eq. (9) can have negative values and values larger than 1, and no longer consists of proper PMFs. While this might not be an issue in practice, it is still a deviation from the principled formulation of the attention being normalized to 1. Can the authors think of an alternative formulation for which the normalization proprerty would still hold (or at least for which negative attention values could be avoided)?\n\nApart from that, I only see minor weaknesses. For example that the analysis is only for the network at initialization, but I agree with the authors that this can be addressed in follow-up work.\n\n\nTypos: Missing space in L25 (“train. One”), L122 (“important properties”)\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive and in-depth review! \n\nThere are a number of ways to motivate this work. One is to come up with practically useful modifications. But our principal motivation was to study Transformers using the approach of proportional limits. However, the Transformer architecture does not scale in a stable way to infinite depth, so we needed to modify the architecture to one we can actually study using this technique. The modification should also ideally not hinder training performance.\n\nWhat we find (in Appendix D) is that, on some small benchmarks, we obtain performance comparable to standard architectures. We recognize that a much more comprehensive empirical study would be required to convince the community that the architecture is ready for deployment, but it is beyond our budget and it was not our main objective. We believe that the experiments we run is ample empirical evidence that our modification is at least worth studying. Based on your suggestion to look at GLUE and some further experiments that look at gradients at initialization, we think there's even some evidence that our modification might eventually be useful to practitioners, but that's very much a cherry on top, and we think even our preliminary results should generate excitement that we've found an architecture worth studying in the proportional limit.\n\nWe now aim to address your remaining concerns. We hope that our response will be sufficient to view our work in an even more positive light. \n\n## GLUE Benchmark \n\nWe evaluated our pre-trained models on three small tasks from GLUE (due to time and resource constraints), and have included the results in Table 2 of our one-page pdf file. Overall, we found that the shaped Transformer performed slightly better on these three tasks than our baseline (best pre-LN Transformer). We also found these scores comparable to or outperforming the one-GPU-day-of-BERT-training results reported in [1]. However, we reiterate that the fairest comparison is against our own baselines, that match exactly the training procedure (data, number of pretraining steps, optimizer, etc..) that we adopt. \nAlthough these experiments are preliminary, we think this is an exciting signal that the shaped Transformers are worthy of a comprehensive empirical evaluation at larger scale.\n\n## Shaped Attention Not Proper PMFs\n\nIndeed, PMFs are not required in practice for intermediate layers. However, we understand that having PMF might be a requirement in certain applications. We would kindly like to bring to your attention an experiment recovering the original Softmax (Eq. 33, Appendix D). Essentially, we initialize the Transformer in the shaped setting, but linearly decay the centering terms of the Softmax after 4000 steps of training, thus recovering a vanilla Transformer without layernorm. Hence, the output of the attention block becomes a proper PMF once again. \n\nThe simulations depicted in Figure 3 show that we cannot simply remove centering, as otherwise depth quickly leads to degenerate correlations and covariance. Treating our shaped architecture as an initialization and letting the shaping parameters evolve through training is an exciting direction for future work. \n\n## Analysis After Initialization \n\nSo far, no work has been able to describe training dynamics in the proportional limit and this is one of the major open problems in the area. This is one of our longer-term goals and we are laying the foundations here: our results are analogous to the derivation of the NNGP in the infinite width limit. The NNGP was a key component of the NTK theory of training, which led to feature learning limits.\n\n## Masking Shaped Attention \n\nYes, we can extend shaped attention to the masked setting. The only modification will be for the centering matrix, which originally is the matrix of all $1/m$ entries, but would now have row $k$ taking the value of $1/k$ in the first $k$ entries, and 0 elsewhere. \n\n## Dependence on the Average Token \n\nThe average token term actually arises as a consequence of normalizing the Softmax output to $1$. If, for example, we were to replace the nonlinearity with say entry-wise sigmoids, then the average token terms would not show up in the Taylor expansion in Eq 8. \n\nAt the same time, we do not see potential issues with long-range dependencies, as we still perform attention to all the tokens, and any attention matrix can be learned through optimization (thus, long-range dependencies can still be learned and strengthened through training).\n\n## Varying Embedding Dimension \n\nYes, we can handle general, non-uniform widths in the network, as each layer’s width $n_\\ell^{-1}$ can be interpreted as a non-uniform step size. In this case, the total time (originally the depth-to-width ratio) would be replaced by the sum $\\sum_\\ell n_\\ell^{-1}$. See earlier work by Hanin and Nica [2,3] which handled this general case similarly.\n\n## Improved Training Stability \n\nFirst, we would like to point to our additional one-page PDF of Figures and Tables for plots on gradient norms. In Figure 1, gradients of the query/key parameters at initialization are vanishing in existing Transformer architectures, albeit milder with pre-LN. The *shaped Transformers do not suffer at all from vanishing gradients*.\n\nWe also observed a couple of other signs of improved training stability: \n\n1. When trained with Adam and learning rate warm-up, it appears that the shaped Transformer can handle larger learning rates without instability. In Figure 5 of the paper, we can see that unshaped Transformers are unstable at learning rate 0.001, while the shaped Transformer remains trainable.\n\n2. The entropy does not seem to collapse as much for shaped Transformers, i.e. the Softmax does not concentrate on one token (Figure 6). The collapse of entropy is known to be correlated with unstable training [4], which the shaped Transformers seem to avoid during training.\n"}]

Author rebuttal:
We want to thank all the reviewers for your detailed and constructive reviews. In particular, we appreciate the overwhelmingly positive feedback towards our theoretical results. Our work is the first key step towards a faithful theory of Transformer-like architectures. Of course, to study these architectures using proportional limits, we needed to develop architectural modifications that produced stable signal propagation. From the perspective of theory building, a merely trainable architecture would be an exciting outcome, as we would have arrived at a Transformer-like architecture worth studying via proportional limits. 

That said, we think it's quite exciting that our preliminary experiments find the Shaped Transformers train equally well as fine-tuned, unshaped Transformers on small (in modern terms) benchmarks. These results suggest that with further fine-tuning, especially in deeper regimes where the signals are even more degenerate, we believe the Shaped Transformer has the potential to improve even further. Still, we think these findings are cherries on top of what is a solid theoretical contribution worth sharing with the community.

To provide some further evidence that the architecture might not be merely interesting but also useful, we have summarized further simulations and experiments in our one-page pdf:

Figure 1. We measure the gradient norm of both shaped and unshaped Transformers at initialization, and find that shaping the attention mechanism successfully prevents gradients from vanishing, while unshaped Transformers suffer from rapidly vanishing gradients. 

Table 1. As requested by a reviewer, we have added more seeds (2 more, hence 4 in total) to our training experiments, replicating the same results as we already reported in Appendix D. 

Table 2. We have added experiments on three datasets from the GLUE benchmark. These demonstrate that our shaped Transformers outperform our pre-ln baseline and even previously reported one-GPU-day results. We plan to expand GLUE results, though we would like to reiterate that we do not need to outperform existing architectures to have identified an interesting model to study.

To summarize, we have demonstrated that the shaped Transformer has stable gradients at initialization and throughout training, and can match the performance of well-tuned unshaped Transformers with the same settings on some small benchmarks. 

# Experiment Details 

Figure 1. We take the same hyperparameters as Figure 1 of the submission, but with width $n=256$ and use causally masked attention for an autoregressive task predicting code data. In this setting we have sequence lengths of 128, and the error bars in the plot correspond to one standard deviation of each parameter matrix’s gradient norm, over a batch of 32 sequences.

Table 1. We add 2 more seeds for each experiment and again report the best-performing model across the considered learning rates, as explained in the paper. We have also increased the number of datapoints for the evaluation of the training statistics and re-performed the evaluation.  

Table 2. We take the checkpoints of the baselines and selected shaped transformers ($\gamma=0.2$ for the shallower model $d=18$, $\gamma=0.05$ for the deeper model $d=24$). Then, as per standard, we replace the final head with a randomly initialized MLP with 1 hidden layer and tanh activation function. We then fine-tune on each of three datasets from GLUE (COLA, MRPC, RTE) for 5 epochs and Adam optimizer $lr=5e-4$ with batch size 16, and report test evaluation metric for the corresponding task.

# References 

1. Geiping and Goldstein. "Cramming: Training a Language Model on a single GPU in one day." (2023).

2. Hanin and Nica. "Products of many large random matrices and gradients in deep neural networks." (2020).

3. Hanin and Nica. "Finite Depth and Width Corrections to the Neural Tangent Kernel." (2019).

4. Zhai, Likhomanenko, Littwin, Busbridge, Ramapuram, Zhang, Gu, and Susskind. "Stabilizing Transformer Training by Preventing Attention Entropy Collapse." (2023).

5. Li, Nica, and Roy. "The Neural Covariance SDE: Shaped Infinite Depth-and-Width Networks at Initialization." (2022).

6. Hayou and Yang. "Width and Depth Limits Commute in Residual Networks." (2023).

7. Bordelon, Blake, and Cengiz Pehlevan. "Self-consistent dynamical field theory of kernel evolution in wide neural networks." (2022):

8. Noci et al. "Signal propagation in transformers: Theoretical perspectives and the role of rank collapse." (2022)


Meta-review: The paper provides an abstract model for the dynamics of Gram matrices after applying a modified version of self-attention in mean field regimes. Such abstraction allows leveraging tools from stochastic calculus to study the inherent bias of random self-attention. The paper is well written and theoretical analyses are well motivated.
Decision: Accept (poster)