Simplifying Transformer Blocks
Authors: Bobby He
Thomas Hofmann
Venue: ICLR 2024 poster (2024)
Keywords: transformers
signal propagation theory
self-attention
initialisation
simpler architectures
skip connections
normalisation
fast training speed
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We use signal propagation and empirical insights to remove various components of the transformer block without loss of training speed, including: skip connections, values and projections parameters.

Abstract: A simple design recipe for deep Transformers is to compose identical building blocks. But standard transformer blocks are far from simple, interweaving attention and MLP sub-blocks with skip connections \& normalisation layers in precise arrangements. This complexity leads to brittle architectures, where seemingly minor changes can significantly reduce training speed, or render models untrainable.

In this work, we ask to what extent the standard transformer block can be simplified? Combining signal propagation theory and empirical observations, we motivate modifications that allow many block components to be removed with no loss of training speed, including skip connections, projection or value parameters, sequential sub-blocks and normalisation layers. In experiments on both autoregressive decoder-only and BERT encoder-only models, our simplified transformers match the per-iteration training speed and performance of standard transformers, while enjoying 16\% faster training throughput, and using 15\% fewer parameters.

Published as a conference paper at ICLR 2024
SIMPLIFYING TRANSFORMER BLOCKS
Bobby He & Thomas Hofmann" *
Department of Computer Science, ETH Zurich
ABSTRACT
A simple design recipe for deep Transformers is to compose identical building blocks. But standard transformer blocks are far from simple, interweaving atten- tion and MLP sub-blocks with skip connections & normalisation layers in precise arrangements. This complexity leads brittle architectures, where seemingly mi- nor changes can significantly reduce training speed, or render models untrainable. In this work, we ask if the standard transformer block can be simplified? Com- bining signal propagation theory and empirical observations, we motivate modi- fications that allow many block components to be removed with no loss of train- ing speed, including skip connections, projection or value parameters, sequen- tial sub-blocks and normalisation layers. In experiments on both autoregressive decoder-only and BERT encoder-only models, our simplified transformers emu- late the per-update convergence speed and performance of standard transformers, while enjoying 16% faster training throughput, & using 15% fewer parameters.
INTRODUCTION
The transformer architecture (Vaswani al.] 2017) is arguably the workhorse behind many recent successes in deep learning. A simple way to construct a deep transformer architecture is by stacking multiple identical transformer "blocks" one after another in sequence. Each block, however, is more complicated and consists of many different components, which need to be combined in specific arrangements in order to achieve good performance. Surprisingly, the base transformer block has changed very little since its inception, despite attracting the interest of many researchers.
In this work, we study whether the standard transformer block can be simplified. More specifically, probe the necessity of several block components, including skip connections, projection/value matrices, sequential sub-blocks and normalisation layers. For each considered component, we ask if it can be removed without loss of training speed (both in terms of per-update step & runtime), and what architectural modifications need to be made to the transformer block in order to do so.
We believe the problem of simplifying transformer blocks without compromising training speed is an interesting research question for several reasons. First, modern neural network (NN) architectures have complex designs with many components, and it is not clear the roles played by these different components in NN training dynamics, nor how they interact with each other. This is particularly pertinent given the existing gap between theory and practice in deep learning, where theorists work- ing to understand the mechanisms of deep learning often only consider simplified architectures due to convenience, not necessarily reflective of modern architectures used in practice. Simplifying the NN architectures used in practice can help towards bridging this divide.
On a related theoretical note, our work highlights both strengths and current limitations of signal propagation: theory that has proven influential due to its ability to motivate practical design choices deep NN architectures. Signal propagation (Poole et al. 2016, Schoenholz et al. 2017] Hayou et al. 2019) studies the evolution of geometric information in NN initialisation, captured through inner products of layerwise representations across inputs, and has inspired many impressive results in training deep NNs (Xiao et al. 2018 Brock et al. 2021 Martens et al. 2021 Zaidi al., 2023). However, the current theory only considers a model at initialisation, and often considers only the initial forward pass. As such, signal propagation at present is unable to shed light on many intricacies of deep NN training dynamics, for example the benefits of skip connections for training speed. Though signal propagation is crucial in motivating our modifications, we would not have arrived at our simplified transformer blocks from theory alone, and relied also on empirical insights.
"Correspondence bobby. www.eeeeeeeeeeeenn .et.h.ch..

---

Published as a conference paper at ICLR 2024
Pre-LN
Ours
Namunii
NonLi
MUP
MLPIn
Parallel
Pray
Attendon
MLP Oul
Attention
Nor
Norm
Figure 1: Comparison between different Transformer blocks. (Left) The standard Pre-LN block. (Top Right) Our most simplified block. (Bottom Right) The parallel block (Zhao et al. 2019 Wang Komatsuzaki 2021). Like the parallel block, our block eschews the need for sequential sub- blocks, but we additionally remove all skip connections and normalisation layers, as well as value and projection parameters. Here, @ denotes matrix multiplication, and denotes (potentially weighted) sum.
Finally, on the practical side, given the cost of training and deploying large transformer models nowadays, any efficiency gains in the training and inference pipelines for the transformer architec- ture represent significant potential savings. Simplifying the transformer block by removing non- essential components both reduces the parameter count and increases throughput in models. In particular, we show that it is possible to remove skip connections, value parameters, projection pa- rameters and sequential sub-blocks, all while matching the standard transformer in terms of training speed and downstream task performance. As a result, we reduce parameter count by up to 16% and observe throughput increases of 16% at both train and inference time.
Our starting point to simplify Transformer blocks is He et al. (2023). who show that respecting signal propagation principles allows one train deep Transformers without skip connections or normalisation layers, but at significantly reduced convergence speeds per parameter update. We first show that regulating the updates to values and projection parameters (Sec. 4.1p. or in fact removing them entirely (Sec. 4.2), improves the performance of skipless attention sub- blocks, and recovers the lost per-update training speed reported by He al. (2023). This removes half of the parameters and matrix-multiplicale
in the attention sub-block. In Sec. 4.3 we show our simplifications combine profitably with parallel sub-blocks (Zhao et al.) 2019; Wang & Komatsuzaki 202I), allowing us to remove all remaining skip connections and sequential sub-blocks without compromising per-update training speed, whilst further boosting the throughput increase to be 16%, in our implementation. Finally, in Sec. 5 show that our simplified blocks improve when scaled to larger depths, work well both encoder-only and decoder-only architectures, and that our findings also hold when scaling training length. We conclude with discussion of limitations and future work in Sec. 6
2 RELATED WORK
Simplifying deep NNs by removing block components has received a lot of attention, both in trans- formers and other architectures. In these works, signal propagation theory often acts as inspiration.
For a pair of inputs x, x', mapped to a pair of representation/actiitttin vectors 2t; T/, e Rd at layer signal propagation theory studies the evolution of activation inner products Jo, x', Loe, æt, _ x. at initialisation, which can be tracked with their large limits (Lee et al, 2018 Matthews et al. 2018 Yang 2019). Several pathologies afflicting poorly designed deep NNS can be identified as result (Shoonholzz et al., 2017: Hayou et al. 2019 Yang et al.) 2019 Dong al. 2021 Martens et al. 2021). For examlee, the activation norms _., m; may blow up or vanish, or the cross products 2 a may converge to a value independent of the inputs x, x' at large 1, in which case deeper layers of the model
unable to identify different inputs. Avoiding such degen- eracies is important to allow for good training dynamics and generalisation in deep NNs (Balduzzi
et al. 2017 Xiao al. 2018 2020 Hayou et al. 2021
Martens al. 2021 Noci al. 2022).
2

---

Published as a conference paper at ICLR 2024
It has been shown that judicious use of weight initialisations and architectural tools, like skip connec- tions and normalisation layers, can improve signal propagation degeneracies and the trainability of deep NNs. Such considerations have motivated principled modifications with simpler architectures. De & Smith (220) show that an implicit mechanism of Pre-LN skip connections is to downweight the residual branch relative the skip branch, leading better signal propagation. They also show that explicitly downweighting the residual branch allows normalisation layers to be removed with- out affecting performance. The idea of downweighting residuals for improved signal propagation & trainability has been studied extensively in the literature (Zhang et al., 20188 Hanin & Rolnick)
2018 Tamowski et al.] [019] Zhang al.] 2019; Arpit et al. Xu et al.) 2020; Bachlechner
et al.
2021:
Touvron et al. 2021 Hayou et al. 2021 Hayou & Yang [023] Martens et al. 2021
Davis et al.
2021 Noci et al. 2022 Wang et al. 2022a Huang et al. 2020; Wang et al. 2022b)
2016). it has been shown that transforming non-linear activation functions in MLPs and CNNs be more linear according to given deep architecture can enable good signal propagation even without skip connections (Martens et al.) 2021] Zhang et al. 2002] Li ct al., 2022). He al. (2023) apply similar considerations to the self-attention mechanism, where the key insight is that attention matrices need to be more identity-like in order to prevent signal degradation in skipless transformers. However, these works find that skipless architectures suffer from significant losses in training speed compared to their residual counterparts, when using stan- dard optimisers like SGD or Adam. Such differences were not observed with stronger optimisers like K-FAC (Martens & Grosse| 2015) on CNNs, and this inability to explain training phenom- ena highlights a current limitation of signal propagation theory. [Ding et al. (2021) 2023) design
For skip connections (He et al.
CNN, RepVGG, that can be trained like residual architecture for fast per-update convergence, but reparameterised to be skipless at test time for significantly higher inference throughput. This reparameterisatioon is related to our considerations of value and projection parameters in Sec. 4
Many works have considered simplifications or improvements specific the transformer. Most relevant to our work is the parallel block (Zhao et al. 2019; Wang & Komatsuzaki 202T) (pictured Fig. | bottom right), which computes the MLP and attention sub-blocks in parallel for efficiency gains, with minimal performance loss. Trockman & Kolter (2023) observe that the product of value and projection parameters often has a large identity component trained transformers, and design an initialisation mimicking this to improve performance in standard transformers on small datasets. We find these matrices can be fixed to the identity without loss of performance, which removes them from our simplified architecture. Other works have considered reducing the frequueccy of MLP sub-blocks (Sridhar et al. 2022, Pires et al. 2023) or efficient replacements to softmax attention Katharopouloos et al. [0200] [Schlag et al.] Choromanski et al. [202 /Skkbbatta et al. (2019) remove the MLP by integrating it into the attention sub-block, augmented with persistent memory.
3 PRELIMINARIES
A deep transformer architecture of depth L is formed by sequentially stacking L transformer blocks. The most common block is Pre-LN, depicted in Fig. (left), which we treat as a baseline for com- paring training speed, both in terms of per-update and runtime. It differs from the original Post-LN block only in the position of the normalisation layers relative to the skip connections, but is more popular as the Post-LN block suffers from poor training stability and signal propagation in deep layers (iong et al. 2020: Liu et al. 2020 Noci et al. 2022 He et al. (003)..
Transformer blocks take representations of sequences as inputs. For an input sequence representa- tion X E R Txx, with T tokens and dimension d, the Pre-LN block outputs Xour where:
our = OFF X + 3Pp ML.((Nrm(()) where X USA X + ASA MHA(Norm(X.)) (1)
with scalar gain weights oFF, BrF, aSA, ASΛ fixed to by default. Here, "MHA" stands for Multi- Head Attention (detailed below), and "Norm" denotes normalisation layer (Ba et al.) 2016 /Zhng & Sennrich, 2019). In words, we see that the Pre-LN transformer block consists of two sequential sub-blocks (one attention and one MLP), with normalisation layers and residual connections for both sub-blocks. and crucially the normalisation layers are placed within the residual branch. The MLP is usually single hidden-layer, with hidden dimension that is some multiple of d (e.g. 4 (Vaswani et al. 2017) or 8/3 (Touvron et al. 2023), )) and acts each token in the sequence independently.
The MHA sub-block allows tokens to share information between one another using self-attention. For input sequence X, the self-attention mechanism outputs:
3

---

Published as a conference paper at ICLR 2024
( xxwewk X x *)) (2)
( - A(X)XWV, where AX) Softmax
where WΩ, W E Rddd and w dddd are trainable query, key and value parameters respec- tively, Here, the attention matrix ( E BTXT can be thought of as allowing different tokens to "mix" with each other. M E T is mask taking values in f0, -o) that depend on the mod- elling task. For causal auto-regressive transformers like GPT, Maj i 2 j. which prevents token from obtaining information from future tokens. bidirectional models like BERT, masking is typically applied at the token level and not in the attention mechanism (i.e. M is the zero matrix).
The Multi-Head Attention name arises because is typical in practice to apply self-attention on H different "heads" (with independent parameters) with dy d2 as follows:
MHA(X) = Concat (oncaa(Attm) (Attn, (X).... A X)) W
(3)
where WP E RRddd denotes trainable square projection matrix that combines different attention heads. If we let wy denote the value parameters for head n, then the concatenated value weights W Concat( Con ....., W)) E Rdddd can be viewed as a square matrix. One of our key findings, in Sec. 4.2 to show that fixing the value and projection parameters. WV and W to the identity matrix significantly improves per-update training speed in skipless transformer blocks (to speeds matching or even outperforming the standard Pre-I.N block), whilst simultaneously significantly reducing the parameter count and matrix- miliiiiuulllliiiiionnn FLOPs required, thus increasing throughput.
4 SIMPLIFYING TRANSFORMER BLOCKS
We now describe how we arrive at our simplest Transformer block, Fig. (top right), starting from the Pre-LN block, using combination of signal propagation theory and empirical observations. Each subsection here will remove one block component at time without compromising training speed, and we aim to provide an intuitive account our progress simplifying the Pre-LN block.
All experiments in this section use an 18-block 768-width causal decoder-only GPT-like model on the CodeParrot dataset which sufficiently large that we are in a single epoch regime with minimal generalisation gap (Fig. 2). allowing to focus training speed. We provide depth scaling, and non-causal encoder-only, experiments, in Sec. 5 We use a linear decay learning rate (LR) schedule") with AdamW (Loshchilov & Hutter 2017), with linear warmup for the first 5% steps. The maximum LR is tuned on training loss, using logarithmic grid. Additional experimental details are in App. D
4.1 REMOVING THE ATTENTION SUB-BLOCK SKIP CONNECTION
We first consider skipless attention sub-block, whose output has the simple interpretation of adding, to each token, other token representations according to the attention matrix. In the nota- tion of Eq. (I) this corresponds asa O. Naively removing the attention skip leads to a signal degeneracy called rank collapse (Dong et al. 2021). which harms trainability (Noci et al. 2022).
Setup He al. (2023) outline modifications needed the self-attention mechanism in order to correct these signal degeneracies at large depths, and train such deep skipless networks for the first time. One method they introduce, Value-Skiplnit, modifies the self-attention matrix to compute:
(X) + (aly + BA()))
with trainable scalars a, 39 initialised to and respectively, and IT E RTT is the identity matrix.
The key insight here is to initialise the self-attention matrix to have dominant identity component that encourages a token to attend to itself more relative to other tokens, much in the same way that a Pre-LN skip upweights the skip branch relative to the residual branch for good signal propagation at large depths (De & Smith 2020). We point out that these considerations only apply at initialisation.
Noci et al.
2023 propose extension, Shaped Attention, also motivated by signal propagation:
A(X) <- (aIr + (A - CC).
(5)
'Our setting is taken from https: / /twgg//uggg////// www.nigginiggg face co/ eearn/nnnn n wt--ccuureecapppee
We found linear decay slightly outperform cosine decay for both our models and haselines (c.f. Fig.
4

---

Published as a conference paper at ICLR 2024
Here, a, 2 are trainable, and C' is constant (not trained) centering matrix, set to be equal to the values of A when the query-key dot product _ X-xWWWK x
T is zero Like He et al. 2023).. we initialise queries wº 0, which exactly zeros the query-key dot product ini- tialisation. Then, / Y means that (X) - 7C 0 at initialisation, and a 1 en- sures a dominant identity component, and good signal propagation. Ali et al. (2023) also cen- tre attention and show it helps prevent oversmoothing in vision transformers and graph NNs.
Vd,
We found Shaped Attention, Eq. (5), to slightly out- perform Eq. (4) (c.f. Fig. [3). and use it in our ex- periments on skipless attention sub-blocks, with & = Y a = 1 at initialisation unless stated other-
V-Skipinit (as. train
ViSkipinit (asx = 0), eval
Pre-LN, train
1
Cross entroyy
Pro-LN, eval
wise. We also use head dependenn scalars in Eq. (
on, 3n and "/h, which provided small additional per- formance boost. One final important implementa- tion detail is that for any skipless block we explic- itly downweight the MLP branch by initialising train- able 6TF = O (E < 1 = ayr. This is moti- vated through signal propagation theory (c.f. Stable ResNet, Hayou et al. (202I)), and accounts for the fact that removing skip connections (in either MLP or MHA sub-block) reduces the implicit downweighting 2020). For the depth L 18 networks in this
22
3BC
Training step
Figure 2: Loss of training speed in trans- formers without attention sub-block skip (Hl et al. (023). even with Shaped Auen- tion, Eq. and MLP skips (uFF = 1).
effect of Pre-LN blocks (De & Smith section, we initialise BFF 0.1.
2.28
Recovering lost training speed Despite allowing skip- less transformers to train for the first time, He et al. (2023) reported a significant loss of training speed per step com- pared to the Pre-LN block. We verify this in Fig. 2
3.44
To recover the lost training speed without attention skips, note that identity attention matrices make deep trans- former with no MLP sub-blocks act like a deep skipless linear NN at initialisation f(x)=XII'= (WY W(".. where WV, W.P are the value and projection weights in layer L. He
122
Orth Win Wi. identity Wy, Wat
328
3.35
05 Residual weights Bu.
64
Figure 3: Restricting updates W W through smaller 3v, recovers training speed in skipless transformers (AsA = 0).
He et al. (2023), they initialise WY, W to be independent random orthogonal matrices to avoid signal degeneracies from Gaussian initialisations (Saxe et al. [0033] Hu et al. 2020 Meterez et al. 2023).
It is known that such deep skipless networks train slower than their residual counterparts MMartens eet al. 2021). Moreover, it is also known that Pre-LN skips downweight residual branches (De & Smith (020), which is equivalent reduced learning rates & downscaled parameter updates from initialisation in linear layers (e.g. Ding et al. (2023); we outline and empirically verify this duality in App. Ap. This motivates to study reparameterisation of the value/projection weights WV W.P
WV av W + By AWV, and W = ap W P + 3p AWP.
(6)
with "skip" WV. mi fixed to be random orthogonal to preserve the signal propagation achieved at initialisation, and "residual" A AW trainable and initialised to zero. We consider downweighting
the residuals with fixed By I av = 1. which biases the matrices WV, W to stay closer to initialisation, and would expect By = O(+))) recover the benefits of skip connections (Hayou
et al. 2021 Similar considerations apply for W init: AWP 3r.33p
In Fig. 3. we find as expected that using smaller By and 3p with this reparameterisation, Eq. (6) already restores much of the training speed loss in skipless attention-blocks using orthogonally initialised A WP To close this gap further, we note that from a signal propagation perspective, init: init.
initialising W W
init: init to be the identity matrix is equivalent orthogonal initialisation when the
For example, when there is no masking,
C' becomes the uniform T x T stochastic matrix: 1111 T
*We set aside the MI. sub-block here for simplicity, but point out that of our experiments use ML.Ps so our findings carry over to the full setting.
A ithough the initial forward pass is identical regardless of By, due to zero initialised A !
5

---

Published as a conference paper at ICLR 2024
attention sub-block is skipless. With identity initialisation W Wi Ia a consistent improvement over orthogonal initialisation, which essentially matches the Pre-LN block. One thing can conclude from this experiment. is that restricting the updates to the values and projections from their initialisation replicates the effects of the attention sub-block skip connection, and recovers the lost per-update training speed. We investigate the difference in Fig. 3 of performance between identity and random orthogonal in the appendix (Fig. 15
4.2 REMOVING VALUE AND PROJECTION PARAMETERS
In fact, we can also conclude from Fig- 3 that it is possible to completely remove the value and projection parameters W WP with minimal loss of per-update training speed. Namely, when 3v = Bp- = 0 and identity-initiaiin performance after equal numbers of training steps. In this case, we have WV W = I through- training, i.e. the values and projection parameters are identity.
W W init = I, we essentially match the Pre-LN block
mi
To further verify this surprising observation, we consider reparameterised WV, W in Eq. 66 with identity WY
8.20
Layer
init W mit but now trainable scalars av, ap, 3p, From an initialisation of ay=ap=l and 3v=Bp=0.2, plot the evolution of "residual-skip" ratios 8v. de in Fig. Weight decay was not applied on av, ap.ap,
0.5
ResidualSkip gain
Prl pelo
Vake Kely
8.0
av ap
We see that the residual-skip weight ratios f Be con-
0.00
av ap verge 0 for the vast majority of layers, which indicates that these reparameterised matrices WV WP converge to the identity during training. As a result, the extra ca- pacity to perform linear projections via WV w is not used. We plot the corresponding trajectories for other scalar parameters like 3rr, in Figs. 17 to 20 which do not tend to 0. The model in Fig. 4 with trainable WV. W achieved worse final evalua- tion loss than the model in Fig. 3 with identity WV, W (1.194 vs. 1.178). Interestingly, this trend is reversed if the attention skip is re-added (Fig. 23
DOC
Training
Figure 4:
Residual-skip gain ratios converge to during training.
fv. Bp ay op
V-Skipintt lHe ct 20231
We thus elect to remove values and projection parameters WV. W in our skipless attention sub-blocks, by setting them to the i We refer to the resulting sub-block as the Simplified Attention Sub-block (SAS). Our full SAS block is depicted in Fig. 10 and we detail the mathemat- ical computation in Eq. (12.. We note that SAS blocks use only half of the parameters as well as half the matrix- multiplications in the attention sub-block: only query and key parameters remain. This results in a 13% reduction in the total number of parameters (146M 167M for 18 blocks) in the models we consider in this section
2.13
SAS (14.2)
Cross-ntruu
SAS (54 3)
SAS-P, no norm (54.4)
Paralicc (Wbng
120
Runtime (hours)
Figure 5: Training speed in terms of run- time. We see our models match (or even slightly outperform) the Pre-LN block.
In Fig. 5 we see that when comparing speed in terms of wall-clock runtime on an A5000 GPU, our SAS block already trains at speeds (slightly) outperforming the default Pre-LN transformer. The corresponding plot comparing speed in terms of training steps taken is provided in Fig.
26 more
detailed analysis of efficiency gains in our simplified blocks can be found in Sec. 5
Though we do not have a rigorous proof for why the training dynamies in skipless transformers forgo additional capacity by converging to identity value and projection parameters (Fig.E) nor why fixing such matrices to the identity results in no performance degradation and in fact trains faster than having trainable values and projections (Fig. 3 we offer some half-explanations. First, the fact that WV W are simply linear projections of the input sequence representations X (as opposed to in the MLP sub-block where elementwise non- oon-iineaistees are places between such matrices), could
'hhe only exception is the first layer's value parameters WY. which is the only ratio above 0.05 in Fig. 4 We very minor performance gains by keeping w (c.f. Fig. 24 so keep it whilst removing all other w!. W for I T.
In general, the % of all parameters removed by W " depends on the ratio of width d to the MLP hidden width dyr, vocabulary size and depth. Here, we have width d = 768, MI.P hidden width dgr 3072 4d, vocabulary size 50K and depth I. 18. In the large depth limit, only the ratio dre /d matters.

---

Published as a conference paper at ICLR 2024
mean that the additional capacity afforded by such matrices is not particularly substantial This is corroborated by Trockman & Kolter (2023) who found in trained transformers the product W V W'w" often has a dominant identity component, Also, from signal propagation perspective, there is no reason why initialising such matrices to be non-identity (e.g. orthogonal or Gaussian) would be preferred to identity initialisation, nor is it clear why they would be necessary in the first place, especially given the additional matrix- matrix-multiplication FLOPs they require.
4.3 REMOVING THE MLP BUB-BLOCK SKIP CONNECTION
So far we have simplified the Pre-LN transformer block by removing, without loss of training speed, three key components: 1) the attention sub-block skip connection, as well as 2) value and 3) projec- tion matrices. We next turn to removing the remaining skip connection in the MLP sub-block.
This proved more challenging. Like previous works (Martens et al. 2021 Zhang et al., 2022 He et al., 2023), we found that making activations more linear, motivated through signal propagation, still resulted in a significant loss of per-update training speed without MLP skips when using Adam, as shown in Fig. [5] We also experimented with variants of the Looks Linear initialisation (Balduzzi ct al. 2017), with Gaussian, orthogonal identity weights, to avail. As such, activations (e.g. ReLU in this section) and initialisations in the MIP sub-block throughout our work.
standard
Instead, we turn the idea of parallel MHA and MLP sub-blocks (Zhao et al., 2019 Wang & Komatsuzaki] (02I), which has proven popular in several recent large transformer models, such as PALM (Chowdhery al. 2022) and ViT-22B (Dehghani et al., (023). The parallel transformer block is depicted in Fig. [ right), and mathematically, given input Xin it outputs Xoar, where:
X out Occomb Xin + 3P MLP((Nom(X.nn) + 95Λ MHA(Norm(X..))
with skip gain Cccmmb 1, and residual gains BFF = ASA = 1 as default.
In the parallel block, the MLP and MHA sub-blocks each take the same (normalised) input, af- fording more parallelisation compared to the standard Pre-LN block, which computes sub-blocks sequentially. The two sub-blocks are combined by summing their outputs, in conjunction with a single skip connection, with weight acamb- This parallelisation, as well as the removal of one skip
connection and one normalisation layer enables efficiency gains: [Chowdhery et al. (2022) report the parallel block has 15% faster training speed compared to the standard "sequential" Pre-I.N block.
It is straightforward to combine our simplifications from Secs, 4.1 and with the parallel block in
Eq. (7: we simply 1) use our SAS attention sub-block, Eq. (2). 2) set fixed Clcomb = 0 to remove all skip connections in the block, and 3) downweight 3FF < 1. The resulting block is pictured in Fig. [1 and we refer to it as in terms of runtime compared to the SAS and Pre-LN blocks, and matches the training speed of the parallel block despite using 13% fewer parameters. Our intuition is that the combination of Shaped Attention and identity values/projections preserves signal between blocks throughout training and replaces the need for a skip connection in either sub-block. Moreover, we note that our attention sub-block is the identity function, Xout Xin. at initialisation, so there is no difference between our sequential SAS (Fig. 10P and parallel SAS-P (Fig. IT) blocks at initialisation.
(SAS-P for short). We see in Fig. thhtt SAS-P trains even faster
4.4 REMOVING NORMALISATION LAYERS
The final simplification we consider is removing normalisation layers, leaving us with our simplest block (Fig. [ top right). From signal propagation initialisation perspective, normalisation has been expendable at all stages our simplifications in this section: the idea is that normalisation in Pre-LN blocks implicitly downweights residual branches, and this beneficial effect can be replicated without normalisation by another mechanism: either explicitly downweighting residual branches when skips are used, or biasing attention matrices to the identiyy/ransffrmiing MLP non- non-linearities to be "more" linear otherwise. As we account for these mechanisms in our modifications (downweighted MLP Orr & Shaped Attention), from initialisation perspective there is no need for normalisation.
Of course, these modifications have effects on training speeds and stability beyond initialisation, which are harder predict from existing theory alone. In Fig. we see that removing normalisa-
FF.g.. . in single-head attention one can reparameterise VV P into one matrix with no expressivity loss.

---

Published as a conference paper at ICLR 2024
tion allows even simplest transformer block, which does not have skips, sequential sub-blocks values, projections or normalisations, to match the training speed of the Pre-LN block in terms of runtime. Having said that, do observe slight degradation in training speed per iteration, seen in Fig. [6] suggesting that normalisation layers have some beneficial properties for training speed beyond what is captured by signal propagation theory. We thus treat our SAS (Fig. [T] and SAS-P (Fig. [T] blocks, with normalisation, as our main approaches. On this note, we point out that De- lighani et al. (2023) found extra normalisation on the queries and keys to provide improved training stability in ViT-22B, going against the recent trend of researchers seeking to remove normalisation.
5 FURTHER EXPERIMENTAL ANALYSIS
Having introduced all of our simplifications in Sec. we now provide further empirical analysis of our simplified blocks across a range of settings, as well as details of the efficiency gains afforded by our simplifications. In interest of space, additional experimental details can found App. D
Depth Scaling Given that signal propagation the- ory often focuses on large depths, where signal de- generacies usually appear, it natural ask whether the improved training speeds of our simplified trans- former blocks also extend to larger depths. In Fig.[6) we see that scaling depth from 18 to 72 blocks leads to an increase in performance in our models as well as the Pre-LN transformer, indicating that our sim- plified models are able to not only train faster but also utilise the extra capacity that more depth pro- vides. Indeed, the per-update trajectories of our sim- plified blocks and Pre-LN are near-indistnn indistngguisaabee across depths, when using normalisation.
Pre-LN
SAS (54
SAS-P (54.3)
Eval Crossrrss--eee Entrs-EEttrpyy
SAS-P, norm (54.4)
V-Skiptnit (He et al. 2023)
Depth=18 Depth=72
25
x
304
15c
201
Trining step Our models improve when deeper (dashed, marked lines) vs. shallower (solid lines), unlike V-Skiplnit (He et al. 2023).
Figure 6:
On the other hand, we see that Value-SkipInit (He et al., 2023) actually trains slower per update at depth 72 compared to 18 despite the increase in capacity and parameter count. Moreover, the gap in performance between Value-Skiplnit and the other models increases with larger depth, which implies poor scalability of the previous method. We note that 72 blocks is already reasonably deep
by publically-available modern standards (Hoffmann et al. 2022 Touvron et al. 2023).
BERT Next, we demonstrate our simplified blocks' performance extends to different datasets and architectures besides autoregressive decoder-only, as well as on downstream tasks. We choose the popular setting of the bidirectional encoder-only BERT model Devlin al. (2018) for masked lan- guage modelling, with downstream GLUE benchmark.
Crammed SEAT (PrlA
In particular, we adopt the "Crammed" BERT setup Geiping & Goldstein 2023), which asks how well one train a BERT model with a modest train- ing budget: 24 hours on a single consumer GPU. The authors provide an architecture, data pipeline and training setup that has been optimised for this low resource setting. We note that the Crammed ar- chitecture uses the Pre-LN block, and describe other setup details in App. D We plug-in our simplified blocks, keeping the existing optimised hyperparam- eters, besides tuning learning rate and weight decay.
Paral oi 1Wanng at al. 2021)
S45 (6421 SAS- (64.3)
SAS-P, no norm (54.4)
ViSAlaira: 2ie al 20231
MLM
3.0
20
Runtime (hours)
Figure 7: Masked language modelling loss vs runtime on a 2080Ti GPU for 24 hours.
In Fig. 7 we see that our simplified blocks (especially with normalisation) match the pre-training speed on the masked language modelling task compared to the (Crammed) Pre-LN baseline within the 24 hour runtime. On the other hand, the removal of skip connections without modifying the values and projections (as Hn et al. (2023) once again leads to a significant loss of training speed. In Fig. 27 we provide the equivalent plot in terms of microbatch steps.
Moreover in Table I we find that our methods match the performance of the Crammed BERT baseline after finetuning the GLUE benchmark. We provide breakdown over the downstream tasks in Table 2 We use the same finetuning protocol as Geiping & Goldstein (2023) (5 epochs, constant hyperparameters across tasks, dropout regularisation) for fair comparison. Interestingly,
8

---

Published as a conference paper at ICLR 2024
Value-Skiplnit is largely able to recover from its poor pre-training in the fine-tuning phase. This, combined with the need for dropout when fine-tuning, suggests that factors besides pre-training speed are also important for fine-tuning. As the focus of our work primarily concerns training speed from random initialisations, we leave this to future work. Relatedly, we found removing normalisa- tions (Sec. (.4) to cause instabilities when fine-tuning, where a small minority of sequences in some downstream datasets had NaN values in the initial forward pass from the pre-trained checkpoint.
Efficiency Gains In Table I we also detail the pa- rameter count and training speeds of models using different Transformers blocks on the masked lan- guage modelling task. We compute the speed as the ratio of the number of microbatch steps taken within the 24 hours of pre-training relative to the baseline Pre-LN Crammed BERT. We see that models use 16% fewer parameters, and SAS-P & SAS are 16% & 9% faster per iteration, respectively, compared to the Pre-LN block in our setting. We note that in our implementation the Parallel block is only 5% faster than the Pre-LN block, whereas Chowdhery et al. (2022) observed 15% faster training speeds, suggesting that further through-
Table 1: GLUE henchmark & efficiency gains. Our SAS & SAS-P match the downstream perfor- mance of the Pre-LN haseline up statistical sig- nificance over seeds, but use 16% fewer param- eters and enjoy up to 16% faster throughput.
Block
GLUE
Params
Speed
Pre-LN (Crammed)
78.9
120M
Parallel
78. 78.5+.6
120M
1.05
V-Skiplnit
78. +..
120M
0.95
SAS (Sec.
4.2
78 78.4+.s
101M
1.09
SAS-P (Sec. 4.3
78.3
101M
1.16
SAS-P, norm
101M
1.20
out increases may be possible with a more optimised implementation. Our implementation, like
Geiping & Goldstein (2023), uses automated operator fusion in PyTorch (Sarofeen et al., 2022)
Longer training Finally, given the current trends of training smaller models for longer on more data (Hoff-
SAS (5+.2)
mann et al.) 2022; Touvron et al., 2023), we investigate
SAS-P 134.3)
if our simplified blocks continue to match the training speeds of the Pre-LN block with longer training. To do this, we take our models from Fig. E on CodeParrot and
Eval Cross-entrppy
2.6
1x token
tokena
train with 3x tokens. To be precise. we train for around 120K (rather than 40K) steps with batch size 128 and se- quence length 128, which results in around 2B tokens. In Fig. 8, we do indeed see that our simplified SAS and SAS-P blocks continue to match outerperform the Pre- LN block in training speed when trained on more tokens.
Rumtime (hours)
Figure 8: Training speeds continue to hold with longer training.
DISCUSSION
Limitations and future work While we have demonstrated the efficacy of our simplifications across architectures, datasets, and tasks, the models we have considered (100-300M parameters) are small relative to the largest transformers. It would be interesting to investigate the performance of our simplified blocks at larger scales, especially because Chowdhery et al. (2022) report parallel blocks improve relative to Pre-LN blocks with scale. Our depth scaling experiments already show promise in this regard. On the theoretical side, though we were able to match the training speed of Pre-LN blocks with normalisation removed (Fig. (5). there are still unanswered questions regarding the benefits of normalisation for training speed and stability, and we were unable to remove normal- isation with good downstream task performance. Moreover, while we tuned key hyperparameters like learning rate, it is possible that many default hyperparameters and choices we inherited, e.g. the AdamW W optimiser, or fine-tuning protocol, are overfit to the default Pre-LN block, and an exhaus- tive hyperparameter search for our simplified blocks would yield further improvements. Finally, on the practical side, we believe that a more sendwaar-specifi. implementation of our simplified blocks could give further improvements to training speed and performance.
Conclusion In this work, we asked whether it is possible to simplify the standard Transformer block by removing unnecessary components. Combining signal propagation theory and empirical insights, we have shown that it is possible to remove skip connections, sequential sub-blocks, value and projection parameters, without loss of training speed or downstream task performance. As a result, our models have around 15% fewer parameters and 16% increased throughput. We believe our work can lead to simpler architectures being used in practice, thereby helping to bridge the gap between theory and practice in deep learning, and reducing the cost of large transformer models.
9

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY STATEMENT
Our code for experiments on auto-regressive transformers can be found at https: //gpt:: com/bobby- L
ACKNOWLEDGMENTS
We would like to thank Sotiris Anagnostidis, Andrei Ivanov & Lorenzo Noci for helpful discussions in the initial stages of this project, and James Martens, John Martinis, Keivan Mohtashami, Tiago Pimentel & Imanol Schlag, as well as the anonymous reviewers, for constructive feedback on an early version of this manuscript.
REFERENCES
Ameen Ali, Tomer Galanti, and Lior Wolf. Centered self-attention layers. arXiv preprint vrXiv:2306.01611 2023.
Devansh Arpit, Víctor Campos, and Yoshua Bengio. How to initialize your network? robust initial- ization for weightnorm & resnets. H. E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Buc,
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 2016.
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence at large depth. In Uncertainty in Artificial Intelligence, pp. 1352-1361. PMLR, 2021.
David Balduzzi, Marcus Frean, Lennox Leary. JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams. The shattered gradients problem: If resnets the answer, then what is the question? In Interna- tional Conference on Machine Learning, pp. 342-350. PMLR, 2017.
Andy Brock, Soham De, Samuel L Smith, and Karen Simonyan. high-performance large-scale image recognition without normalization. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 1059-1071. PMLR, 18-24 Jul 2021.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane. Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with per- formers. In International Conference Learning Representations, 2021. URL https: //openreviest www.tet..oooooo
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, al. Palm: Scaling language modeling with pathways. arXiv preprint urXiv:2204.02311 2022.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. International conference on machine learning, pp. 933 933-941. PMLR, 2017.
Jared Davis, Albert Gu, Krzysztof Choromanski, Tri Dao, Christopher Re, Chelsea Finn, and Percy Liang. Catformer: Designing stable transformers via sensitivity analysis. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 22999, PMLR, 18-24 Jul pta.eesess.........sssssll.....
2021. URL W...URRRttte... http:ts:///ooceedd ngs ml r wwtpppsssss/////.
html
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.
10

---

Published as a conference paper at ICLR 2024
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In International Conference Machine Learning, pp. 7480-7512. PMLR, 2023.
Jacob Devlin, Ming-Wei Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 0r1iv:/810.045000 2018.
Xiaohan Ding, Xiangyu Zhang, Ningning Ma, Jungong Han, Guiguang Ding, and Jian Sun. Repvgg: Making vgg-style convnets great again. Proceedings of the IEEE/CVE Conference on Com- puter Vision and Pattern Recognition (CVPR), pp. 13733--17422 June 2021.
Xiaohan Ding, Honghao Chen, Xiangyu Zhang. Kaiqi Huang, Jungong Han, and Guiguang Ding. Re-parameterizing your optimizers rather than architectures. In The Eleventh International Con- ference on Learning Representations, 2023. URL https:/ /ptp///peeeeeeeeee/e/ee/e./oe..e net forum?
1d B92B92G 7xp
Yihe Dong, Jean- Sann-Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference Machine Learning, pp. 2793-2803, PMLR, 2021.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He. Anish Thite. Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint (XXv:2101.00007 2020.
Jonas Geiping and Tom Goldstein. Cramming: Training a language model on a single GPU in one day, 2023. URL httts::///ppeeeeee
ew. net. wwrm/formmmmmmmm 2=t.1dd.6.0...22222233 aer.
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. 369-579, Red Hook, NY, USA, 2018. Curran Associates Inc.
Soufiane Hayou and Greg Yang. Width and depth limits commute in residual networks. arXiv preprint arXiv:2302.0045; 2023.
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.
Bobby He, James Martens, Guodong Zhang, Aleksandar Botev, Andrew Brock, Samuel L Smith, and Yee Whye Teh. Deep transformers without shorteuts: Modifying self-attention for faithful signal propagation. In The Eleventh International Conference on Learning Representations, 2023.
URL httpps::/ net wtt.:/ommmiii.iPeeeses
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. /5556 2022.
Wei Hu. Lechao Xiao, and Jeffrey Pennington. Provable benefit of orthogonal initialization in op- timizing deep linear networks. In International Conference on Learning Representations, 2020. URL
URI.htttps::/pppereeee http:::///ppeeeeeeeeee wtt.ee/wwwe.eeeeeooooooooorooooa

---

Published as a conference paper at ICLR 2024
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.), Proceddiggs of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learn- ing Research, pp. 4475 4483. PMLR, 13-18 Jul 2020. URL https: /ttpss//proceeeee ngs. mml..........................mm. .mlr.
pTeas/v119/huang200 html
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mnss: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165. PMIR, 2020.
Jaehoon 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. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 5747-5763, 2020.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization.
arXiv preprint
1rXiv:1711.05101, 2017.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, pp. 2408-2417. PMLR, 2015.
James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz. Swirszez, Valentin Dalibard, Jascha Sohl-Dickstein, and Samuel S Schoenholz. Rapid training of deep neural networks without skip connections or normalization layers using deep kernel shaping. arXiv preprint 10.00221 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 on Learning Representations, volume 4, 2018.
Alexandru Meterez, Amir Joudaki, Francesco Orabona, Alexander Immer, Gunnar Rätsch, and Hadi Daneshmand. Towards training without depth limits: Batch normalization without gradient ex- plosion. arXiv preprint arXiv:2310.02012. 2023.
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 wrXiv:206.003166 2022.
Lorenzo Noci, Chuning Li. Mufan Bill Li, Bobby He, Thomas Hofmann, Chris Maddison, and Daniel Roy. The shaped transformer: Attention models in the infinite depth-and and-width limit. arXiv preprint 2ttt:::2306.7777555 2023.
Telmo Pessoa Pires. António V Lopes, Yannick Assogba, and Hendra Setiawan. One wide feedfor- ward is all you need. arXiv preprint arXiv:2309.01826 2023.
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 (eds.), Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. 2019.
Christian Sarofeen, Piotr Bialecki, Jie Jiang, Kevin Stephano, Masaki Kozuki, Introducing nyFuser, deep learning URL attps::/pyttorch.or intrroducing-mnfuserree -eteer-a-deep--e lea htt.r//rrrr..re...eoooooooooooor
Neal Vaidya, and Stas. Bekman.
compiler for PyTorch. 2022.
12

---

Published as a conference paper at ICLR 2024
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynam- ics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers secretly fast weight programmers. In International Conference on Machine Learning, pp. 9355 9366. PMLR, 2021.
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations, 2017.
Sharath Nittur Sridhar, Anthony Sarah, and Sairam Sundaresan. Trimbert: Tailoring bert for trade- offs. arXiv preprint 2rXiv:2202.1244: 2022.
Aleksandar Stanié, Dylan Ashley, Oleg Serikov, Louis Kirsch, Francesco Faccio, Jürgen Schmidhu- ber, Thomas Hofmann, and Imanol Schlag. The languini kitchen: Enabling language modelling research different scales of compute. arXiv preprint a 1997, 2023.
Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Herve Jegou, and Armand Joulin. Aug- menting self-attention with persistent memory. arXiv preprint arXiv: 1ete:0:07007779 2019.
Wojciech Tarnowski, Piotr Warchot, Stanistaw Jastrzebski, Jacek Tabor, and Maciej Nowak. Dy- namical isometry is achieved in residual networks in universal way for any activation function. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 2221-2230.
PMLR, 2019.
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In Proceedings of the IEEE/CVE International Conference on Computer Vision, pp. 32-42, 2021.
Hugo Touvron, Thibaut Lavril, Gautier Izacard. Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
Asher Trockman and Zico Kolter. Mimetic initialization of self-attention layers. arXiv preprint 09880 2023.
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- tion processing systems, 30, 2017.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations, 2019. URL patttss:/oopereviies
Ben Wang and Aran Komatsuzaki. GPT-J-6B: Billion Parameter Autoregressive Language Model. https:://gtthiii com www.kkingoollloooooooooo..oo.ooe wtat:///mmmttttttooooooooeeeeeee jax May 2021.
Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. Deepnet: Scaling transformers 1,000 layers. arXiv preprint arXiv:2203.00553 2022a.
Peihao Wang, Wenqing Zheng, Tianlong Chen, and Zhangyang Wang. Anti-oversmoothing in deep vision transformers via the fourier domain analysis: From theory practice. arXiv preprint arXiv:2203.05962, 2022b.
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, 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 Machine Learning, pp. PMLR, 2018.
Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainability and general- ization deep neural networks. In International Conference on Machine Learning, pp. 10462- 10472. PMLR, 2020.
13

---

Published as a conference paper at ICLR 2024
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 Anmual Meeting of the Association for Computational Linguistics, pp. 397-402, 2020.
Greg Yang. Wide feedforward or 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.
Greg Yang. Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. A mean field theory of batch normalization. In International Conference on Learning Representa- fions, 2019. URL /tRp:/ttpp../ppeeeeeeee... http://wwww.eetttttiiieeeeeeeeee (o
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Baachez-GGnzaler Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via de- noising for molecular property prediction. In The Eleventh International Conference on Learning Representations, 2023. URL https: /ttp::::///pereeiiiiii wee.ew.wweeeeee..e..eeeeeee.e.ee www.:/rruteotuuu/e Petetitttt-tee wtwp:ummiiiitttttttteeeeeee
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, Alcksandar Botev, and James Martens. Deep learning without shortcuts: Shaping the kernel with tailored rectifiers. In International Conference Learning Representations, 2022.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. In International Conference on Learning Representations, 2018.
Guangxiang Zhao, Xu Sun, Jingjig Xu, Zhiyuan Zhang, and Liangchen Luo. Muse: Parallel multi- scale attention for sequence sequence learning. arXiv preprint vXXv:1911.09483 2019.
14

---

Published as a conference paper at ICLR 2024
A DUALITY BETWEEN DOWNWEIGHTED RESIDUALS AND RESTRICTING UPDATES IN LINEAR LAYERS
In Sec. we motivated reparameterisation of the value and projection parameters, Eq. (6) through a duality between downweighted residuals branches and restricting parameter updates (ma- terialised through smaller learning rates) in linear layers. This is relatively simple argument, found elsewhere in the literature e.g. Ding et al. (2033), which we outline here for completeness.
We suppose have (differentiable) loss function L(W). which is function of some parameter matrix W. We consider taking a gradient step to minimise L, with learning rate 1/w' from initialisa- tion Wo. This would give new parameters Wi:
dL
W1 = Wo
(8)
dW
W-=
Now suppose we have a reparameterisation of the parameters W with the same loss L(W') as before:
W = = BV for fixed scalar B. fixed matrix U and trainable parameter matrix V. We let V be initialised to Vo. satisfying Wo U + 3V--
(9)
If we take gradient step in V with learning rate y, then we get new parameters V1:
dL
V = V/
T
dV
V=
dW"
dL
Vo
T dV
dW'
V=Vo
W'=Wo
dL
Vo
yu dW
3
W'=Wa
dL
Vo
Ju 3
(10)
dW
W=W,
where the last line we just relabelled the reparameterisation variable W to W.
Feeding Eq. (1Op back into Eq. (9.
obtain:
W = U + +
= U + +VV% πγ nV 82 dL
dW l-=oo
Wo пu B2 dT.
(11)
dW W-Wo
due to the equivalence of initialisations.
To match W7, Eq. [1]. with Wi, Eq. 8
require:
n/ 32
Thus, any gradient step we take in the reparameterisation, W' = = 3VV corresponds to taking the same gradient step original parameterisation, W, but with a learning rate scaled by B2, If B < 1, as is the case in Pre-LN residual branches, this corresponds to downscaling the learning rate.
In the context of our reparameterisation values and projection parameters Eq. 6 this is then equivalent to downscaling the learning rates of WV W by 32 82. if using (stochastic) gradi- ent descent. With Adam W (Loshchilov & Hutter, 2017), one factor of 3 gets divided out by the preconditioner, the reparameterisation acts as if scale the learning rate by not B2.
15

---

Published as a conference paper at ICLR 2024
To verify this theoretical duality empirically, plot the equivalent of Fig. but where, instead of reparameterisation (Eq. (6) with varied B, we reduce the learning rate of the value and projection parameters (keeping the learning rate of other parameters fixed). As expected, the ratio of learning rate for value/projection parameters compared to other parameters improves the training speed, just like the downweighted residual reparametersiation (Fig.)
that reducing
1.50
Eval Loss after 40K steps
1.45
1.40
1.35
1.30
Id WV W.P
Orth W WP
1.20
Pre-LN loss
V-Skiplnit loss
1.15
Ratio of LRs for WV, WP vs. other params
Figure 9: Equivalent of Fig- that empirically confirms the duality between downweighted residuals and reduced learning rates.
B BLOCK LAYOUTS
In Fig. 10 and Fig. 11 we show the layouts of our SAS block (Sec. 4.2 and parallel SAS-P block
(Sec. 4.3p. These are the equivalent plots to the layouts in Fig [
Mathematically, our SAS attention sub-block computes (in the notation of Eq. 2
Xour MHA(Nomm (X))), where MHA(X) Concat (Attn (X)...., Attn (X)),
(12)
1 Xww.X X XT+M
Attn
and Ax(X) SM
Vdl
(13)
Here, X # are column blocks of XxRTxd i.e. X-Conca((x X & SM is Softmax.
C ADDITIONAL EXPERIMENTS
In this section, provide additional experiments and ablations top of those provided in the main paper. The experiments in this section are ordered to follow the chronological order of where they referenced (or most relevant) in the main paper.
Linear vs Cosine decay LR schedule Fig. compares linear and cosine decay LR schedule. We that linear decay provides better final performance across both our models and baselines, and use linear decay throughout the rest the paper.
Shaped Attention Value-Skiplnit Fig. 13 explains reasons for using Shaped Attention Eq. (5) (Noci et al. 2023 over the modified attention matrix, al + (X(X), Eq. (4). that was in- troduced by He al. 2023) in Value-Skiplnit. We that Shaped Attention gives small but
16

---

Published as a conference paper at ICLR 2024
MLP Out
NonLin
MLP
A
Norm
H x
Shaped Attention
K
Norm
Figure 10: The SAS block that we obtain the end of Sec. 4.2
H
Shaped Attention
MLP Out
NonLin
MLP In
Norm
Figure 11: The SAS-P block, with normalisation, that we obtain at the end of Sec. 4.3
consistent gain throughout training. The experiments here follow the same training and hyperpa- rameter setup as those in Sec. [4
Sensitivity to MI.P block gain initialisation In Sec, 4.1 we motivated downweighting the initial- isation of trainable MLP block weight 8pp (c.f. Eqs. and in skipless architectures replicate the implicit downweighting mechanism of Pre-LN skips. Fig. shows the sensitivity of final loss our initialisation for trainable 9pr.
Figure with tied orthogonals In Fig. we observed that restricting the updates to value and projection parameters recovers nearly all of the lost training speed in Transformer blocks without attention sub-block skips. This phenomenon occurred for both random orthogonal and identity initialisations WV' WP
init inir. but identity initialisation outperformed orthogonal, which may be bit surprising as they should be identical from a signal propagation perspective.
To investigate this, consider two alternatives:
17

---

Published as a conference paper at ICLR 2024
Pre-LN, cosine
Eval Eval Cross entropy Loss
4.0
Pre-LN, linear
Cross-entropy Loss
Pre-LN, linear
3.5
SAS, casine
SAS, linear
1.40
SAS, linear
SAS, cosine
3.0
SAS-P, cosine
1.35
SAS-P, cosine
SAS-P, linear
SAS-P, linear
1.30
Eval
14K
40K
35K
Training step
Training step
Comparing training performance with cosine and linear decay LR schedulers on Code- Parrot. The right plot is a zoomed-in version of the left. We see that linear decay consistently provides a better final performance than cosine decay, despite trailing for most of the steps towards the end of training.
Figure 12:
4.0
SAS-P, V-Skiplnit attn
Eval Cross-entropy Loss
3.5
SAS-P, shaped attn
SAS, V-Skiplnit attn
3.0
SAS, shaped attn
2.5
2.0
1.5
1.0
10K
20K
30K
40K
Training step
Figure 13: Shaped Attention (dashed lines) provides a small performance boost compared to the attention matrix of Value-Skiplnit (solid lines), for both SAS and SAS-P blocks. All transformers are Layer autoregressive GPT models, and the dataset is CodeParrot.
18

---

Published as a conference paper at ICLR 2024
Eval Loss after 40K steps
1.25
1.24
1.23
1.22
1.21
1.20
1.19
0.0
0.5
2.0
2.5
3.0
Initialisation for MLP residual weights BFF
Figure 14: Final test loss achieved as function of the initialisation for trainable MLP block gains OFF on CodeParrot.
1. "LL-tied" or last-layer tied: this is when we initialise all but the final layer projection ma- trix independent random orthogonals. Then for the final layer projection matrix WPP (or rather its transpose), we tie the initialisation to the previous layers, as:
init,
WPP L
II
w W.P init.!
WV L.
init,
The purpose of this so that the combined product over all layers I WV WP the identity, which mimicks the functional output the whole transformer with identity WV init,l I = WP init, (up to the MLP blocks, but we note that the MLP weights independently Gaussian initialised and so are rotationally invariant at initialisation, and are also downweighted as they lie on a downweighted residual branch).
innt. mit. is
2. "All-tied": this is when for each attention layer/sub-block / we have W.Y
init./ W init.7 with random orthogonal initialisation, which makes the value-projection product identity, WV. WP exactly if we had identity values and projections initialisation. This initialisation is similar to that of Trockman & Kolter (2023). although here we are considering skipless attention sub-blocks
init, init, I VI < L, and hence matches the outputs of each attention sub-block
Fig. is the equivalent of Fig. but with LL-tied (yellow line with diamond markers) and all-tied (blue line with star markers) included. In Fig. [5] we see that matching the functional output (as LL tied) with orthogonal initialisation provides slight improvement to close the gap between the random orthogonal (green line) and identity (purple line) initialisations. Matching the attention sub-block outputs (as in all-tied) further improves performance, but does not fully close the gap to identity initialisation. Interestingly. it seems like orthogonally initialised values and projections do benefit from being trainable (i.e. small but non-zero By, Pp). We leave further exploration of these observations to future work.
Further scalar parameter trajectories In Fig. we saw that residual-skip gain ratios on the values and projections By 8p (from the reparameterisation in Eq. (6)) converge to zero during o training for the vast majority of layers, in models without attention sub-block skip connections. In Fig. 16 plot the corresponding plot Fig. but for a parallel block with no skip connections (i.e. SAS-P with reparameterisation Eq. (6) and trainable values and projections. Like before, we also initialise trainable Bv, 3p to 0.2, and av to 1). Again,
that the vast majority of
19

---

Published as a conference paper at ICLR 2024
1.50
Eval Loss after 40K steps
1.45
1.40
1.35
1.30
Orth W W
Orth W W LL tied
1.25
Orth W W tied
Identity Wi., W
Pre-LN loss
V-Skiplnit loss
1.15
0.2
0.4
0.6
0.8
1.0
Residual weights Bv, Bp
Figure 15: Equivalent of Fig. 3 but with tied orthogonal initialisations.
ratios converge to 0, indicating that the value and projection parameters converge to the identity. Ålso, again
that the first value matrix is the only ratio above 0.05 the end of training.
Figs. 17 to 20 plot the corresponding trajectories, the model in Fig. for various other train- able scalar parameters we have: namely, the three shaped attention parameters
the identity
(Fig. 17). 2) on the softmax attention output (Fig. [8). 3) on the centring matrix (Fig. (Fig.T9). as well as 4) OFF on the MLP block (Fig. 200. We see that none of these other scalar parameters converge to zero. The shaped attention parameters have error bars denoting standard deviations across (the 12) heads.
In Fig. 21 we plot the equivalent of Fig. but for an attention skipless model with random orthog- onally initialiased (from the Haar measure) value and projection weights. Again, we see that the vast majority of residual/skip gain ratios converge from their initial value of 0.2 to O during training albeit with slightly more outliers than in Fig.
However, in Fig. 22 we plot the equivalent of Figs. and 21 but for a Pre-LN model that has attention skip connection (and identity initialised values/projections). Now, we see that the introduction of the skip connection encourages many more residual/skip gain ratios in the value/projection weights to increase during training (notice the y-axis), i.e. encouraging the values/projections to leave their identity initialisation. Together, these results reaffirm the complex interactions between skip nections and value/projection weights in the standard transformer architecture highlighted by our work.
observed that moving values and projection parameters by setting them to the identity improves the convergence speed per parameter update of transformer blocks with skipless attention sub-blocks. This raises the question of whether the same would occur in the standard block that uses attention sub-blocks skip connections i.e. the Pre-LN block. In Fig. 23 we compare the default Pre-LN block one with values and projections set to be identity. We see that in this case identity values and projec- tions actually slightly hurt performance in terms of loss reduction per update, in constrast to the skipless case. We also tried to downweight the attention block residual (3sA < in the notation of Eq. with identity values and projections to see if the scale of the attention skip and residual was the reason for this difference, but this did not change the findings of Fig. 23 We do not have satisfying explanation for this, but our intuition is that identity value and projections (as opposed to e.g. Gaussian/random orthogonal initialisation) with attention skip means that the two branches of the skip are no longer independent at initialisation and interfere with each other, though it is unclear that this would continue hold during training.
Identity values and projections with default Pre-LN block In Sec. 4.2
20

---

Published as a conference paper at ICLR 2024
Residual/Skip gain ratios, B/a
0.20
Layer
1
11
2
12
17
0.15
4
14
Proj Brlap
Value By/av
0.05
0.00
5K
10K
15K
20K
25K
30K
35K
40K
Training step
Figure 16: Corresponding plot to Fig.
but for parallel block with no skips.
2.0
Layer
Shaped Attn skip gain, a
1.8
1
6
11
15
16
1.6
3
8
13
17
1.4
9
14
18
10
0
5K
10K
15K
20K
25K
30K
35K
40K
Training step
Figure 17: Trajectories for shaped attention parameter.
21

---

Published as a conference paper at ICLR 2024
2.0
Shaped Attn residual gain, B
Layer
1.8
6
11
15
7
12
16
1.6
8
13
17
1.4
9
14
18
10
0.4
0
5K
10K
15K
20K
25K
30K
35K
40K
Training step
Figure 18: Trajectories for shaped attention | parameter.
2.0
Layer
1.8
6
11
15
Shaped Attn centre gains,
7
12
16
1.6
8
13
17
1.4
9
14
18
10
0
5K
10K
15K
20K
25K
30K
35K
40K
Training step
Figure 19: Trajectories for shaped attention parameter.
22

---

Published as a conference paper at ICLR 2024
0.5
Layer
6
11
15
MLP residual gains, BFF
7
12
16
8
13
17
14
18
0.3
10
0.1
0.0
0
5K
10K
15K
20K
25K
30K
35K
40K
Training step
Figure 20: Trajectories for MLP block 9FF parameter.
Skipless with orthogonal w/ Wht WP
init
Residual/Skip gain ratios, B/a
0.20
Layer
1
11
16
12
0.15
3
13
18
4
14
Proj Balap
5
10
15
Value By/av
0.10
0.05
0.00
5K
10K
15K
20K
25K
30K
35K
40K
Training step
Figure 21: Trajectories for MLP block 3F parameter.
23

---

Published as a conference paper at ICLR 2024
Pre-LN (with skip)
1.0
Residual/Skip gain ratios, B/a
Layer
11
16
0.8
7
12
17
13
0.6
9
14
Proj Ba/ap
10
Value Bylav
0.4
0.2
0.0
0
5K
10K
15K
20K
25K
30K
35K
40K
Training step
Figure 22: Trajectories for MLP block 3r parameter.
WV, WP =1, Pre-LN
Eval Cross-entropy Loss
W", wP #1, Pre-LN
2.00
1.50
1.25
10K
20K
30K
40K
Training step
Figure 23: Pre-LN block performs worse when setting values and projections to identity, unlike in the skipless setting.
24

---

Published as a conference paper at ICLR 2024
Using the first value matrix In Figs. and 16 we that the vast majority of value and projection parameters stay close to the identity during training when initialised to identity, even when they have away from initialisation. The first layer value matrix WY is an exception to this rule. Fig. 24 we see that allowing the first layer value parameters to be trainable provides a very small boost training performance, when all other value and projection weights the identity. Intuitively it makes sense that the first layer value parameters would be more important than others because they directly on the input embedding at the beginning of the model. We thus choose to reincorporate trainable value parameters (using identity initialisation in Eg. (6) in the first
the capacity
fixed to
layer of our models using SAS and SAS-P blocks, but remove all other values W for > and all projections too W 4I 2 by fixing to the identity.
3.50
SAS-P, W -/
SAS-P, W/
Eval Cross-cns entropy Loss
SAS-P
SAS-P
3.5
Wy
SAS, W/ -/
SAS
2.0
2.5
2.0
1.20
1.5
13K
10K
A0K
20K
25K
13K
35K
40K
Training step
Training step
Figure 24: Comparing training performance with trainable first layer values WY with identity first layer values WY for models using our SAS and SAS-P blocks. All other values and projections are set to the identity. The right plot zoomed-in version of the left. We see that having trainable first layer values provides a (very small) boost in performance.
Linearising MLP activations As stated Sec. we tried to use the recent idea of "linearising" activation functions in order to obtain better signal propagation (Martens et al. 2021 Zhang al.) 2022 Li et al. 2022) in deep NNs with skip connections, and recover lost training speed when the MLP skip is removed. In particular, Li et al. (2022) show for Leaky ReLU,
LReLU((x) max Sx),
with negative slope E [0, 1]. we need = 1 O 1 in MIPs large depths T..
VT. to obtain well behaved signal propagation
In Fig. 25 we took our 18-block model trained with SAS block (Fig. 10 and assessed training
performance without the MLP skip, OpF = 0. We tried 3 different activations: standard ReLU, LReLU with slope s = 0.2, and 3) LReLU with 0.8 z
VI8
We see that all blocks without MLP skip train significantly slower than our SAS block (which matches the training speed of the Pre-LN block). fact, linearising ReLU into LReLU seemed to hurt training speed rather than help it. These findings consistent with those of previous works with AdamW optimiser (Martens et al. 2021 Zhang al. 2022 He et al., (023). We note that
big reason behind this is that the architectures with skipless MLP sub-blocks required order of magnitude smaller learning rate (1e - 4 vs le - otherwise training was unstable.
Loss vs training step In Fig. 26 we provide the equivalent plot to Fig, 5 but in terms of loss over the steps taken. Our SAS and SAS-P essentially match the Pre-LN model in terms of loss reduction per step, whilst removing normalisation slightly hurts performance.
Crammed Bert loss training step In Fig. 27 plot the MLM loss in the Crammed Bert setting on the Pile dataset, as function of the number of microbatch steps taken. Because our models have higher throughput they able take more steps within the 24 hour allotted time.
25

---

Published as a conference paper at ICLR 2024
4.0
no MLP skip, LReLU slope=0.2 no MLP skip, LReLU slope=0.8 no MLP skip, ReLU with MLP skip
Eval Cross-entropy Loss
3.5
3.0
2.5
2.0
1.5
10K
20K
30K
40K
Training step
Figure 25: Removing MLP skips results in significant losses of training speed, even when linearising activations.
3.00
V-Skiplnit (He et al. 2023)
2.75
Pre-LN
Eval Cross-enti Pross-entropy
SAS ($4.2)
2.50
SAS-P (54.3)
2.25
SAS-P, no norm ($4.4)
Parallel (Wang et al, 2021)
2.00
1.75
1.50
1.25
10K
20K
30K
40K
Training step
Figure 26: Equivalent of Fig. but with steps on the x-axis.
GLUE breakdown In Table 2 we provide a breakdown of the GLUE results in Table | in terms of different tasks.
Autoregressive Language Modelling We investigate if our findings hold in the next-token pre- diction language modelling domain. This also allows us to test larger sequence lengths (512) than other experiments in this work (128). The task is autoregressive language modelling on the Languini Benchmark books dataset (Stanić et al. 2023)). and we use the same codebase and tokeniser provided by the authors. Models have 12 layers with width 768, which gives 100M parameters (including tied embeddnng/nemmeeddingg) by default when MLP width 3072 (4x 768). Sequence length is 512 and train for 19K steps batch size 128, giving 1.2B training tokens. Learning rate is lin- early warmed up for 500 steps to maximum value that is tuned for models separately (3e-3 for simplified models, 1e-3 for the default blocks), before linear decay. ALiBi positional encoding
26

---

Published as a conference paper at ICLR 2024
Crammed BERT (Pre-LN)
Parallel (Wang et al, 2021)
SAS (54.2)
SAS-P ($4.3)
SAS-P, no norm ($4.4)
MLM loss
V-Skiplnit (He al. 2023)
3.5
3.0
2.5
100K
200K
300K
400K
500K
600K
Microbatch step
Figure 27: MLM loss vs microbatch steps taken. Note that because the LR schedule depends on the total number of steps taken by model, and is different in different models for the same number of steps taken, comparing models in terms of MLM loss at a given step is not so informative.
Table 2: Breakdown of GLUE results different tasks. Results are the mean seeds.
GLUE
MNLI
SST-2
STSB
RTE
QNLI
QQP
MRPC
CoL.A
Pre-LN (Crammed)
78.9
81,2/81.6
89.9
87.6
56.3
88.9
87.1
88.3
49.4
Parallel
78.5
80.8/80.9
91.0
87.5
54.9
88.0
87.0
87.3
49.0
SAS (Sec.
4.2
7
79.7/80.1
90.3
84.7
58.4
87.5
86.8
87.5
50.6
SAS-P (Sec.
4.3
78.3
79.5/79.4
90.8
85.2
59.1
87.5
86.5
86.0
50.5
V-SkipInit
78.0
79.7/80.4
90.4
85.1
54.4
87.0
86.2
87.9
51.5
and GcLU activations used. Training takes place a single RTX-2080Ti (with microbatches of size 16), and we use AdamW with weight decay 0.1.
We plot training speed comparison of simplified blocks against default in Fig. 28 below, in terms of runtime on the x-axis and evaluation perplexity on the y-axis. We again see that our models able match the training speed of default Pre-LN and Parallel blocks. Moreover, (orange curve with star markers) achieves the same final perplexity after 19K steps (22.37 vs 22.39) the Parallel block (pink curve with diamond markers) despite using 15% fewer parameters.
SAS block
D IMPLEMENTATION DETAILS
In this section add remaining implementation details that were not discussed in the main paper. We break down our implementation details into two subsections, one for the next-token prediction task on CodeParrot and one for Crammed BERT (Geiping & Goldstein 2023) masked language modelling experiments pretrained on the Pile dataset (Gao et al. (202)) and fine-tuned to downstream GLUE benchmark (Wang et al.] 2019). To avoid repetition, any details that mentioned in one subsection but not the other are shared between both subsections. All runtime results on CodeParrot were run on single A5000 GPU.
D. CODEPARROT NEXT-TOKEN PREDICTION
As mentioned, much of our setup derived from http::/hhggingffacccc///ee n nlp-ccoure//caateer
27

---

Published as a conference paper at ICLR 2024
SAS-P
50
SAS
Pre-LN
45
Eval Perplexity
Parallel
40
25
2
10
12
Runtime (hours)
Figure 28: Eval perplexity vs runtime on an autoregressive language modelling task.
Model The model is a 18-layer GPT-stylc auto-regressive decoder-only transformer. We use width 768, and H 12 heads in multi-head attention. We remove dropout entirely as our focus is on training speed, and we are always in single-epoch regime regularisation hurts training speed. The MLP uses ReLU activation unless stated otherwise, and we use MLP hidden dimension 3072 4d. The only exception to this is in Fig. 6 where reduce the MLP hidden dimension to 1536 = 2d to account for the increased memory requirements of larger depths.
For any of our simplified model we initialise SFF 0.1 in Eq. to account for the lack of skip, apart from the 18-layer models in Fig. 6 where 3FF = 0.2 due to the narrower width.
We use RMSNorm (Zhang & Sennrich 2019) where applicable with epsilon le 8, and add final normalisation after the decoder. Sinusoidal positional encodings are used and added at the embedding level.
Parameter Initialisation For the Pre-LN and parallel blocks, we initialise all weights (including the embedding layer) to have standard deviation 0.02 is standard in GPT2 (Radford et al.) 2019) and BERT (Devlin et al. 2018). This a choice that prevalent the field, and is on the order of O for d = 768, that would expect from signal propagation.
For our models, we always initialise W = 0 like (He et al. 2023). which as discussed zeros the query-key dot product, and allows shaped attention to have a dominant identity component at initialisation. Also as discussed, we initialise trainable scalar parameters in shaped attention to 1 for simplicity; the same applies for the av By we use in the first layer value parameters WY. All other scalar parameters in the attention and MLP branches ASA, fFF (initialised to and 0.1 resp.) are also trainable our models, which we found give small boost in performance.
Training We use AdamW optimiser (Loshchilov & Hutter 2017) with weight decay 0.1 which we tuned on small grid, and found to work well for both baselines and our models. We not apply weight decay to any scalar gain parameter. We clip gradients with with clipping parameter 1, and use epsilon le 8 and default betas (0.9, 0.999) in AdamW. As discussed, we use linear decay rate with 5% of all steps used for linear warmup. The optimal learning rate was tuned in all cases, and for our best (SAS and SAS-P) models, was found to be le - 3, which exactly matched that of the default Pre-LN. This held true also when scaled to 72 layers. V-Skiplnit needed lower learning rate for the depth scaling experiments (3e 4 and le 4 for depths and 72 respectively). We use batch size of 128 with microbatches of size 32.
28

---

Published as a conference paper at ICLR 2024
Dataset The Codeparrot dataset is large corpus of 20 million python files from GitHub. We take the dataset, pre- pre-prseessing and tokeniser from https: /bwgggiggaacceeeec/////p/a co co/learn/ htt--courseehhappeereeee
We use sequence length T 128 throughout, and our tokeniser has 50K vocabulary size. Our base experiments train for around 43K steps on batch size 128 and sequence length 128 which is around 700M tokens. In Fig- g we scale this to 2B tokens.
Task
The model is trained on next-token prediction using cross-entropy loss.
D.2 BERT ENCODER-ONLY
As discussed in Sec. we inherit much our hyperparameters from the Cramming setup of Geiping & Goldstein 2023), and also base our implementation from their excellent codebase We highlight important implementation details here.
Model We use 16-layer encoder only model, with width 768 and 12 heads. We use MLP width 3072 = 4d, but now we use GLU (Dauphin et al. (017) with GeLU activation, which essen- tially halves the hidden dimension. We use LayerNorm (Ba al., 2016) for normalisation where applicable with epsilon le - 12 as taken from Geiping & Goldstein (2023); we always use a final LN after all layers. Again, we remove all dropout, and use sequence length of 128. We found our simplified skipless models prefered smaller MIP block scales and initialise fTr = 0.05.
Parameter Initialisation The initialisations are identical to those in Codeparrot, and are detailed above.
Datasets Like Geiping & Goldstein (2023). we train on the Pile dataset (Gao et al. 2020 with WordPiece tokeniser of vocabulary size 32768, and a sequence length of 128. Our fastest runs took around 600K steps with microbatch size 64 in 24 hours, which corresponds to around 5B tokens.
Training We again trained with AdamW optimiser, with weight decay 0.1. AdamW had hypa- rameters [81, [2] [0.9, 0.98). and epsilon le 12. We used microbatch of 64 (to fit on RTX-2080Ti), and scale the batch size to reach 8192 linearly after 60% of total training like in Geiping & Goldstein (2023). We the same aggressive learning rate as Geiping & Goldstein (2023), which increase linearly to max value after 75% of all training steps, before linear decay, and tune the maximum learning rate 3e 3 for our SAS and SAS-P models. This was slightly too large for the SAS-P model without normalisation, so we reduce to 2e- 3. We inherit the clipping parameter of 0.5 from Geiping & Goldstein (2023).
Fine-tuning We followed the same protocol Geiping & Goldstein 2023). In particular, we fine- tune for epochs with fixed hyperparameters across tasks. We found dropout to be important for good downstream performane (unlike during pre-trat trainiann and set dropout probability p 0.1. We use batch size 32, with a maximum learning of 1.5e 4. We keep other hyperparameters, e.g. the choice of cosine decay and AdamW epsilon 1e 6, like from Geiping & Goldstein 2023)
Task The model is trained on the masked language modelling task with masking probability 0.25,
as in Geiping & Goldstein (2023)
29


Reviews:
[{"review_id": "YusjJT16o2", "reviewer": "Reviewer_iboi", "summary": "This work presents a thorough investigation into the complexity of standard transformer blocks, questioning the necessity of various components commonly included in their design. Through a combination of signal propagation theory and empirical observations, the authors propose a simplified transformer architecture that maintains performance levels while offering benefits in terms of training speed and model parameter count.\nThe paper challenges the conventional wisdom of transformer design by systematically evaluating the impact of removing certain elements, such as skip connections, projection or value parameters, sequential sub-blocks, and normalization layers. The experimental results, as reported, demonstrate that these simplifications do not detrimentally affect the training speed and can lead to a 15% increase in training throughput and a similar reduction in parameter count.", "questions": "- The authors mention that there is a Pre-LN and a Post-LN block and the Pre-LN work is more popular because Post-LN suffers from training instability. Is this phenomenon observed in the signal propagation literature only or all the current LM works also use this Pre-LN variant?\n- The authors mention “our simplified transformers match the per-update training speed and performance of standard transformers, while enjoying 15% faster training throughput”. If the parameters are being reduced, why isn’t the per update speed faster? And, why is the throughput faster?\n- On page 3, the authors mention skipless transformers are slower than those with skip connections. Why is that?\n- Are there any settings (data or tasks) where this simplified transformer does not work as well compared to the standard transformer?\n- What are the limitations of the proposed simplifications in terms of model expressiveness and capability?\n- What is the size of the model which has 18 blocks and 768 hidden dimensions?\n- In section 4,1 what is the motivation from going from SkipInit to ShapedAttention? Another question is since this work is adding identity addition to value and projection matrices, in what way this is simpler than the residual connection?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- As the work mentions, this is the first work which has simplified the transformer architecture with training throughput gains. Previous works which have tried to simplify transformer architectures have led to increase in training speeds.\n- The paper is overall well written. The authors discuss each design choice in detail, which helps in understanding the motivation behind the simplifications proposed.\n- The authors have appropriately discussed the limitations of their work.", "weaknesses": "- One weakness is that the results are limited to smaller models of size ~100M. It is not clear if the results would scale to bigger models. \n- Moreover, the authors only consider GLUE benchmark for downstream evaluation. It would be nice to include a more comprehensive evaluation across different task categories like classification, generation or reasoning.", "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": "", "extra_scores": {}}, {"review_id": "cWQ1QZVrnY", "reviewer": "Reviewer_W1bM", "summary": "This paper studies how to simplify the self-attention block in transformer models. Some modifications are made based on several interesting and empirical findings:\n- This first is that, while somewhat similar to prior work, this paper points out that the skip connection in the self-attention subblock could be removed, without loss of training speed.\n- The second is that the value and projection parameters in the sub-attention blocks could be fixed to the identity matrix, which reduces two parameterized matrices.\n- Third is that the MLP sub-block skip connection could also be removed in the parallel MHA and MLP structure.\n- The forth is that all normalization layers could be removed given that the proposed block has already taken the effect of normalization into account (to some extent?).", "questions": "Depite the weaknesses, I generally like this paper as it does provide some deep insights of the self-attention mechanism. I also have a few additional questions that expect the authors to reply:\n- The simplification is mostly carried out in the context of skipless self-attention block. Why the block is preferred to be skipless? Is there any benefits? From the view of signal propagation theory, the skip connection benefits propagation, particularly for deeper models, and does not violate the philosophy of simplicity,.\n- From Fig. 19, the pre-LN block performs worse when setting values and projections to identity. An intuition is provided, but I am not totally convinced by this intuition. In fact, I see the performance difference is almost negligible. It seems OK to set the values and projections to identity in the standard attention block with skip connection. This also reduces the number of parameters.\n- Why is the performance SAS-P, no norm not reported? Is the training failed or the performance is significantly lower? In fact, I would suggest that the block can preserve the normalization layers to keep the consistency between analyses and experiments.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Some interesting findings revealing the underlying working mechanism of the attention sub-block in the self-attention is reported. These findings, despite being empirical, also lead to a simplified self-attention block. It provides valuable insignts on how to design an efficient self-attention block. The community would like to see it.\n- The paper did a good job explaining the step-by-step simplification flow and have provided sufficient reasons/experiments to discuss why.\n- The paper is very well written, which is a pleasant read.", "weaknesses": "- Some critical designs in the proposed block are from existing literatures, e.g., the shaped attention (Noci et al. 2023) and the parallel structure (Wang & Komatsuzaki, 2021) of MHA and MLP sub-blocks.\n- While the block is simplified, the benefits of the simplification is not significant. The simplified transformers only match the per-update training speed and performance of standard transformers. While being 15% faster and having 15% fewer parameters, some part of the throughput gain is due to the use of the parallel structure by (Wang & Komatsuzaki, 2021).\n- The removal of normalization somewhat contradicts between the analysis and approaches. In Sec. 4.4, it says that the normalization layers are unnecessary, but it quickly comments that the normalization still has benefits and is still used in experiments. This somehow leads to a confusion of read.\n- The performance of the simplified block slightly falls behind the standard Pre-LN and Parallel in Table 1, which means in practice the community may not use the proposed simplified block.", "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": "", "extra_scores": {}}, {"review_id": "jQPIC1bqqn", "reviewer": "Reviewer_wksT", "summary": "This paper undertake a critical study of several components of the transformer block, with aim of simplifying the block. The authors utilize shaped attention (identity-preserving attention initialization), and scaling the MLP block to remove the value and projection matrices, and all the skip connections from transformers. The authors do detailed ablations of all variations/modifications and scaling factors, as well as studying the impact of layernorms and skip connections. The changes are studied on pre-training for GPT and BERT models, and downstream finetuning for BERT on GLUE. The proposed changes result in models with comparable/better performance compared to vanilla models, with improved training speed and fewer params.", "questions": "1. By initializing attention $\\beta$ as 1, but MLP $\\beta_{FF}$ as $\\frac{1}{\\sqrt{L}}$, this results in the FFN blocks \"contributing\" $\\frac{1}{\\sqrt{L}}$ as much to the output as compared to the attention block, correct? If yes, Perhaps an interesting experiment could be decreasing the dimensions $4d$ of the MLP block when using the proposed SAS-P method - the motivation being that if the MLP block is not contributing as significantly, perhaps it can use fewer params. This could perhaps lead to significant speedups and parameter savings.\n1. Were any downstream fine-tuning experiments conducted for GPT, like GLUE for BERT? It will help make paper stronger if downstream performance also remains similar/better. \n\nMinor presentation suggestions (the authors need not respond to these) - \n1. In Figure 7, some exponential smoothing should be applied to the MLM loss (or ideally, eval loss over multiple batches should so be plotted which will be hopefully much more stable.) \n1. In Figure 6 , it is difficult to distinguish the plots as many of the lines overlap/intersect significantly. Perhaps some of the lines could be omitted, or separate chart provided in supplementary.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The approach is well grounded on related work (SkipInit/Shaped-attention, residual scaling, etc.), and presents strong natural extension of them.\n1. Completely removing the Value and Projection Matrices and skip connections offers strong simplifications.\n1. The authors method achieves parameter saving and training speedup, while also resulting in simpler blocks. The training speedup could perhaps be further optimized with better implementation/variations.\n1.  The experiments presented are thorough and well-motivated. They cover a wide range of variations/ablations, multiple different models, and multiple tasks.", "weaknesses": "1. The authors acknowledge that their experiments are in the range of smaller models (100M-300M). It will be interesting to see if the gains carry to larger models.\n1. Most of the experiments are conducted on training for only 650M tokens. 650M is a low number of tokens, compared to even \"compute-optimal\" tokens suggested by chinchilla for the model sizes the authors use. It will be important to see if the proposed method falls behind on training for longer. That being said, Figure 8 was trained for 2B, and BERT ones for 5B.\n1. The sequence length is somewhat unorthodox at only 128 - A smaller sequence length will \"weaken\" the ability of attention. Perhaps this could be why value projection matrices could be dropped? It is unclear how the proposed method will perform for longer sequence lengths.\n1. No downstream evaluation of the GPT model.\n\n[1] https://arxiv.org/abs/2203.15556", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)