Stabilizing Transformer Training by Preventing Attention Entropy Collapse

Introduction

Transformers (Vaswani et al., 2017) are state-of-the-art models in many application domains. 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-LN fashion). There has since been various works attempting to promote better training stability and robustness. For example, the pre-LN (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 promote better signal propagation. Zhang et al. (2019); Huang et al. (2020) remove LNs with carefully designed initializations.

In this work, we study the training instability of Transformers through the lens of training dynamics. We start by monitoring the entropy of the attention maps averaged over all query positions, heads and examples. We have found that the attention entropy is tightly correlated with the model’s stability and convergence. In particular, small attention entropy is often accompanied with slow convergence, fluctuations in training loss and, in the worst case, divergence. As a motivator, we plot the attention entropy curves of a highly optimized Vision Transformer (ViT) (Dosovitskiy et al., 2021; Touvron et al., 2021) in Figure 2. We observe an initial loss oscillation happening at the same time with sharp dips of the attention entropy curves. When doubling the default learning rate, all attention entropy collapses to near zero and training diverges. In addition, we show in Figures 4 and 7 two sets of experiments of baseline Transformers models with training instability occurring at the same time of entropy collapse. And more generally, similar observations can be made in a wide range of model/task settings if hyperparameters such as learning rate, warmup, initialization are not carefully tuned.

To further demonstrate this connection, we modify the Transformer to have a global temperature by dividing the pre-softmax (logits) matrix of each attention mechanism by a scalar quantity whose default value is 1. Modifying the temperature gives direct control over the attention entropy, enabling the investigation of a causal connection between entropy collapse and training instability (see Figure 2 and Figures 8 and 9 in Appendix B). Here we train a ViT-B/16 on ImageNet1k. At an intervention epoch we modify the temperature from its default value to 0.1. We see that when performing this intervention during warmup, attention entropy drops to near zero and training becomes unstable. A late intervention also causes a drop in entropy and accuracy curves, however, the model is able to recover to a higher attention entropy regime, although yielding a lower accuracy than non-intervened training.

To further understand these phenomena, we computed the sharpness – the largest singular value of the Hessian (the second order derivative of the loss with respect to the model parameters), as its magnitude has implications for training stability Ghorbani et al. (2019); Yao et al. (2020); Cohen et al. (2021, 2022); Gilmer et al. (2021). When sharpness exceeds an algorithm-dependent stability threshold, training iterations diverge Cohen et al. (2021, 2022). We see that interventions inducing the largest drop in attention entropy result in the sharpness exceeding the stability threshold, whereas the later interventions do not cause the threshold to be crossed, explaining how they can recover. For details on the empirical setup and additional results see Appendix B.

The empirical correlation of entropy collapse and training instability leads to the following questions: 1) How do we prevent entropy collapse? 2) Can we improve training stability by doing so? We answer these by showing that entropy collapse can be effectively prevented by controlling the spectral norms of the query and key projections. In particular, we prove a tight lower bound on the attention entropy, which decreases exponentially fast with the growth of the spectral norm of the attention matrix logits. This bound suggests that entropy collapse can occur swiftly when letting the spectral norm of the weights increase uncontrollably. We then provide a simple fix, $\sigma$Reparam, which reparameterizes all weight matrices by sequentially applying Spectral Normalization (Miyato et al., 2018) and a learned multiplicative scalar. Intuitively, $\sigma$Reparam decouples the update of the spectral norms of weights from their dimensionality, which allows them to update smoothly and in a controlled way. Also note that $\sigma$Reparam does not change the model space, which allows one to learn an equally expressive model.

We evaluate five tasks: image classification, self-supervised learning (SSL), machine translation, automatic speech recognition (Appendix E), and language modeling (Appendix G). We highlight the empirical results as follows:

1. We show that entropy collapse is commonly observed in the baseline models of various benchmarks.

2. Image classification: $\sigma$Reparam enables a drastically simplified ViT training recipe by removing pre-LN, learning rate warmup, weight decay and not requiring adaptive optimizers. This recipe leads to equivalent (or slightly better) model performance against baseline training strategies, all the while reducing training duration by 16% .

3. Self-supervised learning: $\sigma$Reparam helps to drastically improve the stability and robustness of the SimCLR training, improving upon existing baselines.

4. Machine translation: $\sigma$Reparam allows us to stabilize very deep post-LN architectures up to 100L-100L encoder-decoder layers.

5. Speech recognition: $\sigma$Reparam allows us to improve training stability and simplify the training recipe for post-LN Transformer by removing learning rate warmup and adaptive optimization.

6. Language modeling: $\sigma$Reparam is compatible with causal Transformer architectures, and achieves results competitive with state-of-the-art without using post-LN.

Related Works

Transformers have relied heavily on LNs to achieve training stability. Besides the popular post-LN and pre-LN configurations, other variants have been proposed (Wang et al., 2022; Shleifer et al., 2021; Liu et al., 2020a). On the one hand, we show empirically that entropy collapse (and its accompanied training instability) happens even equipped with extensive use of normalization layers. On the other hand, $\sigma$Reparam does not rely on specific normalization layers and can even work in the absence of it, while effectively smoothing the attention entropy curves.

There have also been numerous attempts to design better Transformer initialization schemes, including Zhang et al. (2019); Huang et al. (2020); Yang et al. (2022); Bachlechner et al. (2021). While proper initializations are indeed crucial to stable and fast training, we argue that the training dynamics (affected by the optimizer and training hyperparameters) is equally important. $\sigma$Reparam in this sense is an orthogonal approach that specifically targets the entropy collapse problem, which makes it compatible with standard initialization methods and provides robust performance.

$\sigma$Reparam is a special case of weight reparameterization, which has found wide adoption in deep learning. WeightNorm (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, $\sigma$Reparam is the first simple reparameterization technique that provides competitive performance with well optimized baseline models. Normalizing weights by its spectral norm is also inspired by SpectralNorm (Miyato et al., 2018), with the key difference that SpectralNorm explicitly constrains the model’s capacity, which brings significant performance loss.

Another related line of work is the rank collapse of Transformer training, first identified by (Dong et al., 2021). Rank collapse refers to the degenerate state of attention where its output converges to a rank 1 matrix, where all tokens share the same representation. This analysis is further followed up by (Anagnostidis et al., 2022) suggesting that rank collapse causes vanishing gradient of the attention query and keys. Entropy collapse, on the other hand, characterizes a different failure pattern, where the attention matrix remains high rank, and it tends to introduce high gradient norms rather than vanishing gradients (see Figure 4).

Method

We next investigate the properties of attention entropy. We show in the next theorem that $\text{Ent}(A)$ is directly connected to the spectral norm (the largest singular value) of $W_{K}W_{Q}^{\top}$.

Let $\sigma=\|W_{K}W_{Q}^{\top}\|_{2}$, $\sigma_{x}=\|XX^{\top}\|_{2}$, $\bm{\sigma}=\sigma\sigma_{x}$ and $\beta=\exp\left(-\bm{\sigma}\sqrt{\tfrac{T}{T-1}}\right)$. Then it holds that:

Moreover, there exist inputs $X$ and weights $W_{K},W_{Q}$ for which the lower bound in Equation 1 is tight.

Therefore, for large $\sigma,T$, the minimum attainable entropy behaves like $\Omega(T\sigma e^{-\bm{\sigma}})$, hence decreasing exponentially fast with $\bm{\sigma}$. We note that the bound on the entropy in 3.1 is tight in a sense that it is achievable for some inputs $X$. Proofs for Theorem 3.1 and the following Proposition are provided in Appendix A.

Entropy collapse and training stability. Transformers are hard to train, requiring a careful tuning of a variety of hyperparameters. Notably, transformers can exhibit stages of training instability, with loss values oscillating uncontrollably, to the point of divergence. From a loss geometry perspective, we hypothesize that these regions of instability are caused when the weights enter a region of high curvature, a hypothesis supported by Chen et al. (2022), which showed that transformer models tend to converge to extremely sharp local minima. In this paper however, we step away from the loss geometry perspective and identify a novel empirical observation unique to the Transformer architecture. We observe that training instability and attention entropy collapse appear in tandem. Moreover, this observation is consistent across multiple settings and modalities (see Figures 12, 15, 7, 17 and 4). Equipped with this observation, we might ask whether preventing attention collapse might in turn prevent training instability. We highlight that the affirmative answer provided in this paper could prove extremely practical, as attention entropy is easier to compute and potentially manipulate then directly tackling the loss geometry, which typically involves computing second derivatives, as in Foret et al. (2021). We next describe out method for preventing entropy collapse through a simple reparameterization scheme.

2 σ𝜎\sigmaReparam

$\sigma$Reparam is a method to reparameterize the weights of a linear layer with:

The noise second moment $n^{2}$ is typically in the order of $\mu^{2}$, hence Equation 3 indicates that the spectral norm of the ideal update should be large, growing linearly with $\sqrt{w}$. Moreover, for large batch sizes we would have $n^{2}\ll 1$, resulting in $\sigma(\Delta)\sim\sqrt{w}$ This estimation would be exact for full batch optimization.. While such a large spectral norm could be offset by a proper learning rate adjustment, this would be counterproductive since 1) a small learning rate typically induces inferior performance, and 2) architectures with layers of varying sizes, such as the case in Transformers, would require a per layer learning rate tuning. In contrast, $\sigma$Reparam avoids this issue since the spectral norm of each layer is controlled by a single parameter $\gamma$, hence the size of its update does not scale with $w$ and is uniform across layers. This indicates $\sigma$Reparam should provide the models of improved robustness with respect to learning rate and other related hyperparameters, by maintaining the spectral norm of the weights (and as a result the attention entropy) in a healthy regime.

Experiments

Improved robustness. We first start from a well tuned recipe with ViT-B on ImageNet1k (Deng et al., 2009; Touvron et al., 2021), and vary its hyperparameters in the grid \Big{[}\texttt{baseLR}\in\{5\times 10^{-4},10^{-3}\}, $\texttt{batchSize}\in\{1024,\,2048\}$, \texttt{warmupEpochs}\in\{0,5\}\Big{]}. 7/8 configurations lead to divergence except for the default \Big{[}5\times 10^{-4}, $2048,$ 5\Big{]} hyperparameter. We next apply $\sigma$Reparam to all the linear layers (including the initial patch embedding), and remove all the pre-LNs instances. All configurations in the same grid search converge with an average top-1 accuracy of 81.4% ($\pm$0.52%) demonstrating improved robustness with respect to hyperparameters.

Simplified recipe. $\sigma$Reparam also enables a simplified framework for training ViT-B, ViT-L and ViT-H models, in contrast to state-of-the art ImageNet1k ViT training protocols such as the fully supervised MAE recipe (He et al., 2022) and DeiT (Touvron et al., 2021), see Table 1. In the case of ViT-B models, we are able to train for a shorter duration, remove all pre-LNs layers, remove learning rate (LR) warmup, remove cosine scheduling (requiring only a simple step schedule at 210 epochs) and use no weight decay. Furthermore, $\sigma$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 overheadWe observe a 8.2% memory reduction in full fp32 precision (for a 1:1 comparison) with a batch size of 86 per GPU.. For the ViT-L model we relax the LR schedule back to cosine and slightly increase the training interval to 300 epochs. All models use FP32 precision on the attention and $\sigma$Reparam operands and keep mixed precision training for the rest of the network. The full set of hyperparameters is available in Appendix H. We note that for larger models like the ViT-L/16 and ViT-H/14 a slight weight decay cosine schedule from $0.0$ to $10^{-5}$ enables easier training.

To further understand the effect of $\sigma$Reparam, we track both the attention entropy, and the largest singular value of the attention weight matrix over the course of training. In Figure 3, $\sigma$Reparam maintains lower spectral norms for the attention weight matrices and presents a higher, but monotonically decreasing attention entropy throughout training. The benefit of such smooth and bounded attention entropy curves is reinforced by the accelerated performance observed in Test Top 1 and the 50 epoch reduction in training time for the $\sigma$Reparam ViT-B/16 shown in Figure 3.

Finally, we extend $\sigma$Reparam to a much larger 11M sample training dataset, ImageNet21k (Ridnik et al., 2021), and train a ViT-B/16. We then finetune this model with ImageNet1k and report the performance in Table 2. We observe that $\sigma$Reparam presents competitive results against ViT-B/16’s trained on drastically larger datasets such as JFT3B (Zhai et al., 2022) and the 400M sample CLIP pre-training dataset (Dong et al., 2022), all the while presenting stable training and not requiring LayerNorm or LR warmup.

2 Self-Supervised Training of Visual Representations

In computer vision, 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., 2021). These instabilities can be remedied through a combination of frozen patch embedders, initialization 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., 2021).

Here, we demonstrate that $\sigma$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 $\sigma$Reparam and $\sigma$Reparam + pre-LN are our solution.

These methods are detailed in Table 3, and their full hyperparameters are given in Table 6 of Section D.1.

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. $\sigma$Reparam, however, can protect against both types of instability, yielding more reliable training (see Figures 5 and 4).

As noted in Chen et al. (2021), instabilities reduce final performance. We show the instability impact on performance below. in Figure 6. The methods with the best performing individual runs are Frozen Patcher and $\sigma$Reparam + pre-LN, whereas the most stable methods are $\sigma$Reparam + pre-LN and $\sigma$Reparam.

Our main stability experiments use 40 epochs of learning rate warmup, matching the setting of Chen et al. (2021). Using $\sigma$Reparam, as in the supervised setting, gives training stability even at the lower learning rate warmup of 10 epochs. For more details, see Section D.2.

Finally, we look at the performance attainable when training for a longer duration of 300 epochs in Table 3. The best performing method run is given by $\sigma$Reparam + pre-LN, with Frozen Patcher performing almost as well, and both outperforming the reference SimCLR result (Chen et al., 2021).

Ultimately, we see while $\sigma$Reparam produces the lowest degree of instability, the best overall method for stable training of SimCLR ViTs is $\sigma$Reparam + pre-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 trials, both at long and short learning rate warmup.

3 Machine Translation

In machine translation (MT) stable training of deep encoder-decoder post-LN Transformers is an active research area (Wang et al., 2022; Liu et al., 2020a). Vanishing gradients problem has been reported by many works, leading to different solutions including rescaling residual connections: e.g., Wang et al. (2022) trained a 1000-layer Transformer by properly rescaling residual connections and initialization depending on model depth, dubbed DeepNorm. We examined attention entropy collapse for the deep Transformers in MT and found that they suffer not only from vanishing gradients but also from entropy collapse, both for vanilla post-LN and DeepNorm. By injecting $\sigma$Reparam alongside post-LN/DeepNorm, we empirically show that it is able to bound attention entropy and stabilize training without any divergent training loss growth issues. Details on experiments and all findings are in Appendix F.

Empirical setup. We use standard WMT’17 English-German benchmark with newstest2016 as a validation and newstest2017 as test sets. We consider $N$L-$N$L encoder-decoder models with $N$ encoder and $N$ decoder layers, where $N=6,18,50,100$, for both post-LN and DeepNorm configurations. For all models we report BLEU score on validation and test sets across 3 runs with different seeds.

Attention entropy collapse occurs in deep models. While we reproduced stable results for 6L-6L post-LN and observed nicely bounded attention entropy behaviour, for 18L-18L configurations, divergence is observed when varying the random seed. By close inspection we observe no vanishing gradients problem, but attention entropy collapse clearly occurs during training. Deeper models, namely 50L-50L and 100L-100L, are unable to train due to vanishing gradients as well as attention entropy collapse for some of the deep layers (Figure 17). For DeepNorm while we are able to reproduce results for 6L-6L, 18L-18L and 50L-50L depths observing stable training (no any models diverged and training behaved well), yet we observe instability in training of the 100L-100L model, resulting in only 1 over 3 (different seeds) successful run. By closer inspection of the training behaviour we do not see any drastic issue of vanishing gradients, however we see attention entropy collapse, see Figures 7 and 18.

$\sigma$Reparam resolves entropy collapse in deep models. To alleviate attention entropy collapse and confirm $\sigma$Reparam effectiveness for deep models we inject $\sigma$Reparam into post-LN and DeepNorm models. As a result, $\sigma$Reparam nicely bounds attention entropy for 18L-18L and 50L-50L post-LN models (Figure 19), resolving any divergence issues as well as vanishing gradients in the 50L-50L model. $\sigma$Reparam also nicely bounds attention entropy for 18L-18L, 50L-50L, 100L-100L DeepNorm models (Figure 20), resolving any divergence issues for 100L-100L, see Figure 7 (vanishing gradients are not observed as DeepNorm targets it). In terms of performance (Table 4), $\sigma$Reparam with post-LN or DeepNorm matches their baselines for 6L-6L and in the same ballpark for 18L-18L. However, $\sigma$Reparam is inferior to DeepNorm for 50L-50L and 100L-100L.

4 Speech Recognition and Language Modeling

We also conduct empirical analysis of speech recognition in Appendix E and observe attention entropy collapse for different configurations. $\sigma$Reparam alongside with post-LN (a) stabilizes training of post-LN (b) improves robustness with respect to hyperparameters and (c) to the best of our knowledge, for the first time allows model training without an adaptive optimizer achieving stable training and comparable performance. For language modeling, see Appendix G, $\sigma$Reparam simplifies training recipe by removing all LayerNorms and achieves comparable performance to state-of-the-art.

Conclusion

Transformer training stability is a well acknowledged, but still unsolved problem. This problem comes with many facets, and there are multiple necessary conditions that need to be met in order to guarantee stable and robust training. Our work identifies attention entropy collapse as a unique failure pattern that seems to be commonly observed in a wide range of settings and tasks. We also show that $\sigma$Reparam as a simple reparameterization of the weights can effectively address the entropy collapse problem, which often leads to improved training stability and robustness.

There are also limitations of our work. First of all, it is unclear if there is a causal relationship between entropy collapse and training instability of Transformers. We believe that establishing such a connection will enable a deeper understanding of the challenges of Transformer training from the optimization perspective. Second, $\sigma$Reparam, while effective, is not a panacea. In the practical sense, one might still benefit from combining $\sigma$Reparam with many other useful techniques, including initialization, feature normalization, advanced optimizers, etc. We hope that our work opens new perspectives towards inventing new design and training principles in the future.

Acknowledgement

We would like to thank Navdeep Jaitly, Vimal Thilak, Russ Webb for their helpful feedback and critical discussions on the experimental part of the work; Samy Bengio, Andy Keller, Russ Webb, Luca Zappella for their help throughout the process of writing this paper; Hassan Babaie, Mubarak Seyed Ibrahim, Li Li, Evan Samanas, Cindy Liu, Guillaume Seguin, Okan Akalin, and the wider Apple infrastructure team for assistance with developing scalable, fault tolerant code; and Shuming Ma for providing details on the DeepNorm reproduction steps. Names are listed in alphabetical order.

References

Appendix A Proof of 3.1 and 3.2

where $Z=\sum_{k=1}^{T}e^{u_{k}}$ is the partition function. The entropy given $p(u)$ is then:

We wish to solve the following minimization problem:

To find all saddle points, we solve the system of equations:

Giving rise to the following set of equations:

As a first step, assume that for the minimizer $u^{\star}$ of Equation 6 there exists an index $k^{\star}$ such that $u^{\star}_{k^{\star}}=0$. Using Equation 10:

From the first set of equations we arrive at the condition:

This however implies that $u_{1}=u_{2}=...=u_{T}=0$, hence a contradiction to Equation 11. Now, assuming $\forall_{k}~{}u_{k}\neq 0$, we have using Equation 10:

Let $u_{j_{1}},u_{j_{2}}<0$, then according to Equation 17:

Appendix B Relationship Between Entropy Collapse and Training Instability

Here we will investigate the interplay between entropy collapse and training stability by asking: would a model with stable training but not exhibiting entropy collapse have been stable if entropy collapse was induced, all other factors held constant? In do-calculus (Pearl, 2009), this roughly corresponds to checking

Holding ${\bm{u}}$ constant, the entropy is low when $\tau\rightarrow 0$, and is high when $\tau\rightarrow\infty$. As entropy collapse is observed in experiments when $H_{p}({\bm{u}},\tau)\rightarrow 0$, we will attempt to induce entropy collapse by sending $\tau\rightarrow\tau_{\text{target}}$, where $\tau_{\text{target}}\ll 1$.

Concretely, for a Transformer model, we normalize the logits of the attention matrix by temperature. We use the same temperature normalization for every layer, i.e. the Transformer has a global temperature. We start the temperature $\tau=1$ which corresponds to the default Transformer model without temperature normalization. At a prescribed epoch during training, we perform a temperature intervention, where we change the temperature from $\tau=1$ to a target temperature $\tau_{\text{target}}$. The transition is sharp, and happens at the start of the prescribed epoch, which we refer to as the intervention epoch.

We use the MAE ViT-B/16 recipe (see Appendix H) for these experiments, and train for a total of 100 epochs on ImageNet1k. To simplify the analysis, we only use ImageNet1k training augmentations, and use no learning rate decay schedule (i.e. the learning rate is flat after warmup).

Eigenvalues of the Hessian

As properties of the Hessian have been successfully used to gain an understanding of stability of the learning process Ghorbani et al. (2019); Yao et al. (2020); Cohen et al. (2021, 2022); Gilmer et al. (2021), we will also use them in our analysis. Specifically, we will analyze the magnitude $|\lambda_{i}|$ of the largest magnitude eigenvalues $\lambda_{i}$ of the Hessian $H$

where $\theta^{a}$ is the $a$–th parameter, $\mathcal{L}$ is the scalar loss, $P$ is the number of model parameters, and $v_{i}$ is the normalized eigenvector corresponding to the eigenvalue $\lambda_{i}$. We take $|\lambda_{1}|>|\lambda_{2}|>\ldots>|\lambda_{P}|$, and call the largest eigenvalue $|\lambda_{1}|$ the sharpness, in line with the stability literature.

Computing and storing the Hessian explicitly is problematic, as it is $O(P^{2})$ in time and memory. Instead, noting that the Hessian Vector Product (HVP) $Hv$ for any vector $v$ can be computed using the Vector Jacobian Product (VJP) or Jacobian Vector Product (JVP), avoiding explicit computation of $H$. Treating the HVP as a linear operator then allows the use of numerical methods for computing the spectrum Yao et al. (2020); Ghorbani et al. (2019). For our iterative method we use the implementation of Lanczos from CuPy Okuta et al. (2017). We compute the 5 largest eigenvalues of $H$ using 32,768 samples from the ImageNet1k training set, and perform this computation at the end of each training epoch.

The Stability Threshold

Different optimization algorithms have a stability threshold; under a local quadratic assumption, if any Hessian eigenvalue of the loss exceeds this threshold, iterations of the optimization procedure will diverge Cohen et al. (2021, 2022). For AdamW, the stability threshold $\Gamma$ is derived in the case of a short time-horizon frozen (i.e. non-fully adaptive) approximation of AdamW, has been shown empirically as a suitable stability threshold for the full algorithm Cohen et al. (2022), and is given by

where $\beta_{1}=0.9$ is the Adam momentum of the gradient moving average Kingma & Ba (2015). We include this threshold in our analysis.

B.2 Results

Appendix C Implementation of σ𝜎\sigmaReparam

To compute the spectral norm of the current matrix we use the power method as an approximation method to speed up computations. See Algorithm 1 for a sketch implementationBy default we use one step of power iteration per gradient update step, similar to (Miyato et al., 2018). Empirically we found no difference in performance when using multiple power iteration steps.. Note that in practice fp32 precision is typically required for numerical stability. We have experimented with various configurations applying $\sigma$Reparam to key and query weights, and/or in other parts (e.g., all other linear layers in the model). While we found that the performance is robust to the configurations, applying it to all the layers amounts to the simplest implementation and also works well in practice, e.g., allowing the removal of LN layers. $\sigma$Reparam does not bring any overhead compared to pre-LN or post-LN configurations, see Table 5.

Appendix D Self-Supervised Training of Visual Representations

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.

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.

D.2 Reduced Learning Rate Warmup

In Chen et al. (2021) 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 $\sigma$Reparam, either with or without pre-LN, we are able to achieve stable SimCLR+ViT training at the original warmup period of 10 epochs. 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 Figures 10 and 11).

Appendix E Automatic Speech Recognition (ASR)

In this section we focus on empirical investigation of Transformer training stability and attention entropy collapse phenomenon for automatic speech recognition (ASR) task.

Data All experiments are performed on the LibriSpeech dataset Panayotov et al. (2015) where audio paired with transcriptions is available. The standard LibriSpeech validation sets (dev-clean and dev-other) are used to tune all hyperparameters, 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 Model We stick to a vanilla Transformer model trained with Connectionist Temporal Classification (Graves et al., 2006) loss for simplicity of analysis where only encoder is used (no decoder). We use current, to the best of our knowledge, state-of-the-art vanilla Transformer model configuration and training recipe from Likhomanenko et al. (2021a, b): the model consists of (a) 1D convolution to perform striding (kernel of 7 with stride of 3), (b) Transformer encoder with 36 layers, post-LayerNorm (post-LN), 4 heads, embedding dimension of 768 and MLP dimension of 3072, and (c) a final linear layer to map to the output number of tokensThe token set consists of the 26 English alphabet letters augmented with the apostrophe and a word boundary token.. To speed up the model training (2-3x) and decrease memory usage we are using CAPE positional embedding (Likhomanenko et al., 2021c) instead of relative one (Shaw et al., 2018): both models perform in the same ballpark.

Training We follow a training recipe from Likhomanenko et al. (2021a, b). As they, we use SpecAugment (Park et al., 2019) which is activated right at the beginning of the training (no difference is found if it is used after 5k training steps): 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$ are used; time warping is not used. We also use Adagrad (Duchi et al., 2011) if not specified otherwise, and learning rate (LR) decaying by 2 each time the WER reaches a plateau on the validation set. 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 weight decay is used. Default warmup is set to 64k and can be varied if stated so. The default LR is 0.03 and is optimized across models. We also apply gradient clipping of 1.

E.2 Training Stability, Robustness and Generalization

We start with exploring training stability of the baseline model described above using both pre-LayerNorm (pre-LN) and post-LayerNorm (post-LN) configurations trained on small-scale data, namely 100h of LibriSpeech (train-clean-100). By varying different hyperparameters, such as learning rate, warmup, and gradient clipping, post-LN models fail to train. By inspecting the gradient norms per layer and per each parameters’ matrix we find a similar vanishing gradients problem as reported, e.g., by Liu et al. (2020b, a); Wang et al. (2022) for deep Transformers ($>12$ layers) in machine translation domain. At the same time, pre-LN is stable as reported by, e.g., Nguyen & Salazar (2019); Wang et al. (2022); Liu et al. (2020a): we are able to reduce warmup from 64k to 16k, increase learning rate from 0.03 to 0.5, and obtain better results than the training setting from the post-LN baseline. However, stable training of pre-LN leads to a degradation in performance compared to post-LN in ASR, similarly as reported in the aforementioned works: validation WER is worse while training loss is lower, see top of Table 7. By varying, e.g., learning rate and warmup hyperparameters and deeper inspecting training stability of pre-LN models we observe that attention entropy is not bounded and can collapse leading to the model divergence with training loss growing, see Figure 12.

As discussed above in Section 3, we now investigate how $\sigma$Reparam affects the training stability and controls the attention entropy bound. First, by removing all LayerNorms (pre-LN or post-LN) and switching to $\sigma$Reparam for all linear layers in Transformer blocks and in the final linear layer, we observe (a) stable training similar to pre-LN with no vanishing gradients issue; (b) accepting a wider range of hyperparameters (Figure 13) than pre-LN; (c) no attention entropy collapse phenomenon. While $\sigma$Reparam significantly outperforms a pre-LN model with the baseline hyperparameters used for post-LN, it performs worse than an optimized version of a pre-LN model as well as an unstable post-LN model (see top of Table 7). However, combining $\sigma$Reparam with post-LN brings two worlds together: stable training similar to pre-LN and generalization similar to post-LN. In summary, $\sigma$Reparam with post-LN achieves (a) similar performance on the validation and test sets and lower training loss (Table 7); (b) no vanishing gradients are observed as for post-LN; (c) the model accepts a wide range of hyperparameters (Figure 13) compared to unstable post-LN and stable pre-LN.

To demonstrate the necessity of $\sigma$Reparam in the form presented in Section 3, we compare it with spectral normalization (SN) where $\gamma$ is set to 1 and is not learnable, and WeightNorm (Salimans & Kingma, 2016) baselines. Both SN and WN perform poorly compared to $\sigma$Reparam (with or without post-LN), see Table 7.

We further investigate training behaviour if we increase the model depth by 2x resulting in 72 encoder layersThe total batch size is reduced by 2x to use the same amount of computational resources.. In such setting we are unable to train a post-LN model (vanishing gradients are observed) while pre-LN, $\sigma$Reparam and $\sigma$Reparam with post-LN are training out of the boxDeeper models perform worse compared to smaller ones, however we did not optimize deep models and this is out of scope of the current work. and have bounded attention entropy throughout the training with no vanishing gradients problem, see Figure 14.

E.3 Training with SGD

Vanishing gradients and unbalanced gradients can be one of the reasons why the standard SGD fails in training Transformers, especially for deeper architectures, and one needs adaptive optimizers. E.g., Li et al. (2022) report also another issue with SGD – ability for generalization, and propose Transformer components modification to improve generalization with SGD training.

To confirm prior findings, we first experiment with baseline models, pre-LN and post-LN, and SGD optimizer. While post-LN is not training, a pre-LN model can be trained but has a poor generalization. The same holds for $\sigma$Reparam and $\sigma$Reparam with post-LN: the gradient magnitude between the first and last layers can differ not drastically as in post-LN, but generalization is still poor. Similarly to vision experiments, we switch to the LARS (You et al., 2017) (with momentum 0.9) optimizer which normalizes gradients by their magnitudes and thus provides balanced gradients. By carefully tuning only the learning rate from 0.1 to 1.5 (the rest stays the same as for the adaptive optimizer except warmup which is set to 0k) we are able to train pre-LN and post-LN, see bottom of Table 7.

In our experiments post-LN is more unstable (many learning rates are diverging or not training) and gives significantly worse results than pre-LN. Furthermore, pre-LN is still behind the baseline that uses an adaptive optimizer. However, if we switch to $\sigma$Reparam (key, queries and values are represented as one matrix) we observe stable training with respect to learning rate changes, and combined together with post-LN it achieves similar performance to the best results from top of Table 7 while keeping the training loss lowFor the separate reparametrization for (keys, queries) and values, we observe less stable training with LARS and no warmup relative to reparametrizing them together.. To the best of our knowledge, this is the first ASR Transformer model trained without an adaptive optimizer achieving stable training and comparable performance. Regarding attention entropy collapse, we observe it with LARS training also, see Figure 15: $\sigma$Reparam controls the bound resulting in wider range of accepted hyperparameters for stable training (models can be trained with learning rate up to 1, while pre-LN and post-LN result in model divergence).

E.4 Hyperparameters

We present hyperparameters for our ASR experiments on 100h of LibriSpeech in Table 8.

E.5 Large-Scale Experiments: 1k Hours of LibriSpeech

We also evaluate $\sigma$Reparam for large-scale data: for further experiments we take all $\sim$1k hours of LibriSpeech as the training data. We consider again the 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 hyperparameters as for small-scale experiments from top of Table 8 except dropout and layer drop which are decreased to 0.1 to decrease model regularization effect. For all models we tune only the learning rate. As before, spectral reparametrization of keys and queries is done separately from values. We also use the learning rate on gamma to be twice bigger than the main learning rate. Similarly to small-scale experiments, training on LibriSpeech shows (see Table 9) that $\sigma$Reparam accompanied with post-LN can match the post-LN baseline, while having robustness to the hyperparameter changes (e.g. it allows larger learning rate values without any stability issues).

Appendix F Machine Translation (MT)

In this section we focus on empirical investigation of training stability and attention entropy collapse in deep Transformers for machine translation (MT) with an encoder-decoder architecture. We track attention entropy for the encoder self-attention, the decoder cross-attention and the encoder-decoder self-attention separately to study the entropy collapse phenomenon. The goal of this section is to understand how varying the model depth for the well-established recipes affects the training stability.

We build our experiments on top of the open-sourced codehttps://github.com/microsoft/torchscale and baseline recipes provided by Wang et al. (2022). We follow their instructionshttps://github.com/microsoft/torchscale/tree/main/examples/fairseq#example-machine-translation and hyperparameters given in Wang et al. (2022).

Data Following Wang et al. (2022) we perform all experiments on standard WMT’17 English-German benchmarkhttps://www.statmt.org/wmt17/translation-task.html: we use all provided training data for English-German pair, newstest2016 set as a validation set and newstest2017 as a test set for final evaluation purpose only. We use Fairseq (Ott et al., 2019) script to preprocess data: it uses Byte Pair Encoding (BPE) vocabulary jointly for source and target language resulting in 41k subword tokens.

Models We consider both regular and deep configurations for a vanilla encoder-decoder Transformer model with $N$ encoder and $N$ decoder layers where $N$ is taken as 6 (6L-6L), 18 (18L-18L), 50 (50L-50L), and 100 (100L-100L). Every Transformer layer in each configuration has an embedding dimension of 512, MLP dim of 2048, and 8 heads. Sinusoidal absolute positional embedding (Vaswani et al., 2017) is used for both encoder and decoder.

Training We strictly follow the same training recipe from Wang et al. (2022) (without using back-translation or other domain-specific augmentations) with detailed hyperparameters in Table 10. All models are trained on 8 GPUs of A100 80GB with mixed precision computations and dynamic batching resulting in total batch size of 524288 tokens: for each architecture we pack maximum tokens per GPU and use gradient accumulation (4 for 6L-6L and 18L-18L, 8 for 50L-50L and 16 for 100L-100L).

Evaluation As it is not specified in Wang et al. (2022) how the best checkpoint is selected on the validation set, we decided to stick to simple rule: checkpoint with best perplexity on the validation set is selected and further evaluated on both validation and test sets for BLEU score computation which is reported throughout the paper. BLEU is computed by in-built BLEU scripts of Fairseq with the beam size of 5. As reported in prior works we also observe a strong correlation between perplexity and BLEU score: improved perplexity leads to better BLEU score. However BLEU scores on validation and test sets are less correlated and high variation is observed. For that reason we often perform 3 runs with different seeds to estimate standard deviation (std) of the BLEU score.

F.2 Training Stability of Deep Models

We start with exploring training stability of the baseline model described in Wang et al. (2022) with pre-LayerNorm (pre-LN) and post-LayerNorm (post-LN) across different depths (all hyperparameters stay the same except depth is varied). Note that post-LN is a popular design choice for MT tasks due to its good generalization properties.

For pre-LN models, we reproduced stable results and convergence, however the BLEU score we get is better (Table 11) than reported by Wang et al. (2022). We also observed the same trend of decreasing model performance with increasing the model depth. Attention entropy is nicely bounded across all depths similarly to ASRNote, we did not do any hyperparameters search to investigate how models behave with, e.g., wider range of learning rates as we did for ASR models., see Figure 16.

For post-LN models, we reproduced stable results for 6L-6L depth and observe nicely bounded attention entropy behaviour. However for 18L-18L configurations, divergence is observed when varying the random seed. By close inspection we observe no vanishing gradients problem while attention entropy collapse clearly occurs during training (compare top and middle in Figure 17) in the encoder attention and the encoder-decoder cross-attention. Deeper models, namely 50L-50L and 100L-100L, are unable to train and we observe the same vanishing gradients problem as reported by Wang et al. (2022); Liu et al. (2020a) as well as attention entropy collapse for some of the deep layers across the board, see bottom plot in Figure 17.

Wang et al. (2022); Liu et al. (2020a) are recent works that proposed to rescale residual connections. To stabilize training and resolve vanishing gradients problem in deep post-LN models to preserve post-LN generalization properties. We focus in this paper on Wang et al. (2022), DeepNorm, solution (it uses post-LN and rescale residual connections depending on the initial model depth) as they reported ability to train up to 1000-depth Transformer models. We are able to reproduce DeepNorm results for 6L-6L, 18L-18L and 50L-50L depths observing stable training (no any models diverged and training went nicely). However we see no performance gain of a 50L-50L depth model over a 18L-18L model. Furthermore, we observe instability in training of the 100L-100L model resulting in only 1 successful run among 3 (only seed is varied) while 2 others are diverging after some time (training loss is growing). By close inspection of the training behaviour we do not see any drastic issue of vanishing gradients, however we see the attention entropy collapse happening, see Figure 18. First of all, attention entropy is not bounded for DeepNorm even in 18L-18L and 50L-50L similarly to what we observed in post-LN models. Also a tiny attention entropy collapse happens in 50L-50L (see top plot in Figure 18) though it does not lead to any divergence. Second, attention entropy collapse is clearly pronounced for 100L-100L models (second, third, and forth rows of Figure 18) leading to 2/3 seeds divergence and one with worse performance than 50L-50L modelsFrom our empirical observations in other domains it could be that deeper models are worse as any attention entropy collapse degrades optimization process resulting in worse generalization.. Finally, it is interesting to note that attention entropy collapse in 100L-100L can happen for different layers, first and / or last, and with different regimes for the encoder/decoder self-attention and the encoder-decoder cross-attention.

All models performance on validation and test sets across depths as well as the number of successful runs are reported in Table 11.

F.3 σ𝜎\sigmaReparam for Deep Models

We now experiment with injection of $\sigma$Reparam into post-LN and DeepNorm models to alleviate attention entropy collapse and confirm $\sigma$Reparam effectiveness for deep models. $\sigma$Reparam is used for every linear layer in the encoder and decoder Transformer blocks alongside with post-LN. With DeepNorm we also apply its rescaling of initialization and residual connections.

$\sigma$Reparam nicely bounds attention entropy for 18L-18L and 50L-50L post-LN models, resolving any divergence issues as well as vanishing gradient in the 50L-50L model, see Figure 19. However, 100L-100L is still experiencing a vanishing gradient problem and only careful initialization of std for $\sigma$Reparam can resolve it: for that reason we report that model training is not stable. In terms of performance, $\sigma$Reparam with post-LN matches post-LN for 6L-6L, in the same ballpark for 18L-18L and performs the same as 18L-18L for 50L-50L. Note, that we did not do any hyperparameters search except tuning learning rate as $\sigma$Reparam has different learning rate scales.

$\sigma$Reparam also nicely bounds attention entropy for 18L-18L, 50L-50L, 100L-100L DeepNorm models, resolving any divergence issues for 100L-100L (vanishing gradient is not observed as DeepNorm targets it), see Figure 19. In terms of performance $\sigma$Reparam with DeepNorm matches DeepNorm for 6L-6L, in the same ballpark as DeepNorm for 18L-18L and inferior to DeepNorm for 50L-50L and 100L-100L.

Appendix G Language Modeling (LM)

As we discussed above encoder Transformer for vision and speech domains and encoder-decoder for machine translation, in this section we focus on the pure decoder architecture in language model task to verify if $\sigma$Reparam is effective for stable training and can simplify a training recipe there too.

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, 4096 fully connected dimensions and post-LayerNorm (post-LN). 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). We integrate $\sigma$Reparam implementation into the open-sourced code and recipe for the baselinehttps://github.com/facebookresearch/fairseq/blob/main/examples/language_model/README.adaptive_inputs.md. All models are trained in full precision on 8 GPUs of A100 40GB.

G.2 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. This is consistent and in line with observations in machine translation where entropy collapse is observed in the encoder and cross-attention. Nonetheless, we experimented with $\sigma$Reparam to test its generality on a different modality/problem. We apply $\sigma$Reparam to all linear layers of the Transformer while removing all post-LNs, 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, including Nesterov SGD optimizerNote, this is different from other domains where a standard recipe includes only adaptive optimizers.. The results are shown in Table 12. We see that even in the absence of LayerNorm, $\sigma$Reparam shows strong performance in convergence and validation/test performance. With a mild weight decay, $\sigma$Reparam also outperforms the baseline wrt the validation/test PPL. In summary, while there is no observed entropy collapse in language model training, $\sigma$Reparam can simplify a training recipe by removing all post-LNs.

Appendix H Hyperparameters for Supervised Vision

As mentioned in Section 4.1 we compare $\sigma$Reparam against DeiT (Touvron et al., 2021) and MAE (He et al., 2022) supervised training recipes for vision Transformers. In Table 13 we highlight the differences between DeiT, MAE supervised and $\sigma$Reparam. $\sigma$Reparam presents a simplified and stable training objective for ViT-B variants. In Table 14 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.

Appendix I Ablations

First, we found that it is better to initialize $\gamma$ 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. We compared different initializations for the kernel and we did not see any differences in initialization (e.g. uniform, normal). The only factor that influences training behavior is the standard deviation (std) of the initialization pdf, which also influences effective learning rate. In speech recognition we found that training is robust with respect to changes of std (Figure 13), however larger std performs better and sweet spot is 0.2-0.3. In machine translation models are also robust to the choice of std, however some architectures perform better with std of 0.01 while others with 0.1 std. In language modeling we observed robust performance with respect to initialization, and we use the default initialization from the Transformer baseline for all experiments.

In vision we initialize the $\sigma$Reparam $\gamma$ term using the first singular value, computed with the SVD at weight initialization. We then use one power iteration for all further updates. We provide weight and patcher initializations for the ViT-B/16 in Table 13 and the ViT-L/16 and ViT-H/14 in Table 14.

Separate σ𝜎\sigmaReparam for key, queries and values

We found that they behave 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 ASR training with LARS it is better to have joint reparametrization to achieve stable training and comparable results with adaptive optimizers, see Section E.3.

Appendix J Discussion

We believe that our experiments have covered representative domains, architectures and training losses for typical Transformer applications. The key factors that affect training stability are the initial token embedding layer (where for vision and speech tasks they are based on convolution projections, and for MT and language modeling are with word embeddings), topology of architecture (encoder mode for vision and speech, encoder-decoder for MT and decoder for language modeling), and the training loss (cross-entropy loss, contrastive and CTC loss). While each of these design choices may bring their own challenges for training, we show that entropy collapse is a common symptom accompanying instability and $\sigma$Reparam is a general fix compatible with all settings.

Appendix K Contributions

All authors contributed into writing the manuscript, designing experiments and discussion of all results at every stage of the project.

All initial experiments and findings of the attention entropy collapse phenomenon for ViT models on ImageNet are done by Shuangfei Zhai. Preliminary theoretical analysis and proposal to use $\sigma$Reparam as a solution is also done by Shuangfei Zhai.

Theory

All theoretical results, Appendix A, are done by Etai Littwin. Review of proofs is done by Tatiana Likhomanenko.

Causality Analysis

Simulated case-control study investigation and all related experimental work done by Dan Busbridge. Tatiana Likhomanenko, Etai Littwin, Jason Ramapuram, Russ Webb and Shuangfei Zhai helped with designing the experimental setting for intervention methodology.

Supervised Learning in Vision

Shuangfei Zhai conducted the initial $\sigma$Reparam vision experiments with DeiT and made the initial observations of relaxing / removing weight decay from $\sigma$Reparam. Jason Ramapuram scaled and conducted the remaining supervised vision experiments and analysis (including the MAE, weight-norm and spectral-norm baselines) over ImageNet1k (Table 1) and Imagenet21k (Table 2) and enabled the drastically simplified $\sigma$Reparam vision transformer recipe from Table 1. This simplified $\sigma$Reparam recipe enables SGD (LARS) training over fewer epochs and completely removes {weight-decay, pre-LN layers, LR warmup and the LR cosine schedule}.

Self-Supervised Learning in Vision

Known issues with SimCLR (Chen et al., 2020) stability that were observed in Chen et al. (2021) pointed out by Vimal Thilak. All investigations, experiments and related analysis done by Dan Busbridge.

Automatic Speech Recognition

All speech recognition experiments are done by Tatiana Likhomanenko. Shuangfei Zhai and Jason Ramapuram advised to have also large scale results.

Machine Translation

Initial implementation and experiments of $\sigma$Reparam applicability to MT were done on WMT’14 by Jiatao Gu. Later, Tatiana Likhomanenko pushed to investigate deep transformer models and their stability. Jason Ramapuram pointed to the deepnorm (Wang et al., 2022) results to probe for entropy collapse phenomenon. All later experiments, Section 4.3 and Appendix F, with deep transformers and deepnorm are done by Tatiana Likhomanenko.

Language Modeling

Initial implementation and preliminary results on applicability of $\sigma$Reparam to the language modeling, Appendix G, are done by Yizhe Zhang with help from Jiatao Gu. Shuangfei Zhai contributed to the experiments and obtained the final results.

Implementation Details and Ablations

Investigation into how initialization influences $\sigma$Reparam is done in parallel in different domains and experiments by Jason Ramapuram, Shuangfei Zhai and Tatiana Likhomanenko. Investigation of different variants (with stop gradient, with different matrices) of $\sigma$Reparam was done by Shuangfei Zhai, Tatiana Likhomanenko and Jason Ramapuram. Investigation of full precision training vs mixed precision training was done by Tatiana Likhomanenko, Dan Busbridge and Jason Ramapuram.

Implementation is done in 2 frameworks (PyTorch and Jax) and in 5 codebases. The initial implementation of $\sigma$Reparam module is done in PyTorch by Shuangfei Zhai, with further reimplementation in Jax by Tatiana Likhomanenko. Later the implementation was integrated and adopted into other baseline toolboxes by Jason Ramapuram, Dan Busbridge, Yizhe Zhang, Tatiana Likhomanenko and Jiatao Gu.