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

Akhil Kedia  Mohd Abbas Zaidi  Sushil Khyalia  Jungho Jung  Harshith Goka
Haejun Lee

###### Abstract

In spite of their huge success, transformer models remain difficult to scale
in depth. In this work, we develop a unified signal propagation theory and
provide formulae that govern the moments of the forward and backward signal
through the transformer model. Our framework can be used to understand and
mitigate vanishing/exploding gradients, rank collapse, and instability
associated with high attention scores. We also propose DeepScaleLM, an
initialization and scaling scheme that conserves unit output/gradient moments
throughout the model, enabling the training of very deep models with $100$s of
layers. We find that transformer models could be much deeper – our deep models
with fewer parameters outperform shallow models in Language Modeling, Speech
Translation, and Image Classification, across Encoder-only, Decoder-only and
Encoder-Decoder variants, for both Pre-LN and Post-LN transformers, for
multiple datasets and model sizes. These improvements also translate into
improved performance on downstream Question Answering tasks and improved
robustness for image classification.

Machine Learning, ICML

##  1 Introduction

Transformer models are extremely popular across different domains of machine
learning, however, deep transformers are plagued with issues of gradient
explosion/vanishing (Shleifer et al., 2021; Takase et al., 2022; Smith et al.,
2022; Chowdhery et al., 2023) and rank collapse (Zhou et al., 2021; Noci et
al., 2022) that adversely affect training stability. Proposed remedies include
residual scaling, changing initialization, or extra/modified layernorms (Zhang
et al., 2019a; Xiong et al., 2020; Bachlechner et al., 2021; Wang et al.,
2022a; Dehghani et al., 2023).

Theoretical analysis via signal propagation, kernel methods has led to an
improved understanding of these issues. Several works in the signal
propagation domain (Glorot & Bengio, 2010; Arpit et al., 2016; Xu et al.,
2019; Dong et al., 2021; Davis et al., 2021; Wang et al., 2022b) have analysed
the propagation of moments for some components of deep transformers, but often
make simplifying assumptions of IID inputs, uncorrelated outputs, ignoring
effect of query/key initialization, simplifying non-linearity, etc. We
observed break down of each of these assumptions with real world data,
adversely affecting model stability.

These issues highlight the need for a holistic theoretical framework that can
fully explain signal propagation through transformer models with real data. In
this work, we provide a framework to fully explain signal propagation through
transformer models, by deriving closed-form expressions for the first and
second-order moments (mean and variance) of the outputs and gradients of each
of the components of the transformer model (Embeddings, FFN, ReLU/GeLU,
LayerNorm, Dropout, Softmax, Single-Head Attention), Attention and FFN blocks,
and through the entire model. Our derived equations are empirically verified
within strict error bounds with real world data (code in supplementary).

We apply this framework to understand and mitigate instability issues with
deep transformers – vanishing/exploding gradients, rank collapse, and
instability caused by high QK values. To harness the improved complexity of
deeper models (Montúfar et al., 2014; Poole et al., 2016; Raghu et al., 2017),
we propose DeepScaleLM, a novel initialization scheme that augments
residual/output scaling, and ensures the moments of outputs and gradients
remain fully conserved throughout the model. DSLM enables us to break the
depth barrier and train models with $100$s of layers which consistently
outperform shallow models for BERT, GPT, Encoder-Decoder style models on
Question-Answering, Speech-to-text translation and Image Classification tasks
while reducing parameters.

##  2 Related Works

For detailed discussion of prior works, refer to Appendix A.

##### Initialization

Several works (Glorot & Bengio, 2010; He et al., 2015; Brock et al., 2021;
Poole et al., 2016; Schoenholz et al., 2017) improved the initialization of
ResNets/ReLU networks. These works do not consider transformers, and are
unable to handle Softmax/Attention. Others, such as ADMIN (Liu et al., 2020a),
Mishkin & Matas (2016); Liu et al. (2020b) achieve unit variance for faster
convergence by scaling the weights and/or outputs based on empirical profiling
of a forward pass. Blake et al. (2023) also tries to achieve this, but does
not completely handle correlation and non-zero mean of ReLU. We demonstrate
that this profiling is unnecessary, and can instead be done theoretically in
our work.

##### Signal Propagation

Signal propagation in Neural Networks (Neal, 1995; LeCun et al., 1996) has a
long history, such as for ResNets (He et al., 2015; De & Smith, 2020; Brock et
al., 2021; Schoenholz et al., 2017; Hoedt et al., 2022; Labatie et al., 2021;
Marion et al., 2022; Klambauer et al., 2017; Balduzzi et al., 2017), and for
transformers in (Xu et al., 2019; Dong et al., 2021; Davis et al., 2021; Noci
et al., 2022; Martens et al., 2021; He et al., 2023; Shi et al., 2022; Wang et
al., 2022b). Our work considers previously often neglected effects of dropout,
input correlation, activation non-linearity, $QK$ initialization, providing
closed forms with verifiable correctness of signal propagation. This allows us
to constrain the output and gradient to almost exactly unit variance.

##### Moment Control & Residual Scaling

Bounded gradients have been shown to results in better/faster convergence
(Shen et al., 2020; Yu et al., 2017; You et al., 2017, 2020; Takase et al.,
2022; Shleifer et al., 2021; Hayou et al., 2019). Different scaling schemes
for residual networks ($\lambda$ for skip connections and $\beta$ for residual
output) have been explored by prior works, such as $\lambda^{2}+\beta^{2}=1$
for ResNets (Balduzzi et al., 2017; Szegedy et al., 2017; Hanin & Rolnick,
2018; Arpit et al., 2019; Zhang et al., 2019b; Hoedt et al., 2022). Learnable
$\beta\approx 0$ was used in SkipInit (De & Smith, 2020), ReZero (Bachlechner
et al., 2021), LayerScale (Touvron et al., 2021b), Value-SkipInit (He et al.,
2023). Others proposed $\beta^{2}=$$\mathcal{O}(\frac{1}{N})$, where $N$ is
max/current layer was used in Arpit et al. (2019); Brock et al. (2021); Marion
et al. (2022); Zhang et al. (2022); He et al. (2023); Noci et al. (2022); De &
Smith (2020); Liu et al. (2020a, b); Davis et al. (2021); Blake et al. (2023),
while DSInit (Zhang et al., 2019a), T-Fixup (Huang et al., 2020), DeepNorm
(Wang et al., 2022a) used $\beta^{2}<$$\mathcal{O}(\frac{1}{N})$. However, the
optimal initialization/scaling can vary based on data/model characteristics
(Zhang et al., 2022; Marion et al., 2022). Our contribution goes beyond
providing an optimal scaling scheme – our theory enables informed choices
about these initialization/scaling schemes based on their expressivity-
trainability trade-off.

##### Other Network modifications for Deep Networks

Architectural modifications such as Zhai et al. (2023); Zhou et al. (2021);
Shleifer et al. (2021) can only stabilize the model later during training and
not at initialization. They are orthogonal to our approach, and our equations
can be easily extended to cover these. Some works such as DeepNet, ADMIN show
performance improvements by making the model deeper, but much larger. In this
work, we explore a stricter setting of keeping transformer parameters and
compute constant while making the model deeper.

##  3 Moments of Transformer Models

Table 1: Signal propagation for forward and backward passes through components
of a transformer (GeLU in Section B.5). The expressions here are illustrative
simplification of full closed form formulae in Appendices D and B.

Component $\mathbf{\mu_{x_{\text{out}}}}$
$\mathbf{\sigma^{2}_{x_{\text{out}}}}$ $\mathbf{\sigma^{2}_{g_{in}}}$
$\mathbf{r^{l}_{x_{out}}}$ $\mathbf{r^{l}_{g_{in}}}$ Embeddings 0
$\sum\sigma^{2}_{w_{\text{embd}}}$ -
$\dfrac{\pi^{2}}{18*\mathrm{log}(|V|)^{2}}+\dfrac{2}{9}$ - Linear
($d_{\text{in}}\rightarrow d_{\text{out}}$) 0
$d_{\text{in}}\sigma_{w}^{2}(\sigma_{x_{\text{in}}}^{2}+\mu_{x_{\text{in}}}^{2})$
$d_{\text{out}}\sigma_{w}^{2}\sigma_{g_{\text{out}}}^{2}$
$\dfrac{r^{l}_{x_{\text{in}}}+\mu_{x_{\text{in}}}^{2}/\sigma^{2}_{x_{\text{in}}}}{1+\mu_{x_{\text{in}}}^{2}/\sigma^{2}_{x_{\text{in}}}}$
$r^{l}_{g_{out}}$ ReLU $\dfrac{\sigma_{x_{\text{in}}}}{\sqrt{(2\pi)}}$
$\dfrac{(\pi-1)}{(2\pi)}\sigma_{x_{\text{in}}}^{2}$
$\dfrac{1}{2}\sigma_{g_{\text{out}}}^{2}$
$0.7r^{l}_{x_{\text{in}}}+0.3{r^{l}_{x_{\text{in}}}}^{2}$
$(\dfrac{1}{2}+\dfrac{\sin^{-1}{(r^{l}_{x_{\text{in}}})}}{\pi})\mathrm{r^{l}_{g_{\text{out}}}}$
LayerNorm ($d$) 0 1
$\dfrac{\sigma^{2}_{g_{\text{out}}}}{\sigma^{2}_{x_{\text{in}}}}$
$r^{l}_{x_{\text{in}}}$ $r^{l}_{g_{\text{out}}}$ Dropout ($p$) $\mu_{x_{in}}$
$\dfrac{\sigma_{x_{\text{in}}}^{2}+p\mu_{x_{\text{in}}}^{2}}{1-p}$
$\dfrac{1}{1-p}\sigma_{g_{\text{out}}}^{2}$
$\dfrac{r^{l}_{x_{\text{in}}}(1-p)}{1+p\mu_{x_{\text{in}}}^{2}/\sigma_{x_{\text{in}}}^{2}}$
$(1-p)r^{l}_{g_{\text{out}}}$ SHA-without V 0
$r^{l}_{x_{\text{in}}}\sigma^{2}_{x_{\text{in}}}$
$r^{l}_{g_{{\text{out}}}}\sigma^{2}_{g_{{\text{out}}}}$ $1$ $1$ Softmax
$\dfrac{1}{L}$
$\dfrac{e^{(1-r^{d}_{x_{\text{in}}})\sigma^{2}_{x_{\text{in}}}}-1}{L^{2}}$
$\dfrac{e^{(1-r^{d}_{x_{\text{in}}})\sigma_{x_{\text{in}}}^{2}}}{L^{2}}\sigma_{g_{\text{out}}}^{2}$
- -

###  3.1 Moments of Transformer Components

Following an analysis similar to that of Xavier initialization (Glorot &
Bengio, 2010), we derive closed-form expressions for the mean and variance of
the output and of the backpropagated gradient for all the components of the
transformer model in Table 1.

Here $\mu_{x_{\text{in}}}$, $\sigma_{x_{\text{in}}}^{2}$,
$\mu_{x_{\text{out}}}$, $\sigma^{2}_{x_{\text{out}}}$ are the mean and
variance of the input/outputs, $\sigma_{g_{\text{out}}}^{2}$,
$\sigma^{2}_{g_{\text{in}}}$ are the variance of the gradient back-propagated
to/from the component, and $r^{l}$, $r^{d}$ are the correlations across
sequence length and hidden dimension. $p$ is the dropout probability, $L$
sequence length, $d_{\text{in}},d_{\text{out}}$ input/output dimensions of
Linear layer, $\sigma_{w}^{2}$, $\sigma^{2}_{w_{\text{embd}}}$ variances of
the weights of the Linear layer and the Embeddings table. At the input side,
$r^{l}_{x_{\text{in}}}$ originates from repeated tokens. For text, we estimate
input correlation theoretically by assuming that input tokens follow Zipf
(Kingsley, 1935) distribution. Detailed proofs are provided in Appendix B, and
all assumptions are summarized in Section M.2.

###  3.2 Moments of Transformer Blocks

Combining the expressions reported in Table 1, we derive closed-form
expressions for the moment transformation during the forward and backward pass
of the transformer Attention and FFN blocks. The Attention block refers to the
$Q,K,V$ projection, followed by Multi-Head Attention and Output-Projection
Layer. The FFN block refers to the Linear layer followed by non-linearity
(ReLU) and output Linear layer. Table 2 provides our derived equations for
these, where $\sigma_{v}^{2}$, $\sigma_{o}^{2}$, $\sigma_{w_{1}}^{2}$,
$\sigma_{w_{2}}^{2}$ are the variances for $V$ weights, Output-Projection
weights, and weights of FFN block Linear layers, and $d$ is model the hidden
size. These results show that considering correlation $r^{l}$, dropout $p$ and
effects of non-linearity are crucial for correctly modelling signal
propagation through Transformer blocks.

Table 2: Moment Propagation through the blocks of a transformer layer. Exact
closed forms / proofs are provided in Appendices D and C.

Component $\mathbf{\sigma^{2}_{x_{out}}}$ $\mathbf{r^{l}_{x_{out}}}$
$\mathbf{\sigma^{2}_{g_{in}}}$ $\mathbf{r^{l}_{g_{in}}}$ Attention Block
$\dfrac{d^{2}\sigma_{o}^{2}\sigma_{v}^{2}\sigma^{2}_{x_{\text{in}}}*r^{l}_{x_{\text{in}}}}{(1-p)}$
$1-p$
$\dfrac{d^{2}\sigma_{o}^{2}\sigma_{v}^{2}*\sigma^{2}_{g_{\text{out}}}}{(1-p)}r^{l}_{g_{\text{out}}}$
$1-p$ FFN Block
$\dfrac{2d^{2}\sigma_{w_{1}}^{2}\sigma_{w_{2}}^{2}\sigma^{2}_{x_{\text{in}}}}{(1-p)}$
$(1-p)(\dfrac{1}{\pi}+\dfrac{r^{l}_{x_{\text{in}}}}{2}+(\dfrac{1}{2}-\dfrac{1}{\pi}){r^{l}_{x_{\text{in}}}}^{2})$
$\sigma^{2}_{x_{\text{out}}}*\sigma^{2}_{g_{\text{out}}}$
$(1-p)(\frac{1}{2}+\frac{\sin^{-1}{(r^{l}_{x_{\text{in}}})}}{\pi})r^{l}_{g_{\text{out}}}$

Figure 1: Pre-LN: Forward variance increases linearly with number of layers
$N$.

Figure 2: Pre-LN: Backward gradient variance increases hyperbolically with
$N$.

Figure 3: Post-LN: Backward gradient vanishes exponentially with $N$ (log-
scale).

Figure 4: DeepScaleLM: The variances remain conserved for both backward and
forward pass.

###  3.3 Moments of Entire Transformer Model

By repeatedly applying the expressions in Table 2 for each layer, we calculate
the propagation of moments of outputs and gradients through the entire
transformer model. We do this for both Pre-LN style transformers, in which the
skip connection bypasses the LayerNorm, and for Post-LN style transformers, in
which the Layernorm is applied before the skip-connection. The method is fully
detailed in Sections G.1 and G.2. Figures 4, LABEL:, 4, LABEL: and 4 provide
the forward (left to right) and backward (right to left) signal propagation at
initialization through the layers of a very deep $192$-layer model with Xavier
initialization.

###  3.4 Numerical Validation of Theoretical Results

We verify the theoretical formulae of transformer components and blocks by
running simulations with real/synthetic data, (detailed in Appendix E, code
released as supplementary material). Even at $99$ percentile, no error (other
than SHA gradient $\sigma^{2}$) is larger than $10\%$, verifying our
assumptions.

All our derivations are modality-agnostic. We verify our formulae for the
entire transformer model using real textual MLM data, as shown in Figures 4,
LABEL:, 4, LABEL:, 4, LABEL: and 4 (Reproducible in supplementary), and using
ImageNet data (as shown in Section I.1). Our formulae predict the observed
gradient and forward/backward norms with remarkable accuracy, with mean and
median relative errors of $6.8\%$ and $5.2\%$ respectively, and an $R^{2}$ of
$0.998$. We further verify that for model depths in range $[1-768]$, and model
dimensions $[128-6096]$, the reported formulae are within $10\%$ error, even
across $768$ layers of the transformer model.

###  3.5 Validity of Theoretical Predictions even after Training

Interestingly, our theoretical estimates hold approximately even after the
models have been trained for a large number of steps. The model stays in the
regime it is initialized with (as has also been shown in Li & Liang (2018);
Arora et al. (2019a); Lee et al. (2019); Jesus et al. (2020); Arora et al.
(2019b)), highlighting the importance of correct initialization. Further, we
analyze forward explosion in a $48$-layer PreLN model (after $100$k training
steps) and gradient explosion in a $64$-layer PreLN model (after $150$k
training steps) and use our theory to predict the moments. Our linear
estimation for the forward growth and hyperbolic estimation for the gradient
explosion match closely with the observed moments as shown in Appendix in
Figures 8, LABEL: and 8.

##  4 Applications

###  4.1 Explaining Variance Explosion in Transformer

Our approach theoretically proves the gradient vanishing/explosion (Table 3)
for both Pre-LN and Post-LN transformers.

##### Exploding Output and Gradient in Pre-LN

The forward output for Pre-LN transformer increases linearly with increasing
depth $N$ (Section G.1) since each layer’s output is directly added to the
skip connection, as seen in Figure 4. For the backward pass, the gradient
increases hyperbolically with increasing $N$, as seen in Figure 4.
Intuitively, this is because the gradient increases in every layer when a
block’s gradient is added to the skip connection, and the fractional increase
in gradient is inversely proportional to the forward variance (which increases
by $N$) because of LayerNorm.

##### Vanishing/Exploding Gradient in Post-LN

While layernorm solves the explosion in the forward pass of networks with
residual connections (De & Smith, 2020), it has the opposite impact on the
gradient. As proved in Section G.2, the gradient in a Post-LN transformer
grows/decays exponentially with the number of layers (Figure 4).

Intuitively, the gradient is first transformed within the layer and then at
the LayerNorm placed before the layer. The multiplicative factor is applied
repeatedly, and causes gradient to vanish or explode exponentially, as was
also observed in Schoenholz et al. (2017). This explains why Post-LN models
are more challenging to train than Pre-LN for deeper networks (Wang et al.,
2022a; Shleifer et al., 2021; Takase et al., 2022).

###  4.2 Explaining impact of large QK values

In Dehghani et al. (2023), the authors observed large QK values destabilized
the training, and solved this empirically by adding a layernorm after
attention scores. Prior work on signal propagation either ignored the effect
of QK initialization (Wang et al., 2022a) or suggest that the backpropagated
gradients have a linear relation to QK variance (Noci et al., 2022).
Critically, note from our derivations of softmax(Section B.7), the backwards
gradients from Q/K are exponentially related to their variance, highlighting
the critical significance of correct initialization of Q/K. For e.g., by
initializing them to only $2$x the xavier values (keep all other
initializations the same), backwards gradients exploded $10000$x through a
$192$ layer model. Our theory explains these empirical observations and
suggests simple initialization strategy to fix this problem, achieving the
same variance on QK without the overhead of LayerNorm.

###  4.3 Explaining and Mitigating Rank Collapse

Similar to our work, Noci et al. (2022) also analyze moment propagation
through the transformer, and observed the rank collapse of the token’s
representations at initialization after just a few layers, i.e., all the token
representations became the same ($r^{l}_{x}\approx 1$ after just $12$ layers)
at initialization. This has also been reported in Shi et al. (2022); Zhou et
al. (2021); Wang et al. (2022b); He et al. (2023); Bachlechner et al. (2021);
Zhai et al. (2023), and suggested modifications such as adding a skip
connection on attention scores, initializing Q/other weights to $0$, or
normalizing all FFN weights.

Figure 5: Forward $r^{l}_{x_{out}}$ for FFN and Attention blocks with $p=0.1$.
FFN reduces $r^{l}_{x_{out}}$ for $r^{l}_{x_{in}}>0.65$, and attention always
has $r^{l}_{x_{out}}<1$.

Figure 6: No rank collapse is observed with Xavier init and dropout. $r^{l}$
increases slower with $\beta^{2}=\frac{2}{N}$ or for DSLM.

Our theory suggests a very simple solution – Dropout. As our closed form
expressions suggest, both FFN block (because of ReLU) and dropout reduce the
correlation (Figure 6). With dropout, our method shows that such a rank
collapse will not occur, and $r^{l}_{x}$ will quickly reach a stable value
$<1$ (Appendix F), as verified empirically in Figure 6.

Alternatively, scaling the block output by $\beta=\frac{1}{\sqrt{N}}$, or
equivalently initializing the weights very small in Post-LN will also prevent
rank collapse, even without Dropout. For Pre-LN, $\lambda=1$ slows down
increase in $r^{l}$ compared to $\lambda^{2}=1-\frac{1}{N}$ (but the same
slowdown can be achieved by decreasing $\beta$). While similar to Noci et al.
(2022), we highlight some issues in Noci et al. (2022) in Appendix F. This
highlights the criticality of correct initialization, dropout and scaling for
deep transformer models, as well as the explainability power of our
theoretical framework.

Table 3: Comparison of maximum theoretical forward (Fwd) pass and backward
(Bck) pass growth in variance for the entire transformer model across methods
(See Appendix G for proofs). Here $\beta$ is the initial value of residual
scaling for LayerScale.

Method Post-LN Pre-LN Bck Sensitivity Fwd Bck Sensitivity DSLM
$\bm{\mathcal{O}(1)}$ $\mathcal{O}(1)$ 1 $\bm{\mathcal{O}(1)}$
$\mathcal{O}(1)$ Vanilla $\mathcal{O}(c^{\pm N})$ $\bm{\mathcal{O}(N)}$
$\mathcal{O}(N)$ $\mathcal{O}(N)$ $\bm{\mathcal{O}(logN)}$ DSInit
$\mathcal{O}(1)$ $\mathcal{O}(N^{-1})$ $\mathcal{O}(1)$ $\mathcal{O}(1)$
$\mathcal{O}(N^{-1})$ LayerScale $\mathcal{O}(1)$ $\mathcal{O}(\beta N)$
$\mathcal{O}(1)$ $\mathcal{O}(1)$ $\mathcal{O}(\beta N)$ DeepNet
$\mathcal{O}(1)$ $\mathcal{O}(N^{-0.5})$ - - -

###  4.4 DeepScaleLM: Enabling Deep Transformers

We propose DeepScaleLM (DSLM), a new initialization / scaling scheme that
alleviates the issues discussed above.

##### Residual/Skip-Connection Scaling

Let $\sigma^{2}_{\text{skip}}$, $\sigma^{2}_{\text{block}}$,
$\sigma^{2}_{\text{model}}$ be the variances of the skip connection, the
block, and the output of the final layer of the model, respectively. Let
$\sigma^{2}_{\text{skip}}=\sigma^{2}_{\text{block}}$, and we scale them by
scalars $\lambda$ and $\beta$ respectively. Then, as has been proven in
numerous works (Section A.3), if $\lambda^{2}+\beta^{2}=1$, this scaling will
maintain the variance after addition of the residual.

##### Initialization

However while ensuring $\sigma^{2}_{\text{skip}}=\sigma^{2}_{\text{block}}$
(and equal to the variance of model input) has been done for ResNets (Section
A.1), it is difficult to achieve theoretically for transformers. By leveraging
the equations in Table 2, our theory provides us the tools to achieve this. We
modify the initialization of the components of the transformer FFN and
Attention blocks such that the variance of their output is $1$, as further
detailed in Appendix N –

  1. 1.

We set the variance of embedding weights as
$\sigma_{e}^{2}=\frac{1-p}{num_{\text{embd}}}$, where $num_{\text{embd}}$ is
the number of embeddings types. As embeddings are followed by a dropout, this
ensures the input variance to the model is $1$.

  2. 2.

We set
$\sigma_{w_{2}}^{2}=\sigma_{w_{1}}^{2}=\frac{1}{d}*\sqrt{\frac{1-p}{2}}$, to
make the output of the FFN block $1$.

  3. 3.

We iteratively calculate layer-by-layer $r^{l}_{x_{\text{in}}}$,
$r^{l}_{x_{\text{out}}}$ using expressions from Table 2, and calculate the
initial variance of the attention block weights to make the output variance
$1$.

This initialization of transformer blocks, combined with the scaling of the
skip connection and residual, and correct initialization of the embeddings,
results $\sigma^{2}_{\text{model}}=1$, irrespective of the number of layer
$N$. This initialization also preserves the backward gradient, as proved for
Pre-LN and Post-LN, in Sections G.3, LABEL: and G.4. Empirically, we show the
backward gradient being preserved for both Pre-LN and Post-LN even across 192
layers at the start of training (Figure 4).

##### Choice of Scaling Parameters

While any choice of $\beta$ will work at initialization, higher values of
$\beta$, for example $\beta^{2}=0.5$ causes gradients to vanish (Figure 10,
Table 4). This is because covariance between residual and skip connection
increases the forward variance, which causes normalization to decrease
backward gradient (De & Smith, 2020).

Similar to other prior works (Section A.3), we use $\beta^{2}=\frac{k}{N}$ in
all our experiments, where $k$ is some small constant. This enables us to
bound the fall in gradient (Section G.3) for Pre-LN. For Post-LN,
$\beta^{2}\leq\frac{k}{N^{2}}$ is theoretically required to bound the gradient
(Section G.6). In practice, with $\beta^{2}=\frac{2}{N}$, even with $768$
layers, we empirically observed the final output variance from the model does
not exceed $30$, and all our models converge. We hence use
$\beta^{2}=\frac{k}{N}$ (Figure 10), but a practitioner may choose
$\beta^{2}=\frac{k}{N^{\alpha}}$, with $\alpha>1$ if more stability is
required at the expense of performance/“sensitivity” (Refer to Section 5.6 and
comparison to prior works in Section 5.5). While the above analysis assumes
positive covariance (which we always observed experimentally), negative
covariance follows a similar reasoning, and will cause gradient explosion
instead.

##### Preventing Rank Collapse

For DSLM, applying block equations iteratively shows that
$r^{l}_{x}<1-\frac{1}{e^{2}}$ after $N$ layers.

##### Simpler Initialization

Another avenue to handle the covariance between residual and skip connection
could be to set $\lambda^{2}+\beta^{2}<1$. We therefore also consider a
simpler initialization method(Figure 14), in which we modify the
initialization of attention value and output matrices to be the same as those
of FFN block. This decreases the "effective" $\beta$ of the attention block,
but as the attention block has $2$x fewer params than FFN, this change in
weightage seems reasonable. As we show in Sections G.5, LABEL: and G.6 while
variances are no longer unit at initialization, they are still bounded. This
change does not impact performance significantly, as we show in Table 13.

##### Folding Scaling into Weights for Inference

The scaling parameters introduced here can be fully absorbed into the model
checkpoint weights by scaling layernorm gain and output linear weights, hence
and do not require any changes to vanilla transformers inference code.

DeepScaleLM enables training deeper-narrower models with $100$s of layers,
outperforming standard models across transformer variants, tasks and
modalities.

##  5 DeepScaleLM Results

###  5.1 Improvements on Encoder-only (BERT) Models

##### Implementation Details

We test our method on the Masked Language Modelling task with the BERT (Devlin
et al., 2019) model. We use Pile-CC dataset (Gao et al., 2021) to pre-train
our model. We use $k=2$ for $\beta$, and we use all original hyper-parameters
of BERT, except for learning rate (LR). We find that higher LR is needed for
our deeper-narrower models (similar to Yang et al. (2021)). Hence, we search
for LR for all the models. The training steps were decided based on Chinchilla
(Hoffmann et al., 2022), at $6.6$B tokens. Table 25 provides all hyper-
parameter details. When using DSLM, model output was down-scaled by $\sqrt{d}$
before LM-head.

We train different language models with the same number of parameters and
compute – while increasing the depth ($N$), we reduce the hidden dimension $d$
keeping number of transformer parameters ($Nd^{2}$) constant. When changing
from $12$-layer $1024$-d model to $192$-layer $256$-d model, compute
negligibly increases by only $6.6\%$ when keeping $Nd^{2}$ constant (Table
24), while the number of parameters decreases by $5-15\%$ due to decreasing
embedding size.

##### Evaluation Metrics

Pre-training Perplexity (exponential of pre-training test-set loss) is often
used to measure MLM pre-training performance (RoBERTa (Liu et al., 2019),
Megatron-LM (Shoeybi et al., 2019), Tay et al. (2023), or similar variants in
Salazar et al. (2020); Lu et al. (2023)), and is well-correlated with
downstream performance (Geiping & Goldstein, 2023). We use the perplexity as
reported by Megatron-LM here. Calling this measure “perplexity” is a slight
abuse of notation (as previous words which are masked are not available, and
future words are). For downstream fine-tuning, we use accuracy.

##### Pre-Training Performance Improvements

In Table 4, we provide the results obtained on scaling model depth after
applying DSLM to Post-LN. PostLN models often diverge while scaling model
depth. DSLM stabilizes the training of Post-LN models, and even a $768$ layer
Post-LN model (with $2300$ Linear and $768$ attention layers) converges.

Our method is comparable to the baseline for shallow models but starts to
outperform as the model gets deeper. Our $192$-layer model outperforms the
vanilla $12$-layer, and our $96$ layer outperforms the vanilla $24$-layer
model. The $160$M $192$-layer model outperforms the vanilla $24$-layer $336$M
model with more than $2\times$ the params.

Reading Table 4 vertically, we can compare the performance of our approach
with the baseline as we vary the model depth ($N$) while keeping the hidden
dimension ($d$) constant. The baseline models often diverge at larger depths.
By stabilizing the training, DSLM allows training larger models with better
performance, with performance consistently increasing with larger depth.

Table 4: Performance (perplexity) of BERT models with different shapes. Deep-
Thin models provide large improvements with fewer parameters.

Model N/D 12/1024 48/512 192/256 768/128 (# Params) (185M) (168M) (160M)
(156M) Baseline 14.2 14.8 17.2 diverge DSLM 15.5 13.1 12.9 18.4 Model N/D
24/1024 96/512 384/256 - (# Params) (336M) (319M) (311M) - Baseline 13.2
diverge diverge - DSLM 14.0 11.7 12.3 -

##### Pre-training Performance Improvements for Pre-LN

We also applied DSLM to the deep Pre-LN models, trained for $3.3$B tokens.
Table 5 show that DSLM significantly improves the performance of the Pre-LN
model across a range of model depths.

Table 5: DSLM with Pre-LN Models.

Model N/D 12/512 96/512 192/256 768/128 Baseline 29.4 20.6 19.8 26.9 DSLM 26.0
15.4 17.0 25.9

##### Performance Improvements after Longer Pre-training

Due to compute limitations, our experiments were trained for Chinchilla
optimal steps. To ensure reproducibility of our work (scripts provided in
supplementary), and demonstrate sustained improvements for standard models, we
trained the BERT-base model using public Wikipedia data for $64$B tokens
($30x$ chinchilla tokens). We train a $4x$ deeper, $10\%$ smaller model using
DSLM ($N$/$d$ = $48$ / $384$). We finetune these models on the public RACE-M,
RACE-H (Lai et al., 2017), MNLI (Williams et al., 2018) and QQP 111Quora
Question Pairs dataset datasets. As shown in Table 6, our model provides
better pretraining performance which is translated into downstream Question-
Answering tasks’ performance across all datasets.

Table 6: BERT-base (trained for $64$B tokens) pre-training and finetuning
results (mean accuracy across 5 runs with stderr).

Dataset Baseline DSLM Pretraining Performance Validation PPL 8.3 7.8
Finetuning Accuracy MNLI 82.4 $\pm$ 0.1 83.7 $\pm$ 0.1 QQP 90.8 $\pm$ 0.03
91.1 $\pm$ 0.05 RACE-Middle 71.1 $\pm$ 0.2 74.0 $\pm$ 0.3 RACE-High 63.7 $\pm$
0.1 65.7 $\pm$ 0.2

###  5.2 Improvements on Decoder-only Models (GPT)

We applied DSLM to the decoder-only GPT model, trained for $8$B tokens
(slightly more than Chinchilla-optimal). Similar to BERT, increasing model
depth by $4x$ with DSLM while keeping the parameters constant results in
improved performance (Table 7).

Table 7: Application of DSLM to Decoder-only model (GPT), while increasing
model depth to 4x (token-level PPL).

Type Model Shape (N/d) LM PPL Baseline DSLM Baseline DSLM Pre-LN 12/1024
48/512 11.6 11.2 Pre-LN 24/1024 96/512 10.4 10.1 Post-LN 12/1024 48/512 12.7
11.7 Post-LN 24/1024 96/512 11.6 10.8

###  5.3 Improvements on Speech (Encoder-Decoder)

We apply DSLM on encoder/decoder style transformer for Speech-to-Text
translation task. The baseline has been trained on the MuST-C (Di Gangi et
al., 2019) dataset using the fairseq (Ott et al., 2019) library. Using DSLM,
we successfully train a $4$x deeper model which outperforms the $18$-layer
($12$-encoder, $6$-decoder layers) baseline with $9\%$ less parameters as seen
in Table 8.

Table 8: Application of DSLM to Speech-to-Text translation.

Model Model Size BLEU Score Nenc–Ndec / d Pre-LN Post-LN Baseline 12–6 / 256
(31.1M) 24.9 21.9 DSLM 48–24 / 128 (28.4M) 25.6 23.8

###  5.4 Improvements on Vision Modality

We also apply our method to Vision Transformers. As we show in Section I.1
using ImageNet-1k (Russakovsky et al., 2015) data with ViT (Dosovitskiy et
al., 2021) model, our method can also constrain the growth of moments in
Vision Transformers. We train our models on the Image Classification task
using ViT baselines provided by Beyer et al. (2022), and trained a $4$x deeper
model with same params. The deeper DSLM model outperforms the baseline ViT
both in both $90$ and $300$ epoch settings. The improvements also translate to
improved robustness on ImageNet-v2 (Recht et al., 2019), ImageNet-R (Hendrycks
et al., 2021) and ImageNet-Sketch (Wang et al., 2019).

Table 9: Applying DSLM to Image classification using ViT.

Eval Set 90-epoch 300-epoch Baseline DSLM Baseline DSLM ImageNet 76.5 77.2
79.8 80.3 ImageNet-Real 83.2 83.8 85.4 85.5 ImageNet-v2 63.7 65.2 67.9 68.3
ImageNet-R 23.9 24.4 27.8 28.3 ImageNet-Sketch 24.4 25.5 28.7 29.9

###  5.5 Comparison with Prior Methods

In Table 10, we compare DSLM with several prior methods for deep transformers.
DSInit and DeepNet stabilize the model training at the expense of reduced
“sensitivity” (Section 5.6) by using smaller effective values of $\beta^{2}$,
at $\mathcal{O}(N^{-2})$ and $\mathcal{O}(N^{-1.5})$ respectively.
Interesting, $96$-layer model diverges with DSInit, despite DSInit using a
smaller $\beta$ asymptotically – this is because the constants hidden in
$\mathcal{O}(N^{-2})$ are much larger for DSInit. Our method, by analysing
signal propagation, has constants exactly $1$.

Bamboo method is a vanilla Pre-LN transformer, which our method out-performs.
SkipInit, ReZero, LayerScale and Value-Skipinit all initialize $\beta$ to
zero/very small values – this choice may slow down learning initially by
reducing back-propagated gradients, and a learnable $\beta$ under-performs
compared to fixed (Table 12). Noci et al. (2022) initializes Query and Key
matrices to a large value, causing divergence (Section 4.2). ADMIN requires an
extra profiling pass through the model, and more importantly, cannot stop
vanishing gradients ( Section A.1), causing the $192$-Layer model to diverge.

Table 10: Comparison with prior methods for deep Transformers.

Method 192/256 96/512 DSInit (Zhang et al., 2019a) 15.9 diverge ADMIN (Liu et
al., 2020a) diverge 25.2 SkipInit (De & Smith, 2020) 15.1 13.1 ReZero
(Bachlechner et al., 2021) diverge diverge LayerScale (Touvron et al., 2021b)
13.2 14.4 DeepNorm (Wang et al., 2022a) 14.4 13.4 Noci et al. (2022) diverge
diverge Bamboo (Xue et al., 2023) 17.1 diverge Value-SkipInit (He et al.,
2023) 18.8 17.1 DeepScaleLM (ours) 12.9 11.7

###  5.6 Analysis of DSLM

##### Model Quantization

Similar to Unit Scaling (Blake et al., 2023), conserving unit activations and
gradients from our method results in models which lose much less performance
when quantized (via direct casting) to FP8 precision compared to original
models. We apply $8$-bit quantization to the $48$-Layer $512$-dim Bert
baseline model and the model trained with DSLM. Table 11 provides the
performance corresponding to the full precision inference and FP8 inferences
(corresponding to two different FP8 standards, E5M2 and E4M3). DSLM model can
be compressed to $25\%$ of the original size with significantly lower
performance loss.

Table 11: Model performance on direct casting to FP8

Model FP32 E5M2 E4M3 Baseline 14.8 42.5 ($\Delta$ 27.7) 16.5 ($\Delta$ 1.7)
DSLM 13.1 21.4 ($\Delta$ 8.3) 13.9 ($\Delta$ 0.8)

##### Ablation of Residual Scaling

Table 12 provides the results corresponding to the different components of our
proposed DSLM scheme for training $96$-layer $512$-d model Post-LN model. The
model fails to converge without the proposed residual scaling. $\beta$ may
also be set as learnable (similar to BatchNorm (Ioffe & Szegedy, 2015)), after
initializing it with $\beta^{2}=\frac{2}{N}$. We find that this does not
significantly impact performance, and $\beta$ remains within $[0.2-5]\times$
of its initialized values.

Table 12: Ablation of various DeepScaleLM components.

Model Perf Vanilla Xavier (with or w/o $\beta^{2}=0.5$) diverge DSLM-Init
(with or w/o $\beta^{2}=0.5$) diverge DSLM-Init + $\beta^{2}=\frac{2}{N}$
(learnable $\beta$) 12.2 DSLM-Init + $\beta^{2}=\frac{2}{N}$ (fixed $\beta$)
11.7

##### Ablation of Initialization

Table 13 provides ablation results for our proposed initialization. All
experiments in Table 13 were conducted for the Pre-LN model with our proposed
scaling ($\lambda,\beta$), since the Post-LN model diverged with Xavier
initialization. Xavier initialization performs significantly worse for very
deep models, due to higher QK initialization. BERT default initialization with
$\sigma=0.02$ also performs worse. Finally, DSLM simpler initialization
performs comparably to DSLM.

Table 13: Ablation of the initializations.

Model Model Size (N/d) Perf Xavier 192/256 (160M) 38.2 DSLM 192/256 (160M)
17.0 DSLM (simple) 192/256 (160M) 17.9 Fixed $\sigma=0.02$ 96/512 (319M) 20.5
DSLM 96/512 (319M) 17.9

##### Compute

Appendix K provides detailed theoretical and wall-clock compute overheads for
making models deeper. We observe that up to $200$ layers, the theoretical
compute is within $6-7$% and wall-clock times is within $15\%$ of the original
shallow model. While our $192$-layer $256$-d model requires $6$% extra compute
than the $12$-layer $1024$-d model, it manages to outperform the $24$-layer
$1024$-d model, that has $62.5$% more parameters, at equal wall-clock time and
at equal number of tokens.

##### Discussion of Relative Strength

In general, for a $\beta$ of the form $\beta^{2}=\frac{k}{N^{\alpha}}$, we can
choose from a wide range of values for the constant $k$ and exponent $\alpha$.
There is an expressivity-trainability trade-off in training deep networks
(Yang & Schoenholz, 2017) – having lower $\beta$ (smaller $k$ or higher
$\alpha$) will result in networks where observed issues (forward growth or
gradient explosion/vanishing) are mitigated, but they may converge slowly/sub-
optimally.

Davis et al. (2021) defines “sensitivity” as the variance of relative change
in output for small perturbations in parameters, averaged across all
parameters. If $\sigma^{2}_{\text{skip}}=1$, sensitivity can be shown to be
mean across layers of $N*(1/\sigma^{2}_{\text{block}})=N*\beta^{2}$. Mean is
not robust to outliers, and hence we suggest median may provide a more robust
measure. For e.g., for vanilla pre-LN, Davis et al. (2021)’s definition gives
sensitivity as $\mathcal{O}(log(N))$, whereas using median provides a more
robust measure as $\mathcal{O}(1)$. But only the first $N/10$ layers have
$\mathcal{O}(log(N))$ sensitivity, and the last $9N/10$ layers have
$\mathcal{O}(1)$ sensitivity. We will use median in the discussion below.

In Appendix L, we show that the fall in gradient for both pre-LN and post-LN
for $\beta^{2}=k/N^{\alpha}$ is $\mathcal{O}(e^{kN^{1-\alpha}})$. The
sensitivity is hence $kN^{1-\alpha}$. For DSLM, we chose $\alpha=1$, that is
the sweet spot on the stability-expressivity curve where both the gradient
fall bound and sensitivity expressions become independent of model depth. For
higher values of $\alpha$ such as $\alpha=2$ (DS-Init) and, $\alpha=1.5$
(DeepNet), the gradient becomes stable using but the model expressivity
reduces with depth, as shown in Table 3. Such models might not be able to
extract better results when going deeper, as we indeed verify empirically in
the comparison with prior works paragraph in Section 5.5.

##  6 Conclusion

We theoretically derive closed forms for the growth of variances for forward
and backward pass through individual transformer components as well as the
entire transformer model. These formulae enable us to identify and solve the
key reasons for vanishing/exploding gradients and rank collapse in very deep
transformers. Via scaling and correct initialization, we also enable training
very deep transformers up to $768$ layers. Our experiments suggest that deeper
transformers should be explored – using our method, models with $100$s of
layers outperform larger standard models across multiple modalities, tasks,
and transformer variants.

## Ethics Statement

This paper presents work whose goal is to advance the field of Machine
Learning. There are many potential societal consequences of our work, some
which we feel must be specifically highlighted here. Using crawled web data
for pre-training language models is questionable, something which society has
yet to finalize its views on. Language modelling in particular suffers from
hallucinations, and may be used for misinformation.

## References

  * Arora et al. (2019a) Arora, S., Du, S. S., Hu, W., Li, Z., Salakhutdinov, R., and Wang, R.  On exact computation with an infinitely wide neural net.  In Wallach, H. M., Larochelle, H., Beygelzimer, A., d’Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , pp. 8139–8148, 2019a.  URL https://proceedings.neurips.cc/paper/2019/hash/dbc4d84bfcfe2284ba11beffb853a8c4-Abstract.html.
  * Arora et al. (2019b) Arora, S., Du, S. S., Hu, W., Li, Z., and Wang, R.  Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.  In Chaudhuri, K. and Salakhutdinov, R. (eds.), _Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 322–332. PMLR, 2019b.  URL http://proceedings.mlr.press/v97/arora19a.html.
  * Arpit et al. (2016) Arpit, D., Zhou, Y., Kota, B. U., and Govindaraju, V.  Normalization propagation: A parametric technique for removing internal covariate shift in deep networks.  In Balcan, M. and Weinberger, K. Q. (eds.), _Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016_ , volume 48 of _JMLR Workshop and Conference Proceedings_ , pp. 1168–1176. JMLR.org, 2016.  URL http://proceedings.mlr.press/v48/arpitb16.html.
  * Arpit et al. (2019) Arpit, D., Campos, V., and Bengio, Y.  How to initialize your network? robust initialization for weightnorm i& resnets.  In Wallach, H. M., Larochelle, H., Beygelzimer, A., d’Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , pp. 10900–10909, 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/e520f70ac3930490458892665cda6620-Abstract.html.
  * Bachlechner et al. (2021) Bachlechner, T., Majumder, B. P., Mao, H. H., Cottrell, G., and McAuley, J. J.  Rezero is all you need: fast convergence at large depth.  In de Campos, C. P., Maathuis, M. H., and Quaeghebeur, E. (eds.), _Proceedings of the Thirty-Seventh Conference on Uncertainty in Artificial Intelligence, UAI 2021, Virtual Event, 27-30 July 2021_ , volume 161 of _Proceedings of Machine Learning Research_ , pp. 1352–1361. AUAI Press, 2021.  URL https://proceedings.mlr.press/v161/bachlechner21a.html.
  * Balduzzi et al. (2017) Balduzzi, D., Frean, M., Leary, L., Lewis, J. P., Ma, K. W., and McWilliams, B.  The shattered gradients problem: If resnets are the answer, then what is the question?  In Precup, D. and Teh, Y. W. (eds.), _Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017_ , volume 70 of _Proceedings of Machine Learning Research_ , pp. 342–350. PMLR, 2017.  URL http://proceedings.mlr.press/v70/balduzzi17b.html.
  * Beyer et al. (2022) Beyer, L., Zhai, X., and Kolesnikov, A.  Better plain vit baselines for imagenet-1k, 2022.
  * Blake et al. (2023) Blake, C., Orr, D., and Luschi, C.  Unit scaling: Out-of-the-box low-precision training.  In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), _International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 2548–2576. PMLR, 2023.  URL https://proceedings.mlr.press/v202/blake23a.html.
  * Brock et al. (2021) Brock, A., De, S., and Smith, S. L.  Characterizing signal propagation to close the performance gap in unnormalized resnets.  In _9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021_. OpenReview.net, 2021.  URL https://openreview.net/forum?id=IX3Nnir2omJ.
  * Chen et al. (2020) Chen, M., Wei, Z., Huang, Z., Ding, B., and Li, Y.  Simple and deep graph convolutional networks.  In _Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 1725–1735. PMLR, 2020.  URL http://proceedings.mlr.press/v119/chen20v.html.
  * Chowdhery et al. (2023) Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., Schuh, P., Shi, K., Tsvyashchenko, S., Maynez, J., Rao, A., Barnes, P., Tay, Y., Shazeer, N., Prabhakaran, V., Reif, E., Du, N., Hutchinson, B., Pope, R., Bradbury, J., Austin, J., Isard, M., Gur-Ari, G., Yin, P., Duke, T., Levskaya, A., Ghemawat, S., Dev, S., Michalewski, H., Garcia, X., Misra, V., Robinson, K., Fedus, L., Zhou, D., Ippolito, D., Luan, D., Lim, H., Zoph, B., Spiridonov, A., Sepassi, R., Dohan, D., Agrawal, S., Omernick, M., Dai, A. M., Pillai, T. S., Pellat, M., Lewkowycz, A., Moreira, E., Child, R., Polozov, O., Lee, K., Zhou, Z., Wang, X., Saeta, B., Diaz, M., Firat, O., Catasta, M., Wei, J., Meier-Hellstern, K., Eck, D., Dean, J., Petrov, S., and Fiedel, N.  Palm: Scaling language modeling with pathways.  _J. Mach. Learn. Res._ , 24:240:1–240:113, 2023.  URL http://jmlr.org/papers/v24/22-1144.html.
  * Clark et al. (2020) Clark, K., Luong, M., Le, Q. V., and Manning, C. D.  ELECTRA: pre-training text encoders as discriminators rather than generators.  In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net, 2020.  URL https://openreview.net/forum?id=r1xMH1BtvB.
  * Davis et al. (2021) Davis, J. Q., Gu, A., Choromanski, K., Dao, T., Ré, C., Finn, C., and Liang, P.  Catformer: Designing stable transformers via sensitivity analysis.  In Meila, M. and Zhang, T. (eds.), _Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 2489–2499. PMLR, 2021.  URL http://proceedings.mlr.press/v139/davis21a.html.
  * De & Smith (2020) De, S. and Smith, S. L.  Batch normalization biases residual blocks towards the identity function in deep networks.  In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual_ , 2020.  URL https://proceedings.neurips.cc/paper/2020/hash/e6b738eca0e6792ba8a9cbcba6c1881d-Abstract.html.
  * Dehghani et al. (2023) Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., Steiner, A. P., Caron, M., Geirhos, R., Alabdulmohsin, I., Jenatton, R., Beyer, L., Tschannen, M., Arnab, A., Wang, X., Ruiz, C. R., Minderer, M., Puigcerver, J., Evci, U., Kumar, M., van Steenkiste, S., Elsayed, G. F., Mahendran, A., Yu, F., Oliver, A., Huot, F., Bastings, J., Collier, M., Gritsenko, A. A., Birodkar, V., Vasconcelos, C. N., Tay, Y., Mensink, T., Kolesnikov, A., Pavetic, F., Tran, D., Kipf, T., Lucic, M., Zhai, X., Keysers, D., Harmsen, J. J., and Houlsby, N.  Scaling vision transformers to 22 billion parameters.  In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), _International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 7480–7512. PMLR, 2023.  URL https://proceedings.mlr.press/v202/dehghani23a.html.
  * Dettmers et al. (2023) Dettmers, T., Pagnoni, A., Holtzman, A., and Zettlemoyer, L.  QLoRA: Efficient finetuning of quantized LLMs.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.  URL https://openreview.net/forum?id=OUIFPHEgJU.
  * Devlin et al. (2019) Devlin, J., Chang, M., Lee, K., and Toutanova, K.  BERT: pre-training of deep bidirectional transformers for language understanding.  In Burstein, J., Doran, C., and Solorio, T. (eds.), _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers)_ , pp. 4171–4186. Association for Computational Linguistics, 2019.  doi: 10.18653/V1/N19-1423.  URL https://doi.org/10.18653/v1/n19-1423.
  * Di Gangi et al. (2019) Di Gangi, M. A., Cattoni, R., Bentivogli, L., Negri, M., and Turchi, M.  MuST-C: a Multilingual Speech Translation Corpus.  In Burstein, J., Doran, C., and Solorio, T. (eds.), _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pp. 2012–2017, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.  doi: 10.18653/v1/N19-1202.  URL https://aclanthology.org/N19-1202.
  * Dong et al. (2021) Dong, Y., Cordonnier, J., and Loukas, A.  Attention is not all you need: pure attention loses rank doubly exponentially with depth.  In Meila, M. and Zhang, T. (eds.), _Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 2793–2803. PMLR, 2021.  URL http://proceedings.mlr.press/v139/dong21a.html.
  * Dosovitskiy et al. (2021) Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021_. OpenReview.net, 2021.  URL https://openreview.net/forum?id=YicbFdNTTy.
  * Gao et al. (2021) Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C.  The pile: An 800gb dataset of diverse text for language modeling.  _CoRR_ , abs/2101.00027, 2021.  URL https://arxiv.org/abs/2101.00027.
  * Geiping & Goldstein (2023) Geiping, J. and Goldstein, T.  Cramming: Training a language model on a single GPU in one day.  In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), _International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 11117–11143. PMLR, 2023.  URL https://proceedings.mlr.press/v202/geiping23a.html.
  * Glorot & Bengio (2010) Glorot, X. and Bengio, Y.  Understanding the difficulty of training deep feedforward neural networks.  In Teh, Y. W. and Titterington, D. M. (eds.), _Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, AISTATS 2010, Chia Laguna Resort, Sardinia, Italy, May 13-15, 2010_ , volume 9 of _JMLR Proceedings_ , pp. 249–256. JMLR.org, 2010.  URL http://proceedings.mlr.press/v9/glorot10a.html.
  * Hanin & Rolnick (2018) Hanin, B. and Rolnick, D.  How to start training: The effect of initialization and architecture.  In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada_ , pp. 569–579, 2018.  URL https://proceedings.neurips.cc/paper/2018/hash/d81f9c1be2e08964bf9f24b15f0e4900-Abstract.html.
  * Hayou et al. (2019) Hayou, S., Doucet, A., and Rousseau, J.  On the impact of the activation function on deep neural networks training.  In Chaudhuri, K. and Salakhutdinov, R. (eds.), _Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 2672–2680. PMLR, 2019.  URL http://proceedings.mlr.press/v97/hayou19a.html.
  * He et al. (2023) He, B., Martens, J., Zhang, G., Botev, A., Brock, A., Smith, S. L., and Teh, Y. W.  Deep transformers without shortcuts: Modifying self-attention for faithful signal propagation.  In _The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023_. OpenReview.net, 2023.  URL https://openreview.net/pdf?id=NPrsUQgMjKK.
  * He et al. (2015) He, K., Zhang, X., Ren, S., and Sun, J.  Delving deep into rectifiers: Surpassing human-level performance on imagenet classification.  In _2015 IEEE International Conference on Computer Vision, ICCV 2015, Santiago, Chile, December 7-13, 2015_ , pp. 1026–1034. IEEE Computer Society, 2015.  doi: 10.1109/ICCV.2015.123.  URL https://doi.org/10.1109/ICCV.2015.123.
  * Hendrycks et al. (2021) Hendrycks, D., Basart, S., Mu, N., Kadavath, S., Wang, F., Dorundo, E., Desai, R., Zhu, T., Parajuli, S., Guo, M., Song, D., Steinhardt, J., and Gilmer, J.  The many faces of robustness: A critical analysis of out-of-distribution generalization.  In _2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10-17, 2021_ , pp. 8320–8329. IEEE, 2021.  doi: 10.1109/ICCV48922.2021.00823.  URL https://doi.org/10.1109/ICCV48922.2021.00823.
  * Hoedt et al. (2022) Hoedt, P.-J., Hochreiter, S., and Klambauer, G.  Normalisation is dead, long live normalisation!  In _ICLR Blog Track_ , 2022.  URL https://iclr-blog-track.github.io/2022/03/25/unnormalized-resnets/.  https://iclr-blog-track.github.io/2022/03/25/unnormalized-resnets/.
  * Hoffmann et al. (2022) Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Rae, J. W., Vinyals, O., and Sifre, L.  Training compute-optimal large language models.  _CoRR_ , abs/2203.15556, 2022.  doi: 10.48550/ARXIV.2203.15556.  URL https://doi.org/10.48550/arXiv.2203.15556.
  * Huang et al. (2020) Huang, X. S., Perez, F., Ba, J., and Volkovs, M.  Improving Transformer Optimization Through Better Initialization.  In _Proceedings of the 37th International Conference on Machine Learning_ , pp. 4475–4483. PMLR, November 2020.  URL https://proceedings.mlr.press/v119/huang20f.html.
  * Ioffe & Szegedy (2015) Ioffe, S. and Szegedy, C.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In Bach, F. R. and Blei, D. M. (eds.), _Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015_ , volume 37 of _JMLR Workshop and Conference Proceedings_ , pp. 448–456. JMLR.org, 2015.  URL http://proceedings.mlr.press/v37/ioffe15.html.
  * Jesus et al. (2020) Jesus, R. J., Antunes, M. L. P., da Costa, R. A. F., Dorogovtsev, S. N., Mendes, J. F. F., and Aguiar, R. L.  Effect of the initial configuration of weights on the training and function of artificial neural networks.  _CoRR_ , abs/2012.02550, 2020.  URL https://arxiv.org/abs/2012.02550.
  * Kim et al. (2021) Kim, H., Papamakarios, G., and Mnih, A.  The lipschitz constant of self-attention.  In Meila, M. and Zhang, T. (eds.), _Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 5562–5571. PMLR, 2021.  URL http://proceedings.mlr.press/v139/kim21i.html.
  * Kingsley (1935) Kingsley, Z. G.  _The psycho-biology of language: an introduction to dynamic philology_.  Houghton Mifflin, 1935.
  * Klambauer et al. (2017) Klambauer, G., Unterthiner, T., Mayr, A., and Hochreiter, S.  Self-normalizing neural networks.  In Guyon, I., von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , pp. 971–980, 2017.  URL https://proceedings.neurips.cc/paper/2017/hash/5d44ee6f2c3f71b73125876103c8f6c4-Abstract.html.
  * Korotkov & Korotkov (2020) Korotkov, N. E. and Korotkov, A. N.  _Integrals related to the error function_.  Chapman & Hall/CRC, Philadelphia, PA, March 2020.  ISBN 9780367408206.  URL https://www.taylorfrancis.com/books/mono/10.1201/9780367809232/integrals-related-error-function-nikolai-korotkov-alexander-korotkov.
  * Labatie et al. (2021) Labatie, A., Masters, D., Eaton-Rosen, Z., and Luschi, C.  Proxy-normalizing activations to match batch normalization while removing batch dependence.  In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (eds.), _Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual_ , pp. 16990–17006, 2021.  URL https://proceedings.neurips.cc/paper/2021/hash/8d2a5f7d4afa5d0530789d3066945330-Abstract.html.
  * Lai et al. (2017) Lai, G., Xie, Q., Liu, H., Yang, Y., and Hovy, E. H.  RACE: large-scale reading comprehension dataset from examinations.  In Palmer, M., Hwa, R., and Riedel, S. (eds.), _Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, EMNLP 2017, Copenhagen, Denmark, September 9-11, 2017_ , pp. 785–794. Association for Computational Linguistics, 2017.  doi: 10.18653/V1/D17-1082.  URL https://doi.org/10.18653/v1/d17-1082.
  * LeCun et al. (1996) LeCun, Y., Bottou, L., Orr, G. B., and Müller, K.  Effiicient backprop.  In Orr, G. B. and Müller, K. (eds.), _Neural Networks: Tricks of the Trade_ , volume 1524 of _Lecture Notes in Computer Science_ , pp. 9–50. Springer, 1996.  doi: 10.1007/3-540-49430-8\\_2.  URL https://doi.org/10.1007/3-540-49430-8_2.
  * Lee et al. (2019) Lee, J., Xiao, L., Schoenholz, S. S., Bahri, Y., Novak, R., Sohl-Dickstein, J., and Pennington, J.  Wide neural networks of any depth evolve as linear models under gradient descent.  In Wallach, H. M., Larochelle, H., Beygelzimer, A., d’Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , pp. 8570–8581, 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/0d1a9651497a38d8b1c3871c84528bd4-Abstract.html.
  * Li & Liang (2018) Li, Y. and Liang, Y.  Learning overparameterized neural networks via stochastic gradient descent on structured data.  In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada_ , pp. 8168–8177, 2018.  URL https://proceedings.neurips.cc/paper/2018/hash/54fe976ba170c19ebae453679b362263-Abstract.html.
  * Liu et al. (2020a) Liu, L., Liu, X., Gao, J., Chen, W., and Han, J.  Understanding the difficulty of training transformers.  In Webber, B., Cohn, T., He, Y., and Liu, Y. (eds.), _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020_ , pp. 5747–5763. Association for Computational Linguistics, 2020a.  doi: 10.18653/V1/2020.EMNLP-MAIN.463.  URL https://doi.org/10.18653/v1/2020.emnlp-main.463.
  * Liu et al. (2020b) Liu, X., Duh, K., Liu, L., and Gao, J.  Very deep transformers for neural machine translation.  _CoRR_ , abs/2008.07772, 2020b.  URL https://arxiv.org/abs/2008.07772.
  * Liu et al. (2019) Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V.  Roberta: A robustly optimized BERT pretraining approach.  _CoRR_ , abs/1907.11692, 2019.  URL http://arxiv.org/abs/1907.11692.
  * Lo (2013) Lo, C. F.  WKB approximation for the sum of two correlated lognormal random variables.  _Applied Mathematical Sciences_ , 7:6355–6367, 2013.  ISSN 13147552.  doi: 10.12988/ams.2013.39511.  URL http://www.m-hikari.com/ams/ams-2013/ams-125-128-2013/39511.html.
  * Lu et al. (2023) Lu, J., Zhu, D., Han, W., Zhao, R., Namee, B. M., and Tan, F.  What makes pre-trained language models better zero-shot learners?  In Rogers, A., Boyd-Graber, J. L., and Okazaki, N. (eds.), _Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2023, Toronto, Canada, July 9-14, 2023_ , pp. 2288–2303. Association for Computational Linguistics, 2023.  doi: 10.18653/V1/2023.ACL-LONG.128.  URL https://doi.org/10.18653/v1/2023.acl-long.128.
  * Marion et al. (2022) Marion, P., Fermanian, A., Biau, G., and Vert, J.  Scaling resnets in the large-depth regime.  _CoRR_ , abs/2206.06929, 2022.  doi: 10.48550/ARXIV.2206.06929.  URL https://doi.org/10.48550/arXiv.2206.06929.
  * Martens et al. (2021) Martens, J., Ballard, A., Desjardins, G., Swirszcz, G., Dalibard, V., Sohl-Dickstein, J., and Schoenholz, S. S.  Rapid training of deep neural networks without skip connections or normalization layers using deep kernel shaping.  _CoRR_ , abs/2110.01765, 2021.  URL https://arxiv.org/abs/2110.01765.
  * Mishkin & Matas (2016) Mishkin, D. and Matas, J.  All you need is a good init.  In Bengio, Y. and LeCun, Y. (eds.), _4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings_ , 2016.  URL http://arxiv.org/abs/1511.06422.
  * Montúfar et al. (2014) Montúfar, G. F., Pascanu, R., Cho, K., and Bengio, Y.  On the number of linear regions of deep neural networks.  In Ghahramani, Z., Welling, M., Cortes, C., Lawrence, N. D., and Weinberger, K. Q. (eds.), _Advances in Neural Information Processing Systems 27: Annual Conference on Neural Information Processing Systems 2014, December 8-13 2014, Montreal, Quebec, Canada_ , pp. 2924–2932, 2014.  URL https://proceedings.neurips.cc/paper/2014/hash/109d2dd3608f669ca17920c511c2a41e-Abstract.html.
  * Neal (1995) Neal, R. M.  _Bayesian learning for neural networks_.  PhD thesis, University of Toronto, Canada, 1995.  URL https://librarysearch.library.utoronto.ca/permalink/01UTORONTO_INST/14bjeso/alma991106438365706196.
  * Ng & Geller (1969) Ng, E. W. and Geller, M.  A table of integrals of the Error functions.  _Journal of Research of the National Bureau of Standards, Section B: Mathematical Sciences_ , 73B(1):1, 1969.  ISSN 0098-8979.  doi: 10/gdtk9p.  URL https://nvlpubs.nist.gov/nistpubs/jres/73B/jresv73Bn1p1_A1b.pdf.
  * Nguyen & Salazar (2019) Nguyen, T. Q. and Salazar, J.  Transformers without tears: Improving the normalization of self-attention.  In Niehues, J., Cattoni, R., Stüker, S., Negri, M., Turchi, M., Ha, T., Salesky, E., Sanabria, R., Barrault, L., Specia, L., and Federico, M. (eds.), _Proceedings of the 16th International Conference on Spoken Language Translation, IWSLT 2019, Hong Kong, November 2-3, 2019_. Association for Computational Linguistics, 2019.  URL https://aclanthology.org/2019.iwslt-1.17.
  * Noci et al. (2022) Noci, L., Anagnostidis, S., Biggio, L., Orvieto, A., Singh, S. P., and Lucchi, A.  Signal propagation in transformers: Theoretical perspectives and the role of rank collapse.  In _NeurIPS_ , 2022.  URL http://papers.nips.cc/paper
_files/paper/2022/hash/ae0cba715b60c4052359b3d52a2cff7f-Abstract-
Conference.html.

  * Ott et al. (2019) Ott, M., Edunov, S., Baevski, A., Fan, A., Gross, S., Ng, N., Grangier, D., and Auli, M.  fairseq: A fast, extensible toolkit for sequence modeling, 2019.
  * Poole et al. (2016) Poole, B., Lahiri, S., Raghu, M., Sohl-Dickstein, J., and Ganguli, S.  Exponential expressivity in deep neural networks through transient chaos.  In Lee, D. D., Sugiyama, M., von Luxburg, U., Guyon, I., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain_ , pp. 3360–3368, 2016.  URL https://proceedings.neurips.cc/paper/2016/hash/148510031349642de5ca0c544f31b2ef-Abstract.html.
  * Raghu et al. (2017) Raghu, M., Poole, B., Kleinberg, J. M., Ganguli, S., and Sohl-Dickstein, J.  On the expressive power of deep neural networks.  In Precup, D. and Teh, Y. W. (eds.), _Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017_ , volume 70 of _Proceedings of Machine Learning Research_ , pp. 2847–2854. PMLR, 2017.  URL http://proceedings.mlr.press/v70/raghu17a.html.
  * Recht et al. (2019) Recht, B., Roelofs, R., Schmidt, L., and Shankar, V.  Do imagenet classifiers generalize to imagenet?  _CoRR_ , abs/1902.10811, 2019.  URL http://arxiv.org/abs/1902.10811.
  * Rong et al. (2020) Rong, Y., Huang, W., Xu, T., and Huang, J.  Dropedge: Towards deep graph convolutional networks on node classification.  In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net, 2020.  URL https://openreview.net/forum?id=Hkx1qkrKPr.
  * Russakovsky et al. (2015) Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M. S., Berg, A. C., and Fei-Fei, L.  Imagenet large scale visual recognition challenge.  _Int. J. Comput. Vis._ , 115(3):211–252, 2015.  doi: 10.1007/S11263-015-0816-Y.  URL https://doi.org/10.1007/s11263-015-0816-y.
  * Salazar et al. (2020) Salazar, J., Liang, D., Nguyen, T. Q., and Kirchhoff, K.  Masked language model scoring.  In Jurafsky, D., Chai, J., Schluter, N., and Tetreault, J. R. (eds.), _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020_ , pp. 2699–2712. Association for Computational Linguistics, 2020.  doi: 10.18653/V1/2020.ACL-MAIN.240.  URL https://doi.org/10.18653/v1/2020.acl-main.240.
  * Schoenholz et al. (2017) Schoenholz, S. S., Gilmer, J., Ganguli, S., and Sohl-Dickstein, J.  Deep information propagation.  In _5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings_. OpenReview.net, 2017.  URL https://openreview.net/forum?id=H1W1UN9gg.
  * Shen et al. (2020) Shen, S., Yao, Z., Gholami, A., Mahoney, M. W., and Keutzer, K.  Powernorm: Rethinking batch normalization in transformers.  In _Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 8741–8751. PMLR, 2020.  URL http://proceedings.mlr.press/v119/shen20e.html.
  * Shi et al. (2022) Shi, H., Gao, J., Xu, H., Liang, X., Li, Z., Kong, L., Lee, S. M. S., and Kwok, J. T.  Revisiting over-smoothing in BERT from the perspective of graph.  In _The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022_. OpenReview.net, 2022.  URL https://openreview.net/forum?id=dUV91uaXm3.
  * Shleifer et al. (2021) Shleifer, S., Weston, J., and Ott, M.  Normformer: Improved transformer pretraining with extra normalization.  _CoRR_ , abs/2110.09456, 2021.  URL https://arxiv.org/abs/2110.09456.
  * Shoeybi et al. (2019) Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B.  Megatron-lm: Training multi-billion parameter language models using model parallelism.  _CoRR_ , abs/1909.08053, 2019.  URL http://arxiv.org/abs/1909.08053.
  * Smith et al. (2022) Smith, S., Patwary, M., Norick, B., LeGresley, P., Rajbhandari, S., Casper, J., Liu, Z., Prabhumoye, S., Zerveas, G., Korthikanti, V., Zheng, E., Child, R., Aminabadi, R. Y., Bernauer, J., Song, X., Shoeybi, M., He, Y., Houston, M., Tiwary, S., and Catanzaro, B.  Using deepspeed and megatron to train megatron-turing NLG 530b, A large-scale generative language model.  _CoRR_ , abs/2201.11990, 2022.  URL https://arxiv.org/abs/2201.11990.
  * Szegedy et al. (2017) Szegedy, C., Ioffe, S., Vanhoucke, V., and Alemi, A. A.  Inception-v4, inception-resnet and the impact of residual connections on learning.  In Singh, S. P. and Markovitch, S. (eds.), _Proceedings of the Thirty-First AAAI Conference on Artificial Intelligence, February 4-9, 2017, San Francisco, California, USA_ , pp. 4278–4284. AAAI Press, 2017.  URL http://aaai.org/ocs/index.php/AAAI/AAAI17/paper/view/14806.
  * Takase et al. (2022) Takase, S., Kiyono, S., Kobayashi, S., and Suzuki, J.  On layer normalizations and residual connections in transformers.  _CoRR_ , abs/2206.00330, 2022.  doi: 10.48550/ARXIV.2206.00330.  URL https://doi.org/10.48550/arXiv.2206.00330.
  * Tay et al. (2023) Tay, Y., Dehghani, M., Abnar, S., Chung, H. W., Fedus, W., Rao, J., Narang, S., Tran, V. Q., Yogatama, D., and Metzler, D.  Scaling laws vs model architectures: How does inductive bias influence scaling?  In Bouamor, H., Pino, J., and Bali, K. (eds.), _Findings of the Association for Computational Linguistics: EMNLP 2023, Singapore, December 6-10, 2023_ , pp. 12342–12364. Association for Computational Linguistics, 2023.  URL https://aclanthology.org/2023.findings-emnlp.825.
  * Touvron et al. (2021a) Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and Jégou, H.  Training data-efficient image transformers & distillation through attention.  In Meila, M. and Zhang, T. (eds.), _Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 10347–10357. PMLR, 2021a.  URL http://proceedings.mlr.press/v139/touvron21a.html.
  * Touvron et al. (2021b) Touvron, H., Cord, M., Sablayrolles, A., Synnaeve, G., and Jégou, H.  Going deeper with image transformers.  In _2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10-17, 2021_ , pp. 32–42. IEEE, 2021b.  doi: 10.1109/ICCV48922.2021.00010.  URL https://doi.org/10.1109/ICCV48922.2021.00010.
  * Wang et al. (2019) Wang, H., Ge, S., Lipton, Z. C., and Xing, E. P.  Learning robust global representations by penalizing local predictive power.  In Wallach, H. M., Larochelle, H., Beygelzimer, A., d’Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , pp. 10506–10518, 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/3eefceb8087e964f89c2d59e8a249915-Abstract.html.
  * Wang et al. (2022a) Wang, H., Ma, S., Dong, L., Huang, S., Zhang, D., and Wei, F.  Deepnet: Scaling transformers to 1, 000 layers.  _CoRR_ , abs/2203.00555, 2022a.  doi: 10.48550/ARXIV.2203.00555.  URL https://doi.org/10.48550/arXiv.2203.00555.
  * Wang et al. (2022b) Wang, P., Zheng, W., Chen, T., and Wang, Z.  Anti-oversmoothing in deep vision transformers via the fourier domain analysis: From theory to practice.  In _The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022_. OpenReview.net, 2022b.  URL https://openreview.net/forum?id=O476oWmiNNp.
  * Williams et al. (2018) Williams, A., Nangia, N., and Bowman, S. R.  A broad-coverage challenge corpus for sentence understanding through inference.  In Walker, M. A., Ji, H., and Stent, A. (eds.), _Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2018, New Orleans, Louisiana, USA, June 1-6, 2018, Volume 1 (Long Papers)_ , pp. 1112–1122. Association for Computational Linguistics, 2018.  doi: 10.18653/V1/N18-1101.  URL https://doi.org/10.18653/v1/n18-1101.
  * Xiong et al. (2020) Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L., and Liu, T.  On layer normalization in the transformer architecture.  In _Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 10524–10533. PMLR, 2020.  URL http://proceedings.mlr.press/v119/xiong20b.html.
  * Xu et al. (2019) Xu, J., Sun, X., Zhang, Z., Zhao, G., and Lin, J.  Understanding and improving layer normalization.  In Wallach, H. M., Larochelle, H., Beygelzimer, A., d’Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , pp. 4383–4393, 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/2f4fe03d77724a7217006e5d16728874-Abstract.html.
  * Xue et al. (2023) Xue, F., Chen, J., Sun, A., Ren, X., Zheng, Z., He, X., Chen, Y., Jiang, X., and You, Y.  A study on transformer configuration and training objective.  In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), _International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 38913–38925. PMLR, 2023.  URL https://proceedings.mlr.press/v202/xue23b.html.
  * Yang & Schoenholz (2017) Yang, G. and Schoenholz, S. S.  Mean field residual networks: On the edge of chaos.  In Guyon, I., von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , pp. 7103–7114, 2017.  URL https://proceedings.neurips.cc/paper/2017/hash/81c650caac28cdefce4de5ddc18befa0-Abstract.html.
  * Yang et al. (2021) Yang, G., Hu, E. J., Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, J.  Tuning large neural networks via zero-shot hyperparameter transfer.  In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (eds.), _Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual_ , pp. 17084–17097, 2021.  URL https://proceedings.neurips.cc/paper/2021/hash/8df7c2e3c3c3be098ef7b382bd2c37ba-Abstract.html.
  * You et al. (2017) You, Y., Gitman, I., and Ginsburg, B.  Scaling SGD batch size to 32k for imagenet training.  _CoRR_ , abs/1708.03888, 2017.  URL http://arxiv.org/abs/1708.03888.
  * You et al. (2020) You, Y., Li, J., Reddi, S. J., Hseu, J., Kumar, S., Bhojanapalli, S., Song, X., Demmel, J., Keutzer, K., and Hsieh, C.  Large batch optimization for deep learning: Training BERT in 76 minutes.  In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net, 2020.  URL https://openreview.net/forum?id=Syx4wnEtvH.
  * Yu et al. (2017) Yu, A. W., Lin, Q., Salakhutdinov, R., and Carbonell, J. G.  Normalized gradient with adaptive stepsize method for deep neural network training.  _CoRR_ , abs/1707.04822, 2017.  URL http://arxiv.org/abs/1707.04822.
  * Zhai et al. (2023) Zhai, S., Likhomanenko, T., Littwin, E., Busbridge, D., Ramapuram, J., Zhang, Y., Gu, J., and Susskind, J. M.  Stabilizing transformer training by preventing attention entropy collapse.  In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), _International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 40770–40803. PMLR, 2023.  URL https://proceedings.mlr.press/v202/zhai23a.html.
  * Zhang et al. (2019a) Zhang, B., Titov, I., and Sennrich, R.  Improving deep transformer with depth-scaled initialization and merged attention.  In Inui, K., Jiang, J., Ng, V., and Wan, X. (eds.), _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019_ , pp. 898–909. Association for Computational Linguistics, 2019a.  doi: 10.18653/V1/D19-1083.  URL https://doi.org/10.18653/v1/D19-1083.
  * Zhang et al. (2019b) Zhang, H., Dauphin, Y. N., and Ma, T.  Fixup initialization: Residual learning without normalization.  In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_. OpenReview.net, 2019b.  URL https://openreview.net/forum?id=H1gsz30cKX.
  * Zhang et al. (2022) Zhang, H., Yu, D., Yi, M., Chen, W., and Liu, T.  Stabilize deep resnet with a sharp scaling factor $\tau$.  _Mach. Learn._ , 111(9):3359–3392, 2022.  doi: 10.1007/S10994-022-06192-X.  URL https://doi.org/10.1007/s10994-022-06192-x.
  * Zhao et al. (2023) Zhao, H., Ma, S., Zhang, D., Deng, Z., and Wei, F.  Are more layers beneficial to graph transformers?  In _The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023_. OpenReview.net, 2023.  URL https://openreview.net/pdf?id=uagC-X9XMi8.
  * Zhou et al. (2021) Zhou, D., Kang, B., Jin, X., Yang, L., Lian, X., Hou, Q., and Feng, J.  Deepvit: Towards deeper vision transformer.  _CoRR_ , abs/2103.11886, 2021.  URL https://arxiv.org/abs/2103.11886.

###### Contents

  1. 1 Introduction
  2. 2 Related Works
  3. 3 Moments of Transformer Models
    1. 3.1 Moments of Transformer Components
    2. 3.2 Moments of Transformer Blocks
    3. 3.3 Moments of Entire Transformer Model
    4. 3.4 Numerical Validation of Theoretical Results
    5. 3.5 Validity of Theoretical Predictions even after Training
  4. 4 Applications
    1. 4.1 Explaining Variance Explosion in Transformer
    2. 4.2 Explaining impact of large QK values
    3. 4.3 Explaining and Mitigating Rank Collapse
    4. 4.4 DeepScaleLM: Enabling Deep Transformers
  5. 5 DeepScaleLM Results
    1. 5.1 Improvements on Encoder-only (BERT) Models
    2. 5.2 Improvements on Decoder-only Models (GPT)
    3. 5.3 Improvements on Speech (Encoder-Decoder)
    4. 5.4 Improvements on Vision Modality
    5. 5.5 Comparison with Prior Methods
    6. 5.6 Analysis of DSLM
  6. 6 Conclusion
  7. A Related Works
    1. A.1 Initialization
    2. A.2 Signal Propagation
    3. A.3 Moment Control & Residual Scaling
    4. A.4 Other Network modifications for Deep Networks
  8. B Moment Propagation through Transformer Components
    1. B.1 Embeddings
    2. B.2 Linear
    3. B.3 Dropout
    4. B.4 ReLU
    5. B.5 GeLU
    6. B.6 LayerNorm
    7. B.7 Softmax
    8. B.8 Scaled Dot-Product Attention
  9. C Moment Propagation through Transformer Blocks
    1. C.1 Transformer Attention Block
    2. C.2 Transformer FFN Block
  10. D Summary Table of Moment Propagation through Transformer Components
  11. E Numerical Verification
  12. F Rank Collapse and Correlation Analysis
  13. G Moment Propagation through the Entire Transformer Model
    1. G.1 Vanilla Pre-LN
      1. G.1.1 Forward Pass
      2. G.1.2 Backward Pass
    2. G.2 Vanilla Post-LN
      1. G.2.1 Forward Pass
      2. G.2.2 Backward Pass
    3. G.3 DeepScaleLM Pre-LN
      1. G.3.1 Forward Pass
      2. G.3.2 Backward Pass
    4. G.4 DeepScaleLM Post-LN
      1. G.4.1 Forward Pass
      2. G.4.2 Backward Pass
    5. G.5 DeepScaleLM (Simplified) Pre-LN
      1. G.5.1 Forward Pass
      2. G.5.2 Backward Pass
    6. G.6 DeepScaleLM (Simplified) Post-LN
      1. G.6.1 Forward Pass
      2. G.6.2 Backward Pass
  14. H Model Moment Figures
    1. H.1 Validity of Theoretical Predictions and Exploding Gradients even after Training
    2. H.2 Importance of Residual Scaling
  15. I Other Modalities: Vision and Speech
    1. I.1 Vision Transformers
    2. I.2 Transformers for Speech Modality
  16. J Statistical Significance
    1. J.1 Error Bars for Pre-Training Experiments
    2. J.2 Statistical Significance for Fine-tuning Experiments
  17. K Compute
    1. K.1 Theoretical compute
    2. K.2 Wall Clock times
  18. L Discussion of Relative Strength
  19. M Discussion of Approximations and Assumptions
    1. M.1 Illustrative Approximations of Full Formulae in Main Paper
    2. M.2 Assumptions and Approximations in Derivations
  20. N DeepScaleLM Pseudocode
  21. O Hyper-parameters

##  Appendix A Related Works

###  A.1 Initialization

Several works, such as Glorot & Bengio (2010); He et al. (2015); Brock et al.
(2021) improved the initialization of ResNets/ReLU networks, but crucially
these works do not consider the impact of correlation in the input, which is
large in Transformer models. Poole et al. (2016) takes correlation into
account, and Schoenholz et al. (2017) initializes weights for networks with
bounded activations so that correlation reaches $1$ asymptotically.

Some works, such as Mishkin & Matas (2016), sequentially profile each layer
empirically by running forward passes through the model, and scaling the
weights and/or output to achieve unit variance, and Liu et al. (2020a, b)
applied the same method for Transformers. Blake et al. (2023) also tries to
achieve unit variance, but does not consider correlation in input or across
tokens, and ignores the non-zero mean of ReLU, resulting in incorrect scale.
Bachlechner et al. (2021) shows unit variance leads to faster convergence at
the start of the training.

We demonstrate that this profiling is unnecessary, and can instead be done
theoretically in DeepScaleLM. Furthermore, where output or gradient increases
in some prior works with more layers (eg. for ADMIN (Liu et al., 2020a), grad
decreases by $\mathcal{O}(N)$ (increases by $\mathcal{O}(log(N))$ for Pre-
LN)), our method allows maintaining both unit output and equal gradient across
all layers at initialization, and bounded during training.

###  A.2 Signal Propagation

Signal propagation in Neural Networks has a long history, such as Neal (1995);
LeCun et al. (1996). More recently, several works have focused on signal
propagation for ResNets, such as He et al. (2015); De & Smith (2020); Brock et
al. (2021); Schoenholz et al. (2017); Hoedt et al. (2022); Labatie et al.
(2021); Marion et al. (2022); Klambauer et al. (2017); Balduzzi et al. (2017).

For transformers, signal propagation was studied in Xu et al. (2019); Dong et
al. (2021); Davis et al. (2021); Noci et al. (2022). Our work also considers
previously neglected effects of dropout, input correlation between tokens,
non-linearity, $QK$ initialization, and provides closed forms with verifiable
correctness of this signal propagation. Ours is the first work to
theoretically constrain the output and gradient to almost exactly unit without
any profiling passes, showing the validity of our formulae and of our
assumptions. See the section Section 4.3 for more discussion on Noci et al.
(2022) specifically.

He et al. (2023) extends neural kernel methods of DKS (Martens et al., 2021)
to Transformers to model network behaviour, assuming the MLP to be linear in
its effect on attention with respect to correlation. $Q/C$ maps in kernel
methods are similar to signal propagation, as expected moments are equivalent
to $q$ and $m$ values of kernels (Martens et al., 2021). Our method relaxes
these assumptions, and we show that considering the impact of ReLU/GeLU on
correlation is critical to correctly modelling attention. In particular, our
formulae show that an MLP block with GeLU will also increase correlation in
the absence of dropout (the same setting as used in He et al. (2023) ). At
large depths, He et al. (2023)’s method suffers from rank collapse (with their
deeper models under-performing shallower ones), which our method successfully
prevents.

We also account for cases with non-IID inputs that may occur due to
segment/position embeddings or due to non-uniform token distributions in real
data (that are distributed approximately per Zipf’s law Kingsley (1935)) – and
find that this strongly affects output variance of the attention block.

###  A.3 Moment Control & Residual Scaling

Bounded gradients, or normalizing per-layer gradients, have been shown to
results in better/faster convergence (Shen et al., 2020; Yu et al., 2017; You
et al., 2017, 2020). Woks such as Takase et al. (2022); Shleifer et al.
(2021); Hayou et al. (2019) also achieved improved training by empirically
mitigating the gradient explosion.

Scaling with $\lambda^{2}+\beta^{2}=1$ to control moments have often been used
for ResNets (Balduzzi et al., 2017; Szegedy et al., 2017; Hanin & Rolnick,
2018; Arpit et al., 2019; Zhang et al., 2019b; Hoedt et al., 2022). Szegedy et
al. (2017) proposed to use any small $\beta$, Balduzzi et al. (2017) proposed
to set $\beta^{2}=0.5$, Bachlechner et al. (2021) sets $\beta=0$ and
learnable. De & Smith (2020) showed that $\lambda^{2}=0.5$ is not sufficient
to solve vanishing gradients.

$\beta^{2}=\frac{k}{N}$ was used to control growth of moments in Arpit et al.
(2019); Brock et al. (2021); Marion et al. (2022); Zhang et al. (2022); He et
al. (2023); Noci et al. (2022)) . $\beta^{2}=\frac{k}{n}$, where $n$ is the
current layer, was used in De & Smith (2020); Liu et al. (2020a, b); Davis et
al. (2021); Blake et al. (2023), but this results in logarithmic bounds
instead of constant for forward propagation if $\lambda=1$ is used, and
vanishing gradient for backward propagation otherwise.

Values of $\beta^{2}<\frac{k}{N}$, such as (effectively) $\frac{1}{N^{2}}$ for
DSInit (Zhang et al., 2019a) or $\frac{1}{N^{1.5}}$ for DeepNet (Wang et al.,
2022a) decrease sensitivity of the model, and may result in the model becoming
“too linear”. DeepNet shows performance improvements by making the model
deeper, but keeping the hidden dimension constant. Our setting is much more
strict - we keep the number of parameters (and hence compute) constant, and
our method still show performance improves on making the model deeper. For
example, DeepNet‘s $200$ layer model is $3.2$B params, whereas our $192$ layer
model is $160$M params ($20$x smaller).

Sometimes, these $\beta$ values are used in conjunction with $\lambda=1$, such
as in Liu et al. (2020a, b), but as shown in He et al. (2023), fully
normalized residual connections with $\lambda^{2}+\beta^{2}=1$ often perform
better than those with $\lambda=1$. We also observed lower performance with
$\lambda=1$ in our initial experiments, and hence we fully normalize the
residual connections.

Our contribution goes beyond providing an optimal scaling scheme. Using the
theoretical framework and closed-form expressions for moment propagation
through both Pre-LN and Post-LN developed in this work, practitioners can make
informed choices about using any of the scaling factors above based on the
stability-performance tradeoffs, such as using a lower $\beta$ for scenarios
with high correlation, or using higher $\beta$ with uncorrelated inputs.

###  A.4 Other Network modifications for Deep Networks

Shi et al. (2022); Zhou et al. (2021); Wang et al. (2022b); Dong et al. (2021)
showed that attention causes rank collapse in deeper models, and Chen et al.
(2020); Zhao et al. (2023) showed the same for graphs. Takase et al. (2022)
added some extra skip connections from the input of the model, Nguyen &
Salazar (2019) modified layernorm slightly, Zhai et al. (2023) normalized all
linear layers by their spectral norm, and Shleifer et al. (2021) added extra
layer norms. Some works in particular, such as Zhai et al. (2023); Zhou et al.
(2021) can only prevent attention entropy collapse later during training, but
our work will also 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. The methods in these works are
orthogonal to our approach, and our equations can be easily extended to cover
the architectural modifications suggested in these.

##  Appendix B Moment Propagation through Transformer Components

We provide detailed proofs of the closed-form expression for each of the
transformer component – Linear layer, Dropout, ReLU, GeLU, LayerNorm, and
Softmax.

For any component, input is represented as $\mathbf{x_{\text{in}}}$ and
$\mathbf{x_{\text{out}}}$ is the output. The gradient flowing in into the
component from the output side is represented as $\mathbf{g_{\text{out}}}$ and
the backpropagated gradient towards the input is $\mathbf{g_{\text{in}}}$. We
switch from vector to matrix notation ($\mathbf{X_{\text{in}}}$,
$\mathbf{X_{\text{out}}}$) whenever needed. We assume that the input is
distributed normally $\mathcal{N}(0,\sigma_{x_{in}})$. No assumptions are made
regarding the covariance of the input – it is not assumed to be IID, and it
may/may-not have covariance both along the sequence length and hidden
dimension. Additional assumptions needed to derive the proofs for softmax and
attention can be found in the respective proofs.

###  B.1 Embeddings

The BERT model’s embedding component 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. Other transformer models, such as decoder-only GPT lack
some (eg. segment) of these, but the derivations remain similar. We derive
formulae for each of these embedding types below.

##### Token Embeddings

We do not assume the input embeddings to be IID. Repetition of same token
introduces correlation across the sequence length. We assume that the input
tokens have been sampled from a multinomial distribution. The words / token
ids are distributed almost according to Zipf’s law (Kingsley, 1935). Assuming
we initialize all the embeddings with variance $\sigma_{w_{embd}}^{2}$, the
relevant statistics for word embeddings output ${x_{\text{out}_{we}}}$ are as
follows

$\displaystyle\mu_{x_{\text{out}_{we}}}$ $\displaystyle=0$
$\displaystyle\sigma_{x_{\text{out}_{we}}}^{2}$
$\displaystyle=\sigma_{w_{\text{embd}}}^{2}$
$\displaystyle\mathrm{Cov}^{l}({x_{\text{out}_{we}}})$
$\displaystyle=\sum{\frac{N_{i}*(N_{i}-1)}{L*(L-1)}}*\sigma_{w_{\text{embd}}}^{2}$
$\displaystyle\mathrm{r}^{l}({x_{\text{out}_{we}}})$
$\displaystyle=\sum{\frac{N_{i}*(N_{i}-1)}{L*(L-1)}}$
$\displaystyle\mathrm{Cov}^{d}({x_{\text{out}_{we}}})$ $\displaystyle=0$

Assume $i$th word occurs $N_{i}$ times, it contributes
$\frac{N_{i}*(N_{i}-1)}{L*(L-1))}$ to the covariance along sequence length.
Similarly, we can calculate the correlation for segment-type embeddings output
${x_{\text{out}_{se}}}$. Zipf’s law states that the probability for each token
is inversely proportional to its rank. For the word with rank $i$,
$p_{i}=\frac{c}{i}$, where
$c=\frac{1}{\sum_{i}\frac{1}{i}}=\frac{1}{\gamma+\mathrm{log}(|V|)}$, where
$\gamma\approx 0.58$ is the Euler’s constant.

For a sentence of length $L$, the token with probability $p_{i}$ is expected
to occur $p_{i}.L$ times. Hence, for a given vocabulary size $|V|$, we can
calculate the correlation as follows

$\displaystyle r^{l}({x_{\text{out}_{we}}})$
$\displaystyle=\sum{\frac{N_{i}*(N_{i}-1)}{L*(L-1)}}$
$\displaystyle=\sum_{i}^{|V|}{\frac{p_{i}L*(p_{i}L-1)}{L*(L-1)}}$
$\displaystyle=\frac{\sum_{i}p_{i}^{2}*L-1}{L-1}$
$\displaystyle=\frac{\sum_{i}\frac{c^{2}}{i^{2}}*L-1}{L-1}$
$\displaystyle\approx\frac{\frac{L\pi^{2}}{6.(\gamma+\mathrm{log}(|V|))^{2}}-1}{L-1}$
$\displaystyle\approx\frac{\pi^{2}}{6.\mathrm{log}(|V|)^{2}}~{}~{}\text{,assuming
$\gamma\approx 0.58<<\mathrm{log}(|V|)\approx 10.4$, $L>>1$}$

##### Segment Type Embeddings

Similarly, the segment type embeddings have two possible values denoting the
sentence order. If first sentence has length $x$, we can consider this as a
special case of the analysis performed above with two possible tokens, where
$N_{1}=x$ and $N_{2}=L-x$. Assuming $x$ is distributed uniformly between $0$
to $L$, $L-x$ also has the same distribution. Hence,

$\displaystyle r^{l}({x_{\text{out}_{se}}},N_{1},N_{2})$
$\displaystyle=\frac{N_{1}^{2}+N_{2}^{2}-L}{L*(L-1)}$

Taking expectation, we get

$\displaystyle r^{l}({x_{\text{out}_{se}}})$
$\displaystyle=\frac{\frac{2}{3}*L^{2}-L}{L*(L-1)}$
$\displaystyle\approx\frac{2}{3}$

##### Position Embeddings

Since learnt position embeddings are lookup tables with unique inputs, the
correlation from position embeddings is $0$.

##### Final Model Input Embeddings

Each of the above embeddings are added before being passed to the transformer
model. Since the variance is same for all embedding types, the final
correlation is the average of the three. Hence:

$\displaystyle r^{l}({x_{\text{out}}})$
$\displaystyle=\frac{1}{3}(r^{l}({x_{\text{out}_{we}}})+r^{l}({x_{\text{out}_{se}}}))$
$\displaystyle=\frac{\pi^{2}}{18*\mathrm{log}(|V|)^{2}}+\frac{2}{9}$

For our case, $|V|=32000$ and sequence length $L=256$, the theoretically
predicted correlation $r^{l}_{x_{in}}=0.227$ which is within $3\%$ of the
empirically observed correlation ($0.221$).

Hence, the final moments for the embedding output are

$\displaystyle\mu_{x_{\text{out}}}$ $\displaystyle=0$
$\displaystyle\sigma_{x_{\text{out}}}^{2}$
$\displaystyle=3*\sigma_{w_{\text{embd}}}^{2}$
$\displaystyle\mathrm{Cov}^{l}_{x_{\text{out}}}$
$\displaystyle=(\frac{\pi^{2}}{18*\mathrm{log}(|V|)^{2}}+\frac{2}{9})\sigma_{x_{\text{out}}}^{2}$
$\displaystyle\mathrm{Cov}^{d}_{x_{\text{out}}}$ $\displaystyle=0$

###  B.2 Linear

For linear layer with $d_{in}$ dimensional input $\mathbf{x_{\text{in}}}$, and
$d_{out}$ dimensional output $\mathbf{x_{\text{out}}}$, we can define the
forward pass mathematically as,

$\displaystyle\mathbf{x_{\text{out}}}$
$\displaystyle=\mathbf{x_{\text{in}}}\mathbf{W}$
$\displaystyle\implies{x_{\text{out}_{j}}}$
$\displaystyle=\displaystyle\sum_{i=1}^{d_{\text{in}}}{x_{\text{in}_{i}}}W_{i,j}$

Similarly, we define the backward pass as,

$\displaystyle\mathbf{g_{{\text{in}}}}$
$\displaystyle=\mathbf{g_{{\text{out}}}}\mathbf{W^{T}}$ $\displaystyle\implies
g_{{\text{in}_{j}}}$
$\displaystyle=\displaystyle\sum_{i=1}^{d_{\text{out}}}g_{{{\text{out}_{i}}}}W_{j,i}$

For expectation of output we have,

$\displaystyle\mathbb{E}[{x_{\text{out}_{j}}}]$
$\displaystyle=\mathbb{E}[\displaystyle\sum_{i=1}^{d_{\text{in}}}{x_{\text{in}_{i}}}W_{i,j}]=\displaystyle\sum_{i=1}^{d_{\text{in}}}\mathbb{E}[{x_{\text{in}_{i}}}W_{i,j}]$
$\displaystyle=\displaystyle\sum_{i=1}^{d_{\text{in}}}\mathbb{E}[{x_{\text{in}_{i}}}]\mathbb{E}[W_{i,j}]=\mu_{x_{\text{in}}}\mu_{w}$
$\displaystyle\mu_{x_{\text{out}}}$ $\displaystyle=0$ ( $\forall j$ )

To get variance of the output of forward pass we have,

$\displaystyle\mathrm{Var}({x_{\text{out}_{j}}})$
$\displaystyle=\mathrm{Var}(\displaystyle\sum_{i=1}^{d_{\text{in}}}{x_{\text{in}_{i}}}W_{i,j})$
As the weights are initialized independently each term in summation is
independent of each other
$\displaystyle=\displaystyle\sum_{i=1}^{d_{\text{in}}}(\mathrm{Var}({x_{\text{in}_{i}}}W_{i,j}))$
$\displaystyle=\displaystyle\sum_{i=1}^{d_{\text{in}}}((\sigma^{2}_{x_{\text{in}}}+\mu^{2}_{x_{\text{in}}})(\sigma_{w}^{2}+\mu_{w}^{2})-\mu^{2}_{x_{\text{in}}}\mu_{w}^{2})$
$\displaystyle=\displaystyle\sum_{i=1}^{d_{\text{in}}}(\sigma^{2}_{x_{\text{in}}}+\mu^{2}_{x_{\text{in}}})\sigma_{w}^{2}$
$\displaystyle\mathrm{Var}({x_{\text{out}_{j}}})$
$\displaystyle=d_{\text{in}}(\sigma^{2}_{x_{\text{in}}}+\mu^{2}_{x_{\text{in}}})\sigma_{w}^{2}$
( $\forall j$ ) $\displaystyle\sigma_{x_{\text{out}}}^{2}$
$\displaystyle=d_{\text{in}}(\sigma^{2}_{x_{\text{in}}}+\mu^{2}_{x_{\text{in}}})\sigma_{w}^{2}$

If we have two inputs $\mathbf{x_{\text{in}}}$ and $\mathbf{y_{\text{in}}}$
such that for all $i$ we have
$\mathrm{Corr}(x_{\text{in}_{i}},y_{\text{in}_{i}})=r^{l}_{x_{\text{in}}}$,
and $\mathbf{{x_{\text{out}}}}=\mathbf{x_{\text{in}}}\mathbf{W}$ and
$\mathbf{{y_{\text{out}}}}=\mathbf{y_{\text{in}}}\mathbf{W}$. Then for any $j$
we have

$\displaystyle\mathrm{Corr}({x_{\text{out}_{j}}},{y_{\text{out}_{j}}})$
$\displaystyle=\frac{\mathbb{E}[{x_{\text{out}_{j}}}{y_{\text{out}_{j}}}]-\mathbb{E}[{x_{\text{out}_{j}}}]\mathbb{E}[{y_{\text{out}_{j}}}]}{\sqrt{\mathrm{Var}({x_{\text{out}_{j}}})\mathrm{Var}({y_{\text{out}_{j}}})}}$
$\displaystyle=\frac{\mathbb{E}[{x_{\text{out}_{j}}}{y_{\text{out}_{j}}}]}{\sqrt{\sigma_{x_{\text{out}}}^{2}\sigma_{x_{\text{out}}}^{2}}}$
$\displaystyle=\frac{\mathbb{E}[\sum_{i=1}^{d_{\text{in}}}{x_{\text{in}_{i}}}W_{i,j}\sum_{k=1}^{d_{in}}{y_{\text{in}_{k}}}W_{k,j}]}{\sigma_{x_{\text{out}}}^{2}}$
$\displaystyle=\frac{\mathbb{E}[\sum_{i=1}^{d_{\text{in}}}{x_{\text{in}_{i}}}{y_{\text{in}_{i}}}W_{i,j}^{2}+\sum_{k=1,k\neq
i}^{d_{in}}\sum_{i=1}^{d_{\text{in}}}{x_{\text{in}_{i}}}{y_{\text{in}_{k}}}W_{i,j}W_{k,j}]}{\sigma_{x_{\text{out}}}^{2}}$
In second summation all terms are independent of each other and as the
expectation of weights is 0 we have
$\displaystyle\mathrm{Corr}({x_{\text{out}_{j}}},{y_{\text{out}_{j}}})$
$\displaystyle=\frac{\mathbb{E}[\sum_{i=1}^{d_{\text{in}}}{x_{\text{in}_{i}}}{y_{\text{in}_{i}}}W_{i,j}^{2}]}{\sigma_{x_{\text{out}}}^{2}}$
$\displaystyle=\frac{\sum_{i=1}^{d_{\text{in}}}\mathbb{E}[{x_{\text{in}_{i}}}{y_{\text{in}_{i}}}W_{i,j}^{2}]}{\sigma_{x_{\text{out}}}^{2}}$
(Independence of weight initialization)
$\displaystyle=\frac{\sum_{i=1}^{d_{\text{in}}}\mathbb{E}[{x_{\text{in}_{i}}}{y_{\text{in}_{i}}}]\mathbb{E}[W_{i,j}^{2}]}{\sigma_{x_{\text{out}}}^{2}}$
$\displaystyle=\frac{\sum_{i=1}^{d_{\text{in}}}(r^{l}_{x_{\text{in}}}\sigma_{x_{\text{in}}}^{2}+\mu_{x_{\text{in}}}^{2})\sigma_{w}^{2}}{\sigma_{x_{\text{out}}}^{2}}\text{~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}(Definition
of correlation)}$
$\displaystyle=\frac{d_{\text{in}}(r^{l}_{x_{\text{in}}}\sigma_{x_{\text{in}}}^{2}+\mu_{x_{\text{in}}}^{2})\sigma_{w}^{2}}{d_{\text{in}}(\sigma^{2}_{x_{\text{in}}}+\mu^{2}_{x_{\text{in}}})\sigma_{w}^{2}}$
$\displaystyle\mathrm{Corr}({x_{\text{out}_{j}}},{y_{\text{out}_{j}}})$
$\displaystyle=\frac{r^{l}_{x_{\text{in}}}\sigma_{x_{\text{in}}}^{2}+\mu_{x_{\text{in}}}^{2}}{\sigma^{2}_{x_{\text{in}}}+\mu^{2}_{x_{\text{in}}}}$
$\displaystyle r^{l}_{x_{\text{out}}}$
$\displaystyle=\frac{r^{l}_{x_{\text{in}}}\sigma_{x_{\text{in}}}^{2}+\mu_{x_{\text{in}}}^{2}}{\sigma^{2}_{x_{\text{in}}}+\mu^{2}_{x_{\text{in}}}}$

As the backward pass has similar structure, assuming
$\mu_{g_{{\text{out}}}}=0$ we can use the same analysis to get,

$\displaystyle\mu_{g_{{\text{in}}}}$ $\displaystyle=0$
$\displaystyle\sigma_{g_{{\text{in}}}}^{2}$
$\displaystyle={d_{\text{out}}}\sigma_{g_{{\text{out}}}}^{2}\sigma_{w}^{2}$

###  B.3 Dropout

We can define Dropout mathematically as,

$\displaystyle\mathbf{x_{\text{out}}}$
$\displaystyle=\mathrm{Dropout}(\mathbf{x_{\text{in}}})$
$\displaystyle\implies{x_{\text{out}_{i}}}$
$\displaystyle=\begin{cases}\frac{x_{\text{in}_{i}}}{(1-p)}&\text{with
probability }1-p\\\ 0&\text{else}\end{cases}$

To calculate expectation of dropout,

$\displaystyle\mathbb{E}[{x_{\text{out}_{i}}}]$
$\displaystyle=0*p+(1-p)*\mathbb{E}[\frac{{x_{\text{in}_{i}}}}{(1-p)}]$
$\displaystyle\mu_{x_{\text{out}}}$ $\displaystyle=\mu_{x_{\text{in}}}$

For variance,

$\displaystyle\mathrm{Var}({x_{\text{out}_{i}}})$
$\displaystyle=\mathbb{E}[{x^{2}_{\text{out}_{i}}}]-\mathbb{E}[{x_{\text{out}_{i}}}]^{2}$
$\displaystyle=0*p+(1-p)*\mathbb{E}[\frac{x^{2}_{\text{in}_{i}}}{(1-p)^{2}}]-\mu_{x_{\text{in}}}^{2}$
$\displaystyle=\frac{\mathbb{E}[{x^{2}_{\text{in}_{i}}}]}{(1-p)}-\mu_{x}^{2}$
$\displaystyle=\frac{\sigma_{x_{\text{in}}}^{2}+\mu_{x_{\text{in}}}^{2}}{(1-p)}-\mu_{x_{\text{in}}}^{2}$
$\displaystyle\sigma_{x_{\text{out}}}^{2}$
$\displaystyle=\frac{\sigma_{x_{\text{in}}}^{2}+p\mu_{x_{\text{in}}}^{2}}{(1-p)}$

If we have two inputs $\mathbf{x_{\text{in}}}$ and $\mathbf{y_{\text{in}}}$
such that for all $i$ we have
$\mathrm{Corr}(x_{\text{in}_{i}},y_{\text{in}_{i}})=r^{l}_{x_{\text{in}}}$,
and $\mathbf{{x_{\text{out}}}}=\mathrm{Dropout}(\mathbf{x_{\text{in}}})$ and
$\mathbf{{y_{\text{out}}}}=\mathrm{Dropout}(\mathbf{y_{\text{in}}})$. Then for
any $j$ we have

$\displaystyle\mathrm{Corr}({x_{\text{out}_{j}}},{y_{\text{out}_{j}}})$
$\displaystyle=\frac{\mathbb{E}[{x_{\text{out}_{j}}}{y_{\text{out}_{j}}}]-\mathbb{E}[{x_{\text{out}_{j}}}]\mathbb{E}[{y_{\text{out}_{j}}}]}{\sqrt{\mathrm{Var}({x_{\text{out}_{j}}})\mathrm{Var}({y_{\text{out}_{j}}})}}$
$\displaystyle=\frac{\mathbb{E}[{x_{\text{out}_{j}}}{y_{\text{out}_{j}}}]-\mu_{x_{\text{out}}}\mu_{x_{\text{out}}}}{\sqrt{\sigma_{x_{\text{out}}}^{2}\sigma_{x_{\text{out}}}^{2}}}$
$\displaystyle=\frac{p^{2}*0+2*p*(1-p)*0+(1-p)^{2}*\mathbb{E}[\frac{{x_{\text{in}_{j}}}{y_{\text{in}_{j}}}}{(1-p)*(1-p)}]-\mu_{x_{\text{out}}}^{2}}{\sigma_{x_{\text{out}}}^{2}}$
$\displaystyle=\frac{\mathbb{E}[{x_{\text{in}_{j}}}{y_{\text{in}_{j}}}]-\mu_{x_{\text{out}}}^{2}}{\sigma_{x_{\text{out}}}^{2}}$
$\displaystyle\mathrm{Corr}({x_{\text{out}_{j}}},{y_{\text{out}_{j}}})$
$\displaystyle=\frac{(r^{l}_{x_{\text{in}}}\sigma_{x_{\text{in}}}^{2})(1-p)}{\sigma_{x_{\text{in}}}^{2}+p\mu_{x_{\text{in}}}^{2}}=r^{l}_{x_{\text{out}}}$

We can define the backward pass of Dropout as,

$\displaystyle g_{{{\text{in}}}_{i}}$
$\displaystyle=\begin{cases}\frac{g_{{{\text{out}}}_{i}}}{(1-p)}&\text{if
$x_{i}$ isn't dropped out (which has probability $(1-p)$)}\\\
0&\text{else}\end{cases}$

Again we can see that backward has similar definition to that of forward pass.
Assuming $\mu_{g_{x_{\text{out}}}}=0$ and using similar analysis we get,

$\displaystyle\mu_{g_{{\text{in}}}}$ $\displaystyle=0$
$\displaystyle\sigma_{g_{{\text{in}}}}^{2}$
$\displaystyle=\frac{\sigma_{g_{{\text{out}}}}^{2}}{(1-p)}$

###  B.4 ReLU

Formuale functionally equivalent to ours for $\mu_{x}$, $\sigma^{2}_{x}$, and
$\sigma^{2}_{g}$ have also been derived in Arpit et al. (2016).

We can define ReLU mathematically as,

$\displaystyle\mathbf{x_{\text{out}}}$
$\displaystyle=\mathrm{ReLU}(\mathbf{x_{\text{in}}})$
$\displaystyle\implies{x_{\text{out}_{i}}}$
$\displaystyle=\begin{cases}{x_{\text{in}_{i}}}&\text{if
}{x_{\text{in}_{i}}}>0\\\ 0&\text{else}\end{cases}$

For getting expectation of output of ReLU for normally distributed input we
have,

$\displaystyle\mathbb{E}[{x_{\text{out}_{i}}}]$
$\displaystyle=\int_{-\infty}^{\infty}\frac{\mathrm{ReLU}({x_{\text{in}_{i}}})\exp{(\frac{-{x^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}d{x_{\text{in}_{i}}}$
$\displaystyle=\int_{-\infty}^{0}\frac{0*\exp{(\frac{-{x^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}d{x_{\text{in}_{i}}}+\int_{0}^{\infty}\frac{{x_{\text{in}_{i}}}\exp{(\frac{-{x^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}d{x_{\text{in}_{i}}}$
$\displaystyle=\int_{0}^{\infty}\frac{{x_{\text{in}_{i}}}\exp{(\frac{-{x^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}d{x_{\text{in}_{i}}}$
Substituting
$t=\displaystyle\frac{{x^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}}$
we have
$dt=\displaystyle\frac{{x_{\text{in}_{i}}}d{x_{\text{in}_{i}}}}{\sigma_{x_{\text{in}}}^{2}}$
we get, $\displaystyle\mathbb{E}[{x_{\text{out}_{i}}}]$
$\displaystyle=\int_{0}^{\infty}\frac{\sigma_{x_{\text{in}}}\exp{(-t)}dt}{\sqrt{2\pi}}$
$\displaystyle=\frac{\sigma_{x_{\text{in}}}}{\sqrt{2\pi}}[-\exp{(-t)}]_{0}^{\infty}=\frac{\sigma_{x_{\text{in}}}}{\sqrt{2\pi}}$

Hence, the mean of output

$\displaystyle\mu_{x_{\text{out}}}=\frac{\sigma_{x_{\text{in}}}}{\sqrt{2\pi}}$
(1)

Variance of output can be calculated by,

$\displaystyle\mathrm{Var}({x_{\text{out}_{i}}})$
$\displaystyle=\mathbb{E}[{x_{\text{out}_{i}}}^{2}]-\mathbb{E}[{x_{\text{out}_{i}}}]^{2}$
$\displaystyle=\int_{-\infty}^{\infty}\frac{(\mathrm{ReLU}({x_{\text{in}_{i}}}))^{2}\exp{(\frac{-{x^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}d{x_{\text{in}_{i}}}-\frac{\sigma_{x_{\text{in}}}^{2}}{2\pi}$
$\displaystyle=\int_{-\infty}^{0}\frac{0*\exp{(\frac{-{x^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}d{x_{\text{in}_{i}}}+\int_{0}^{\infty}\frac{{x^{2}_{\text{in}_{i}}}\exp{(\frac{-{x^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}d{x_{\text{in}_{i}}}-\frac{\sigma_{x_{\text{in}}}^{2}}{2\pi}$
$\displaystyle=\int_{0}^{\infty}\frac{{x^{2}_{\text{in}_{i}}}\exp{(\frac{-{x^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}d{x_{\text{in}_{i}}}-\frac{\sigma_{x_{\text{in}}}^{2}}{2\pi}$
Let
$I=\displaystyle\int_{0}^{\infty}\frac{{x^{2}_{\text{in}_{i}}}\exp{(\frac{-{x^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}d{x_{\text{in}_{i}}}$,
then substituting $t=-{x_{\text{in}_{i}}}$ we have, $\displaystyle I$
$\displaystyle=\int_{0}^{-\infty}\frac{-t^{2}\exp{(\frac{-t^{2}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}dt$
$\displaystyle=\int_{-\infty}^{0}\frac{t^{2}\exp{(\frac{-t^{2}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}dt$
$\displaystyle\implies I+I$
$\displaystyle=\int_{-\infty}^{0}\frac{t^{2}\exp{(\frac{-t^{2}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}dt+\int_{0}^{\infty}\frac{{x^{2}_{\text{in}_{i}}}\exp{(\frac{-{x^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}d{x_{\text{in}_{i}}}$
$\displaystyle 2I$
$\displaystyle=\int_{-\infty}^{\infty}\frac{{x^{2}_{\text{in}_{i}}}\exp{(\frac{-{x^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}d{x_{\text{in}_{i}}}=\sigma_{x_{\text{in}}}^{2}$
$\displaystyle\implies\mathrm{Var}({x_{\text{out}_{i}}})$
$\displaystyle=\frac{\sigma_{x_{\text{in}}}^{2}}{2}-\frac{\sigma_{x_{\text{in}}}^{2}}{2\pi}=\frac{\sigma_{x_{\text{in}}}^{2}}{2}(1-\frac{1}{\pi})$
$\displaystyle\sigma_{x_{\text{out}}}^{2}$
$\displaystyle=\frac{\sigma_{x_{\text{in}}}^{2}}{2}(1-\frac{1}{\pi})$

Now for two inputs $\mathbf{{x_{\text{in}}}}$ and $\mathbf{{y_{\text{in}}}}$
such that for all $i$ we have
$\mathrm{Corr({x_{\text{in}_{i}}},{y_{\text{in}_{i}}})}=r^{l}_{x_{\text{in}}}$,
and $\mathbf{{x_{\text{out}}}}=\mathrm{ReLU}(\mathbf{{x_{\text{in}}}})$ and
$\mathbf{{y_{\text{out}}}}=\mathrm{ReLU}(\mathbf{{y_{\text{in}}}})$. Then for
any $j$ we have,

$\displaystyle\mathrm{Corr}({x_{\text{out}_{j}}},{y_{\text{out}_{j}}})=\frac{\mathbb{E}[{x_{\text{out}_{j}}}{y_{\text{out}_{j}}}]-\mathbb{E}[{x_{\text{out}_{j}}}]\mathbb{E}[{y_{\text{out}_{j}}}]}{\sqrt{\mathrm{Var}({x_{\text{out}_{j}}})\mathrm{Var}({y_{\text{out}_{j}}})}}$
$\displaystyle\mathbb{E}[{x_{\text{out}_{j}}}{y_{\text{out}_{j}}}]=\displaystyle\int_{0}^{\infty}\int_{0}^{\infty}\frac{{x_{\text{in}_{j}}}{y_{\text{in}_{j}}}}{2\pi\sigma_{x_{\text{in}}}^{2}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-({x^{2}_{\text{in}_{j}}}+{y^{2}_{\text{in}_{j}}}-2r^{l}_{x_{\text{in}}}{x_{\text{in}_{j}}}{y_{\text{in}_{j}}})}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}d{x_{\text{in}_{j}}}d{y_{\text{in}_{j}}}$
$\displaystyle=\displaystyle\int_{0}^{\infty}\int_{0}^{\infty}\frac{{x_{\text{in}_{j}}}{y_{\text{in}_{j}}}}{2\pi\sigma_{x_{\text{in}}}^{2}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-({x_{\text{in}_{j}}}-r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}})^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}d{x_{\text{in}_{j}}}d{y_{\text{in}_{j}}}$

Substituting $t={x_{\text{in}_{j}}}-r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}$,
and assuming $y_{\text{in}_{j}}$ is constant for the inner
integral,$dx_{\text{in}_{j}}=dt$

$\displaystyle\mathbb{E}[{x_{\text{out}_{j}}}{y_{\text{out}_{j}}}]=$
$\displaystyle=\int_{0}^{\infty}\frac{{y_{\text{in}_{j}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\displaystyle\int_{-r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}^{\infty}\frac{t+r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-t^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dtd{y_{\text{in}_{j}}}$
$\displaystyle=\int_{0}^{\infty}\frac{{y_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\int_{-r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}^{\infty}\frac{t}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-t^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dtd{y_{\text{in}_{j}}}$
$\displaystyle+\int_{0}^{\infty}\frac{{y_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\int_{-r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}^{\infty}\frac{r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-t^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dtd{y_{\text{in}_{j}}}$

Let us first define $I_{1}$ and $I_{2}$ as:

$\displaystyle I_{1}$
$\displaystyle=\displaystyle\int_{0}^{\infty}\frac{{y_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\int_{-r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}^{\infty}\frac{t}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-t^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dtd{y_{\text{in}_{j}}}$
$\displaystyle I_{2}$
$\displaystyle=\displaystyle\int_{0}^{\infty}\frac{{y_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\int_{-r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}^{\infty}\frac{r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-t^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dtd{y_{\text{in}_{j}}}$
$\displaystyle I_{1}$
$\displaystyle=\displaystyle\int_{0}^{\infty}\frac{{y_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\int_{-r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}^{\infty}\frac{t}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-t^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dtd{y_{\text{in}_{j}}}$

Substituting
$p=\displaystyle\frac{t^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})}$
we have
$dp=\displaystyle\frac{tdt}{\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})}$

$\displaystyle I_{1}$
$\displaystyle=\int_{0}^{\infty}\frac{{y_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\int_{\frac{(r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}})^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})}}^{\infty}\frac{\sigma_{x_{\text{in}}}\sqrt{(1-({r^{l}_{x_{\text{in}}}})^{2})}}{\sqrt{2\pi}}\exp{(-p)}dpd{y_{\text{in}_{j}}}$
$\displaystyle=\int_{0}^{\infty}\frac{{y_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\frac{\sigma_{x_{\text{in}}}\sqrt{(1-({r^{l}_{x_{\text{in}}}})^{2})}}{\sqrt{2\pi}}\exp{(\frac{-(r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}})^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}d{y_{\text{in}_{j}}}$
$\displaystyle=\int_{0}^{\infty}\frac{{y_{\text{in}_{j}}}\sqrt{(1-({r^{l}_{x_{\text{in}}}})^{2})}}{2\pi}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}d{y_{\text{in}_{j}}}$

Substituting
$m=\displaystyle\frac{{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})}$,
$dm=\displaystyle\frac{{y_{\text{in}_{j}}}d{y_{\text{in}_{j}}}}{\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})}$,

$\displaystyle I_{1}$
$\displaystyle=\int_{0}^{\infty}\frac{\sqrt{(1-({r^{l}_{x_{\text{in}}}})^{2})}}{2\pi}(1-({r^{l}_{x_{\text{in}}}})^{2})\sigma_{x_{\text{in}}}^{2}\exp{(-m)}dm$
$\displaystyle=\frac{(1-({r^{l}_{x_{\text{in}}}})^{2})^{\frac{3}{2}}\sigma_{x_{\text{in}}}^{2}}{2\pi}$
$\displaystyle I_{2}$
$\displaystyle=\int_{0}^{\infty}\frac{{y_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\int_{-r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}^{\infty}\frac{r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-t^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dtd{y_{\text{in}_{j}}}$
$\displaystyle=\int_{0}^{\infty}\frac{r^{l}_{x_{\text{in}}}{y^{2}_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\int_{-r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}^{\infty}\frac{1}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-t^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dtd{y_{\text{in}_{j}}}$

Substituting $p=-t$, where $\Phi$ is CDF of Standard Normal Distribution

$\displaystyle I_{2}$
$\displaystyle=\int_{0}^{\infty}\frac{r^{l}_{x_{\text{in}}}{y^{2}_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\int_{r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}^{-\infty}\frac{-1}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-p^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dpd{y_{\text{in}_{j}}}$
$\displaystyle=\int_{0}^{\infty}\frac{r^{l}_{x_{\text{in}}}{y^{2}_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\int^{r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}_{-\infty}\frac{1}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-p^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dpd{y_{\text{in}_{j}}}$
$\displaystyle=\int_{0}^{\infty}\frac{r^{l}_{x_{\text{in}}}{y^{2}_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\Phi(\frac{r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}{\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}})d{y_{\text{in}_{j}}}$
$\displaystyle=\int_{0}^{\infty}\frac{r^{l}_{x_{\text{in}}}{y^{2}_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}[\frac{1}{2}(1+\mathrm{erf}(\frac{r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}{\sigma_{x_{\text{in}}}\sqrt{2(1-({r^{l}_{x_{\text{in}}}})^{2})}}))]d{y_{\text{in}_{j}}}$
$\displaystyle=\frac{r^{l}_{x_{\text{in}}}}{2}\int_{0}^{\infty}\frac{{y^{2}_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}d{y_{\text{in}_{j}}}+$
$\displaystyle\frac{r^{l}_{x_{\text{in}}}}{2\sqrt{2\pi}\sigma_{x_{\text{in}}}}\int_{0}^{\infty}{y^{2}_{\text{in}_{j}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\mathrm{erf}(\frac{r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}{\sigma_{x_{\text{in}}}\sqrt{2(1-({r^{l}_{x_{\text{in}}}})^{2})}})d{y_{\text{in}_{j}}}$

Let us define $I_{2,1}$ and $I_{2,2}$ as

$\displaystyle I_{2,1}$
$\displaystyle=\displaystyle\frac{r^{l}_{x_{\text{in}}}}{2}\int_{0}^{\infty}\frac{{y^{2}_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}d{y_{\text{in}_{j}}}$
$\displaystyle I_{2,2}$
$\displaystyle=\displaystyle\frac{r^{l}_{x_{\text{in}}}}{2\sqrt{2\pi}\sigma_{x_{\text{in}}}}\int_{0}^{\infty}{y^{2}_{\text{in}_{j}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}\mathrm{erf}(\frac{r^{l}_{x_{\text{in}}}{y_{\text{in}_{j}}}}{\sigma_{x_{\text{in}}}\sqrt{2(1-({r^{l}_{x_{\text{in}}}})^{2})}})d{y_{\text{in}_{j}}}$
$\displaystyle I_{2,1}$
$\displaystyle=\displaystyle\frac{r^{l}_{x_{\text{in}}}}{2}\int_{0}^{\infty}\frac{{y^{2}_{\text{in}_{j}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{j}}}}{2\sigma_{x_{\text{in}}}^{2}})}d{y_{\text{in}_{j}}}$
$\displaystyle I_{2,1}$
$\displaystyle=\frac{r^{l}_{x_{\text{in}}}\sigma_{x_{\text{in}}}^{2}}{4}$
(Same integral as in variance calculation)

From Ng & Geller (1969) we have
$\displaystyle\int_{0}^{\infty}x^{2}\exp{(-b^{2}x^{2})\mathrm{erf}(ax)dx}=\frac{\sqrt{\pi}}{4b^{3}}-\frac{\tan^{-1}(\frac{b}{a})}{2\sqrt{\pi}b^{3}}+\frac{a}{2\sqrt{\pi}b^{2}(a^{2}+b^{2})}$.

Hence, putting
$a=\displaystyle\frac{r^{l}_{x_{\text{in}}}}{\sigma_{x_{\text{in}}}\sqrt{2(1-({r^{l}_{x_{\text{in}}}})^{2})}}$
and $b=\displaystyle\frac{1}{\sigma_{x_{\text{in}}}\sqrt{2}}$ we get,

$\displaystyle
I_{2,2}=\frac{r^{l}_{x_{\text{in}}}}{2\sqrt{2\pi}\sigma_{x_{\text{in}}}}[\frac{2\sqrt{2}\sigma_{x_{\text{in}}}^{3}}{4}-\frac{\tan^{-1}(\frac{\sqrt{(1-({r^{l}_{x_{\text{in}}}})^{2})}}{r^{l}_{x_{\text{in}}}})2\sqrt{2}\sigma_{x_{\text{in}}}^{3}}{2\sqrt{\pi}}+\frac{\sqrt{2}r^{l}_{x_{\text{in}}}\sigma_{x_{\text{in}}}^{3}\sqrt{(1-({r^{l}_{x_{\text{in}}}})^{2})}}{\sqrt{\pi}}]$
$\displaystyle=\frac{r^{l}_{x_{\text{in}}}\sigma_{x_{\text{in}}}^{2}}{4}-\frac{r^{l}_{x_{\text{in}}}\cos^{-1}{(r^{l}_{x_{\text{in}}})}\sigma_{x_{\text{in}}}^{2}}{2\pi}+\frac{({r^{l}_{x_{\text{in}}}})^{2}\sqrt{(1-({r^{l}_{x_{\text{in}}}})^{2})}\sigma_{x_{\text{in}}}^{2}}{2\pi}$
$\displaystyle\mathbb{E}[{x_{\text{out}_{j}}}{y_{\text{out}_{j}}}]=I_{1}+I_{2,1}+I_{2,2}$
$\displaystyle=\frac{(1-({r^{l}_{x_{\text{in}}}})^{2})^{\frac{3}{2}}\sigma_{x_{\text{in}}}^{2}}{2\pi}+2*\frac{r^{l}_{x_{\text{in}}}\sigma_{x_{\text{in}}}^{2}}{4}-\frac{r^{l}_{x_{\text{in}}}\cos^{-1}{(r^{l}_{x_{\text{in}}})}\sigma_{x_{\text{in}}}^{2}}{2\pi}+\frac{{(r^{l}_{x_{\text{in}}}})^{2}\sqrt{(1-({r^{l}_{x_{\text{in}}}})^{2})}\sigma_{x_{\text{in}}}^{2}}{2\pi}$
$\displaystyle=\frac{r^{l}_{x_{\text{in}}}\sigma_{x_{\text{in}}}^{2}}{2}-\frac{r^{l}_{x_{\text{in}}}\cos^{-1}{(r^{l}_{x_{\text{in}}})}\sigma_{x}^{2}}{2\pi}+\frac{\sqrt{(1-({r^{l}_{x_{\text{in}}})}^{2})}\sigma_{x_{\text{in}}}^{2}}{2\pi}$
$\displaystyle\mathrm{Corr}({x_{\text{out}_{j}}},{y_{\text{out}_{j}}})=\frac{\mathbb{E}[{x_{\text{out}_{j}}}{y_{\text{out}_{j}}}]-\mathbb{E}[{x_{\text{out}_{j}}}]\mathbb{E}[{y_{\text{out}_{j}}}]}{\sqrt{\mathrm{Var}({x_{\text{out}_{j}}})\mathrm{Var}({y_{\text{out}_{j}}})}}$
$\displaystyle=\frac{\displaystyle\frac{r^{l}_{x_{\text{in}}}\sigma_{x_{\text{in}}}^{2}}{2}-\frac{r^{l}_{x_{\text{in}}}\cos^{-1}{(r^{l}_{x_{\text{in}}})}\sigma_{x_{\text{in}}}^{2}}{2\pi}+\frac{\sqrt{(1-({r^{l}_{x_{\text{in}}}})^{2})}\sigma_{x}^{2}}{2\pi}-\frac{\sigma_{x_{\text{in}}}^{2}}{2\pi}}{\displaystyle\frac{\sigma_{x_{\text{in}}}^{2}}{2}(1-\frac{1}{\pi})}$
$\displaystyle r^{l}_{x_{\text{out}}}$ $\displaystyle=\frac{\frac{\pi
r^{l}_{x_{\text{in}}}}{2}+r^{l}_{x_{\text{in}}}\sin^{-1}{(r^{l}_{x_{\text{in}}})}+\sqrt{(1-({r^{l}_{x_{\text{in}}}})^{2})}-1}{\pi-1}$

Backward pass on ReLU can be defined as,

$\displaystyle g_{{\text{in}_{i}}}$
$\displaystyle=\begin{cases}g_{{\text{out}_{i}}}&\text{if
${x_{\text{in}_{i}}}>0$ (which has probability $\frac{1}{2}$)}\\\
0&\text{else}\end{cases}$

Assuming $\mu_{g_{{\text{out}}}}=0$,

$\displaystyle\mathbb{E}[g_{{{\text{in}_{i}}}}]$
$\displaystyle=\frac{1}{2}*0+\frac{1}{2}*\mathbb{E}[g_{{{\text{out}_{i}}}}]$
$\displaystyle\mu_{g_{{\text{in}}}}$ $\displaystyle=0$
$\displaystyle\mathrm{Var}(g_{{{\text{in}_{i}}}})$
$\displaystyle=\mathbb{E}[g_{{{\text{in}_{i}}}}^{2}]-\mathbb{E}[g_{{\text{in}_{i}}}]^{2}=\mathbb{E}[g_{{\text{in}_{i}}}^{2}]$
$\displaystyle=\frac{1}{2}*0+\frac{1}{2}*\mathbb{E}[g_{{\text{out}}}^{2}]$
$\displaystyle\sigma_{g_{{\text{in}}}}^{2}$
$\displaystyle=\frac{\sigma_{g_{{\text{out}}}}^{2}}{2}$

If for two inputs $\mathbf{x}_{\text{in}}$ and $\mathbf{y}_{\text{in}}$ for
all $i$ we have
$\mathrm{Corr}(g_{\text{out}_{x_{i}}},g_{\text{out}_{y_{i}}})=r^{l}_{g_{\text{out}}}$,
and $g_{\text{in}_{x_{i}}},g_{\text{in}_{y_{i}}}$ be the gradient after
passing through ReLU layer. Then we have,

$\displaystyle\mathbb{E}[g_{\text{in}_{x_{i}}}g_{\text{in}_{y_{i}}}]$
$\displaystyle=\mathbb{P}(x_{\text{in}_{i}}>0,y_{\text{in}_{i}}>0)\mathbb{E}[g_{\text{out}_{x_{i}}}g_{\text{out}_{y_{i}}}]$
$\displaystyle=\mathbb{P}(x_{\text{in}_{i}}>0,y_{\text{in}_{i}}>0)r^{l}_{g_{\text{out}}}\sigma_{g_{{\text{out}}}}^{2}$
$\displaystyle\mathbb{P}(x_{\text{in}_{i}}>0,y_{\text{in}_{i}}>0)=$
$\displaystyle=\int_{0}^{\infty}\int_{0}^{\infty}\frac{{x_{\text{in}_{i}}}{y_{\text{in}_{i}}}}{2\pi\sigma_{x_{\text{in}}}^{2}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-({x^{2}_{\text{in}_{i}}}+{y^{2}_{\text{in}_{i}}}-2r^{l}_{x_{\text{in}}}{x_{\text{in}_{i}}}{y_{\text{in}_{i}}})}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}d{x_{\text{in}_{i}}}d{y_{\text{in}_{i}}}$
$\displaystyle=\int_{0}^{\infty}\int_{0}^{\infty}\frac{{x_{\text{in}_{i}}}{y_{\text{in}_{i}}}}{2\pi\sigma_{x_{\text{in}}}^{2}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-({x_{\text{in}_{i}}}-r^{l}_{x_{\text{in}}}{y_{\text{in}_{i}}})^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}\exp{(\frac{-{y^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}d{x_{\text{in}_{i}}}d{y_{\text{in}_{i}}}$

Substituting $t={x_{\text{in}_{i}}}-r^{l}_{x_{\text{in}}}{y_{\text{in}_{i}}}$,
and assuming $y_{\text{in}_{i}}$ is constant for the inner
integral,$dx_{\text{in}_{i}}=dt$

$\displaystyle\mathbb{P}(x_{\text{in}_{i}}>0,y_{\text{in}_{i}}>0)=$
$\displaystyle\int_{0}^{\infty}\frac{1}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}\displaystyle\int_{-r^{l}_{x_{\text{in}}}{y_{\text{in}_{i}}}}^{\infty}\frac{1}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-t^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dtd{y_{\text{in}_{i}}}$

Substituting $p=-t$, where $\Phi$ is CDF of Standard Normal Distribution

$\displaystyle\mathbb{P}(x_{\text{in}_{i}}>0,y_{\text{in}_{i}}>0)=$
$\displaystyle=\int_{0}^{\infty}\frac{1}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}\int_{r^{l}_{x_{\text{in}}}{y_{\text{in}_{i}}}}^{-\infty}\frac{-1}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-p^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dpd{y_{\text{in}_{i}}}$
$\displaystyle=\int_{0}^{\infty}\frac{1}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}\int^{r^{l}_{x_{\text{in}}}{y_{\text{in}_{i}}}}_{-\infty}\frac{1}{\sqrt{2\pi}\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}}\exp{(\frac{-p^{2}}{2\sigma_{x_{\text{in}}}^{2}(1-({r^{l}_{x_{\text{in}}}})^{2})})}dpd{y_{\text{in}_{i}}}$
$\displaystyle=\int_{0}^{\infty}\frac{1}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}\Phi(\frac{r^{l}_{x_{\text{in}}}{y_{\text{in}_{i}}}}{\sigma_{x_{\text{in}}}\sqrt{1-({r^{l}_{x_{\text{in}}}})^{2}}})d{y_{\text{in}_{i}}}$
$\displaystyle=\int_{0}^{\infty}\frac{1}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}[\frac{1}{2}(1+\mathrm{erf}(\frac{r^{l}_{x_{\text{in}}}{y_{\text{in}_{i}}}}{\sigma_{x_{\text{in}}}\sqrt{2(1-({r^{l}_{x_{\text{in}}}})^{2})}}))]d{y_{\text{in}_{i}}}$
$\displaystyle=\frac{1}{2}\int_{0}^{\infty}\frac{1}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-{y^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}d{y_{\text{in}_{i}}}+\frac{1}{2\sqrt{2\pi}\sigma_{x_{\text{in}}}}\int_{0}^{\infty}\exp{(\frac{-{y^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}\mathrm{erf}(\frac{r^{l}_{x_{\text{in}}}{y_{\text{in}_{i}}}}{\sigma_{x_{\text{in}}}\sqrt{2(1-({r^{l}_{x_{\text{in}}}})^{2})}})d{y_{\text{in}_{i}}}$
$\displaystyle=\frac{1}{4}+\frac{1}{2\sqrt{2\pi}\sigma_{x_{\text{in}}}}\int_{0}^{\infty}\exp{(\frac{-{y^{2}_{\text{in}_{i}}}}{2\sigma_{x_{\text{in}}}^{2}})}\mathrm{erf}(\frac{r^{l}_{x_{\text{in}}}{y_{\text{in}_{i}}}}{\sigma_{x_{\text{in}}}\sqrt{2(1-({r^{l}_{x_{\text{in}}}})^{2})}})d{y_{\text{in}_{i}}}$

From Ng & Geller (1969) we have
$\displaystyle\int_{0}^{\infty}\exp{(-b^{2}x^{2})\mathrm{erf}(ax)dx}=\frac{\sqrt{\pi}}{2b}-\frac{1}{b\sqrt{\pi}}\tan^{-1}(\frac{b}{a})$

Putting
$a=\displaystyle\frac{r^{l}_{x_{\text{in}}}}{\sigma_{x_{\text{in}}}\sqrt{2(1-({r^{l}_{x_{\text{in}}}})^{2})}}$
and $b=\displaystyle\frac{1}{\sigma_{x_{\text{in}}}\sqrt{2}}$ we get,

$\displaystyle\mathbb{P}(x_{\text{in}_{i}}>0,y_{\text{in}_{i}}>0)$
$\displaystyle=\frac{1}{4}+\frac{1}{2\sqrt{2\pi}\sigma_{x_{\text{in}}}}[\frac{\sqrt{\pi}\sigma_{x_{\text{in}}}\sqrt{2}}{2}-\frac{\sigma_{x_{\text{in}}}\sqrt{2}}{\sqrt{\pi}}\tan^{-1}(\frac{\sqrt{(1-({r^{l}_{x_{\text{in}}}})^{2})}}{r^{l}_{x_{\text{in}}}})]$
$\displaystyle=\frac{1}{4}+\frac{1}{2\pi}[\frac{\pi}{2}-\cos^{-1}{(r^{l}_{x_{\text{in}}})}]$
$\displaystyle=\frac{1}{4}+\frac{\sin^{-1}{(r^{l}_{x_{\text{in}}})}}{2\pi}$
$\displaystyle\implies\mathbb{E}[g_{\text{in}_{x_{i}}}g_{\text{in}_{y_{i}}}]$
$\displaystyle=(\frac{1}{4}+\frac{\sin^{-1}{(r^{l}_{x_{\text{in}}})}}{2\pi})r^{l}_{g_{\text{out}}}\sigma_{g_{{\text{out}}}}^{2}$
$\displaystyle\mathrm{Corr}(g_{\text{in}_{x_{i}}},g_{\text{in}_{y_{i}}})$
$\displaystyle=\frac{(\frac{1}{4}+\frac{\sin^{-1}{(r^{l}_{x_{\text{in}}})}}{2\pi})r^{l}_{g_{\text{out}}}\sigma_{g_{{\text{out}}}}^{2}}{\frac{\sigma_{g_{{\text{out}}}}^{2}}{2}}$
$\displaystyle r^{l}_{g_{\text{out}}}$
$\displaystyle=(\frac{1}{2}+\frac{\sin^{-1}{(r^{l}_{x_{\text{in}}})}}{\pi})r^{l}_{g_{\text{out}}}$

###  B.5 GeLU

Forward pass through GeLU is defined as,

$\displaystyle\mathbf{x_{\text{out}}}$
$\displaystyle=\mathrm{GeLU}(\mathbf{x_{\text{in}}})$
$\displaystyle\implies{x_{\text{out}_{i}}}$
$\displaystyle={x_{\text{in}_{i}}}\Phi({x_{\text{in}_{i}}})$ where $\Phi(x)$
is CDF of Standard Normal Distribution at $x$
$\displaystyle=\frac{{x_{\text{in}_{i}}}}{2}\left(1+\mathrm{erf}(\frac{{x_{\text{in}_{i}}}}{\sqrt{2}})\right)$

To get the mean of output of GeLU, we have

$\displaystyle\mathbb{E}[{x_{\text{out}_{i}}}]$
$\displaystyle=\int_{-\infty}^{\infty}\frac{{x_{\text{out}_{i}}}}{\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-x_{\text{in}_{i}}^{2}}{2\sigma_{x_{\text{in}}}^{2}})}dx_{\text{in}_{i}}$
$\displaystyle=\int_{-\infty}^{\infty}\frac{{x_{\text{in}_{i}}}(1+\mathrm{erf}(\frac{{x_{\text{in}_{i}}}}{\sqrt{2}}))}{2\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-x_{\text{in}_{i}}^{2}}{2\sigma_{x_{\text{in}}}^{2}})}dx_{\text{in}_{i}}$
$\displaystyle=\int_{-\infty}^{\infty}\frac{{x_{\text{in}_{i}}}}{2\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-x_{\text{in}_{i}}^{2}}{2\sigma_{x_{\text{in}}}^{2}})}dx_{\text{in}_{i}}+\int_{-\infty}^{\infty}\frac{{x_{\text{in}_{i}}}\mathrm{erf}(\frac{{x_{\text{in}_{i}}}}{\sqrt{2}})}{2\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-x_{\text{in}_{i}}^{2}}{2\sigma_{x_{\text{in}}}^{2}})}dx_{\text{in}_{i}}$
$\displaystyle=\int_{-\infty}^{\infty}\frac{{x_{\text{in}_{i}}}\mathrm{erf}(\frac{{x_{\text{in}_{i}}}}{\sqrt{2}})}{2\sqrt{2\pi}\sigma_{x_{\text{in}}}}\exp{(\frac{-x_{\text{in}_{i}}^{2}}{2\sigma_{x_{\text{in}}}^{2}})}dx_{\text{in}_{i}}$
(Integral of odd function)