$\sigma$Reparam: Stable Transformer Training with Spectral Reparametrization
Authors: Shuangfei Zhai
Tatiana Likhomanenko
Etai Littwin
Jason Ramapuram
Dan Busbridge
Yizhe Zhang
Jiatao Gu
Joshua M. Susskind
Venue: Submitted to ICLR 2023 (2023)
Keywords: Transformers
self-attention
optimization
stability
spectral normalization
self-supervised learning
vision
speech
language
contrastive learning

Abstract: Training stability is of great importance to Transformers. In this work, we investigate the training dynamics of Transformers by examining the evolution of the attention layers. In particular, we track the "attention entropy" for each attention head during the course of training, which is a proxy of the attention's sharpness. We observe a common, non monotonic evolution of attention entropy across different settings: the attention entropy first quickly decreases in the initial phase of training, followed by quickly increasing, and finally entering a long stable phase. While the exact shape can be affected by hyperparameters such as warmup, initialization, learning rate etc., we found that there is a close correlation between the minima of attention entropy and the model's training stability. To this end, we propose a simple and efficient solution dubbed $\sigma$Reparam, where we reparametrize all linear layers with Spectral Normalization and an additional learned scalar. We provide a lower bound on the attention entropy as a function of the spectral norms of the query and key projections, which suggests that small attention entropy can be obtained with large spectral norms. $\sigma$Reparam decouples the growth rate of a weight matrix's spectral norm from its dimensionality, which we verify empirically. We conduct experiments with $\sigma$Reparam on image classification, image self supervised learning, automatic speech recognition and language modeling tasks. We show that $\sigma$Reparam provides great stability and robustness with respect to the choice of hyperparameters.

Under review as a conference paper at ICLR 2023
O REPARAM: STABLE TRANSFORMER TRAINING WITH SPECTRAL REPARAMETRIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
Training stability is of great importance to Transformers, In this work, we inves- tigate the training dynamics of Transformers by examining the evolution of the attention layers. In particular, we track the "attention entropy" for each attention head during the course of training, which is proxy of the attention's 's sharpness. We observe common, non monotonic evolution of attention entropy across different settings: the attention entropy first quickly decreases in the initial phase training, followed by quickly increasing, and finally entering long stable phase. While the exact shape can be affected by hyperparameters such as warmup, initialization, learning rate etc., found that there is a close correlation between the minima of attention entropy and the model's training stability. To this end, we propose a simple and efficient solution dubbed a Reparam, where reparametrize all linear layers with Spectral Normalization and an additional learned scalar. We provide lower bound on the attention entropy as a function of the spectral norms of the query and key projections, which suggests that small attention entropy can be ob- tained with large spectral norms, oReparam decouples the growth rate of weight matrix's 's spectral norm from its dimensionality, which we verify empirically. We conduct experiments with o Reparam on image classification, image self supervised learning, automatic speech recognition and language modeling tasks. We show that oReparam provides great stability and robustness with respect to the choice hyperparameters.
INTRODUCTION
Transformers (Vaswani et al. 2017) are state-of-the-ar models in many application domains. However, despite their empirical success and wide adoption, great care often needs to be taken in order to achieve good training stability and convergence. In the original paper (Vaswani et al. 2017), residual connections and Layer Normalizations (LNs) (Ba et al. 2016) are extensively used for each Attention and MLP block (specifically, in the "Post Norm" fashion). There has since been various works attempting to promote better training stability and robustness. For example, the "Pre Norm" (Radford let al. 2019) scheme has gained wide popularity, where one moves the placement of LNs to the beginning of each residual block. Others have argued that it is important to properly condition the residual connections. [Bachlechner et al. (2021) proposes to initialize the residual connections to zero to promoter better signal propagation. Zhang et al. (2018): Huang et al. (2020) remove LNs with carefully designed initialization schemes.
In this work, we study the training instability of Transformers from the lens of training dynamics. We start by monitoring the average entropy of the attention heads (by treating each attention head as a multinomial distribution) over all query positions and examples. Interestingly, the average attention entropy often evolves a pattern consisting of three phases. In the beginning, attention entropy starts high (corresponding to uniform attention scores) and quickly drops to small value; This is then followed by a second stage where it quickly increases to a relatively high entropy regime; Lastly the attention entropy curve stabilizes and smoothly evolves to convergence. See the top left plot of Figure [ for an illustration, which is a Vision Transformer (Touvron et al. 2021) (VTT) trained on ImageNet classification, using well optimized hyper parameters.
Empirically, we have found that the attention entropy is directly correlated with the model's 's stability and convergence. In particular, small attention entropy reached in the initial phase often causes slow

---

Under review as a conference paper at ICLR 2023
Baseline
0.2x Learning Rate
7.1
7.0
6.5
6.0
6.0
Loss
1.5
Trirtng
Art cxmpps layer
Entropy
5.5
Teaking
entogy Ayer
Entropy
Art erappy Syard Alta catrep layer
Arm eneogy lyerr 6
Alto ealoggy lays
4.5
4.5
3.5
3.
10
20
401
101
Training Epochs
Training Epochs
2x Learning Rate Warmup
2x Learning Rate
7.0
7.0
0.1
6.0
Losss
Aus cxtrpp laoc
Entropy
Loss
5.5
caloayy
Entropy
Als catopp laye
Ala estoge laya
Amr layer II
Am entogy lyer
4.5
4.5
1.5
3.5
10
20
20
40
30
Training Epochs
Training Epochs
Figure 1: The training loss curves of ViT-B on ImageNet, together with the attention entropy for three layers. From top left to bottom right: baseline with default hyper parameters from Touvron et al.) (202T); 0.2x learning rate; 2 warmup epochs; 2x learning rate. We see close correlation between the dip of the attention entropy and the convergence and stability of the training loss.
convergence, fluctuations in training loss and, in the worst case, divergence. This is shown in Figure where we vary the learning rate and warmup epochs of the baseline ViT model. We see that both decreased the learning rate and increased warmup epochs provide smoothing effects to the attention entropy curves, which in turn yield lower training losses. On the other hand, increasing learning rate brings a detrimental impact training where the attention entropy collapses to zero and training diverges. We denote the rapid dip of attention entropy to a near zero value and its resulting pathological optimization dynamics as "entropy collapse".
The remaining questions are: 1) How do we get rid of entropy collapse? 2) Can we improve training stability by doing so? We answer them by showing that attention entropy is closely related to the spectral norms of the query and key projections. In particular, we show a lower bound of the attention entropy, which suggests that large spectral norms of the projections collapse. We then provide a simple fix, dubbed oReparam which reparameterizes all weight matrices by sequentially applying Spectral Normalization (Miyato et al. (2018) and learned multiplicative scalar. Intuitively, Reparam decouples the update of the spectral norms of weights from their dimensionality, which allows them update smoothly controlled way. Also note that oReparam does not change the model space, which allows one to learn an arbitrarily expressive model.
more easily lead entropy
We validate oReparam on 4 tasks: image classification, image self supervised learning, automatic speech recognition and language modelling. We show that Reparam effectively slows down the growth of each layer's spectral norms, and as a result, their attention entropy eurves are greatly smoothed. This allows us to achieve great robustness with respect to the choice of hyper parameters. In certain cases, we are able to remove Layer Norms and still achieve competitive results.
2 RELATED WORKS
Transformers have relied heavily on LNs to achieve training stability. Besides the popular Post Norm and Pre Norm configurations, other variants have been proposed (Wang al. 2022 Shleifer et al.
2

---

Under review as a conference paper at ICLR 2023
2021). RReparam does not rely on LN and can even work in the absence of it, which avoids the computational over head of explicit activation normalization.
There have also been numerous attempts to design better Transformer initialization schemes, includ- ing Zhang et al. (2018); Huang et al. (2020); Yang et al. (2022); [Bachlechner et al. (202T). oReparam addresses the training dynamics of attention layers, which makes it
is an orthogonal approach compatible with standard initialization methods and provides robust performance.
Reparam is special case of weight reparameterization, which has found wide adoption in Deep Learning. Weight Norm (WN) (Salimans & Kingma 2016) is a well known example of such methods, but its effectiveness in Transformers is limited. In ConvNets, simple additive weight reparameter- ization (Ding et al.) 2021) has been demonstrated useful in speeding up training convergence. To the best of our knowledge, &Reparam is the first simple reparamterization technique that provides competitive performance with well optimized baseline models.
3 METHOD
3.1 ATTENTION ENTROPY
At the core of Transformers is dot product attention. Let X E Rxx denote input sequence an attention layer (we assume self attention for simplicity of presentation), where T, d are the number of tokens and the token dimension, respectively; and let Wx. We Rdxms Wv E Rdx., denote the key, query and value matrices. A simple attention layer then computes Att(X) = AXWv where A (a))
XWR x W X and & the row-wise softmax function. We define the attention entropy of row i of A by Ent(A,) - Ar log(A..)) We also overload the notation and let Em(AA + Σ Ent(A.) denote the average attention entropy of A. As shown in Figure 1 the attention entropy (and the entropy collapse phenomenon) a strong indicator of training stability of Transformers. Our goal is to alleviate the entropy collapse problem and achieve smooth evolution of the attention entropy through training.
We next investigate the properties of attention entropy. We show in the the next theorem that Ent(A) is directly connected to the Spectral norm (the largest singular value) of W k W
Theorem 3.1 (Attention entropy lower bound). Assume without loss of generality || S 1, and let spectral norm [| 112. Then it holds that:
Ent( (A,) 2 log (1+ (T (T- 1)e oV +
o T -Ie
(1)
1 (T 1)e
Moreover, there exists imputs and weights WK, Wq for which the lower bound in Eq. (1) is tight.
Therefore, for large T, the minimum attainable entropy behaves like (((Tr We note that the bound on the entropy in Theorem i. tight in sense that it achievable for some inputs X. Moreover, the typical Frobenious (L2) regularization would not ensure a small (a small Frobenious norm much more restrictive than a small Spectral norm), hence it would not be as effective in preventing an "entropy collapse". Proofs for Theorem 3.1 3.1land and the following Proposition are provided in Appendix A
3.2 a REPARAM
We then present &Reparam, method to re-parameterize the weights of linear layer with:
W
o(W) W.
(2)
where (W) E RR is the spectral norm of W and 7 E IR is learnable parameter, initialized to 1. In practice, or (W) can be computed via power iteration (Mises & Dellazzek-Geiringe 1929) as in Spectral Normalization (SN) (Miyato et al. 2018), see Algorithm ] for sketch implementation. Note that oReparam brings little extra overhead as the power iteration mainly consists of two matrix vector products and is only performed on the parameters rather than activations. During inference,
can compute W once and freeze it, which means that has the same cost as regular linear layer.
3

---

Under review as a conference paper at ICLR 2023
Table 1: Supervised Image Classification on ImageNetIk. The B/L refer to ViT-B/16 and ViT-L/16 variants respectively. SN corresponds to the Spectral Norm baseline without the learnable scalar. Also note that the WN configuration leads immediate divergence without using Layer Norm, and here we only report the result with WN LN.
DeiT (B)
&Reparam (B)
SN (B)
WN (B)
MAE (B/L)
oReparam (B/L.)
Top-1
81.8%
82.2%
69.81%
78.25%
82.1%
81.88% / 82.41 % 81.5%
EMA Top-1
68.41% 1%
76.95%
82.37% / 82.48% 82.3% 82.6%
Training Epochs
300
300
250
250
300
250 / 300 No Yes (LARS)
Layer Norm
Yes
No
No
Yes
Yes
SGD
No
Yes (LARS)
No
No
Cosine Schedule
Yes
Yes
No
No
Yes
No / Yes
LR Warmup
Yes
Yes
No
No
Yes
No
Weight Decay
Yes
Yes
No
No
Yes
No
Why Reparam?? Unlike the standard SN, &Reparam introduces an additional multiplier "y which explicitly controls the SN of the weights, and there is explicit pressure regularize the SN. The additional multiplier is necessary to avoid restricting the capacity of the network, and we find that training and overall performance significantly degraded in its absence. Since the representational capacity of the layer remains unchanged, it is not immediately clear why &Reparam would effectively regularize the SN of the weights. While full theoretical characterization beyond the scope of this paper, we identify a property of adaptive optimizers which, if left unchecked, causes the spectral norm of weight matrices to grow rapidly for large weight matrices. To illustrate this, we adopt common assumptions in stochastic optimization, and model the stochastic gradients at some point in the optimization by g = μ + e IR R=x where z is the mean and e is random variable with E(e] E(c² n E "ox= A typical Adam optimizer update attempts approximate the
following ideal update: A E[gl The following proposition lower bounds the spectral norm of VEL=a
the ideal update oA): Proposition 3.2. It holds that:
n
a(A)
W
+ n
(3)
i
Note that the noise second moment n² typically in the order of 222, hence Eq. (3) (3) indicates that the spectral norm of the ideal update should be large, growing linearly with Vw. Moreover, for large batch sizes we would have n²² < 1, resulting in (A) /w While such large spectral norm could be offset by proper learning rate adjustment, this would be counter productive since 1) a small learning rate typically induces inferior performance, and 2) architectures with layers of varying sizes, such attention layers, would require a per layer learning rate tuning. In contrast, oReparam avoids this issue since the spectral norm of each layer is controlled by a single parameter y, hence the size of its update does not scale with w and is uniform across layers.
EXPERIMENTS
4.1 SUPERVISED IMAGE CLASSIFICATION
Improved robustness. We first start from a well tuned recipe with ViT-B on ImageNet-Ik (Touvron)
et al. 2021), and vary its hyper parameters in the grid [base Jr E (5e - 4, le - 3}, butch_size E {1024, 2048], }. wwrmup.
[0,5}] 7/8 configurations lead to divergence except for the default [5e- 4, 2048, 5) hyper parameter. We next apply RReaara to all the linear layers (including the initial patch embedding). and removed all the LayerNorm instances. All configurations in the same grid search converge with an average top-1 accuracy of 81.4% (max 82.2%, shown in Table I This suggests improved robustness with respect to hyperparameters.
This would he exact for full batch optimization.
4

---

Under review as a conference paper at ICLR 2023
Simplified recipe. RRepaam also enables simplified framework for training ViT-B and ViT-L models, in contrast to sfate-of-the art ImageNet- Ik ViT training protocols such as the fully supervised MAE recipe (He et al. 2022) and DeiT (Touvron et al. 2021 (Table DD. In the of ViT-B models, we are able to train for shorter duration, remove all LayerNorm layers, remove LR warmup. remove cosine scheduling (requiring only simple step schedule at 210 epochs) and Furthermore, &Reparam enables SGD training via LARS (You et al., (017) (with momentum 0.9) something possible with traditional ViT training protocols (Touvron al. 2021 He et 2022). These simplifications also have the added benefit of reducing GPU memory overhead For the VIT-L model we relax the LR schedule back to cosine and match the baseline model's training interval. Both models use FP32 precision on the attention operands and keep mixed precision training for the rest the network. The full set of hyperparameters available in Appendix
weight decay.
80
Test Top (%)
00
40
ollauara
Buscine
20
(2 entropy (nats)
10
CNyy C nooor
(0)
100 150
200
250
300
Training Epochs
Figure 2: Test performance, attention entropy, and largest singular value of attention weights of supervised oReparam ViT-B/16 alongside supervised MAE VIT-B/16 line) trials of each method are presented. The MAE ViT-B/16 presents a more constrained attention entropy in contrast to the DeiT formulation from Figure T due the longer warmup, lower learning rate and stronger weight decay.
baselines. Best (solid line) and worst (dashed
To further understand the effect of oReparam, we track both the attention entropy, and the largest singular value of the attention weight matrix over the course of training. In Figure 2 oReparam maintains a lower largest attention weight singular value and presents a higher, but monotonically decreasing attention entropy throughout training. As previously discussed, smaller bounded singular value helps with stable training, whereas a higher attention entropy encourages exploration more diverse solutions. This is reinforced by the accelerated performance observed in Test Top and the 50 epoch reduction in training time for the oReparam ViT-B/16 shown Figure
4.2 SELF-SUPERVISED TRAINING OF VISUAL REPRESENTATIONS
In computer vision, self-supervised learning (SSL) has been effective in enabling efficient training downstream tasks (Assran et al. 2022). Most of this progress has been made using convolutional architectures, while works using VITss often require specialized training recipes (Caron et al. 2021).
Recently, it was found that ViTs suffer from training instabilities in SSL tasks Chen et al. 2021
These instabilities can be remedied through a combination of frozen patch embedders, initialization
We observe 8.2% memory reduction in full FP32 (for 1:1 comparison) with batch size of 86 per GPU.
5

---

Under review as a conference paper at ICLR 2023
Table 2: (top) Best SimCLR ImageNetlk trial top linear probing performance training for 300 epochs. RReparam + LN yields the highest performing run, with Frozen Patcher performing competitively. (bottom) Configuration of the variants used in our stability analysis. The MoCo v3 weight initialization and patch initialization scheme descrihed in Chen et al. (202T). For full hyperparameters see Table 6/ Appendix C.
Baseline
Frozen Patcher
Reparam
oReparam + LN
Top @ (ours)
72.4
74.4
73.7
74.5
Weight Init
MoCo
MoCo
trunc.norm 02)
trunc_norm (. 02)
Patcher Init
MoCo
MoCo trunc_nerm 02)
trunc_norm (. (.02)
Frozen Patcher
No
Yes
No
No
<Reparam
No
No
Yes
Yes
Layer Norm
Yes
Yes
No
Yes
Test Top (S)
Packe
00 entrone (eats)
Test Top
grad -000 2010
Eackiz
From Rd
clopun- +1x
cilipuee
Training Epochs
Training Epochs
(a) Statistics of best and worst trials per method.
(b) Stability over 10 trials per method.
Figure Ten trials of SimCLR each method on ImageNetlk with 40 epochs learning rate warmup. (a) Linear probe performance for the best (solid line) and worst (dashed line) trials of each method, against relevant metries from the first attention layer (top to hottom); attention entropy, the spectral norm of the attention weights, and the -radieet norm of the attention weights. We see that the Frozen Patcher method functions intended, regulating its gradient norm, protecting it from the large gradient norms inducing instability Baxeline. We also observe a second form of instability during training: the growing spectral norm leads to poorly behaved attention mechanism, entropy collapse, and a drop in performance as described in SScciion This affects Baseline, well as Frozen Patcher, neither method gives specific protection against this second type of instability (solid and dashed red, and dashed green lines), Finally, we see that &Reparam with and without layer normalization regulate both gradient norms, as well as the spectral norms, giving defense against both types of instability. (b) Linear probe performance of every trial. We see that RReparam is the most stable method. oReparam LN is also quite stable. In the case where it experiences instabilities, able
that it is much quicker than Baseline and Frozen Patcher. This due the regularization of the spectral norm which 1) prevents any arising instability pushing the model too far away from the current solution, and 2) keeps the attention mechanism useful, such that gradients are available for any required correction.
schemes, and longer learning rate warmups; however, there is an open question whether a general solution providing stable SSL ViT training exists Chen et al. (221).
Here, we demonstrate that Reearam is a ViT SSL stabilizer. Taking SimCLR as our SSL method,
investigate four variants. Baseline and Frozen Patcher were studied in Chen et al. (2021), whereas Reparam and Reearamm LN are our solution. These methods are detailed in Table and their full hyperparameters given in Table of Appendix C.1

---

Under review as a conference paper at ICLR 2023
Figure 4: Linear probe performance on Ima- geNetlk at the end of training over 10 trials for each method. Trials are ordered by decreasing per- formance, with run rank 1 (10) correspending to the best (worst) trial. Frozen Patcher and Reeaarm IN produce the best individual runs, with Repparam marginally lower. a Reparam + LN and oReparam are the methods most reliably giving good perfor- mance, with Baxeline and Frozen Patcher each sus- coptible to at least one instability type.
Test Top 1 (%)
55
50
Poszline
45
Frum Paacher
dReparum LN
40
dReparum
Run rank
We observe two types of instability. The first, as observed in Chen et al. (2021). is induced by large gradient norms in early layers. The second, described in |Section 3] relates to entropy collapse. We find that Frozen Patcher protects against the first type. but is still susceptible to the second. &Reparam, however, can protect against both types of instability, yielding more reliable training (see|Figure> 3).
As noted in Chen et al. (2021), instabilities reduce final performance. We show instability impact on performance in Figure 4 The methods with the best performing individual runs are Frozen Patcher and or Reparam L.V, whereas the most stable methods are Reparam LN and o Reparam.
Our main stability experiments use 40 epochs of learning rate warmup. matching the setting of Chen et al. 20221). Using oReparam, as in the supervised setting, gives training stability even at the lower learning rate warmup of 10 epochs. For more details, see (eppendix C.2
Finally, we look at the performance attainable when training for longer duration of 300 epochs in T The best performing method run is given by with Reparam + I.N, with Frozen Patcher performing almost well, and both outperforming the reference SimCLR result (Chen et al. 2021)
Ultimately, we see while Reparam produces the lowest degree of instability, the best overall method for stable training of SimCLR ViTs is oReparam + LN, producing both the highest ImageNet linear probe performance at 100 epochs (69.6 %) and 300 epochs (74.5 %) epochs, as well as very stable training many trails, both at long and short learning rate warmup.
4.3 SPEECH
In this section we focus experiments for automatic speech recognition (ASR).
Data All experiments are performed on the subset of 100h audio paired with transcriptions (train- clean-100) of LibriSpeech dataset Panayotov et al. (015). The standard LibriSpeech validation sets (dev-clean and dev-other) are used to tune all hyper parameters, as well as to select the best models. Test sets (test-clean and test-other) are used only to report final word error rate (WER) performance without an external language model. We keep the original 16kHz sampling rate and compute log-mel filterbanks with 80 coefficients for 25ms sliding window, strided by 10ms, later normalized to zero mean and unit variance per input sequence.
Acoustic Models We are using current, to the best of our knowledge, state-of-ther the-art model on 100h of LibriSpeech (Likhomanenko et al.) (2021a). The model consists of ID convolution to perform striding, Transformer encoder with post-L.N and final linear layer to map to the output number of
tokeng The model is trained with Connectionist Temporal Classification (Graves et al. 2006) loss.
To speed up the model training (-3x) and decrease memory usage we are using CAPE positional embeddings (Likhomanenko et al. 2021c instead of relative embeddings Shaw et al. (2018).
Data augmentation We use SpecAugment (Park et al. 2019 activated right at the beginning of training. We use two frequency masks with frequency mask parameter F 30, ten time masks with maximum time-mask ratio 0.1 and time mask parameter = 50; time warping is not used.
Training We use Adagrad (Duchi et al. 2011) if not specified otherwise, and LR decaying by 2 each time the WER reaches plateau on the validation. We use dynamic batching of 240s audio per GPU and train with tensor cores fp32 on 8 Ampere A100 (40GB) GPUs for 350-500k updates. No
The taken set consists of the 26 English alphahet letters augmented with the apostrophe and a word houndary taken.
7

---

Under review as a conference paper at ICLR 2023
Table 3: Comparison between different normalizations and our re-parametrizatioo for speech domain: training loss and word error rate are reported for the best models.
post-LN
pre-LN (same)
pre-LN (optimized)
SN
SN
WN
WN
oReparam
ReReprrmm +post-LN
+post-LN
+post-LN
Train Loss
37.7
35.3
37.2
160.4
120.3
35.6
35.4
37.5
34.9
dev-clean WER
5.9
42.6
20.3
7.0
6.3
6.4
dev-other WER
17.7
21.3
19.1
62.9
42.7
22.3
19.4
20.5
17.8
lest-clean WER
6.2
7.1
6.3
42.4
20.4
7.3
6.7
6.8
6.4
lest-nther WER
17.8
21.6
19.3
63.6
43.6
22.6
19.5
21.0
18.0
weight decay is used. Default warmup set to 64k for the baselines and varied for different models. The default LR is 0.03 and also optimized across models. We also apply gradient clipping of 1.
4.3.1 TRAINING STABILITY, ROBUSTNESS AND GENERALIZATION
First, we experiment with stability of training for the baselines using both "Pre Norm" (pre-LN) and "Post Norm" (post-LNN architectures. If we vary LR, warmup, and gradient clipping, all post-LN experiments either diverge reduce warmup from 64k to 16k, increase learning rate from 0.03 to 0.5, and obtain better results than before. While LNe is more stable than post-LN, it generalizes worse: validation WER is worse while training loss is lower, see Table When we switch to o Reparam we observe the same stability for pre-LN, Figure 5 while having better generalization than not optimized pre-LN. We are not able to match the post-LN results until we combine post-LN together with oReparam, which allows us achieve similar performance on the dev and test sets and lower training loss. In Figure both oReparm and o &Reparam with post-LN demonstrate robustness with respect to training hyperparameters. We also compare with Spectral Norm (SN) where y set to I and is not learnable and WN baselines. Both SN and WN perform poor compared to Reeaaam,, see Table
training is observed. At the same time, pre-LN is stable: we can
In prior works it was reported that post-LN can be impossible train with very deep architectures, see e.g. Liu et al. (220b)aa a). We reproduced similar results: if we increase the encoder size to 2x then post-LN does not train, while pre-LN works out of the box and improves over the smaller architecture. We applied the same settings to a Reparam and combination of oReparm and post-LN: for both cases out of the box models train well and achieve similar results as pre-Norm. This confirms o Reearamm's ability for stable training even with post-LN.
40000
40000
03
100.0
23h
25%
Figure 5: Robustness of oReparam with respect to different hyperparameters: learning rate (left), warmup (middle) and initialization std value (right).
4.3.2 TRAINING WITH SGD
Prior works report different problems training transformers with SGD (see e.g. (Li et al., 2022) First, we experimented with the baselines, pre-LN and post-LN and observed similar issues. is hard to find hyperparameters that enable the model to train. Following vision experiments we switch to the LARS (You et al.) 2017) (with momentum 0.9) optimizer, and are able to train pre-LN and post-LN by carefully tuning the LR (the rest stays the same, including gradient clipping) which is varied from 0.1 to 1.5, see Table First, we observe that Post-LN is more unstable (many LRs are diverging or not training) and gives significantly worse results than pre-LN. Second, pre-LN is still behind the baseline that uses an adaptive optimizer. However, if we switch to &eearam (key, queries matrix) we observe stable training with respect to LR changes,
and values represented
8

---

Under review as a conference paper at ICLR 2023
Table 4: Comparison between different normalizations and our re-parametrization for speech domain when no warmup and LARS optimizer are used: training loss and word error rate are reported for the best models. oReparam performs re- re-parametrization for joint matrix for key, queries and values in self-attention. DV denotes model divergence: we are not able to train SN with post-LN configuration.
post-LN
pre-LN
SN
SN
WN
WN
oReparam
Reepaam +post-LN
+post-LN
+post-LN
Train Loss
64.5
29.4
160.0
DV
59.1
63.2
51.1
34,2
dev-clean WER
8.1
5.9
49.8
DV
8.3
7.1
7.2
5.8
dev-other WER
25.0
18.9
69.6
DV
25.9
22.0
22.8
18.1
test-clean WER
8.6
6.4
49.4
DV
8.7
7.5
7.5
6.2
test-other WER
25.6
19.2
70.9
DV
26.4
22.1
23.2
18.7
and combined together with post-LN it achieves similar performance to the best results from Table while keeping the train loss low
4.4 LANGUAGE
Setup. We use the WikiText-103 language model (LM) benchmark, which consists of 103M tokens sampled from English Wikipedia (Merity et al. (017)) Our baseline is highly optimized Transformer (Baevski & Auli 2019) with 32 layers. heads, 128 head dimensions, 1024 model dimensions, 4096 fully connected dimensions and post LayerNorm. The word embedding and softmax matrices are tied (Press & Wolf 2017). We partition the training data into oonoolelappinng blocks of 512 contiguous tokens and train the model to autoregressively predict each token (Baevski & Auli 2019). Validation and test perplexity is measured by predicting the last 256 words out of the input of 512 consecutive words to avoid evaluating tokens in the beginning with limited context (early loken curse, Press et al.
2021
Table 5: WikiText- language modeling results in perplexity.
PPLJ
Model
train
dev.
test
Reparam w/ weight decay
16.5
17.9
18.6
o Reparam w/o weight decay
12.9
18.5
19.3
Baseline Transformer Baevski & Auli (2019)
15.4
18.1
18.7
Results. We do not experience training instability with the baseline Transformer, likely because the masked attention in autoregressive models makes entropy collapse less likely we experimented with oReparam to test its generality on a different modality/problem. We apply oReparam all linear layers of the Transformer while removing all LayerNorms, and search for learning rate in a grid [1, 1.5, 2, 2.5] and weight decay in the grid [le-3, 1e-4, O]. All other hyperparameters are kept the same as the baseline. The results are shown in Table [ We see that even in the absence of LayerNorm, o Reparam shows strong performance in convergence and dev/test performance. With a mild weight decay, o RReaaram also outperforms the baseline wrt the dev/test PPL.
Nonetheless,
5 CONCLUSION
We analyze the training stability of Transformers from the lens of the attention entropy. We show that training instability or divergence is often accompanied by the entropy collapse phenomenon, and provide a simple fix named oReparam. We demonstrate over wide set of benchmarks, domains, and training methodologies, that &Reparam provides great stability and robustness, often leading to simplified model design and/or better performance.
For the separate reparametrization for (keys, queries) and values observe less stable training with LARSS and warmup relative reparametrizing them together.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Mahmoud Assran, Mathilde Caron, Ishan Misra, Piotr Bojanowski, Florian Bordes, Pascal Vincent, Armand Joulin, Michael G. Rabbat, and Nicolas Ballas. Masked siamese networks for label- efficient learning. CoRR, abs/2204.07144 141, 2022. doi: 10448550/aXxx.2244001144. URL https
//doi 48550/arrxxxv 2204. 07141
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv: 1607.06600 2016.
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley, Rezero is all you need: Fast convergence at large depth. In Uncertainty in Artificial Intelligence, pp. 1352-1361. PMLR, 2021.
Alexei Baevski and Michael Auli. Adaptive input representationn for neural language modeling. In Proc. of ICLR, 2019. BtRLbttps://
10853.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10-17, 2021, pp. 9630 9640. IEEE, 2021. doi: 10.1 ww/1///CV4992..20211.... 100511. URL https? 7/doi www.oogg/.00..
/ 110. 1tt.l..or//...11//C///////////// (ww.099/CCCV4888 2tw.:22222222000000
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 9640-9649, 2021.
Xiaohan Ding, Xiangyu Zhang, Ningning Ma, Jungong Han, Guiguang Ding, and Jian Sun. Repvgg: Making vgg-style convnets great again. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 13733-13742, 2021.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(111):1211-2159, 2011.
Alex Graves, Santiago Fernández, Faustino Gomez, and Jürgen Schmidhuber. Connectionist temporal classification: labelling unsegmented sequence data with recurrent neural networks. In Proceedings of the 23rd international conference on Machine learning, pp. 369-376, 2006.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 16000-16009, 2022.
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimization through better initialization. In International Conference on Machine Learning, pp. 4475 4488. PMLR, 2020.
Zhiyuan Li, Srinadh Bhojanapalli, Manzil Zaheer, Sashank Reddi, and Sanjiv Kumar. Robust training of neural networks using scale invariant architectures. In International Conference on Machine Learning, pp. 12656-12684. PMLR, 2022.
Tatiana Likhomanenko, Qiantong Xu, Jacob Kahn, Gabriel Synnaeve, and Ronan Collobert. slimipl: Language-modellfree iterative pseudo-labeling, Proc. Interspeech, 2021a.
Tatiana Likhomanenko, Qiantong Xu, Vineel Pratap. Paden Tomasello, Jacob Kahn, Gilad Avidov, Ronan Collobert, and Gabriel Synnaeve. Rethinking evaluation in asr: Are our models robust enough? Proc. Interspeech, 2021b.
Tatiana Likhomanenko, Qiantong Xu, Gabriel Synnaeve, Ronan Collobert, and Alex Rogozhnikov. Cape: Encoding relative positions with continuous augmented positional embeddings. Advances in Neural Information Processing Systems, 34, 2021c.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. In Proceedings of the 2020 Conference on Empirical Methods Natural Language Processing (EMNLP 2020), 2020a.
10

---

Under review as a conference paper at ICLR 2023
Xiaodong Liu, Kevin Duh, Liyuan Liu, and Jianfeng Gao. Very deep transformers for neural machine translation. In arXiv:2008. 07772 [cs]. 2020b.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In Proc. of ICLR, 2017. URLhttts://axxxy wwtp:///tg/.......////// 07843
RV Mises and Hilda Pollacrek- Praktische verfahren der gleichungsauflösung. ZAMM- Journal of Applied Mathematics and Mechamics/ZZitschriff für Angewandte Mathematik und Mechanik, 9(1):88-77 77, 1929.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference Learning Representations, 2018. 8.
Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur. Librispeech: an asr corpus based on public domain audio books. In 2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). pp. 5206-5210, IEEE, 2015.
Daniel S Park, William Chan, Yu Zhang, Chung-Cheng Chiu, Barret Zoph, Ekin D Cubuk, and Quoc V Le. Specaugment: A simple data augmentation method for automatic speech recognition. Proc. Interspeech 2019, pp. 2613-2617, 2019.
Ofir Press and Lior Wolf. Using the output embedding to improve language models. In Proc. of EACL, 2017. Btt.htts::/// wtwpi/xxiiivv//... http:///aaaavvvvv.aaaaa//aa
05859.
Ofir Press, Noah A. Smith, and Mike Lewis. Shortformer: Better language modeling using shorter inputs, 2021.
httpp::/aaax.vvorrg/aa///// 2502//
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural information processing systems, 29, 2016.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume 2 (Short Papers), pp. 464-463 2018.
Sam Shleifer, Jason Weston, and Myle Ott. Normformer: Improved transformer pretraining with extra normalization. arXiv preprint arXiv:21 10.02450 2021.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning, pp. 10347-10357. PMLR, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Hongyu Wang, Shuming Ma. Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. Deepnet: Scaling transformers to 1,000 layers. arXiv preprint arXiv:2203.005: 00555, 2022.
Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint (rXiv:2203.03466 2022.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arxiv: /00 03888, 2017.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. In International Conference on Learning Representattons, 2018.
11

---

Under review as a conference paper at ICLR 2023
A PROOF OF THEOREM 3. AND PROPOSITION 3.2
Theorem 3.1 (Attention entropy lower bound), Assume without loss of generality |||| S 1, and let spectral norm || 2|2. Then it holds that:
T(T 1
Ent( 2 log
(T 1 +
(1)
1 (T 11-
Moreover, there exists inputs X and weights WK, WQ for which the lower bound in Eq. (1) is fight.
Proof. WLOG let RT denote the j'th row of From the condition that ||x||2 S it holds that ||l|| I a. Let p(u) denote the softmax probabilities given by:
c Ps
(4)
Z
where Z = =E e is the partition function. The entropy given p(u) is then:
T
u
Ent(u)
log's
+ log( Z)).
(5)
Z
7
Z
We wish solve the following constrained minimization problem:
min Ent(u) s.t ||u|2 o²
(6)
where > 0. Define the lagrangian:
C(u, A) =nt(u))) A[llu|1-
To find all saddle points, we solve the system of equations:
aL(u, A)
ØL(u, A)
0,
=0
(8)
Ou
ax
Giving rise to the following set of equations:
VICKST Aua Σ 24 (di.
[1 log( e
Z
Z
7
(log((p))) Ent(u))
(10)
|||| 2 o²
(11)
As first step, assume that for the minimizer w* Eq. (6) there exists index & such that ut Using Eq. (7)
log(ps) Ent(u)
Σ
Pi pi log
Σ ps log(e")) Σ p.u
Eu
(12)
Pk
From the first set of equations we arrive the condition:
log(p,) Ent(u)
log(pp) Ent(u)
V #40 Pj
Py
(13)
4y
Eu
Eu
pj
+
Ppr
(14)
Uj
uy
Pj'
(15)
This however implies that ui un Now, assuming VA us A O. we have that:
ui 0. hence a contradiction
Eq. (9)
" u
Va
Vuy
ZEu const
(16)
12

---

Under review as a conference paper at ICLR 2023
The monotonicity of the LHS of Eq. (16) implies that contains only 2 distinct values. WLOG assume ui ax V, W;
D2 -2
Then we have:
V _
e
ae + (1 - 0² - a
zz
T-
(17)
7-1
With solution:
a
7- 7: B= T(T-1)
(18)
With the corresponding entropy:
Ent(u") log (1 (T (T-1)e - + o T 1)e V
(19)
1+ (T -o
Proposition A.1. It holds that:
w
n
a(A) 2 W
(3)
w
P + n
1-=1
Proof. We have that:
a
A
Trace A
ATA)
1 w
Σ
Vi
12
(20)
n
w2
P
+
n
1,j=1
1.j=1
B IMPLEMENTATION OF dREPARAM
To compute spectral norm of the current matrix we use the power method as approximation method to speed up computations. See Algorithm| for sketch implementation.
Algorithm 1 Pseudo code of oReearam in PyTorch-like style.
parameters. Ms weight matrix, shape (d, gammes the learned apectral shape (1,) butters, ut shape 1o,), vi snape (o,), the lett and right singular vectors ot if init: # initialize u, v as rardom unit. vectors and garma to
zar.d (d)
i1 norm (dim=0)
v Fardn (c)
genma v.norm (dim-0)
if traiming: $ if in the treining mode, perform one step of power iteration first w.mv (v)
u.norm (dim=0)
V W.T.mv (t)
v v. no.nom (dim=0))))
sigma einsum (' de, >', hat
gamma sigme the effective spectral norn of What would be ganma
13

---

Under review as a conference paper at ICLR 2023
Table 6: Default hyperparameters of the variants of SimCLR used our stability analysis. The MoCo v3 weight initialization and patch initialization scheme are described in Chen et al (202)). SinCos refers to stacked 2D SinCos positional encodings Vaswani al. (2017). The table is divided vertically into hyperparameters that differ across methods (top) and hyperparameters shared across methods (bottom).
Baseline
Frozen Patcher
RReparam
Reepaam + LN
Reparam
No
No Yes
Yes No No trunc norm 02)
Yes
Frozen Patcher
No
No
Layer Norm
Yes
Yes
Yes
Patcher Init
MoCo v3
MoCo v3
trunc (.)
Weight Init
MoCo v3
MoCo
trune_norm C2)
trune_norm (. 02)
Architecture
ViT-B/16
ViT-B/16
ViT-B/16
VIT-B/16
Batch Size
4096
4096
4096
4096
ColorJitter Strength
0.5
0.5
0.5
0.5
Learning Rate
2 10
2 10
2 10
2 x 10
Learning Rate Sched
Cosine
Cosine
Cosine
Cosine
Learning Rate Warmup
Epochs AdamW SinCos
40 Epochs
40 Epochs
Epochs
Optimizer
AdamW
AdamW
Adam W
Positional Encoding
SinCos
SinCos
SinCos
Weight Decay
0.1
0.1
0.1
0.1
C SELF-SUPERVISED TRAINING OF VISUAL REPRESENTATIONS
C.I HYPERPARAMETERS
Here we outline the hyperparameters of our experimental setup for SimCLR+ViT stability. For variations, alongside their default hyperparameters see Table 6 These hyperparameters are used in all SimCLR runs unless stated otherwise.
Augmentations We SimCLR augmentations throughout, however, ter strength, equal to the Colorlitter strength of MoCo v3. For completeness, we provide our training augmentation here, our testing augmentation is the standard resize, center crop and normalize. Half color strength corresponds to col wwjjitte_sstreeeee coloz.jitter.strengt
at half ColorJit-
= 0.5. Setting recovers the base SimCLR training augmentations.
transforme.u Raw.oneesizeddCrgr mmmee_aize_overreee ccalescrop_ccala interpclat peteppllation-Imag.... BICUBIC BICUBC
transforma forma dandomiiorizontaaraa.. (p- (p=0.5), cransforms RandomApply
trans
transforms. Colorditter brigh brightness=0.e
contrast-0.3 color_jiiter. stee_eghtte_stteeghhh
saturation=0.8 polor_ polor_jicter_ senggt hue=0.2 oolor__itterr__trenthh
1 p=0.8, }, transforms RandemSrayscale (p=0.2),
framsforms... RandomApply [M. (M..aausi.io anslur ([0.1, 2.0)),, p=0.5),
transtorma ToTenaor (),
IMAGINET IMAGINDT_KORVALI222
C.2 REDUCED LEARNING RATE WARMUP
In Chen et al. 12021 the authors noted that the learning rate warmup period needed extending from its typical ImageNetlk default of 10 epochs to 40 epochs, enhancing the stability of the method. We observe that using Reparam, either with or without Layer Norm, we are able to achieve stable SimCLR+ VIT training the original warmup period of 10 epochs (see Rigure 6).
As with our analysis at the longer warmup period, we also investigate the performance distribution across the trials, giving sense of how instability impacts the final model (see (eeeFeg 6).
14

---

Under review as a conference paper at ICLR 2023
(55
Bucta:
Text Top
Banka
Aappaaeb clapaanti
0-noem ge enntopy (buts)
Drace Fodar
Test Top 1 (56)
grad Ear nOOT
Training Epochs
Training Epochs
(a) Statistics of best and worst trials per method.
(b) Stability over 8 trials per method.
Figure 6: Eight trials of SimCLR for each method on ImageNetlk with 10 epochs of learning rate warmup. (a) Linear probe performance for the best (solid line) and worst (dshed line) trials of each method, against relevant metrics from the first attention layer (top to bottom): attention entropy. the spectral norm of the attention weights, and 6oo -sodgiedi norm of the attention weights. Our observations are consistent longer warmup 40 epochs investigated in Fiigure 3] except that here, Frozen Potcher
those of able to tame carly layer gradient norms than it was in the longer warmup (dashed green linc). (b) Linear probe performance of every trial. Observations Norm are the most stable methods. RReaarm (0.01) refers oReparam with initialization scheme of crunc_norm rmal ( (.01) instead of trunc.no .nommal ( 0.02), with former showing some signs of instability. Understanding the source of this instability will be the subject future work. oReparum LN uses the default t Crunc_norma (. (2).
again consistent with the longer warmup; Reparam with and without Layer
Figure 7: Linear probe performance on ImageNet at the end of training over 8 trials for each method. Trials ordered by decreasing performance, with run rank (8) corresponding to best (worst) trial. Frozen Patcher produce the best individual, with all other methods marginally lower. &Reparam + LN and a Reparam are the methods most reliably giving good performance, with Baseline and Frozen Patcher each susceptible to least one instability type.
Test Top (%)
90
40
Froze Putcher
oReqanm
dRuquaa (0.00)
Run rank
D SPEECH EXPERIMENTS
D.1 ABLATIONS ON THE INITIALIZATION FOR 3REPARAM
First, we found that it is better to initialize 7 as and not compute it from the initialized kernel as there could be different values for spectral norm depending on the initialization of the kernel. In this case we observed values greater than for the spectral norm which cause divergence no training. From practical point it is native to keep 7 1. We compared different initializations for kernel and we didn't i see any differences in initialization (e.g. uniform, normal). The only thing influences the std of the initialization pdf which influences also effective LR. In speech we found that training is robust with respect to changes of std (Figure 5). however larger std performs better and sweet spot is 0.2-0.3.
15

---

Under review as a conference paper at ICLR 2023
D.2 FULL LIBRISPEECH EXPERIMENTS
We also evaluate Reparam for large scale data in speech domain: we take now the whole LibriSpeech the training data. We consider again Adagrad optimizer with two schedules learning rate: cosine (with phase of 500k iterations) and step-wise decaying as before for Irain-clean- 100 experiments. We use exactly the same architecture and hyper-papse parameters as Table except dropout and layer drop which are decreased to 0.1 to decrease model regularization. For all models we tune only learning rate. Keys and queries spectral reparametrization is done separately from values, also we use learning rate gamma to be twice bigger than the main learning rate. Our experiments for train-clean-100 show, see Tables 7 and g that Reparam accompanied with post-LN can match the post-LN baseline, while having robustness to the hyper- hyper-parameter changes (e.g. allows larger learning rate values without any issues).
Table 7: Comparison between different normalizations and our re- p--paamtttizationn for speech domain on full LibriSpeech with step-wise LR schedule: word error rate reported for the best models.
posi-LN 2021b
posi-LN
pre-LN (same)
pre-LN (optimized)
oReparam
a Reparrm +post-LN
Likhomanenko et al.
dev-clean WER dev-other WER
2.6
7.0
6.9
7.7
7.1
test-clean WER
2.7
2.7
3.0
2.8
2.9
2.9
lest-other WER
6.9
6.9
7.8
6.8
7.3
7.0
Table 8: Comparison between different normalizations and our re-parametrization for speech domain on full LibriSpeech with cosine LR schedule: word eπror rate are reported for the best models.
post-LN
pre-LN (same)
dReparam
Reparam +post-LL LN
dev-clean WER
2.6
2.6
2.8
2.7
dev-other WER
7.1
7.6
test-clean WER
2.9
2.8
3.0
2.9
test-other WER
7.2
7.0
7.7
7.2
D.3 ABLATIONS ON SEPARATE FREPARAM FOR KEY, QUERIES AND VALUES
We found that in the end they behaves achieve lower training loss due to larger capacity ability which provides potential to scale. However, for training with LARS is better to have joint re- pe-aarametrzzatoon achieve stable training and comparable results with adaptive optimizers, see Section 4.3.2
or less similar while separate normalization allows to
D.4 HYPERPARAMETERS
We present hyperparameters for our speech experiments in Table 9 and speech experiments with LARS in Table 10
16

---

Under review as a conference paper at ICLR 2023
Table Training hyperparameter comparison for speech domain, Table
post-LNN
pre-LN
RReparam
Reparam + post-LN
dev-clean
5.9
6.2
6.4
6.1
dev-other
17.7
19.1
20.5
17.8
Weight Init
uni fniifomm. (.036)
uniform (.036)
tranc_ncrmal (.1)
arunc_normal (.
oReparam
No
No
Yes
Yes Yes
Layer Norm
Yes
Yes
No
Basc LR
0.03
0.5
Optimizer
Adagrad step(330k, 0.5)
I.R schedule
Batch size
240s
Weight decay
none
Warmup steps
64k
Training steps
500k
Dropout
0.3
Stoch. Depth
0.3
SpecAugment
- 30, T 50, 0.1, fmask 2, tmask 10
Grad. clip
Table 10: Training hyperparameter comparison for speech domain trained with LARS, Table
post-LN
pre-LN
oReparam
dReparrm post-LNN
dev-clcan dev-other
8.1
5.9
7.2
5.8
25
18.9
22.8
18.1
Weight Init
ani form (. (.036)
form (.036)
tranc_ncrmal (.1)
trunc_no mmr.a (.1)
oReparam
No
No
Yes
Yes
Layer Norm
Yes
Yes
No
Yes
Base I.R
0.1
0.5
0.3
Optimizer
LARS
Momentum
0.9
IR schedule
step(300k, 0.5)
Batch size
240s x
Weight decay
none
Warmup steps
Training steps
Dropout
Stoch. Depth
0.3
SpecAugment
= 30, 50. p 0.1, maak 2, tmask 10
Grad. clip
E HYPERPARAMETERS FOR SUPERVISED VISION
As mentioned in Section 4.1 we compare oReparam against DeiT (Touvron et al. 2021 and the MAE (He et al. 2022) (Appendix A.2) supervised training recipes for vision transformers. In Table I1 we highlight the differences between DeiT, MAE supervised and &Reparam. 3Reparam presents simplified and stable training objective for ViT-B variants. Table 12 we present the same comparing the ViT-I variants. There is no exact 1:1 comparison for VIT-L with the DeiT training framework so we only compare against the MAE supervised model.
17

---

Under review as a conference paper at ICLR 2023
Table 11: Training hyper-parameter comparison for supervised ViT-B/16.
DeiT
MAE
o Reearam
Top-1 EMA Top-1
81.8%
82.1%
81.88%
82.3%
82.37%
Weight Init
trunc.normal (. 02)
trunc.normal (, 02)
trunc_norma 1 ( 02)
Patcher Init
trunc.normal (.02)
trunc.normal (.02)
trunc_norma (.02)
oReparam
No
No
Yes
Layer Norm
Yes
No LARS(mom=009)
Optimizer
AdamW($3 92-0.5
Ad;mW(3;=0.9, 82=0.95)
Base LR
5 x 10
x 10
0.1
LR schedule
cosine
cosine
step(210, 0.1)
Batch size
1024
4096
4096
Weight decay
0.05
0.3
0.0
Warmup epochs
5
20
0
Training epochs
300
300
250
Label smoothing
0.1
0.1
0.1
Stoch. Depth
0.1
0.1
0.1
Repeated Aug.
2
2
2
RandAug
9/0.5
9/0.5
9/0.5
Mixup prob.
0.8
0.8
0.8
Cutmix prob.
Erasing prob.
0.25
0.25
0,25
Table 12: Training hyperparameter comparison for supervised ViT-L/16.
MAE
oReparam
Top-1
81.5%
82.41
EMA Top-1
82.6%
82.48 %
Weight Init
trunc normal (. 02)
trunc normal (.01)
Patcher Init
trunc normal (.02)
trunc normal (.0025)
oReparam
No
Yes
Layer Norm
Yes
Optimizer
Adam (00,= 32=0.55)
LARS(mom=0.9)
Base LR
1 10
0.15
LR schedule
cosine
cosine
Batch size
4096
4096
Weight decay
0.3
0.0
Warmup epochs
20
0
Training epochs
300
300
Label smoothing
0.1
Stoch. Depth
0.2
0.2
Repeated Aug.
2
RandAug
9/0.5
9/0.5
Mixup proh.
0.8
0.8
Cutmix prob.
1.0
1.0
Erasing prob.
0.25
0.25
18


Reviews:
[{"review_id": "HU2j0mQEUz", "reviewer": "Reviewer_3X2m", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the instability of Transformers by studying the progression of attention entropy: first a decrease followed by a quick increase then a long stable phase. A strong correlation exists between the minima of attention entropy and training stability. In the study, decreasing learning rate and increasing the warmup epoch leads to lower training loss. In addition, the following observations were made: Increasing the learning rate can lead to attention entropy collapse.To lower attention entropy,  reparametrizement of all linear layers with spectral normalization is considered.  The proposed algorithm is evaluated on 4 tasks: image classification, image self-supervised learning, ASR, and language modeling.", "summary_of_the_review": "See strength and weakness.", "strength_and_weaknesses": "(-) The readability of the paper needs much to be desired. It seems that instability is measured by the entropy of the attention head. Decreasing the learning rate or increasing warmup epoch leads to more stability. Without providing any explanation, the paper mentions that the training loss decreases with increase in learning rate. The paper should state whether decrease in training loss comes from stability. What is the consequence of entropy collapse? Explain how entropy collapse relates with the stability and accuracy? Is  stability and accuracy related? It is not clear how the lower bound of the attention entropy is related to stability.\n\n(-) The paper is proposing a re-parameterization of the weights of the linear layer using a spectral norm. It is not clear how this is related to regularization and performance increase? Many normalization methods have been proposed in the past, and it is not clear how the proposed is different previous normalization methods- a table should compare the performance of previous methods.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "5dLt90KIiw", "reviewer": "Reviewer_fFS8", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Setting: Transformer model, test on speech recognition and other tasks.\n\nMeasure attention entropy.\nFind correlation between the minima of attention entropy and the model’s training stability.\n(How is training stability measured?)\nReparametrize all linear layers with Spectral Normalization and an additional learned scalar.\n\nExperiments with σReparam on image classification, image self supervised learning, automatic speech recognition and language modeling.\n", "summary_of_the_review": "Unfortunately there are too many weaknesses which should be addressed. If those are addressed, and it can indeed yield improvements over good SOTA models, on actual relevant benchmarks, then this would be very interesting. But this requires much more work.\n\nIf the argument is only about more stable training and not about better performance, this anyway needs to be better justified and tested.\n", "strength_and_weaknesses": "Strength:\n- Somewhat novel reparameterization of the weights.\n- Some analysis on attention entropy and its correlation to the training stability.\n\nWeaknesses:\n- Very similar to weight norm, so the novelty is limited.\n- No actual comparisons to weight norm?\n- Source code not released?\n- Results are not good?\n- Librispeech train-clean-100 is not representative as a speech recognition task.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "4q4Q7_4E-g", "reviewer": "Reviewer_janJ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes the $\\sigma$-reparameterised transformer. The authors note a phenomenon in during training of transformers where the 'attention entropy' (essentially, how sharply peaked the attention scores are) decreases and then increases again. In the cases where the transformer fails to learn, this attention entropy does not recover. In order to combat this, the $\\sigma$-reparameterisation is introduced, essentially a version of spectral normalization where the matrices are rescaled by a learned parameter (instead of having unit norm). This provides impressive gains in stability of training and empirical performance on a variety of tasks. ", "summary_of_the_review": "The paper is a good empirical contribution to the state of the art of training transformer-based models, giving a new form of parameterisation of the weight matrices that dramatically improves the robustness to hyperparameter choices. This will lead to improved reliability of training transformers. The scientific argument analysing the attention collapse is a bit weak, and could be improved. But since that is not the focus of the paper, it shouldn't be required for an accept. \n\nUpdate after rebuttal:\nAfter reading the authors' rebuttal, I believe they have addressed my queries and some of the insightful remarks from the other reviewers.\n\nUpdate after end of rebuttal period: \nGiven the lack of response by the reviewers to the rebuttal to the experimental queries from the other reviewers, I'm unsure to what extent the issues raised by the other reviewers have been satisfied. Since I'm not completely familiar with the state of the art in stable transformer architectures, I've reduced my confidence. ", "strength_and_weaknesses": "Strengths\n+ The empirical results are strong, showing increased robustness to changing hyperparameters and similar or better final performance. This is a good result for democratisation of ML, since it reduces the amount of time and money that must be spent on fiddly hyperparameter tuning. Figure 4 is particularly compelling. \n+ The method is extremely simple to implement, yet has (to my knowledge) not been proposed before.\n+ The experimental results are fairly comprehensive. \n\nWeaknesses\n+ While I believe the final experimental results stand on their own, on a scientific level I am not convinced by the argument relating the failure to train to the collapse of the attention-entropy. In the case in figure 1, the two phenomena are definitely correlated, but this doesn't appear to continue in figure 2, where the baseline attention entropy 'collapses' but still gives good downstream performance. Similarly, none of the successful baseline models in figure 3 seem to have this 'dip' phenomenon where the entropy decreases then increases as the model begins to learn properly. Similarly, the connection between the singular values and attention entropy doesn't seem to hold practically: look at figure 3 where we have the green solid and dotted line have similar singular values but completely different attention entropy. It seems plausible that the collapse of the attention entropy and the failure of downstream performance could both be symptoms of some other problem, which the $\\sigma$-reparameterisation avoids. I realise the authors are careful to not say that entropy collapse causes training instability, but I believe it could be possible to get some direct evidence for this claim, which would strengthen the scientific contribution of the paper. \n  \n  + Nitpicks:\n    + Could you try a baseline of dividing the initalization by a factor of the spectral norm of a random matrix?\n    + In the theory you use $\\|W_KW_Q^\\top\\|_2$, but in the implementation it seems that you apply the reparameterisation on all of the weight matrices. Did you try only reparameterising the $W_KW_Q^\\top$ matrix directly?\n    + How does the overall memory usage compare when using LARS and fp32 on attention compared to adam and full mixed precision?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "1IWuJsFdXf", "reviewer": "Reviewer_qd4k", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The author has discussed the transformers models' training stability on various tasks. The author has investigated the sharpness of attention concerning attention entropy on each attention head during training. The author observes that the attention entropy first decreases then increases, and then enters into the long stable stage. To address the issue, the author has proposed a simple technique \\sigma_reparametrization, which reparametrize all the linear layers with spectral normalization and learnable scalar multiplication. Finally, the author has evaluated the proposed techniques on several tasks, such as supervised image classification, self-supervised learning, automatic speech recognition, and language modeling tasks, and claimed better stability and robustness concerning the choice of hyperparameters.", "summary_of_the_review": "\nJustification :\nThe overall writing quality is ok, and the proposed method is simple and beneficial. However, the experiment comparison lacks justification, and the technical contribution is plain.\n", "strength_and_weaknesses": "Strength:\nThe main reason to accept this paper is empirical results and showing performance o various tasks. The author has provided more quantitative results on various tasks.\n\n\nWeaknesses:\n\nThe current literature survey does not discuss the advantages and disadvantages of the paper. The author should validate why I should include this paper in this current work. \n\nHow did you get this equation-2? Is there any intuition of this equation?\n\n\nThe generalization of the method. The equation-2 is valid for the linear layer. What happened to the convolution layer? Is this valuable reparametrization for the convolution network?\n\nThe paper does not explain the relation between attention entropy and \\sigma_reparameterisation.\n\nHow does this learnable parameter contribute to the stability of the transformer?\n\nInstead of evaluating the proposed method on multiple tasks, the author could have done an ablation analysis on different kinds of networks and various types of transformer architecture. The author should provide ablation analysis results to clarify the proposed method's work principle.\n\nWhy do we need such reparametrization, learning another parameter? Did you analyze other kinds of techniques to tackle the attention entropy collapse problem? This requires a substantial literature survey and motivation for the proposed method. Can you justify this? How is the attention entropy collapse solved by the proposed reparametrization technique?\n\n\nThe author should compare the stability analysis results of this transformer with the existing state-of-the-art stable transformer. Also need to add a small paragraph in the related work section about stability/robustness in the transformers.\n\nCan you elaborate on the number of parameters and FLOPS required for this? What about the complexity of the proposed method?\n\nThe name spectral reparametrization in the title makes it confusing here. It seems it works on frequency domain reparametrization. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject