Anonymous authors
Paper under double-blind review
A BSTRACT
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 σ 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.
σ Reparam decouples the growth rate of a weight matrix’s spectral norm from its dimensionality, which we verify empirically.
We conduct experiments with σ Reparam on image classification, image self supervised learning, automatic speech recognition and language modeling tasks.
We show that σ Reparam provides great stability and robustness with respect to the choice of hyperparameters.
I NTRODUCTION
Transformers (Vaswani et al., 2017) are state-of-the-art 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 et 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.
; 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 in a pattern consisting of three phases.
In the beginning, attention entropy starts high (corresponding to uniform attention scores) and quickly drops to a 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 1 for an illustration, which is a Vision Transformer (Touvron et al., 2021) (ViT) trained on
ImageNet classification, using well optimized hyper parameters.
Empirically, we have found that the attention entropy is directly correlated with the model’s stability and convergence.
In particular, small attention entropy reached in the initial phase often causes slow
[Figure]
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 1 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 on training where the attention entropy collapses to near 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 can more easily lead to entropy collapse.
We then provide a simple fix, dubbed σ Reparam, which reparameterizes all weight matrices by sequentially applying Spectral Normalization (Miyato et al., 2018) and a learned multiplicative scalar.
Intuitively, σ Reparam decouples the update of the spectral norms of weights from their dimensionality, which allows them to update smoothly in a controlled way.
Also note that σ Reparam does not change the model space, which allows one to learn an arbitrarily expressive model.
We validate σ Reparam 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 curves 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.
R ELATED W ORKS
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 et al., 2022; Shleifer et al.,
σ Reparam 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, including Zhang et al.
; Huang et al.
; Yang et al.
; Bachlechner et al.
σ Reparam is an orthogonal approach as it addresses the training dynamics of attention layers, which makes it compatible with standard initialization methods and provides robust performance.
σ Reparam is a 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 reparameterization (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.
M ETHOD 3. A TTENTION E NTROPY
At the core of Transformers is dot product attention.
Let X ∈ R T ×
d denote an input sequence to 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 W K ,
W Q ∈ R d × n a , W V ∈ R d × n v denote the key, query and value matrices.
A simple attention layer then computes Att ( X ) =
AXW V where
A = ψ ( a ) , a = XW K W ⊤ Q X ⊤ and ψ is the row-wise softmax function.
We define the attention entropy of a row i of A by Ent ( A i ) =
− (cid:80)
Tj =1 A i,j log( A i,j ) .
We also overload the notation and let
Ent ( A ) =
1 T (cid:80)
Ti =1 Ent ( A i ) denote the average attention entropy of A .
As shown in Figure 1, the attention entropy (and the entropy collapse phenomenon) is a strong indicator of training stability of
Transformers.
Our goal is to alleviate the entropy collapse problem and achieve a 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 ⊤ Q .
Theorem 3.1 (Attention entropy lower bound) .
Assume without loss of generality ∥
X ∥ 2 ≤ 1 , and let spectral norm σ = ∥
W K W ⊤ Q
∥ 2 .
Then it holds that:
[Equation]
Moreover, there exists inputs X and weights W K , W Q for which the lower bound in Eq.
(1) is tight.
Therefore, for large σ, T , the minimum attainable entropy behaves like Ω( Tσe − σ ) .
We note that the bound on the entropy in Theorem 3.1 is tight in a sense that it is achievable for some inputs X .
Moreover, the typical Frobenious (L2) regularization would not ensure a small σ (a small Frobenious norm is much more restrictive than a small Spectral norm), hence it would not be as effective in preventing an “entropy collapse”.
Proofs for Theorem3.1 and the following Proposition are provided in Appendix A.
3. σ R EPARAM
We then present σ Reparam, a method to re-parameterize the weights of a linear layer with:
[Equation]
where σ ( W ) ∈ R is the spectral norm of W and γ ∈ R is a learnable parameter, initialized to 1.
In practice, σ ( W ) can be computed via power iteration (Mises & Pollaczek-Geiringer, 1929) as in
Spectral Normalization (SN) (Miyato et al., 2018), see Algorithm 1 for a sketch implementation.
Note that σ Reparam 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, one can compute (cid:99) W once and freeze it, which means that it has the same cost as a regular linear layer.
Also note that the WN configuration leads to immediate divergence without using Layer and here we only report the result with
+ [Table]
Why σ Reparam?
Unlike the standard SN, σ Reparam introduces an additional multiplier γ which explicitly controls the SN of the weights [Table] and there is no explicit pressure to regularize the SN.
The additional multiplier is necessary to avoid restricting the capacity of the network [Table] and we find that training and overall performance is 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 a full theoretical characterization is 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 = µ + ϵ ∈ R w × w , where µ is the mean and ϵ is a random variable with E
[ ϵ ] = 0 , E [ ϵ 2 ]
= n 2 ∈ R w × w .
A typical Adam optimizer update attempts to approximate the following ideal update: ∆ =
E [ g ] √ E [ g 2 ] .
The following proposition lower bounds the spectral norm of the ideal update σ (∆) :
Proposition 3.2.
It holds that:
[Equation]
Note that the noise second moment n 2 is typically in the order of µ 2 , hence Eq.
(3) indicates that the spectral norm of the ideal update should be large, growing linearly with √ w .
Moreover, for large batch sizes we would have n 2 ≪ 1 , resulting in σ (∆) ∼ √ w 1 .
While such a large spectral norm could be offset by a 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 as attention layers, would require a per layer learning rate tuning.
In contrast, σ Reparam avoids this issue since the spectral norm of each layer is controlled by a single parameter γ , hence the size of its update does not scale with w and is uniform across layers.
E XPERIMENTS 4. S UPERVISED I MAGE C LASSIFICATION
Improved robustness.
We first start from a well tuned recipe with ViT-B on ImageNet-1k (Touvron et al., 2021), and vary its hyper parameters in the grid [ base lr ∈ { 5 e − 4 , 1 e − 3 } , batch size ∈ { 1024 , 2048 } , warmup epochs ∈ { 0 , 5 } ] .
7/8 configurations lead to divergence except for the default [5 e − 4 , 2048 , 5] hyper parameter.
We next apply σ Reparam 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 1).
This suggests improved robustness with respect to hyperparameters.
Simplified recipe.
σ Reparam also enables a simplified framework for training ViT-B and ViT-L models, in contrast to state-of-the art ImageNet-1k ViT training protocols such as the fully supervised
MAE recipe (He et al., 2022) and DeiT (Touvron et al., 2021), (Table 1).
In the case of ViT-B models, we are able to train for a shorter duration, remove all LayerNorm layers, remove LR warmup, remove cosine scheduling (requiring only a simple step schedule at 210 epochs) and use no weight decay.
Furthermore, σ Reparam enables SGD training via LARS (You et al., 2017) (with momentum 0.9) – something not possible with traditional ViT training protocols (Touvron et al., 2021;
He et al., 2022).
These simplifications also have the added benefit of reducing GPU memory overhead 2 .
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 of the network.
The full set of hyperparameters is available in Appendix E.
[Figure]
To further understand the effect of σ Reparam, we track both the attention entropy, and the largest singular value of the attention weight matrix over the course of training.
In Figure 2, σ Reparam maintains a lower largest attention weight singular value and presents a higher, but monotonically decreasing attention entropy throughout training.
As previously discussed, a smaller bounded singular value helps with stable training, whereas a higher attention entropy encourages exploration of more diverse solutions.
This is reinforced by the accelerated performance observed in Test Top 1 and the 50 epoch reduction in training time for the σ Reparam ViT-B/16 shown in Figure 2.
4. S ELF - SUPERVISED TRAINING OF VISUAL REPRESENTATIONS
In computer vision, self-supervised learning (SSL) has been effective in enabling efficient training on downstream tasks (Assran et al., 2022).
Most of this progress has been made using convolutional architectures, while works using ViTs 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.
These instabilities can be remedied through a combination of frozen patch embedders, initialization [Table] [Figure]
[Table]
[Figure] [Figure] [Figure]
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.
Here, we demonstrate that σ Reparam is a ViT SSL stabilizer.
Taking SimCLR as our SSL method, we investigate four variants.
Baseline and Frozen Patcher were studied in Chen et al.
(2021), whereas σ Reparam and σ Reparam + LN are our solution.
These methods are detailed in Table 2 and their full hyperparameters given in Table 6 of Appendix C.1. [Figure]
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 σ Reparam + LN , whereas the most stable methods are σ Reparam + LN and σ Reparam .
Our main stability experiments use 40 epochs of learning rate warmup, matching the setting of Chen et al.
Using σ Reparam, as in the supervised setting, gives training stability even at the lower learning rate warmup of 10 epochs.
For more details, see Appendix C.2.
Finally, we look at the performance attainable when training for a longer duration of 300 epochs in
Table 2.
The best performing method run is given by with σ Reparam + LN , with Frozen Patcher performing almost as 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 σ Reparam + LN , producing both the highest ImageNet1k linear probe performance at 100 epochs (69.6 %) and 300 epochs (74.5 %) epochs, as well as very stable training over many trails, both at long and short learning rate warmup.
4. S PEECH
In this section we focus on experiments for automatic speech recognition (ASR).
Data All experiments are performed on the subset of 100h audio paired with transcriptions ( trainclean-100 ) of LibriSpeech dataset Panayotov et al.
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 a 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-the-art model on 100h of LibriSpeech (Likhomanenko et al., 2021a).
The model consists of 1D convolution to perform striding, Transformer encoder with post-LN and a final linear layer to map to the output number of tokens 3 .
The model is trained with Connectionist Temporal Classification (Graves et al., 2006) loss.
To speed up the model training (2-3x) and decrease memory usage we are using CAPE positional embeddings (Likhomanenko et al., 2021c) instead of relative embeddings Shaw et al.
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 p = 0 .
1 and time mask parameter T = 50 ; time warping is not used.
Training We use Adagrad (Duchi et al., 2011) if not specified otherwise, and LR decaying byeach time the WER reaches a 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 [Table] weight decay is used.
Default warmup is 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.
[Table]
T RAINING STABILITY , ROBUSTNESS AND GENERALIZATION
First, we experiment with stability of training for the baselines using both “Pre Norm” (pre-LN) and “Post Norm” (post-LN) architectures.
If we vary LR, warmup
[Table]
and gradient clipping, all post-LN experiments either diverge or no training is observed.
At the same time, pre-LN is stable: we can reduce warmup from 64k to 16k, increase learning rate from 0.03 to 0.
[Table]
and obtain better results than before.
While pre-LN is more stable than post-LN, it generalizes worse: validation WER is worse while training loss is lower, see Table 3.
When we switch to σ Reparam we observe the same stability as 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 σ Reparam, which allows us to achieve similar performance on the dev and test sets and lower training loss.
In
Figure 5 both σ Reparam and σ Reparam with post-LN demonstrate robustness with respect to training hyperparameters.
We also compare with Spectral Norm (SN) where γ is set to 1 and is not learnable and WN baselines.
Both SN and WN perform poor compared to σ Reparam, see Table 3.
In prior works it was reported that post-LN can be impossible to train with very deep architectures, see e.g. Liu et al.
(2020b;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 σ Reparam and combination of σ Reparm and post-LN: for both cases out of the box models train well and achieve similar results as pre-Norm.
This confirms σ Reparam’s ability for stable training even with post-LN. [Figure] [Figure] [Figure]
[Figure]
L e a r n i n g R a t e
[Figure]
S t d o f i n i t i a li z a t i o n ( kq , v ) + post-LN ( kq , v ) ( kqv )
+ post-LN ( kqv )
4.3. T RAINING 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.
It 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 4.
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 σ Reparam (key, queries and values are represented as one matrix)
we observe stable training with respect to LR changes, [Table] and combined together with post-LN it achieves similar performance to the best results from Tablewhile keeping the train loss low 4 .
4. L ANGUAGE
Setup.
We use the WikiText-103 language model (LM) benchmark, which consists of 103M tokens sampled from English Wikipedia (Merity et al., 2017).
Our baseline is a highly optimized Transformer (Baevski & Auli, 2019) with 32 layers, 8 heads, 128 head dimensions, 1024 model dimensions,fully connected dimensions and post LayerNorm.
The word embedding and softmax matrices are tied (Press & Wolf, 2017).
We partition the training data into non-overlapping 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 token curse , Press et al., 2021).
[Table]
Results. We do not experience training instability with the baseline Transformer, likely because the masked attention in autoregressive models makes entropy collapse less likely to occur.
Nonetheless, we experimented with σ Reparam to test its generality on a different modality/problem.
We apply σ Reparam to 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 [1e-3, 1e-4, 0].
All other hyperparameters are kept the same as the baseline.
The results are shown in Table 5.
We see that even in the absence of LayerNorm, σ Reparam shows strong performance in convergence and dev/test performance.
With a mild weight decay, σ Reparam also outperforms the baseline wrt the dev/test
PPL.
C ONCLUSION
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 σ Reparam.
We demonstrate over a 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.
A P ROOF OF T HEOREM 3.1 AND P ROPOSITION 3.
Theorem 3.1 (Attention entropy lower bound) .
Assume without loss of generality ∥
X ∥ 2 ≤ 1 , and let spectral norm σ = ∥
W K W ⊤ Q
∥ 2 .
Then it holds that:
[Equation]
Moreover, there exists inputs X and weights W K , W Q for which the lower bound in Eq.
(1) is tight.
Proof.
WLOG let u ∈ R T denote the j ’th row of a .
From the condition that ∥ X ∥ 2 ≤ 1 it holds that ∥ u ∥ ≤ σ .
Let p = p ( u ) denote the softmax probabilities given by:
[Equation]
where Z = (cid:80) Tj =
e u j is the partition function.
The entropy given p ( u ) is then:
[Equation]
We wish to solve the following constrained minimization problem:
[Equation]
where D > 0 .
Define the lagrangian:
[Equation]
To find all saddle points, we solve the system of equations:
[Equation]
Giving rise to the following set of equations:
[Equation]
As a first step, assume that for the minimizer u ⋆ of Eq.
(6) there exists an index k
such that u ⋆ k = 0 .
Using Eq.
[Equation]
From the first set of equations we arrive at the condition:
[Equation] [Equation]
This however implies that u ⋆ 1 = u ⋆ 2 = ...
=
u ⋆T = 0 , hence a contradiction to Eq.
Now, assuming ∀ k u k ̸ = 0 , we have that:
[Equation]
[Equation]
The monotonicity of the LHS of Eq.
(16) implies that u contains only 2 distinct values.
WLOG assume u ⋆ 1 =
α, ∀ i> 1 , u ⋆i = − (cid:113) D 2 − α 2 T − 1 .
Then we have:
[Equation]
With a solution:
[Equation]
With the corresponding entropy:
[Equation]
Proposition A.1.
It holds that:
[Equation]
Proof.
We have that:
[Equation]
B I MPLEMENTATION OF σ R EPARAM
To compute spectral norm of the current matrix we use the power method as approximation method to speed up computations.
See Algorithm 1 for a sketch implementation.
[Table]
[Table] [Table]
[Table]
C S ELF - SUPERVISED TRAINING OF VISUAL REPRESENTATIONS C. H YPERPARAMETERS
Here we outline the hyperparameters of our experimental setup for SimCLR+ViT stability.
For the variations, alongside their default hyperparameters see Table 6.
These hyperparameters are used in all SimCLR runs unless stated otherwise.
Augmentations We use SimCLR augmentations throughout, however, we run at half ColorJitter strength, equal to the ColorJitter 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 color jitter strength = 0.5 .
Setting color jitter strength = 1.0 recovers the base SimCLR training augmentations.
[Table]
C. R EDUCED LEARNING RATE WARMUP
In Chen et al.
the authors noted that the learning rate warmup period needed extending from its typical ImageNet1k 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 at the original warmup period of 10 epochs (see Figure 6).
As with our analysis at the longer warmup period, we also investigate the performance distribution across the trials, giving a sense of how instability impacts the final model (see Figure 6).
[Figure]
Understanding the source of this instability will be the subject of future
[Figure]
D S PEECH E XPERIMENTS D. A BLATIONS ON THE INITIALIZATION FOR σ R EPARAM
First, we found that it is better to initialize γ as 1 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 1 for the spectral norm which cause divergence / no training.
From practical point it is native to keep γ = 1 .
We compared different initializations for kernel and we didn’t see any differences in initialization (e.g. uniform, normal).
The only thing influences is 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.
D. F ULL L IBRI S PEECH E XPERIMENTS
We also evaluate σ Reparam for large scale data in speech domain: we take now the whole LibriSpeech as the training data.
We consider again Adagrad optimizer with two schedules on learning rate: cosine (with 1 phase of 500k iterations) and step-wise decaying as before for train-clean-100 experiments.
We use exactly the same architecture and hyper-parameters as in Table 9 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 on gamma to be twice bigger than the main learning rate. Our experiments as for train-clean-100 show, see Tables 7 and 8, that σ Reparam accompanied with post-LN can match the post-LN baseline, while having robustness to the hyper-parameter changes (e.g. allows larger learning rate values without any issues) [Table] [Table]
[Table]
D. A BLATIONS ON SEPARATE σ R EPARAM FOR KEY , QUERIES AND VALUES
We found that in the end they behaves more or less similar while separate normalization allows to achieve lower training loss due to larger capacity ability which provides potential to scale.
However, for training with LARS it is better to have joint re-parametrization to achieve stable training and comparable results with adaptive optimizers, see Section 4.3.2.
D. H YPERPARAMETERS
We present hyperparameters for our speech experiments in Table 9 and speech experiments with
LARS in Table 10.
[Table]
[Table]
E H YPERPARAMETERS FOR SUPERVISED VISION
As mentioned in Section 4.1 we compare σ Reparam against DeiT (Touvron et al., 2021) and the
MAE (He et al., 2022) (Appendix A.2) supervised training recipes for vision transformers.
In Table 11 we highlight the differences between DeiT, MAE supervised and σ Reparam.
σ Reparam presents a simplified and stable training objective for ViT-B variants.
In Table 12 we present the same comparing the ViT-L variants.
There is no exact 1:1 comparison for a ViT-L with the DeiT training framework so we only compare against the MAE supervised model.
[Table]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table] [Figure] [Table] [Figure]