Softmax Gradient Tampering: Decoupling the Backward Pass for Improved Fitting

1 INTRODUCTION . 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 . 2 METHOD : SOFTMAX GRADIENT TAMPERING . 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 . 3 EXPERIMENTS . 3.1 GRID SEARCH EXPERIMENTS . 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 ) . 3.2 EXPERIMENT ON THE IMAGENET DATASET . 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 . 3.3 EXPERIMENT ON CIFAR DATASETS . 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 . 3.4 TOWARDS REMOVING NORMALIZATION . 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 . 3.5 EFFECT ON LOSS , LOGITS AND OTHER METRICS . 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 .

Summary: 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.