Transformers Get Stable: An End-to-End Signal Propagation Theory for Language Models

Response to Reviewer PFEM: Part 1/2

### **For more Finetuning Experiments, modern architectures and new results**: Please see common 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). ### **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. ### **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. ### **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. ### **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.

Response to Reviewer PFEM: Part 2/2

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

Request for Rebuttal Feedback

Dear Reviewer, We would like to thank you for reviewing our paper. With less than 1 day left, we would appreciate if you could provide feedback regarding whether our response has addressed your concerns. Kindly let us know if there are any other details you would like us to clarify! Best regards, The Authors

Response to Reviewer wRZy

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

Request for Rebuttal Feedback

Dear Reviewer, We would like to thank you for reviewing our paper. With less than 1 day left, we would appreciate if you could provide feedback regarding whether our response has addressed your concerns. Kindly let us know if there are any other details you would like us to clarify! Best regards, The Authors

Response to Reviewer eeiy : Part 1/4

### **Regarding applying DSLM to Vision Transformers**: Kindly refer to common response. ### **Regarding using the same architecture for the initialisation method**: Please see common responses for results with decoder-only (GPT) architecture. ### **Regarding comparison with Arpit et al., 2016**: 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. ### **Regarding reformatting the Related Works section**: 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. ### **Regarding “linearity” being bad for performance**: 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). ### **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$. 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.

Response to Reviewer eeiy : Part 2/4

### **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, our models 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. ### **Regarding Perplexity**: The confusion regarding weak performance arose from two different definitions of perplexity. We calculate perplexity as the exponential of the model’s loss, (as provided in Megatron-LM paper[1] and code [link](https://github.com/NVIDIA/Megatron-LM/blob/443ce9f3f98fdc5a53c6b480c6e21b79944d198e/megatron/training.py#L975). Regarding GPT perplexity < 1, we estimate that the reviewer is referring to bits-per-byte as perplexity, whereas we calculate on a token level. In hindsight, we should have clarified this in the main paper. The chinchilla paper (in Figure A5) shows expected loss given our compute limit of 1e-19 flops at approx 2.8, and our GPT models achieve better performance than this. Furthermore in [2], 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 24L 1024D model was at 11.8 PPL. [1] Shoeybi, Mohammad, et al. "Megatron-lm: Training multi-billion parameter language models using model parallelism." arXiv preprint arXiv:1909.08053 (2019). [2] He, Bobby, et al. "Deep Transformers without Shortcuts: Modifying Self-attention for Faithful Signal Propagation." The Eleventh International Conference on Learning Representations. 2022. ### **Regarding consistency of improvements and error bars**: 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. ### **Regarding non-IID inputs being the only case that is accounted for in B.2**: 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. ### **Regarding the computation contained in the embedding component**: 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.

Response to Reviewer eeiy : Part 3/4

### **Regarding missing the Euler constant in the approximation for the correlation in the embedding layer**: 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. ### **Regarding stating which formulae can be found in literature**: 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. ### **Regarding more explanation in the derivation of LayerNorm**: 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. ### **Regarding limitations of different approximations**: 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. ### **Regarding the proof of scaled dot-product attention being different from softmax**: 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. ### **Regarding using perplexity for masked language modeling**: 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.

Response to Reviewer eeiy : Part 4/4

### **Regarding accuracy of numerical results**: Kindly refer to the response regarding 'approximation for the correlation in the embedding layer'. Our correlation estimates were within 3% of the ground truth value. Furthermore, as can be clearly seen from our formulae and from simulations, correlation is indeed very critical to correctly model signal propagation in transformers. ### **Regarding comparisons with competing methods**: 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

Response to Reviewer vaGo

### **Regarding more experiments with decoder-only Transformer, and new finetuning datasets**: Kindly refer to the common response section to find new results. ### **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. ### **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. ### **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.

Request for Rebuttal Feedback

Dear Reviewer, We would like to thank you for reviewing our paper. With less than 1 day left, we would appreciate if you could provide feedback regarding whether our response has addressed your concerns. Kindly let us know if there are any other details you would like us to clarify! Best regards, The Authors