TEMPERATURE CHECK: THEORY AND PRACTICE FOR TRAINING MODELS WITH SOFTMAX-CROSS-ENTROPY LOSSES
Year: 2020

Abstract:
The softmax function combined with a cross-entropy loss is a principled approach to modeling probability distributions that has become ubiquitous in deep learning. The softmax function is defined by a lone hyperparameter, the temperature, that is commonly set to one or regarded as a way to tune model confidence after training; however, less is known about how the temperature impacts training dynamics or generalization performance. In this work we develop a theory of early learning for models trained with softmax-cross-entropy loss and show that the learning dynamics depend crucially on the inverse-temperature β as well as the magnitude of the logits at initialization, ||βz||2. We follow up these analytic results with a large-scale empirical study of a variety of model architectures trained on CIFAR10, ImageNet, and IMDB sentiment analysis. We find that generalization performance depends strongly on the temperature, but only weakly on the initial logit magnitude. We provide evidence that the dependence of generalization on β is not due to changes in model confidence, but is a dynamical phenomenon. It follows that the addition of β as a tunable hyperparameter is key to maximizing model performance. Although we find the optimal β to be sensitive to the architecture, our results suggest that tuning β over the range 10−2 to 10 improves performance over all architectures studied. We find that smaller β may lead to better peak performance at the cost of learning stability.

heading: 1 INTRODUCTION
text: 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].

heading: 2 THEORY
text: 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.

heading: 2.1 BASIC MODEL AND NOTATION
text: 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 Θ̂θ.

heading: 2.2 LINEARIZED DYNAMICS
text: 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 β.

heading: 2.3 EARLY LEARNING TIMESCALE
text: 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.

heading: 2.4 NONLINEAR TIMESCALE
text: While linearized dynamics are useful to understand some features of learning, the best performing networks often reside in the nonlinear regime (Novak et al., 2019a). Here we define the nonlinear timescale, τnl, corresponding to the time over which the network deviates appreciably from the linearized equations. We will show that τnl ∝ β/η̃. Therefore, in terms of β and ‖Z0‖F , networks with small β will access the nonlinear regime early in learning, while networks with large β will be effectively linearized throughout training. We note that a similar point was raised in Chizat et al. (2019), primarily in the context of MSE loss.
We define τnl to be the timescale over which the change in Θ̂θ (which contributes to the second order term in Equation 6) can no longer be neglected. Examining the second time derivative of Z, we have
d2Zi dt2 = η̃ K∑ j=1 −(Θ̂θ)ij(x,X ) ddtσ(Z(X ))︸ ︷︷ ︸ linearized dynamics + d dt [ (Θ̂θ)ij(x,X ) ] (Yj − σ(Zj(X ))︸ ︷︷ ︸
nonlinearized dynamics≡ Z̈nl  (8) The first term is the second derivative under a fixed kernel, while the second term is due to the change in the kernel (neglected in the linearized limit). A direct calculation shows that the second term, which we denote Z̈nl, can be written as
(Z̈nl)i = β −1η̃2 K∑ j=1 K∑ k=1 (Yk(X )−σ(Zk(X ))T ( ∂zk(X ) ∂θ · ∂ ∂θ [Θ̂θ]ij ) (Yj(X )−σ(Zj(X ))) (9)
This gives us a nonlinear timescale τnl defined, at initialization, by τnl ≡ ‖Ż(X , 0)‖F /‖Z̈nl(X , 0)‖F . We can interpret τnl as the time it takes for changes in the kernel to contribute to learning.
Though computing ‖Z̈nl(X , 0)‖F in exactly is analytically intractable, its basic scaling in terms of β and ‖Z0‖F (and therefore, that of τnl) is computable. We first note the explicit β−1η̃2 dependence. The remaining terms are independent of β and vary by at mostO(1) with ‖Z0‖F ; indeed as described above, ‖Y(X ) − σ(Z(X , 0))‖F saturates for large and small ‖Z0‖F . Morevoer, the derivative, ∂z(X ,0) ∂θ , is the square root of the NTK and, at initialization, it is independent of ‖Z
0‖F . Together with our analysis of τz we have that, up to some O(1) dependence on ‖Z0‖F , τnl ∝ β/η̃. Therefore, the degree of nonlinearity early in learning is controlled via β alone.
Once again we can confirm the quantitative and conceptual understanding of τnl numerically. Qualitatively, we see that for fixed ‖Z0‖F , models with smaller β deviate sooner from the linearized dynamics when learning curves are plotted against η̃t (Figure 1). We compute τnl explicitly by taking ‖[Z(t)−Zlin(t)]/t2‖ as small times, where Zlin(t) is the solution to Equation 5. τnl/β has an O(1) dependence on ‖Z0‖F only (Figure 3).

heading: 2.5 LEARNING RATES AND LEARNING SPEED
text: The timescales, τz and τnl, can be combined to gain information about training with SGD. Indeed, the largest allowable learning rate is controlled by the curvature of the loss function Du et al. (2019); Allen-Zhu et al. (2019). A necessary condition for training is that the curvature be small compared with the step size, which happens when τnl ≤ c, where c is an architecture-dependent constant of O(1) which has not yet been calculated theoretically (see Lewkowycz et al. (2020) for empirical calculations of c). This predicts a maximum effective learning rate η̃ of O(β) which in turn implies a raw learning rate of O(β−1) (as η = η̃/β2).
However, if η̃ is O(β), then τz is O(β−1). This means that, for early learning, networks with smaller β will take more SGD steps to reach an appreciable change in logits. Therefore, we predict that networks with small β take longer to train (which we see in practice as well).

heading: 3 EXPERIMENTAL RESULTS

heading: 3.1 OPTIMAL LEARNING RATE
text: We begin our empirical investigation by training wide resnets (Szegedy et al., 2016) without batch normalization on CIFAR10, as this architecture is well within the regime where our theory applies. In order to understand the effects of the different timescales on learning, we control β and ‖Z0‖F independently by using a correlated initialization strategy outlined in Appendix D.1.
Before considering model performance, it is first useful to understand the scaling of the learning rate with β. We define the optimal learning rate η∗ as the learning rate with the best generalization performance. To do this, we initialize networks with different β and conduct learning rate sweeps for each β. The optimal learning rate η∗ has a clear 1/β dependence (Figure 4). This matches the prediction in Section 2.5, suggesting the maximum learning rate corresponds to the regime where the non-linear effects become important at the fastest rate for which training still converges. Again, as predicted, networks with smaller β also learn more slowly in terms of number of SGD steps.
Thus, at small β we expect learning to take place slowly and nonlinear effects to become important by the time the function has changed appreciably. At large β, by contrast, our results suggest that the network will have learned a significant amount before the dynamics become appreciably nonlinear.

heading: 3.2 PHASE PLANE
text: In the preceding discussion two quantities emerged that control the behavior of early-time dynamics: the inverse-temperature, β, and the rescaled logits ‖Z0‖F . In attempting to understand the behavior of real neural networks trained using softmax-cross-entropy loss, it therefore makes sense to try to reason about this behavior by considering neural networks that span the β − ‖Z0‖F phase plane, the space of allowable pairs (β, ‖Z0‖F ). By construction, the phase plane is characterized by the timescales involved in early learning. To summarize, τz ∼ ‖Z0‖F /η̃ sets the timescale for early learning, with larger values of ‖Z0‖F leading to longer time before significant accuracy gains are made (Section 2.3). Meanwhile, τnl ∼ β/η̃ controls the timescale for learning dynamics to leave the linearized regime - with small β leading to immediate departures from linearity, while models with large β may stay linearized throughout their learning trajectories (Section 2.4).
In Figure 5 (a), we show a schematic of the phase plane. The colormap shows the test performance of a wide residual network (Zagoruyko and Komodakis, 2017), without batch normalization, trained on CIFAR10 in different parts of the phase plane. The value of β makes a large difference in generalization, with optimal performance achieved at β ≈ 10−2. In general, larger β performed worse than small β as expected. Moreover, we observe similar generalization for all sufficiently large β; this is to be expected since models in this regime are close to their linearization throughout training (see Figure 1) and we expect the linearized models to have β-independent performance. Generalization was largely insensitive to ‖Z0‖F so long as the network was sufficiently well-conditioned to be trainable. This suggests that long term learning is insensitive to τz .
In Figure 5 (b), we plot the accuracy after 20 steps of optimization (with the optimal learning rate). For fixed ‖Z0‖F , the training speed was slow for the smallest β and then became faster with increasing β. For fixed β the training speed was fastest for small ‖Z0‖F and slowed as ‖Z0‖F increased. Both these phenomena were predicted by our theory and shows that both parameters are important in determining the early-time dynamics. However, we note that the relative accuracy across the phase plane at early times did not correlate with the performance at late times.
This highlights that differences in generalization are a dynamical phenomenon. Another indication of this fact is that at the end of training, at time tf , the final training set logit values ‖Zf‖F ≡ ‖Z(X , tf )‖F tend towards 1 independent of the initial β and ‖Z0‖F (Figure 5, (c)). With the exception of the poorly-performing large ‖Z0‖F regime, the different models reach similar levels of certainty by the end of training, despite having different generalization performances. Therefore generalization is not well correlated with the final model certainty (a typical motivation for tuning β).
3.3 ARCHITECTURE DEPENDENCE OF THE OPTIMAL β
Having demonstrated that β controls the generalization performance of neural networks with softmaxcross-entropy loss, we now discuss the question of choosing the optimal β. Here we investigate this question through the lens of a number of different architectures. We find the optimal choice of β to be strongly architecture dependent. Whether or not the optimal β can be predicted analytically is an open question that we leave for future work. Nonetheless, we show that all architectures considered display optimal β between approximately 10−2 and 101. We observe that by taking the time to tune β it is often the case that performance can be improved over the naive setting of β = 1.

heading: 3.3.1 WIDE RESNET ON CIFAR10
text: In Figure 6 (a) we show the accuracy against β for several wide residual networks whose weights are drawn from normal distributions of different variances, σ2w, trained without batchnorm, as well as a network with σ2w = 1 trained with batchnorm (averaged over 10 seeds). The best average performance is attained for β < 1, σw = 1 without batchnorm, and in particular networks with large σw are dramatically improved with β tuning. The network with batchnorm is better at all β, with optimal β ≈ 10. However, we see that the best performing seed is often at a lower β (Figure 6 (b)), with larger σw networks competitive with σw = 1, and even with batchnorm at fixed β (though batchnorm with β = 10 still performs the best). This suggests that small β can improve best case performance, at the cost of stability. Our results emphasize the importance of tuning β, especially for models that have not otherwise been optimized.

heading: 3.3.2 RESNET50 ON IMAGENET
text: Motivated by our results on CIFAR10, we experimentally explored the effects of β as a tunable hyperparameter for ResNet-50 trained on Imagenet. We follow the experimental protocol established by (Ghiasi et al., 2018). A key difference between this procedure and standard training is that we train for substantially longer: the number of training epochs is increased from 90 to 270. Ghiasi et al. (2018) found that this longer training regimen was beneficial when using additional regularization. Table 1 shows that scaling β improves accuracy for ResNet-50 with batchnorm. However, we did not
find that using β < 1 was optimal for ResNet-50 without normalization. This further emphasizes the subtle architecture dependence that warrants further study.

heading: 3.3.3 GRUS ON IMDB SENTIMENT ANALYSIS
text: To further explore the architecture dependence of optimal β, we train GRUs (from Maheswaranathan et al. (2019)) whose weights are drawn from two different distributions on an IMDB sentiment analysis task that has been widely studied (Maas et al., 2011). We plot the results in Figure 6 (c) and observe that the results look qualitatively similar to the results on CIFAR10 without batch normalization. We observe a peak performance near β ∼ 1 averaged over an ensemble of networks, but we observe that smaller β can give better optimal performance at the expense of stability.

heading: 3.4 PROPOSED TUNING PROCEDURE
text: Our results suggest the following tuning procedure for networks trained with SGD/momentum:
• Train/tune a model as normal. Note the optimal learning rate η0. • For the best parameter set, sweep over β ∈ [10−2, 101], scaling learning rate as η0/β. • If best performing β is at an endpoint of the range, continue tuning.
If there is sufficient compute, the β search can instead be folded into the overall hyperparameter tuning.
We note that our observation of optimal β < 1 in classification settings stands in contrast to the observation of optimal β > 1 in the Bayesian inference setting (Wenzel et al., 2020). We speculate that the differences are related to the fact that in Bayesian inference, the uncertainty of estimates are important, and that the objective (learn an SDE which reproduces a target distribution) may have qualitatively different properties than classification tasks.

heading: 4 CONCLUSIONS
text: Our empirical results show that tuning β can yield sometimes significant improvements to model performance. Perhaps most surprisingly, we observe gains on ImageNet even with the highlyoptimized ResNet50 model. Our results on CIFAR10 suggest that the effect of β may be even stronger in networks which are not yet highly-optimized, and results on IMDB show that this effect holds beyond the image classification setting. It is possible that even more gains can be made by more carefully tuning β jointly with other hyperparameters, in particular the learning rate schedule and batch size.
One key lesson of our theoretical work is that properties of learning dynamics must be compared using the right units. For example, τnl ∝ 1/βη, which at first glance suggests that models with smaller β will become nonlinear more slowly than their large β counterparts. However, analyzing τnl with respect to the effective learning rate η̃ = β2η yields τnl ∝ β/η̃. Thus we see that, in fact, networks with smaller β tend to become more non-linearized before much learning has occurred, compared to networks with large β which can remain in the linearized regime throughout training. Our numerical results confirm this intuition developed using the theoretical analysis.
As discussed above, our analysis suggests a range of good β, but does not predict the optimal value. Architecture-dependent, non-scaling multiplicative factors to key learning parameters have been observed in other contexts (Lewkowycz et al., 2020), and their numerical estimation is in general difficult. Extending the theoretical results to make predictions about these quantities is an interesting avenue for future work. Another area that warrants further study is the instability in training at small β.

heading: A LINEARIZED LEARNING DYNAMICS

heading: A.1 FIXED POINTS
text: For the linearized learning dynamics, the trajectory z(x, t) can be written in terms of the trajectories of the training set as
z(x, t)− z(x, 0) = Θ̂(x,X )Θ̂+(X ,X )(z(X , t)− z(X , 0)) (10) where + is the pseudo-inverse. Therefore, if one can solve for z(X , t), then in principle properties of generalization are computable.
However, in general Equation 3 does not admit an analytic solution even for fixed Θ̂, in contrast to the case of mean squared loss. It not even have an equilibrium - if the model can achieve perfect training accuracy, the logits will grow indefinitely. However, there is a guaranteed fixed point if the appropriate L2 regularization is added to the training objective. Given a regularizer 12λθ‖δθ‖
2 on the change in parameters δθ = θ(t)− θ(0), the dynamics in the linearized regime are given by
ż(x) = βηΘ̂(x,X )(Y(X )− σ(βz(X )))− λθδz(x) (11)
where the last term comes from the fact that ∂z∂θ δθ = δz(x) in the linearized limit.
We can write down self-consistent equations for equilibria, which are approximately solvable in certain limits. For an arbitrary input x, the equilibrium solution z∗(x) is
0 = βΘ̂(x,X )(Y(X )− σ(βz∗(X )))− λθδz∗(x) (12) This can be rewritten in terms of the training set as
δz∗(x) = Θ̂(x,X )Θ̂+(X ,X )z∗(X ) (13) similar to kernel learning.
It remains then to solve for z∗(X ). We have:
δz∗(X ) = β λθ Θ̂(X ,X )[Y(X )− σ(βz∗(X ))] (14)
We immediately note that the solution depends on the initialization. We assume z(x, 0) = 0, so δz = z in order to simplify the analysis. The easiest case to analyze is when ‖βz∗(X )‖F 1. Then we have:
z∗(X ) = β λθ
Θ̂(X ,X ) [ Y(X )− 1
K (1 + βz∗(X ))
] (15)
which gives us
z∗(X ) = β λθ
[ 1 + β
Kλθ Θ̂(X ,X )
]−1 Θ̂(X ,X )(Y(X )− 1/K) (16)
Therefore the self-consistency condition for this solution is ‖ βλθ Θ̂‖F 1, which simplifies the solution to
z∗(X ) = β λθ Θ̂(X ,X )(Y(X )− 1/K) (17)
This is equivalent to the solution after a single step of (full-batch) SGD with appropriate learning rate. We note that unlike linearized dynamics with L2 loss and a full-rank kernel, there is no guarantee that the solution converges to 0 training error.
The other natural limit is ‖βz∗(X )‖2 1. We focus on the 2 class case, in order to take advantage of the conserved quantity of learning with cross-entropy loss. We note that the vector on the right hand side of Equation 3 sums to 1 for every training point. Suppose at initialization, Θ̂θ has no logit-logit interactions, as is the case for most architectures in the infinite width limit with random initialization. More formally, we can write Θ̂θ = IdK×K ⊗ Θ̂x where Θ̂x is M ×M . Then, the sum of the logits for any input x is conserved during linearized training, as we have:
1Tż(x) = ηβ1T [ IdK×K ⊗ Θ̂x ] (Y − σ(βz(X ))) (18)
Multiplying the right hand side through, we get 1Tż(x) = ηβΘ̂x [ 1T(Y − σ(βz(X ))) ] = 0 (19)
(Note that if Θ̂θ has explicit dependence on the logits, there still is a conserved quantity, which is more complicated to compute.)
Now we can analyze ‖βz∗(X )‖F 1. With two classes, and z(X ) = 0 at initialization, we have z∗1 = −z∗2. Therefore, without loss of generality, we focus on z∗1, the logit of the first class. In this limit, the leading order correction to the softmax is approximately:
σ(βz∗1) ≈ 1z∗1>0 − sign(z ∗ 1)e −2β|z∗1 | (20)
The self-consistency equation is then:
z∗1(X ) = β
λθ Θ̂(X ,X )
[ Y(X )− 1z∗1>0 + sign(z ∗ 1)e −2β|z∗1 | ] (21)
The vector on the right hand side has entries that are O(e−2β|z ∗ 1 |) for correct classifications, and O(1) for incorrect ones. If we assume that the training error is 0, then we have:
z∗1(X ) = β
λθ Θ̂(X ,X )sign(z∗1)e−2β|z ∗ 1 | (22)
This is still non-trivial to solve, but we see that the self consistency condition is that ln(β||Θ̂||F /λθ) 1. Here also it may be difficult to train and generalize well. The individual elements of the right-handside vector are broadly distributed due to the exponential - so the outputs of the model are sensitive to/may only depend on a small number of datapoints. Even if the equilibrium solution has no training loss, generalization error may be high for the same reasons.
This suggests that even for NTK learning (with L2 regularization), the scale of ||βz|| plays an important role in both good training accuracy and good generalization. In the NTK regime, there is one unique solution so (in the continuous time limit) the initialization doesn’t matter; rather, the ratio of β and λθ (compared to the appropriate norm of Θ̂) needs to be balanced to prevent falling into the small βz regime (where training error might be large) or the large βz regime (where a few datapoints might dominate and reduce generalization).

heading: A.2 DYNAMICS NEAR EQUILIBRIUM
text: The dynamics near the equilibrium can be analyzed by expanding around the fixed point equation. We focus on the dynamics on the training set. The dynamics of the difference z̃(X ) = z(X )− z∗(X ) for small perturbations is given by
˙̃z(X ) = −η [ β2[Idz ⊗ Θ̂(X ,X )]σz(βz∗(X )) + λθ ] z̃(X ) (23)
where σz is the derivative of the softmax matrix
σz(z) ≡ ∂σ(z)
∂z′ = diag(σ(z))− σ(z)σ(z′)T (24)
We can perform some analysis in the large and small β cases (once again ignoring λz). For small ‖βz∗(X )‖F , we have ‖ βλθ Θ̂‖F 1 which leads to:
σz(βz ∗(X )) = (1/K − 11T/K2) +O(βΘ̂) (25)
This matrix has K − 1 eigenvalues with value 1/K, and one zero eigenvalue (corresponding to the conservation of probability). Therefore ‖β2[Idz ⊗ Θ̂(X ,X )]σz(βz∗(X ))‖F λθ, and the well-conditioned regularizer dominates the approach to equilibrium.
In the large β case (ln(β||Θ̂||/λθ) 1), the values of σ(βz(X )) are exponentially close to 0 (K − 1 values) or 1 (the value corresponding to the largest logit). This means that σz(βz(X )) has exponentially small values in ‖βz(X )‖F - if any one of σ(βzi(X )) and σ(βzj(X )) is exponentially
small, the corresponding element of σz(βz(X )) is as well; for the largest logit i the diagonal is σ(βzi(X ))(1− σ(βzi(X ))) which is also exponentially small.
From Equation 22, we have λθ β2e2β|z ∗ 1 |; therefore, though the σz term of H is exponentially small, it dominates the linearized dynamics near the fixed point, and the approach to equilibrium is slow. We will analyze the conditioning of the dynamics in the remainder of this section.

heading: A.3 CONDITIONING OF DYNAMICS
text: Understanding the conditioning of the linearized dynamics requires understanding the spectrum of the Hessian matrix H = ( Idz ⊗ Θ̂(X ,X ) ) σz(βz
∗(X )). In the limit of large model size, the first factor is block-diagonal with training set by training set blocks (no logit-logit interactions), and the second term is block-diagonal with K ×K blocks (no datapoint-datapoint interactions). We will use the following lemma to get bounds on the conditioning:
Lemma: Let M = AB be a matrix that is the product of two matrices. The condition number κ(M) ≡ λM,maxλM,min has bound
κ(B)/κ(A) ≤ κ(M) ≤ κ(A)κ(B) (26)
Proof: Consider the vector v that is the eigenvector of B associated with λB,min. Note that ||Av||/||v|| ≤ λA,max. Analogously, for w, the eigenvector associated with λB,max, ||Aw||/||w|| ≥ λA,min. This gives us the two bounds:
λM,min ≤ λA,maxλB,min, λM,max ≥ λA,minλB,max (27)
This means that the condition number κ(H) ≡ λM,maxλM,min is bounded by
κ(M) ≥ λA,maxλB,min λA,minλB,max = κ(B)/κ(A) (28)
In total, we have the bound of Equation 26, where the upper bound is trivial to prove.
In particular, this means that a poorly conditioned σz(βz∗(X )) will lead to poor conditioning of the linearized dynamics if the NTK Θ̂(X ,X ) is (relatively) well conditioned. This bound will be important in establishing the poor conditioning of the linearized dynamics for the large logit regime ||βz|| 1.

heading: A.3.1 SMALL LOGIT CONDITIONING
text: For ‖βz∗(X )‖F 1, the Hessian H is
H = 1
K
( 1− 1
K 11T
) ⊗ Θ̂(X ,X ) (29)
Since H is the Kroenecker product of two matrices, the condition numbers multiply, and we have
κ(H) = κ(Θ̂) (30)
which is well-conditioned so long as the NTK is. Regardless, the well-conditioned regularization due to λθ dominates the approach to equilibrium.

heading: A.3.2 LARGE LOGIT CONDITIONING
text: Now consider ‖βz∗(X )‖F 1. Here we will show that the linearized dynamics is poorly conditioned, and that κ(H) is exponentially large in β.
We first try to understand σz(βz∗(x)) for an individual x ∈ X . To 0th order (in an as-of-yetundefined expansion), σz is zero - at large temperature the softmax returns either 0 or 1, which by Equation 24 gives 0 in all entries. The size of the corrections end up being exponentially dependent on ‖βz∗‖F ; the entries will have broad, log-normal type distributions with magnitudes which scale
as exp(−β|z∗1|). There will be two scaling regimes one with a small number of labels in the sense√ β √ ln(K), where the largest logit dominates the statistics, and one where the number of labels is large (and the central limit theorem applies to the partition function). In both cases, however, there is still exponential dependence on β; we will focus on the first which is easier to analyze and more realistic (e.g. for 106 labels “large” β is only ∼ 15). Let z1 be the largest of K logits, z2 the second largest, and so on. Then using Equation 24 we have:
(σz)i1 = −e−β(z1−zi) (31)
for i 6= 1, (σz)ij = δije
−β(z1−zi) − e−β(2z1−zi−zj) (32) for i 6= j and
(σz)11 = e −β(z1−z2) (33)
The eigenvectors and eigenvalues can be approximately computed as:
(v1)1 = 1√ 2 , (v1)2 = − 1√ 2 , λ1 = 2e −β(z1−z2) (34)
(v2)1 = 1√ 2 , (v2)2 = 1√ 2 , λ2 = − 1 2 e−2β(z1−z2) (35)
and for i > 2, (vi)i = 1, (vi)1 = e −β(z2−zi), λi = e −β(z1−zi) (36)
with all non-explicit eigenvector components 0. This expansion is valid provided that β/K 1 (so that eβ(z1−zi) eβ(z1−zi+1)). Therefore the spectrum of any individual block σz(βz(x)) is exponentially small in β. Using the bound in the Lemma, we have:
κ(H) ≥ eO(β|z ∗ 1 |)/κ(Θ̂(X ,X )) (37)
This is a very loose bound, as it assumes that the largest eigendirections of σz are aligned with the smallest eigendirections of Idz ⊗ Θ̂, and vice versa. It is possible κ(H) is closer in magnitude to the upper bound eβ(z2−zK)κ(Θ̂(X ,X )). Regardless, κ(H) is exponentially large in β - meaning that the conditioning is exponentially poor for large ‖βz∗‖F .

heading: B SGD AND MOMENTUM RESCALINGS

heading: B.1 DISCRETE EQUATIONS
text: Consider full-batch SGD training. The update equations for the parameters θ are:
θt+1 = θt − η∇θL (38)
We will denote gt ≡ ∇θL for ease of notation. Training with momentum, the equations of motion are given by:
vt+1 = (1− γ)vt − gt (39)
θt+1 = θt + ηvt+1 (40)
where γ ∈ [0, 1]. One key point to consider later will be the relative magnitude ∆θ of updates to the parameters. For SGD, the magnitude of updates is η||g||. For momentum with slowly-varying gradients the magnitude is η||g||/γ.

heading: B.2 CONTINUOUS TIME EQUATIONS
text: We can write down the continuous time version of the learning dynamics as follows. For SGD, for small learning rates we have:
dθ dt = −ηg (41)
For the momentum equations we have
dv dt = −γv − g (42)
dθ dt = ηv (43)
From these equations, we can see that in the continuous time limit, there are coordinate transformations which can be used to cause sets of trajectories with different parameters to collapse to a single trajectory. SGD is the simplest, where rescaling time to τ ≡ ηt causes learning curves to be identical for all learning rates.
For momentum, instead of a single universal learning curve, there is a one-parameter family of curves controlled by the ratio Tmom ≡ η/γ2. Consider rescaling time to τ = at and ν = bv, where a and b will be chosen to put the equations in a canonical form. In our new coordinates, we have
dν dτ = −(γ/a)ν − (b/a)g (44)
dθ dτ = ην/(ab) (45)
The canonical form we choose is dν
dτ = −λν − g (46)
dθ dτ = ν (47)
From which we arrive at a = b = √ η, which gives us λ = γ/ √ η.
Note that this is not a unique canonical form; for example, if we fix a coefficient of −1 on ν, we end up with
dν dτ = −ν − (η/γ2)g (48)
dθ dτ = ν (49)
with a = γ. This is a different time rescaling, but still controlled by Tmom.
Working in the canonical form of Equations 46 and 47, we can analyze the dynamics. One immediate question is the difference between λ 1 and λ 1. We note that the integral equation
ν(τ) = ν(0) + ∫ τ 0 e−λ(τ−τ ′)g(τ ′)dτ ′ (50)
solves the differential equation for ν. Therefore, for λ 1, ν(t) only depends on the current value g(t) and we have ν(τ) ≈ g(τ)/λ. Therefore, we have, approximately:
dθ dτ ≈ 1 λ g (51)
This means that for large λ all the curves will approximately collapse, with timescale given by√ ηλ−1 = γη (dynamics similar to SGD).
For λ 1, the momentum is essentially the integrated gradient across all time. If ν(0) = 0, then we have
dθ dτ ≈ ∫ τ
0
g(τ ′)dτ ′ (52)
In this limit, θ(τ) is the double integral of the gradient with respect to time. Given the form of the controlling parameter Tmom, we can choose to parameterize γ = γ̃ √ η. Under this parameterization, we have Tmom = γ̃2. The dynamical equations then become:
dν dτ = −γ̃ν − g (53)
dθ dτ = ν (54)
which automatically removes explicit dependence on η.
One particular regime of interest is the early-time dynamics, starting from ν(0) = 0. Integrating directly, we have:
θ(τ) = −1 2 gτ2 + 1 6 γ̃gτ3 + . . . (55)
This means that τ alone is the correct timescale for early learning, at least until τ γ̃ ∼ 1 - which in the original parameters corresponds to t ∼ 1/γ (the time it takes for the momentum to be first “fully integrated”).

heading: B.3 DETAILED ANALYSIS OF MOMENTUM TIMESCALES
text: One important subtlety is that 1 is not the correct value to compare λ to. The real timescale involved is the one over which g changes significantly. We can approximate this in the following way. Suppose that there is some relative change ∆θ||θ|| ∼ c of the parameters that leads to an appreciable relative change in g. Then the timescale over which θ changes by that amount is the one we must compare λ to.
We can compute that timescale in the following way. We assume g fixed for what follows. Therefore, Equation 51 approximately holds. The timescale τc of the change is then given by:
∆θ ||θ|| = 1 λ ||g|| ||θ|| τc ∼ c (56)
which gives τc ∼ cλ||θ||/||g|| (57)
In particular, this means that the approximation is good when λτc 1, which gives γ2/η ||g||||θ|| - the former being a function of the dynamical parameters, the latter being a function of the geometry of L with respect to θ. One consequence of this analysis is that if the ||θ|| remains roughly constant, for fixed η and γ, late in learning when the gradients become small the dynamics shifts into the regime where λ is large, and we effectively have SGD.

heading: B.4 CONNECTING DISCRETE AND CONTINUOUS TIME
text: One use for the form of the continuous time rescalings is to use them to compare learning curves for the actual discrete optimization that is performed with different learning rates. For small learning rates, the curves are expected to collapse, while for larger learning rates the deviations from the continuous expectation can be informative.
With momentum, we only have perfect collapse when γ and η are scaled together. However, one typical use case for momentum is to fix the parameter γ, and sweep through different learning rates. With this setup, if g is changing slowly compared to γ (more precisely, γ2/η ||g||/||θ||), as may be the case at later training times, the change in parameters from a single step is ∆θ ∼ (η/γ))||g|| and the rescaling of taking t to ηt (as for SGD) collapses the dynamics. Therefore given a collection of varying η, but fixed γ curves, it is possible to get intermediate and late time dynamics on the same scale.
However, at early times, while the momentum is still getting “up to speed" (i.e. in the first 1/γ steps), the appropriate timescale is η−1/2. Therefore, in order to get learning curves to collapse across
different η at early times, we need to rescale γ with η as implied by Equations 46 and 47. Namely, one must fix γ̃ and rescale γ = γ̃ √ η. We note that, since γ < 1, this gives us a restriction η < γ̃−2 for the maximum learning that can be supported by the rescaled momentum parameter.

heading: B.5 MOMENTUM EQUATIONS WITH SOFTMAX-CROSS-ENTROPY LEARNING
text: For cross-entropy learning with softmax inputs βz, all the scales acquire dependence on β. If we define `z ≡ ∣∣∣∣∣∣ ∂L∂βz ∣∣∣∣∣∣ and gz ≡ ∣∣∣∣ ∂z∂θ ∣∣∣∣F , then we have, approximately, ||g|| ≈ β`zgz . Consider the goal of obtaining identical early-time learning curves for different values of β. (The curves are only globally consistent across β in the linearized regime.) In order to get learning curves to collapse, we want dLdτ to be independent of β in the rescaled time units. We note that the change in the loss function ∆L from a single step of SGD goes as
∆L ∼ ηβ2`2zg2z (58)
This suggests that one way to collapse learning curves is to plot them against the rescaled learning rate η̃t, where η̃ = ηβ2. While hyperparameter tuning across β, one could use η = η̃/β2, sweeping over η̃ in order to easily obtain comparable learning curves.
However, a better goal for a learning rate rescaling is to try and stay within the continuous time limit - that is, to control the change in parameters ∆θ for a single step to be small across β. We have
∆θ ∼ ηβ`zgz (59)
which suggests that maximum allowable learning rates will scale as 1/β. This suggests setting η = η̂β−1, and rescaling time as η̂β in order to best explore the continuous learning dynamics.
We can perform a similar analysis for the momentum optimizers. We begin by analyzing the continuous time equations for the dynamics of the loss. Starting with the rescalings from Equations 53 and 54 we have
dν dτ = −γ̃ν − βg (60)
dθ dτ = ν (61)
dL dτ = βg · ν (62)
where g = ∂L∂βz ∂z ∂θ . Rescaling τ by β gets us:
dν dβτ = − γ̃ β ν − g (63)
dθ
dβτ =
1 β ν (64)
dL dβτ = g · ν (65)
This rescaling causes a collapse of the trajectories of the L at early times if γ̃/β is constant for varying β.
One scheme to arrive at the final canonical form, across β, is by the following definitions of η, γ, and τ :
• η = η̃β−2 • γ ≡ β√ηγ̃ = √ η̃γ̃ • τ ≡ β√ηt = √ η̃t
where curves with fixed γ̃ will collapse. The latter two equations are similar to before, except with η replaced with η̃. The dynamical equations are then:
dν dτ = −γ̃ν − g (66)
dθ dτ = 1 β ν (67)
dL dτ = g · ν (68)
The change in parameters from a single step (assuming constant g and saturation) is
∆θ = ||g|| γ̃β
√ η̂ (69)
If we instead want the change in parameters from a single step to be invariant of β so the continuous time approximation holds, while maintaining collapse of trajectories, we first note that
∆θ ∼ √ η
γ̃ β`zgz (70)
from a single step of the momentum optimizer. To keep ∆θ invariant of β, we can set:
• η = η̂β−1 • γ ≡ β√ηγ̃ = √ η̂γ̃ = β1/2 √ η̂γ̃ • τ ≡ β√ηt = β1/2 √ η̂t
Note that the relationship between γ and γ̃ is the same in both schemes when measured with respect to the raw learning rate η.

heading: C SOFTMAX-CROSS-ENTROPY GRADIENT MAGNITUDE

heading: C.1 MAGNITUDE OF GRADIENTS IN FULLY-CONNECTED NETWORKS
text: The value of τz has nontrivial (but bounded) dependence on ‖Z0‖F via the ‖Θ̂(Y − σ(Z0(X )))‖F term in Equation 7. We can confirm the dependence for highly overparameterized models by using the theoretical Θ̂. In particular, for wide neural networks, the tangent kernel is block-diagonal in the logits, and easily computable.
The numerically computed τz/‖Z0‖F correlates well with ‖Θ̂(Y − σ(Z0(X )))‖−1F for wide (2000 hidden units/layer) fully connected networks (Figure 7). The ratio depends on details like the nonlinearities in the network; for example, Relu units tend to have a larger ratio than Erf units (left and middle). The ratio also depends on the properties of the dataset. For example, the ratio increases on CIFAR10 when the training labels are randomly shuffled (right).
Therefore in general the ratio of τz/‖Z0‖F at large and small ‖Z0‖F depends subtly on the relationship between the NTK and the properties of the data distribution. A full analysis of this relationship is beyond the scope of this work. The exact form of the transition is likely even more sensitive to these properties and is therefore harder to analyze than the ratio alone.

heading: D EXPERIMENTAL DETAILS

heading: D.1 CORRELATED INITIALIZATION
text: In order to avoid confounding effects of changing β and ‖Z0‖F with changes to Θ̂, we use a correlated initialization strategy (similar to (Chizat et al., 2019)) which fixes Θ̂ while allowing for independent variation of β and ‖Z0‖F . Given a network with final hidden layer h(x,θ) and output weights WO, we define a combined network zc(x, θ̃) explicitly as
zc(x,WO,1,WO,2,θ1,θ2) = WO,1h(x,θ1) + WO,2h(x,θ2) (71)
where, at initialization, θ1 = θ2, and the elements of WO,a have statistics
E[(WO,a)ij(WO,b)kl] = { δikδjl for a = b c · δikδjl for a 6= b
(72)
for correlation coefficient c ∈ [−1, 1], where δij is the Kronecker-delta which is 1 is i = j and 0 otherwise. Under this approach, the initial magnitude of the training set logits is given by ‖Z0‖F = β √ (1 + c)‖z0‖F , where ‖z0‖F is the initial magnitude of the logits of the base model. By manipulating β and c, we can independently change β and ‖Z0‖F with the caveat that ‖Z0‖F ≤ √ 2β‖z0‖F since c ≤ 1. It follows that the small β, large ‖Z0‖F region of the phase plane (upper left in Figure 5) is inaccessible with most well-conditioned models where ‖z0‖F ∼ 1 at initialization. If we only train one set of weights, the Θ̂ is independent of c.
Unless otherwise noted, all empirical studies in Sections 2 and 3.2 involve training a wide resnet on CIFAR10 with SGD, using GPUs, using the above correlated initialization strategy to fix Θ̂. All experiments used JAX (Bradbury et al., 2018) and experiments involving linearization or direct computation of the NTK used Neural Tangents (Novak et al., 2019b).

Review summaries:
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.

This paper investigates how the inverse temperature parameter $\beta$ in the softmax-cross-entropy loss impacts the learning and the generalization. In the theory part, this paper introduces the concepts of early learning timescale and nonlinear timescale, and shows how the learning dynamics depend on the parameter $\beta$. The empirical investigations are carried out with wide Resnets on CIFAR10, Resnet50 on ImageNet, and GRUs on IMDB. The results suggest that the optimal $\beta$ is architecture sensitive.