Deep Transformers without Shortcuts: Modifying Self-attention for Faithful Signal Propagation

Introduction

Despite numerous impressive successes, the practice of training deep neural networks (DNNs) has progressed to a large extent independently of theoretical justification. Most successful modern DNN architectures rely on particular arrangements of skip connections and normalisation layers, but a general principle for how to use these components in new architectures (assuming they are even applicable) remains unknown, and their roles in existing ones are still not completely understood.

The residual architecture, arguably the most popular and successful of these, was first developed in the context of convolutional networks (CNNs) (He et al., 2016), and later in self-attention networks yielding the ubiquitous transformer architecture (Vaswani et al., 2017). One proposed explanation for the success of residual architectures is that they have superior signal propagation compared to vanilla DNNs (e.g. Balduzzi et al., 2017; Xiao et al., 2018; Hayou et al., 2019; De & Smith, 2020; Martens et al., 2021), where signal propagation refers to the transmission of geometric information through the layers of a DNN, as represented by a kernel function (Daniely et al., 2016; Poole et al., 2016; Schoenholz et al., 2017).

Recently, using signal propagation principles to train DNNs at high depths, without the skip connections and/or normalisation layers found in residual architectures, has become an area of interest in the community. The reasons are two-fold. First, it would validate the signal propagation hypothesis for the effectiveness of residual architectures, thus clarifying our understanding of DNN trainability. And second, it could lead to general principles and techniques for achieving trainability in DNNs beyond the residual paradigm, with the potential for improved or more efficient architectures.

For CNNs, Xiao et al. (2018) showed that improved signal propagation from better initialisation enables very deep vanilla networks to be effectively trained, although at significantly reduced speeds compared to residual networks. Martens et al. (2021) later proposed Deep Kernel Shaping (DKS) which uses activation function transformations to control signal propagation, achieving training speed parity between vanilla and residual networks on ImageNet assuming the use of strong 2nd-order optimizers like K-FAC (Martens & Grosse, 2015). Zhang et al. (2022) extended ideas from DKS to a larger class of activation functions, achieving near parity in terms of generalisation as well.

The key quantity that is analysed in signal propagation is the DNN’s initialisation-time kernel, or more precisely, the approximate kernel given by the infinite width limit (Neal, 2012; Matthews et al., 2018; Lee et al., 2018; Yang, 2019). For MLPs, and for CNNs that use a Delta-initialisation (Balduzzi et al., 2017; Xiao et al., 2018), this kernel can be written as a simple recursion over layers that involves only 2D functions, facilitating a straightforward analysis.

Analogously to the case of MLPs, where signal propagation is judged by looking at the behavior of the (one-dimensional) kernel, signal propagation in transformers can be judged by looking at the evolution of these (high-dimensional) kernel matrices through the layers of the network. One situation we must avoid is where the diagonal entries rapidly grow or shrink with depth, which corresponds to uncontrolled activation norms and can lead to saturated losses or numerical issues. A more subtle form of signal degradation can occur where $\Sigma_{l}$ converges to a rank-1 matrix, which is known as rank collapse (Dong et al., 2021). Dong et al. (2021) showed that skip connections are essential to avoid the collapsed state: skipless transformers quickly converge to rank collapse at large depths, which we corroborate in Fig. 1 (top). Moreover, Noci et al. (2022) showed that rank collapse may lead to zero gradients for certain parameters in attention layers, hindering the trainablility of deep transformers. Thus, avoiding rank collapse is necessary for deep transformers to be trainable, and the question of whether one can train deep skipless transformers remains open.

In the present work we address this question, demonstrating for the first time that it is possible to successfully train deep transformers without skip connections or normalisation layers. To do so, we study the problem of signal propagation and rank collapse in deep skipless transformers, and derive three approaches to prevent it in Section 3. Our methods use combinations of: 1) parameter initialisations, 2) bias matrices, and 3) location-dependent rescaling, and highlight several intricacies specific to signal propagation in transformers, including the interaction with positional encoding and causal masking. In Section 4, we empirically demonstrate that our approaches result in trainable deep skipless transformers. On WikiText-103 and C4 datasets we show that using our main approach, Exponential Signal Preserving Attention (E-SPA), it is possible to match the training loss of standard transformers with our skipless ones by training for around 5 times longer. Moreover, by combining this approach with skip connections, we show that transformers without normalisation layers are able to match the training speed of standard ones.

Problem setting

where the shortcut and residual weights $\alpha,\beta$ are typically both $1$. In this work we will focus on skipless transformers with $\alpha=0$ and $\beta=1$, and on vanilla transformers, which are skipless transformers without normalisation layers. For simplicity we will also devote much of our analysis to attention-only models, which are transformers without MLP blocks (so that $\hat{{\mathbf{X}}}_{l}={\mathbf{X}}_{l}$).

Note that we restrict our analysis to decoder-only transformers in this work, as they have a simpler structure which is easier to analyse, and are widely used in practice. Also note that Eq. 1 corresponds to the “Pre-LN” (Baevski & Auli, 2018; Child et al., 2019) rather than the original “Post-LN” transformer (Wang et al., 2019). In Post-LN transformers, the normalisation operation is applied at the output of each MLP and attention block instead of at the beginning of each residual branch.

where $\Gamma$ is a large positive constant that zeros the attention coefficients corresponding to future tokens, making ${\mathbf{A}}$ a lower triangular matrix.

As discussed in Section 1, Dong et al. (2021) showed that deep skipless transformers suffer from rank collapse, where the kernel matrix converges in depth to have rank 1, and Noci et al. (2022) showed that rank collapse can prevent trainability.

Moreover, Noci et al. (2022) demonstrated that rank collapse in transformers can occur in the absence of normalisation layers even with skip connections, and that downscaling the residual branch by setting $\beta=\frac{1}{\sqrt{L}}$ can alleviate this issue. This latter observation is in line with previous findings concerning the benefits of downscaling residual weights in ResNets (Hanin & Rolnick, 2018; Zhang et al., 2018; Arpit et al., 2019; Hayou et al., 2021; Bachlechner et al., 2021) and transformers (Zhang et al., 2019; Xu et al., 2020; Huang et al., 2020; Touvron et al., 2021; Wang et al., 2022). Davis et al. (2021) showed that concatenation acts similarly to a downweighted skip as an alternative way to connect skip and residual branches. De & Smith (2020) noted that the interaction of standard skip connections and normalisations can also effectively downweight the residual branch to give better signal propagation properties, but only if the normalisation layer is placed on the residual branch, like in Pre-LN transformers. Such an effect does not occur for Post-LN transformers, where the normalisation layer is after the residual branch, and we indeed observe in Fig. 7 that Post-LN attention-only transformers also suffer from rank collapse at large depths. This may explain some of the training instabilities of Post-LN transformers that have been observed in practice (Xiong et al., 2020; Liu et al., 2020).

Constructing trainable deep transformers without shortcuts

To date, the only strategy for rectifying rank collapse in transformers relies on skip/shortcut connections, which “skip” around the trainability issues intrinsic to self-attention layers. We seek instead to tackle this issue directly. To do so, we first develop a better understanding of signal propagation through attention layers, then derive modifications from our insights to achieve faithful signal propagation in deep transformers, allowing them to be trained regardless of the use of skip connections.

To start with, we consider the simplified setting of a deep attention-only vanilla transformer, and suppose we are in a single-head setting ($h=1$) or a multi-head setting where the attention matrix ${\mathbf{A}}$ does not vary across heads. If block $l\leq L$ has attention matrix ${\mathbf{A}}_{l}$ at initialisation, then the final block’s representation ${\mathbf{X}}_{L}$ takes the following form:

From this simplified formula for kernel matrices in deep attention-only transformers, we identify three requirements on $({\mathbf{A}}_{l})_{l}$:

$\Sigma_{l}=\Pi_{l}\cdot\Sigma_{0}\cdot\Pi_{l}^{\top}$ must be well-behaved at each block, avoiding degenerate situations such as rank collapse and exploding/vanishing diagonal values.

${\mathbf{A}}_{l}$ must be elementwise non-negative $\forall l$ (recalling that ${\mathbf{A}}_{l}$ is constructed through the softmax operation Eq. 4).

${\mathbf{A}}_{l}$ should be lower triangular $\forall l$, for compatibility with causal masked attention.We describe compatibility of our methods with non-causal attention in Appendix A.

In Sections 3.1 and 3.2, we focus on finding attention matrices that satisfy our desiderata above, and demonstrate how to modify softmax attention to achieve these attention matrices in Section 3.3.

An obvious solution to our above requirements on $({\mathbf{A}}_{l})_{l}$ is the trivial one: ${\mathbf{A}}_{l}={\mathbf{I}}_{T}\hskip 5.0pt\forall l$, where each sequence location attends only to itself. In this case, $\Sigma_{L}=\Sigma_{0}$ perfectly preserves the input kernel matrix, and will avoid rank collapse assuming $\Sigma_{0}$ is non-degenerate. Unfortunately, identity attention isn’t compatible with a viable solution to obtain trainable vanilla transformers. This is because to achieve ${\mathbf{A}}_{l}={\mathbf{I}}$ we would need to saturate the softmax operation (Eq. 2) so that gradients do not pass to the query and key parameters, and the attention matrix stays close to identity during training.

To provide a partial solution to this that achieves identity attention matrix at initialisation yet is still trainable, we introduce our first approach, Value-SkipInit, based on SkipInit (De & Smith, 2020), and the related ReZero method (Bachlechner et al., 2021). In Value-SkipInit, we modify the attention operation $\text{Attn}({\mathbf{X}})={\mathbf{A}}({\mathbf{X}}){\mathbf{V}}({\mathbf{X}})$, to

with trainable parameters $\alpha$ and $\beta$ that are initialised to $1$ and respectively. Thus, at initialisation, the attention matrix is the identity. For transformers with MLPs blocks, this yields identical behaviour to a standard MLP acting on each sequence location independently at initialisation, so we can apply the DKS or TAT frameworks (Martens et al., 2021; Zhang et al., 2022) to achieve well-behaved signal propagation in the entire model.This is similar in principle to the Delta-initialisation for CNNs (Balduzzi et al., 2017; Xiao et al., 2018), and modifications to GNNs to enable compatibility with TAT (Zaidi et al., 2022). We note that Value-Skipinit is an approach to remove the standard skip connections in transformer blocks, Eq. 1, but can be construed as adding a skip connection around the value computation, and in that respect isn’t strictly “skipless”. Moreover, there is useful information contained in the positions of sequence locations that we do not employ when all attention matrices are identity at initialisation. As a result, we treat Value-Skipinit as a baseline for our main two methods, which we describe next.

2 Signal Preserving Attention methods

Returning to our requirements on $({\mathbf{A}}_{l})_{l}$ in Section 3, we see that controlling their product is key to achieving faithful signal propagation. Given that we are now interested in non-identity $({\mathbf{A}}_{l})_{l}$, this becomes more difficult. To overcome this, we consider ${\mathbf{A}}_{l}$ of the form

such that individual ${\mathbf{L}}_{i}$’s cancel in the product, giving $\Pi_{l}={\mathbf{L}}_{l}{\mathbf{L}}_{0}^{-1}$. Then if ${\mathbf{L}}_{0}$ satisfies ${\mathbf{L}}_{0}^{-1}\Sigma_{0}{\mathbf{L}}_{0}^{-1^{\top}}={\mathbf{I}}_{T}$, we thus have

Assuming input embeddings are initialised independently, and no repeated tokens in the input sequence, we have $\Sigma_{0}={\mathbf{I}}_{T}$ in the large width limit,Because the embedding matrix ${\mathbf{E}}$ is initialised with variance $1/\text{fan-out}$ we rescale the embeddings it produces by $\sqrt{d_{0}}$ to get a $\Sigma_{0}$ with ones on the diagonal. and can thus take ${\mathbf{L}}_{0}={\mathbf{I}}_{T}$. In practice, we will apply slight modifications to our methods to account for repeated tokens, as detailed in Appendix B.

Now if ${\mathbf{L}}_{l}$ is lower triangular, then from Eq. 9 we see it is simply a Cholesky factor for the kernel matrix $\Sigma_{l}$. By the uniqueness of Cholesky factors for PSD matrices up to sign, this means we just need to choose $\{\Sigma_{l}\}_{l=0}^{L}$ to be a family of well-behaved kernel matrices that satisfyConstraint (iii) is satisfied as lower triangular matrices are closed under multiplication and inversion. our non-negativity constraint (ii) on ${\mathbf{A}}_{l}={\mathbf{L}}_{l}{\mathbf{L}}_{l-1}^{-1}$. We identify two such families, which give rise to our main method, Signal Preserving Attention (SPA):

Uniform (U-SPA): $\Sigma_{l}(\rho_{l})=(1-\rho_{l}){\mathbf{I}}_{T}{+}\rho_{l}\bm{1}\bm{1}^{\top}$. Here, the kernel matrix $\Sigma_{l}(\rho_{l})$ has diagonal entries equal to 1 and off-diagonal entries equal to $\rho_{l}$. The condition $\rho_{l}\leq\rho_{l+1}$ is required for elementwise non-negativity of the ${\mathbf{A}}$’s, as shown in Theorem 1. Setting $\rho_{0}=0$ yields identity input kernel matrix, and as long as $\rho_{L}<1$, we avoid rank collapse in the skipless attention-only setting.

Exponential (E-SPA): \big{(}\Sigma_{l}(\gamma_{l})\big{)}_{i,j}=\text{exp}(-\gamma_{l}|i-j|). Here, diagonal entries are again 1, but now off-diagonals decay exponentially for more distant locations, with decay rate $\gamma_{l}$. Thus, unlike U-SPA, E-SPA captures the notion of positional encoding, as the vector representations for nearby locations have larger inner products (i.e. are more similar to each other). The condition $\gamma_{l}\geq\gamma_{l+1}$ is required for elementwise non-negativity of the ${\mathbf{A}}$’s, as established in Theorem 2. Setting $\gamma_{0}=\infty$ yields the identity input kernel matrix, and rank collapse will be prevented as long as $\gamma_{L}>0$.

We state and prove Theorems 1 and 2 in Appendix I, and in particular provide a closed-form solution for ${\mathbf{L}}_{l}{\mathbf{L}}_{l-1}^{-1}$ in the case of E-SPA in Theorem 3, enabling cheap computation. From these theorems, we see that our proposed SPA approaches are viable as long as the kernel matrix values become progressively larger (in an elementwise fashion) as depth increases, as dictated by $(\rho_{l})_{l=1}^{L}$ and $(\gamma_{l})_{l=1}^{L}$. We find $\rho_{L}=0.8$ and $\gamma_{L}=0.005$ to be good default choices for the final block, and describe how to vary $\rho_{l}$ and $\gamma_{l}$ across depth in Appendix C. In Alg. 1 in Appendix D, we summarise how to construct a trainable self-attention layer with our main method E-SPA using ideas from Section 3.3, given an input-output decay rate pair $\gamma_{\text{in}},\gamma_{\text{out}}$ (in the notation of Alg. 1). Given a decreasing set of decay rates $(\gamma_{l})_{l=0}^{L}$, at block $l$ we set $\gamma_{\text{in}}=\gamma_{l-1}$ and $\gamma_{\text{out}}=\gamma_{l}$.

In Fig. 1, we verify that our two proposed SPA schemes, U-SPA and E-SPA, successfully avoid rank collapse in attention-only vanilla transformers even at large depths. Moreover, because $\Sigma_{l}$ has diagonals equal to 1 for all $l$, there is an implicit mechanism in these two schemes to control the representation vector norms across all sequence locations at deep layers. This means that they can be used with or without normalisation layers, as we will verify empirically in Section 4. Moreover, in Fig. 1 we see that E-SPA observes a recency bias as expected, where representation vectors for nearby locations have larger cosine similarity, as seen with positional encoding schemes like ALiBi (Press et al., 2022) (Fig. 7). As a result, even though all three of our approaches successfully avoid rank collapse, we expect E-SPA to outperform both U-SPA and Value-SkipInit.

3 Reverse engineering self-attention layers at initialisation

In Section 3.2 we identified two families of lower-triangular non-negative attention matrices $({\mathbf{A}}_{l})_{l}$ which enable us to obtain well-behaved kernel matrices $\Sigma_{L}$ at large depth $L$, and hence faithful signal propagation. It remains to show how we can actually realise these attention matrices via parameter initialisation and minor modifications of the attention mechanism.

which reduces to $\text{Attn}({\mathbf{X}})={\mathbf{A}}{\mathbf{V}}({\mathbf{X}})$ (with a data-independent attention matrix) at initialisation if the query-key dot product $\frac{1}{\sqrt{d^{k}}}{\mathbf{Q}}({\mathbf{X}}){\mathbf{K}}({\mathbf{X}})^{\top}=0$. We note that zero query-key dot products occur if one uses a $\frac{1}{d^{k}}$ scaling rather $\frac{1}{\sqrt{d^{k}}}$ in the infinite width limit for independently initialised ${\mathbf{W}}^{Q},{\mathbf{W}}^{K}$ (Yang, 2019). In practice, to achieve zero initial query-key dot product, we can initialise either: 1) ${\mathbf{W}}^{Q}=0$, 2) ${\mathbf{W}}^{K}=0$, or 3) both ${\mathbf{W}}^{Q}$ and ${\mathbf{W}}^{K}$ to have small initial scale (which achieves approximately zero dot product). In our experiments we found these three options to all perform similarly, and decided to use option 1): ${\mathbf{W}}^{Q}=0$. An empirical evaluation of the sensitivity to the choice of initial scale in option 3) is provided in Fig. 8.

In Eq. 10, ${\mathbf{D}}$ acts as a non-trainable location-dependent rescaling, and is needed to realise arbitrary attention matrices since the softmax output ${\mathbf{P}}({\mathbf{X}})$ is constrained to have row sums equal to 1.In SPA, ${\mathbf{D}}$ also corrects for the fact that masked softmax attention tends to reduce the norms of representation vectors for locations at the end of the sequence. To see this, if we have kernel matrix $\Sigma={\mathbf{X}}{\mathbf{X}}^{\top}/d$ with $\Sigma_{ii}=1,\hskip 3.00003pt\forall i$, and softmax attention matrix ${\mathbf{A}}$ with row $i$ ${\bm{a}}_{i}$, then $({\mathbf{A}}\Sigma{\mathbf{A}}^{\top})_{ii}=\frac{1}{d}\lVert{\bm{a}}_{i}{\mathbf{X}}\rVert^{2}_{2}\leq\frac{1}{d}\lVert{\mathbf{X}}\rVert_{2}^{2}\lVert{\bm{a}}_{i}\rVert_{2}^{2}=\lVert{\bm{a}}_{i}\rVert_{2}^{2}$. But ${\bm{a}}_{i}$ sums to 1 (as a softmax output), so $\lVert{\bm{a}}_{i}\rVert_{2}\leq 1$ with equality if and only if ${\bm{a}}_{i}$ has exactly one non-zero entry (equal to 1). This holds for the first token (which can only attend to itself in masked attention) so $({\mathbf{A}}\Sigma{\mathbf{A}}^{\top})_{11}=1$, but not in general for later tokens, so $({\mathbf{A}}\Sigma{\mathbf{A}}^{\top})_{ii}$ will usually be less than 1, for $i>1$. Given target kernel matrices with 1’s on the diagonal (which we have in SPA) its inclusion is akin to using RMSNorm at initialisation. However, this will gradually stop being true during training, leading to a (slightly) different model class. The additive pre-softmax biases ${\mathbf{B}}$ will also have an effect on the model class, but this will be similar to that of the ALiBi positional encoder (Press et al., 2022), which too involves a non-trainable bias matrix being added to the logits. In our early experiments we tried including a trainable gain parameter on the bias matrix, initialised to 1, in order to better preserve the model class, but found that this didn’t have a significant impact on training performance.

We also note that this approach to controlling the attention matrix by zeroing the query-key dot product at initialisation is compatible with popular positional encodings like Relative (Shaw et al., 2018; Huang et al., 2019; Dai et al., 2019) and RoPE (Su et al., 2021), as discussed in Appendix E. Unless stated otherwise, we use RoPE in our experiments, and provide an ablation in Fig. 10.

4 Addressing MLP blocks and skip connections

Up until this point we have restricted our focus to attention-only skipless transformers, for the sake of simplicity. However, MLP blocks are an important part of the transformer architecture that we must also address. And we would like for our approaches to be compatible with skip connections too, both for the sake of generality, and because they might combine profitably.

Because MLP blocks operate on the representation vectors independently across locations, their effect on the kernel matrix can be easily computed using the standard limiting kernel formulas for MLPs (Neal, 2012) which are exact in the infinite width limit. In particular, there is a known $f$ that maps the kernel matrix $\Sigma_{l}$ of an MLP block’s input sequence to the kernel matrix $\hat{\Sigma}_{l}$ of its output sequence. In principle, one could modify SPA to account for this change to the kernel matrix by taking ${\mathbf{A}}_{l}={\mathbf{L}}_{l}\hat{{\mathbf{L}}}_{l-1}^{-1}$, where ${\mathbf{L}}_{l}$ and $\hat{{\mathbf{L}}}_{l}$ are the Cholesky factors of $\Sigma_{l}$ and $\hat{\Sigma}_{l}$ respectively. Unfortunately, there is no guarantee that the resulting ${\mathbf{A}}_{l}$ would be elementwise non-negative in general. In our experiments we thus elected to ignore the effect on the kernel matrices of the MLP blocks, approximating $f$ as the identity function, which we found to work well enough in practice. As discussed in Martens et al. (2021), this approximation becomes better as the MLP blocks tend to linear functions (for which $f$ is exactly identity), which will happen as we increase the shortcut weight $\alpha$ relative to the residual weight $\beta$ (see Eq. 1), or when the MLP’s activation functions are close to linear. Notably, the latter condition will tend to be true when using DKS or TAT to transform the activation functions in deep networks. Note, when using DKS, $f$ decreases the size of off-diagonal elements of the kernel matrix, which intuitively will help to combat rank collapse.

When using skip connections as in Eq. 1, the output kernel matrix of a block is given by $\Sigma_{\text{block}}=\alpha^{2}\Sigma_{\text{shortcut}}+\beta^{2}\cdot\Sigma_{\text{residual}}$. One of the main ways that kernel matrices degenerate in DNNs is when their diagonal entries either explode or shrink with depth. As shown by Martens et al. (2021), this can be prevented in fully-connected and convolutional DNN architectures by rescaling the output of the activation functions (which happens automatically as part of DKS or TAT), and by replacing weighted sums with “normalised sums”, which yields normalised skip connections defined by the condition $\alpha^{2}+\beta^{2}=1$. When using normalised skip connections with U-SPA, both $\Sigma_{\text{shortcut}}$ and $\Sigma_{\text{residual}}$ will have the form $\Sigma(\rho)=(1-\rho){\mathbf{I}}_{T}{+}\rho\bm{1}\bm{1}^{\top}$ for some $\rho$, and thus so will $\Sigma_{\text{block}}$. Moreover, we will have $\rho_{\text{block}}=\alpha^{2}\rho_{\text{shortcut}}+\beta^{2}\cdot\rho_{\text{residual}}$, which is less than $\rho_{\text{residual}}$. This means we can easily adjust U-SPA to be compatible with normalised skip connections by replacing ${\mathbf{L}}_{l}$ in the formula ${\mathbf{A}}_{l}={\mathbf{L}}_{l}{\mathbf{L}}_{l-1}^{-1}$ with the Cholesky factor of $\Sigma((\rho_{l}-\alpha^{2}\rho_{l-1})/\beta^{2})$. In the case of E-SPA, we note that $\Sigma_{\text{block}}$ won’t be of the form \big{(}\Sigma(\gamma)\big{)}_{i,j}=\text{exp}(-\gamma|i-j|) for some $\gamma$, even when both $\Sigma_{\text{shortcut}}$ and $\Sigma_{\text{residual}}$ are. To work around this, we use an approximation described in Appendix F, which seeks to make the combined effect on the kernel matrix of the normalised skip and attention block approximately equal to the effect of a skipless attention block.

Experiments

We now assess the capabilities of our proposed methods in training deep skipless and/or normaliser-free transformers. Our main experiment setting uses a transformer with 36 transformer blocks, which is deep enough that the effects of poor signal propagation and rank collapse render skipless training without modifications impossible. We begin our investigation on WikiText-103 (Merity et al., 2017) focusing primarily on training performance of our various methods, before moving to the larger C4 dataset (Raffel et al., 2019), where overfitting is not an issue. We use Adam optimiser (Kingma & Ba, 2014), as is standard for transformers, and train without dropout (Srivastava et al., 2014). Additional results and experimental details are provided in Appendices G and H respectively.

To start with, we verify that a standard deep transformer without skip connections is untrainable even with normalisation layers (LN) and transformed activations, and that our approaches remedy this. Fig. 2 compares vanilla transformers using our proposed SPA methods and Value-Skipinit to standard transformers both with and without skips, on a 36 block transformer. We clearly see that removing skip connections from a standard transformer makes it untrainable, with training loss plateauing around 7.5. This holds true too even when we use DKS to transform the GeLU MLPs, highlighting that the issue lies with the attention layers, which suffer from rank collapse as shown in Fig. 1.

On the other hand, all three of our approaches train even for vanilla deep transformers, with our E-SPA method outperforming U-SPA and Value-Skipinit. However, the default transformer with skips and LN still retains a training speed advantage compared to our skipless methods, mirroring the situation for CNNs without powerful second-order optimisers (Martens et al., 2021; Zhang et al., 2022).

In Table 1, we assess the effect of different activation functions in the MLP blocks, as well as the use of LN, in skipless transformers using our proposed methods. We see that at depth 36 we achieve good training performance for a range of activations: DKS-transformed GeLU, TAT-transformed Leaky ReLU, as well untransformed GeLU (Hendrycks & Gimpel, 2016), but not untransformed Sigmoid. We also see that layer normalisation is relatively unimportant for training speed, and can even be harmful with transformed activations when using SPA, which already has an inbuilt mechanism to control activation norms (as discussed at the end of Section 3.2).

In Fig. 3, we see that one way to match the training loss of the default transformer, without more iterations, is by using normalised skip connections. While this is perhaps unsurprising, we observe that our E-SPA method (left) matches standard training both with and without normalisation, whereas a standard Transformer with normalised skip connections (right) requires normalisation in order to match the training speed of the default Pre-LN.

So far, we tested our proposed methods on WikiText-103 (Merity et al., 2017), on which we observed overfitting without the use of extra regularisation. Therefore, we further compare our methods to standard transformers on a larger dataset, C4 (Raffel et al., 2019), where overfitting isn’t an issue. Importantly, we see similar trends across validationArgued by Nakkiran et al. (2021), the validation curves measure the “training speed” of the online setting. (Fig. 4(a)), training (Fig. 13) and downstream task (Table 6) performance, and so the benefits of our methods do extend beyond training. Due to the memory overhead of longer sequences, we use a 32-block transformer. One can notice that E-SPA performs the best among skipless transformers on all settings: training, validation and downstream tasks.

Moreover, in Table 2 we find that E-SPA with normalised skips and LN outperforms the default Pre-LN transformer, achieving 24.0 (vs 24.7) validation perplexity on C4 after 50K steps. Even without LN, E-SPA exactly matches the Pre-LN transformer in training speed, and outperforms a range of baselines designed to remove LN, including Stable ResNet (Hayou et al., 2021; Noci et al., 2022) and SkipInit (De & Smith, 2020), both with & without LN.

In Figs. 2 and 4(a), we observed that while our methods are able to train deep skipless transformers (a result which is unprecedented in the literature), there is a significant gap in training speed compared to a standard Pre-LN transformer (with skips). Martens et al. (2021) and Zhang et al. (2022) observed a similar training speed gap for skipless CNNs, and showed that such a gap can be closed by using more sophisticated second order optimisers like K-FAC (Martens & Grosse, 2015) or Shampoo (Gupta et al., 2018; Anil et al., 2020). As second order optimisers for transformers are not well established, we instead demonstrate that the training loss gap can be closed by simply training for longer in Fig. 4(b). We observe that our E-SPA method matches the training loss of a standard pre-LN transformer on C4 if one trains for around 5 times longer with Adam, in line with the findings from the convolutional case (Zhang et al., 2022). An equivalent plot for WikiText-103 is provided in Fig. 14.

Finally, as unmodified networks suffer from worse signal propagation properties at larger depths, it is natural to ask how our modified vanilla transformers perform as depth increases. In Table 3, we compare the training performance of different depths (36, 72, and 108) at a range of training step budgets on WikiText-103. We find that whilst the shallower depth 36 network trains fastest initially over 100K steps, it is matched at 400K steps and then surpassed at 1000K steps by the larger capacity depth 72 network. Moreover, our depth 108 vanilla transformer is able to close the gap in performance to its depth 36 counterpart with longer training, going from 0.3 at 100K steps to just 0.01 at 1000K steps.

Conclusion

We have shown for the first time that it is possible to successfully train deep transformers without skip connections or normalisation layers. To do so, we have proposed 3 approaches: E-SPA, U-SPA and Value-Skipinit, each of which control the attention matrices of a transformer to enable faithful signal propagation even at large depths. Our best approach, E-SPA enables deep vanilla transformers to match their standard counterparts with around 5 times more iterations, and also deep transformers without normalisation to match the training speed of standard ones. We hope that our work may potentially pave the way to new and improved architectures, and more research into improving the capabilities of deep learning in practice using insights from theory.

Reproducibility Statement

Pseudocode for our main approach, E-SPA, can be found in Alg. 1, using the notation and setup provided in Section 2. All experimental details can be found in Appendix H, including general and experiment-specific implementation details.

Acknowledgements

We thank Christos Kaplanis for helpful discussions during initial stages of this project, as well as the anonymous reviewers for their feedback. BH is supported by the EPSRC and MRC through the OxWaSP CDT programme (EP/L016710/1).

References

Appendix A Compatibility with non-causal attention

In Section 3, we focus on causal masked self-attention for two reasons. First, next-token prediction using causal masked self-attention is arguably the most popular setting for self-attention. And second, it is a more challenging setting to work with in terms of controlling signal propagation, due to the additional constraint that attention matrices must be lower triangular. In this section we describe how our methods can be made compatible with non-causal masked attention, where the attention matrices are no longer required to be lower triangular.

To start with, Value-SkipInit does not modify the softmax-attention computation and hence is already compatible with any form of attention. For our SPA methods, it is straightforward to extend to non-causal attention by changing ${\mathbf{L}}_{l}$ in Eqs. 8 and 9 from being the Cholesky decomposition of $\Sigma_{l}$ to being the (symmetric) matrix square root of $\Sigma_{l}$. In this case, for U-SPA it is possible to analytically calculate that ${\mathbf{A}}_{l}$ in Eq. 8 will be element-wise non-negative if $\rho_{l}\geq\rho_{l-1}$ (exactly like the Cholesky case in Theorem 1). This is easy to see because the matrix square-root, inverses and products of uniform kernel matrices are all still uniform of the form $\Sigma(\rho)=(1-\rho){\mathbf{I}}_{T}{+}\rho\bm{1}\bm{1}^{\top}$ (up to positive rescaling), so that ${\mathbf{A}}_{l}$ will be too, and one simply needs to track $\rho$ and verify that it is positive. For E-SPA. we have verified empirically that the resulting ${\mathbf{A}}_{l}={\mathbf{L}}_{l}{\mathbf{L}}_{l-1}^{-1}$ will be non-negative if $\gamma_{l}\leq\gamma_{l-1}$, just like the Cholesky case in Theorem 2.

Appendix B Modifications to SPA methods for repeated tokens

For simplicity, we assumed that our input sequences had no repeated tokens when presenting SPA in Section 3. This meant that we could take the input kernel matrix $\Sigma_{0}$ to be the identity, with zero off-diagonals, which was convenient for our construction of SPA. The effect of repeated tokens, e.g. if the word ‘cat’ occurs multiple times in the same sentence, is that our input kernel matrices $\Sigma_{0}$ will have non-zero off-diagonal entries, corresponding to entries where a token is repeated. This will impact the kernel matrices $\Sigma_{l}$ at deeper layers with SPA, particularly the diagonal values of $\Sigma_{l}$, which we would like to able to control. In this section we discuss how we can modify our SPA approaches to account for the fact that we will often be working with sequences where a fraction of the tokens are repeated.

We stress that the general principle that all our methods (Value-SkipInit and SPA methods) follow is independent of the input kernel matrix (i.e. independent of duplicate input tokens): we seek to prevent the product of attention matrices from deviating away from the identity matrix and degenerating to a rank-1 matrix. From Eq. 6, we see that if the product of attention-matrices is rank-1 then regardless of the input kernel we will have a rank-1 output kernel i.e. rank collapse (Dong et al., 2021). On the other hand, if we control the deviation of the attention matrix product from the identity then no matter the input kernel, the output kernel will bear some similarity to the input kernel. So as long as the input kernel is non-degenerate and has full rank (regardless of duplicate tokens) so too will be the output kernel.

where $\Pi_{l}={\mathbf{A}}_{l}{\mathbf{A}}_{l-1}\dots\mathbf{A}_{1}$ is the product of attention matrices.

In SPA, we parameterise ${\mathbf{A}}_{l}={\mathbf{L}}_{l}{\mathbf{L}}_{l-1}^{-1}$ for $({\mathbf{L}}_{l})_{l}$ corresponding to the Cholesky factors of some family of kernel matrices $(\Sigma_{l})_{l}$, with either uniform or exponentially decaying off diagonals:

U-SPA: $\Sigma_{l}(\rho_{l}){=}(1-\rho_{l}){\mathbf{I}}_{T}{+}\rho_{l}\bm{1}\bm{1}^{\top}$ for $0\leq\rho\leq 1$

E-SPA: \big{(}\Sigma_{l}(\gamma_{l})\big{)}_{i,j}=\text{exp}(-\gamma_{l}|i-j|)) for $\gamma\geq 0$.

It is important that ${\mathbf{A}}_{l}={\mathbf{L}}_{l}{\mathbf{L}}_{l-1}^{-1}$ is constructed from two Cholesky matrices belonging to the same family, because Theorems 1 and 2 show in that case we will satisfy our non-negativity constraint on ${\mathbf{A}}_{l}$ (which is computed through a softmax operation), and otherwise there is no prior reason to suppose that non-negativity will be satisfied.

Therefore, in an ideal world, $\Sigma_{0}$ would be an element of our family of kernel matrices, so that we can set ${\mathbf{L}}_{0}$ to be the Cholesky factor of $\Sigma_{0}$, satisfying ${\mathbf{L}}_{0}^{-1}\Sigma_{0}{\mathbf{L}}_{0}^{-1^{\top}}={\mathbf{I}}_{T}$.

Clearly, when $r=0$, we have $\Sigma_{0}={\mathbf{I}}$ is a member of both the uniform and exponential families of kernel matrices, corresponding to uniform off-diagonals of $\rho_{0}=0$ or exponentially decaying off-diagonals with rate $\gamma_{0}=\infty$. In this case we can set ${\mathbf{L}}_{0}={\mathbf{I}}$ too.

On the other hand, for $r>0$, it is in general difficult to say much more about an individual sequence’s $\Sigma_{0}$, given that different sequences will have repeated tokens in different locations. Moreover, the naive approach of ignoring the repeated tokens and treating $\Sigma_{0}={\mathbf{I}}$ leads to increasing diagonal values of $\Sigma_{l}$ (i.e. activation norms) at large depth without corrections, as shown in Fig. 6.We find that $r\approx 0.008$ for sentencepiece tokenisations like we use in WikiText-103 and C4, but $r\approx 0.05$ for character-level prediction like in EnWiki-8. This imbalance between blocks could be problematic for training dynamics, and also is incompatible with frameworks like DKS and TAT which suppose that diagonal values of $\Sigma_{l}$ are constant across blocks and locations, usually set to 1.

To circumvent this, we will derive our modifications by considering the average input kernel matrix $\bar{\Sigma}_{0}$ (averaged over different sequences), under the assumption that repeated tokens occur independent of location:

By linearity of Eq. 11 in $\Sigma_{0}$ (and because we are controlling ${\mathbf{A}}_{l}$ to be input-independent at initialisation in Section 3), it also follows that the average depth l kernel matrix $\bar{\Sigma}_{l}$ satisfies:

so we can modify our SPA approaches to control the average kernel matrix $\bar{\Sigma}_{l}$ instead.

For U-SPA, the situation is more straightforward as $\bar{\Sigma}_{0}$ is a uniform kernel matrix with $\rho=r$, so it suffices to simply let $\rho_{0}=r$ instead of $\rho_{0}=0$.

For E-SPA, we are unable to view $\bar{\Sigma}_{0}$ as having exponentially decaying off-diagonals, and so we keep $\gamma_{0}=\infty$ which translates to ${\mathbf{L}}_{0}={\mathbf{I}}$ and $\Pi_{l}={\mathbf{L}}_{l}$. Instead, to help us understand the effect of repeated tokens (to motivate our modifications), we first expand on Eq. 13 to simplify a little:

Note that all terms in Eq. 14 are easily computable (given that \big{(}\Sigma_{l}(\gamma_{l})\big{)}_{i,j}=\text{exp}(-\gamma_{l}|i-j|)) and ${\mathbf{L}}_{l}$ has an analytic form provided in Lemma 1). So we can use Eq. 14 to compute the expected diagonal $(\bar{\Sigma}_{l})_{i,i}$ for each location $i$ (where the expectation is taken over different sequences), which we denote by a diagonal matrix $\bar{D}_{l}$:

Thus, we propose to replace ${\mathbf{A}}_{l}={\mathbf{L}}_{l}{\mathbf{L}}_{l-1}^{-1}$ with ${\mathbf{A}}_{l}=\bar{D}_{l}^{-\frac{1}{2}}{\mathbf{L}}_{l}{\mathbf{L}}_{l-1}^{-1}\bar{D}_{l-1}^{\frac{1}{2}}$, setting $\bar{D}_{0}={\mathbf{I}}$ by default. This means that our product $\Pi_{l}=\prod_{i=1}^{l}{\mathbf{A}}_{i}=\bar{D}_{l}^{-\frac{1}{2}}{\mathbf{L}}_{l}$, and Eq. 13 is updated to:

Though our modifications for repeated tokens only consider averages across different sequences, we find that for individual sequences they still lead to well behaved diagonal values of $\Sigma_{l}$, as shown in Fig. 6. Moreover, we find that their effect on off-diagonals is still favourable for individual sequences, as shown in Fig. 1.

In this section we describe how we set the uniform off-diagonals $(\rho_{l})_{l}$ and the exponential decay rates $(\gamma_{l})_{l}$ in SPA, at different depths.

Recall that Theorems 1 and 2 show that we are free to choose $(\rho_{l})_{l=0}^{L}$ and $(\gamma_{l})_{l=0}^{L}$ such that $(\rho_{l})_{l=0}^{L}$ increase with depth and $(\gamma_{l})_{l=0}^{L}$ decrease with depth. Moreover, $\rho_{0}=0$ (or the shared token fraction $r$, as per Appendix B) and $\gamma_{0}=\infty$ at the input layer, whilst we have found $\rho_{L}=0.8$ and $\gamma_{L}=0.005$ to be good default values for the last block.

For U-SPA, in terms of setting how $(\rho_{l})_{l}$ vary with depth, we tried different polynomial rates of increase with depth from $\rho_{0}$ to $\rho_{L}$, but did not observe a noticeable difference in performance across different rates so chose to increase from $\rho_{0}$ to $\rho_{L}$ linearly in depth.

For E-SPA, we choose $(\gamma_{l})_{l}$ so that the diagonal elements of the attention matrices ${\mathbf{A}}_{l}$ are constant across blocks, akin to using a constant shortcut weight $\alpha$ over different blocks. From Theorem 3, we have that the diagonal entries of ${\mathbf{A}}_{l}$ satisfy:

where $a(\gamma)=\sqrt{1-\text{exp}(-2\gamma)}$ with inverse $\gamma(a)=-\frac{1}{2}\text{log}(1-a^{2})$. This means that for a set of positive decreasing $(\gamma_{l})_{l}$, there exist a corresponding set of decreasing $(a_{l})_{l}$ with values between and $1$.

Because $\gamma_{0}=\infty$, we have $a_{0}=1$, and likewise for a given $\gamma_{L}$ we can compute $a_{L}$.

Thus to have constant diagonal values of ${\mathbf{A}}_{l}$ over different blocks $l$, we choose to set

We found this scheme to work well empirically, and note the similarity to other works which have discussed the choice of how to scale branches in residual architectures (Hayou et al., 2021). We leave further study of choosing how to set $(\rho_{l})_{l}$ and $(\gamma_{l})_{l}$ to future work.

Appendix D E-SPA algorithm

In Alg. 1 we present pseudocode to construct a trainable E-SPA masked attention layer.

Appendix E Compatibility of SPA with existing positional encodings

In Section 3.3, we showed how to control the attention matrix at initialisation, using bias matrices and location-dependent rescaling, as well as making the query-key dot product, $\frac{1}{\sqrt{d^{k}}}{\mathbf{Q}}({\mathbf{X}}){\mathbf{K}}({\mathbf{X}})^{\top}{=}\frac{1}{\sqrt{d^{k}}}{\mathbf{X}}{\mathbf{W}}^{Q}({\mathbf{X}}{\mathbf{W}}^{K})^{\top}$, zero at initialisation. This scheme is used in our SPA methods. We detailed several ways to achieve this, and in practice we chose to initialise ${\mathbf{W}}^{Q}$ to zero, and initialise ${\mathbf{W}}^{K}$ as usual (Gaussian fan-in or orthogonal).

In this section we show how zero initialising the query-key dot product is also possible when using two standard positional encodings: RoPE (Su et al., 2021) and Relative (Shaw et al., 2018; Huang et al., 2019; Dai et al., 2019). This means that we can use our methods in Section 3.3 in combination with these positional encoders.

Let us denote the unscaled query-key dot product as ${\mathbf{S}}={\mathbf{X}}{\mathbf{W}}^{Q}({\mathbf{X}}{\mathbf{W}}^{K})^{\top}$.

For RoPE, the $(i,j)$ query-key dot product, ${\mathbf{S}}_{i,j}$, is modified from

For Relative positional encoding, we take the scheme from (Dai et al., 2019). In that case, the $(i,j)$ query-key dot product, ${\mathbf{S}}_{i,j}$, is modified from

Appendix F Using normalised skip connections with E-SPA

In this section, we describe how to combine our E-SPA method with normalised skip connections in our attention blocks:

where $\alpha=0$ is the skipless setting that our methods are originally designed for.

The general gist is that we will look at the combined effect, on the kernel matrix after the residual attention block, of the residual branch and the diagonal terms in the attention matrices for preserving signal propagation, and approximate the combination to match the setting where we are without skips, i.e. $\alpha=0$.

To combine E-SPA with normalised skip connections, we consider how the cosine-similarity between two locations $T_{1}$, $T_{2}\leq T$ is affected through the normalised skip connection. Suppose we have an input kernel matrix:

where $(\Sigma)_{i,i}=1,\forall i$ and $d$ is the width. Then after the residual attention block, Eq. 18, we now have:

Either we have ${\mathbf{W}}^{V}$ to be an orthogonal matrix sampled uniformly at random from the Haar measure, or ${\mathbf{W}}^{V}\overset{i.i.d.}{\sim}\mathcal{N}(0,\frac{1}{d})$. In both cases we have $\frac{1}{d}{\mathbf{A}}{\mathbf{X}}{\mathbf{W}}^{V}{\mathbf{W}}^{V^{\top}}{\mathbf{X}}^{\top}{\mathbf{A}}^{\top}$ going towards $\frac{1}{d}{\mathbf{A}}{\mathbf{X}}{\mathbf{X}}^{\top}{\mathbf{A}}^{\top}$, and the cross terms Eq. 20 converging to 0 for large $d$.

Thus, we can consider the large $d$ approximation:

Then, if we look at the inner product between the $T_{1}$ and $T_{2}$ locations for locations $T_{1}\neq T_{2}$ such that $T_{1},T_{2}>1$, we have:

where $\lambda={\mathbf{A}}_{T_{1},T_{1}}={\mathbf{A}}_{T_{2},T_{2}}$ is the constant diagonal of the attention matrices in E-SPA, c.f. Theorem 3, and $\nu=\alpha^{2}+(1-\alpha^{2})\lambda^{2}$.

Moreover, we have defined $\delta=\sum_{{i\neq T_{1}\cup j\neq T_{2}}}{\mathbf{A}}_{T_{1},i}\Sigma_{i,j}{\mathbf{A}}_{T_{2},j}$ as it is a term that is not possible to control with only knowledge of $\Sigma_{T_{1},T_{2}}$ and will vary from sequence to sequence, and hence we argue can be discarded from a signal propagation perspective. For example, one could consider a sequence where the kernel matrix $\Sigma$ has all off-diagonals equal to 0 apart from $\Sigma_{T_{1},T_{2}}$, and $T_{1}$ could be sufficiently distant from $T_{2}$ such that there is no $i$ for which ${\mathbf{A}}_{T_{1},i}$ and ${\mathbf{A}}_{T_{2},i}$ are large at the same time (which can be seen using the analytic form of ${\mathbf{A}}$ given in Theorem 3).

Thus, from Eq. 23, we see that after an residual attention block, an input cosine similarity of $\Sigma_{T_{1},T_{2}}$ between locations $T_{1},T_{2}$ is diluted by a factor of

with shortcut weight $\alpha$ and attention diagonal probability $\lambda$. Our approximation then seeks to preserve this factor $\nu$ when $\alpha>0$ to match the skipless case.

Now, in the skipless case $\alpha=0$ described in Section 3, at block $l$ we suppose that the incoming kernel matrix $\Sigma$ has exponentially decaying off-diagonals with rate $\gamma_{l-1}$, and that we construct the attention matrix ${\mathbf{A}}$ so that the output kernel matrix has exponentially decaying off diagonals with rate $\gamma_{l}$. From Theorem 3, we see that this gives diagonal entries of ${\mathbf{A}}$ to be:

where $a(\gamma)=\sqrt{1-\text{exp}(-2\gamma)}$.

Thus, to preserve $\nu$ to match the case for $\alpha=0$, if we have shortcut weight $\alpha>0$ we need the attention matrix diagonal probability $\lambda_{\alpha}$ to satisfy:

In turn, when we have shortcut weight $\alpha$, this means that we need to choose our outgoing decay rate at block $l$, $\gamma_{l,\alpha}$, such that:

Inverting the definition of $a(\gamma)$, we see we need to set $\gamma_{l,\alpha}$:

To summarise, if we have a sequence $(\gamma_{l})_{l}$ of decreasing exponential decay rates, our proposed approximation when using shortcut weights $\alpha$ at block $l$ is, using the notation of Alg. 1, to set $\gamma_{\text{in}}=\gamma_{l-1}$ as normal, and to set $\gamma_{\text{out}}=\gamma_{l,\alpha}$ from Eq. 25 in order to preserve the signal propagation from the combined residual attention block Eq. 23. We see that this approximation reduces to the standard skipless setting when $\alpha=0$, as a sanity check.

Appendix G Additional results

In this section we present additional results and ablations that were not included in Section 4.

In Fig. 7 we plot the evolution of normalised kernel matrices for transformers with skips and or RMSNorm normalisation, in addition to those for vanilla transformers in Fig. 1. We see that both skipless with RMSNorm (fourth row) and Post-LN (bottom row) converge to rank collapse at larger depths. The degeneration of skipless with RMSNorm is expected from the results of (Dong et al., 2021), and while the convergence to rank collapse is slower for Post-LN, it is still expected (Hayou et al., 2021; Noci et al., 2022). This is because the residual and shortcut branches are effectively given a constant weighting at all blocks in Post-LN, even as the network’s depth increases.

On the other hand, Pre-LN observes sensible signal propagation even at depth 100, as the positioning of the LN in the residual branch effectively downweights the residual branch at later blocks (De & Smith, 2020). Likewise, Pre-LN with skip weight $\alpha=0.98$ also observes faithful signal propagation, because each block is effectively downweighted. This effect means that standard Pre-LN’s (fifth row) kernel matrix increases elementwise faster with depth than Pre-LN with normalised skips (sixth row), despite both kernel matrices being qualitatively similar at block 100. Note, all methods besides our SPA methods used ALiBi positional encoder (Press et al., 2022) (detailed in Section H.2) and we observe both Pre-LN kernel matrices also have a recency bias, like our main method, E-SPA (third row).

In Table 1 we compared the training speeds for our skipless methods, and found that E-SPA outperforms both U-SPA and Value-Skipinit. In particular, we found E-SPA with a DKS-transformed GeLU activation without LN to perform best. In Table 4, we present the corresponding results but for validation perplexity. Again, we see that E-SPA is the best performing of our attention modifications, but in this case TAT with Leaky ReLU and no LN matches or outperforms DKS with GeLU.

Recall in Section 3.3 that for attention layers using SPA we seek to initialise weights such that the query-key dot product, $\frac{1}{\sqrt{d^{k}}}{\mathbf{X}}{\mathbf{W}}^{Q}({\mathbf{X}}{\mathbf{W}}^{K})^{\top}$, is zero or small at initialisation. In our main experiments, we achieve this by initialising ${\mathbf{W}}^{Q}=0$, and letting ${\mathbf{W}}^{K}$ to be initialised as normal. In Fig. 8, we assess the sensitivity of our E-SPA scheme to the scale of non-zero attention dot product, when ${\mathbf{W}}^{K},{\mathbf{W}}^{Q}$ are both orthogonally initialised but with scale $\sigma$ (i.e. at each layer, both ${\mathbf{W}}^{K},{\mathbf{W}}^{Q}$ are initialised as two independent uniform orthgonal matrices multiplied by $\sigma$). We see that for small initialisation scales, there is little effect of varying initial scale, but that training performance degrades at larger scales, when our attention matrix reverse engineering in Section 3.3 (which expects small or zero query-key dot product at initialisation) is less precise.

Recall in Section 3 that our kernel matrix evolution for attention-only transformers is exact at finite widths using orthogonally initialised weight matrices, and will be approximate using standard (Gaussian) fan-in initialised In Fig. 9, we ablate over using orthogonally initialised weight matrices, compared to Gaussian fan-in initialisation. Across activations, we see that E-SPA with orthogonal initialisation slightly outperforms Gaussian fan-in initialisation (by around 0.15 train loss).

As noted in Appendix E, our methods are compatible with several standard positional encodings, and by default all our experiments use the popular RoPE (Su et al., 2021) positional encoding. In Fig. 10 and Table 5, we assess the effect of removing positional encodings, on training and validation performance respectively, in our skipless methods. We see that all methods are improved when combined with RoPE, however the improvement is most mild in E-SPA, which as discussed has an in-built recency bias akin to a positional encoder. Moreover, E-SPA without additional positional encoding still outperforms all other approaches, including U-SPA (which on its own has no notion of position) with RoPE.

In Fig. 11, we provide an equivalent plot to Fig. 3 using a Stable ResNet (Hayou et al., 2021) rescaling of the shortcut weights ($\alpha=1,\beta=O(\frac{1}{\sqrt{L}})$). In this case, the shortcut weight is always $1$ and the residual weight $\beta$ is uniform across blocks and scales as $O(\frac{1}{\sqrt{L}})$ in depth $L$. Hayou et al. (2021) showed that such a scaling leads to non-degenerate signal propagation in large depth MLPs/CNNs without normalisation, and Noci et al. (2022) showed that the stable scaling prevents rank collapse in transformers without normalisation. We see in Fig. 11 that for large enough $\beta$, the stable residual weighting with normalisation matches the training speed of the default transformer (which is unsurprising given that once $\beta=1$ it is exactly default Pre-LN). However, there is a small but consistent gap without normalisation (with optimal $\beta=0.1$). Here, $L=36$, so $\frac{1}{\sqrt{L}}=\frac{1}{6}\approx 0.17$, or alternatively if we count the 72 nonlinear layers (one self-attention and element-wise nonlinearity for each transformer block), we have $\frac{1}{\sqrt{72}}\approx 0.12$.

Recall that in a transformer block Eq. 1, there are two distinct skip connections: one for the self-attention block and one for the MLP block. Moreover, we observed a training speed gap when we remove both skip connections in Fig. 2. This leads us to ask if it is possible that only one of the skips, MLP or attention, is causing this gap in training speed. In Fig. 12, we investigate this by varying the MLP shortcut weight for skipless attention blocks (left) and varying the attention shortcut weight for skipless MLP blocks (right). For all skip connections (for both MLP and self-attention blocks), we use a normalised skip connection ($\beta=\sqrt{1-\alpha^{2}})$. We observe that removing either skip connection results in a comparable loss of training speed (the default Pre-LN on WikiText-103 obtained train loss of 1.76 after 100K steps), although having one is still better than having neither. Moreover, we observe that the attention and MLP blocks prefer slightly different shortcut weights, with dense shortcuts performing better on slightly lower weightings ($\alpha=0.9$ or $0.97$) compared to attention shortcuts ($\alpha=0.98$ or $0.99$), whereas our experiments in Figs. 3 and 11 use a joint weighting for both.

In Fig. 13 we compare the training performance of our various vanilla transfomers to the default Pre-LN transformer on C4. This is akin to Fig. 4(a), which compared validation performance.

Typically, transformers are pre-trained on a large corpus of data before evaluation on a set of downstream tasks. To assess whether our conclusions about training performance in pre-training transfer over to downstream tasks, we assess models trained on C4 on 5 common sense downstream tasks: BoolQ (Clark et al., 2019), HellaSwag (Zellers et al., 2019), Winogrande (Sakaguchi et al., 2020), PIQA (Bisk et al., 2020), and SIQA (Sap et al., 2019). These datasets are commonly used to evaluate large pre-trained transformers (Brown et al., 2020; Rae et al., 2021; Smith et al., 2022; Hoffmann et al., 2022).

In Table 6, we see that the conclusions of pre-training on C4 largely carry over to the downstream tasks:

Among transformers trained for the same number of steps (50k), E-SPA beats U-SPA and Value-SkipInit each on 4 out of 5 downstream tasks. However, the default transformer outperforms skipless transformers on all tasks with the same amount of training.

With around 5 times longer training (200K and 300K steps), E-SPA achieves similar performance on downstream tasks to the standard transformer, outperforming on 2 out of 5 tasks (Winogrande and PIQA).

In Fig. 14 we see that 4.5x more training allows a vanilla E-SPA transformer to match the validation performance of a standard transformer on WikiText-103. This mirrors our findings on C4 in Fig. 4(b).

Fig. 15 shows the evolution of the empirically-computed normalized kernel matrix during training of a (finite width) vanilla E-SPA transformer on WikiText-103. The network has depicted has both attention and MLP blocks, which use DKS-transformer GeLU activations. We note that the untrained network shows good agreement with Fig. 1, despite the fact that Fig. 1 is computed for an attention-only network in the infinite width limit. We also note that while significant changes to the kernel matrix occur during training, it retains the property of being larger close the diagonal.

Appendix H Implementation details

We first describe all additional general implementation details for our experiments, before going into details relevant for individual results.

We present experiments on WikiText-103 (Merity et al., 2017) and C4 (Raffel et al., 2019). For both datasets, we use the SentencePiece tokeniser (Kudo & Richardson, 2018) with vocabulary size $|V|=32,000$. For WikiText-103 we use sequence length of 512 for both training and validation. For C4, the sequence length is 2048 for both.

In all our experiments apart from Table 3, the model width $d=1024$ across all blocks. We use 8 head multi-head attention, so that $d^{k}=128$. The MLP block consists of a single hidden layer with width $4d$, with input and output dimensions both equal to $d$. All our experiments use the Pre-LN transformer block Eq. 1, rather than Post-LN. On WikiText-103, we use a 36 block transformer by default for all experiments apart from Table 3. On C4, we use a 32 block transformer due to memory constraints. Any normalisation layer we consider is RMSNorm (Zhang & Sennrich, 2019), which is simpler than Layer Normalisation (Ba et al., 2016) and is commonly used in transformers (Rae et al., 2021). By default, our models use RoPE positional encoder (Su et al., 2021) (apart from the ablation in Fig. 10).

By default, all weight matrices are fan-in initialised $\mathcal{N}(0,\frac{\sigma^{2}}{\text{fan-in}})$ with $\sigma=1$. The two exceptions for this are: 1) when using orthogonal intiailisation, we use the scaled-corrected uniform orthogonal initialisation (Martens et al., 2021) with scale $\sigma=1$ (which for square matrices is just an orthogonal matrix sampled from the Haar measure Meckes (2019) multiplied by $\sigma$), and 2) for the parameter matrix immediately after the activation we set $\sigma$ to take input activation norm (“q-values” or diagonals of Gram matrices) of 1 to output activation norm of 1. In the latter case, $\sigma=1$ by construction for activations transformed by DKS/TAT Martens et al. (2021); Zhang et al. (2022). All bias parameters are initialised to 0.

For DKS (Martens et al., 2021) we set slope parameter $\xi=5$, and for TAT (Zhang et al., 2022) with leaky ReLU, we set $\eta=0.9$. Both values were chosen by a small hyperparameter sweep on WikiText-103. The DKS and TAT transformations are chosen without consideration of the attention blocks, where the transformer can be viewed as an MLP (potentially with residual connections). Unless stated otherwise, all skipless transformers used a DKS-transformed GeLU as the nonlinearity in the MLP block by default.

We use Adam optimiser (Kingma & Ba, 2014) with global gradient clipping of 0.1 by default (Pascanu et al., 2013). We do not use weight decay in our experiments.

We use mini-batching of 16 sequences for WikiText-103 and 8 for C4, due to memory constraints. Unless stated otherwise, we train for 100K steps on WikiText-103 and 50K steps for C4.

H.2 Additional implementation details for individual experiments

We calculate the kernel matrix $\Sigma$ evolution directly in $T\times T$ kernel matrix-space, where $T=100$. Our input kernel matrix $\Sigma$ is sampled assuming a fraction $r=0.02$ of repeated tokens, with value of $1$ if the token is repeated, and 0 else. For all configurations of skip/normalisation/attention modifications corresponding to a row of Fig. 7, we use 8 heads.

We now detail how each operation in any configuration of Fig. 7 affects the kernel matrix. From this it should be possible to reconstruct the kernel evolution for any row in Fig. 7. The 3 possible operations are: 1) attention, 2) skip connection, or 3) LN/RMSNorm operation.

Attention Because our SPA methods (second and third rows) are agnostic to the number of heads at initialisation (all attention matrices in a self-attention block are the same across heads at initialisation), we apply Eq. 6 directly, so that a single attention block amounts to $\Sigma\leftarrow{\mathbf{A}}\Sigma{\mathbf{A}}^{\top}$ for attention matrix ${\mathbf{A}}$ and incoming kernel matrix $\Sigma$. Our E-SPA method uses $\gamma_{L}=0.005$ and our U-SPA method uses $\rho_{L}=0.8$.

For all other rows, the self-attention operation uses ALiBi (Press et al., 2022), a popular positional encoder which uses head-dependent pre-softmax bias matrices. More specifically, from the default pre-softmax bias matrices given by ALiBi (for 8 heads), we obtain 8 attention matrices $\{{\mathbf{A}}_{h}\}_{h=1}^{8}$ using the softmax operation (which is exact assuming zero query-key dot product at initialisation). Because the different heads in transformers are typically concatenated along 8 equal size fractions of the total width $d$, the kernel evolution of an attention block on kernel matrix $\Sigma$ with 8-head ALiBi corresponds to (Martens et al., 2021):

For a skip connection with shortcut weight $\alpha$ and residual weight $\beta$, if $\Sigma_{\text{attn}}(\Sigma)$ denotes the output of a kernel matrix $\Sigma$ after an self-attention operation (i.e. from point 1. above), then an incoming kernel matrix $\Sigma$ gets mapped to:

For a normalisation operation, the incoming kernel matrix $\Sigma$ gets mapped to:

For experiments on WikiText-103 with 100K steps, for our U-SPA transformers we tuned $\rho_{L}\in\{0.6,0.8\}$, and for our E-SPA transformers we tuned $\gamma_{L}\in\{0.005,0.2\}$. For all other settings (longer/deeper training on WikiText-103 or any C4 experiment), we used the default $\gamma_{L}=0.005$ and $\rho_{L}=0.8$. All hyperparameters throughout our work are tuned based on training loss.

All experiments with skip connections (i.e. shortcut weight $\alpha\neq 0$ for either the MLP or attention block) use untransformed GeLU activation in the MLP blocks. We combine E-SPA with normalised skips as described in Appendix F. We note that for high shortcut weights $\alpha$ and small final decay rate $\gamma_{L}$, then the value of attention matrix diagonal $\lambda_{\alpha}$, Eq. 24, may not be real. This is because the input to the square root in Eq. 24 may be negative. To get by this, we tune $\gamma_{L}\in\{0.005,0.2,0.4,0.6\}$ when using a 36 block transformer on WikiText-103 and $\gamma_{L}\in\{0.005,0.2,0.4,0.7,1\}$ for Table 2, which used a 32 block transformer on C4.

For Table 2, the normalised skip connections had separately tuned attention and MLP shortcut weights, as we observed a difference in the optimal shortcut weight for self-attention vs MLP blocks in Fig. 12. We tuned the attention shortcut weight in the range $\alpha\in\{0.98,0.99,0.995,0.9975,0.999\}$ and the MLP shortcut weight in the range $\alpha\in\{0.98,0.99,0.995\}$. Likewise, both the stable residual weights were tuned in $\beta\in\{0.05,0.1,0.15,0.2\}$ separately for self-attention and MLP skips. The selected shortcut/residual weights (using validation performance) are presented in Table 7.

Due to memory constraints, for our deeper networks in Table 3 we use width $d=512$ rather than $1024$, with 8 heads to give $d^{k}=64$. All depth scaling runs used a DKS-transformed GeLU with $\xi=5$.

Appendix I Theoretical results

In this section we state and prove our theoretical results, including Theorems 1 and 2:

(Non-negativity for U-SPA) Let $\Sigma{=}(1-\rho){\mathbf{I}}_{T}{+}\rho\bm{1}\bm{1}^{\top}$ and $\Sigma^{\prime}{=}(1-\rho^{\prime}){\mathbf{I}}_{T}{+}\rho^{\prime}\bm{1}\bm{1}^{\top}$, with respective (positive) Cholesky factors ${\mathbf{L}}$ and ${\mathbf{L}}^{\prime}$. Then if $\rho\leq\rho^{\prime}$, we have ${\mathbf{L}}^{\prime}{\mathbf{L}}^{-1}$ is elementwise non-negative.

(Non-negativity for E-SPA) Let matrices $(\Sigma)_{i,j}=\emph{exp}(-\gamma|i-j|))$ and $(\Sigma^{\prime})_{i,j}=\emph{exp}(-\gamma^{\prime}|i-j|))$ with respective (positive) Cholesky factors ${\mathbf{L}}$ and ${\mathbf{L}}^{\prime}$. Then if $\gamma\geq\gamma^{\prime}$, we have ${\mathbf{L}}^{\prime}{\mathbf{L}}^{-1}$ is elementwise non-negative.

We prove Theorem 1 second as it is more involved.

We actually prove Theorem 2 as a corollary of Theorem 3, which provides the analytic form for ${\mathbf{L}}^{\prime}{\mathbf{L}}^{-1}$.

Let matrices $(\Sigma)_{i,j}=\emph{exp}(-\gamma|i-j|))$ and $(\Sigma^{\prime})_{i,j}=\emph{exp}(-\gamma^{\prime}|i-j|))$ with respective (positive) Cholesky factors ${\mathbf{L}}$ and ${\mathbf{L}}^{\prime}$. Then if $\gamma\geq\gamma^{\prime}>0$, we have ${\mathbf{A}}={\mathbf{L}}^{\prime}{\mathbf{L}}^{-1}$ takes the following form:

where $a(\gamma)=\sqrt{1-\emph{exp}(-2\gamma)}$, and likewise $a(\gamma^{\prime})=\sqrt{1-\emph{exp}(-2\gamma^{\prime})}$

From Theorem 3, we have the analytic form of ${\mathbf{A}}$. Clearly $a(\gamma),a(\gamma^{\prime})$ are positive, and moreover if $\gamma^{\prime}\leq\gamma$, then $\text{exp}(-\gamma^{\prime})-\text{exp}(-\gamma)$ is non-negative.

Finally, because $\frac{a(\gamma^{\prime})}{a(\gamma)}\leq 1$ when $\gamma^{\prime}\leq\gamma$, then $\text{exp}(-\gamma^{\prime})-\frac{a(\gamma^{\prime})}{a(\gamma)}\text{exp}(-\gamma)$ is non-negative too. ∎

To prove Theorem 3, we first compute what the analytic form of Cholesky factor ${\mathbf{L}}$ for $(\Sigma)_{i,j}=\text{exp}(-\gamma|i-j|))$ takes in Lemma 1.

Let $(\Sigma)_{i,j}=\emph{exp}(-\gamma|i-j|))$ with (positive) Cholesky factors ${\mathbf{L}}$ such that ${\mathbf{L}}{\mathbf{L}}^{\top}=\Sigma$. Then, we have:

It is clear that $\Sigma$ is positive semi definite, as it is the covariance matrix of a stationary Ornstein-Uhlenbeck process, hence a Cholesky factor must exist. We now show that it is ${\mathbf{L}}$, Eq. 28.

If we define $l=\text{min}(m,n)$, then we have:

By the uniqueness of (positive) Cholesky factors, the proof is complete. ∎

We want to show $({\mathbf{A}}{\mathbf{L}})_{k,l}=({\mathbf{L}}^{\prime})_{k,l},\forall k,l$. This is clearly true for the top diagonal $k=l=1$.

We now show this for the rest of the first column, when $l=1,k>1$:

Applying the geometric sum to Eq. 29 yields:

I.2 Proof of Theorem 1

Before diving into the actual proof of the theorem we will derive several useful properties and notations. First we make a slight notational change from the main text of the theorem by replacing $\Sigma$, which depends on $\rho$, with $C_{n}(x)$, where $n$ represents the size of the matrix, while $x$ replaces $\rho$. In addition, since the case of $\rho=\rho^{\prime}$ (or $x=y$ in the rest of the proof) is trivial, since then the resulting matrix is the identity, we will restrict ourselves to dealing with the case where $0<x<y<1$. In any of the mathematical derivations we will denote with capital English letters (e.g. $A,B,C,T$) any temporary expressions, that will be expanded on the following lines. Note that these are never general definitions, so they might be used multiple times for different expressions.

We will denote vectors and vector functions in bold and scalars and scalar function in standard font.

Let $\bm{1}_{n}$ be the $n$-dimensional vector with only ones:

We will denote with $\bm{x}_{n}$ the $n$-dimensional vector with only $x$’s:

First, we define the linear map $p_{n}(x)$ as:

The vector $\bm{1}_{n}$ is an eigenvector of $C_{n}(x)$ with an eigenvalue $p_{n-1}(x)$.

Directly calculating the $k$-th entry of the product $C_{n}(x)\bm{1}_{n}$ gives:

$[C_{n}(x)\bm{1}_{n}]_{k}=\sum_{i=1}^{n}[C_{n}(x)]_{k,i}=\sum_{i=1}^{n}\delta_{i}^{k}+(1-\delta_{i}^{k})x=1+(n-1)x$

$C_{n}(x)\bm{1}_{n}=p_{n-1}(x)\bm{1}_{n}$. ∎

The vector $\bm{1}_{n}$ is an eigenvector of $C_{n}(x)^{-1}$ with an eigenvalue $\frac{1}{p_{n-1}(x)}$.

$C_{n}(x)^{-1}\bm{x}_{n}=\frac{1}{p_{n-1}(x)}\bm{x}_{n}$.

$\bm{x}_{n}^{T}C_{n}(x)^{-1}\bm{x}_{n}=\frac{nx^{2}}{p_{n-1}(x)}$.

Further we define the following useful functions:

I.2.2 Lemmas

If $0<x<y<1$ then all entries of the vector $\bm{v}_{n}(x,y)$ are non-negative - $[\bm{v}_{n}(x,y)]_{i}\geqslant 0$.

If $0<x<y<1$ then the function $g_{n}(x,y)=\frac{y(1-y)p_{n-1}(x)}{r_{n}(x)r_{n}(y)}-\sqrt{\frac{d_{n+1}(y)}{d_{n+1}(x)}}\frac{x}{p_{n}(x)}$ is non-negative.

If $0<x<y<1$ then the function $f_{n}(x,y)=\frac{xy(1-y)}{r_{n}(x)r_{n}(y)}+\sqrt{\frac{d_{n+1}(y)}{d_{n+1}(x)}}\frac{x}{p_{n}(x)}-\sqrt{\frac{d_{n}(y)}{d_{n}(x)}}\frac{x}{p_{n-1}(x)}$ is negative.

The partial sum the functions $f_{n}(x,y)$ from $k+1$ to $n-1$ will be denote by $h_{k}^{n}(x,y)=\sum_{i=k+1}^{n-1}f_{i}(x,y)$, which from Lemma 4 follow are always negative.

I.2.3 Main proof

First for $n=1$ we have that $C_{1}(x)=C_{1}(y)=$, which implies that $L_{1}(x)=L_{1}(y)=$ and the condition is trivially satisfied.

Now assuming that the statement is true for all integers up to $n$, we will prove that it also holds for $n+1$:

Using the fact that $L_{n}(y)L_{n}(x)^{-1}$ is a lower triangular and non-negative matrix by the inductive assumption combined with Lemma 2 and the fact that $d_{n}(x)\geqslant 0$ it follows that $L_{n+1}(y)L_{n+1}(x)^{-1}$ is also a lower triangular non-negative matrix.

I.2.4 Proof of Lemma 2

First we will inspect the evolution of $\bm{v}_{n}(x,y)$ as we increase $n$:

Using the definition of from Lemma 3, Lemma 4 and Definition 2 we have that:

Thus proving the lemma reduces to proving that:

Expanding on the equation that we need to prove is positive:

Now we turn our attention to the sum in the middle:

Let’s define the partial sum in the brackets as:

First for $n=k$, we have that $I^{n}_{n}(x)=\frac{1}{p_{n}(x)p_{n-1}(x)}$ which clearly satisfies the above equation. Assuming this is true for $n$, we will now show it holds for $n+1$:

Which concludes the proof. This now means that:

From the fact that $h^{n}_{k}(x,y)\geqslant 0$ it impliest that $\mathcal{L}^{n}_{k}$ is a decreasing function of $n$, hence to we only need to show that for any fixed $x$ and $y$ $\mathcal{L}^{\infty}_{k}(x,y)$ is positive.

For this, we need to take the limit of the second and third term in the above equation.

where the last line is true since the denominator is of higher degree in $n$.

Taking the limit of the second term corresponds to computing the limit:

All of the limits must be taken for fixed $x$ and $y$ (e.g. we can’t have them approach or $1$ simultanously with $n$).

I.2.5 Proof of Lemma 3

From the definition of $g_{n}(x,y)$ and the fact that $p_{n}(x)$ and $r_{n}(x)$ are non-negative functions we can conclude that:

Denoting with $A$ the left hand side of the second equation we have:

Hence for $x\leqslant y$ we have that $A\geqslant 0\Leftrightarrow g_{n}(x,y)\geqslant 0$.

I.2.6 Proof of Lemma 4

First we will prove that $C-B\geqslant 0$

With this we can now conclude that $f_{n}(x,y)\geqslant 0\Leftrightarrow A^{2}-(C-B)^{2}\geqslant 0$

Hence with this we can conclude that $f_{n}(x,y)\geqslant 0\quad\forall n$.