Paper: Transformers Get Stable: An End-to-End Signal Propagation Theory for Language Models (ICLR-2024)

Reviewer comment (Novelty):
Lack of novelty: scaling Residual/Skip-Connection is a well-known trick to stabilize training of deep neural networks, very similar ideas can be found at [1, 2]

Author response:
### **Regarding scaling in prior works**:
Scaling the residual/skip connection is one of the features we use in DSLM to stabilize the model, and it indeed has a long history, as we discuss in B.3. However, as our ablations of our initialization in Table 10, and comparisons to DSInit (Table 5) shows, it is crucial to account for the exact constants in model variance propagation.

Without these constants derived using our formulae (Table 1-2, 11-16), the models perform significantly worse.

Reviewer comment (Novelty):
Similar signal propagation idea is presented in [3]

Author response:
### **Regarding ‘similar signal propagation idea presented in [3]’**:
As we discuss in our related works (Section B.2 Paragraph 3), [3] (Deep transformers without shortcuts) assumes MLP to be linear in the effect it has on attention. Our analysis does not make any such assumptions about the interactions between various model components, and in particular, we use closed-form expression for the output correlation of GeLU (the nonlinearity used in [3]). In particular, our formulae can be used to show that an MLP block with GeLU will also increase correlation, in the absence of dropout (the same setting as used in [3]). As such, at deeper depths, the method in [3] will still exhibit rank collapse.

Furthermore, their 72 and 108 layer models underperform compared to a 36 layer model with same hidden dimension, in spite of having many times extra parameters - this clearly highlights deficiencies in their modeling of signal propagation, in particular the impact of non-linearities.

Reviewer comment (Novelty):
The results are not surprising [4] shows exactly the same conclusion with similar model configs.

Author response:
### **Regarding similar observations as in [4]**:
[4] (A Study on Transformer Configuration and Training Objective (Bamboo)) shows that Masked autoencoder objective may help stabilize a vanilla narrow-but-deep transformer model. Their configurations are vanilla transformers without any modifications, similar to our “baseline” models. The performance of their model falls at 96 layers, and they found 48 layers to be optimal.

As can be seen in Table 4, we also observe this drop in performance, for both vanilla Pre-LN and Post-LN models for our experiments on 165M models, where deeper models (such as 192 layers) underperform. Using our method however, one can train much deeper models (such as 192 or 384 layer models) which not only converge, but outperform the vanilla transformers of [4]. These differences are even more pronounced for larger models, where vanilla post-LN BERT from [4] did not converge for deeper models.

Reviewer comment (Novelty):
The idea of preventing rank collapse has been thoroughly explored in [5,6]

Author response:
### **Regarding rank collapse work in [5,6]**:
We discussed in section 3.4 and B.4 that [5] ($\sigma$Reparam) is one of the many works that have reported/studied rank collapse, and that they propose scaling all the linear weight matrices in the model by their spectral norm. Our method theoretically prevents rank collapse at initialization, whereas [5] empirically finds that bounded weights of attention will prevent rank collapse later during training. Similarly, [6] is also one of the papers we mentioned in 3.4 and B.4, and is again a method that targets attention collapse that occurs dynamically during training, and not at initialization.

Both these methods will not prevent rank collapse at initialization caused by the very structure of the transformer model, in particular increase in correlation caused by both attention and ReLU/GeLU as shown by our formulae. This is clearly demonstrated by the fact that [5] was unable to train a 200L post-LN model with $\sigma$ Reparam.

Both these methods are orthogonal to our method, to further stabilize training dynamics of the model - though in none of the training of our models did we observe divergence, even upto 768 layers (which is significantly deeper than the maximum depths studied in these works).

Reviewer comment (Experiments):
Experiments on downstream tasks seem lack of diversity. More downstream tasks with different characteristics should be included.

Author response:
### **Regarding Wall Clock Time and  benefit of using deeper and narrow config**:
As we detailed in Section 4.4 and Appendix J, our deeper models have small overheads in compute and wall clock times compared to shallow models. The benefits of using deeper and narrower configs with our method are -
1. Improved perplexity - allowing 165M model to outperforms the baseline 330M model with twice the params. At equal params, we achieve 1.0 improvements in perplexity.
2. Large improvements in downstream performance - Our deeper configs outperform original models by 2.2 points on RACE and 1.0 points on MNLI. The 165M model again outperforms original model with twice the params
3. Improved model quantization - Similar to Unit Scaling [1], our method results in models which lose much less performance, when quantized (via direct casting) to FP8 precision compared to original models, allowing our model to be served with much less GPU VRAM requirements.

Reviewer comment (Novelty):
Overall, most of the tricks are already well-studied and published, the paper generally feels incremental and results are not surprising. In combination with weak empirical studies on models with trivial sizes, this paper doesn't seem to be significant enough to be presented at the ICLR venue.

Author response:
### **Regarding results of the pretraining experiments compared to other works**:
See common response to all reviewers regarding model performance.
Regarding counting the embedding parameter size when reshaping the networks:
We do not count the embedding parameter size when reshaping the model for fair comparison. As our thinner-deeper models have smaller hidden dim, their embeddings are much smaller - for example, the 192L-256D model has 25M fewer params than the baseline 12L-1024D model, and still outperforms them.

### **Regarding pre-training performance**:

Original BERT model was trained for 128B tokens, and its performance was reported in Table 6 in Devlin et al. Our models are trained for Chinchilla optimal 3B/6B tokens, 40-20x fewer than the original BERT pre-training tokens. Also, the original Devlin et al models are trained on Book Corpus and Wikipedia, which are much more cleaner datasets than Pile-CC, which is derived from Common Crawl.

Our baseline model are trained for 1e-19 flops, and 165M models for 2.5e-18 flops. As can be seen in figure A5 of Chinchilla paper, their experiments predict a loss of approximately 3.2 to 2.8 given this compute budget - which our models clearly outperform.  Our GPT models achieve better performance than this at a loss of $ln(11.6)=2.45$, our 330M BERT achieves a loss of $ln(13.2)=2.58$, and the 165M model achieves $ln(14.2)=2.65$. Furthermore in [1], for a 32-layer 1024D model trained on C4 (another common crawl derived dataset) for 800M tokens, they report a perplexity of 24.7 for vanilla pre-LN transformer. At the same number of tokens, our baseline GPT 24L 1024D model was at 11.8 PPL.

We would like to reiterate that we used all the original hyper-parameters of BERT/GPT, and our baseline models directly used the original Megatron-LM codebase. Furthermore, we did a sweep of LR to find the best LR. We will release the training scripts to enable direct reproduction/comparison of our training.

[1] He, Bobby, et al. "Deep Transformers without Shortcuts: Modifying Self-attention for Faithful Signal Propagation." The Eleventh International Conference on Learning Representations. 2022.

### **Regarding Downstream performance**:
We add the performance of our model on the MNLI task as well (see common response). Figure 5 in Devlin et al. provides the expected MNLI dev accuracy based on the number of pretraining steps. From Figure 5 of BERT paper, this would correspond to around 25k steps and the expected accuracy is under 80% for 110M Bert Base. Our models achieve better accuracy than expected based on this Figure.

Reviewer comment (Presentation):
In Figure 2, there seems to be a discrepancy. The author mentions that the backward gradient variance rises hyperbolically with N, but the depicted curve suggests a decline as N grows. This is somewhat perplexing.

Author response:
### **Regarding gradient variance growth discrepancy in Figure 2**:
We would like to clarify Figure1-4. These figures depict the output (forward) and gradient (backward) variance after the $N^{th}$ layer, for a single model with 192 layers. The layers with higher numbers are near the output, lowest numbers are near the input, with $N$ denoting the current layer. As can be seen from Figure 2, as the gradient is backpropagated to shallower layers, it increases hyperbolically as a function of the current layer number $N$.

We should have used $n$ in Figure 1-4 to ensure consistency with our derivation in Appendix H. We will add the above clarifications to Figure 1-4. We will release the scripts to reproduce these figures 1-4.

Reviewer comment (Presentation):
The representation and caption for Figure 5 lack clarity, making it challenging to decipher the conveyed information.

Author response:
### **Regarding lack of clarity in Figure 5**:
Figure 5 depicts the input correlation $r^l_x$ (along x axis), vs the output correlation (along y axis) of a transformer layer, for both FFN block (Red line) and Attention Block (Blue line). The grey line is the $y=x$ line, added to show more clearly that the FFN block decreases correlation (compared to its input) after  $r^l_x > 0.65$, and that the Attention block always outputs correlation 0.9. These output correlations are plotted using Formulae for $r^l_{x_{out}}$ from Table 2.

This figure shows that with dropout, rank collapse (i.e., a situation where $r^l_x \approx 1$) will not occur in a standard transformer. This result is then empirically verified in Figure 6, where it can be clearly seen that no rank collapse is observed, contrary to the findings of Noci. et. al. 2022, as further discussed in Appendix G. We will explain these figures more clearly in the manuscript.

Reviewer comment (Evaluation):
For Table 4 and Figure 7, it would be beneficial if the author could elucidate why the thinnest and deepest transformers utilizing DSLM yield the most optimal results.

Author response:
### **Regarding explanation about why the thin/deep transformers utilizing DSLM yield the most optimal results (Table 4 and Figure 7)**:
Table 4 and Figure 7 show that standard transformers should be much more deep, at the expense of width - for example for 165M params, 48-192 layers deep, and for 330M params, 96-384 layers deep. Montúfar et al. (2014); Raghu et al. (2017) show that the complexity of Deep Neural Nets increases polynomially with width and exponentially with depth. Given a fixed parameter budget, there is a tradeoff between having richer representations from more deep layers stacked one after another, versus the representation capacity of model width.

Note that the 768 layer model underperforms compared to the 192 model - after a very large depth, the width becomes too narrow for the model to have enough representation capacity. The thinnest and deepest model (768 layer) does not perform the best - But our experiments show that, with correct initialization and scaling, the optimal depth is much more deeper (and the width narrower) than is usually used in standard models.

Reviewer comment (Novelty):
(originality) A very similar idea for regular networks has been presented in (Arpit et al., 2016). The connections with this work should definitely be discussed.

Author response:
Normalization Propagation (Arpit et al., 2016), proposes theoretical value for the batch statistics of mean and variance after ReLU, and uses these to normalize the output after each Relu. Their theoretical formulae for $\mu_x$, $\sigma^2_x$, and $\sigma^2_g$ are the same as in this work. This normalization successfully conserves the mean and variance of the forward propagation of the model as zero mean and unit variance.

However, this results in the backward gradient variance increasing by $0.5*\frac{1}{\sigma^2_x} = 1.47$, and to counteract this, they scale the inputs before applying relu by $\sqrt{\frac{1}{1.47}}$, and this successfully conserves the gradient as unit variance.

If our understanding of Arpit et al 2016 is correct, this scaling operation of the input results in the mean and the variance no longer being conserved - something seemingly ignored in Arpit et al. In fact, the variance falls to $30\%$ of initial value after just 4 layers. Applying an analysis similar to Klambauer et al. 2017 shows that Arpit et al.’s method converges to a fixed stable point of approximately $\mu=-0.34$ and $\sigma^2_x=0.24$, as we confirmed using numerical simulations. While we skip this proof here for brevity, kindly let us know if you would like us to expand on its proof and simulation results.

Furthermore, Arpit et. al does not handle attention, which we show requires careful consideration of covariances.

Reviewer comment (Novelty):
(originality) By moving all related work to the appendix and not citing much in the main text it becomes unclear which concepts are new and which already exist. I did notice the references to the appendix section, but I would argue that a citation puts more emphasis on the fact that something already exists and is not a contribution of this work. Citing one (seminal) paper for each aspect should suffice to counter this issue.

Author response:
Since our work uses proper initialization, residual scaling and signal propagation, we believed a detailed discussion of related works covering all these topics was essential. We felt that a short related works section in the main paper would not do justice to the large number of prior works we wanted to discuss. As such, we dedicated 2 pages to the related works, and it was difficult to have this entire section in the main paper due to space constraints. Instead, throughout our main paper, we tried to refer to prior works in-line. Being cognizant of the importance of this section, we made it the first section of our appendix. We will add a summarized related works section in the main paper as suggested. We had missed Poole et al., 2016 and the seminal references (Neal, 1996, & LeCun et al., 1998), thank you for pointing it out. We will add these to our manuscript.

Reviewer comment (Theory):
(clarity) In Section 4.4 an initialisation that renders a model more linear is claimed to be bad for performance. However, there is quite a bit of work where this kind of linearity is considered desirable (e.g. Hardt & Moritz, 2017; Zhang et al., 2019).

Author response:
Hardt & Moritz, 2017 suggests that a network should be able to express the identity transformation, in order to optimize well. We would like to clarify that by a “more linear” network, we are referring to a network which behaves almost like a single linear projection during training, because the residuals are negligibly small. As a thought experiment, if $\beta$ is reduced to a very small value (such as 1e-20), and assuming layernorm gain is fixed to 1, the transformer network will return the same output as input - and even training for a long time will not affect the model’s output significantly. As [1] shows, there is an expressivity-trainability tradeoff in training deep networks, and while having lower $\beta$ will result in networks whose gradients don’t explode/vanish, they will converge slowly/suboptimally.


[1] Yang, Ge, and Samuel Schoenholz. "Mean field residual networks: On the edge of chaos." Advances in neural information processing systems 30 (2017).

Reviewer comment (Evaluation):
(clarity) I am not sure if it makes sense to report perplexity in the context of masked language modelling. I am no expert in NLP, but I thought perplexity is only meaningful for autoregressive language models.

Author response:
The perplexity definition used here is simply exponential of the pre-training test-set loss, and hence a direct measure of model pre-training performance. Perplexity for MLM has been used in other works, such as Roberta [1], Deberta[2], Megatron LM[3], Scaling Laws vs Model Architectures[4], or a perplexity variant in [5], and in [6].

[1] Liu, Yinhan, et al. "Roberta: A robustly optimized bert pretraining approach." arXiv preprint arXiv:1907.11692 (2019).

[2] He, Pengcheng, et al. "DEBERTA: DECODING-ENHANCED BERT WITH DISENTANGLED ATTENTION." International Conference on Learning Representations. 2020.

[3] Shoeybi, Mohammad, et al. "Megatron-lm: Training multi-billion parameter language models using model parallelism." arXiv preprint arXiv:1909.08053 (2019).

[4] Tay, Yi, et al. "Scaling laws vs model architectures: How does inductive bias influence scaling?." arXiv preprint arXiv:2207.10551 (2022).

[5] Salazar, Julian, et al. "Masked Language Model Scoring." Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics. 2020.

[6] Lu, Jinghui, et al. "What Makes Pre-trained Language Models Better Zero-shot Learners?." Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2023.

Reviewer comment (Experiments):
(quality) The baseline performances seem to be remarkably weak. Table 6 in (Delvin et al., 2019) reports perplexities in the range 3-5. Results for GPT models on the pile go below 1. Also, Yang et al. (2019) report results in the range of 70-76% accuracy for BERT on RACE.

Author response:
Original BERT model was trained for 128B tokens, and its performance was reported in Table 6 in Devlin et al. Our models are trained for Chinchilla optimal 3B/6B tokens, 40-20x fewer than the original BERT pre-training tokens. Also, the original Devlin et al models are trained on Book Corpus and Wikipedia, which are much more cleaner datasets than Pile-CC, which is derived from Common Crawl.

Our baseline model are trained for 1e-19 flops, and 165M models for 2.5e-18 flops. As can be seen in figure A5 of Chinchilla paper, their experiments predict a loss of approximately 3.2 to 2.8 given this compute budget - which our models clearly outperform.  Our GPT models achieve better performance than this at a loss of $ln(11.6)=2.45$, our 330M BERT achieves a loss of $ln(13.2)=2.58$, and the 165M model achieves $ln(14.2)=2.65$.

We would like to reiterate that we used all the original hyper-parameters of BERT/GPT, and our baseline models directly used the original Megatron-LM codebase. Furthermore, we did a sweep of LR to find the best LR. We will release the training scripts to enable direct reproduction/comparison of our training.

Reviewer comment (Experiments):
(significance) There are too little direct comparisons with competing methods (e.g. Zhang et al., 2019; Noci et al., 2022; Wang et al., 2022a; He et al., 2023). Apart from Table 5, these alternative methods seem to be ignored completely. Also, it would be easier to compare if some of the experiments from these papers would have been adopted.

Author response:
We further discuss the methods the reviewer mentioned:

**Zhang et al 2019**: We discuss DSInit’s signal propagation in 3.3, their scaling in B.3, and compare performance for two deep models in Table 5. Crucially, DSInit’s initialization fails to converge for the 96-layer model despite a more conservative scaling of $O(N^{-2})$, highlighting the impact of carefully considering the constants (from initialization) hidden away in asymptotic formulations.

**Noci et al 2022**: We provide a detailed discussion of the Rank Collapse of Noci. et. al. in Section 3.4,  Appendix G highlighting some issues with their rank collapse. We also discuss some deficiencies in their signal propagation in 3.3 and B.2 - By not assuming constant attention, we arrive at exponential back-propagated gradients to Q,K, which we show can affect model stability, both theoretically and empirically. We also refer to their scaling in 3.4 and B.3. Their scaling corresponds to $\lambda=1$ and $\beta=k/N$. As we mentioned in B.3, we used $lambda^2 +\beta^2 = 1$ instead, as He. et al 2023 showed fully normalized residual connections often result in better performance. Our initial experiments also showed lower scores with $\lambda=1$.

**Wang et al. 2022a**: We discuss DeepNorm’s signal propagation in 3.3, and compare performance for two deep models in Table 5. As we further discuss in B.3, DeepNorm shows performance improvements on making the model deeper but keeping the hidden dimension constant, with a larger model - whereas our method shows performance improvements on making the model deeper while keeping the parameters constant.

**He et al. 2023**: We discuss the signal propagation theory of He et al in B.2, scaling in B.3. While He et al. considers the MLP to be linear, but by also considering the non-linearity in MLP we can correctly model how the nonlinearity of MLP strongly affects attention via correlation. We arrive at formulae which model the propagation accurately across the entire model, which cannot be done without considering the non-linearity. Their 72 and 108 layer models underperform compared to a 36 layer model with same hidden dimension, in spite of having many times extra parameters - compared to our method where deeper models outperform standard models without extra params, we hence did not benchmark this method. [2] also found He et al to underperform compared to baselines.

We would like to further discuss 3 more methods for deeper models, which we investigated:

**Admin (Liu et. al. 2020a)**: Admin requires additional “profiling” passes over the dataset to achieve unit variance for forward propagation, as we discuss in B.1. - while we can do this fully theoretically. We attempted to train our 192L-256D models of Table 5 with Admin, but all our experiments failed to converge, even after decreasing LR by 20x. As Admin paper only trained upto 72-layer models, we did not consider it for inclusion in Table 5.

**SkipInit**: SkipInit (De & Smith, 2020) underperforms compared to baseline Pre-LN models for Language Modelling, as observed by He et. al. 2023. As our method outperforms baseline Pre-LN model, we did not perform comparisons to this method.

**Shaped Attention**:  As per [2], Shaped Attention[1] does not perform as well as baseline models, resulting in significantly worse performance (Figure 2, 3, 5).

[1] Noci, Lorenzo, et al. "The shaped transformer: Attention models in the infinite depth-and-width limit." arXiv preprint arXiv:2306.17759 (2023).

[2] Anonymous, Simplifying Transformer Blocks, 2023. https://openreview.net/forum?id=RtDok9eS3s

Reviewer comment (Evaluation):
(quality) Experiments do not have error bars. Especially for a random initialisation strategy, error bars would be helpful to assess how consistent the improvements are. If error bars would make the experiments prohibitively expensive, it would be nice to include at least one small-scale experiment to provide some insights on the variability of the proposed method.

Author response:
In our initial experiments, we observed very little variation in performance across different runs - we conjecture that the model is trained on a large enough number of tokens for differences in initialization/data seed to not matter. We provide numbers below for 12L-1024D Post-LN and DSLM models from Table 4 below -

| Model | Mean | Standard Error |
| ------ | ------ |  ------ |
|Post-LN Baseline |14.33 | 0.14 |
| DSLM | 15.56 | 0.08  |

As the variation was so small, and due to compute limitations, we did not run multiple runs for other experiments thereafter. We also reported the best score for Baseline Post-LN, and the worst score for DSLM for the 12L-1024D models Table 4 for a conservative comparison. We will add these details in the paper.

Reviewer comment (Novelty):
(quality) It seems like two fundamental signal propagation papers are missing in the related work. The foundations for signal propagation analysis can be found in (Neal, 1996) and a popular reference is (LeCun et al., 1998).

Author response:
We had missed Poole et al., 2016 and the seminal references (Neal, 1996, & LeCun et al., 1998), thank you for pointing it out. We will add these to our manuscript.

Reviewer comment (Writing):
(clarity) Section B.2§4 states that non-IID inputs are "also" accounted for, but I would argue that this is the only case that is accounted for.

Author response:
Our exact formulae for blocks and components also account for IID cases - as can be verified by our simulations, in which we do cover cases IID inputs with exactly 0 correlation, as noted in $Corr^l_{x_{in}}$ column in Table 18. In the simplified formulae for Table 19, and in DeepScaleLM initialization and model, we simplified our formulae so that they only remain accurate for non-IID inputs. This was because of three considerations:
1. In NLP domain, most text will inevitably be non-IID due to repeated common words. This was encountered in all our experiments.
2. Even in Vision, for ViT in particular, there will be correlation among pixel intensities across patch embeddings, as discussed in common response section.
3. Even if there is exactly 0 correlation in input, the very first attention layer and the first FFN layer in particular, will add correlations to the output, ensuring our simplified formulae hold reasonably accurately.

Reviewer comment (Writing):
(clarity) It is unclear what computation is contained in the embedding component. Originally, I suspected this to be only about the token embeddings, but it seems to include other computations as well.

Author response:
We will update our paper with clear description and details of the embedding component. Our embedding component is same as that of the original BERT model, except for position embeddings. The embeddings component of BERT consists of 3 look up tables - token embeddings, position embeddings, and segment embeddings. For a given input token, each of these 3 embeddings are added before being passed to the transformer model. We describe each of these below -

- Word Embeddings - A lookup table mapping input word ids, of size (Vocabulary X Hidden_dim). Output denoted by $x_{out_{we}}$ in Section 3.1.

- Position Embeddings -  A lookup table mapping input positions, of size (Sequence_length X Hidden_dim). We use the more commonly used trainable position embeddings (compared to BERT’s original sinusoidal embeddings)

- Segment Embeddings - A lookup table mapping segments (denoting Sentence A or Sentence B for Next Sentence Prediction task of BERT), of size (2 X Hidden_dim). Output denoted by $x_{out_{se}}$ in Section 3.1.

Reviewer comment (Theory):
(quality) The approximation for the correlation in the embedding layer seems to be missing the Euler constant:  
$$\sum_{i=1}^{|V|} p_i^2 = \frac{\sum_{i=1}^{|V|} 1 / i^2}{\Big(\sum_{i=1}^{|V|} 1 / i\Big)^2} \approx \frac{\zeta(2)}{(\ln |V| + \gamma)^2}.$$  
Without the constant, the approximation is pretty bad. Also, it could be made clearer in the derivation where this approximation comes from and that the final simplification step assumes large $L$.

Author response:
We would like to apologize for a minor error here, where we wrote our predicted $r^l_{x_{in}} = 0.247$. In writing this value for the paper, we had mistakenly used Log to the base 10, instead of natural logarithm. Using natural logarithm gives a predicted value of 0.2273 (instead of the 0.247), which is within 3% of the empirical correlation 0.221. Adding the euler’s constant term gives an even closer value of 0.2267, but not significantly different from our simplified version. In our experiments, we had used the natural logarithm value of 0.227.

We had skipped the Euler constant, as it was relatively small (0.58) compared to $Log_e(|V|)=10.37$ for our vocabulary of size 32000. We will clarify this, as well as mention the simplification from large L.

Reviewer comment (Novelty):
(originality) It should be more clearly stated for each derivation where it can be found in literature. In its current form, it is hard to distinguish which derivations are really new.

Author response:
We will more clearly specify which formulae (or equivalent forms) cannot be found in prior literature, and also color-code Tables 1, 11-16 for the same.

Reviewer comment (Theory):
(clarity) The derivation of LayerNorm could use some more explanation. Also, the dependencies between samples, sample mean and sample variance should be discussed a bit more. Finally, the affine transformation that is typically included after each normalisation layer is not addressed at all.

Author response:
We will add further explanations to the LayerNorm derivations. The affine transformation for layernorm are typically initialized with 1 scale and 0 bias, hence do not change any of our derivations and were hence ignored. We will clarify this in our paper.

Reviewer comment (Theory):
(clarity) It should be more clear what the limitations of the different approximations for the softmax derivation are. A quick numerical check verifies that the approximation breaks down quite quickly for larger variances. Although this is probably not practically relevant, it would be nice to provide some insights when the analysis breaks down.

Author response:
We will add a discussion on the effective range of softmax approximation. Our approximations become inaccurate for larger variances, particularly when softmax starts to “degenerate” to a single dimension 1 and others 0. It is reasonably accurate in more realistic ranges however.

Reviewer comment (Theory):
I noticed that the derivation for scaled dot-product attention starts with some rough shortcuts. Also, a completely different approach is used to handle the softmax. This seems suspicious. I do not have time to check the math any further, but the numerical results should provide enough proof that these results, if not correct, are at least useful.

Author response:
We had initially attempted a much more simpler proof, where one assumed attention scores to be independent of values. This then enabled the use of our previous LogNormal-based softmax derivation to easily derive the forward variances. But the theoretically predicted values strongly disagreed with empirical values from simulations. This is because SHA is -
$\mathrm{Dropout}(\mathrm{SoftMax}(\frac{\mathbf{X_{\text{in}}}\mathbf{W_Q}\mathbf{{W_K}^T}\mathbf{X^T_{\text{in}}}}{\sqrt{d_k}}))\mathbf{X_{\text{in}}}\mathbf{W_V}$ , and the $\mathbf{{W_K}^T}\mathbf{X^T_{\text{in}}}$ term cannot be treated independently of the $\mathbf{X_{\text{in}}}\mathbf{W_V}$ term.

A simple verification of this can be checked by simply simulating $(XW)^T*X$, and verifying that the variances of the results do not match that of $L * \sigma^2(XW)$, but do if the second X is replaced by another random tensor.

This necessitates an alternate methodology to derive SHA, where the components are treated as a unified whole.

Reviewer comment (Reproducibility):
The concrete picture of the adjusted model architecture is missing to the reader. The proposed DeepScaleLM scheme scales residual connections, and seems like it adds extra dropout layers. In which component are these adjustments applied? Part of a layer or every component in a layer? Partial layers or every layer?

Author response:
### **Regarding details about adjusted model architecture**:
We do not add any new dropout over the original standard transformer architecture of BERT - we show the critical importance of dropout (which are already present in standard transformer) in stabilizing very deep transformers by preventing rank collapse. The only changes we made were

1) Scaling of residual and skip connection, and
2) Changing the standard deviation of initialization of weights of model params.

Reviewer comment (Evaluation):
Plotting log-scale in Figure 3 does not help the reader understand deeply about the gap between empirical simulation and theoretical prediction. Providing statistics such as min, max, mean, variance of the gap would be better.

Author response:
### **Regarding clarity and statistics for Figure 3**:
To clarify, Figure 3 shows the variance of gradient across different layers for a given model with 192 layers- it shows the gradient flowing backwards decreases exponentially. We will clarify this in the paper. We provide percentage relative error of our predicted gradient variance compared to the empirical values below:

| Mean Error | Median Error | Std Error |
| ------------------|---------------------|-----------------------------------|
|    6.8\% | 5.2\% | 7.8\% |

Furthermore,  we also measure the $R^2$ of our predicted gradient as a goodness of fit measure, and we find and $R^2$ of 0.998 - as can be seen, our predictions provide a good fit. We will add these goodness of fit measures to the main paper.

Reviewer comment (Theory):
In Section 2, the paper assumes normal distribution of inputs, weights, and gradients while deriving the formulas in Table 1. To what extent this assumption holds since in Section 1 paragraph 2 the paper mentioned that some of the assumptions in prior works broke down on real world data?

Author response:
### **Regarding the assumption of normally distributed inputs, weights, and gradients**:
**Inputs**: As the embeddings are lookup tables of token-ids, and embedding weights are initialized from normal distribution in xavier, the inputs to the transformer are normally distributed. We will state this clearly in the paper. However, while the inputs are normal, due to the repetition of the tokens / segment ids, the inputs are not IID, as many previous works have assumed - our work handles this explicitly.

**Weights**: Weights are initialized from normal distribution in xavier, and are hence normal.

**Gradients**: As the model outputs are gaussian, the softmax of the classification head results in a Log-Normal distribution for probabilities $p$, as shown in C.7. Since the cross-entropy loss is $-log(p)$, this results in the loss (and hence the final gradient being back-propagated) being log(Log-normal distribution), which is normal distribution.

We further verify this empirically by checking the normality of the backpropagated gradients to the deepest transformer layer, and the gradients match the best-fit normal distribution with an $R^2$ of $0.999$, showing that the gradients are indeed normally distributed. We will add this discussion on normality to the main paper.