Year: 2021

Abstract:
We introduce Softmax Gradient Tampering, a technique for modifying the gradients in the backward pass of neural networks in order to enhance their accuracy. Our approach transforms the predicted probability values using a powerbased probability transformation and then recomputes the gradients in the backward pass. This modification results in a smoother gradient profile, which we demonstrate empirically and theoretically. We do a grid search for the transform parameters on residual networks. We demonstrate that modifying the softmax gradients in ConvNets may result in increased training accuracy, thus increasing the fit across the training data and maximally utilizing the learning capacity of neural networks. We get better test metrics and lower generalization gaps when combined with regularization techniques such as label smoothing. Softmax gradient tampering improves ResNet-50’s test accuracy by 0.52% over the baseline on the ImageNet dataset. Our approach is very generic and may be used across a wide range of different network architectures and datasets.

heading: 1 INTRODUCTION
text: Smooth gradient flow is the key to successful convergence of deep neural networks. Batch Normalizing Ioffe & Szegedy (2015), Weight Standardization Qiao et al. (2019), and Group Normalization Wu & He (2018) are all types of normalization techniques that smooth the gradient landscape in the backward pass. Batch normalization smoothes the loss and gradient surfaces by limiting their lipschitzness Santurkar et al. (2018). Methods of smoothing gradients by manipulating the labels distribution include label smoothing Szegedy et al. (2016), mixup Zhang et al. (2017).
In this paper, we manipulate the gradients in an unconventional way to impose smooth gradients in the backward pass. This technique is one of a family of Gradient Tampering methods, called Softmax Gradient Tampering since it is conducted at the softmax stage. We demonstrate that Softmax Gradient Tampering improves convergence fit and therefore improves training and testing accuracy.
Our main contributions in this paper are the following:
1. We propose Softmax Gradient Tampering that improves the learning capacity of neural networks and allows for a much better fit on the training dataset. This reflects both in training and validation accuracy metrics.
2. We theoretically analyze how softmax gradient tampering works from the perspective of gradient smoothness.
3. We provide comprehensive findings from different models, and datasets.

heading: 2 METHOD: SOFTMAX GRADIENT TAMPERING
text: The softmax gradient tampering technique is described in this section. Softmax gradient tampering modifies the gradient in the backward pass. It does not alter the behaviour in the forward pass. Before we describe our method, we briefly explain the forward pass itself. A neural network with parameters W generates C logits denoted by z for every input vector x. z is given as z = Wx. Then a set of probability values pi are generated from the logits using a softmax function which
is defined as pi = e zi∑C
j=1 zj . pi and zi represent the predicted probability values and the logits for
the ith class respectively. Following this step, the loss function is invoked and the loss between the predicted probability values the ground truth labels (which is also a probability distribution) is calculated. If the loss function is cross-entropy loss, then the value of the loss is given as L = − ∑C
i=1 qi log (pi) where qi is the ground truth label of the i th class for a particular training example. By backpropagation rule, we can calculate the gradient of the loss with respect to the logits which takes the form ∂L∂zi = pi − qi.
The softmax gradient tampering technique is now described. We introduce a hyperparameter α, which takes a value between [0, 1] and regulates the degree of tampering. The softmax gradient tampering method modifies the predicted probability values in the backward pass as follows:
p′i = pαi∑C j=1 p α j
i = 1, . . . , C 0 ≤ α ≤ 1 (1)
The above transformation changes the gradient of the loss with respect to the logits as follows:
∂L ∂zi = p′i − qi (2)
The rest of the backward pass proceeds as usual. We denote the original probability distribution as P (with values pi at the ith index) and the transformed distribution as P ′ (with values p′i at the i th index).
A code snippet of the algorithm implemented using PyTorch Paszke et al. (2019) is given in the appendix. In PyTorch, the available method to modify gradients in the backward pass is using the register_hook function which is called in the forward function.
The backward pass gradient profile is smoothed by the transformation of the probability distribution using equation 1. Because the gradient at the softmax layer’s input (∂L/∂zi) is directly proportional to the predicted probabilities pi, smoothing out the probability distribution also smooths out the gradient profile. We explore the effect of this change from a theoretical standpoint in section 5. In section 3, we offer empirical data and experiments to support the proposed method.

heading: 3 EXPERIMENTS

heading: 3.1 GRID SEARCH EXPERIMENTS
text: We do a thorough hyperparameter search on the ResNet-18 He et al. (2016) architecture and the ImageNet-1K dataset, changing the value of α from 0.1 to 0.9. We begin with a coarse grid search
in 0.1 increments. The optimum values of α are found to be between 0.2 and 0.4. Then we do a fine grid search in 0.05 increments between 0.2 and 0.4. We find that the optimum performance for ResNet-18 is achieved at a α = 0.25. We discover that the optimum value of α is 0.3 by repeating the grid search for ResNet-50 He et al. (2016) in the range of 0.2 to 0.4 with increments of 0.05. We utilise a batch size of 1024 over four V100 GPUs for each grid search experiment. Each experiment is carried out with a 50 epoch budget. In the first 2 epochs, we raise the learning rate linearly from 4× 10−4 to 0.4. The learning rate is then decayed using a cosine scheduler (single schedule with no restarts) Loshchilov & Hutter (2016) from 0.4 at the start of the 2nd epoch to 4× 10−4 at the end of the 46th epoch. We continue training for the next 4 epochs at a constant learning rate of 4 × 10−4 (cooldown phase). The plots are shown in Figure (2). We do a grid search on the ResNet-20 He et al. (2016) architecture and the CIFAR-10 dataset, with a epoch budget of 360. However, we get no advantage for any value of α. This is covered in more detail in section 3.3.
Figure (2a,b) shows a pattern in which training accuracy rises when α is lowered from 1 to 0, indicating the method’s potential as an algorithm to drive the network towards better learnt feature representations and therefore a better fit on the training dataset. This is due to the transform’s nature, which smoothes the gradients. This line of reasoning is expanded upon in section 5.
The generalization gap is lowest when α = 0.2 but it comes at the expense of test accuracy, so we disregard it. Values of α < 0.2 cause issues with half precision training, therefore we choose α = 0.25. In our trails with ResNet-18, we find optimal values of α = 0.25, while in our experiments with ResNet-50, we get optimal values of α = 0.3. We port these values for all 100 epoch schedules, as well as the Squeeze-and-Excitation versions of the ResNets. We indeed find a peak at α = 0.1, but we found that such low α values are typically unstable in most subsequent tests, especially with mixed precision or half precision arithmetic. According to our observations in section 3.5, low values of α frequently result in high logit norms, which directly cause floating point problems Micikevicius et al. (2017).

heading: 3.2 EXPERIMENT ON THE IMAGENET DATASET
text: We perform experiments on the ImageNet-1K dataset Deng et al. (2009), which has 1.28 million images divided into 1000 classes. ImageNet-1K’s test set consists of 5000 images per class. We train models based on several ResNets, including the Squeeze-and-Excitation versions. We observe a constant 0.1% to 0.15% rise in ResNet-18’s test accuracy. What’s more interesting is that we’re seeing comparable improvements in training accuracy. Softmax gradient tampering is not a regularization technique, as shown by this. It just allows the network to better fit the dataset. In section 4, we show how combining softmax gradient tampering with regularization techniques such as label smoothing may improve network performance and reduce the generalization gap.
We find significantly larger improvements in experiments using ResNet-50, a 0.5% increase over the step scheduling baseline and a 0.3% increase over the cosine scheduling baseline. In studies with Squeeze-and-Excitation networks, we use the cosine scheduler. In these experiments, we find consistent improvements across the board, with SEResNet-18 Hu et al. (2018) gaining 0.4% and
SEResNet-50 Hu et al. (2018) gaining 0.7%. As we saw with ResNet-18, we find continuous improvements in training performance because the networks are just training better and arriving to better optimas at convergence. ResNet-101 training provides a modest advantage of 0.35%.
All models are trained on four V100 GPUs with a batch size of 1024. We utilise the same set of hyperparameters in each experiment, which are as follows: 100 epoch budget, 5 epochs linear warmup phase beginning with a learning rate of 4 × 10−4 and ending with a peak learning rate of 0.4. In our studies, we employ either a step scheduler (dividing the learning rate by 10 at the 30th, 60th, and 90th epochs) or a cosine decay scheduler. We observe a consistent improvement in test metrics when utilising cosine scheduling across all of our experiments. Finally, towards the conclusion of training, we impose a 10 epoch cooling phase in which the network is trained with a constant learning rate schedule set at 4×10−4. Other hyperparameters include a weight decay value of 5× 10−4 and a momentum of 0.9 for the SGD Nesterov optimizer. All of our models are trained using mixed floating point precision. The findings are shown in Table 1.

heading: 3.3 EXPERIMENT ON CIFAR DATASETS
text: The CIFAR-10 dataset Krizhevsky et al. (2014) is divided into 10 classes, each with 5000 datapoints for training and 1000 datapoints for testing. In comparison to ImageNet-1K, it is a much smaller dataset. The most frequent situation seen when training models on CIFAR-10 is that the models fully overfit the training dataset, achieving almost 100% accuracy. Given that the proposed method (section 2) is not intended as a regularization technique, no improvements in the training of such datasets are anticipated. We see the same thing in Figure 2c, where we show the results of a hyperparameter grid search. Different values of α do not impair ResNet-20’s performance on CIFAR-10, but neither do they improve on the baseline.

heading: 3.4 TOWARDS REMOVING NORMALIZATION
text: Experiments are carried out on non-Batch Normalized networks. We specifically utilize the ResNet18 model and train it on ImageNet-1K without any of the BatchNorm layers. We see a substantial increase in test accuracy of about 1%. Similarly to prior findings, we get a 1.1% increase in training accuracy, leading us to infer that when softmax gradient tampering is used, the network’s convergence optima is significantly superior. Table 2 shows the findings.

heading: 3.5 EFFECT ON LOSS, LOGITS AND OTHER METRICS
text: As seen in Figure (3a), the logit norm increases as α falls from 1 to 0. In the initial training phase, softmax gradient tampering causes the gradient in both the valid and wrong classes to rise as the
network misclassifies majority of the examples. Due to the non-linear nature of the softmax function, the predicted probability distribution may frequently be close to a delta distribution, even in instances of misclassifications. But softmax gradient tampering reshapes the predicted probability distribution. It redistributes probability from the confident classes to the less confident ones. Gradient rises in all classes. We also see that the final value of the loss rises as α drops, and that the logit norm and the final value of the loss are linearly correlated Figure (3c).
ResNet-18 and ResNet-50 per epoch training and test accuracies with and without softmax gradient manipulation are shown in Figure (4a,b). In all instances, we observe a gain. This means that softmax gradient tampering forces the network to learn better representations.

heading: 4 ABLATION STUDY
text: We conduct an ablation study to investigate the consequences of softmax gradient tampering. We have previously demonstrated from the experiments depicted in section 3 that softmax gradient tampering causes the network to arrive at a better optima and fit a particular training dataset better than the baseline. Furthermore, based on our CIFAR-10 tests, we have experimentally shown that the technique performs best on varied datasets. Because most networks on the CIFAR-10 dataset begin overfitting extremely early, and therefore a non-regularizing technique like softmax gradient tampering has little effect on the network’s performance. As a result, we combine our approach with other regularization techniques to see whether this is really the case. If softmax gradient tampering is not a regularizer, the generalization gap will widen. The generalization gap would be reduced if we added a regularizer on top of it. The findings are shown in the Table 3. We see that adding softmax gradient tampering increases the gap between the training and test accuracies, while introducing a label smoothing regularizer decreases the gap. When both softmax gradient tampering and label smoothing are applied concurrently, the best accuracies are achieved.

heading: 5 WHY DOES SOFTMAX GRADIENT TAMPERING WORKS?
text: In this section, we describe why softmax gradient tampering work. We use the same setup as described in section 2. We start by investigating the effect of the transform on the softmax probabilities, and then we prove that the transform smoothens the gradient profile.
Lemma 1. For any arbitrary probability distribution P with probability values given by pi for i = 1, . . . , C, the corresponding transformed probability values p′i given by [1] has a threshold(∑C
j=1 p α j
) 1 α−1
and
p′i ≥ pi, if pi ≤ ( C∑
j=1
pαj
) 1 α−1
p′i < pi, if pi > ( C∑
j=1
pαj
) 1 α−1
(3)
Proof. The proof follows directly from the definition of p′i. We notice that if p ′ i < pi, we get:
pαi∑C j=1 p α j < pi (4)
⇒ pα−1i < C∑
j=1
pαj (5)
⇒ pi > ( C∑
j=1
pαj
) 1 α−1
(6)
We call this threshold, the stationary threshold. The stationary threshold is that value of pi that does not change after the transformation.
Proposition 1. At α = 0, the stationary threshold equals 1/C and all values of the transformed distribution p′i reduces reduces to the uniform distribution for i = 1, . . . , C,.
Proof. From equation (3), we see that the stationary threshold at α = 0 is 1/C. Also, following from the definition of the transformed probabilities (1) we conclude that if α = 0, then all values of p′i are 1/C. Therefore the transformed distribution at α = 0 is a uniform distribution.
Since we have established that values of pi which are greater than the stationary threshold decreases and move down towards the stationary threshold, and values in pi lower than the stationary threshold moves up towards the stationary threshold, this transformation makes the distribution more uniform (i.e. it smooths out the actual distribution) as α is decreased from 1 and down towards 0. This final observation we prove in the following theorem.
Theorem 1. For any arbitrary probability distribution P with probability values pi for i = 1, . . . , C, the stationary threshold of the transformed distribution P ′ with probability values p′i = pαi∑C j=1 p α j
, 0 ≤ α ≤ 1 is a monotonically non-decreasing function with respect to α. Proof. To prove monotonicity, we first compute the gradient of the stationary threshold with respect to the variable in concern, α.
∂
∂α  c∑ j=1 pαj  1α−1 =  c∑ j=1 pαj  1α−1 ∑cj=1 pαj log (pj) (α− 1) ∑c j=1 p α j − log (∑c j=1 p α j ) (α− 1)2  (7) = 1
α(α− 1)2  c∑ j=1 pαj  1α−1  (α− 1)∑Cj=1 pαj log (pαj )∑c j=1 p α j − α log  c∑ j=1 pαj  (8) If a1, . . . , an and b1, . . . , bn are non-negative numbers, then using the log sum inequality, we get∑n
j=1 aj log ( aj bj ) ≥ (∑n j=1 aj ) log (∑n j=1 aj∑n j=1 bj ) . Setting aj = pαj and bj = 1, we get the following
lower bound C∑
j=1
pαj log ( pαj ) ≥  C∑ j=1 pαj  log  1 C C∑ j=1 pαj  (9) Substituting (9) in (8), we get:
∂
∂α  c∑ j=1 pαj  1α−1 ≥ 1 α(α− 1)2  c∑ j=1 pαj  1α−1 (1− α) log(C)− log  C∑ j=1 pαj  (10) pα is concave, and so by Jensen’s inequality we get the following upper bound for the second term: 1
C C∑ j=1 pj  α ≥ 1 C C∑ j=1 pαj (11)
⇒ log  C∑ j=1 pαj  ≤ (1− α) log(C) (12) Substituting (12) in (10),
∂
∂α  c∑ j=1 pαj  1α−1 ≥ 0 (13)
We conclude from the analysis that the stationary threshold is a monotonic non-decreasing function with respect to α, which in turn means that the transformation is an order-preserving map. All values greater than the threshold move towards the threshold after transformation and all values below the
threshold also move towards the threshold, and the threshold itself moves monotonically towards 1/C as α is decreased from 1 to 0. This concludes that the transformation smooths the original distribution. Since the gradient with respect to logits (∂L/∂zi) is a direct function of the smoothed out distribution, this smoothes out the gradients as well.

heading: 6 RELATED WORK
text: In this section, we examine techniques that are related to gradient tampering and emphasize the distinctions between the methods.
Relationship to Label Smoothing: Label smoothing, introduced by Szegedy et al. Szegedy et al. (2016), utilizes smoothing of the ground truth labels as a method to impose regularization on the logits and the weights of the neural network.
The formulation is as follows. If qi is the value at the ith index of the one hot encoded vector of the ground truth label, then the transformed distribution of the labels q′i is:
q′i =
{ 1− ϵ if i = y,
ϵ K−1 otherwise
(14)
where ϵ is a hyperparameter with a value in [0, 1].
Müller et al. investigates how label smoothing work in Müller et al. (2019). Using visualizations of penultimate layer activations, they demonstrate that label smoothing calibrates the networks predictions and aligns them with the accuracies of the predictions. Parallel works in label smoothing include Reed et al. (2015); Zhang et al. (2021); Reed et al. (2014); Xie et al. (2016); Dubey et al. (2018); Li et al. (2019).
Softmax gradient tampering works in a much different way. As we proved in section 5 that our proposed transformation of probabilities smooths the prediction vector which in turn has the effect of smoothing the gradient landscape, since the gradient vector arising from softmax is p′i − qi. This has the opposite effect of label smoothing, because the gradient vector arising from softmax in a scenario with label smoothing is pi − q′i. Hence, softmax gradient tampering is not a regularization method. From our studies, we show that it in fact leads to larger generalization gap (Table 1, 3). Coupled with softmax gradient tampering, label smoothing closes the generalization gap.
Relationship to Knowledge Distillation: The formulation of softmax gradient tampering is similar to the temperature based logit suppression formulation of Hinton et al. Hinton et al. (2015). Knowledge distillation Ba & Caruana (2014) is a process in which two networks are trained. First a teacher network is trained on a given dataset and then the soft-labels (the predicted probabilities) of the teacher network are used to train another network, called the student network. In similar style of label smoothing, knowledge distillation aims to generate smooth gradient behaviour by forcing the network not to overpredict on any single image. As in case of label smoothing, the student network’s weights are automatically penalized if the network assigns a probability value higher than the soft labels generated by the teacher network.
The formulation of temperature based logit suppression introduced as a special case of logit matching, so far has only been used in the context of distillation. Even though the formulation of temperature based logit formulation leads to the same formulation of softmax gradient tampering, there is
one key difference, and that is the aforementioned logit suppression is applied in the forward pass. This means that the temperature variable becomes a part of the computation graph. In the case of softmax gradient tampering, where we directly tamper the gradients without introducing any change in the forward pass is a more generalized way of arriving at the smoothness effects of the gradients in the backward pass. The introduction of temperature based logit suppression in the forward pass also leads to making inverse changes in the learning rate or in the formulation of the loss function. However, in case of softmax gradient tampering, no such changes are required. In section 3, we demonstrate that networks using softmax gradient tampering method can be trained using the same hyperparameters as the ones used in the baseline models.
Variants of knowledge distillation include Zhang et al. (2019); Furlanello et al. (2018); Passalis & Tefas (2018); Ge et al. (2019); Yun et al. (2020); Wang et al. (2021); Aguilar et al. (2020); Yao et al. (2018); Wang et al. (2020); Peng et al. (2019); Ge et al. (2018); Zhang et al. (2018b); Liu et al. (2018); Cho & Hariharan (2019); Yuan et al. (2020). Aside from that, the smoothing effects on gradients have never been studied before, and this is the first time, to the best of our knowledge, that this has been done in this study.
Relationship to Gradient Clipping (and variants thereof): Gradient clipping is often used as means of limiting exploding gradients in the backward pass. It normalizes and scales the gradient, if the norm of the gradient exceeds a predetermined threshold.
The formulation is as follows: If G is a gradient vector which is ∂L/∂θ (the gradient of the loss L with respect to the parameters θ)
G → { λ G∥G∥ if ∥G∥ > λ, G otherwise
(15)
It works in practice as it helps the network to avoid gradient explosion during training, but it doesn’t really make the loss landscape smoother than it already is. It simply alters the scale of the gradients with respect to the norm. Brock et al. Brock et al. (2021) finds that the clipping threshold λ is an extremely sensitive hyperparameter.
Adaptive gradient clipping is a variant of gradient clipping. The adaptiveness of adaptive gradient clipping comes from the use of the Frobenius norm of the weight matrix. NF-Nets Brock et al. (2021) uses Adaptive Gradient Clipping (AGC) and achieves close to state-of-the-art accuracies in ImageNet benchmark without the use of traditional normalization methods.
Gradient clipping methods and the variants thereof are also gradient tampering methods, i.e. they alter the behaviour of the backward pass independently from the forward pass. Our proposed method softmax gradient tampering works directly at the softmax level, whereas gradient clipping methods work at the parameter level, requiring much more compute. Softmax gradient tampering on the other hand requires only a small modification and significantly lower computation overhead.
Similar lines of theoretical analysis as in section 5 on logits and temperatures include Zhang et al. (2018a); Pereyra et al. (2017).

heading: 7 DISCUSSION
text: We showed the application of softmax gradient tampering, a technique for providing a smooth gradient during neural network training. With the use of different network topologies and datasets, we were able to show its potential and explore its impacts from a theoretical standpoint.

heading: 8 REPRODUCIBILITY
text: Reproducible code, training recipes for experiments in section 3, pretrained checkpoints and training logs are provided at: https://github.com/bishshoy/ softmax-gradient-tampering. Code adapted from Wightman (2019).
Instability caused by half precision floating point operations is frequent in neural network training procedures, and it is significant when the logits reach high values Micikevicius et al. (2017). When
training ResNet-50 with high batch sizes (1024), we experienced sporadic problems. We were able to easily stabilize the training by training the first epoch without softmax gradient tampering and then training the other epochs with softmax gradient tampering enabled. However, when we trained ResNet-50 with smaller batch sizes (512), we had no problems with reproducibility.

heading: A APPENDIX
text: In PyTorch, the available method to modify gradients in the backward pass is using the register_hook function which is called in the forward function.

Review summaries:
In this paper, the authors propose a technique called Softmax Gradient Tampering, which transforms the predicted output class probabilities to improve training performance of neural networks. The authors show that the proposed technique results in a smoother output probability distribution for lower values of the hyperparameter $\alpha$. On standard benchmarks on the ImageNet and CIFAR-10 datasets, the authors show that the proposed method results in better training accuracy performance as well as improved generalization performance.

This paper proposes the softmax gradient tampering to modify the gradients in the backward pass to enhance the accuracy. The predicted probability value is transformed using a power-based probability transformation, and the gradient profile is more smooth. The experimental results show the slight accuracy increase.