Deep Transformers without Shortcuts: Modifying Self-attention for Faithful Signal Propagation
Authors: Bobby He
James Martens
Guodong Zhang
Aleksandar Botev
Andrew Brock
Samuel L Smith
Yee Whye Teh
Venue: ICLR 2023 poster (2023)
Keywords: signal propagation
neural networks and kernels
deep transformers
self-attention
residual connections
layer normalisation
rank collapse
positional encoding

Abstract: Skip connections and normalisation layers form two standard architectural components that are ubiquitous for the training of Deep Neural Networks (DNNs), but whose precise roles are poorly understood. Recent approaches such as Deep Kernel Shaping have made progress towards reducing our reliance on them, using insights from wide NN kernel theory to improve signal propagation in vanilla DNNs (which we define as networks without skips or normalisation). However, these approaches are incompatible with the self-attention layers present in transformers, whose kernels are intrinsically more complicated to analyse and control.  And so the question remains: \emph{is it possible to train deep vanilla transformers?} We answer this question in the affirmative by designing several approaches that use combinations of parameter initialisations, bias matrices and location-dependent rescaling to achieve faithful signal propagation in vanilla transformers. Our methods address various intricacies specific to signal propagation in transformers, including the interaction with positional encoding and causal masking. In experiments on WikiText-103 and C4, our approaches enable deep transformers without normalisation to train at speeds matching their standard counterparts, and deep vanilla transformers to reach the same performance as standard ones after about 5 times more iterations.

Published as a conference paper at ICLR 2023
DEEP TRANSFORMERS WITHOUT SHORTCUTS: MODIFYING SELF- SELF-ATTENTIO FOR FAITHFUL SIGNAL PROPAGATION
Bobby He' James Martens Guodong Zhang
Aleksandar Botey"
Andrew Brock² Samuel L. Smith Yee Whye Teh 1.2
'University Oxford, *DeepMind
Correspondence to: babby. www.thts ox uk, jtnnsssstttttteeggggggll
ABSTRACT
Skip connections and normalisation layers form two standard architectural com- ponents that are ubiquitous for the training of Deep Neural Networks (DNNNs, but whose precise roles are poorly understood. Recent approaches such as Deep Ker- nel Shaping have made progress towards reducing our reliance on them, using in- sights from wide NN kernel theory to improve signal propagation in vanilla DNNs (which we define as networks without skips or normalisation layers). However, these approaches are incompatible with the self-attention layers present in trans- formers, whose kernels are intrinsically more complicated to analyse and control. And so the question remains: is it possible to train deep vanilla transformers? We answer this question in the affirmative by designing several approaches that use combinations of parameter initialisations, bias matrices and location-dependent rescaling to achieve faithful signal propagation in vanilla transformers. Our meth- ods address several intricacies specific to signal propagation in transformers, in- cluding the interaction with positional encoding and causal masking. In experi- ments on WikiText- 103 and C4, our approaches enable deep transformers without normalisation to train at speeds matching their standard counterparts, and deep vanilla transformers to reach the same performance as standard ones after about 5 times more iterations.
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 general principle for how to use these components in new architectures (assuming they are even applicable) remains unknown, and their roles in existing
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 for the success of residual architectures is that they have superior signal propagation compared to vanilla DNNs (c.g. Balduzzi al., 2017: Xiao al., 2018; Hayou 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 al., 2016; Schoenholz et al., 2017).
al., 2017). One proposed explanation
Recently, using signal propagation principles to train DNNs at high depths, without the skip connec- tions 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 en- ables very deep vanilla networks to be effectively trained, although at significantly reduced speeds

---

Published as a conference paper at ICLR 2023
Layer 1
Layer
Layer 10
Layer 20
Layer 50
Layer 100
Standart
Unitarm SPM
Exparential
Figure 1: Normalised kernel matrices diag(F")) Eq 2i--diagg Ex)) (which are like kernel matrices except with cosine similarities instead of duner-produce at various depths for standard attention-only vanilla transformers and two of our proposed alternatives (Section 3). Standard attention-only vanilla transformers (top) quickly suffer from rank collapse where all entries of the normalised kemnel converge to whereas our approaches, U-SPA and E-SPA, maintain controlled signal propagation even at large depths. Morcover, our main method E-SPA (bottom) exhibits a recency bias, where cosine similarities corresponding to nearby pairs of locations are larger, akin to positional encoding. Equivalent plots for attention only transformers with skips and normal- isation can be found in Fig. 7.
compared to residual networks. Martens 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 of strong 2nd- order optimizers like K-FAC (Martens & Grosse, 2015). Zhang et al. (2022) extended ideas from DKS to larger class of activation functions, achieving near parity in terms of generalisation as well.
The key quantity that analysed in signal propagation is the DNN's initialisation-tine kernel, or more precisely, the approximate kernel given by the infinite width limit (Neal, 2012; Matthews al., 2018; Lee al., 2018; Yang. 2019). For MLPs, and for CNNs that use Delta-initialisation (Balduzzi et al., 2017: 7; Xiao et al., 2018), this kernel can be written as a simple recursion over layers that involves only 2D functions, facilitating straightforward analysis.
Unfortunately, the evolution of the kernel across layers of transformer is more complicated, and as result, existing approaches like DKS are not applicable to transformers (or indeed any architecture that contains self-attention layers). More concretely, if X, E R denotes length-7 sequence of activations at layer / of transformer, then the kernel matrix 2; = X,X, /d E RRTx for layer (or more precisely its limit as + oo), can be written as a function of the kernel matrix Σn-1 of the previous layer (Hron et al., 2020). In the case of self-attention layers, the relationship of Σ, on Σt-1 cannot be simplified decomposed into lower dimensional functions, leading recursion that is intrinsically high dimensional and harder to analyse or control.
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 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. rank-1 matrix, which
more subtle form of signal degradation can occur where Σ, converges is known as rank collapse (Dong et al., 2021). Dong et al. (2021) showed that skip connections essential to avoid the collapsed state: skipless transformers quickly converge rank collapse at large depths, which we corroborate in Fig. (top). Moreover, Noci et al. (2022) showed that rank collapse may lead to 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 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 ini- tialisations, 2) bias matrices, and 3) location-dependent rescaling, and highlight several intricacies
2

---

Published as a conference paper at ICLR 2023
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 train- able 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.
2 PROBLEM SETTING
Transformer models The input to transformer consists of sequence x==xx)][= over 'T' loca- tions consisting of tokens from a vocabulary V: H; E ( [V.. The model takes this sequence, and using a trainable embedding matrix E giVixd creates matrix of vector representations Xo E Rº xd by performing a direct look-up for each location: [Xo]e En, E RR.
After this, the sequence is successively transformed via series of L "transformer blocks", with independently initialised parameters. We will denote by X₁ E Txd the output sequence for block I, and will sometimes refer to the rows of X₁ representation vectors. Each transformer block consists of two component blocks that both employ a standard residual structure (He et al., 2016), computing the sum of residual branch which performs the main computation, and a shortcut branch (aka skip connection) which just copies the block's 's inputs to its output. The attention block, which is the first of these component blocks, applies layer normalisation (LN) (Ba et al., 2016) or RM- SNorm (Zhang & Sennrich, 2019), followed by multi-head attention (MHA), on its residual branch. The MLP block, which comes next, applies LN followed by a standard (typically shallow) MLP applied to each sequence element inde-
on its residual branch. The LN and MLP operations pendently (using shared parameters). so that information is only communicated between sequence elements via the MHA operation (which we will define later). In summary we have and
X - 1 + B MHA (RMSNNMM
X = aX, + (M
where the shorteut and residual weights a: are typically both 1. In this work we will focus on skipless transformers with a = 0 and / = 1, and on vanilla transformers, which are skipless transformers without normalisation layers. For simplicity will also devote much of our analysis to altention-only models, which are transformers without MIP blocks (so that X, X.).
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) corre- sponds 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.
Self-attention
Given an input sequence X E RTxxd, the self-attention mechanism computes
Attn(X)= V(X), with A AX) softmax (1
dE Q(X)K(X))
(2)
where the softmax function is applied row-wise. Q(X) = XW K(X) XWK and V(X) XWV denote the queries, keys and values respectively, with trainable parameters Wº, w Rdxd" and WV e ddd"", In practice, the attention mechanism, Eq. (2), is applied over "heads" (with independent parameters), giving rise to so-called multi-head attention:
MHA(X) & Concat (Attat (X).... Attms, (t)tms()))o (X))W) W
(3)
where wo E Rhd xd are trainable parameters and usually d² d d. In this case, we can define V Concat(WW,,... W where W n denotes the value parameters for head n.
We focus our investigation to models that perform next-token prediction: at location i-1 the model outputs prediction over the identity of the 7th target token, but using only information from input tokens through i- This corresponds to using a causal masked attention with mask M E RT satisfying M.j fi 2 j}, where the attention matrix A in Eq. (2) is replaced with
M Q())(() T r(( 1 M)) ((1-M)
A (X) softmax
(4)
V th
3

---

Published as a conference paper at ICLR 2023
where I is large positive constant that zeros the attention coefficients corresponding to future tokens, making A a lower triangular matrix.
Signal propagation in transformers 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 ab- sence of normalisation layers even with skip connections, and that downscaling the residual branch by setting 3 _ can alleviate this issue. This latter observation in line with previous find- ings concerning the benefits of downscaling residual weights in ResNets (Hanin & Rolnick, 2018; Zhang al., 2018 8; 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 downweighted skip as an alter- native 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 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 al., 2020; Liu et al., 2020).
VT
3 CONSTRUCTING TRAINABLE DEEP TRANSFORMERS WITHOUT SHORTCUTS
To date, the only strategy for rectifying rank collapse in transformers relies on skip/shorteut connec- tions, which "skip" around the trainability issues intrinsic 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 prop- agation 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 A does not vary across heads. If block / < L has attention matrix A; at initialisation, then the final block's ' representation X,, takes the following form:
X =
...AL-1...AA. -A.AA--...AA.))
(5)
where W- I' W/ E Rdxd can be made to be orthogonal at initialisation (so that W TWW-- Ia), if each wY and w are orthogonally initialised (assuming d' d). Going forward we will assume such an orthogonal initialisation, providing an ablation study in Fig. 9.
In that case, if we denote by Eo I1 A,Ar-1... A₁ to be the product of attention matrices up to the 1h block, then the location- wise kernel matrix Σ = X,XI E TRT
SoXJERTxx the input kernel matrix across locations, and
xT at block / simplifies to
(6) From this simplified formula for kernel matrices in deep attention-only transformers, we identify three requirements on (At);:
Σ = 11, En -1,T.
(i) Σn II Eg II, must be well-behaved at each block, avoiding degenerate situations such as rank collapse and exploding/vanishing diagonal values.
(ii) A; must be elementwise non-negative vZ (recalling that A constructed through the soft- max operation Eq. (4)).
(iii) A should be lower triangular VI, for compatibility with causal masked attention. 2
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.
We note that this formula for the kernel matrix exact even at finite widths, W is orthogonal at initiali- sation. This is contrast to the standard kernel correspondence of NNs at initialisation, where the usual kernel equations are only approximations at finite width that become increasingly accurate as width increases (Neal, 2012; Daniely ct al., 2016; Yang, 2019; Martens, 2021; Li et al., 2022).
We describe compatibility of methods with non-causal attention in Appendix A.
4

---

Published as a conference paper at ICLR 2023
3.1 IDENTITY ATTENTION: ISSUES AND 'ALUE-SKIPINGS
An obvious solution to our above requirements on (At), is the trivial one: A = Ir VI, where each sequence location attends only to itself. In this case, Σ,, = Eo perfectly preserves the input kernel matrix, and will avoid rank collapse assuming Eg is non-degenerate. Unfortunately, identity atten- tion isn't 't compatible with a viable solution to obtain trainable vanilla transformers. This is because to achieve A, = 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-Skiplni, based on SkipInit (De & Smith, 2020), and the related ReZero method (Bachlechner et al., 2021). In Value-Skiplnit, we modify the attention operation Attn(X) VX)), to
Attn(X) (aI 3A(X)) V(X)
with trainable parameters a and 3 that are initialised to 1 and 0 respectively. Thus, at initialisa- tion, 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 can apply the DKS or TAT frameworks (Martens et al., 2021; Zhang et al., 2022) achieve well-behaved signal propagation in the entire model.3 3 We note that Value-Skipinit is an approach to remove the standard skip connections in transformer blocks, Eq. (I), 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-Skipini as baseline for our main two methods, which we describe next.
3.2 SIGNAL PRESERVING ATTENTION METHOSS
Returning our requirements on (Ai); in Section 3, that controlling their product is key to achieving faithful signal propagation. Given that we are now interested in non-identity (Ai)., this becomes more difficult. To overcome this, consider A of the form
A L,
(8)
such that individual L cancel in the product, giving 11 = L, Then if Lo satisfies EoLL" Ir, we thus have
Σ LL
Assuming input embeddings are initialised independently, and no repeated tokens in the input se- quence. we have En Ir in the large width limit," and can thus take Lo = Ir. In practice, we will apply slight modifications to our methods account for repeated tokens, as detailed in Appendix B.
Now L is lower triangular, then from Eq. (9) we see it is simply a Cholesky factor for the kernel matrix Σn. By the uniqueness of Cholesky factors for PSD matrices up to sign, this means we just need to choose (2i to be a family of well-behaved kernel matrices that satisfy² our non- negativity constraint (ii) on A, = -, We identify two such families, which give rise to our main method, Signal Preserving Attention (SPA):
1. Uniform (U-SPA): En(p:) = (1 -x)1r=pp111, Here, the kernel matrix Ee(pt) has diagonal entries equal 1 and off-diagonal entries equal to pt. The condition Pt S pt+1 required for elementwise non negaatiivy of the A 's,, as shown in Theorem 1. Setting Pu = 0 yields identity input kernel matrix, and as long as PL < 1, we avoid rank collapse in the skipless attention-only setting.
jl). Here, diagonal entries are again 1, but now off- dif-oiigsssls decay exponentially for more distant locations, with decay rate n- 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 2 7+1 required for elementwise nog-negayy of the A 's, as established in Theorem 2. Setting no = 0o yields the identity input kernel matrix, and rank collapse will be prevented as long as 7 > O.
2. Exponential (E-SPA): (Ex(7)),, = exp(
"This is similar in principle to the Delta-initialisationn for CNNs (Balduzzi et al., 2017; Xiao et al., 2018), and modifications to GNNs to enable compatibility with TAT (Zaidi et al., 2022).
'Because the embedding matrix E is initialised with variance 1/fan-out we rescale the embeddings it produces by vā to get a Eg with ones on the diagonal.
"Constraint (iii) is satisfied as lower triangular matrices are closed under multiplication and inversion.
5

---

Published as a conference paper at ICLR 2023
We state and prove Theorems and 2 in Appendix I, and in particular provide closed-form solution for L 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 pro- gressively larger (in an elementwise fashion) as depth increases, as dictated by (pt)f21 and (n))= We find PL 0.8 and 7 0.005 to be good default choices for the final block, and describe how to vary Pr and "/ across depth in Appendix C. In Alg. in Appendix D, we summarise how to construct trainable self-attention layer with our main method E-SPA using ideas from Section 3.3, given an input-output decay rate pair Yin Your (in the notation of Alg. 1). Given a decreasing set of decay rates (m at block I we set Yin 7-- and Yout
71
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 Σ, has diagonals equal to for all I, 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, 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 outperform both U-SPA and Value-Skiplnit.
3.3 REVERSE ENGINEERING SELF-ATTENTION LAYERS AT INITIALISATION
In Section 3.2 we identified two families of lower-triangular non-negative attention matrices (Aı); which enable us to obtain well-behaved kernel matrices Σg at large depth L, and hence faithful signal propagation. It remains show how we can actually realise these attention matrices via parameter initialisation and minor modifications of the attention mechanism.
More precisely, we will show how for any given trianuurra A E RTXT with non-negative entries, the masked softmax self-attention layer Eq. (4) can be initialised and augmented such that its output is exactly Attn(X)
(X) at initialisation. To do so, we first define matrices D, P e R such that A- = DP, where D is diagonal with positive entries and P is lower triangular with row sums 1, Then
log(P and M the causal mask M.,, = i 2 j}. we set"
Attn -DP(X & P (X) =softmax
M
Q())(()))) +B th
-r(1 -I(- M) (10)
which reduces to Attn(X (X) = AV (with a data-independent attention matrix) at initialisation
the query-key dot product v ( = O. We note that zero query-key dot products occur
one uses a 1 scaling rather _ in the infinite width limit for independently initialised Wº, w (Yang, 2019). In practice, to achieve zero initial query-key dot product, we can initialise either: 1) e = 0, 2) W 0, 3) both w and W to have small initial scale (which achieves approximately zero dot product). similarly, and decided to use option 1): we = 0. An empirical evaluation of the sensitivity to the choice of initial scale in option 3) is provided in Fig. 8.
our experiments we found these three options to all perform e
In Eq. (10), D acts as a non-trainable location-depender rescaling, and is needed to realise arbitrary attention matrices since the softmax output P(X) is constrained to have row sums equal to 1.? Given target kernel matrices with I'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 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 non-trainable bias matrix being added to the logits. In our early experiments we tried including trainable gain parameter on the bias matrix, initialised to 1, in order to better preserve the model class, but found that this didn't 't have a significant impact on training performance.
"We take log(0)
or large negative constant, e.g. 10 in practice.
"In SPA, D also corrects for the fact that masked softmax attention tends to reduce the norms of represen- tation vectors for locations at the end of the sequence. To see this, if we have kemel matrix Σ XX /d with En, = 1, Vi, and softmax attention matrix A with row i a,, then (AEATH ),r =|a,x X| < || = ||a 1|3 But a; sums to I (as a softmax output), so ||a:]la < 1 with equality if and only if a, has exactly one non-zero entry (equal to 1). This holds for the first token (which can only attend to itself in masked attention)
(AEA )11 1, but not in general for later tokens, so (AEAA) )as will usually be less than 1, for i > 1.

---

Published as a conference paper at ICLR 2023
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 Fig. 10.
3.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 be compatible with skip connections too, both for the sake of generality, and because they might combine profitably.
MLP blocks Because MLP blocks operate on the representation vectors independently across locations, their effect on the kemne 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 Σn of an MLP block's input sequence to the kernel matrix Σ of its output sequence. In principle, one could modify SPA to account for this change the kernel matrix by taking A, LL where L, and L, are the Cholesky factors of Σn and En respectively. Unfortunately, there is no guarantee that the resulting A would be elementwise non-negative in general. In our experiments 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 exactly identity), which will happen as we increase the shortcut weight relative to the residual weight / (see Eq. (I)), or when the MLP's activation functions are close 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 combat rank collapse.
thus elected ignore the effect the kernel matrices of the MLP
Skip connections When using skip connections as in Eq. (1), the output kernel matrix of a block is given by Vibuck = a² Sshuntuu + 1. EEesidual- 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 a² + 82 = 1. When using normalised skip connections with U-SPA, both Easoneut and Eesidual will have the form E(p) = (1 [1-p)[++pp11 for some p. and thus so will E Moreover, we will have Phlock
pahomeett + B2 Presiduul, which is less than Presidual- This means can easily adjust U-SPA to be compatible with normalised skip connections by replacing L, in the formula A -,Lي- with the Cholesky factor of [((p a² pt-1)/133 3²). In the case of E-SPA, we note that Ebtock won't be of the form (E(r)),, exp(- - jl) for some y, even when both Eshoncul and Eresidual are. To work around this, we use an approximation described in Appendix F, which seeks to make the combined effect the kernel matrix of the normalised skip and attention block approximately equal to the effect of a skipless attention block.
4 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 respectively.
WikiText-103 baselines
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 proposed SPA

---

Published as a conference paper at ICLR 2023
Exponential SPA
Standard
FN=
WT103 Train loss
4.0-
LN=True
default (100k)
3.0-
2.5-
13-
0.900.7
0:90 Shortcut weight
Shortcut weight a
Figure 3: Transformers with normalised skip connections, trained for 100K steps. E-SPA (left) without nor- malisation matches the training speed of standard transformer. Results denote the mean of 2 seeds.
methods and Value-Skipinit to standard transformers both with and without skips, on a block transformer. We clearly see that removing skip connections from 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.
On the other hand, all three of our approaches train for vanilla deep transformers, with our E-SPA method outperforning U-SPA and Value-Skipinit, However, the default transformer with skips and IN still retains training speed advantage com- pared to our skipless methods, mirroring the situa- tion for CNNs without powerful second-order opti- misers (Martens et al., 2021; Zhang et al., 2022).
7-
skip+LN
WT103 Train loss
e-spa
value skipinit
skipless
skipless
skipless LNK
In Table 1, assess the effect of different acti- vation functions in the MLP blocks, as well as the of LN, in skipless transformers using our pro- posed methods. We see that at depth 36 we achieve good training performance for range of activations: OKS-transformed GeLU, TAT-transformed Leaky ReLU, 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).
2-
20K
40K
60K
BOK
Training step
Figure 2: Comparison of skipless transformers. Vanilla transformers are not trainable without our modifications. All curves average over 3 seeds.
Normalised skip connections In Fig. 3, see that one way to match the training loss of the default transformer, without more itera- tions, by using normalised skip connections. While this is perhaps unsurprising, we observe that our E-SPA method (left) matches stan- dard training both with and without normalisa- tion, whereas standard Transformer with nor- malised skip connections (right) requires nor- malisation in order match the training speed of the default Pre-LN.
Table 1: WT103 train loss of skipless transformers for different activations, with without LN. Mean and standard deviation are computed across random seeds.
Attention
Activation
LN
w/o LN
E-SPA
DKS-GeLU
2.99
2.79
TAT-LReLU
3.00
2.894.02
GeLU
2.86
2.93..00
Sigmoid
7.42
7.43..01
U-SPA
DKS-GeLU
3.17 1.17..03
2.99
TAT-LReLU
3.24
3.04
C4 baseline 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 fur- ther compare our methods to standard trans- formers on larger dataset, C4 (Raffel al., 2019), where overfitting isn't an issue. Importantly, we see similar trends across validation (Fig. 4a), training (Fig. 13) and downstream task (Table 6) performance, and so the benefits of methods do extend beyond training. Due the memory
GeLU
3.10
3.25
V-SkipInit
DKS GeLU
3.20
3.26...2
TAT-LReLU
3.23=.02
3.32+.01
GeLU
3.18
3.17+.02
Arrgued by Nakkiran al. (2021), the validation curves the "training speed" of the online setting.
8

---

Published as a conference paper at ICLR 2023
4.2-
Skip+LN
50k steps
e-spa
4.0-
100k steps
C4 Eval loss
u-spa
Ca Eval loss
200k steps
value skipinit
3.0
300k steps
400k steps
3.6
default (50k)
OK
10K
20K
30K
40K
S0K
0K
50K
100K
150K 200K
250K 300k 350K 400K
Training step
Training step
(a) Validation loss of vamilla transformers.
(b) Longer training.
Figure 4: Results on C4. (a): E-SPA again outperforms our other approaches. default Pre-LN after 5x more iterations.
(b): Vanilla E-SPA matches
overhead of longer sequences, we use a 32-block transformer. One can notice that E-SPA performs the best among skipless transformers all settings: training, validation and downstream tasks.
Moreover, in Table 2 we find that E-SPA with nor- malised 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 remove LN, including Stable ResNet (Hayou et al., 2021; Noci et al., 2022) and Skiplnit (De & Smith, 2020), both with & without LN.
Table Validation perplexity for different skips on C4 with without I.N after 50K iterations.
Attention
Skip
LN
w/o LN
Default
Pre-LN
24.7
Normalised
24.5
25.1
Stable (a=1) 1)
24.4
25.3
Skiplnit
25.9
35.3
U-SPA E-SPA
Normalised
27.6
28.0
Normalised
24.0
24.7
Longer training matches default transformers In Figs. 2 and 4a, we observed that while our methods are able to train deep skipless transformers (a result which unprecedented in the litera- ture). there is significant gap in training speed compared to standard Pre-LN transformer (with skips). Martens al. (2021) and Zhang al. (2022) observed a similar training speed gap for skip- less 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 al., 2018; Anil al., 2020). As second order optimisers for transformers are not well established, we instead demonstrate that the training loss gap be closed by simply training for longer in Fig. 4b. We observe that 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 provided in Fig. 14.
Depth Scaling Finally, as unmodified networks suffer from worse signal propagation properties at larger depths, it is nat- ural ask how our modified vanilla transformers perform as depth increases. In Table 3, we compare the training perfor- mance of different depths (36, 72, and 108) at range of train- ing step budgets on WikiText-103. We find that whilst the shal- lower depth 36 network trains fastest initially over 100K steps, 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 close the gap in performance to its depth 36 counterpart with longer training, going from 0.3 100K steps to just 0.01 at 1000K steps.
Table WT-103 training loss depths. Deeper vanilla E-SPA trans- formers improve with more training.
Training steps
Depth
100K
400K
1000K
3.12
2.66
2.20
3.25
2.67
2.11
108
3.42
2.78
2.21
CONCLUSION
We have shown for the first time that it is possible 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 transformer to enable faithful signal propagation even 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 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.
match the training speed of standard ones. We hope that our work may

---

Published as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
Pseudocode for our main approach, E-SPA, can be found Alg. 1, using the notation and setup provided in Section 2. All experimental details can be found in Appendix H, including general and experiment-speccffc 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 Ox WaSP CDT programme (EP//011710/1)..
REFERENCES
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning, arXiv preprint arXhv:2002.09018. 2020.
Devansh Arpit, Victor Campos, and Yoshua Bengio. How to initialize your network? robust initial- ization for weightnorm & resnets. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32.
Curran Associates, Inc., 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arX1v:1607.06450, 2016.
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence large depth. In Uncertainty in Artificial Intelligence, pp. 1352-1366 1361. PMLR, 2021.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. arXiv preprint arviv: 1809. 10853, 2018
David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams. The shattered gradients problem: If resnets are the answer, then what is the question? In Interna- tional Conference on Machine Learning, pp. 342-350. PMLR, 2017.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical com- monsense natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 1904.1090040000 2019.
Christopher Clark, Kenton Lee, Ming-Was Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq:
Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 2924-2936, 2019.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint arXiv: 19111991..00 02860, 2019.
Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding of neural networks: The power of initialization and dual view on expressivity. D. Lee, M. Sugiyama, U. Luxburg, Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
10

---

Published as a conference paper at ICLR 2023
Jared Davis, Albert Gu, Krzysztof Choromanski, Tri Dao, Christopher Re, Chelsea Finn, and Percy Liang. Catformer: Designing stable transformers via sensitivity analysis. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 2489-2499. PMLR, 18-24 Jul 2021. URL https: http://oocee/eeeeee// mlz press/ wtwppppsssss..peep.. wwwp://vis/v1lleeeeaaaaaa htmll.
Soham De and Sam Smith. Batch normalization biases residual blocks towards the identity function in deep networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 19964-19975. Curran As- sociates, Inc., 2020.
Yihe Dong, Jean-Bappaeeeeee---e Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning. pp. 2793-2803. PMLR, 2021.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In International Conference on Machine Learning, pp. 1842-1851 1850. PMLR, 2018.
Boris Hanin and David Rolnick. How to start training: The effect of initialization and architecture. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS'18 18, pp. 569-579, Red Hook, NY, USA, 2018. Curran Associates Inc.
Soufiane Hayou, Arnaud Doucet, and Judith Rousseau. On the impact of the activation function on deep neural networks training. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Pro- ceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pp. 2672-2680. PMLR, 09-15 Jun 2019.
Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis. Arnaud Doucet, and Judith Rousseau. Stable resnet. In International Conference on Artificial Intelligence and Statistics, pp- 1324-1332. PMLR, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision, pp. 630-645. Springer, 2016.
Dan Hendryeks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint 16Xiv:1606.08411 2016.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, al. Train- ing compute-optimal large language models. arXiv preprint 2rXiv:2203.15553 2022.
Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nngp and ntk for deep attention networks. In International Conference on Machine Learning, pp. 4376- 4386. PMLR, 2020.
Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Ian Simon, Curtis Hawthorne, Noam Shazeer, Andrew M. Dai, Matthew D. Hoffman, Monica Dinculescu, and Douglas Eck. Music transformer. In International Conference on Learning Representations, 2019.
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimiza- tion through better initialization. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learn- ing Research, pp. 4475-4483, PMLR, 13-18 Jul 2020. URL https: tttttttss/ppcceeeiinoe pttppppiccceg...ee////ppp mlr.. press/ v1 1 9/ hww.aggnn000....m html.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 69800, 2014.
Taku Kudo and John Richardson. SentencePiece: A simple and language independent sub- word tokenizer and detokenizer for neural text processing. In Proceedings of the 2018 Con- ference on Empirical Methods Natural Language Processing: System Demonstrations, pp. 66-71, Brussels, Belgium, November 2018. Association for Computational Linguistics. doi: 10.18653/v1/D18-22
11

---

Published as a conference paper at ICLR 2023
Jachoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and
Yasaman Bahri. Deep Neural Networks as Gaussian Processes. In International Conference on Learning Representations, 2018.
Mufan Bill Li, Mihai Nica, and Daniel M Roy. The neural covariance sde: Shaped infinite depth- and-width networks at initialization. arXiv preprint arXiv:2206.02766. 2022.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the diffi- culty of training transformers. Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 5747-5763, 2020.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv: 00388, 2016.
James Martens. On the validity of kernel approximations for orthogonally-inttallizee neural net- works. arXiv preprint (0:::2204.0875 2021.
James Martens and Roger Grosse. Optimizing neural networks with moonecker-factorue approximate curvature. In International conference on machine learning, pp. 2408-24111 PMLR, 2015.
James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszez, Valentin Dalibard, Jascha Sohl-Dickstein, and Samuel Schoenholz. Rapid training deep neural networks without skip connections or normalization layers using deep kernel shaping. urXiv preprint 10X00:21 2021.
Alexander G de G Matthews, Mark Rowland, Jiri Hron, Richard Turner, and Zoubin Ghahramani. Gaussian Process Behaviour in Wide Deep Neural Networks. In International Conference Learning Representations, volume 4, 2018.
Elizabeth S. Meckes. The Random Matrix Theory of the Classical Compact Groups. Cambridge Tracts in Mathematics. Cambridge University Press, 2019, doi:
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2017.
Preetum Nakkiran, Behnam Neyshabur, and Hanie Sedghi. The deep bootstrap framework: Good online learners are good offline generalizers. In International Conference on Learning Represen- tations, 2021.
Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. arXiv preprint urXiv:2206.03/222 2022.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pp. 1310-1318. PMLR, 2013.
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Expo- nential expressivity in deep neural networks through transient chaos. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022.
Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint
2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv e-prints, 2019.
12

---

Published as a conference paper at ICLR 2023
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver- sarial winograd schema challenge at scale. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 8732-8740, 2020.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan Le Bras, and Yejin Choi. Social iqa: Common- sense reasoning about social interactions. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 4463-4473, 2019.
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations, 2017.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representa- tions. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Pupers), pp- 464-468, 2018.
Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, al. Using deep- speed and megatron to train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint ar arXiv:2201.11990 2022.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting, Journal of Machine Learning Research, 15(56))1929-199) 1958, 2014.
Jianlin Su. Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.00 09864, 2021.
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 32 42, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. Deepnet: Scaling transformers 1,000 layers. arXiv preprint arXiv:2203.00555 2022.
Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F Wong, and Lidia S Chao. Learning deep transformer models for machine translation. arXiv preprint 2019.
017877,
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dicksteri Samuel Schoenholz, and Jeffrey Pennington. Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla con- volutional neural networks. In International Conference on Machine Learning, pp. 5393-5402. PMLR, 2018
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pp. 10524-10533. PMLR, 2020.
Hongfei Xu, Qiuhui Liu, Josef van Genabith, Deyi Xiong, and Jingyi Zhang. Lipschitz constrained parameter initialization for deep transformers. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 397-402, 2020.
Greg Yang. Wide feedforward recurrent neural networks of any architecture are gaussian pro- cesses. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee whye Teh, Alvaro Sanchez-Gonzalez, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via de- noising for molecular property prediction. arXiv preprint 2rXiv:2206.00/33 2022.
13

---

Published as a conference paper at ICLR 2023
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Compulational Linguisties, pp. 4791-4800, 2019.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Infor- mation Processing Systems, 32, 2019.
Biao Zhang, Ivan Titov, and Rico Sennrich. Improving deep transformer with depth-scaled initial- ization and merged attention. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 898-909, 2019.
Guodong Zhang, Aleksandar Botev, and James Martens. Deep learning without shortcuts: Shaping the kernel with tailored rectifiers. In International Conference on Learning Representations, 2022.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. In International Conference on Learning Representations, 2018.
14

---

Published as a conference paper at ICLR 2023
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 sec- ond, it 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, Skiplitt does not modify the offtaa--attention computation and hence is al- ready compatible with any form of attention. For our SPA methods, it is straightforward to extend to non-causal attention by changing L, in Eqs. (8) and (9) from being the Cholesky decomposition of Σ, to being the (symmetric) matrix square root of Σn In this case, for U-SPA it is possible to analytically calculate that A in Eq. (8) will be element-wise non-negative if /t 2 pr-1 (exactly like the Cholesky case in Theorem 1). This easy to see because the matrix square-root, inverses and products of uniform kernel matrices are all still uniform of the form (p) = (1 p)1r+pp11 (up to positive rescaling), so that A will be too, and one simply needs to track p and verify that is pos-
itive. For E-SPA. we have verified empirically that the resulting A, L,L,-1 will be non-negative n I R-1, just like the Cholesky case in Theorem 2.
B MODIFICATIONS TO SPA METHODS FOR REPEATED TOKENS
For simplicity, we assumed that our input sequences had repeated tokens when presenting SPA in Section 3. This meant that we could take the input kernel matrix Eg to be the identity, with zero off-diagonals, which was convenient for our construction of SPA. The effect of repeated tokens, e.g. the word 'cat" occurs multiple times the same sentence, is that our input kernel matrices Eo will have non-zero aff-diagonal entries, corresponding to entries where a token is repeated. This will impact the kernel matrices Σ, at deeper layers with SPA, particularly the diagonal values of Σ, 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 will often be working with sequences where fraction of the tokens are repeated.
We stress that the general principle that all our methods (Value-Skiphng and SPA methods) follow is independent of the input kernel matrix (i.c. independent of duplicate input tokens): we seek to prevent the product of attention matrices from deviating away from the identity matrix and degen- erating to rank-1 matrix. From Eq. (6), we see that the product of attention-matrices is rank-1 then regardless of the input kernel we will have a rank-1 1 output kernel i.e. rank collapse (Dong 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 long as the input kernel is non-degenerate and has full rank (regardless of duplicate tokens) so too will be the output kernel.
Recall that with attention matrices (Ai),, input kernel matrix across T locations Eo RTXT gets mapped, at depth I, to
Σ 11, Co II
(11)
where II, = A,A,1 A₁ is the product of attention matrices.
In SPA, we parameterise A LL for (Li), corresponding to the Cholesky factors of some family of kernel matrices (Ei),. with either uniform or exponentially decaying off diagonals:
1. U-SPA: [1(p)== p1111 for 0 <
2. E-SPA: (Ei(1))),, exp(-nli- jI)) for 7 2 0.
It is important that A, = L is constructed from two Cholesky matrices belonging to the same family, because Theorems and 2 show in that case we will satisfy our negaativvy constraint on A₁ (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, E. would be an element of our family of kernel matrices, so that we can set Lo to be the Cholesky factor of Eg, satisfying I Ir-
15

---

Published as a conference paper at ICLR 2023
Repeated prob=0.0
Repeated prob=0.00 0.01
Repeated prob=0.05
Repeated prob=0.
22
Contin Similarin
Figure 5: Input token kernel matrices for different proportion of repeated tokens.
corrected
Diagonal entries o
Σ1
uncorrected
4-
3-
2-
25
Layer
Figure 6: Diagonal entries of E₁ for single sequence of length T = 100 across blocks for E-SPA in the presence of 0.05 shared tokens, with and without modifications. We see that without modifications and simply assuming Eg by default (green) the average diagonal diverges deeper blocks, when 7 smaller and off dif-diago 2, are larger. On the other hand, our modifications (purple) keep the average diagonal around with controlled standard deviations. Here, the last kernel matrix (at depth 36) has decay rate 7y. = 0.02 and the standard deviation is computed over the 100 locations.
Typically, En = XXT /o for input token embeddings X E Txxdo which are independently initialised for different tokens identites. Thus, in the wide limit d oc, (up to rescaling) Eg will have diagonals equal to 1, and off-diagonals equal to if there is a repeated token and 0 else. We plot examples of such kernel matrices in Fig. for different, fractions of repeated tokens r, under the assumption that repeated tokens occur independently of location.
Clearly, when 0, have Mg I is member of both the uniform and exponential families of kernel matrices, corresponding to uniform off-diagonals of po = or exponentially decaying off-diagonals with rate To = co. In this case we can set L. I too.
On the other hand, for r quence's Eo. given that different sequences will have repeated tokens in different locations. More- over, the naive approach of ignoring the repeated tokens and treating Ee I leads to increasing diagonal values of Σ, (i.e. activation norms) at large depth without corrections, as shown in Fig. 6. 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 Σj are constant across blocks and locations, usually set to 1.
is in general difficult say much more about an individual
To circumvent this, we will derive our modifications by considering the average input kernel matrix Eo (averaged over different sequences), under the assumption that repeated tokens occur indepen- dent of location:
Σn = -1 (1- r11
(12)
We find that 0.008 for sentencepiece tokenisations like in WikiText- 103 and C4, but z 0.05 for character-level prediction like in En Wiki-8.
16

---

Published as a conference paper at ICLR 2023
By linearity of Eq. (11) in Ey (and because controlling A, be input-independent at initial-
isation in Section 3), also follows that average depth kernel matrix Σt satisfies:
(X, x) - 11, Eo 11
(13)
modify our SPA approaches to control the average kernel matrix Eg instead.
For U-SPA, the situation is more straightforward as En a uniform kernel matrix with - r, so suffices to simply let Po r instead of po O.
For E-SPA, we are unable to view Eg as having exponentially decaying off- difg-iggssass, and so we keep "oo = oo which translates to Lo = I and II, = Lt Instead, to help understand the effect of repeated tokens (to motivate our modifications), we first expand on Eq. (13) to simplify little:
I1. En IJ =4 E. L
= -1-r) +rL,11'u
-Ei(n) + I)L
=I(n) +L
(14)
where Or = 11T I7 E RXT is 0 on the diagonal and off the diagonal.
Note that all terms in Eq. (14) are casily computable (given that (Ei(1)) i.j exp(->yli-) j()) and Lj has an analytic form provided in Lemma 1). So we can use Eq. (14) compute the expected
diagonal (E)),,: for each location (where the expectation is taken over different sequences), which we denote by a diagonal matrix Dj:
D. = D (L: Eo L,) ERTT
Thus, we propose to replace A -اL, with A = DT L; D setting Do by default. This means that our product II, = IT'=1 A; = D $ $L, and Eq. (13) updated to:
EX = D, L Eo L. D
(15)
which has diagonals controlled to 1.
Though our modifications for repeated tokens only consider averages different sequences, find that for individual sequences they still lead to well behaved diagonal values of Σ,, as shown in Fig. 6. Moreover, we find that their effect off-diagonals is still favourable for individual sequences, as shown in Fig. 1.
C SETTING (pt): AND (71)
In this section we describe how we set the uniform off-diagonals (pt): and the exponential decay rates (r)) in SPA, at different depths.
Recall that Theorems and show that we are free to choose (p)) and (r))²n such that (p)) increase with depth and (nn)f decrease with depth. Moreover, Po = 0 (or the shared token fraction r. as per Appendix and 7o oo the input layer, whilst we have found pr 0.8 and 72 0.005 to be good default values for the last block.
For U-SPA, in terms of setting how (Pi); vary with depth, tried different polynomial rates of increase with depth from pa to Pr. but did not observe a noticeable difference in performance across different rates chose to increase from Po to pr. linearly in depth.
For E-SPA, we choose (n)) so that the diagonal elements of the attention matrices A are constant across blocks, akin to using constant shortcut weight a over different blocks. From Theorem 3, we have that the diagonal entries of A, satisfy:
(Ai), = a(v)
Vi 1
a(n-1))
17

---

Published as a conference paper at ICLR 2023
where a(7) /1
cxp(-2y) with inverse 7(a) - glog(1 a²).
This means that for set of positive decreasing (7)),, there exist corresponding set of decreasing (at) with values between O and
Because = oc, we have ag = 1, and likewise for given we can compute az.
Thus to have constant diagonal values of A over different blocks we choose to set
1/L log(1 -00g aL 21/2)
ar.
for S L, and thus
1
21 = ((i))
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 al., 2021). We leave further study of choosing how to set (Pt)z and (%) to future work.
D E-SPA ALGORITHM
In Alg. 1 we present pseudocode to construct a trainable E-SPA masked attention layer.
Algorithm 1: Modified E-SPA masked attention layer.
Input:
Input sequence representation X E RTxx Output: Updated sequence representation Attn(X) E Rx
Hyperparameters: Input "/n & output Yout exponential decay rates, Tin 2 Your Number of heads h with head dimension d d/h. Causal mask M E TxT
M.j = 1fi 2 j}. Large constant I, e.g. 1000, to enforce causal mask. Trainable parameters: W W n WY E Rddd, for S n S h. W E Raxd
Compute Lin and Low to be the Cholesky factors of (E.) }.j = exp( Tim j]) and
(22g )Jj = exp( Tour (1--]] j|) respectively.
Set A Lo (using analytic form Theorem 3) decompose A DP where D diagonal & P has row sums 1. Denote B = IPg(
for head E (h do Initialise WK WV i.i.d N(O, 4) (alternatively orthogonally). Initialise W 0. Set Q. XX) XWQ, (X) XWK and (X) XWV.
Compute P (X) softmax M ( va. Qn (X),, XX) B) T(1 111-My M)
Set Attn, (X) DP,, (X)
Initialise wo iidd NN 4)) (alternatively an orthogonal matrix)
return Concat (Attn; (X)...., Attn, (X) Wo
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, VaT ((X))(()) vith / X 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 w to zero, and initialise WK usual (Gaussian fan-in 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 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 S XW
18

---

Published as a conference paper at ICLR 2023
For RoPE, the (i.j) query-key key dot product, Say, is modified from
(w (w*x)) X,)
to
(R, (R, (R, w
(16)
where R are some location-dependen rotation matrices defined in Su et (2021), and X; denotes
row of the incoming representation X RT'x to the attention block. Clearly, Eq. (16) is 0 for all j if WQ is zero.
For Relative positional encoding, we take the scheme from (Dai et al., 2019). In that case, the (i.j) query-key dot product, St,j. is modified from
(w2x))
to
(WQX,) (W*)())*)))))))))))) + (W*X,) (W.. + (W XW) e T & (17)
where R Rd fixed, and W Rd4xd u gd are trainable. Thus, can achieve zero query-key dot products initialisation we initialise u, v 0. in addition W = 0.
F USING NORMALISED SKIP CONNECTIONS WITH E-SPA
In this section, we describe how combine our E-SPA method with normalised skip connections in our attention blocks:
Z aX + V1
(18)
where a = 0 is the skipless setting that our methods are originally designed for.
The general gist that we will look 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 match the setting where skips, i.e. a = O.
without
To combine E-SPA with normalised skip connections, we consider how the cosine-similarity be- tween two locations Ti, T2 T affected through the normalised skip connection. Suppose we have an input kernel matrix:
Xx =
where (E),, 1, Vi and d is the width. Then after the residual attention block. Eq. (18), we now have:
az) 1 zz =0² xx (1-
AXW W X +avi oz XW A AXW
(19)
(20)
Either we have V to be an orthogonal matrix sampled uniformly at random from the Haar measure, or WV Lidd (0, }). In both cases we have 1 / WVTXT XTAT A going towards } AXX AT, and the eross terms Eq. (20) converging to 0 for large d.
Thus, we can consider the large d approximation:
7 zZ
=a 7 XXT + (1 a² ) A
(21)
=0-x + (1 111-22)AE AAAA
(22)
19

---

Published as a conference paper at ICLR 2023
Then, if we look at the inner product between the Ti and T2 locations for locations Ti # T2 such that Ti, T2 > 1, we have:
1 d ZZ TT,T2 = + (1- TT,,
=a + +11-a²) Σ
AT E
ET + +(1 (O?AA,,, Dr,,T ATI,T a1)
E
inTIviATs
=Ta² + (1 - -aa))a²] [>) + a1-
Σ
ji AT
(TTUU
=Ia² + (1- aa22 LT.T + a2-6
(23)
=VLT,,T, + (1--0 a1-a y0
where AT.T, ATaTT is the constant diagonal of the attention matrices in E-SPA, c.f. Theorem 3, and a + (1-a2)
Moreover, we have defined = EITT,UUHTT AT, A as it is term that is not possible to control with only knowledge of ΣT; and will vary from sequence to sequence, and hence we argue be discarded from signal propagation perspective. For example,
could consider sequence where the kernel matrix Σ has all off-diagonals equal to 0 apart from Σr, and T1 could be sufficiently distant from T22 such that there is for which AT, and AT., are large at the same time (which can be seen using the analytic form of A given in Theorem 3).
Thus, from Eq. (23), we see that after an residual attention block, an input cosine similarity of En,22, between locations Ti, T3 diluted by a factor of
v == + - 2²) 11
with shortcut weight a and attention diagonal probability X. Our approximation then seeks to pre- serve this factor when a to match the skipless case.
Now, in the skipless case a = 0 described in Section 3, at block we suppose that the incoming kernel matrix M has exponentially decaying off-diagonals with rate 7M-1, and that we construct the attention matrix so that the output kernel matrix has exponentially decaying off diagonals with rate m,. From Theorem 3, we see that this gives diagonal entries of A to be:
a(n) a(7-1)
Ao
where a(7) /1 exp(-27).
Thus, to preserve y to match the case for a 0, if we have shortcut weight a > 0 we need the attention matrix diagonal prohability Aa satisfy:
X a
Aa
(24)
1- a²
In turn, when have shorteut weight a, this means that need choose our outgoing decay rate at block 1, N... such that:
a(n,a) = da)(n-1))
Inverting the definition of a(7), we see we need to set n?,:
n. = - 1111 (22
(25)
= log 1 x3- a²
a (1
-expl- (-p(-29-11))
(26)
20

---

Published as a conference paper at ICLR 2023
Table 4: Validation perplexity equivalent of Table
Attention
Activation
With LN
Without LN
E-SPA
DKS-GeLU
21.0
20.0
TAT-LReLU
21.5=.1
19.7
GeLU
20.4
23.1
U-SPA
DKS-GeLU
24.1
23.2
TAT-LReLU
26.0
21.8
GeLU
24.7=.2
29.2
V-Skiplnit
DKS-GeLU
24.7=.1
27.5+.1
TAT-LReLU
25.7=.2
27.4+.1
GeLU
25.5_,3
26.14.7
To summarise, if we have a sequence (7t): of decreasing exponential decay rates, our proposed approximation when using shortcut weights a at block is, using the notation of Alg. 1, to set Tin 7 normal, and Your = Voo from Eq. (25) in order preserve the signal propagation from the combined residual attention block Eq. (23). We see that this approximation reduces to the standard skipless setting when O, sanity check.
G ADDITIONAL RESULTS
In this section we present additional results and ablations that were not included Section 4.
Normalised kernel matrix evolution for attention-only transformers with skips and normalisa- tion 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, 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 position- ing of the LN in the residual branch effectively downweights the residual branch later blocks (De & Smith, 2020). Likewise, Pre-LN with skip weight a = 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 Appendix H.2) and we observe both Pre-LN kernel matrices also have a recency bias, like our main method, E-SPA (third row).
Validation perplexity across different skipless methods In Table I 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-tranformmd GeLU activation without LN to perform best. In Table we present the corresponding results but for validation perplexity. Again, that E-SPA is the best performing of our attention modifications, but in this case TAT with Leaky ReLU and no LN matches outperforms DKS with GeLU.
Sensitivity to query-key initialisation scale Recall in Section 3.3 that for attention layers using SPA seek to initialise weights such that the query- dot product, XX (XW*) (TT iss zero
or small at initialisation. In our main experiments, we achieve this by initialising Wº = 0. and letting W 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 WK' Wº are both orthogonally initialised but with scale (i.e. at each layer, both WK WO initialised
independent uniform
21

---

Published as a conference paper at ICLR 2023
Layer 1
Layer 5
Layer 10
Layer 20
Layer 50
Layer 100
Standard
Skipless Exporeetial
Presu skoweight--
PostUN
Figure Equivalent of Fig. with four additional configurations Transformers: 1) Skipless+LN; 2) Pre-LN; 3) Pre-LN with normalised skips = 0.98 and 3 = V1-aa* and 4) Post-LN, All rows except SPA methods (i.e all rows except second and third) pre-softmax bias matrices from the ALiBi (Press et al., 2022) positional encoding compute attention matrices, assuming query-ket dot product.
22

---

Published as a conference paper at ICLR 2023
orthgonal matrices multiplied by d). We 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 zero query-key dot product at initialisation) less precise.
3.10
zero query weight init (100k)
3.05
WT103 Train loss
3.00
2.95
2.90
2.85
2.80
2.75
2.70 -1l
0.0001
0.001
0.01
0.1
Query & Key init scale a
Figure 8: The sensitivity of training performance to the initialisation scale of WQ and w using vanilla E-SPA. Mean and errrr bars are over seeds.
Ablation orthogonal initialisation Recall in Section 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 ing 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).
Dass
DKS-GeLL,orth=-True
4.0
TAT-LReLL,ortt=Faale
WT103 Train loss
TAT-LReLL,orthhTTrue
3.8
GeLUlorth==alse
GeLU,orth=
3.6
3.4
3.2
3.0
20K
40K
60K
100K
Training step
Figure 9: Ablation over using orthogonal parameter initialisations, compared to standard Gaussian fan-in, for main method E-SPA vanilla transformers over range of activation functions. We see that orthogonal intialisation leads small improvement in training speed. Curves denote mean over seeds
23

---

Published as a conference paper at ICLR 2023
Ablation positional encodings As noted in Appendix E. methods 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, that all methods are improved when combined with RoPE, however the improvement is most mild E-SPA, which as discussed has in-built recency bias akin to positional encoder. Moreover, E- SPA without additional positional encoding still outperforms all other approaches, including U-SPA (which on its own has notion of position) with RoPE.
our skipless methods. We see
Table 5: Validation perplexity equivalent of Fig. 10
Attention
PosEnc
With LN
Without LN
E-SPA
None
22.4
22.1
RoPE
21.0+.1
20.0
U-SPA
None
34.7+1.2
28.1+.4
RoPE
24.1+.3
23.2+.3
V-SkipInit
None
29.34.1
34.4
RoPE
24.7+.1
27.5+.1
6.0
e-spa.posenc=mone
e-spa.posenccrope
WT103 Train loss
5.5
u-spa.posenc=none
5.0
u-spa.posenc=rope value skipini,,posenc=none
4.5
value skip
4.0
3.5
3.0
OK
20K
40K
60K
8ÓK
100K
Training step
Figure Ablation over positional encoding using vanilla transformer with E-SPA, averaged over 3 random seeds.
Stable residual connection (Hayou et al., 2021; Noci et al., 2022) WikiText-103 In Fig. 11, we provide an equivalent plot to Fig. 3 using a Stable RexNer (Hayou et al., 2021) rescaling of
the shortcut weights (a = 1, 0 + )). In this case, the shortcut weight is always 1 and the
residual weight 3 is uniform across blocks and scales as 0 (/)) in depth L. Hayou et al. (2021)
showed that such 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 B, the stable residual weighting with normalisation matches the training speed of the default transformer (which is unsur- prising given that once / = 1 it is exactly default Pre-LN), However, there is small but consistent gap without normalisation (with optimal 3 0.1). Here, L 36, so 1 1 s 0.17, or altema- VL tively we count the 72 nonlinear layers (one self-attention and element- wise nonlincarity for each transformer block), we have / z 0.12.
24

---

Published as a conference paper at ICLR 2023
Stable Resnet, a =
2.8
In=none
In=rmsnorm default (100k)
2.6
WT103 Train loss
2.4
2.2
2.0
1.8
1.6
0.01
0.02
0.05
0.2
0.5
Residual weight B
Figure 11: Performance of a Stable (Hayou et al., 2021; Noci al., 2022) residual connection weighting on WikiText- 103 for different residual weights 3. Mean over 2 seeds.
Different shortcut weights for MIP and self-attention blocks Recall that in a transformer block Eq. (I), there are two distinct skip connections: one for self-attention block and one for the MLP block. Moreover, we observed training speed gap when we remove both skip connections in Fig. 2. This leads us ask if is possible that only one of the skips, MLP attention, 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 shorteut weight for skipless MLP blocks (right). For all skip connections (for both MI. and self-attention blocks), we use normalised skip connection (8=vl VV= a²). 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 slightly lower weightings (a 0.9 or 0.97) compared to attention shortcuts (a 0.98 or 0.99), whereas our experiments in Figs. and 11 use joint weighting for both.
No attn shortcut
No dense shortcut
2.3
WT103 Train loss
2.7
2.6-
2.4
0.995
0.99
0.98
0.9
0.7
0.99
0.98 0.97
0.9
0.7
Dense shortcut weight
Attn shortcut weight
Figure 12: Effeet of using different shortcut weights for the attention and ML. blocks, trained on WikiText- 103 for 100K steps. Mean and standard deviation over 2 seeds.
25

---

Published as a conference paper at ICLR 2023
Training performance C4 In Fig. 13 compare the training performance of various vanilla transfomers to the default Pre-LN transformer on C4. This is akin to Fig. 4a, which compared validation performance.
7.0
Skip+LN
6.5
e-spa
C4 Train loss
6.0
u-spa
5.5
value skipinit
5.0
4.5
4.0
3.5
3.0
OK
10K
20K
30K
40K
50K
Training step
Figure 13: Training performance C4, equivalent Fig. 4a
Downstream tasks Typically, transformers are pre-trained on large corpus of data before evalu- ation on 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 al., 2019), HellaSwag (Zellers al., 2019), Winogrande (Sakaguchi et al., 2020), PIQA (Bisk et al., 2020), and SIQA (Sap et al., 2019). These datasets commonly used to evaluate large pre-trained transformers (Brown al., 2020; Rae 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 4 out of 5 downstream tasks. However, the default transformer outperforms skipless transformers on all tasks with same amount training.
With around 5 times longer training (200K and 300K steps), E-SPA achieves similar per- formance on downstream tasks to the standard transformer, outperforming on 2 out of tasks (Winogrande and PIQA).
Table 6: Downstream evaluation of various models (default Pre-LN and also our skipless models) from Figs. 4a and 4b, (pre-)trained C4, We evaluate zero-shot on common sense tasks: BoolQ, HellaSwag, Winogrande, SIQA. In brackets shows the number of steps each model trained for. Reported values percentage accuracies: higher better.
Model
BoolQ
HellaSwag
Winogrande
PIQA
SIQA
Default (50K)
60.9
29.1
52.5
62.2
41.9
V-Skiplnit (50K)
38.7
25.8
51.6
56.0
39.0
U-SPA (50K)
39.2
25.8
50.9
56.7
40.6
E-SPA (50K)
59.5
26.2
51.3
59.7
39.8
E-SPA (200K)
60.2
27.7
52.8
63.1
40.9
E-SPA (300K)
56.8
28.4
53.0
62.5
41.1
Longer training on WikiText-103 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 mirros findings on C4 in Fig. 4b.
26

---

Published as a conference paper at ICLR 2023
4.0
100k steps
WT103 Train loss
3.5
200k steps
400k steps
3.0
600k steps
default (100k)
2.5
2.0
1.5
0K
100K
200K
300K
400K
500K
600K
Training step
Figure 14: Longer training WikiText- 103 with E-SPA vanilla transformer. It matches the performance of standard transformer after 4.5x more iterations of training.
Kernel matrix evolution during training Fig. 15 shows the evolution of the empirically- computed normalized kernel matrix during training of (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. I, despite the fact that Fig. I is computed for an attention-only network the infinite width limit. We also note that while significant changes to the kernel matrix occur during training, it retains the property being larger close the diagonal.
H IMPLEMENTATION DETAILS
We first describe all additional general implementation details for experiments, before going into details relevant for individual results.
H.1 GENERAL IMPLEMENTATION DETAILS
Datasets We present experiments WikiText-103 (Merity 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 2048 for both.
Embedding By default the embedding layer E E IR Vix fan-out initialised at initialisation, so that E, ltd. (O, 2 This means that each embedding entry has scale do at initialisation, whereas Gram matrices (which dimension normalised) expect each entry have scale 1 at initiali- sation. To get by this, in our proposed vanilla methods we rescale the embedding matrix by a fixed factor of dg. so that E < dd which is then treated as lookup table. At initialisation, this has the same effect using RMSNorm layer weights are shared with the embedding weights. We not use the embedding layer positional encoding introduced by Vaswani al. (2017).
the embedding layer. In all our experiments, the unembedding
Model our experiments apart from Table 3, model width d = 1024 across blocks. We use 8 head multi-head attention, so that d = 128. The MLP block consists of 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 32 block transformer due memory constraints. Any normalisation layer we consider is RMSNorm (Zhang & Sennrich, 2019), which is simpler than Layer Normalisation (Ba al., 2016) and is commonly used in transformers (Rae
27

---

Published as a conference paper at ICLR 2023
Block 0
Block
Block 10
Block 15
Block 25
Block 35
Stop 0
DOS dlags
Cosine similsrity
0000 dlogs
0002 days
000S doss
Figure 15: The normalized kernel matrix for various transformer blocks at three different stages of training of vanilla E-SPA transformer on WikiText- 103. Step 0 corresponds to the untrained network. Note that this is computed with actual finite-width network with randomly sampled (and then trained) parameters, unlike Fig. which corresponds the infinite-width limit (at initialisation..
et al., 2021). By default, our models use RoPE positional encoder (Su et al., 2021) (apart from the ablation in Fig. 10).
Parameter initialisations By default, all weight matrices are fan-in initialised N(O0, Ton-in with The two exceptions for this are: 1) when using orthogonal intiailisation, we use scaled- corrected uniform orthogonal initialisation (Martens et al., 2021) with scale = (which for square matrices is just an orthogonal matrix sampled from the Haar measure Meckes (2019) multiplied by σ), and 2) for the parameter matrix immediately after the activation we set a to take input activation norm ("q-values" diagonals of Gram matrices) of I to output activation norm of 1. In the latter
28

---

Published as a conference paper at ICLR 2023
case, 1 by construction for activations transformed by DKS/TAT Martens et al. (2021); Zhang al. (2022). All bias parameters are initialised to
(Transformed) Activations For DKS (Martens et al., 2021) we set slope parameter ç and for TAT (Zhang et al., 2022) with leaky ReLU, we set 7 = 0.9. Both values were chosen by tikiiText-133. The DKS and TAT transformations
small hyperparameter sweep 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 DKS-transformed GeLU as the nonlinearity in the MLP block by default.
chosen
Loss We perform next-token prediction, with softmax cross entropy loss. More concretely, X R denotes our final layer representation and E RIVixd denotes our embedding layer, then obtain logits E Tx|V| for each location for [V|-wa classification. The loss is softmax cross entropy, obtained by using each location i to predict the identity, in [V], of the token at location i + 1. In our training loss results, we use an exponential moving average with smoothing factor 0.01, to reduce randomness from mini-batching.
Optimiser 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.
Training We use mini-batching of 16 sequences for WikiText- 103 and for C4, due to memory constraints. Unless stated otherwise, we train for 100K steps on WikiText-103 and 50K steps for C4.
Learning rate In all experiments we use linear learning rate warm-up period of 4000 steps which increases from 0 to a maximum value (which is tunable hyperparameter). In the remaining iterations, we use a "one-cycle" cosine learning rate schedule (Loshchilov & Hutter, 2016) which reaches 0 at the final iteration. The maximum learning rate was tuned in {1, 2, 3.5) x 10"" over Z for all experiments. As the embedding layer does not change across the different depths and skip/normalisation settings we considered, we use a different, fixed maximum learning rate for the embedding layer, which was chosen to match the optimal learning rate for the default transformer (2 x 10 for WikiText- 103 and 5 x 10-4 for C4) and not tuned further.
H.2 ADDITIONAL IMPLEMENTATION DETAILS FOR INDIVIDUAL EXPERIMENTS
Attention-only kernel matrix evolution: Figs. 1 and 7 We calculate the kernel matrix Σ evolu- tion directly in T x T kernel matrix-space where T 100. Our input kernel matrix Σ is sampled assuming a fraction r = 0.02 of repeated tokens, with value of if the token is repeated, and else. For all configurations of skiphhormallistion/attttttoo modifications corresponding Fig. 7, we use 8 heads.
row of
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. The 3 possible operations are: 1) attention, 2) skip connection, or 3) LN/RMSNorm operation.
1. Attention Because our SPA methods (second and third rows) are agnostic to the number of heads at initialisation (all attention matrices a self-attention block are the same across heads at initialisation). we apply Eq. (6) directly, so that a single attention block amounts + AEAT for attention matrix A and incoming kernel matrix Σ. Our E-SPA method uses 7L = 0.005 and our U-SPA method uses PL 0.8.
For all other rows, the self-attention operation uses ALiBi (Press et al., 2022), a popular positional encoder which uses head- deadddennddent pre-softmax bias matrices. More specifi- cally, from the default pre-softmax bias matrices given by ALiBi (for 8 heads), we obtain 8 attention matrices { using the softmax operation (which is exact assuming zero query-key dot product 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 E with 8-head ALiB corresponds to (Martens al.,
29

---

Published as a conference paper at ICLR 2023
Table 7: Shortcut and residual weights for Table 2. Weights presented (Aurn, Smna) / (aMI.P fMIP). where Olann denotes the shorteut weight for attention block and OMLP denotes the shorteut weight for the ML. black. Residual weights 3 are defined similarly. denotes that initial value a trainable parameter.
Attention
Skip
w/o
Default
Default
(1.1)/(1,1)
Normalised
0.999'((099,90),,00000000 M(0.99,vI -.999))))
(0.999. 0.999 ((0.99,
/ 0.99
Stable (a==)) I)
(1.0.2)(1.0.15)
(1.0.05)<(.0.))
SkipInit
(1,05)(1,09)
(1,09)(1,0))
U-SPA E-SPA
Normalised
(0.98.v1 V0,99.0
0.99*)
(0.995.v1 0.995 v1 0.995
Normalised
(0.9975, - 0.9975 M(0.995,v
VT 0.995
(0.9975,v 1 0.9975 M(0.99,v 0.99*)
2021):
1 8 ++
EA
8 h=1
2. For skip connection with shorteut weight and residual weight 88 if Egg (E) denotes the output of kernel matrix Σ after an self-attention operation (i.e. from point above), then incoming kernel matrix Σ gets mapped to:
Σ a B Emn
3. For normalisation operation, the incoming kernel matrix E gets mapped to:
Σ t diag(E) 3 Σ diag (Σ)
Hyperparameter tuning of skipless transformers For experiments WikiText- 103 with 100K steps, for our U-SPA transformers we tuned fr. E {0.6, 0.8), and for our E-SPA transformers we tuned 7u E {0.005, 0.2}. For all other settings (longer/deeper training on WikiText-103 or any C4 experiment), we used the default 7, 0.005 and PL 0.8. All hyperparameters throughout our work are tuned based on training loss.
Hyperparameter tuning of transformers with normalised skips All experiments with skip con- nections (i.e. shortcut weight o 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 and small final decay rate V/z then the value of attention matrix diagonal A.. 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 7r. E (0.005, 0.2, 0.4, 0.6) when using 36 block transformer on WikiText-103 and 7/2. (0.005, 0.2, 0.4, 0.7, 1) for Table 2, which used 32 block transformer on C4.
For Table 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 MLP blocks in Fig. 12. We tuned the attention shortcut weight in the range a E {0.98, 0.99, 0.995, 0.9975, 0.999} and the MLP shortcut weight in the range f0.98, 0.99. 0.995). Likewise, both the stable residual weights were tuned 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.
Depth scaling Due to memory constraints, for our deeper networks in Table we use width d = 512 rather than 1024, with 8 heads to give dh 64. All depth scaling runs used DKS -t-rranfforred GeLU with = 5.
I THEORETICAL RESULTS
In this section state and prove theoretical results, including Theorems 1 and 2:
30

---

Published as a conference paper at ICLR 2023
Theorem 1. (Non-negativity for U-SPA) Let Σ=(1 - p)I and E'=(1 +I with respective (positive) fuctors and L' Then if p S p, we have L'L. -1 is elementwise non-negative.
Theorem 2. (Non-negativity for E-SPA) Let matrices (E),, exp( /// j()) and (E),. = xxp(->li--l)) with respective (positive) Cholesky factors and L'. Then i 2 y we have L'L is elementwise non-negative.
We prove Theorem second is more involved.
I.1 PROOF THEOREM 2
We actually prove Theorem 2 as corollary of Theorem 3, which provides the analytic form for L'L-1.
Theorem 3. Let matrices (E),, exp(-7li- j()) and (E')s, exp(->'li- jl)) with respective
(positive) Cholesky factors L and L'. Then i 2 r O, we have A = L'L takes the following form:
if i i=j=1
a(r')
if i=j$1
a(>)
AI.
(exp(-7')= el7')
a(7) elpllexp( exp( exp((>))]expp 7/ - jjj11 1)),
if
(27)
a(22) a(r)
exp(-xx) exp(->)) exp( -i- 1)),
if j
else
where a(7) = vI - e 22),, and likewise a(7') = /1 -1- ei-eexp(22")
Proof of Theorem 2. From Theorem 3, have the analytic form of A. Clearly a(y), a(r') positive, and moreover if 7 S 7, then exp(->') exp(-7) is non-negative.
Finally, because alr') < when 7 S 7, then exp( (07 <17) xx7pexp(-1 exp(->) is non-negative too. a(7)
To prove Theorem 3, we first compute what the analytic form of Cholesky factor for (E)..j exp(-7li - ji)) takes in Lemma I.
Lemma 1. Let (E)i. exp(->lt - -I)) with (positive) Cholesky factors L such that LL Then, we have:
Σ.
exp(-rli i)), Laj
if j=1,121
v - (1--22)xxp -oopp-rl---- - if f 0
exp(-
(28)
else
Proof of Lemma 1. It is clear that Σ positive semi definite, as it is the covariance matrix a stationary Ornstein-Uhlenbecck process, hence Cholesky factor must exist. We now show that is L, Eq. (28).
If we define 1 = min(m, n)), then we have:
LLL mm Lm L
Lm,
= =exp(->(m + 2)) + exp(-27)) Σ exp( (xp(->/mm+0) 21))
1-1
= exp(->(mm + 2)) [1 + exp(-27)) Σ
cxp 2yi)]
1=1
= exp(->(m + - 2)) [1 + exp(-27)) exp(27) (x=x exp(2> - 1))
exp(27)
31

---

Published as a conference paper at ICLR 2023
exp( (mp(-7(m--)) 2)) (1 exp(2>(1 - 1)))
exp( (xp(-7(m--)) + n 21))
exp(-7/m n!)
Em,n
By the uniqueness of (positive) Cholesky factors, the proof is complete.
Proof of Theorem 3. We want to show (AL)R.) = (L')R,1, Vk,1. I. This is clearly true for the top diagonal k = 1= 1.
We now show this for the rest of the first column, when ( 1, 1:
(AL) AA,1 + AL +
A Li,1
a(7')
=(exp(->>)
exp( (7) expp - 7/: (ks 2))
a(7)
a(r')
exp( 7 exp(->(h: - -1))
a(7)
exp( - - 7 (x
-11))=1)ppppp 1))exp(- 7 1)) (29)
Applying the geometric sum to Eq. (29) yields:
a(>')
(AL) =lexp(->')) exp(<>)) exp( - 7 - 22))
a(>)
a(7')
exp( (> 11)
{(>)
+ a(7')
exp( (2((-2))- -epp exp(->)] exp( -1 ((k3))) - -2) -xxe((((k 2)( - -7))
a(>)
1 1 --exp(?' - 77)
=exp( (k -1)
((21)
[exxp( - >/(k: 2) 7) + exp(->(k 1))
((7)
+ +exp( -k'(k - 3) exxp((>>)( -xp(px-ppxx 22/72222022220 7)))
- = 17714 111
-
((71)
-exp( 7/(k:
-k- - 2) -7) + +exp(->(e 1))
+ eexa - - 7) (1 exp((k: 2)(r' +)))
=exp(->/ - (t]
((11)
[eexp( - -/(k (k= - 2) -7) + expp - -1(1 - 2) -7)
a(7)
+ eexp(>>(k 1)) exp(->(k: - 1))]
=exp( =exp(->(kk - 1))
-L/,1
as desired. Likewise, if > 1,1 S k:
k-1
(AL) ki.1 =A + Σ Li
32

---

Published as a conference paper at ICLR 2023
R-1
=a(1')
exp( (k - 1) + Σ
- - i 1)) 1)) 1))
exp( (1 1) + [exp(_>> - -xp exp( exp( - -2 - 1)) 1 - exp((2' - 7) - 11) 1 1-exp(? -
=a(>'
=a(>')
exp( ( () + <ep(~>/)epp( - 7 - 1 1)) (1 - exp((n' - -7) ( 11))
exp( 71: 1) exp( (k: 1)) (1 - exp((>' -): - 11)
=a(>')
exp( 7(k - 1)) + exp( - ->( 1)) - exp( - -7(k - (1))
=a(>")
=a(7')exp( 7'k --1) 1))
33

---

Published as a conference paper at ICLR 2023
1.2 PROOF OF THEOREM 1
Theorem 1. (Non-negativity for U-SPA) Ler E-(1 - pppp+1 and E'= t'It1'' with respective (positive) Cholesky factors L and L' Then if p S pl, we have L'L is elementwise non-negative.
1.2.1 PRELIMINARIES
Before diving into the actual proof of the theorem we will derive several useful properties and notations. First we make slight notational change from the main text of the theorem by replacing I. which depends on with (z) where n represents the size of the matrix, while replaces In addition, since the case of p = p' (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 < I < < 1. In any of the mathematical derivations any temporary expressions, that will be expanded on the following lines. Note that these are never general definitions, they might be used multiple times for different expressions.
will denote with capital English letters (e.g. A, B, C, T)
We will denote vectors and vector functions in bold and scalars and scalar function in standard font. Let 1, be the n-dimensional vector with only ones:
1, =
E ER"
We will denote with In the n-dimensional vector with only r's:
Un = rlp
First, we define the linear map Pr (x) as:
(z) nE +1
Proposition The vector 1, is an eigenvector of C, (x) with an eigenvalue Pn- (I).
Proof. Directly calculating the k-th entry of the product C Cn(r),, gives:
[C)(2)ln] = (C, Eim of + (1 )² = 1 + (n - 1)x
which directly implies that:
1, = (
Corollary 1. The vector 1, an eigenvector of (n with an eigenvalue Pa-1
Corollary 2. (x( Ixn Pa p==1(x) En.
Corollary 3. Cx -
na² Pa 1(m)
Further we define the following useful functions:
nx.
+ nx..
(1 x) x) 1)
d,(x) 1- (z) =1
1 (n 1)x
(n 1)x
(n 1)x
Pn(x)
(1 x) Pa- Pa-1(2)) (2))
(x)
(d(())pa-1(( V(1 -x)pp ())p((r))--1(0)))
Tn(x) v (x)
Pn-1(x)
TT(
(x) Pn-1(x)
Definition 1. The function vn [[0, 1] T[0, Th maps the bounded scalars and y the
vector:
da(y)
(2, = Ln - yn - dz(x)
Cx)(x) I...
34

---

Published as a conference paper at ICLR 2023
1.2.2 LEMMAS
Lemma 2. If0 < x y 1 then all entries of the vector vn y) are non-negative (0 y)
0.
Lemma 3. 1f0 < I < y < 1 then the function gm y)
y(1 2(1-g)pa-r(a) (x(2)ra())
dn+ily) du+t Pa Pa( is
non-negative.
Lemma 4. 0 < x y < 1 then the function fn(I,y)
20(1-1) (x(()),(1)
+
dn+i(y) 5 dxt1(2) P.la)
(dn(t) V (m) 1(m) is negative.
Definition 2. The partial sum the functions fn v) from k to n will be denote by (2 EiERh f((x, y). which from Lemma 4 follow are always negative.
1.2.3 MAIN PROOF
The theorem is proven by induction.
First for n we have that (xr) = (1 [1], which implies that Li(x) Li(y) [1] and the condition is trivially satisfied.
Now assuming that the statement is true for all integers up to will prove that it also holds for n +
Cx)1(x)
Cn(r) x
In 1
Ln(x) o (n))) -T
On Vdn(z)
(x))
Ln(x)
On
Ln(x)
Ln(x)- E (d(
Lx+1(z) (m+z
2 (n)x) -T
V (z)
n
L (z) Ln (z)
In
A = x, (x)
-T Inn (x)
Lan + dz(x)
T Cn())
-lon,
on + dh dn(x) =
[n(x)-1
On
(x+)(x)= =
Tn(2) 2
dd(())
[n(())
On
L
(x))
0
(x+1(x))nnt(())
x Ln(x) -T
V (x)
(m)
(x)
In
x Ln( -T Ln (n(x))
(da(a) (x)
x
1
In
0₂
In
O,,
Cx
0
Pn-1(x)
Ln(y)
On
I (n((a)-)
O,,
Tn+t1 (y) [m+1(2))
3 Ln(y) -T
dd)(y)
(n(x)) x
Vd, (m)
0
y Ln -T (x)
(x)
T
d.
du (x) 12)
(2)
Ln Ln(y) Ln)))
On
(x
(2)
Using the fact that L, (z) is lower triangular and non-negative matrix by the inductive assumption combined with Lemma 2 and the fact that d, (xx > 0 follows that Lu (y) Jon+ (xx) also a lower triangular non-negative matrix.
35

---

Published as a conference paper at ICLR 2023
1.2.4 PROOF OF LEMMA 2
First we will inspect the evolution of (2, y) increase n
d V d2+1 (2)
(2,,(1)))
Ln+2(x) -T T (n++il))
Un++
dn+i(y)
A V dn-1(2) (I) Pa(
2nt1
A
Ln+1(x) T (n+ily) yn++
Ln(x)
Tm(1)
(n(
O,
0%
(.) yn
1
yn+1
(dn(x)
(
(n(x)
(1) Yn
Un Tn(
0%
ny² (1)
(d)
V Vd.
(x)
Tn
T yy yn
Vdn(z)9n-166 (x)
0
1
Vdr(z)
(w(
B
VPn-1 no
v((n
1)y +
ny
y y
y(1 y)
Ln(x)
In ra(w)
Ton (1) (1)))=+ Yn
A =
0%
y(l-vl vl (w)
(x)
(x) Ln (3))
8(1-8) Yr (w)
((1-))
(2)r((()))
-(1-v) v) ddn (2)
y(1 y) (1(x)r 8p
B = [x(x) -1 Ln(y)) -y
dn.(y)
y(1 - v)
(,, + V dn dx)
(z(z) -'xnn
Tn (x) En
y(1 y)
dy(y)) 1
(2,
In +
dn (z) Pn (xx) En
Tn (2) (3)
Un (,
T. (x)r, (v) Fn
V
(1(
(x) 12)) Fn
A
a(1 w)p
(v)
du+
(n+1(w) Ln En
Un (2, = A
(n) d, 1(m) Pm)(
U, (2.8)
(- x(1-v)
(n()))r.)))) V dat1(z)
dxt1(y) (v) 1
(x) V
(x) p., 1(z)
de (v)
En
((1-))p--1(1 (x)
dn+i(y)
(.(
dn)1(a)) (x)
Using the definition of from Lemma 3, Lemma and Definition 2 we have that:
[oo]
= Igo fi.g1l"
= Igo fi - f2.91 92
36

---

Published as a conference paper at ICLR 2023
n-
n-1
n-
Un
90
Σ
f 91
>
fe...... 9k Σ fis...
im2
im
Wo 1
11 ;
9k
11 1..
Or in other words we have that:
[u, 8/]a = 9k v) 1/ v)
Thus proving the lemma reduces proving that:
(x v) - hi (z, v) > 0 vk, n
Expanding on the equation that we need to prove is positive:
[2 9x(1, v) ng (x, = gb
Σ M(
i=k+1
3(1 1-1-1/0 (()
dx+1(y) (2)
(x(
1 dx+1(a)) (x) (t(
12
xy(11 y)
d1+1(y) x
dz(y)
+ V
(1(1))(1)) i=++1
d (3) Pt(ar)
dx(x))
(
01 1())
a sy(1 y) B
Σ
(1)
k=++
dxt1(y) x
di+1(y) x
n
d:(1) X
B
Σ
V dx+1(x) (x(
d++1(z) (z) (x(x)
dx(z) (
i
k+1
di+1 x
di-1 (g) x
d++1 (z) (x)
d:+1 (z) pz(z))
d (y) x
V dn+1 (x) Pm (x)
(T v) = y1 y) Pk 1(x)
ry(1 9) dn+ily)
Σ
V ( Pa(x)
(x)
#++1
Now we turn our attention to the sum in the middle:
1y(1 y) xyy(1 y) Σ
1
B
Σ
(x(
(1))))()))
k=
V(- (1-x)(1-9) 4)
2
T i=k+1 ((()))(())
Using Cacuhy- -Schwartz incquality:
Σ
n
V(1 11-=)(1199) y)
Σ
"
1-x x rT(2)2
Σ
n
1-y ry)3
(1)))(()))
Let's define the partial sum in the brackets as:
(x) = ΣE
x
n
1
x (x)p (z)
pe(x) (=(x)p:-1(2)
t=k
We will prove by induction that:
[xxx) = n n-k+1
Pz-1(x)pm.22
37

---

Published as a conference paper at ICLR 2023
Pr, (=)pp---11 (a) which clearly satisfies the above equation. Assuming this is true for n, we will now show holds for 1:
First for k, have that [2
k
(x)
Σ
1
(z) + pa(e))n41(0)
Pa Pa(c)pnn
n ++1
1
Pk- (x)pm (I) Pr pr(x))n+1(z)
(n)1(z)(n - k + 1) - pz-1
((n + 1)x 1)(n 1 ((k 1)x + +)
+) ()
Pk- (z) pm Pn+ (z)
Pk Pk-1 (2-1
x(n² - nk
k
1) + k+1+1
Pk (z) (r))p+1 (x)
x(n² - nk 2n) + n2 - k + xn(n k + 2) + k22
Pk- Pn (1) Pn+1
Pk- (1) Pr P2)
(n 1: + ( nE + 1)
n- k
Pk-1(x)) Pn+1(())
(2-1(2)pn+1l
Which concludes the proof. This now means that:
1 - B < mV - /TL =zy/ V Pk (x)) (v)
(n h)2
Thus we can conclude that
y(1 - ((1-v)pe-((2) (x(x)r(())
y x
(n - k)² pk (f))))(())p (M))
du+1(y) x dn+ (x) pn (x)
TV
From the fact that (2, y) 3 0 it impliest that CE is a decreasing function of n, hence to we only need to show that for any fixed and y 2 (x, y) is positive.
For this, we need to take the limit of the second and third term the above equation.
d x
(1 3)
1 [v] (a(
T
A =
lim n=
V
(x:)(x)) Pr(
lim n
(1 20
1 (x))
Pa (I)
x 21 1)
lim n
(1 x J)Pn+t (x) (z) (1)
x (1 v)
Pn (v)
lim
0
x)
n-
Pa+1 (I) (x)p (v)
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:
k 1
A
lim +00 Pr (x(
nx + 1 x
Hence, this means that:
M(1 - y 1) 1y 1 x V
1-
C (x, v)
>
(1(
Pa (x()) (v))y
v(1 v)pk-1(xx)
11
x x)pa (x) ) (xr)
(y)pkp)kk-11 (v)
1 V
pa px(z)px(1
y
V (P)
y
VV
vay - x
(x(p))p())
38

---

Published as a conference paper at ICLR 2023
V8 1
1
Ph (P)))) (3) (VUPk-1(2) - v /xpx-1(1)
A
yP 2pk-1 yo(k 1) + y zy(k 1) y x 3 0
(x,1) >> 0
Note:
All of the limits must be taken for fixed x and y (e.g. we can't have them approach 0 or simul- tanously with n).
1 x V (T) V > y(1 - y)pk-1(a) ry
(n - 1)2
dn+1 (3) x
(2(2,0)
(x(2)r((1)
d (x) Pr (z)
A- C)
- (1)))))-1(z))
y(1 - y)px-
A
(x(x))rx())
V(1 x)pkk 2 Pk 1 (x) / u)p:
1
Pk- (x)
VV 1
V
Pk-1(v) (v)
(1 y) Pm++ll))
dn-1(y) x
c
Pr.
V
1 (2) pn(x))
(1 11-rr x P=+1(2)
(x)
Pr.
(
y
Puti(U)
p) (y) (z))
Pa (
1
(e(r))(((UPPrt (v)
v Pk (r (v)
Pr (n( (r) Py)Pn+1(3))
y/
Pk- (I) (I) VV
(n kn= (r)())pp()))
(x) (v) Pn (z) (x))
A- A-B-C
1y
x (P)(
(y)
B
=
I V (2)
A = /V Pk (1) Pn)
B zy(n 1b) / -
C zV
2 (kr - x 1) 1) + 1) 1)
B² z²y² - 1) +1)
C lkk + 1 () +1)
(n k)
A = Ty 1
Pa(z) (x)Pk (p) (2)
B = dx+1(y)
(1 - 4)
Pail(y) (v)
Pn())
x
(x) Pr(lr))
\
(1 - x (
Pr (1:)
(r)
(1 v)Pu+110 (1)
Pn+ily)
/
()pp(V)Pp (1) (p) (x)
Pro ( (I)P (I)
1-
C = Vzy
1 (x(xppx(v)
39

---

Published as a conference paper at ICLR 2023
TTy( k)
C - A = rV
Pa (x(
Px (2)())))())
1 R² = zy(n - 1:)2
Pro (M) + zoyln k:)2
(n (2) Pm (1)
(m))
n²xy + no + Tyy + 1 + zy(n - k)²
syk(n2 + (n 1)) +1 2 0
Pn (n)(()
(a(v))n(()
(1 R2)2
1 R² - -22
1 1-2 y
1 (A(())p(()) (1 + R² 1 2R) - x²2 121-1 Pn+ (v)
(C - A)2 B² = Fy
1 Pn Pa
y zyk(n2 (n k)))
1- Pa+1(y)
T
1
(1( (y)p (2)
Pr P (x)pp+10)
1 4y
2R
x Pk (x())pp(()
1 y
1 1 2R
K - 2y 1 ph (x))x(()
x Pr (2)
K Pa+1())(yk((( + (n h) 3pn+1 (1) (3) p
(nz + 1)y(2n2 2nk + 12) 2nny - 1kx + 1ky +
1 Vzy(n k)
Ppat (y)Pk (I) (x)
C C CA+ABB) = Veyy 1
Pk (x( (v)
v Pn (T)
yPn (y)n()Pp (1) (z)
(A B))
V PK (x(x))x(()
2
(A+ + B)²
(Valn k:) + (va(n=k)+
Pa+ (w)
(P++1(.)
a(n k))2 Putz(V)p((()pp(()) + +2/z(n k)
Pa+1 (v+)(V)px(T)paa
yPn+ (3)
zy(n k)² Pr+t Pa+1 (a)))
(p)+1
21/2y(n k) / Pa+ (y) Pact (1)
Wp++1
yP (x) zy(n - k) P Pa+ (1 2E
1 - (A+ B)²
(
E
Vzy(n k) P (
1.2.5 PROOF OF LEMMA
From the definition of gn (I, v) and the fact that pm (x) and rn (2) are non-negative functions we can conclude that:
2
ga (x, (y(1 ((1-)))-((00))))
dnailt) x dnti(() pn(x))
2 0
(n(x)r
Denoting with A the left hand side of the second equation we have:
A = y2 - y) Pm- (x)2
dx+1(y) x²
rr( dn+i(x) (x) (x)
40

---

Published as a conference paper at ICLR 2023
(1 P) 122)
(1 - y) pnttll))
Pro (3))
(1 pn (z) (1) v)Pn (1)Pn (y) (1
tu (xx)
(2)
(x)
(1 - Pn-
(1 1)= Putt (v)
(1 p) Pn) (1)Pn-1(1)
(1
P) (2) Pr) (0n(()))p0 (z)
11)
(PPn-1(1)
x p
(1 2)Dn (x) (z)pn(M)) (1)
Pn- Pn-1(1)
(a)
(1-y)B
(1 x)pn (x) (y)P (100)/))))11()PPt1(C0 (1))r+1((r)
B = (z) x² Pro+1 (9) Pn-1
v'((nc +1) x)((nx 1) x) (ny + 1) g)((ny +) + v)
yn + -1²² - 2² (ny + 1)² x²y²
(y(nz +) ())((() +1) <nn + 1))
(2nry + + uyy
(1- y1 + (1)
A (y 2) (1 (2) Pm
Hence for x & y we have that A > 0 gn 9) > O.
1.2.6 PROOF OF LEMMA 4
xy(1 - 9) (n(x))((())
dx+1(y) x
dn(w)
(n(
A+ B C
dn+1(x) Pn(x)
(x( Pn-1(x)
First we will prove that C - B 3 0
(1 - 8) Pr
x.2
(1 - y) Pn+1(y)
x
C²-B² B² =
Pn- 199)
Pr (1)
(1 x) pn)
(x)
Pn-1(z)²
(1 - x)
pu+1(z)
Pn (r))
Pm(a) (xr)
x² y
Pe (y)
Pn++1 (3)
x
Pn-1 (y)P Pn-11 (2) Pn Pn Pr. (z) (x)
y
(1) Pn+ (z) Pn (1)))
x (x(
Pn (5) Pn-
A
Pn Pn+))))n11x)) - Pn+ -
(p()))()n(t))))))))))))))))))))) + - -) () -
y)
(x)
(pn(z) - x)
(x) (x) Pn( (n (x) (2)² + xn²
= (x) + +yy² 20
C> B
With this we can conclude that fa v) 2 2 - (C B)2 > 0
L = 2 2 2 + 2BC
2h = 2 C - B²
2 y2 (1 y)
(1 - x P) (x p
11 (1 y Pm (3)
_ y
(n(
+
1 I Pn (x(
Pn- Pa-1 (y) Pnly)pu+1 (I)
y
Pn pn(
Pa+i(y)
1 I (z)
Pn- Pn-1 12) (2) (PP(()pn-1(0
Pn-1 (n-1(y)pn-163)
Pa (n())pu41(2)
y
y22 Pn(2))
Pn+1(y)
4²
E Pm. (x)
Pn- (n)
Pr. Pr
41

---

Published as a conference paper at ICLR 2023
y
Pn+ (n)
Pn+i(y)
I. Pm (x)
Pa- Pr (2) (3) Pn Pr(
4 1- y
Pn- (1))
1
Pn Pn-1 (z) (1) +
x Pr (x)
Pn (1)
y Pn+1ly)
(po+1 (I) + PPn-1(x))
x Pm (m))))m(())
Pn-1 (1)
(2nx + 2) y Pn+1ly)
-2x y Pa+1ly)
Pn(x))
x Pm (2 (z) (V) (x-1(2)))+1(z))
1 x pn(x))p(()) Pn- (2)
y
Pn+1l1)
2x
x Pa( Pn 1 (z)P+11(2)
L> 0
C²C² -}>0
122 1-- 1-1ーy y/
Pr P1( (g)
B²
1
Pn (x)² Pn (1)Pn (x: (1)Ppttt))))
x2
1/
Pn 1 (1)2
1 x
Pr (v) (u))) Pmn 1 (x+))(=(
Put1(y)
Pa+1 (2)
1-x
Pn-1 (1)1 (x) (x)²
Pn(y) Pr+ (x))²
121-1 1/
Pa+1 (1)
W 1-x
Pn (r) (x)
D = Pa (E- Pn+1 (1)
- = (p) (2) (pxx(v)² y2²)
- Pm (z) y (x) v'pn(x)² - 2y) 2 + +1) x² (ny + 1)²
2 +y² -nxzyy² 2nx y 12
2nry(y - x) + (y-))(v
x) +x)
(u (2noy + v) 2
Hence with this we can conclude that fn, v) > 0 Vn.
42


Reviews:
[{"review_id": "gSuc4HR5lg", "reviewer": "Reviewer_HAdV", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper tried to train the deep transformer without skip connection and/or normalisation layers.\nThis paper tried to train the deep transformer directly by combining parameter initialisation, bias matrices, and location-dependent rescaling in the signal propagation view.\nThe performance indicates that though the proposed methods converge slowly compared to the previous method of using skip connection and normalisation layers, it could potentially reach the matching performance with enough training cost.", "summary_of_the_review": "This submission provides an interesting way of training deep transformers without skip connection and normalisation layers. ", "strength_and_weaknesses": "Strength:\nThis paper challenged the well-known skip connection and normalisation layers and tried to provide a new way of training deep transformers.\n\nWeakness:\nConverge slow, 5 times more iterations in fact huge in practice.\nThe proposed training framework is more complex compared with skip connection and normalisation layers, means more inductive bias may be introduced to reach the same performance.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "SZifYR6y8o", "reviewer": "Reviewer_2AHB", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "**Note: Score updated from 5 to 8 after author response**\n\nThe authors develop methods to train deep transformers that lack skip connections and/or normalization layers. They achieve this through the following theoretically-motivated interventions: modified initialization, bias matrices, and location-dependent scaling. They are able to train skip connection, normalization-free networks to the same quality as standard transformers (though requiring 5x the number of steps), to within 0.3 perplexity of a standard transformer in the same amount of time, and to depths of 108 layers without divergence.", "summary_of_the_review": "Meaningful progress towards stable training of normalization- and skip-free transformers, but not well-justified why we want normalization- and skip-free transformers. No downstream evaluation, and related work section feels light.", "strength_and_weaknesses": "**Strengths**\n\nExtensive theoretical groundwork to justify design choices\n\nPromising incremental advance for understanding and improving transformer stability\n\nAble to train deep skipless transformers!\n\n**Weaknesses**\n\nNo downstream evaluation\n\nThe motivation for skip and normalization-free transformers could be better justified\n\nNo immediate practical use\n\nLight on related work\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "PDUpAIGvfF", "reviewer": "Reviewer_sotB", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors show that with a careful initialization of parameters, we can train skipless attention-only models without skip connections or normalization layers when trained longer. The major idea is to maintain the signal propagation at initialization in an attention-only model. The authors corroborate their claims with an experimental study on WikiText-103 and C4.", "summary_of_the_review": "Overall, my score is on the positive side. The experiments show that the model can achieve similar training loss as the default transformer, even without skip connections and layer normalization but with longer training. However, the question remains on whether good pretraining performance transfers to downstream tasks.  Hence, I would ask the authors to think about the questions I have posted above during the rebuttal.", "strength_and_weaknesses": "The major strength of the paper is the extensive set of experimentations to showcase that transformers can indeed be trained without skip connections and normalization layers. Such a study is important for the community to understand the necessity of each component inside the transformer. Overall, this paper shows that we simply need a good initialization to train a transformer.\n\n\nHowever, I have the following questions:\n\na) How do the trained transformer (on C4) with E-SPA perform on simple downstream task like SQUAD or SuperGLUE tasks? Does achieving the same pre-training loss with E-SPA also lead to similar performance as a default transformer in downstream tasks? \n\n\nb) During the training of vanilla transformer with E-SPA, how do the attention matrices behave?  Do the position biases stay dominant term for most of the attention matrices in the transformer? If so, then it might explain why the vanilla transformer continues to train well during training, without the usage of skip connections.\n\nc) How fragile is the performance of vanilla transformer with E-SPA ? Is a small learning rate necessary for the transformer to train? For all figures, do you plot the best performing model across the different hyperparameters? How do you select the best performing model across different hyperparameters?\n\nd) What does \"In\" denote in the legend of Figure 3? Moreover, how should I read the plot? Is it that at $\\alpha=0.99$ E-SPA has the same train loss as the default transformer after $100K$ steps? Also, why did the authors decide to compare the train loss at $100K$ step?\n\n\nFurthermore, what is the motivation to measure the performance of a standard transformer with normalized skip connections? Can the authors comment on the necessity of normalization for normalized skip connection in the standard transformer?\n\ne) Does the standard transformer (transformer with skip connections and layer normalization) work with E-SPA initialization? Is there a change in the rate of decrease of the training loss? \n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "KDtO8i6bdxW", "reviewer": "Reviewer_QFbf", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Paper tries to answer: what do skip connections and normalisation layers do? Can we train without them?\n\nThe authors attempt to analyze how our dependency on normalization and skip connections can be removed in deep NNs (specifically, in the self-attention operation of transformers). They do this by porting the “Deep Kernel Shaping” (DKS) technique to Transformer attention, with a few modifications that result in 3 proposed methods: *Value-SkipInit*, Uniform SPA,  and *Exponential SPA*. All of these help train more vanilla Transformers to varying degrees of effectiveness.", "summary_of_the_review": "The authors show that you can theoretically train the Transformer (attention + MLP) without skip connections. However, their methods and assumptions are difficult to apply in the multitude of domains that Transformers are used, and there is no evidence that skip-less or normalization free Transformers are more meaningful than the original Transformer. I recommend this paper as a borderline accept.\n\nI did not read the appendices in depth, nor am I particularly familiar with the literature surrounding DKS methods in Transformers.", "strength_and_weaknesses": "**Strengths**\n\n- (+ +) The proposed solutions truly do enable transformers without skip connections to train at larger depths\n- (+) The paper is clearly laid out and explained, though I did not leave understanding all the math.\n\n**Weaknesses**\n\n- (- -) From a practical point of view, there is little incentive to use the proposed methods. They cannot handle duplicate tokens in input without hacks. It is not evident that the proposed methods offer any real advantages over training with skip connections either from the computational efficiency or interpretability standpoint.\n- (-) The authors only consider causal masked attention which ensures a lower-triangular Attention matrix, a requisite for their analysis, which does not cover the broad spectrum of data on which Transformers are applied.\n- (-) Doubling and tripling the number of learnable parameters (in depth by adding more blocks) of a skipless Transformer does not bring noticeable increase in performance at even 1000k training steps.\n- (-) No code is mentioned to be released with this submission, though the implementation details in Appendix G are thorough.\n\nHaving stated these weaknesses, the authors are not trying to create a practical alternative but instead to provide theoretical proof and empirical justification that it is at least possible to train Transformers without skip connections. To this end the authors succeed, though the paper does not improve our understanding of the roles that normalization and skip connections play.\n\nQuestions\n\n- Pg 4: The notation $(\\mathbf{A}_l)_l$ is not described when it is introduced. Why the double index?\n- Why does Figure 2 come after Figure 3 in the text?\n- Are Fig 1 and 7 created from a Transformer *******trained******* on a task, or a randomly initialized Transformer?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster