Paper: Simplifying Transformer Blocks (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
**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.

Reviewer comment (Experiments):
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.

Author response:
**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.

Reviewer comment (Experiments):
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.

Author response:
**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.

Reviewer comment (Experiments):
No downstream evaluation of the GPT model.

Author response:
**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.

Reviewer comment (Experiments):
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.

Author response:
**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.

Reviewer comment (Experiments):
Were any downstream fine-tuning experiments conducted for GPT, like GLUE for BERT? It will help make the paper stronger if downstream performance also remains similar/better.

Author response:
**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.

Reviewer comment (Novelty):
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.

Author response:
**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.”

Reviewer comment (Writing):
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.

Author response:
**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?

Reviewer comment (Evaluation):
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.

Author response:
**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.

Reviewer comment (Theory):
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,.

Author response:
**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.

Reviewer comment (Evaluation):
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.

Author response:
**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.

Reviewer comment (Experiments):
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.

Author response:
**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.

Reviewer comment (Experiments):
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?

Author response:
**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.

Reviewer comment (Evaluation):
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?

Author response:
**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.

Reviewer comment (Theory):
On page 3, the authors mention skipless transformers are slower than those with skip connections. Why is that?

Author response:
**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.

Reviewer comment (Experiments):
Are there any settings (data or tasks) where this simplified transformer does not work as well compared to the standard transformer?

Author response:
**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.

Reviewer comment (Theory):
What are the limitations of the proposed simplifications in terms of model expressiveness and capability?

Author response:
**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.

Reviewer comment (Reproducibility):
What is the size of the model which has 18 blocks and 768 hidden dimensions?

Author response:
**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.

Reviewer comment (Novelty):
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?

Author response:
**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).

**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.