Title: The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit
Venue: NeurIPS 2023
Accepted: yes
Keywords: Deep Learning Theory

Covariance SDE

Attention Mechanism

Infinite-Depth-and-Width

Scaling Limit
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

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 Transformers, 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 diffusion 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.

1 Introduction

Pre-trained large language models have experienced a remarkable increase in popularity due to their eerily 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, is 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 to 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 a necessary condition for stable training [10-15]. In fact,

0

50

100

150

Layer

0.2

0.4

0.6

0.8

1

Mean Correlation

Shaped (Ours)

Unshaped

Unshaped+Pre-LN

1

0.5

0

0.5

1

Correlation

d

0

0.25

0.50

0.75

1

Density

SDE (Ours)

Shaped NN

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 on 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 a 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-width 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-41, 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, at initialization, has been shown to closely resemble the behaviour of finite architectures, yet still yield a relatively simple limiting description, expressible in terms of stochastic differential equations [40, 18]. In this work, we fully characterize the initial output distributions of a network with skip connections and Softmax-based attention mechanisms, in the proportional infinite-depth-and-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 linearizes and reduces the saturation of the Softmax, a known source of training instability in Transformers [47]. In order to model the feedforward layer of a Transformer's block, we extend existing results [18] to derive an SDE for the proportional limit of shaped-ReLU feedforward multi-layer perceptrons (MLPs) with skip connections. Combined, we fully characterize the output distribution of a Transformer with shaped non-linearities (Corollary 4.3).

0

20

40

60

80

100

120

140

Layer

0

5

10

15

20

25

Gradient Norm of Value Weights

W V

Shaped

Pre-LN

Post-LN

(a) Value Weights

W

V

ℓ

0

20

40

60

80

100

120

140

Layer

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

Gradient Norm of Query Weights

W Q

Shaped

Pre-LN

Post-LN

(b) Query Weights

W

Q

ℓ

Notably, our modification successfully prevents a poorly conditioned covariance matrix, whereas the vanilla Softmax-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 finitesize 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 new architecture in practice (see Section 5 and Appendix D).

In summary, our contributions are as follows:

We study the effect of skip connections in the proportional limit, showing that under a precise relation between the scaling parameters of the shortcut and residual branches, the feature covariance converges to the solution of a weighted version of the neural covariance SDE for MLPs (Theorem 3.2). The dependence on the depth-to-width ratio implies the existence of a stable non-commutative limit for residual networks, complementing the commutative limit studied in Hayou and Yang [50].

We propose shaped attention , where we modify the Softmax-based attention mechanism to be a perturbation of the identity. We demonstrate that shaped attention successfully prevents the degeneracy of correlation in contrast to existing Transformer architectures (Figure 1). The enhanced stability in the forward pass is reflected in the gradients, which are also stable with depth, as we empirically show in Figure 2.

For the proposed shaped attention architecture, we derive the neural covariance SDE characterizing the initial distribution in the proportional limit (Theorem 4.2). Consequently, we provide the first characterization of Transformer-type architectures, i.e. the shaped Transformer, in the large depth-and-width regime (Corollary 4.3).

We provide simulations to validate the theory and to interpret the effects of network hyperparamaters on the covariance matrix of the shaped Transformer. Specifically, we study finite time stability of the SDE and provide explicit guidance on hyperparameters to prevent numerical instability.

The paper is organized as 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 a 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 to 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 ℓ ∈ R m × n be the data matrix representing a sequence of m tokens embedded in n dimensions at layer ℓ ∈ [ 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 α ℓ ∈ R n is the α -th row of X ℓ . We consider the following attention model with residual connections:

where γ, λ ∈ [0 , 1] are parameters that control the strength of the shortcut and residual branch, respectively, W V ℓ ∈ R n × n is the weight matrix of the values, and A ℓ ∈ R m × m is the attention matrix. We consider Softmax-based scaled dot-product attention, where A ℓ has the form:

where the Softmax is applied row-wise, W Q ℓ , W K ℓ ∈ R n × n k are additional random weights, and τ is a temperature parameter, which controls the entropy of the distribution. Here we let all the weight matrices W Q ℓ , W K ℓ , W V ℓ have N (0 , 1) -iid entries. In the case where λ, γ = 1 , with the application of LayerNorm on the residual branch [48], and with τ = √ n k , we recover the attention block of the vanilla "Pre-LN" Transformer architecture [2]. Here we note that we pull the conventional n -1 / 2 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 λ 2 + γ 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 alignment of different inputs evolves across the network, as measured by the neural covariance V αβ ℓ := 1 n ⟨ x α ℓ , x β ℓ ⟩ (or ρ αβ := ( V αα ℓ V ββ ℓ ) -1 / 2 V αβ ℓ 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-widththen -depth limit, the correlation ρ αβ d of MLPs is known to converge to a fixed point independent of the input [11, 16]. 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-based attention models, Dong et al. [8] proved that the feature matrix X ℓ loses rank doubly exponentially fast with depth, and Noci et 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).

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 γ = 1 / √ depth, and postulates this scaling is sufficient to stabilize the neural

covariance with minimal assumptions on the activation function. Noci et al. [9] adopts this scaling to give precise formulas on the expected covariance of a Transformer at initialization. In this work, we consider γ constant in width and depth, and derive a complementary limiting result.

The Proportional Infinite-Depth-and-Width Limit. In the context of feed-forward MLPs, the output distribution with respect to a single input was studied in [37, 41], where it was shown that for the ReLU nonlinearity, the norm of the activations V αα converges to a log-normal random variable. To resolve the degeneracy of covariance and provide a characterization of output distributions for multiple inputs , M. Li et al. [18] shapes the ReLU by setting its slope 1 / √ width-away from linearity. In the proportional limit, the effect of the non-linearity accumulates over the d layers, and the covariance matrix V ℓ = [ V αβ ℓ ] αβ converges weakly to the solution of the SDE

where the formulae for coefficients b ReLU , Σ lin can be found in Theorem 3.2.

We note that the output neuron distributions are directly recovered as a conditional Gaussian with covariance V T for T = d n , in a similar spirit as the neural network Gaussian process (NNGP) results [20-22]. For example, the i -th output X out ,i conditioned on V d are asymptotically iid. N (0 , V T ) as d, n →∞ . The reader is referred to Appendix A for more technical background on the covariance SDE and the convergence result.

While the existing results are limited to initialization, we remind the reader that this is a necessary step before we can study training dynamics. In particular, the NNGP techniques developed for infinite-width 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 a simplified model composed of a shaped ReLU-activated layer and skip connections:

where σ s ( x ) := s + max( x, 0) + s -min( x, 0) is the shaped ReLU with slopes s ± := 1 + c ± n -1 / 2 for some constants c + , c -∈ R . We assume i.i.d weights ( W pre ℓ ) ij , ( W post ℓ ) ij iid ∼ N (0 , 1) , and c -1 = E σ s ( g ) 2 for g ∼ N (0 , 1) is 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 our 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 ℓ ∈ R m × n be the ℓ -th layer matrix of representations, and define the feature covariance as V ℓ = 1 n X ℓ X ⊤ ℓ . Let V ( n ) t = V ⌊ tn ⌋ ∈ R m ( m +1) / 2 be the the continuous time interpolation of the upper triangular entries as a vector. We say the covariance V ( n ) converges to V , if in the limit as n, d →∞ , d n → T , the process { V ( n ) t } t ∈ [0 ,T ] converges to { V t } t ∈ [0 ,T ] weakly in the Skorohod topology.

Theorem 3.2. Let X ℓ be the hidden layers of a ResNet defined in Eq. 4 with λ 2 + γ 2 = 1 , where both λ and γ do not depend on d, n . Then the feature covariance V ℓ converges to the solution of the following SDE (in the sense of Definition 3.1)

where b res ( V ) = γ 2 b ReLU ( V ) = γ 2 [ ν ( ρ αβ ) √ V αα V ββ ] α ≤ β with ρ αβ = V αβ ( V αα V ββ ) -1 / 2 and

furthermore, Σ res ( V ) = 2 γ 2 Σ lin ( V ) = 2 γ 2 [ V αδ V βω + V αω V βδ ] α ≤ β,δ ≤ ω

1.0

Residual Y

SDE

NN

一1/√depth

percentile

MLP

0.25

0.8

0.5

=y=1/v2

1

Density

95th

0.6

—NN

--SDE

Correlation

0.4

0.2

1/√d-scaling

0.0

-1

0.5

0

0.5

0.0

0.2

0.4

0.6

0.8

1.0

Residual y

Notice how the limiting SDE closely resembles the MLP case (Eq. 3), which is recovered exactly when γ = 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 γ < 1 dampens both the drift and the variance of the Brownian motion by γ 2 , thus it can be interpreted as a time change. In other words, the effect of γ at initialization is equivalent to reducing depth-to-width ratio, inline with existing intuitions that ResNets have a lower 'effective-depth' [56]. To visualize the stabilizing effect of γ on the distribution, in Figure 3 (right) we plot the 95th percentile correlation as a function of γ . The increasing trend indicates a larger probability of perfect alignment between two tokens. In Figure 3 (left) we plot the densities of both the residual SDE and the corresponding residual network for various values of γ . Notice how the samples from the SDE well-approximates the histogram of a 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 λ = 1 , γ = 1 / √ d , the depth and width limits commute . More precisely, the covariance V αβ converges to the same limit regardless of the order with respect to 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 / √ d suppresses the random fluctuations enough to vanish in the limit. On the other hand, our results show that for λ, γ constant in n, d , the random fluctuations are on the right order of O ( n -1 / 2 ) as in the MLP case (Eq. 3), hence they do not vanish in the simultaneous limit. The most notable difference is that our limiting regime is non-commutative as it depends on the depth-to-width ratio of the network. We remark that both regimes prevents degeneracy of covariance for residual architectures, forming two theories that complement each other.

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

Euler-discretization

where b, Σ are deterministic coefficients, and ξ ℓ are random vectors with zero mean and identity covariance. From here onwards, we use O ( n -p ) to denote a random variable Z such that n p Z has all moments bounded by universal constants (i.e. independent of n ). Since the update can be interpreted as discretization with step size n -1 , naturally the Markov chain converges to an SDE. We again note that a stable SDE implies a stable covariance structure for finite size networks.

To achieve the same goal for modified attention mechanisms, we consider a 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 Y ℓ to be the matrix of dot-products between queries, and keys, i.e. Y ℓ := X ℓ 1 √ n W Q ℓ 1 √ n W K, ⊤ ℓ X ⊤ ℓ .

More specifically, given a row y α ∈ R 1 × m of the logits Y ℓ ∈ R m × m , we can Taylor expand the row-wise Softmax in terms of τ -1 :

where y α := 1 m ∑ β y αβ 1 ⊤ and ( y α ) 2 is the vector with squared entries of y α , and 1 ∈ R m × 1 is the (column) vector of ones. We note in practice τ is often set to √ n k , which is often quite large and allows for asymptotic analysis [9].

We observe that the zero-th order term m -1 1 ⊤ is independent of τ . Considering the form of the attention block as A ℓ X ℓ 1 √ n W V ℓ , this yields an update that is no longer a small perturbation of V ℓ , regardless of how τ is chosen. Therefore, to form a Markov chain like Eq. 7, we actually require A ℓ to be approximately the identity matrix.

4.2 Shaped Attention

To shape the Softmax-attention mechanism as a perturbation of the identity matrix, we propose the following modifications which we call the shaped attention 5

The shaped attention presents three modifications to the Softmax attention in Eq. 2. Firstly, the zero-order term m -1 11 ⊤ of the Taylor expansion (Eq. 8) is removed as it causes a non-infinitesimal drift in the Markov Chain that ultimately leads to instability in the covariance (see Section 4.1). Secondly, we also observe that when τ is very large, the centered Softmax is a perturbation around zero. To recover an approximate Euler-update like in 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 B. He et al. [46]. Finally, the Softmax's temperature is chosen to scale as τ = τ 0 √ nn k , for some constant τ 0 > 0 , which guarantees a non-degenerate limit as ( d, n ) →∞ (Theorem 4.2). Note that the extra √ n 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 can be extended to include masking (e.g. casual masking) by centering each i-th row of the Softmax matrix by a different factor 1 /m i , where m i is the number of un-masked tokens in the i-th row.

1.0

Attention model

0.9

108

shaped (ours)

0.8

105

t²=nnk,center

Mean correlation

covariance

id,center

0.7

102

id, t² = nnk

0.6

10-1

only t2 = nnk

0.5

ean

only center

0.4

Mea

10-4

only id

0.3

10-7

0.2

10-10

0

20

40

60

80

100

120

140

0

20

40

60

80

100

120

140

layer

layer

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 are 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 ( n ) converges locally to V if the stopped process { V ( n ) t ∧ T r } t ≥ 0 converges to { V t ∧ T r } t ≥ 0 in the sense of Definition 3.1 for all stopping times of the form T r = inf { t > 0 : ∥ V t ∥ ≥ r } with r > 0 .

Let the covariance with respect to the average token be defined as V α ¯ x := m -1 ∑ m ν =1 V αν , and the average trace be ¯ V := m -1 ∑ m ν =1 V νν . We will need to compute a couple of important moments from the Taylor expansion terms of the Softmax (Lemma C.2)

We are now ready to state our main result.

Theorem 4.2. Let X ℓ be the hidden layers of a residual attention network defined in Eq. 1 with shaped attention in Eq. 9, parameters λ 2 + γ 2 = 1 and τ = τ 0 √ nn k , where λ, γ, τ 0 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)

where the drift has the following form

the diffusion coefficient is defined by Σ( V ) = γ 2 (2 -γ 2 )Σ lin ( V ) + γ 4 τ -2 0 [ A αβ,δω ] α ≤ β,δ ≤ ω , and

The drift depends on the shaped attention mechanism through S αδ,βω 1 and S αδ 2 , the moments of the first and second order terms of the Softmax's Taylor expansion. On the other hand, the diffusion term depends on the attention solely through S 1 , present in the additional term A αβ,δω . The presence of A αβ,δω 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 is compared through dot products to form the attention weights. Finally, notice how the residual branch parameter γ 2 has a dampening effect on the scale of both the drift and the diffusion in a 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:

where A ℓ is 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 X ℓ be the hidden layers of a shaped transformer defined in Eq. 11 with parameters λ 2 + γ 2 = 1 and τ = τ 0 √ nn k , where λ, γ, τ 0 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)

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 t ) in Theorem 4.2 is cubic in the entries of V t . Whenever the drift is not Lipschitz as 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.29]). In fact, MLPs with smooth activations also yield non-Lipschitz drift coefficients as seen in M. Li et al. [18].

However, locally Lipschitz coefficients are sufficient to guarantee the existence of local solutions, in the sense of up to a stopping time [59, Proposition 6.9]. Not only does this fact help us establish a 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 a stopping time.

To demonstrate the potential numerical instabilities, we had to choose an adversarial set of parameters: in particular, an unrealistically large norm (approx. 10 √ n ) for the initial tokens X 0 , which enlarges the eigenvalues of V 0 to the order of 100 . Given these initial conditions and a large residual connection weight γ , 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 γ, τ 0 and the depth-to-width ratio of the network. We demonstrate the effect of tuning γ on both sample trajectories of the maximum eigenvalue of V ℓ and the stopping time in Figure 5. As we may intuitively expect, tuning γ smaller will delay the time scale of numerical instabilities, hence allowing for larger depth networks to remain stable.

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 a novel attention mechanism, which successfully stabilizes the covariance structure in arbitrarily deep Transformers (e.g. Figure 1). The natural next

1019 .

1018.

y= 0.735

1.0

median

>

1017.

1016 .

y=0.214

1

10thpercentile

Eigenvalue

1015

Time (Capped

0.8

1014

1013

1012.

0.6-

1011

1010.

109 .

Maximum

108.

Stopping

0.4

107.

106

105

0.2

104.

103.

102.

0.0*

0

25

50

75

100

125

150

175

200

0.0

0.2

0.4

0.6

0.8

1.0

layer

step is to investigate the scaling of gradients in the infinite-depth-and-width limit. As Yang et 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 cost in the process. Given the existence of an infinite-depth-and-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.

Preliminary Experiments. Although this work is primarily theoretical, it is important to consider whether or not the proposed architecture is useful in practice. Given limited computing resources, we chose to only briefly test the feasibility of training the shaped Transformer. Nevertheless, our preliminary experiments show promising results when it comes to training stability. In particular, the shaped Transformer (without LayerNorm) does indeed train at 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. As mentioned in the introduction, the limitations of infinite-width NTK theories motivates our study of the proportional infinite-depth-and-width limit. In particular, to address many of the open problems in 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 of training as well. Another promising indicator is 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 us 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 a theory of generalization, as demonstrated by the infinite-width approach [67]. In light of our results, we believe that our theory sets the stage for future work on training and generalization in deep learning.

References

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.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need . 2017. arXiv: 1706.03762 .

J. Hestness, S. Narang, N. Ardalani, G. Diamos, H. Jun, H. Kianinejad, M. Patwary, M. Ali, Y. Yang, and Y. Zhou. Deep learning scaling is predictable, empirically . 2017. arXiv: 1712.00409 .

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 .

J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al. Training compute-optimal large language models . 2022. arXiv: 2203.15556 .

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 .

A. Aghajanyan, L. Yu, A. Conneau, W.-N. Hsu, K. Hambardzumyan, S. Zhang, S. Roller, N. Goyal, O. Levy, and L. Zettlemoyer. Scaling Laws for Generative Mixed-Modal Language Models . 2023. arXiv: 2301.03728 .

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.

L. Noci, S. Anagnostidis, 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 .

B. Poole, S. Lahiri, M. Raghu, J. Sohl-Dickstein, and S. Ganguli. 'Exponential expressivity in deep neural networks through transient chaos'. In: Advances in Neural Information Processing Systems 29 (2016).

S. S. Schoenholz, J. Gilmer, S. Ganguli, and J. Sohl-Dickstein. 'Deep Information Propagation'. ICLR . 2017.

G. Yang and S. S. Schoenholz. 'Mean field residual networks: on the edge of chaos'. Advances in Neural Information Processing Systems . 2017, pp. 2865-2873.

S. Hayou, A. Doucet, and J. Rousseau. 'On the impact of the activation function on deep neural networks training'. International Conference on Machine Learning . PMLR. 2019, pp. 2672-2680.

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.

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.

J. Martens, A. Ballard, G. Desjardins, G. Swirszcz, 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 .

G. Zhang, A. Botev, and J. Martens. Deep Learning without Shortcuts: Shaping the Kernel with Tailored Rectifiers . 2022. arXiv: 2203.08120 .

M. Li, M. Nica, and D. Roy. 'The neural covariance SDE: Shaped infinite depth-and-width networks at initialization'. In: Advances in Neural Information Processing Systems 35 (2022), pp. 10795-10808.

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: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer . 2022. arXiv: 2203.03466 .

R. M. Neal. 'Bayesian learning for neural networks'. PhD thesis. University of Toronto, 1995.

J. Lee, Y. Bahri, R. Novak, S. S. Schoenholz, J. Pennington, and J. Sohl-Dickstein. 'Deep Neural Networks as Gaussian Processes'. ICLR . 2018.

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 .

J. Hron, Y. Bahri, J. Sohl-Dickstein, and R. Novak. 'Infinite attention: NNGP and NTK for deep attention networks'. International Conference on Machine Learning . PMLR. 2020, pp. 4376-4386.

A. Jacot, F. Gabriel, and C. Hongler. 'Neural tangent kernel: Convergence and generalization in neural networks'. Advances in Information Processing Systems (NeurIPS) . 2018. arXiv: 1806.07572 .

G. Yang and E. J. Hu. Feature learning in infinite-width neural networks . 2020. arXiv: 2011. 14522 .

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.

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), E7665E7671.

L. 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).

G. M. Rotskoff and E. Vanden-Eijnden. Trainability and Accuracy of Neural Networks: An Interacting Particle System Approach . 2018. arXiv: 1805.00915 .

S. Yaida. 'Non-Gaussian processes and neural networks at finite widths'. Mathematical and Scientific Machine Learning . PMLR. 2020, pp. 165-192.

E. Dyer and G. Gur-Ari. 'Asymptotics of wide networks from feynman diagrams'. In: arXiv preprint arXiv:1909.11304 (2019).

D. A. Roberts, S. Yaida, and B. Hanin. The principles of deep learning theory . Cambridge University Press, 2022.

J. Zavatone-Veth and C. Pehlevan. 'Exact marginal prior distributions of finite Bayesian neural networks'. In: Advances in Neural Information Processing Systems 34 (2021).

J. Zavatone-Veth, 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.

B. Hanin. Correlation Functions in Random Fully Connected Neural Networks at Finite Width . 2022. arXiv: 2204.01058 .

E. Dinan, S. Yaida, and S. Zhang. Effective Theory of Transformers at Initialization . 2023. arXiv: 2304.02034 .

B. Hanin and M. Nica. 'Products of many large random matrices and gradients in deep neural networks'. In: Communications in Mathematical Physics (2019), pp. 1-36.

B. Hanin and M. Nica. Finite depth and width corrections to the neural tangent kernel . 2019. arXiv: 1909.05989 .

Z. Hu and H. Huang. 'On the Random Conjugate Kernel and Neural Tangent Kernel'. International Conference on Machine Learning . PMLR. 2021, pp. 4359-4368.

M. B. Li, M. Nica, and D. Roy. 'The future is log-Gaussian: ResNets and their infinite-depthand-width limit at initialization'. In: Advances in Neural Information Processing Systems 34 (2021), pp. 7852-7864.

L. 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).

B. Hanin and G. Paouris. 'Non-asymptotic results for singular values of Gaussian matrix products'. In: Geometric and Functional Analysis 31.2 (2021), pp. 268-324.

G. Akemann, Z. Burda, and M. Kieburg. 'From integrable to chaotic systems: Universal local statistics of Lyapunov exponents'. In: Europhysics Letters 126.4 (2019), p. 40001.

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.

T. Jiang and Y. Qi. 'Spectral radii of large non-Hermitian random matrices'. In: Journal of Theoretical Probability 30.1 (2017), pp. 326-364.

B. He, J. Martens, G. Zhang, A. Botev, A. Brock, S. L. Smith, and Y. W. Teh. Deep Transformers without Shortcuts: Modifying Self-attention for Faithful Signal Propagation . 2023. arXiv: 2302.10322 .

M. Dehghani, J. Djolonga, B. Mustafa, P. Padlewski, J. Heek, J. Gilmer, A. Steiner, M. Caron, R. Geirhos, I. Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters . 2023. arXiv: 2302.05442 .

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.

J. L. Ba, J. R. Kiros, and G. E. Hinton. 'Layer normalization'. In: (2016). arXiv: 1607.06450 .

S. Hayou and G. Yang. Width and Depth Limits Commute in Residual Networks . 2023. arXiv: 2302.00453 .

S. Hayou, E. Clerico, B. He, G. Deligiannidis, A. Doucet, and J. Rousseau. 'Stable resnet'. International Conference on Artificial Intelligence and Statistics . PMLR. 2021, pp. 13241332.

W. Tarnowski, P. Warchoł, S. Jastrzebski, J. Tabor, and M. Nowak. 'Dynamical isometry is achieved in residual networks in a universal way for any activation function'. The 22nd International Conference on Artificial Intelligence and Statistics . PMLR. 2019, pp. 22212230.

S. Hayou. On the infinite-depth limit of finite-width neural networks . 2022. arXiv: 2210. 00688 .

G. Yang. 'Tensor programs ii: Neural tangent kernel for any architecture'. In: arXiv preprint arXiv:2006.14548 (2020).

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.

A. Veit, M. J. Wilber, and S. Belongie. 'Residual networks behave like ensembles of relatively shallow networks'. In: Advances in neural information processing systems 29 (2016).

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 .

I. Karatzas and S. Shreve. Brownian motion and stochastic calculus . V ol. 113. Springer Science &Business Media, 2012.

J. Miller. Stochastic Calculus (Lecture Notes) . http://www.statslab.cam.ac.uk/ /~jpm205/teaching/lent2016/lecture_notes.pdf . 2015.

B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. 'When do neural networks outperform kernel methods?' In: Advances in Neural Information Processing Systems 33 (2020), pp. 14820-14830.

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.

J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, and G. Yang. 'High-dimensional asymptotics of feature learning: How one gradient step improves the representation'. Int. Conf. Learning Representations (ICLR) . 2022. URL: https://openreview.net/forum?id= akddwRG6EGi .

A. Damian, J. Lee, and M. Soltanolkotabi. 'Neural networks can learn representations with gradient descent'. Conference on Learning Theory . PMLR. 2022, pp. 5413-5452.

A. Mousavi-Hosseini, S. Park, M. Girotti, I. Mitliagkas, and M. A. Erdogdu. Neural Networks Efficiently Learn Low-Dimensional Representations with SGD . 2022. arXiv: 2209.14863 .

E. Abbe, E. Boix-Adsera, and T. Misiakiewicz. SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics . 2023. arXiv: 2302.11055 .

R. Berthier, A. Montanari, and K. Zhou. Learning time-scales in two-layers neural networks . 2023. arXiv: 2303.00055 .

P. L. Bartlett, A. Montanari, and A. Rakhlin. 'Deep learning: a statistical viewpoint'. In: Acta numerica 30 (2021), pp. 87-201.

O. Kallenberg. Foundations of Modern Probability . Probability theory and stochastic modelling. Springer, 2021. ISBN: 9783030618728.

HuggingFace. Wikipedia data set . https://huggingface.co/datasets/wikipedia .

HuggingFace. Bookcorpus data set . https://huggingface.co/datasets/bookcorpus .

D. P. Kingma and J. Ba. 'Adam: A method for stochastic optimization'. In: arXiv preprint arXiv:1412.6980 (2014).

A. Wang, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman. 'GLUE: A multi-task benchmark and analysis platform for natural language understanding'. In: arXiv preprint arXiv:1804.07461 (2018).

HuggingFace. Hugging face Bert implementation . https://github.com/huggingface/ transformers/blob/main/src/transformers/models/bert/modeling_bert.py .

A Preliminaries: Covariance SDE Framework

In this section, we will review existing results on Markov chain convergence to an 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 J 1 -topology, or just the Skorohod topology for short [68, Appendix 5]. The topology is used to 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 D R + ,S be the space of càdlàg functions (right continuous with left limits) from R + → S . We use x n ul - → x to denote locally uniform convergence (uniform on compact subsets of R + ), and consider the class of bijections λ on R + so that λ is strictly increasing with λ 0 = 0 (can be interpreted as a time change).

Definition A.1. We define Skorohod convergence x n s - → x on D R + ,S if there exists a sequence of bijections λ n satisfying above conditions and

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 n that satisfy the following type of Euler updates

where ξ ℓ are iid random variables with zero mean and identity covariance, then we can interpolate this process in continuous time with X n t = Y n ⌊ tn ⌋ and show that as n →∞ , we have that X n converges to the solution of the following SDE (weakly with respect to the Skorohod topology)

Our next theorem essentially weakens this result in several ways. Firstly, we don't need to take the step size n -1 , but instead replace it with n -2 p for all p > 0 . Next, we can allow the update to contain higher order terms that vanish, in particular, any terms of order O ( n -2 p -δ ) for all δ > 0 . Here, we remind the readers that O ( n -p ) denotes a random variable Z such that n p Z has all moment bounded by a universal constant (independent of n ). Thirdly, we can allow b ( y ) to be random, or more precisely replace it with ̂ b ( y, ω n ) such that E ̂ b ( y, ω n ) = b ( y ) . Fourthly, we can also allow b, σ to be a sequence b n , σ n such that they converge to b, σ on uniformly on compact sets. Finally, we can also weaken the topology of convergence to locally, that is all processes are stopped by a stopping time as in Definition 4.1.

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 n be a discrete time Markov chain on R N defined by the following update for p, δ > 0

where ξ n ℓ ∈ R N are iid random variables with zero mean, identity covariance, and moments uniformly bounded in n . Furthermore, ω n ℓ are also iid random variables such that E [ ̂ b n ( Y n ℓ , ω n ℓ ) | Y n ℓ = y ] = b n ( y ) and ̂ b n ( y, ω n ℓ ) has uniformly bounded moments in n . Finally, σ n is a deterministic function, and the remainder terms in O ( n -2 p -δ ) have uniformly bounded moments in n .

Suppose b n , σ n are uniformly Lipschitz functions in n and converges to b, σ uniformly on compact sets, then in the limit as n →∞ , the process X n t = Y n ⌊ tn 2 p ⌋ converges in distribution to the solution of the following SDE in the Skorohod topology of D R + , R N

Suppose otherwise b n , σ n are only locally Lipschitz (but still uniform in n ), then X n converges locally to X in the same topology (see Definition 4.1). More precisely, for any fixed r > 0 , we consider the stopping times

then the stopped process X n t ∧ τ n converges in distribution to the stopped solution X t ∧ τ of the above SDE in the same topology.

We will briefly recall the main result of M. Li et al. [18] next. As mentioned earlier in the text, the setting is for MLPs defined as follows

where σ s ( x ) = s + max( x, 0) + s -min( x, 0) with s ± = 1 + c ± √ n , c -1 = E σ s ( g ) 2 for g ∼ N (0 , 1) , n is the width of the network, and W ℓ,ij ∼ N (0 , 1) are iid random weights.

Then it was shown that in the limit as n, d → ∞ with d n = T > 0 , the covariance matrix V ℓ = 1 n X ℓ X ⊤ ℓ (or equivalent the post activation c n σ s ( X ℓ ) σ s ( X ℓ ) ⊤ since the activation becomes infinitesimal) converges the solution of the following SDE in the Skorohod topology [18, Theorem 3.2]

where the coefficients were given in Theorem 3.2.

We remark that if the output is defined as X out = 1 √ n W out X d ∈ R n out , then we can recover the output distribution as

where we treated X out as a vector in R out .

B SDE for Residual Network

Recall that we adopt the following model:

where σ s ( x ) := s + max( x, 0) + s -min( x, 0) is the shaped ReLU with slopes s ± := 1 + c ± √ n for some constants c + , c -∈ R . Welet the weights be ( W pre ℓ ) ij ( W post ℓ ) ij iid ∼ N (0 , 1) and c -1 = E σ s ( g ) 2 for g ∼ N (0 , 1) is the He initialization constant [55].

From here onwards, we will define the filtration F ℓ = σ ( { X k } k ∈ [ ℓ ] ) , where σ ( · ) denotes the sigmaalgebra generated by the random variable. Furthermore, we will define the conditional expectation E ℓ [ · ] = E [ · |F ℓ ] .

We are interested in studying the neural covariance , i.e. V αβ ℓ := c n ⟨ x α ℓ , x β ℓ ⟩ where ℓ ∈ [ d ] indexes the layers and d is the depth. In particular, we are interested in understanding the simultaneous limit ( n, d ) →∞ , 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 ) := 1 √ n σ s (√ c n XW pre ) W post , we can write V αβ as:

For the cross product terms, we get:

For the term in γ 2 :

We will define the following terms

and

Hence we get the following update for V αβ ℓ := c n ⟨ x α ℓ , x β ℓ ⟩ :

It is easy to see that the cross product terms have (conditional) mean zero. For the term with γ 2 :

where K 1 ( ρ αβ ) := E [ σ s ( g α ) σ s ( g β )] where g α , g β ∼ N (0 , 1) and E [ g α g β ] = ⟨ x α ,x β ⟩ ∥ x α ∥∥ x β ∥ . Here we recall E ℓ [ · ] = E [ · |F ℓ ] is the conditional expectation given the sigma-algebra generated by F ℓ = σ ( { X k } k ∈ [ ℓ ] ) .

M. Li et al. [18] showed that if the non linearity scaling exponent is p = 1 / 2 , then:

where ν ( ρ ) = ( c + + c -) 2 2 π ( √ 1 -ρ 2 + ρ arccos ( ρ ) ) . Using this result, and summing and subtracting the mean of T 2 :

Using λ 2 + γ 2 = 1 :

Furthermore, we need second order moments of T and T . We derive them in the following Lemmas:

1 2 Lemma B.1.

Proof. Recall the definition of T αβ 1 :

We have that:

Lemma B.2.

Proof. Recall the definition of T αβ 2 :

For the mean, using the independence between W pre and W post , we have that:

which is the desired result. The final expression is the result of the aforementioned expansion for K 1 [18]. For the covariance, we have that:

Let's look at each term of the sum separately. For the first term we have that:

̸

where we have defined the fourth moment of the shaped activation: K 2 ( ρ αβδω ) := E [ σ s ( g α ) σ s ( g β ) σ s ( g δ ) σ s ( g ω ) ] for which it holds that [18, Lemma C.2]:

The other two summands can be solved similarly:

Hence, summing the three terms:

Now, subtracting E ℓ [ T αβ 2 ] E ℓ [ T δω 2 ] , using the aforementioned expansions for K 1 and K 2 :

where in the final step we have used the fact that c = 1 + O ( n -1 / 2 ) . This completes the proof.

Finally we also have that T 1 and T 2 are uncorrelated, i.e.

Lemma B.3.

Proof. It is easy to see that E ℓ T αβ 1 T δω 2 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 λ 2 + γ 2 = 1 , where both λ and γ do not depend on d, n . Then the feature covariance V ℓ converges to the solution of the following SDE (in the sense of Definition 3.1)

where b res ( V ) = γ 2 b ReLU ( V ) = γ 2 [ ν ( ρ αβ ) √ V αα V ββ ] α ≤ β with ρ αβ = V αβ ( V αα V ββ ) -1 / 2 and

Proof. We know that:

Proof.

Using λ 2 + γ 2 = 1 , and summing and subtracting the mean of T 2 , we have that:

Using Lemma B.2 for the mean of T 2 , we have that:

which gives us an expression in the right Markov Chain form with drift term √ V αα V ββ ν ( ρ ) . For the diffusion term, we need to compute the covariance between two different neural covariances V αβ ℓ +1 , V δω ℓ +1 . Using Lemmas B.1 to B.3, we have

where we use Cov ℓ to denote the covariance conditioned on the sigma-algebra F ℓ = σ ( { X k } k ∈ [ ℓ ] ) . 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:

where W K ℓ , W Q ℓ ∈ R n × n k are Gaussian matrices with unit variance entries, and n k is the queries and keys' dimension. Here we study the first two moments of Y ℓ . In particular, note that E ℓ [ Y ℓ ] = 0 , where E ℓ [ · ] = E [ · |F ℓ ] denotes the conditional expectation given the sigma-algebra generated by F ℓ = σ ( { X k } k ∈ [ ℓ ] ) .

For the second moment we have the following Lemma.

Lemma C.1. Let

be the dot-product attention parametrized by W K ℓ , W Q ℓ ∈ R n × n k . Then:

For the expectation, we have that:

where we recall W B = W K W Q, ⊤ .

Hence:

C.2 Shaping the Softmax

Recall that we have the following model:

For the above model, V ℓ has the following form:

In order to have infinitesimal updates in V ℓ , we would intuitively like the attention matrix to be of the form A αβ = δ αβ + O ( n -1 ) . In the case of the usual Softmax based attention, we have:

where τ -1 is a temperature parameter that regulates the entropy of the resulting categorical distribution; τ is often chosen to be large (scale as a power of n or d ), as low temperature results in unstable training.

To transform the Softmax into the desired form of A αβ = δ αβ + O ( n -1 ) , we first center the attention matrix around the identity matrix:

and then examine the Taylor-expansion of the Softmax w.r.t τ -1 :

The first few terms of the expansion provides a good approximation of A ℓ when τ is large. Observe that for fixed m , the zero order term 1 m 11 ⊤ is not of O ( n -1 ) . In particular, suppose that we use the attention model in Eq. 25. We can show that the expectation of the term in γ 2 (responsible for the drift) has the following form:

where the terms scaled by 1 m leads to large incremental updates in expected value of V ℓ w.r.t the previous layer, and precludes the possibility of convergence to an SDE. Hence, we choose to re-center the Softmax by removing the term 1 m 11 ⊤ , as follows:

which admits the following Taylor-expansion:

Hence, up to third order:

For sufficiently large τ , the formulation above allows infinitesimal updates in V ℓ , and as we will show rigorously in the rest of this section, permits convergence to an SDE.

C.3 Lemmas on Moments of Shaped Attention

Define:

Hence, A ℓ can be written as:

We now compute the moments of A ℓ . We define the following quantities:

where we recall E ℓ [ · ] = E [ · |F ℓ ] is the conditional expectation given the sigma-algebra generated by F ℓ = σ ( { X k } k ∈ [ ℓ ] ) .

Lemma C.2 (Moments of Taylor Expansion) .

where ¯ V = 1 m ∑ ν V νν and ¯ x = 1 m ∑ ν x ν is the average token.

Proof. Using Lemma C.1 and linearity of expectation:

and:

where ¯ V = 1 ∑ V ββ .

Lemma C.3.

Proof.

Lemma C.4.

Proof.

Using the linearity of expectation:

Note that the terms above can be computed using Lemma C.2.

Lemma C.5.

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:

For the above model, V ℓ has the following form:

We define:

and

Proof.

Let's look at the first summand:

Hence, the expression for V ℓ simplifies to:

We need to compute the moments for these two quantities:

Lemma C.6 (Moments of T 1 ) .

Hence, in expectation with respect to W :

An identical argument can be made for the remaining three summands. Hence, taking expectation with respect to the Softmax weights:

Now, using Lemma C.3:

Proof.

Taking expectation with respect to W :

Taking expectation w.r.t the Softmax weights, we get the desired result.

For second moment, we can take the conditional expectation:

where we recall E ℓ [ · ] = E [ · |F ℓ ] is the conditional expectation given the sigma-algebra generated by F ℓ = σ ( { X k } k ∈ [ ℓ ] ) .

Using Isserlis Theorem, we have that:

Hence:

By taking expectation w.r.t the Softmax parameters, we get the desired result.

Proof. Using Lemma C.7, we have that:

Now we can use Lemma C.3 and Lemma C.4 to compute the moments of A . For the second summand, we simply have:

hence:

For the first summand, recall from Lemma C.5 that:

Hence:

We are now ready to re-state and proof of Theorem 4.2.

Theorem 4.2. Let X ℓ be the hidden layers of a residual attention network defined in Eq. 1 with shaped attention in Eq. 9, parameters λ 2 + γ 2 = 1 and τ = τ 0 √ nn k , where λ, γ, τ 0 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)

where the drift has the following form

the diffusion coefficient is defined by Σ( V ) = γ 2 (2 -γ 2 )Σ lin ( V ) + γ 4 τ -2 0 [ A αβ,δω ] α ≤ β,δ ≤ ω , and

Proof. Recall that:

Drift. Summing and subtracting E ℓ [ T 2 ] (and using Lemma C.7), we have that:

where we recall E ℓ [ · ] = E [ · |F ℓ ] is the conditional expectation given the sigma-algebra generated by F ℓ = σ ( { X k } k ∈ [ ℓ ] ) .

Re-stating Lemma C.3, we have that:

Plugging in the expression, and using λ 2 + γ 2 = 1 , we have that the drift is:

From here, it is evident that in order to have the drift scaling as O (1 /n ) we need to choose:

where τ 0 > 0 is a constant.

Covariance. Recall that:

Furthermore, we have set λ 2 + γ 2 = 1 and τ 2 = τ 2 0 nn k to have the right scaling for the drift.

What's left to is to compute the conditional covariance for V ℓ +1 given F ℓ . Noting that T αβ 1 , T δω 2 are uncorrelated, i.e. E ℓ [ T αβ 1 T δω 2 ] = 0 (similarly to the case of Resnet with shaped ReLU Lemma B.3), we have that:

where we use Cov ℓ to denote the conditional covariance given the sigma-algebra generated by F ℓ = σ ( { X k } k ∈ [ ℓ ] ) .

Using Lemma C.6 and Lemma C.8, we have that:

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 ℓ be the hidden layers of a shaped transformer defined in Eq. 11 with parameters λ 2 + γ 2 = 1 and τ = τ 0 √ nn k , where λ, γ, τ 0 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)

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

where ξ ℓ , ξ ′ ℓ are independent zero mean and identity covariance random vectors.

Since in the limit, we have that either updates are infinitesimal, i.e.

then we can write b ReLU ( U ℓ ) = ̂ b ReLU ,n ( V ℓ , ω ℓ ) , where eventually we have that

so it will not contribute towards affecting the limit. Here we recall E ℓ [ · ] = E [ · |F ℓ ] is the conditional expectation given the sigma-algebra generated by F ℓ = σ ( { X k } k ∈ [ ℓ ] ) . We can treat Σ lin ( U ℓ ) similarly to get the Markov chain update

which converges to the following SDE with two Brownian motions using Proposition A.2

Observe that since the two Brownian motions B t , B ′ t are independent, it's equivalent to write

We recover the desired results from considering a more general form of 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. In particular, we consider a pre-training masked language modeling task, where we mask 15% of the tokens. We use a subset of the English Wikipedia 20220301.en and English bookcorpus datasets [69, 70]. As a baseline, we adopt a 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 to s -= 1 -1 / √ n instead of 0 . We then incorporate our shaped attention by replacing the attention mechanism as dictated in Eq. 9. We also add scalar multipliers γ 1 , γ 2 ∈ R both to the identity and centering terms of the shaped attention:

and propose two ways to set γ 1 , γ 2 during training. In both alternatives we initialize γ 1 , γ 2 = 1 , thus leveraging the stability properties of shaped attention at initialization. During training, we either:

Recover . Linearly decrease γ 1 , γ 2 to zero with in the first 4000 steps, thus recovering the standard attention layer. Apply the same schedule for the shaped-ReLU slope s -, recovering the usual ReLU activation. This approach recovers the vanilla Transformer architecture (without LayerNorm).

Learn . Learn all the shaping parameters γ 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 or Learn . We report that without either adjustment, shaped attention is still trainable but at much slower rates.

We also incorporate the 1 √ n factor into the initialization of the queries and keys weights by decreasing the variance of their entries by a factor of 1 n . This allows us to not re-tune the learning rate for the queries and keys. We stress that at at initialization, the two formulations ( τ = √ nn k and τ = √ n k with decreased weight's variance) are equivalent. In both alternatives we train all the skip connection parameters λ , γ , and initialize γ in the grid (0 . 05 , 0 . 1 , 0 . 2) and set τ 0 = 1 . We report training instabilities (loss divergence) for larger values of γ . All models -including 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 100 K optimization steps, averaged over 4 random seeds. All the other experimental details can be found in Appendix D.1.

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

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, as it has been observed that the Transformer's training is unstable in the low-entropy regime [57]. The entropy is calculated for each row of the Softmax matrix, and it is 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 a 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 a subset of the English Wikipedia 20220301.en and English bookcorpus datasets [69, 70]. The sentences are tokenized using by pre-training a tokenizer on the training set. We use a vocabulary size of 32000 , and a maximum sequence length of 128 tokens.

Model parameters. We use an embedding size of n = 768 and 8 multi-attention heads. The batch size is fixed to 32 sequences. All the initial weights are sampled from N (0 , n -1 ) , with the exception of the queries and keys' weights W K , W Q in the shaped attention case, that are sampled from N (0 , n -3 / 2 ) (as explained in Appendix D). The feedforward layer maps the n = 768 -dimensional embedding to the larger dimension 3072 , as in the Hugging face implementation of Bert [73].

101

101

Entropy, Layer 0

100

res_gamma

Entropy, Layer 6

0.05

0.1

0.2

Ir

100

res_gamma

10-1

0.0001

0.05

0.001

0.1

0.2

Ir

0.0001

10-2

0.001

0

10000

20000

30000

40000

50000

60000

0

10000

20000

30000

40000

50000

60000

step

step

101

101

10-1

10-1

Entropy, Layer 0

Entropy, Layer 6

10-3

10-3

10-5

Ir

10-5

Ir

0.0001

0.0001

0.0005

0.0005

0.001

0.001

10-7

10-7

0

10000

20000

30000

40000

50000

60000

0

10000

20000

30000

40000

50000

60000

step step

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 on Nvidia DGX-1 GPU nodes equipped with 4 20-core Xeon E5-2698v4 processors, 512 GB of memory and 8 Nvidia V100 GPUs.

E Additional Figures

1.2

1.0 -

0.8-

Density

0.6-

0.4·

0.2

0.0

0

1

2

SDE (Ours)

Shaped NN

100

(t te paddes) aw! bu!ddois ue!paw pateinw!s

26.0

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

2-9.0

2-10.0

21

2-3.0

2-2.8

2-2.6

2-2.4

2-22

2-2.0

2-1.8

2-1.6

2-1.4

2-1.2

2-1.0

2-0.8

2-0.6

2-0.4

2-0.2

20.0

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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\nQUESTIONS:\n* 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\nLIMITATIONS:\nThere 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\nMy 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.\n\nQUESTIONS:\nIs there any reason to prefer an SDE limit over an ODE one (assuming a feature learning limit rather than the NTK)?\n\nLIMITATIONS:\nLimitations are adequately addressed", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

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

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.

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

2 There are several typos, e.g.,  lines 25, ''.One'' -> ''. One'', line 122 , ''importantproperties"-> ''important properties", line 242 "remains -> remain".

QUESTIONS:
Please see weaknesses.

LIMITATIONS:
The authors have adequately addressed the limitations, and there is no potential negative societal impact of their work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n**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.\n\nSTRENGTHS:\nIn 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.\n\nWEAKNESSES:\nTo 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nI think the authors adequately discuss the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-15"}