Temperature check: theory and practice for training models with softmax-cross-entropy losses

1 INTRODUCTION . Deep learning has led to breakthroughs across a slew of classification tasks ( LeCun et al. , 1989 ; Krizhevsky et al. , 2012 ; Zagoruyko and Komodakis , 2017 ) . Crucial components of this success have been the use of the softmax function to model predicted class-probabilities combined with the cross-entropy loss function as a measure of distance between the predicted distribution and the label ( Kline and Berardi , 2005 ; Golik et al. , 2013 ) . Significant work has gone into improving the generalization performance of softmax-cross-entropy learning . A particularly successful approach has been to improve overfitting by reducing model confidence ; this has been done by regularizing outputs using confidence regularization ( Pereyra et al. , 2017 ) or by augmenting data using label smoothing ( Müller et al. , 2019 ; Szegedy et al. , 2016 ) . Another way to manipulate model confidence is to tune the temperature of the softmax function , which is otherwise commonly set to one . Adjusting the softmax temperature during training has been shown to be important in metric learning ( Wu et al. , 2018 ; Zhai and Wu , 2019 ) and when performing distillation ( Hinton et al. , 2015 ) ; as well as for post-training calibration of prediction probabilities ( Platt , 2000 ; Guo et al. , 2017 ) . The interplay between temperature , learning , and generalization is complex and not well-understood in the general case . Although significant recent theoretical progress has been made understanding generalization and learning in wide neural networks approximated as linear models , analysis of linearized learning dynamics has largely focused on the case of squared error losses ( Jacot et al. , 2018 ; Du et al. , 2019 ; Lee et al. , 2019 ; Novak et al. , 2019a ; Xiao et al. , 2019 ) . Infinitely-wide networks trained with softmax-cross-entropy loss have been shown to converge to max-margin classifiers in a particular function space norm ( Chizat and Bach , 2020 ) , but timescales of convergence are not known . Additionally , many well-performing models operate best away from the linearized regime ( Novak et al. , 2019a ; Aitchison , 2019 ) . This means that understanding the deviations of models from their linearization around initialization is important for understanding generalization ( Lee et al. , 2019 ; Chizat et al. , 2019 ) . In this paper , we investigate the training of neural networks with softmax-cross-entropy losses . In general this problem is analytically intractable ; to make progress we pursue a strategy that combines analytic insights at short times with a comprehensive set of experiments that capture the entirety of training . At short times , models can be understood in terms of a linearization about their initial parameters along with nonlinear corrections . In the linear regime we find that networks trained with different inverse-temperatures , β = 1/T , behave identically provided the learning rate is scaled as η̃ = ηβ2 . Here , networks begin to learn over a timescale τz ∼ ‖Z0‖2/η̃ where Z0 are the initial logits of the network after being multiplied by β . This implies that we expect learning to begin faster for networks with smaller logits . The learning dynamics begin to become nonlinear over another , independent , timescale τnl ∼ β/η̃ , suggesting more nonlinear learning for small β . From previous results we expect that neural networks will perform best in this regime where they quickly exit the linear regime ( Chizat et al. , 2019 ; Lee et al. , 2020 ; Lewkowycz et al. , 2020 ) . We combine these analytic results with extensive experiments on competitive neural networks across a range of architectures and domains including : Wide Residual networks ( Zagoruyko and Komodakis , 2017 ) on CIFAR10 ( Krizhevsky , 2009 ) , ResNet-50 ( He et al. , 2016 ) on ImageNet ( Deng et al. , 2009 ) , and GRUs ( Chung et al. , 2014 ) on the IMDB sentiment analysis task ( Maas et al. , 2011 ) . In the case of residual networks , we consider architectures with and without batch normalization , which can appreciably change the learning dynamics ( Ioffe and Szegedy , 2015 ) . For all models studied , we find that generalization performance is poor at ‖Z0‖2 1 but otherwise largely independent of ‖Z0‖2 . Moreover , learning becomes slower and less stable at very small β ; indeed , the optimal learning rate scales like η∗ ∼ 1/β and the resulting early learning timescale can be written as τ∗z ∼ ‖Z0‖2/β . For all models studied , we observe strong performance for β ∈ [ 10−2 , 101 ] although the specific optimal β is architecture dependent . Emphatically , the optimal β is often far from 1 . For models without batch normalization , smaller β can give stronger results on some training runs , with others failing to train due to instability . Overall , these results suggest that model performance can often be improved by tuning β over the range of [ 10−2 , 101 ] . 2 THEORY . We begin with a precise description of the problem setting before discussing a theory of learning at short times . We will show the following : • The inverse temperature β and logit scale ‖Z0‖ control timescales which determine the rate of change of the loss , the relative change in logits , and the time for learning to leave the linear learning regime . • Small β causes training to access the non-linear learning regime . We will see empirically that increasing access to the non-linear regime can improve generalization . • The largest allowable learning rate is set by the timescale to leave the linearized learning regime , which suggests that networks with small β will train more slowly . All numerical results in this section are using a Wide Resnet ( Zagoruyko and Komodakis , 2017 ) trained on CIFAR10 . 2.1 BASIC MODEL AND NOTATION . We consider a classification task with K classes . For an N dimensional input x , let z ( x , θ ) be the pre-softmax output of a classification model parameterized by θ ∈ RP , such that the classifier predicts the class i corresponding to the largest output value zi . We will mainly consider θ trained by SGD on a training set ( X , Y ) of M input-label pairs . We focus on models trained with cross-entropy loss with a non-trivial inverse temperature β . The softmax-cross-entropy loss can be written as L ( θ , X , Y ) = K∑ i=1 Yi · ln ( σ ( βzi ( X , θ ) ) ) = K∑ i=1 Yi · ln ( σ ( Zi ( X , θ ) ) ) ( 1 ) where we define the rescaled logits Z = βz and σ ( Z ) i = eZi/ ∑ j e Zj is the softmax function . Here Z ( X , θ ) is the M ×K dimensional matrix of rescaled logits on the training set . As we will see later , the statistics of individual σ ( Z ) i will have a strong influence on the learning dynamics . While the statistics of σ ( Z ) i are intractable for intermediate magnitudes , ‖Z‖2 , they can be understood in the limits of large and small ‖Z‖2 . For a fixed model z ( x , θ ) , β controls the certainty of the predicted probabilities . Values of β such that β 1/‖z‖2 will give small values of ‖Z‖2 1 , and the outputs of the softmax will be close to 1/K independent of i ( the maximum-entropy distribution on K classes ) . Larger values of β such that β 1/‖z‖2 will lead to large values of ‖Z‖2 1 ; the resulting distribution has probability close to 1 on one label , and ( exponentially ) close to 0 on the others . The continuous time learning dynamics ( exact in the limit of small learning rate ) are given by : θ̇ = ηβ K∑ i=1 ( ∂zi ( X , θ ( t ) ) ∂θ ) T ( Yi − σ ( Zi ( X , θ ( t ) ) ) ( 2 ) for learning rate η . We will drop the explicit dependence of Zi on θ from here onward , and we will denote time dependence as Zi ( X , t ) explicitly where needed . In function space , the dynamics of the model outputs on an input x are given by dzi ( x ) dt = ηβ K∑ j=1 ( Θ̂θ ) ij ( x , X ) ( Yj − σ ( Zj ( X ) ) ) ( 3 ) where we define the M ×M ×K ×K dimensional tensor Θ̂θ , the empirical neural tangent kernel ( NTK ) , as ( Θ̂θ ) ij ( x , X ′ ) ≡ ∂zi ( x ) ∂θ ( ∂zj ( X ) ∂θ ) T ( 4 ) for class indices i and j , which is block-diagonal in the infinite-width limit . From Equation 3 , we see that the early-time dynamics in function space depend on β , the initial softmax input Z ( X , 0 ) on the training set , and the initial Θ̂θ . Changing these observables across a model family will lead to different learning trajectories early in learning . Since significant work has already studied the effects of the NTK , here we focus on the effects of changing β and ‖Z0‖F ≡ ‖Z ( X , 0 ) ‖F ( the norm of the M ×K dimensional matrix of training logits ) , independent of Θ̂θ . 2.2 LINEARIZED DYNAMICS . For small changes in θ , the tangent kernel is approximately constant throughout learning ( Jacot et al. , 2018 ) , and we drop the explicit θ dependence in this subsection . The linearized dynamics of z ( x , t ) only depend on the initial value of Θ̂ and the β-scaled logit values Z ( X , t ) . This suggests that there is a universal timescale across β and η which can be used to compare linearized trajectories with different parameter values . Indeed , if we define an effective learning rate η̃ ≡ ηβ2 , we have dZi ( x ) dt = η̃ K∑ j=1 ( Θ̂ ) ij ( x , X ) ( Yj − σ ( Zj ( X ) ) ) ( 5 ) which removes explicit β dependence of the dynamics . We note that a similar rescaling exists for the continuous time versions of other optimizers like momentum ( Appendix B ) . The effective learning rate η̃ is useful for understanding the nonlinear dynamics , as plotting learning curves versus η̃t causes early-time collapse for fixed Z0 across β and η ( Figure 1 ) . We see that there is a strong , monotonic , dependence of the time at which the nonlinear model begins to deviate from its linearization on β . We will return to and explain this phenomenon in Section 2.4 . Unless otherwise noted , we will analyze all timescales in units of η̃ instead of η , as it will allow for the appropriate early-time comparisons between models with different β . 2.3 EARLY LEARNING TIMESCALE . We now define and compute the early learning timescale , τz , that measures the time it takes for the logits to change significantly from their initial value . Specifically , we define τz such that for t τz we expect ‖Z ( x , t ) −Z ( x , 0 ) ‖F ‖Z ( x , 0 ) ‖F and for t τz , ‖Z ( x , t ) −Z ( x , 0 ) ‖F ∼ ‖Z ( x , 0 ) ‖F ( or larger ) . This is synonymous with the timescale over which the model begins to learn . As we will show below , τz ∝ ‖Z0‖F /η̃ . Therefore in units of η̃ , τz only depends on ‖Z0‖F and not β . To see this , note that at very short times it follows from Equation 5 that Zi ( x , t ) − Zi ( x , 0 ) ≈ η̃ K∑ j=1 ( Θ̂ ) ij ( x , X ) ( Yj − σ ( Zj ( X ) ) ) t+O ( t2 ) ( 6 ) It follows that we can define a timescale over which the logits ( on the training set ) change appreciably from their initial value as τz ≡ 1 η̃ ‖Z0‖F ‖Θ̂ ( X , X ) ( Y − σ ( Z0 ) ) ‖F . ( 7 ) where the norms are once again taken across all classes as well as training points . This definition has the desired properties for t τz and t τz . In units of η̃ , τz depends only on ‖Z0‖F , in two ways . The first is a linear scaling in ‖Z0‖F ; the second comes from the contribution from the gradient ‖Θ̂ ( X , X ) ( Y−σ ( Z ( X , 0 ) ) ) ‖F . As previously discussed , since σ ( Z0 ) saturates at small and large values of ‖Z0‖F , it follows that the gradient term will also saturate for large and small ‖Z0‖F , and the ratio of saturating values is some O ( 1 ) constant independent of ‖Z0‖F and β . The quantitative and conceptual nature of τz can both be confirmed numerically . We compute τz explicitly as || [ Zi ( x , t ) − Zi ( x , 0 ) ] /t|| for short times ( Figure 2 , left ) . When plotted over a wide range of ‖Z0‖F and β , the ratio τz/‖Z0‖F ( in rescaled time units ) undergoes a saturating , O ( 1 ) variation from small to large ‖Z0‖F . The quantitative dependence of the transition on the NTK is confirmed in Appendix C. Additionally , for fixed β and varying ‖Z0‖F , rescaling time by 1/τz causes accuracy curves to collapse at early times ( Figure 2 , middle ) , even if they are very different at early times without the rescaling ( right ) . We note here that the late time accuracy curves seem similar across ‖Z0‖F without rescaling , a point which we will return to in Section 3.2 .

Summary: This paper studies how temperature scaling affects training dynamics of neural networks (with softmax layer and cross-entropy loss). The theoretical analysis shows that neural networks trained with smaller inverse temperatures (beta) exit the linear regime faster, which implies better performance. Experiments on image classification and sentiment analysis confirm that tuning temperature improves neural network generalization, even for state-of-the-art models.