NormSoftmax: Normalize the Input of Softmax to Accelerate and Stabilize Training
Authors: Zixuan Jiang
Jiaqi Gu
David Z. Pan
Venue: Submitted to ICLR 2023 (2023)

Abstract: Softmax is a basic function that normalizes a vector to a probability distribution and is widely used in machine learning, most notably in cross-entropy loss function and dot product attention operations. However, optimization of softmax-based models is sensitive to the input statistics change. We observe that the input of softmax changes significantly during the initial training stage, causing slow and unstable convergence when training the model from scratch. To remedy the optimization difficulty of softmax, we propose a simple yet effective substitution, named NormSoftmax, where the input vector is first normalized to unit variance and then fed to the standard softmax function. Similar to other existing normalization layers in machine learning models, NormSoftmax can stabilize and accelerate the training process, and also increase the robustness of the training procedure against hyperparameters. Experiments on Transformer-based models and convolutional neural networks validate that our proposed NormSoftmax is an effective plug-and-play module to stabilize and speed up the optimization of neural networks with cross-entropy loss or dot-product attention operations.

Under review as a conference paper at ICLR 2023
NORMSOFTMAX: NORMALIZE THE INPUT OF SOFT- MAX TO ACCELERATE AND STABILIZE TRAINING
Anonymous authors Paper under double-blind review
ABSTRACT
Softmax is a basic function that normalizes a vector to a probability distribu- tion and is widely used in machine learning, most notably in cross-entropy loss function and dot product attention operations. However, optimization of softmax- based models sensitive the input statístics change. We observe that the input of softmax changes significantly during the initial training stage, causing slow and unstable convergence when training the model from scratch. To remedy the op- timization difficulty of softmax, we propose simple yet effective substitution, named NormSoftmax, where the input vector is first normalized to unit variance and then fed to the standard softmax function. Similar to other existing normaliza- tion layers in machine learning models, NormSoftmax can stabilize and acceler- ate the training process, and also increase the robustmess of the training procedure against hyperparameters. Experiments on Transformer- based models and convo- lutional neural networks validate that proposed NormSoftmax is an effective plug-and-play module to stabilize and speed up the optimization of neural net- works with cross-entropy loss dot-product attention operations.
INTRODUCTION
Softmax is a critical and widely used function machine learning algorithms, which takes a vector input and generate standard simplex. It is usually used to generate categorical probability distribution. The most notable applications of softmax are cross-entropy loss function for classi- fication tasks and attention map generation in dot product attention operations. By importing the temperature in softmax, we can control the information entropy and sharpness of its output.
However, gradient-based optimization of softmax-based models often suffers from slow and unstable convergence and is sensitive to optimization hyperparameters. Transformer-basee models (Vaswani al., 2017) are known to be hard optimize. A lot of efforts have been devoted to solving this optimization difficulty (Liu et al., 2020). For instance, Bolya et al. (2022) reports that softmax attention may crash with too many heads and proposes new attention functions. Chen et al. (2021) show that the Vision Transformer's (Dosovitskiy et al., 2021) loss landscape is very sharp, and requires advanced optimizers to facilitate its training (Foret et al., 2020). Huang et al. (2020) propose better initialization to improve the Transformer optimization. Xiong et al. (2020) show that the location of layer normalization (LN) has a remarkable impact on the gradients and claim that the Pre-LN Transformer has better training stability.
Among comprehensive reasons for the optimization difficulty of Transformers, cascaded softmax functions are one of them that leads to the training instability. However, limited prior work has discussed the impacts of softmax on optimization. Based on our experimentation, we find that the training difficulty can be attributed to the rapid change in the variance of the softmax inputs and the information entropy of its outputs. In dot-product attention, where the softmax is used to generate weight distribution for key-value pairs, we observe significant statistical fluctuation in softmax inputs. The rapid and extensive variance change in the initial learning stage can lead to unstable training. Moreover, for the softmax used in cross-entropy loss for classification problems, the input of the softmax usually has a lower variance at the initial training stage since the model has less knowledge of the problem (Wei et al., 2022). The model likely to stay in the low- confidence zone, implying that it is difficult to train (Pearce et al., 2021). We need a specially designed mechanism to push the model out of this low-confidence zone for stable and fast learning.

---

Under review as a conference paper at ICLR 2023
In the two cases above, the significant change of the softmax input variance one of the reasons for optimization difficulty. In this paper, propose NormSoftmax to stabilize and accel- erate training by simply re-scaling the softmax inputs, espe- cially in the early stage optimization.
Attention/
Attention/
Linear
Linear
Replace
NormSoftmax
Softmax
Softmax(Scale(x)
With NormSoftmax, we dynamically calculate vector-specine factors to scale the inputs before being fed the standard soft- max. Specifically, when the input variance is too small, Soft- max will generate small gradients that hinder the learning pro-
Cls.
Cls.
Unstable
Stable
Converge slowly
Converge fast
cess. In contrast, our proposed NormSoftmax can help re-scale the input distribution such that the information entropy of the output becomes stable without fluctuation during the training process, which boosts and stabilizes the early stage training.
Sensitive
Robust
Figure I: The standard softmax
and proposed NormSoftmax.
NormSoftmax shares similar properties with the existing normalization techniques in machine learn- ing. We summarize its advantages below.
NormSoftmax can re-scale gradients to stabilize the training process, making the training robust model architectures and optimization recipes (such as optimizers and weight decay schedules).
NormSoftmax can accelerate the early training stage without hurting the model representability.
NormSoftmax is an easy-to-use and low-cost module to replace standard softmax. The induced computation and memory cost overhead negligible.
NormSoftmax has a regularization effect since the re-scaling can slightly restrict the representa- tion space of the input vectors.
In this paper, we focus on two applications of the softmax functions: (1) the activation function in dot-product attention, and (2) cross-entropy loss of the classification problem. ViT-B with our NormSoftmax shows significantly higher robustness to different head settings, showing an average of +4.63% % higher test accuracy on CIFAR-10 10 than its fettmax-basedd counterpart. When training for 100 epochs on ImageNet-IK, ViT with our NormSoftmax can achieve +0.91 higher test accuracy over its softmax baseline.
2 BACKGROUND
We briefly introduce the softmax function and normalization in machine learning. Then discuss the two cases we focus on this paper: softmax dot product attention and cross entropy loss. Throughout this paper, use p(a), (a) to represent the mean and standard deviation (square root of the variance) of a vector a.
2.1 SOFTMAX
The standard softmax function softmax(x)) where az IR. is defined by Equation
for =1, 1,.....,
The output of softmax can be seen as categorical probability distribution since < 2, 1 and Σx = 1. Instead of e, we can also use different base in softmax. A temperature parameter T' > imported to adjust the base.
softmax7 = softmax
Given the same input vector æ. the higher temperature smooths the difference of the input vector and generates a probability distribution with high information entropy H(x) Σ, 24 log(2;). On the contrary, the lower temperature sharpens the output distribution with low entropy. (Agarwala et al., 2020) claim that the temperature has crucial impact the initial learning process.
2

---

Under review as a conference paper at ICLR 2023
2.2 NORMALIZATION
Normalization layers (loffe & Szegedy, 2015; Ba et al., 2016; Ulyanov et al., 2016; Wu & He, 2018) usually normalize the input tensor such that the result has zero-mean and unit-variance along specific dimensions, We can optionally scale and shift the normalized tensor further. Normalization can accelerate and stabilize the optimization by smoothing the loss landscape (Santurkar al., 2018; Xu et al., 2019; Bjorek et al., 2018). Hence, normalization allows for larger learning rate and increases the robustness against hyperparameters. Also, normalization helps generalization since the sharpness of the loss surface is decreased effectively (Lyu et al., 2022). However, we find that normalization is usually used in the intermediate linear transformation layers, and it is rarely applied to the input of softmax functions.
2.3 DOT PRODUCT ATTENTION
The scaled dot product attention (Vaswani et al., 2017) is defined by the following equation, where Q, K, V are query, key, and value matrices and d the dimension of the key vector.
QK?
Attention( Q, K, V) = softmax
(3)
vd
For every query vector, the softmax function calculates the weight for all key-value pairs. The scaling factor of d'1/2 is proposed attempt the normalize the dot product q" k, whose variance is d if the components of q, k are independent random variables with a variance of O. The scaling factor is applied to address the issue that the variance of dot product q" k will likely increase as the length of the vector d increases. The large variance makes the gradient of softmax extremely small, thus making the attention-basee model hard to train. Based on self-attention, Transformer has achieved great success in many areas, especially natural language processing (Devlin et al., 2019) and computer vision (Khan al., 2021).
2.4 CROSS-ENTROPY LOSS
Another important application of softmax is in classification problems, where minimizing the cross- entropy loss is equivalent to maximizing the likelihood. The cross-entropy function takes the esti- mated probability distribution q sofumax(x)) and the probability distribution p as input and compute the result by H (p, q) = E. log qt. æ is the predicted logits, usually generated by classification model. æ can be any vector in R without restrictions, where K is the number of classes.
METHOD
In this section, we first analyze the behavior of the softmax input during the training process. Further, we define NormSoftmax and discuss its advantage in the two cases.
3.1 UNDERSTANDING THE BEHAVIOR OF SOFTMAX IN DIFFERENT TRAINING STAGES
We split the whole training process into three stages.
Initial stage. Starting from scratch, we usually need careful design for initialization and hy- perparameters. The training may fail due to exploding or vanishing gradients. For example, we usually apply learning rate warmup during this stage.
Intermediate stage. Once the parameters are sufficiently warmed up, it is relatively easy and stable to explore the solution space with higher learning rate.
Final stage. The model attempts to converge to a local minimum with a lower learning rate.
The initial stage the most unstable one among these three stages. If we do not encounter severe issues in the initial stage, is likely that proceed to the final results.
Please refer to Footnote of the original paper (Vaswani et al., 2017).
3

---

Under review as a conference paper at ICLR 2023
Distraction
Concentration
Distraction
Learning slowly in the early stage
20
Rapid statistic change
L2
10
std of sfftinnn imput
L3
std o0'solmm floox input
10
1.4
L55
Rapid variarice change
L6.
0.5
4
1.7
2
wrongly
L8
0.1
Stuck in low-confidence region
estimates the std.
sqr(()
0.0
0
20
40
epoch
60
80
0
40
cpoch 60
80
(a) Softmax in Attention of an 8-layer VIT
(b) Softmax for the cross entropy loss in ResNetl 202
The standard deviation of softmax input has significant change during the training process. Note that the vertical axes are in the logarithmic scale.
Figure 2:
Figure 2 illustrates two cases where we observe a significant change in the softmax input variance throughout the training process. We discuss the behavior of the softmax function in these two cases.
Softmax in dot product attention. The input variance of softmax is related to the attended region of this attention layer. With a high variance on the imput, the output of softmax has low information entropy, meaning that only few keys are attended. On the contrary, the smaller input variance implies that much more keys are attended. Figure 2a demonstrates that the standard deviation of softmax input (before scaled with Vd) increases rapidly at the initial stage and then decreases gradually during the intermediate and final stages. In the beginning, the vision Transformer (Dosovitskiy al., 2021) attempts attend almost all the keys since the initialized model has not learned how to extract features and is still in the exploration stage. During training, different keys are learned different paces. Keys with smaller semantic distances from nearby queries are much easier to learn. This imbalanced learning pace among keys drives the model to shrink its receptive field and focus on small region. That is why the variance of softmax input increases rapidly at the initial stage. Afterwards, as training proceeds, Transformer will explore the query-key pairs with longer semantic distance, implied by a gradual reduction in the variance of softmax input. We name the effect
Further, different layers have different input variance. Except for the first layer (LI), the lower layer has high variance than the deeper layers. In the early layers, only small region is attended to. As depth grows, the model attends to larger region, which is similar to the trend of receptive fields in convolutional neural networks. We also draw the line of vd as a reference in Figure 2a, which the scaling factor used normalize the softmax input. Since the input variance has a significant change across different layers and training steps, this constant scaling acaling.acttor of vil might not be the most suitable value to normalize the softmax input (Lee et al., 2021).
Softmax in cross-entropy loss function. For classification model training, the softmax may only appear in the cross-entropy loss function, e.g., ResNet. However, the gradients through softmax can have a large impact on the model training. To verify this claim, we train a deep R (He et al., 2016) on CIFAR10 10 dataset (Krizhevsky et al., 2009) and plot the standard deviation of the softmax input in Figure 2b. At the initial stage, the standard deviation is relatively low since the model is less confident, and the predicted distribution is similar to a uniform distribution. There a leap at the 15-th epoch, where the standard deviation increases from 0.14 to 1.31. Afterward, the standard deviation of the softmax input gradually increases, and the information entropy of the softmax output becomes smaller since the model becomes more confident in predictions as training continues. Hence, we conclude that the variance of softmax input experiences rapid and huge change during training, especially in the initial stage, which explains why training from scratch is difficult.
Moreover, the observed trend in Figure 2 is averaged all the data points. Actually, different data points or training examples are in different learning stages. For instance, in the image classification problem, clear image can easily escape the less-confident zone quickly and stably, while vague
4

---

Under review as a conference paper at ICLR 2023
image may be stuck in this zone where the softmax output has a high information entropy. Therefore, this training difficulty needs to be tackled in a data-specific fashion.
3.2 THE PROPOSED NORMSOFTMAX
We propose NormSoftmax as substitution for softmax, as defined below,
x (x
NormSoftmax(az (x, 7) softmax
(4)
min(a(x=),7)
x
softmax
(5)
min(o(2)),))
2/7
= =softmax
(6)
min(a(21/1),1)
where y > is a pre-defined scalar. Namely, we define the temperature T mi) in NormSoftmax. Since softmax is invariant under translation by the same value, Equations 4 and are equivalent. We do not need shift the input vector to zero-mean. If the standard derivation is smaller than the threshold or (x) S NormSoftmax will normalize the input vector to obtain unit- variance vector before applying the standard softmax function. If (xr > 7 we use the temperature T= y. y +oo, then we will always normalize the input vector. The temperature is dynamically
calculated per vector. For a batch of vectors, each one has its individual temperature.
Lemma
Given y o(w) z softmax(y) y): Oz
R.n
we have
A Dy
0l
( 01 0y
al
8y
O1
0,
S
dx
Ox:
o
(
aa
(xx)
Lemma Given 32
kw
21
softmax /o(1)((()))))))))))))))))))))))) 22 (/tmax((22//(00
R", we have
22. Us
01
a1
01
1
01
Similar to the theorem in Xu et al. (2019), we refer to gradient re-certering and re-scaling as gradient normalization. Similar to Theorem 4.1 in Santurkar et al. (2018), we demonstrate that the normal- ization improves the Lipschitz continuity, indicated by the gradient magnitudes. Specifically, || dæ
at
and
at By can be treated as the continuity of the loss function Given input vectors 31, 222 and
(xx) < (x22) <, NormSoftmax apply different temperature on them, paying much attention to the low-variance vector æ1. In Lemma 2, if ( = 1, then = o(x2), it is clear that the gradients are rescaled by its variance
at
81
[
Logit Normalization (LogitNorm) (Komblith et al., 2021; Wei et al., 2022) uses the & norm to normalize the input vector, as shown in the equation below, where 7 is the temperature parameter modulating the magnitude of the logits.
LogitNorm(z, r) = softmax
-||20||
LogitNorm is proposed to mitigate overconfidence when cross entropy loss is used. It is questionable that LogitNorm does not shift the input vector zero-mean, since the mean has an impact on the l2 norm. Unlike the standard softmax function, LogitNorm is not invariant under translation LogitNorm(az,7) * LogitNorm(1: + cl,7), where is constant.
We demonstrate the relationship between LogitNorm and NormSoftmax in the following equation.
(Vr p(=)1 ((x)
NormSofmaax2,,7 7 +00) softmax
LogitNorm(z p(=)1,1111))) -1/2)
NormSoftmax first shifts the input vector zero-mean and normalizes the shifted vector by its & norm. Hence, NormSoftmax keeps the invariance under translation and can be reduced to Logit- Norm with input shifting.
5

---

Under review as a conference paper at ICLR 2023
3.3 EFFECTS OF NORMSOFTMAX IN THREE STAGES
NormSoftmax can accelerate and stabilize training in the initial stage. With NormSoftmax, the
standard deviation of the softmax input is at least since a min((r(a))))) 2 1. If the softmax input has low variance, which is common in the initial stage for both two applications, we use self- adapted low temperature to magnify the slight difference. The normalization can help the model escape the zone with high information entropy quickly and stably since normalize the gradients as demonstrated in Theorem 1. Specifically, for attention layers, NormSoftmax can accelerate the transition from distraction to concentration.
NormSoftmax can regularize training in the intermediate and final stages. For cross entropy loss. NormSoftmax increases the temperature of the softmax in the training process since the vari- ance of the softmax input will increase gradually. The high temperature can regularize the training. For attention layers, NormSoftmax regularizes the softmax input, thus restricting the attentive areas. NormSoftmax encourages the attention to have slight change on the attended regions, which can be treated as an inductive bias we add.
However, we find that by simply normalizing inputs to unit-variance vectors, i.e., ///(z)), the train- ing process can be impeded due to overly restricted representation space. Our variance clipping technique can effectively solve this issue with pre-defined threshold 7
EXPERIMENTS
Detailed experiment settings can be found in Appendix. Since we primarily use the cosine annealing learning rate scheduler, we train from scratch when the total number of epochs is different. Namely, the learning rate schedule is updated according to the number of epochs.
4.1 DOT-PRODUCT ATTENTION
Settings. We train the vision transformer (ViT) on the CIFARIO dataset from scratch with Adam W (Loshchilov & Hutter, 2019) optimizer for 100 epochs (50,000 iterations with 100 mini-batch size). The resolution of an input image is 3 x 32 x 32, and the patch size is 4. The hidden size, MLP size, number of heads, and the dimension of heads are 256, 1024, 8, and 32, respectively. We discard the original scaling factor Vd and replace the standard softmax with NormSoftmax, setting y = vd or y = +00. The learning rate linearly increased with warmup epochs and then decays with the cosine annealing scheduler (Loshchilov & Hutter, 2017). Following the original ViT paper (Dosovitskiy et al., 202 1), the learning rate is 1e- -3, and a strong weight decay 1e- 1 is applied. We also enable label smoothing (Szegedy et al., 2016) and strong data augmentation (random erasing (Zhong et al., 2020), mixup (Zhang et al., 2018), cutmix (Yun et al., 2019), and TrimialAugment (Müller & Hutter, 2021)). Figure 3 shows the detailed results.
Acceleration. Figure 3a demonstrates the result with different training epochs. When the number of epochs is small, NormSoftmax can achieve better test accuracy than the standard softmax function. When we train with more iterations, NormSoftmax performs similarly to the baseline. The larger pre-defined threshold y may translate into higher acceleration, which pushes the model escape the initial distraction stage more quickly. That is why NSM- inf performs better than NSM- sqrto when the number of epochs is small. However, the normalization has strict regularization effect on the softmax input, which impedes the training in the intermediate and final stages. Hence, the NSM-inf is exceeded by WSM-sqrta when the number of training iterations large. In short, NormSoftmax can accelerate the training process without sacrificing the representation ability, sim- ilar to curriculum learning (Wu et al., 2020).
Stabilization. We investigate the role of each component in the training recipe, with results listed in Figure 3b. We conduct ablation studies by removing (I) learning rate warmup, (2) weight de- cay, (3) label smoothing, and (4) strong data augmentation (random erasing, mixup, cutmix, and TrivialAugment) separately. We also replace the default Adam optimizer with stochastic gradi- ent descent with momentum (SGDM). The results indicate that the techniques above are critical to Transformer no matter what softmax function we use. However, Transformer with standard softmax and scaling factor d11/2 is more sensitive to the training techniques that are related to optimization.

---

Under review as a conference paper at ICLR 2023
95
90
90
sm
nsm--s
nsm-inf
sm nsm--qqrtd
nsm-inf
test accuracy (%)
88
test accuracy
85
86
hostine
we varmop
weight decay
we smouting
we sugmormation
SObM
10
50
epoch
(a) Results with different training epochs
(b) Results with different training recipes
90
90
nsm--sqrid
nsm-inn
88
85
accuracy (%)
86
Test accuraey
75
nsm-sqrid nsm-inf
70
80
65
78
128.
128.
128.
256.
256
256.
512.
512,
$12,
4.32
16
10,
4. 64
8,32
128
64
2
10 # layers
20
hidden size, Ai heads, head dimension
(c) Results on ViT with different depths
(d) ViT with different dimensions and heads
Figure 3: The test accuracy of VIT on CIFARIO dataset with different settings. sm and nsm are short for softmax and NormSoftmax. sqrtd and inf f are the value of y in NormSoftmax.
Without weight decay, the test accuracy of the baseline degrades from 87 87.96% to 84.15%, while the DSM-inf has a small drop from 88.01% % to 87.46%. They share the same robustness against the training recipe for data augmentation.
We also alter the hyperparameters of the ViT and list the results in Figures 3c and 3d. The three methods share similar results when the depth or the hidden dimension is small. However, large depth and hidden dimension impose a challenge for training. 1fff is more robust and provides much better results than the standard softmax.
4.5
Ll
1.2
L33
144
1.5
einapy soltmas qulput
L6
L7
L8
3.5
3.5
2.5
2.5
2.5
1.5
0
20
40
60
80
20
60
0
20
40
80
epoch
epoch
epoch
(a) Baseline with standard softmax
(b) NormSoftmax with 7 = vd
(c) NormSoftmax with = +o0
Figure 4: NormSoftmax reduces the information entropy of softmax output in the 8-layer ViT.
We plot the information entropy of softmax output this VIT in Figure 4. In baseline, the large standard deviation of softmax input induces small information entropy in the softmax output. The effect is visualized Figure 4a. The attended areas of Trans- former undergo significant change during the training process. On the other extreme, NSM-inf

---

Under review as a conference paper at ICLR 2023
generates the output, whose information entropy only slightly fluctuates since the input variance is always 1. The small change in the entropy of the output contributes to the acceleration and stability of the initial training process. The NSM-sqrtd d is interpolation between these two extremes.
Variants of NormSoftmax. We define several variants of NormSoftmax. (1) Adding Iearnable parameters for affine transformation defined by fi fi(x) softmax (wx)) where E R is learnable parameter. (2) Inverting the NormSoftmax defined by f2(2,, 7) softmax mmx (ax))))
(3) NormSoftmax with different values, (4) Logit Normalization.
baseline
fi f2.7 vd
nsm-1
nsm- nsm-Vd/2
nsm Vd
nsm 2vd
nsm-co
logitnorm
87.96
87.97
86.99
81.04
87.16
88.1 13
88.21
88.01
87.72
Table 1: Train a ViT on CIFARIO with different NormSoftmax variants.
For the first variant fi. we find that the learnable parameter is not necessary. Indicated by (Xu et al., 2019), the learnable weight may induce the risk of overfitting. The second variant fz has an inverse clipping and worse than our proposed NormSoftmax. cannot accelerate the training since it encourages distraction in the Vision Transformer. In the third variant, the is critical hyperparameter. ? = 1, vd/2 2 is too small, and the input vector is not effectively scaled. y Vd, 2 obtains similar result. For Logit Normalizaiton, and find that n -1/2 is almost the best one. With this temperature, the only difference between LogitNorm and NormSoftmax is whether the input is shifted to zero-mean, as shown Equation equation 9. The accuracy with this temperature is 87. 37.72%, which is even worse than the baseline, demonstrating the importance of shifting.
sweep the parameter of temperature т,
Other benchmarks. We follow the reference implementation provided by PyTorch (Paszke et al., 2019) to train different ViT's on ImageNet (Deng et al., 2009) from seratch. Strong data augmenta- tions and many techniques are adopted. Results are listed in Table 3. NormSoftmax achieve better performance with a small number of epochs and similar performance with 300-epoch training.
100 epochs
300 epochs
nsm-sqrtd
nsm-inf
sm
nsm-inf
ViT-B-32
71.40
72.01
72.64
75.91
75.92
75.95
ViT-L-32
72.54
73.49
73.46
76.97
77.01
76.92
ViT-B-16
76.52
77.01
77.10
81.07
81.05
81.09
Swin-T
76.91
77.50
77.42
81.41
81.52
81.45
Table 2: Test accuracy of three VIT variants trained on ImageNet-1K from scratch.
We conduct experiments machine translation with Transformers following the settings in (Xu et al., 2019). The benchmarks are WMT English-German Translation (en-de), IWSLT 2014 German-English Translation (de-en), and IWSLT 2015 Einlish--ietmanenn Translation (en-vi) (Cet- tolo et al., 2015). We replace the standard softmax function in both encoder and decoder with our proposed NormSoftmax. The evaluation metric is BLEU (Papineni al., 2002). Similar to the re- sults in computer vision, NormSoftmax can also accelerate the learning process in natural language processing.
4.2 CROSS ENTROPY LOSS OF THE CLASSIFICATION PROBLEM.
We follow the example in the JAX framework (Bradbury al., 2018) train ResNets (He al., 2016) on ImageNet, which has 1,000 classes and about 1.3 million training images. We use SGDM with linear warmup and cosine annealing learning rate scheduler, accompanied by large
The implementation is at this link.
Thee implementation is available this link.
8

---

Under review as a conference paper at ICLR 2023
45k steps
90k steps
sm
nsm-squtd
nsm-inf
sm
nsm-sqrrd
nsm-inf
en-de
24.2
25.6
25.5
28.3
28.3
28.2
de-en
30.1
30.9
31.5
35.4
35.5
35.4
en-vi
26.7
27.1
27.2
31.2
31.3
31.4
Table 3: The BLEU on three machine translation benchmarks with Transformers.
mini-batch size of 8,192 and a large learning rate 3.2. We only enable the horizontal flip and input normalization as data augmentation techniques. We set as 1 and +o0 in NormSoftnax for the cross entropy loss since is the temperature in the baseline. We also apply Logit Normalization and compare it with our proposed method. For Logit Normalization, we conduct grid search to find the optimal hyperparameter of temperature.
epochs
20
45
90
180
sm
69.18
74.71
76.26
76.88
RS0
nsm-1
70.18
74.83
76.41
76.91
nsm-inf
71.42
75.32
76.64
77.17
76
lognorm
69.70
74.79
76.42
76.92
sm 72.44
76.21
77.79
77.99
74
R101
72.48
76.18
77.80
78.08
nsm-inf 72.69
76.56
78.00
78.13
72
sm nsm-l nsm-inf logit-norm
lognorm
72.50
76.19
77.82
78.01
70
72.79
76.78
78.18
78.51
R152
nsm-1 72.80
76.85
78.32
78.50
nsm-inf 72.60
76.73
78.17
78.44
25
50
75
100 opoch
125
150
175
lognorm
72.75
76.81
78.15
78.40
Figure 5:
Test accuracy of ResNet50 Table 4: Test accuracy on ImageNet with ResNets.
ImageNet- 1K with different training epochs lognorm short for Logit Normalization.
Figure 5 and Table 4 show that the NormSoftmax boost the initial training. With 20 epochs, NSM-ine inf can achieve the test accuracy of 1.42% while the baseline with standard softmax obtains 69.18%. With sufficiently long epochs, SM and NSM achieve similar test accuracy, implying that NSM does not impede the representation learning ability of the models. With larger pre-defined threshold y, NSM-inf is faster than NSM-1. NormSoftmax achieves better results than Logit Normalization, demonstrating that it is necessary to shift the input vector before applying Logit Normalization.
5 CONCLUSION
We deeply investigate the behavior softmax in neural network training and discuss its impacts on training stability and convergence. We find that one of the reasons for the optimization difficulty is the significant change in the variance of softmax input during the early training process. To rem- edy the optimization difficulty of softmax, we propose simple yet effective substitution, named NormSoftmax, where the input vectors are first re-scaled by dynamically calculated vector-specific factors and then fed to the standard softmax function. Similar other existing normalization layers in machine learning models, NormSoftmax can stabilize and accelerate the training process and also increase the robustness of the training procedure to hyperparameters. Experiments on Transformers in computer vision and natural language processing benchmarks validate that our proposed Norm- Softmax an effective plug-and-play module to stabilize and speed up the optimization of neural networks with cross-entropy loss or dot-product attention operations.
Limitations. Transformer based models are data hungry and its performance is limited given small intermediate training data. Their representation ability be unleashed with larger datasets. In this work, we only present the results on small- or mid- size datasets. The training behavior on larger datasets may be different from what observe smaller ones.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Atish Agarwala, Jeffrey Pennington, Yann Dauphin, and Sam Schoenholz. Temperature check: theory and practice for training models with softmax-cross-entropy losses. arXiv preprint arXiv:2010.07344 07.44444. 2020.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv:1607.06450, 2016.
arXiv preprint
Nils Bjorck, Carla P Gomes, Bart Selman, and Kilian Q Weinberger. Understanding batch normal- ization. Advances in neural information processing systems. 31, 2018.
Daniel Bolya, Cheng- Yang Fu, Xiaoliang Dai, Peizhao Zhang, and Judy Hoffman. Hydra attention: Efficient attention with many heads. arXiv preprint arXiv:220.004884 2022.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman Minne,, and Qiao Zhang. JAX: composable transformations Python+NumPy programs, 2018. URL http: //github.../ com/ pepoollejaas
Mauro Cettolo, Jan Niehues, Sebastian Stüker, Luisa Bentivogli, Roldano Cattoni, and Marcello The IWSLT 2015 evaluation campaign. In Proceedings of the 12th International Work- shop on Spoken Language Translation: Evaluation Campaign, pp. 2-14, Da Nang, Vietnam, De- cember 3-4 2015. URL https: /wwlaataaaannta.c.aa.aaaa org 2015 iws evaluation. 1.
Federico.
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training strong data augmentations. arXiv preprint arXiv:2106.01548, 2021.
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with a reduced search space. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition workshops, pp. 702-703, 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248-255, 2009. doi: 10. Sttp:C/VR.00095200000
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume 1 (Long and Short Papers), pp. 4171-4188 Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 188553/11/N19914223 URL https: / //aclaanate daclanthologg org/N1ggg/ 9-11199-4233
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An im- age is worth 16x16 words: Transformers for image recognition at scale. International Conference on Learning Representations (ICLR), 2021.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. arXiv preprint arXiv:2010.01412, 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimiza- tion through better initialization. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learn- ing Research, pp. 4475 4483. PMLR, 13-18 Jul 2020. URL https: //proceedinggsseeeeeeeee/ee mlr. press/vll /9/hhagg200 html.
10

---

Under review as a conference paper at ICLR 2023
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference
Machine Learning, volume 37 of Proceedings of Ma- chine Learning Research, pp. 448 448-456, -566,456,11 Lille, France, 07-09 Jul 2015. PMLR. URL https: / //proceeding ml press /v37/offfe1555. httml.
Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak Shah. Transformers in vision: A survey. ACM Computing Surveys (CSUR), 2021.
Simon Kornblith, Ting Chen, Honglak Lee, and Mohammad Norouzi. Why do better loss functions lead to less transferable features? Advances in Neural Information Processing Systems, 34: 18648-286666 2021.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Seung Hoon Lee, Seunghyun Lee, and Byung Cheol Song. Vision transformer for small-size datasets. arXiv preprint arXiv:21 2021.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the dif- ficulty of training transformers. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing (EMNLP), pp. 5747-5763, Online, November 2020. Associ- ation for Computational Linguistics. doi: 10.18653/v12022.emmm 1et.:0eemmlmmmmi4....... URL https:
//aclanthologg org/2020. em maiin... 463.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. Interna- tional Conference on Learning Representations (ICLR), 2017.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Conference on Learning Representations (ICLR), 2019.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normal-
ization layers: Sharpness reduction. arXiv preprint arXiv:2206.0700 2022.
Samuel G Müller and Frank Hutter. Trivialaugment: Tuning-free free yet state-of-the of-the-art data augmen- tation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 774. 782, 2021.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: A method for automatic evaluation machine translation. In Proceedings of the 40th Annual Meeting on Association for Computational Linguistics, ACL '02. pp. 311-318, USA, 2002. Association for Computa- tional Linguistics. doi: 10.3115/107708831077333. URL https: /twii http://iiio.oo...//////o 3115/ 1073083. .100833.007300 1073135.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vito,, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019. URL http: /wtpppaaappppe/eeeeeeeeeaeeeee.. wee.einerrpp. ctt cc/paper/ pdf.
Beygelzimer, d'Alché-Bue,
Tim Pearce, Alexandra Brintrup, and Jun Zhu. Understanding softmax confidence and uncertainty. arXiv preprint arXiv:2106.049/ 2021.
Shibani Santurkar, Dimitris Tsipras, Andrew Ilyas, and Aleksander Madry. How does batch normal- ization help optimization? Advances in neural information processing systems, 31, 2018
Christian Szegedy. Vincent Vanhoucke. Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2818-2826, 2016. doi: 10. 1...10/2......0..000...

---

Under review as a conference paper at ICLR 2023
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing in- gredient for fast stylization. arXiv preprint arXiv:1607.08022. 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 30, 2017.
Hongxin Wei, Renchunzi Xie, Hao Cheng, Lei Feng. Bo An, and Yixuan Li. Mitigating neural network overconfidence with logit normalization. International Conference on Machine Learning (ICML). 2022.
Xiaoxia Wu, Ethan Dyer, and Behnam Neyshabur, When do curricula work? arXiv preprint arXiv:2012.03107, 2020.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference computer vision (ECCV), pp. 3-19, 2018.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng. Shuxin Zheng, Chen Xing. Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In Hal Daumé III and Aarti Singh (eds.). Proceedings of the 37th International Conference on Machine Learning, volume 19 of Proceedings of Machine Learning Research, pp. 10524-10533. PMLR, 13-18 Jul 2020. URL https: /ttpe//eecccee...////o////....se mmr..prese http/vvl//////// vv119/xiog20b.. html.
Jingjing Xu, Xu Sun, Zhiyuan Zhang, Guangxiang Zhao, and Junyang Lin. Understanding and improving layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In Proceed- ings of the IEEE/CVF international conference on computer vision, pp. 6023-6032, 2019.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. International Conference Learning Representations (ICLR), 2018.
Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yang. Random erasing data augmen- tation. In Proceedings of the AAAl conference on artificial intelligence, volume 34, pp. 13001- 13008, 2020.
A PROOF OF LEMMAS
A.1 LEMMA
Let y α(m) softmax(y). Following the chain rules, have
01
0l 0z
(diag(z) (z) xz²") T)
(10)
8y
01
1
xx u(x)
01
I (1- yy"
p(y)1y"
Øl
(11)
dx
(xx)
no22 (x)
Oy
(ar)
n Oy
With 1 1, we can verify that
1 01
0l
-(1" (1'diag(z)
1 =x²"
(12)
Oy
0z
01
(2T x²)
(13)
0z
0
(14)
12

---

Under review as a conference paper at ICLR 2023
With nu(x)) 1"a x; can prove that
2 01
1
1 xxx" - [x(x) 1x
0l
1²
(15)
dz
(xx)
no² (x)
dy
1
1
nu(x) -
0l
(16)
o (xx)
no (x)
dy
1
1" 0l
(17)
o a
dy
(18)
Hence, we prove that = μ ay
0.
We then prove the scaling for L norm by leveraging
8m
the definition of variance (²(a) = |a
/ n - pa))).
0L
01
01
(19)
da:
da
1
T
xx. - ((x))a1
II- (1-
xx.""
((x))
Øl
I
(20)
x2(x)
Oy
no²
no²
(x)
dy
T
ax."" T
1x)) + p(x)ar1
Øl
(21)
(2(x)
(x²(x))
no?
dy
T
1
01
0l
xx T
O
m(x)
8y
By
uux))
x:
8y
(2())
Oy
no 2 (ar)
Oy
n
(x)
1
(22)
1
01
2
01
T
xx
01
(23)
(2
0y
Oy
no (x) By
2
01
1
Øl By
2
a:
(24)
(xx)
0y
not
(x))
1
01
(25)
0 (m²)
3y
12
From the definition of variance, we obtain that
Øl
01
2
1
0l
2
0l
0²
in
/n
o²
(26)
dx
dx
2
(2)
x
Öy
12
x²(x))
By
Above all, we prove that
Ø1
7y
01
0
Oy
<
(27)
dæ
(xx)
0
o (
A.2 LEMMA 2
Given 062 ka., we have 0m/o(m)) (2/2)(;2). Thus, we obtain that 2 softmax (z)))((;))) 2g = softmax (x)2)) Now that = Z2, we can obtain the same
0t st loss f(z1) = f(22) and generate the same gradient Uzi dzz
Based on Equation equation 11, we
conclude that Jaz & Jmm
B EXPERIMENTAL DETAILS.
B. VISION TRANSFORMERS ON CIFAR10.
We train from scratch with AdamW optimizer for 100 epochs (50,000 iterations with 100 mini-batch size) in mixed precision. The resolution of input image is 3 x 32 32, and the patch size is 4.
13

---

Under review as a conference paper at ICLR 2023
100
75
test accuracy
50
nsm-1
25
0
20
40
60
80
epoch
Figure 6: Test accuracy of a ResNet 202 on CIFAR-10
The hidden size, MIP size, number of heads, dimension of heads, and number of layers are 256, 1024, 8, 32, and 8, respectively. We use the following attention function,
QKT
Attention K, V) = softmax min (Qiin((QQ"))))) V
(28)
where the softmax and the standard deviation are calculated along the same axes. We set vd and +0o. The learning rate is le-3, and a strong weight decay le-1 is applied. The learning rate linearly increases from 2e-4 with 5 warmup epochs and then decays to O with the cosine annealing scheduler. We also enable label smoothing (0.1), random erasing with the probability of 0.1, mixup with a = 0.2, cutmix with a = 1.0, and TrimialAugment. The code is attached in the supplementary material.
B.2 VISION TRANSFORMERS ON IMAGENET.
We follow the torchvision's 's reference implementation. The batch size is 512 x 8 4096. We train from scratch with AdamW W optimizer mixed precision. The learning rate is 3e-3, and the weight decay is 3e-1. The learning rate is linearly increased from 3e- 3 0.033 with 30 warmup epochs and then decays to with the cosine annealing scheduler. We also enable label smoothing (0.1) mixup with a 0.2, cutmix with a 1.0, clipping gradient norm with RandAugment (Cubuk al., 2020), repeated augmentation with repetitions, exponential moving average for model parameters.
B.3 RESNET ON IMAGENET.
We follow the example in the JAX framework. We only enable horizontal flip data augmentation. We use SGDM with mini-batch size of 8,192, a learning rate of 3.2, momentum of 0.9, and weight decay of le-4. The learning rate warmup epochs.
B.4 MACHINE TRANSLATION.
We exactly the settings Xu al. (2019). Please refer the original paper for reference.
C EXTENDED EXPERIMENTS
C.1 TRAINING PERFORMANCE OF DEEP RESNET
We train a 2esNet1202, which is investigated in the original ResNet paper (He et al., 2016). Figure 6 compares the test accuracy during the training process. NormSoftmax significantly acceler- ate the training process and achieve better test accuracy than the standard softmax. With standard softmax, the model is stuck in the low confidence zone in the initial stage.
14

---

Under review as a conference paper at ICLR 2023
90
88
lest accuracy
86
84
82
nsm-sqrid nsm-inf
80
78
0.4
0.6
temperature
Figure 7: Test accuracy of a VIT with different temperatures after normalization.
C.2 SCALING FACTORS IN ATTENTION
The default scaling factor in attention d=1/2 We sweep the scaling factors for the standard softmax. We apply the corresponding 7 in NormSoftmax. We follow the same experiment settings as in Section 4.1.
scaling factor for sm
2d 1/2
d
d-1/2
1/2
d' d-1/2//
d=
test accuracy for sm 80.98
86.38
87.96
87.97
86.98
for nsm
d//3//2 /2
d//2
2d
d
0C
test accuracy for nsm
81.04
87.16
88.13
88.21
87.1
88.01
Table 5: Results for (1) the standard softmax (sm) with different scaling factors and (2) NormSoft- max (nsm) with different y
Table 5 shows that scaling factors of d-1/2 /2 and d-1 achieve the best performance for the stan- dard softmax, similar to (Lee et al., 2021). NormSoftmax shares the same trend with the standard softmax, achieving the best performance with y of 2d2/2 and d'//2, NormSoftmax achieves better results than the standard softmax with the corresponding pair of scaling factor and
C.3 TEMPERATURE AFTER NORMALIZATION IN NORMSOFTMAX
We add temperature parameter 7 in NormSoftmax.
NormSoftmax(2; 7, 7) = softmax (a)
(29)
If = +00, the input of softmax will always have standard deviation of --1 For finite y, we have a mmin(r(w)).))
--
Figure illustrates the results with different temperature parameters. both NSM-img inf and NSK-sqrtd. When 7 < 1, NSM-Msinn inf is much more stable than NSM-sqrta since the softmax input may have high variance for WIMM-sqrte On the contrary, when 1, NSM-sqrtd is better than NSM-inf inf. The softmax input of NSM-int has low standard deviation, restricting the performance.
is the best choice for
D CoST ANALYSIS OF NORMSOFTMAX
Since import normalization in the softmax, inevitably introduce the extra computation and memory cost of NormSofumax. However, we show that the overhead of the proposed lightweight NormSoftmax is negligible for large machine learning models.
Memory cost. Equation indicates that we can preprocess the input vector x and pass it to Norm- Softmax. We can fuse the factor with the ascending layer of NormSoftmax if possible during
15

---

Under review as a conference paper at ICLR 2023
inference. For a linear layer Ay+ can always fuse the with A and b, as we do not need to save y during inference. For training, we need to pay the extra memory cost for normalization since we have to save intermediate results for gradient computation.
Computation cost. For training, we have to pay the cost of calculating the variance and the cor- responding gradients. For inference, may discard the normalization if the softmax is in the last layer. In a classification model with cross entropy loss, we can discard the normalization since the normalization has no impact the classification result.
In our experiments on Transformer, NormSoftmax induces 0.2% - 1% extra execution time. For the experiments for cross-entropy loss, the extra execution time is less than 0.05% since we only add one normalization layer.
16


Reviews:
[{"review_id": "tjBFCp-l3_", "reviewer": "Reviewer_Eoot", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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 extends the Logit Normalization technique by centering and introducing  a threshold  gamma in the denominator, which is referred to as NormSoftMax. Experimental results indicate that NormSoftMax works well in a number of DNN models with either croess-entropy loss or attention modules.  ", "summary_of_the_review": "I think NormSoftMax need to be properly motivated in introduction by properly introducing  Logit Normalization. The experiments need to be elaborated to demonstrate the advantage of  NormSoftMax  over Logit Normalization.", "strength_and_weaknesses": "Strength: \nNormSoftMax extends Logit Normalization by introducing centering and a threshold  gamma in the denominator. \n\nWeaknesses:\n1. The introduction does not clearly describe the close relationship between Logit Normalization and NormSoftMax. The authors should first mention the drawbacks of Logit Normalization and then motivate NormSoftMax. Note that Logit Normalization is designed to address a similar issue as NormSoftMax. That is, to reduce the dynamics of the inputs to softmax to make the training easier. \n2.  I don't get why Logit Normalization is not compared with NormSoftMax  in their experiments. Or NormSoftMax with gamma=infty is reduced to  Logit Normalization. If it is true, I don't see the advantage of  NormSoftMax with finite gamma. Please correct me if I am wrong. \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": {}}, {"review_id": "tPh13KBenA", "reviewer": "Reviewer_7Jaw", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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 propose a normalization method, NormSoftmax, to remedy the optimization difficulty of softmax when using attention model or cross entropy loss function. Their method is easy to implement, its effectiveness is verified on several data tasks including CIFAR10, Imagenet, and three NLP translation datasets.", "summary_of_the_review": "I think two major weaknesses needed to be solved to improve the quality of the paper: first, justify the reasoning that why Normsoftmax can improve the training of neural network by stabilizing the softmax input; second, provide more convincing empirical evidence to verify the effectiveness.", "strength_and_weaknesses": "Strengths: writing is good, I can easily get the main points of this paper. \n\nWeakness: \n\n1) The motivation of this paper seems not convincing. As far as I can see, the motivation of Normsoftmax is based on the reasoning in sec 3.1, that \"the variance of softmax input experiences rapid and huge change during training, especially in the initial stage, which explains why training from scratch is difficult.\". I agree with the observation that the softmax inputs suffer from high variance during early training (fig 2 shows that),  but it's not the unique property of softmax, the statistics of batch normalization also fluctuate dramatically at early stage (see fig 2(a) in [1]), and these phenomenons can be well interpreted by the findings of [2], that the training trajectory usually starts from a rugged place to a flat local optimum on loss landscape. The main question is: is this feature bad? Can it really slow down the training process or harm the final performance of the trained model? Current paper is lack of justification on this fundamental reasoning.\n\n2) Current experimental results are not sufficient to verify the effectiveness of the proposed method. The form of Normsoftmax is very simple, which I don't regard as a disadvantage, but very convincing experimental results are required to prove its effectiveness. However, the gap between Normsoftmax and its baseline is too minor: VIT on CIFAR10 is not a good setting to show empirical evidence, let alone the performance is not in SOTA level; As for imagenet and NLP dataset, the improvement on training speed and final performance are both subtle.\n\n3)  Cost analysis (sec 3.4) is ambiguous. Fusing $\\gamma$ into linear layer is feasible, but computing $\\sigma (x/\\gamma )$ still requires extra memory/computation cost during inference, right? I suggest conducting thorough ablation study to study the extra cost brought by Normsoftmax.\n\n[1] Yan, Junjie, et al. \"Towards Stabilizing Batch Statistics in Backward Propagation of Batch Normalization.\" ICLR. 2019.\n\n[2] Li, Hao, et al. \"Visualizing the loss landscape of neural nets.\" NIPS (2018).", "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": "KCACyaLdL7", "reviewer": "Reviewer_ax8S", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "In this work the authors attempt to tackle the issue of large softmax inputs, particularly in attention layers of transformers. They propose to normalize the input to the softmax by min(\\sigma, \\gamma), where \\gamma is a constant, and \\sigma is the unscaled norm of the input logits. They show that this normalization seems to improve performance on a variety of tasks, and in addition makes hyperparameter transfer (specifically, changing the attention dimension) more robust.", "summary_of_the_review": "Overall the method is simple and a good attempt at dealing with the large softmax input problem that plagues attention layers in large models. The ability to transfer hyperparameters more easily across geometry is particularly impressive. If some of the details from Figure 3 were more fleshed out, giving stronger results from the method, I would be comfortable switching my review result to an 8 - the paper is almost there.", "strength_and_weaknesses": "The method is a very simple and elegant solution to the problem of logit normalization in the attention layer of transformer networks. The method seems to lead to small performance gains in many settings, and some performance loss in other settings. The most striking result in my mind is figure 3 (d) - that hyperparameter transfer, in terms of hidden size/head dimension, is more robust with the softmax norm.\n\nFigures 3a and 3c are interesting in particular. It seems that for a small number of epochs the method is helpful, but for a large number of epochs it may not be (or even is slightly worse than the baseline). It does seem that the method helps with increasing depth. For figure 3 (a), what happens if the normalization is used early but not later? It seems that it may hurt more than be helpful at later training epochs. For figure 3 (c), how was hyperparameter tuning done for large depth? If both methods were tuned well, then the result seems quite significant.\n\nOne experiment that would greatly strengthen the paper is the comparison of the method with the 1/d scaling for the attention layers (with no norm-softmax) - and possibly an experiment where gamma is given by d instead of \\sqrt{d}. The 1/d scaling has also been brought up as a method for fixing anomalously large softmax inputs in attention layers, and I'm curious about the relationship between the two methods.\n\nAnother interesting experiment would be tuning the inverse-temperature of the softmax after normalization, as in https://arxiv.org/abs/2010.07344. It's possible that the normalization scale of 1 is sub-optimal for attention - there may be cases where a slightly larger softmax norm allows for more useful non-linearities in the attention mask.", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "7eVwiH4hFMK", "reviewer": "Reviewer_syPR", "summary": "", "questions": "", "limitations": "", "rating": 5, "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 proposes a variant of the softmax function used in many modules of recently developed deep neural networks.\nThe authors first point out the shortcoming of the softmax function, which is unstable in the initial stage of model training.\nThe key idea of the proposed method is to borrow the normalization technique, e.g., layer-normalization, for stabilizing the training to incorporate the variance term in softmax computation.\nThe experiments are conducted on simple image classification and machine translation benchmark datasets.\nThe results show that the proposed method performed better in the initial training phase stage.\n", "summary_of_the_review": "This paper has certain amount of contributions to the community.\nHowever, the advantage of the proposed method is only limited to the initial stage of the training phase.\nIt would be much better to consider and show a scenario that the proposed method can show much better advantages.\n", "strength_and_weaknesses": "Strength:\n* Motivation of this study is understandable.\n* The method is reasonable.\n* This paper is well-organized, easy to follow the main points of this paper.\n\n\nWeaknesses:\n* While the effectiveness of the proposed method only shown in the initial stage of the training phase, it is hard to feel that this method is really effective in the actual use case. I understand the effectiveness, but broader impact in the community seems limited.\n* The deep normalization may lead the gradient vanishing problem. This paper does not discuss this point. Therefore, it may have a potential drawback when the model becomes much deeper.\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