Simplifying Transformer Blocks

We thank Reviewer wksT for their time in reviewing our paper, and are glad to read that they found our work thorough and well-motivated. We address the raised concerns and questions in order as follows:

1. Scale: We are actively exploring ways to scale our methods up, though it is not straightforward for us as we have relatively few GPU resources available to us. We agree though that testing at larger scales will be interesting, and believe our depth scaling experiments (Figure 6) show promise in this regard.

2. Training time: Thanks for this point. Indeed, the motivation for Figure 8 (as well as our BERT experiments) was to demonstrate that our findings hold at longer training times. We agree that it would be interesting to scale training time further along with model size, as is commonly done in practice e.g. Chinchilla, LLaMa.

3. Sequence Length: Thanks for this point. We have run additional experiments with larger sequence lengths using the recent Languini Benchmark (https://arxiv.org/abs/2309.11197), which provides an open-source codebase and a large corpus of high-quality books data (based on books3 from the Pile) to run autoregressive language modelling pre-training experiments at various scales. In these experiments, we use a sequence length of 512, compared to 128 in our submission, and train on 1.2B tokens. We present evaluation perplexity curves and more experiment details in the anonymous link: https://docs.google.com/document/d/1SF8F3CXBQ_9xqTicGOjKC-UWBTKbZDTiT1zp3gqfD5c, and find that our conclusions carry over to this setting with longer sequence lengths.

4. GPT downstream evaluation: thanks for this point! Would you have any particular downstream benchmarks in mind, particularly for researchers with low compute resources who need to train from scratch (not using pre-trained models)? We note that the Languini Benchmark, designed for GPT language modelling research at various smaller scales, evaluates purely on pre-training performance on its books dataset https://arxiv.org/abs/2309.11197.

5. Downscaled MLP branches: Thanks for this point. We certainly agree that finding an optimal “shape” for our blocks is an interesting direction, and it may be that the conventional MLP widths are not (compute) optimal for our simplified blocks.

   Actually, the motivation behind downscaling the MLP block is primarily to replace the lost “implicit downweighting” effect that occurs in Pre-LN blocks (https://arxiv.org/abs/2002.10444). In this paper, the authors show that each layer $l$ in a Pre-LN architecture actually is already downweighted by a factor of $\frac{1}{\sqrt{l}}$. As our models do not have Pre-LN skips, we need to explicitly downweight the MLP otherwise we will suffer from signal propagation degeneracies at large depths. In “Stable ResNet” https://arxiv.org/abs/2010.12859, the authors show that scaling by $O(1/\sqrt L)$ in order to obtain a well-behaved infinite depth limit in terms of signal propagation theory, which is how we motivate this “stable” square-root scaling. We note two other ICLR 2024 submissions (https://openreview.net/forum?id=KZJehvRKGD and https://openreview.net/forum?id=17pVDnpwwl) have concurrently shown this stable scaling enables hyperparameter transfer in residual architectures scaled to large depths.

   Having said that, inspired by your comment, we ran an additional experiment in our rebuttal experiments https://docs.google.com/document/d/1SF8F3CXBQ_9xqTicGOjKC-UWBTKbZDTiT1zp3gqfD5c to test different MLP widths in different blocks. We find reducing MLP widths (in both default and simplified blocks) is an alternative way to reduce parameter count like our removal of values and projections, but also at a slight loss in per-update convergence speed, which is counteracted through the increase in throughput. However, we note that reduced MLP widths in default blocks does not lead to simplified architectures, in terms of conceptual understanding/interpretability, unlike our proposals.

Thanks again for your time and engagement, and also for the reference, which we will look into incorporating in our work!

5. Skipless attention blocks: Thank you for raising this point. By removing the skip connection, the interpretation of the attention block is much clearer in that it simply adds to each token other tokens according to the attention matrix value, whereas previously there was a skip connection and the values/projection weights act as some unknown rotation/linear transformation. Moreover the signal propagation in the attention block has a new interpretation in terms of balancing how much a token attends to itself relative to others, which is explicitly controlled through the attention matrix. We believe these cleaner interpretations are muddied if there is a skip connection.

   Furthermore, though we motivated simplified blocks for improved fundamental understanding of the roles of different architectural components (like skips) and efficiency gains, there have been other works that suggest specific disadvantages to residual architectures, e.g. more transferable adversarial examples (https://openreview.net/forum?id=BJlRs34Fvr), that may provide independent motivation to consider the skipless setting.

6. Figure 19: Setting the values and projections to be identity in the Pre-LN block can be thought of as being “morally” similar to removing the skip and increasing the diagonal components of the attention matrix. In this sense, we agree with the reviewer that explicitly removing the skip connection appears to not be a necessary simplification in the context of our work. However, this is not quite true in practice because of the normalisation layers acting on the residual branch in Pre-LN. This distinction is important and gives us a preference for the skipless setting for three reasons: First, as described in response to the previous question we believe it is conceptually simpler to interpret a skipless block with identity val/proj compared to a Pre-LN block because the entire sub-block (and signal propagation) is determined through the attention matrix, whereas with the skip you have to juggle the skip branch with the diagonals vs off-diagonals of the attention matrix (as well as the non-linear normalisation layer). Secondly, we note that the Pre-LN model in Figure 19 with identity values and projections achieves worse final validation loss (1.213) than both our SAS and SAS-P models (1.178 and 1.187), which suggests that empirically there is a preference for skipless attention sub-blocks (we offered our intuitive explanation for this in the text accompanying Figure 19, but would be curious if the reviewer had an explanation they found more convincing?). Finally, it is cheaper to avoid implementing a full skip connection if possible, as one does not need to store activations of both branches in memory e.g. Ding et al 2021 https://arxiv.org/abs/2101.03697.

   As an aside, regarding reduced parameter counts in Pre-LN blocks, the reviewer may be interested in an additional experiment we ran for the rebuttal in response to wksT, where we compare reduced widths in the MLP layers for different blocks, which can be found at this anonymous link https://docs.google.com/document/d/1SF8F3CXBQ_9xqTicGOjKC-UWBTKbZDTiT1zp3gqfD5c.

7. SAS-P, no norm in Table 1. As discussed in our submission (first paragraph of page 9), we found that “removing normalisation caused instabilities at fine-tuning time, where a small minority of sequences in some downstream GLUE datasets had NaN values in the initial forward pass from the pre-trained checkpoint”. This is despite the fact that we were able to pre-train SAS-P without normalisation from a random initialisation fine on the Pile (e.g. in Figure 7). We are actively working towards a better understanding fine-tuning in general (and the role of normalisation in fine-tuning), but leave this for the future as our primary focus in this present work is on pre-training speed.

We thank Reviewer W1bM for their time in reviewing our paper, and are pleased to see that they enjoyed reading our work. We address the raised concerns and questions in order, below:

1. Critical designs from existing literature: Though several of the tools (like shaped attention or the parallel block) we build on can be found in existing literature, it is not obvious how one would combine them in order to create simplified blocks. Moreover, we extend these existing works and make specific changes to them that are crucial for our best performance. For example, we use head dependent scalars (in shaped attention), or downweighting the MLP block $O(1/\sqrt{\text{depth}})$ in the parallel structure (which is needed to avoid the signal propagation degeneracies of the parallel block if skips are not present). Moreover, neither previous work considers removing residual connections or value/projection parameters.

   We tend to agree with reviewer wksT on this point, who writes that “the approach is well grounded on related work (SkipInit/Shaped-attention, residual scaling, etc.), and presents strong natural extension of them.”

2. Significance of simplifications: Table 1 shows that in our implementation our parallel block only gains in 5% throughput, so 10% (out of 15%) of the throughput gain is through our further simplifications. In the era of large language models with exorbitant training and deployment costs, even 1 or 2 percent of throughput savings is significant.

   Moreover, one of our key motivations for simplifications was to help close the gap between theory and practice in DL. Currently, theorists working in DL work towards understanding simplified architectures (due to convenience) and run the risk of missing out on crucial details that influence performance in practice as a result. Demonstrating that conceptually simpler architectures, with many unnecessary components removed e.g. without skip connections or values or projection parameters, can match performance of standard architectures is significant in our eyes, as it enables theorists to focus on the key factors that dictate performance in practice. Hopefully, an improved understanding acquired through future theory can then feed back into motivating improved methods to train and finetune NNs in practice.

3. Clarity over removing Normalisation: Thank you for raising this point. We made a conscious effort to stress that in section 4.4, the redundancy of normalisation layers is only from a signal propagation theory at initialisation perspective. There are many empirical findings beyond the current theory, e.g. regarding normalisation and training dynamics, that the current theory is unable to predict, which is one of the motivations for this paper and presents opportunities for future work to improve upon the current theory. Notably, our most simplified block (in figure 1, top right), without normalisation, is able to match the training speed of Pre-LN transformers (e.g. in Figure 5) in terms of runtime, but we do observe slight decreases in per-step convergence in e.g. Figure 6.

   Would the reviewer suggest any changes to the wording in particular to improve the clarity? Perhaps we can add "from an initialisation perspective" to the end of the first paragraph in section 4.4?

4. Table 1: The reviewer is correct in that the means in table 1 for our methods are very slightly worse. Having said that, up to statistical significance our simplified models match the downstream performance of standard blocks in Table 1, so we do not believe our models perform worse. Our main focus in this submission was on improving pre-training speed from random initialisations in simplified blocks, and we believe that the factors that allow for fast pre-training are not necessarily the same as good downstream performance (e.g. we need to use dropout regularisation and also much smaller learning rates to avoid overfitting at downstream time, whereas dropout only ever hurts pre-training speed). We followed exactly the fine-tuning protocol of Geiping et al in Crammed BERT (5 epochs, fixed hyperparameters across tasks etc) with minimal hyperparameter tuning besides learning rate. We are actively working towards improving our understanding of fine-tuning from pre-trained checkpoints in order to better test, and improve, the fine-tuning capabilities of our simplified models.

6. Other data/task settings: In all the settings we have tried so far we have been able to emulate the default Pre-LN performance with our simplified blocks. We also present new autoregressive exps on language data (with longer context length of 512) in the rebuttal (in response to Reviewer wksT), finding similar conclusions. These experiments can be found at the following anonymous link https://docs.google.com/document/d/1SF8F3CXBQ_9xqTicGOjKC-UWBTKbZDTiT1zp3gqfD5c/. As you say, we need to test at larger scales in order to be sure that our findings are relevant in practice, but they definitely seem to be promising at present.

7. Limitations of the proposed simplifications: For normalisation, as discussed in the submission, though we are able to remove normalisation by following signal propagation principles, we observe slightly reduced training speeds (e.g. in Figure 6), which we plan to explore further in future work. Note that both of our main methods SAS and SAS-P include normalisation. By setting values and projections parameters to identity and removing skip connections, our models do not linearly project/rotate the values in the residual branch to a basis that is different to the identity projection of the skip branch, before they are summed. From our experiments and intuitions, this appears not to be a notable loss in capacity, especially given the matmul FLOPs saved by avoiding the value/projection weights. We discuss this question of expressivity loss from identity values and projections in more detail in the paragraph just before section 4.3 of our submission.

8. Model size: For nearly all settings our depth 18 Pre-LN model has 170M parameters, which is when it has MLP hidden width=3072. For the depth scaling experiments in Figure 6, we needed to reduce the MLP hidden width to 1536 to account for the 400M depth-72 model, and so our 18 layer model has 120M parameters. Our simplified models do not have values or projection parameters, which corresponds to around 15% reduction.

9. RE V-SkipInit vs Shaped Attention: both are motivated through ensuring good signal propagation when one scales the network to deep layers, and both achieve this by a dominant identity component in the attention matrix at initialisation, but with slightly different parameterisations. We tried both and chose shaped attention over v-skipinit because it provided very slightly better performance (figure 12 of submission).

10. Simplification compared to residual connection: We believe there has been a misunderstanding here. Both of our main methods, SAS and SAS-P, fix the values and projection matrices to be exactly identity. We added identity to the value and projection matrices in section 4.1 as a stepping stone towards removing them entirely later, because we wanted to present an intuitive progression for our simplifications. We believe that skipless blocks with identity values and projections are simpler than the residual blocks as it gives the simple interpretation of the (now removed) skip connection and signal propagation purely within the attention matrix itself, in terms of how much a token attends to itself relative to other tokens (e.g. the diagonals vs the off-diagonals of the attention matrix). On the other hand, with an explicit residual connection the value and projection weights are "projection matrices" which linearly transform/rotate the values, complicating how the skip relates to the attention mechanism.

We thank Reviewer iboi for taking the time and effort to review our work! We will address the raised concerns and questions in order below:

1. Scale: We agree that it would be interesting to test our models at larger scales. However, we would like to point out that we actually experiment up to near 400M parameters, not 100M, in our depth scaling experiment in Figure 6. In Figure 6 we demonstrate that our methods scale well to 72 layers, which is already deep by modern standards (e.g. LLaMa-33B/70B have 60/80 layers respectively). These 72-layer models have 380M parameters (hidden_dim=768, mlp_dim=1536), compared to our other experiments at 18 layers, which are ~120-170M depending on the MLP dim (either 1536 or 3072). Being in academia, it is not easy for us to obtain the compute resources that would allow us to experiment at larger scales (e.g. in the billions of parameters); we are actively pursuing avenues that would allow us to do so.

2. Other downstream evaluations besides GLUE benchmark: Thanks for raising this point: would the reviewer happen to have any particular downstream benchmarks in mind, particularly those suitable for researchers with low GPU resources who need to train from scratch (not using pre-trained models)? Our focus in this work is primarily pre-training speed, though we agree that it is important in practice that these benefits translate to the fine-tuning setting. For that reason, we included experiments on the Crammed-BERT GLUE setting, which is popular in the community particularly for those on academic compute budgets.

3. Pre-LN vs Post-LN: To our knowledge, Pre-LN is used in all LMs these days, rather than the original Post-LN block. The following LM papers explicitly state this: LLaMa section 2.2 https://arxiv.org/abs/2302.13971, GPT2 section 2.3 https://d4mucfpksywv.cloudfront.net/better-language-models/language-models.pdf and other works which followed GPT-2 design e.g. Chinchilla/Gopher. One could argue that the parallel block (used in GPT-J, Palm, ViT-22B) is different to both, but we point out that the standard parallel block (e.g. Figure 1 of submission, bottom right) has a skip connection with normalisation within the residual branch, so can be thought of as the “Pre-LN parallel block”. The “Post-LN parallel block” would not scale to large depths due to signal propagation issues (e.g. rank collapse); we avoid these issues by explicitly down–weighting the MLP branch, as discussed in sections 4.3 and 4.4 of our submission.

4. Questions about efficiency: We believe there has been a misunderstanding here. When we use the term “per-update training speed” we mean “how fast the loss is decreasing per training step”. We show that our simplified transformers essentially match the Pre-LN arch in training speed per update in e.g. Figure 6 (note the x-axis is number of steps) or Figure 3 (where we plot the final eval loss, so lower is better training speed per update). Because we have fewer parameters and a simplified block, this translates to higher throughput and so we actually match or have slightly faster training speed by wall-clock e.g. in Figure 5 (note the x-axis is runtime). Thank you for raising this, we will reword “per-update training speed” to “per-update convergence speed”, which we hope is clearer.

5. Why are skipless transformers slower than skip ones: This is the question that we address in section 4.1 of our submission. We show that the issue lies in the value and projection matrices, and that one needs to restrict the updates to value and projection matrices (or equivalently reduce their learning rate, c.f. Appendix A), to recover this lost training speed (Fig 3). In the most extreme case, this means that we fix the values and projections to be the identity matrix to recover the lost training speed, which leads to our SAS block. We describe our intuitions for this in the second and third paragraphs after the bolded paragraph “Recovering lost training speed” in section 4.1.