SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lorenzo Noci*†
Chuning  $\mathbf{L}\mathbf{i}^{*\ddagger}$
Mufan (Bill)  $\mathbf{L}\mathbf{i}^{*\dagger}$
Bobby He
Thomas Hofmann†
Chris Maddison†
Daniel M. Roy
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.
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,
*Equal contribution. Correspondence to:
lorenzo.noci@inf.ethz.ch, chuning.li@mail.utoronto.ca, mufan.li@mail.utoronto.ca
$\dagger$ ETH Zurich
$^{\ddagger}$ University of Toronto and Vector Institute
$^{\S}$ University of Oxford
Figure 1: Our shaped Transformer prevents token representations from becoming perfectly aligned, i.e. rank collapse. Left: mean correlation  $\rho_{\ell}^{\alpha \beta}$  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 a poorly conditioned covariance matrix. Simulated with  $n = 200$ ,  $d = 150$ ,  $\gamma = 1 / \sqrt{8}$ ,  $\tau_0 = 1$ ,  $\rho_0^{\alpha \beta} = 0.2$ , SDE step size 0.01, and  $2^{12}$  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 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  $\tau$ , 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).
(a) Value Weights  $W_{\ell}^{V}$
(b) Query Weights  $W_{\ell}^{Q}$
Figure 2: Comparing gradients norms at initialization for different parameters as a function of depth, with and without shaped attention. The architecture is the same as in Figure 1 but with autoregressive causal masking, and the task is next-token prediction on code data. Left: Value weights  $W_{\ell}^{V}$  for shaped attention, standard Pre-LN, and the original Post-LN block [2]. Right: the same gradient norm plot but for Query weights  $W_{l}^{Q}$ . 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.
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 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 new architecture in practice (see Section 5 and Appendix D).
1. 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].
2. 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.
3. 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).
4. We provide simulations to validate the theory and to interpret the effects of network hyperparameters 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.
Setup. Let  $X_{\ell} \in \mathbb{R}^{m \times n}$  be the data matrix representing a sequence of  $m$  tokens embedded in  $n$  dimensions at layer  $\ell \in [d]$ , where  $d$  is the depth of the network. We elide the explicit dependence on  $\ell$  when it is clear from the context, and use superscript Greek letters to indicate specific tokens' representations, for instance  $x_{\ell}^{\alpha} \in \mathbb{R}^{n}$  is the  $\alpha$ -th row of  $X_{\ell}$ . We consider the following attention model with residual connections:
$$
X _ {\ell + 1} = \lambda X _ {\ell} + \gamma A _ {\ell} X _ {\ell} \frac {1}{\sqrt {n}} W _ {\ell} ^ {V} \tag {1}
$$
where  $\gamma, \lambda \in [0,1]$  are parameters that control the strength of the shortcut and residual branch, respectively,  $W_{\ell}^{V} \in \mathbb{R}^{n \times n}$  is the weight matrix of the values, and  $A_{\ell} \in \mathbb{R}^{m \times m}$  is the attention matrix. We consider Softmax-based scaled dot-product attention, where  $A_{\ell}$  has the form:
$$
A _ {\ell} = \operatorname {S o f t m a x} \left(\frac {1}{\tau} X _ {\ell} \frac {1}{\sqrt {n}} W _ {\ell} ^ {Q} \frac {1}{\sqrt {n}} W _ {\ell} ^ {K, \top} X _ {\ell} ^ {\top}\right), \tag {2}
$$
where the Softmax is applied row-wise,  $W_{\ell}^{Q}, W_{\ell}^{K} \in \mathbb{R}^{n \times n_k}$  are additional random weights, and  $\tau$  is a temperature parameter, which controls the entropy of the distribution. Here we let all the weight matrices  $W_{\ell}^{Q}, W_{\ell}^{K}, W_{\ell}^{V}$  have  $\mathcal{N}(0,1)$ -iid entries. In the case where  $\lambda, \gamma = 1$ , with the application of LayerNorm on the residual branch [48], and with  $\tau = \sqrt{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  $\lambda^2 + \gamma^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_{\ell}^{\alpha \beta} \coloneqq \frac{1}{n}\langle x_{\ell}^{\alpha},x_{\ell}^{\beta}\rangle$  (or  $\rho^{\alpha \beta} \coloneqq (V_{\ell}^{\alpha \alpha}V_{\ell}^{\beta \beta})^{-1 / 2}V_{\ell}^{\alpha \beta}$  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-then-depth limit, the correlation  $\rho_d^{\alpha \beta}$  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_{\ell}$  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  $\gamma = 1 / \sqrt{\mathrm{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  $\gamma$  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^{\alpha \alpha}$  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 / \sqrt{\mathrm{width}}$ -away from linearity. In the proportional limit, the effect of the non-linearity accumulates over the  $d$  layers, and the covariance matrix  $V_{\ell} = [V_{\ell}^{\alpha \beta}]_{\alpha \beta}$  converges weakly to the solution of the SDE
$$
d V _ {t} = b _ {\text {R e L U}} \left(V _ {t}\right) d t + \Sigma_ {\text {l i n}} ^ {1 / 2} \left(V _ {t}\right) d B _ {t}, \tag {3}
$$
where the formulae for coefficients  $b_{\mathrm{ReLU}}$ ,  $\Sigma_{\mathrm{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 = \frac{d}{n}$ , in a similar spirit as the neural network Gaussian process (NNGP) results [20-22]. For example, the  $i$ -th output  $X_{\mathrm{out},i}$  conditioned on  $V_{d}$  are asymptotically iid.  $\mathcal{N}(0,V_T)$  as  $d,n\to \infty$ . 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:
$$
X _ {\ell + 1} = \lambda X _ {\ell} + \gamma \sigma_ {s} \left(X _ {\ell} \frac {1}{\sqrt {n}} W _ {\ell} ^ {\text {p r e}}\right) \sqrt {\frac {c}{n}} W _ {\ell} ^ {\text {p o s t}}, \tag {4}
$$
where  $\sigma_s(x) := s_+ \max(x, 0) + s_- \min(x, 0)$  is the shaped ReLU with slopes  $s_\pm := 1 + c_\pm n^{-1/2}$  for some constants  $c_+, c_- \in \mathbb{R}$ . We assume i.i.d weights  $(W_\ell^{\mathrm{pre}})_{ij}, (W_\ell^{\mathrm{post}})_{ij} \stackrel{\mathrm{iid}}{\sim} \mathcal{N}(0, 1)$ , and  $c^{-1} = \mathbb{E} \sigma_s(g)^2$  for  $g \sim 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_{\ell} \in \mathbb{R}^{m \times n}$  be the  $\ell$ -th layer matrix of representations, and define the feature covariance as  $V_{\ell} = \frac{1}{n} X_{\ell} X_{\ell}^{\top}$ . Let  $V_{t}^{(n)} = V_{[tn]} \in \mathbb{R}^{m(m+1)/2}$  be 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 \to \infty$ ,  $\frac{d}{n} \to T$ , the process  $\{V_{t}^{(n)}\}_{t \in [0,T]}$  converges to  $\{V_{t}\}_{t \in [0,T]}$  weakly in the Skorohod topology.
Theorem 3.2. Let  $X_{\ell}$  be the hidden layers of a ResNet defined in Eq. 4 with  $\lambda^2 + \gamma^2 = 1$ , where both  $\lambda$  and  $\gamma$  do not depend on  $d, n$ . Then the feature covariance  $V_{\ell}$  converges to the solution of the following SDE (in the sense of Definition 3.1)
$$
d V _ {t} = b _ {r e s} \left(V _ {t}\right) d t + \Sigma_ {r e s} \left(V _ {t}\right) ^ {1 / 2} d B _ {t}, \quad V _ {0} = \frac {1}{n} X _ {0} X _ {0} ^ {\top}, \tag {5}
$$
where  $b_{res}(V) = \gamma^2 b_{ReLU}(V) = \gamma^2 [\nu (\rho^{\alpha \beta})\sqrt{V^{\alpha\alpha}V^{\beta\beta}} ]_{\alpha \leq \beta}$  with  $\rho^{\alpha \beta} = V^{\alpha \beta}(V^{\alpha \alpha}V^{\beta \beta})^{-1 / 2}$  and
$$
\nu (\rho) = \frac {\left(c _ {+} - c _ {-}\right) ^ {2}}{2 \pi} \left(\sqrt {1 - \rho^ {2}} - \rho \operatorname {a r c c o s} \rho\right), \tag {6}
$$
furthermore,  $\Sigma_{res}(V) = 2\gamma^2\Sigma_{lin}(V) = 2\gamma^2 [V^{\alpha \delta}V^{\beta \omega} + V^{\alpha \omega}V^{\beta \delta}]_{\alpha \leq \beta ,\delta \leq \omega}$ .
Figure 3: Left: Kernel density estimates of correlation  $\rho_d^{\alpha \beta}$  for various values of the residual strength parameter  $\gamma$ . In particular,  $\gamma = 1$  recovers a shaped-ReLU MLP without skip connections, and  $\gamma = 1 / \sqrt{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 in dashed lines. Right: 95th percentile of the absolute value of the correlation distribution as a function of  $\gamma$ . Note reducing  $\gamma$  reduces the concentration around  $\rho^{\alpha \beta} = 1$ , and our SDE reliably approximates finite size networks. Simulated with  $n = 300$ ,  $d = 100$ ,  $\rho_0^{\alpha \beta} = 0.2$ ,  $c_{+} = 0$ ,  $c_{-} = -1$ , and  $2^{13}$  samples.
Notice how the limiting SDE closely resembles the MLP case (Eq. 3), which is recovered exactly when  $\gamma = 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  $\gamma < 1$  dampens both the drift and the variance of the Brownian motion by  $\gamma^2$ , thus it can be interpreted as a time change. In other words, the effect of  $\gamma$  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  $\gamma$  on the distribution, in Figure 3 (right) we plot the 95th percentile correlation as a function of  $\gamma$ . 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  $\gamma$ . 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  $\lambda = 1, \gamma = 1 / \sqrt{d}$ , the depth and width limits commute. More precisely, the covariance  $V^{\alpha \beta}$  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  $\gamma = 1 / \sqrt{d}$  suppresses the random fluctuations enough to vanish in the limit. On the other hand, our results show that for  $\lambda, \gamma$  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 prevent 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_{\ell}$  must be on the same order as the random weights in an MLP, thus forming an approximate