Temperature check: theory and practice for training models with softmax-cross-entropy losses
Authors: Atish Agarwala
Samuel Stern Schoenholz
Jeffrey Pennington
Yann Dauphin
Venue:  (2021)
Keywords: theory
learning dynamics
temperature
theory of deep learning
generalization

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 $\beta$ as well as the magnitude of the logits at initialization, $||\beta\bf{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 $\beta$ is not due to changes in model confidence, but is a dynamical phenomenon. It follows that the addition of $\beta$ as a tunable hyperparameter is key to maximizing model performance. Although we find the optimal $\beta$ to be sensitive to the architecture, our results suggest that tuning $\beta$ over the range $10^{-2}$ to $10^1$  improves performance over all architectures studied. We find that smaller $\beta$ may lead to better peak performance at the cost of learning stability.

Under review as a conference paper at ICLR 2021
TEMPERATURE CHECK: THEORY AND PRACTICE FOR TRAINING MODELS WITH SOFTMAX--ROSS-ENTROOY LOSSES
Anonymous authors Paper under double-blind review
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 generalization performance. In this work we develop a theory of early learning for models trained with cattmax-cross-entree
loss and show that the learning dynamics depend crucially on the inverse-temperature / as well as the magnitude of the logits at initialization, ||3z||2. We follow up these analytic results with a large-scale empirical study of a variety of model architectures trained on CIFARIO I0, 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 3 is not due to changes in model confidence, but is dynamical phenomenon. It follows that the addition of as a tunable hyperparameter is key to maximizing model performance. Although find the optimal B to be sensitive to the architecture, our results suggest that tuning 3 over the range 10-2 to 10 improves performance over all architectures studied. We find that smaller 33 may lead better peak performance at the cost of learning stability.
INTRODUCTION
Deep learning has led to breakthroughs across slew of classification tasks (LeCun et al. 1989 Krizhevsky 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 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.) 22019] (Szegedy et al. (0116). Another way 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 be important metric learning (Wu et al.)
generalization performance of softmax wwwscrrrssssrreeeee
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 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 corimax-cros--ae www---roosseeeeeee cross-yetrrppy loss have been shown to converge to max margin classifiers in particular function space norm (Chizat and Bach 2020), but timescales of convergence are not known. Additionally, many well cell-pefforming models operate best away from the lincarized
regime (Novak al. 2019a Aitchison 2019). This means that understanding the deviations of

---

Under review as a conference paper at ICLR 2021
models from their linearization around initialization is important for understanding generalization (Lee et al.) 2019 Chizat et al. 2019).
In this paper, investigate the training of neural networks with softmax-cros crosss-coosss-pppy losses. In general this problem is analytically intractable; to make progress we pursue a strategy that combines analytic insights at short times with 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-femperatures, 3 1/T, behave identically provided the learning rate is scaled as n = n/3 Here, networks begin to learn over a timescale T; |22ll//i where Z are the initial logits of the network after being multiplied by B. 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 Tal ~ ///nj suggesting more nonlinear learning for small 3.. 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 range of architectures and domains including: Wide Residual networks (Zagoruyko and Komodakis, 2017) on CIFARIO Krizhevsky 2009), ResNet-50 (He et al. 2016) on ImageNet (Deng et al. 2009). and GRUs (Chung et al., (0014) the IMDB sentiment analysis yii task (Maas et al. 2011). Db. In the case of residual networks, we consider architectures with and without batch normalization which can appreciably change the learning dynamics (loffe and Szegedy 2015). For all models studied, we find that generalization performance is poor at |z'||2 > 1 but otherwise largely independent of ||Z" 1|2- Moreover, learning becomes slower and less stable at very small 3; indeed, the optimal learning rate scales like 'y" ~ 1/B and the resulting early learning timescale can be written as T ~ | For all models studied, we observe strong performance for E [10 2, 10'] although the specific optimal 3 is architecture dependent. Emphatically, the optimal /3 is often far from 1. For models without batch normalization, smaller /3 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 3 over the range of [10-2 10*].
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 || 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 3 causes training to access the non- 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 (Zagoryyko and Komodakis 2017) trained on CIFARIO.
2.1 BASIC MODEL AND NOTATION
We consider a classification task with K classes. For an N dimensional input x, let 2(x, 0) be the pre-softmax output of classification model parameterized by 0 E IR.,, such that the classifier predicts the class corresponding to the largest output value z, We will mainly consider 0 trained by SGD on a training set (X,)) ))) of M input-label pairs. We focus on models trained with crtss-yntropy loss with non-trive non-trivial inverse temperature B. The softmax-crov wwwee-cooosssoeppppe loss can be written as
K
K
((,,X,)) = Σ 1n-nn(((32((2, 0))) = Σ y, (zo((Z)
(1)
2

---

Under review as a conference paper at ICLR 2021
where define the rescaled logits Z 3z and a (Z), [=)==c= = E is the softmax function. Here Z(X, 0) is the M K dimensional matrix of rescaled logits on the training set.
As we will see later, the statistics individual (Z); will have a strong influence on the learning dynamics. While the statisties of o (Z); are intractable for intermediate magnitudes, [| they can be understood in the limits of large and small || For fixed model z(x, 0), 8 controls the certainty of the predicted probabilities. Values of such that < <//|2 ///z will give small values of || < I. and the outputs of the softmax will be close to 1/K independent of i (the maximum- maximum-entropy distribution on K classes). Larger values of 3 such that > 1/ ||z||2 will lead to large values of || |Z||2 2 > 1; the resulting distribution has probability close to 1 on one label, and (exponentially) close on the others.
The continuous time learning dynamics (exact in the limit of small learning rate) are given by:
K
(2,
T (>; - (2,_,(((
θ = n1.3
Σ
00
for learning rate 7). We will drop the explicit dependence of Z; 0 from here onward, and we will denote time dependence as (X,, t) explicitly where needed.
In function space, the dynamics of the model outputs input given by
dz,(x)
K
m
Σ
(Oa)s (x, (x,) -o(Z((()))
(3)
dt
where we define the M x K x K dimensional tensor Ôg, the empirical neural tangent kernel (NTK), as
Oz,(x)
(0z;(A)
(Oo (x,
(4)
00
00
for class indices 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 99, the initial softmax input Z(X', 0) on the training set, and the initial Ôg. Changing these observables across 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 3 and || = [Z O] (the norm of the M x K dimensional matrix of training logits), independent of Óg.
2.2 LINEARIZED DYNAMICS
For small changes in 9, 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, f) only depend on the initial value of and the 3-scaled logit values ((X',), This suggests that there universal timescale across B and which can be used to compare linearized trajectories with different parameter values. Indeed, if we define an effective learning rate r = n132, we have
K
dz, dt
Σ -02)
(5)
which removes explicit 3 dependence 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 7 is useful for understanding the nonlinear dynamics, as plotting learning curves versus fjt causes early-time collapse for fixed Z across / and 77 (Figure I We that there a strong, monotonic, dependence of the time at which the nonlinear model begins to deviate from its linearization on 3. We will return and explain this phenomenon in Section 2.4
Unless otherwise noted, we will analyze all timescales in units of fj instead of 17. as it will allow for the appropriate early-time comparisons between models with different B3.
3

---

Under review as a conference paper at ICLR 2021
20
D-
0= 10
Loss
Loss
0
10
10"
7f 10F
Figure 1: For fixed initial training set logits Z" plotting learning curves against 7jt B2t causes the learning curves to collapse to the learning curve of the linearized model at early times (right), in contrast un-scaled curves (left). Models with large follow linearized dynamics the longest.
2.3 EARLY LEARNING TIMESCALE
We now define and compute the early learning timescale, Tz, that measures the time it takes for the logits to change significantly from their initial value. Specifically, we define тN such that for / < Tz expect [Z( t) Z(x, 0) < ||Z(x, 0)ll and for > Tz. |[ t) Z(x, 0)llæ [ 0) (or larger). This is synonymous with the timescale over which the model begins to learn. As we will show below,
| r/jj Therefore in units of ij, тz only depends || and not 33.
To see this, note that at very short times it follows from Equation 5 that
K
Z,(x,t) Z, (x, = Σ ((),,(x,))))) - (Z) O(())
(6)
It follows that we can define timescale over which the logits (on the training set) change appreciably from their initial value
1r
7 |6(X, X) a (Z)))l
where the norms are once again taken across all classes as well as training points. This definition has the desired properties for < T3 and Tz-
In units of îj. T. depends only on || IZ second comes from the contribution from the gradient [ (') - a (ZX ()) F. As previously discussed, since oZZ) saturates small and large values of || Fr follows that the gradient term will also saturate for large and small | and the ratio of saturating values is some ((1) constant
| . in two ways. The first is linear scaling in | [|=; the
independent [| and 3.
Tast accuracy Vs. $mo,3 1.00e
accuracy
104
0.18
[22")-=10""
1//122
3-000 51
Accuracy
[R 10
Acsuracy
0.16
[R
[R2)
22
10
2-1.00 82
1.00
| lrr
0.14
B3-
0.12
12R
0.10
0.10
107
107
107
10"
10F
10T
107
101
|1
6/72
Figure 2: The timescale т. depends only on || in units of 7 88 (left. inset). т. depends linearly on |r. to an (11) coefficient which saturates at large and small || r (left, main). Accuracy increases (center). Rescaling time to t/T, causes early accuracy curves to collapse (right).
quickly for small initial || F. though late time dynamics similar
The quantitative and conceptual nature of explicitly as [x, t) (x, 0)]/t|| for short times (Figure 2 left). When plotted over a wide range of ||z and (3, the ratio Tz/ Tz/||2||r (in rescaled time units) undergoes saturating, ((1)
both be confirmed numerically. We compute T;
4

---

Under review as a conference paper at ICLR 2021
variation from small to large F.. The quantitative dependence of the transition on the NTK confirmed in Appendix C Additionally, for fixed 3 and varying | [| rescaling time by 1/T; causes accuracy curves to collapse at early times (Figure 2 middle), even if they very different at early times without the rescaling (right). We note here that the late time accuracy curves seem similar without rescaling, point which will return in Section 3.2
across |
2.4 NONLINEAR TIMESCALE
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, Tnt corresponding to the time over which the network deviates appreciably from the linearized equations. We will show that Tx. x B/ūj. Therefore, in terms of 3 and |Z |1 networks with small 3 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 the context of MSE loss.
We define 7xd be the timescale over which the change in θg (which contributes the second order term in Equation (6) can no longer be neglected. Examining the second time derivative of Z, we have
d
dt = ij
Σ j=1
K
-(00) x) dt (Z(x)) dt [00))((x))) ()) - (Z, -
d
(8)
Lincarized dynamics
nonlineurized dynamics Z...
The first term is the second derivative under a fixed kernel, while the second term due to the change in the kernel (neglected in the linearized limit). A direct calculation shows that the second term, which we denote Zni. can be written
(Zni);= 8- EE o(((())) (2z)(2)) a (8oliy) (2)(()- (9)
00 0g
This gives us a nonlinear timescale Tnl defined, at initialization, by Tnl || 0)[l=/ (X (X, 0)[[F.
can interpret Tul as the time it takes for changes in the
kernel contribute to learning.
Though computing (x)(( in exactly is analytically intractable, its basic scaling in terms of B and || (and therefore, that of Twt)) is computable. We first note the explicit 8-2 dependence. The remaining terms are independent of [3 and vary by most O(1) with |Z indeed as described above, [|)((') - o (Z, 0)) |F saturates for large and small || Morevoer, the derivative, dx(X.0)
30 is the square root of the NTK and, at initialization, it is independent of | Together with our analysis of we have that, up to some O(1) dependence on || Taa x 8/ Therefore, the degree of nonlinearity early in learning is controlled via B alone.
10
10a
6x101
10a
A-
4x101
8- 10-
10"
A= 10-1
3x10
10
A- 10
-- 202
A- 30
101
101
10
10a
10'
| 1|
Figure 3: The time deviation from linearized dynamics, Txl, has large deviation over 3 and | (left), which can be largely explained by linear dependence on / (right), in units of r = 1T There an O(1) dependence on ||Z"| [| which is consistent across varying for fixed |
5

---

Under review as a conference paper at ICLR 2021
Optimal learning rate vs. ,
10
0.01/8
10a
10
10
10a
102
10
10P
101
107
3
Figure 4: Optimal learning rate ry" for WRN on CIFARIO scales as 1/B.
Once again confirm the quantitative and conceptual understanding of Tnt numerically. Qual- itatively, we see that for fixed (Zz||1= models with smaller 3 deviate sooner from the linearized dynamics when learning curves are plotted against ntt (Figure I We compute Tut explicitly by taking
|| [() Zran (1) [t)| as small times, where Zisn (1) is the solution to Equation T has an O(1)
dependence on |/ only (Figure
2.5 LEARNING RATES AND LEARNING SPEED
The timescales, T; and Tet, 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 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 Tyd S G, where is an architecture-dependee (1)) which has not yet been calculated theoretically (see Lewkowycz al. (2020) for empirical
constant of
calculations of c). This predicts a maximum effective learning rate 7j of 0(3) which in turn implies a
raw learning rate of ((3-1)))))))) 1) (as 7) (.)
However, if i is (33). then T; is 0(3-1). This means that, for early learning, networks with smaller will take more SGD steps to reach appreciable change in logits. Therefore, predict that networks with small 3 take longer to train (which we see in practice as well).
EXPERIMENTAL RESULTS
3.1 OPTIMAL LEARNING RATE
We begin our empirical investigation by training wide resnets (Szegedy al. (016) without batch normalization on CIFARIO 10, 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 B and || independently by using a correlated initialization strategy outlined in Appendix [.I
Before considering model performance, it is first useful to understand the scaling the learning rate with B. We define the optimal learning rate 17" as the learning rate with the best generalization performance. To do this, we initialize networks with different ,3 and conduct learning rate sweeps for each 3. The optimal learning rate 77" has clear 1/3 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 3 also learn more slowly in terms of number of SGD steps.
Thus, small we expect learning to take place slowly and nonlinear effects to become important by the time the function has changed appreciably. At large 3, by contrast, our results suggest that the network will have learned significant amount before the dynamics become appreciably nonlinear.

---

Under review as a conference paper at ICLR 2021
3.2 PHASE PLANE
In the preceding discussion two quantities emerged that control the behavior of carly-time dynamics: the inverse-temperature, fi and the rescaled logits |/ In attempting understand the behavior of real neural networks trained using softmax crosseerroosss loss, it therefore makes sense to try to reason about this behavior by considering neural networks that span the B || phase plane, the space of allowable pairs (3, || F). By construction, the phase plane is characterized by the timescales involved in early learning. To summarize, Tz /l/ sets the timescale for early learning, with larger values of | leading to longer time before significant accuracy gains are made (Section 2.3 Meanwhile, Tyd 3/ controls the timescale for learning dynamics to leave the linearized regime- with small [3 leading to immediate departures from linearity, while models with large may stay linearized throughout their learning trajectories (Section 2.4).
(b)
Aoou Soumany USD clepl)
(a)
10F
||2" la 3e Poor conditioning
Slow
19'
16
Rx pellla Plla Fr
Early learning
0.85
Accuracy
10
(c)
10-2
-075
10
0.60
10
Fast
9
10-3
10-2
10-1
10"
10
10a
162
0.40
/9
Tux 9/9
More
Linearily
Figure 5: Properties of early learning dynamics, which affect generalization, can be determined by location in the 3-|| Z phase plane (a). At optimal learning rate 17" small 3 and larger || leads to slower early learning (b), and larger 3 increases time before nonlinear dynamics contributes to learning. Large | has poorly conditioned linearized dynamics. Generalization for wide resnet trained on CIFAR 10 is highly sensitive to B, and relatively insensitive to || outside poor conditioning regime. Final logit variance is relatively insensitive to parameters (c).
In Figure (a), we show schematic of the phase plane. The colormap shows the test performance of wide residual network (Zagoruyko and Komodakis 2017). without batch normalization, trained on CIFAR10 10 in different parts of the phase plane. The value of makes a large difference in generalization, with optimal performance achieved at 3 x 10-2, In general, larger performed worse than small / as expected. Moreover, observe similar generalization for all sufficiently large 3; this to be expected since models this regime are close to their linearization throughout training (see Figure [ and expect the linearized models to have 3-independent performance. Generalization was largely insensitive to | F so long as the network was sufficiently well-conditioned to be trainable. This suggests that long term learning is insensitive T2-
In Figure (b)) we plot the accuracy after 20 steps of optimization (with the optimal learning rate). For fixed [re,,, Ir the training speed slow for the smallest and then became faster with increasing 3. For fixed 3 the training speed was fastest for small | F and slowed as ||z [| 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 early times did correlate with the performance at late times.
This highlights that differences in generalization are dynamical phenomenon. Another indication of this fact is that the end of training, at time ty. the final training set logit values || |r = (Z(xtty)) tend towards 1 independent of the initial 3 and || (Figure 5 (c)). With the exception of the wworly-performing large || 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 (

---

Under review as a conference paper at ICLR 2021
3.3 ARCHITECTURE DEPENDENCE OF THE OPTIMAL 3
Having demonstrated that 3 controls the generalization performance of neural networks with softmax- cross-entropy loss, we discuss the question of choosing the optimal 3. 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 3 can be predicted analytically an open question that we leave for future work. Nonetheless, we show that all architectures considered display optimal / between approximately 10-2 and 10 We observe that by taking the time to tune /3 it is often the case that performance can be improved over the naive setting of B = 1.
3.3.1
WIDE RESNET ON CIFARIO
(a)
Accuracy vs. average
(b)
Accumcy vs. A, max
(c)
1.00
100
0:0
0.8
Accuracy
0.90
0.0
Accuraci
0.4
0.90
-..11
0.00
6,-
1.11
082
0.75
1.21
1.37
1.21
1.37
0.0
ovet sunds Aessage over sneds
0.70
1.52
0.70
0.78
665
10
Figure 6: Dependence of test accuracy for various architectures with 3 tuning. (a) For WRN with batchnorm, trained on CIFARIO, 10, the optimal |3 = 10. Without batchnorm, the performance of the network be nearly recovered with 3-scaling alone with / 10-2 Even poorly conditioned networks (achieved by increasing weight scale aw,) recover performance. (b) For 3 < 10-2 learning is less stable, as evidenced by low average performance but high maximum performance (over 10 random seeds). (c) We see similar phenomenology on the IMDB sentiment analysis task trained with GRUs where average-case best performance is near
but peak performance is at small 3.
In Figure 6 (a) we show the accuracy against 3 for several wide residual networks whose weights are drawn from normal distributions of different variances, oi trained without batchnorm, well as a network with of = trained with batchnorm (averaged over 10 seeds). The best average performance is attained for 8 < 1, aw = 1 without batchnorm, and in particular networks with large ow are dramatically improved with tuning. The network with batchnorm is better at all 3, with optimal / R 10. However,
that the best performing seed is often at a lower 3 (Figure 6 (b)), with larger ow networks competitive with o, 1, and even with batchnorm at fixed / (though batchnorm with 3 10 still performs the best). This suggests that small 3 can improve best case performance, at the cost of stability. Our results emphasize the importance of tuning B, especially for models that have not otherwise been optimized.
3.3.2 RESNET50 ON IMAGENET
Table I: Accuracy on Imagenet dataset for ResNet-50. Tuning B significantly improves accuracy.
Method
Accuracy (%)
ResNet-50 GGiassi et al. 2018
76.51 + 0.07
ResNet-50 + Dropout (Ghasi al. 2018
76.80 + 0.04
ResNet-50 + Label Smoothing (Ghaasi et al. 2018
77.17 + 0.05
ResNet-50 + Temperature check ( 0.3)
77.37 + 0.02
Motivated by our results on CIFARIO, 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 T shows that scaling improves accuracy for ResNet-50 with batchnorm. However, we did not
8

---

Under review as a conference paper at ICLR 2021
find that using < was optimal for ResNet-50 without normalization. This further emphasizes the subtle architecture dependence that warrants further study.
3.3.3 GRUS IMDB SENTIMENT ANALYSIS
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 the results CIFAR 10 without batch
normalization. We observe peak performance near averaged over an ensemble of networks, but we observe that smaller can give better optimal performance at the expense of stability.
3.4 PROPOSED TUNING PROCEDURE
Our results suggest the following tuning procedure for networks trained with SGD/momentum:
Train/tune a model as normal. Note the optimal learning rate Yju-
For the best parameter set, sweep over 3 E [10-2, 10 scaling learning rate as /// /..
If best performing is at an endpoint the range, continue tuning.
If there is sufficient compute, the 3 search can instead be folded into the overall hyperparameter tuning.
We note that our observation of optimal 3 < 1 in classification settings stands in contrast to the observation of optimal > in the Bayesian inference setting (Wenzel et al. 2020p. We speculate that the differences important, and that the objective (learn an SDE which reproduces a target distribution) may have qualitatively different properties than classification tasks.
related to the fact that in Bayesian inference, the uncertainty of estimates
CONCLUSIONS
Our empirical results show that tuning 3 can yield sometimes significant improvements to model performance. Perhaps most surprisingly, we observe gains on ImageNet even with the highly- optimized ResNet50 model. Our results on CIFARIO suggest that the effect of 3 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, Tut & 1/3n, which at first glance suggests that models with smaller B will become nonlinear more slowly than their large 3 counterparts. However, analyzing Tn with respect to effective learning rate n = 821 yields Twt x 3/7j. 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 33 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 B, but does not predict the optimal value. cehniectuur-dependeen non-scaling multiplicative factors to key learning parameters have been observed in other contexts (Lewkowycz et al.] (020), and their numerical estimation is in general difficult. Extending the theoretical results make predictions about these quantities is interesting avenue for future work. Another area that warrants further study is the instability in training at small 3.
9

---

Under review as a conference paper at ICLR 2021
REFERENCES
Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel. Backpropagation Applied to Handwritten Zip Code Recognition. Neural Computation, 1(4): 541-551, December 1989. ISSN 0899-7667. doi: 10. 10.11622/eco.1990 1989.1.4.544.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. ImageNet Classification with Deep Convolutional Neural Networks. In Advances Neural Information Processing Systems 25, pages 1097-1105. Curran Associates, Inc., 2012.
Sergey Zagoruyko and Nikos Komodakis. Wide Residual Networks. arXiv:1605.07146 [cs], June 2017.
Douglas M. Kline and Victor L. Berardi. Revisiting squared-error and cross-entropy functions for training neural network classifiers. Neural Computing & Applications, 14(4):310-318, December 2005. ISSN 1433-3058.
doi: 10. 04.00577.005210005.44667....
Pavel Golik, Patrick Doetsch, and Hermann Ney. Cross-Entropy Squared Error Training: A Theoretical and Experimental Comparison. In Interspeech, volume 13, pages 1756-1760, August 2013.
Gabriel Pereyra, George Tucker, Jan Chorowski, Lukasz Kaiser, and Geoffrey Hinton. Regularizing Neural Networks by Penalizing Confident Output Distributions. arXiv:1701.06548 [cs], January 2017.
Rafael Müller, Simon Korbllihh, and Geoffrey E Hinton. When does label smoothing help? In Advances Neural Information Processing Systems 32, pages 4694-4703. Curran Associates, Inc., 2019.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethinking the Inception Architecture for Computer Vision. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 2818-2826, 2016.
Zhirong Wu, Alexei A. Efros, and Stella X. Yu. Improving Generalization via Scalable Neighborhood Component Analysis. OrXiv:/805) 04699 [cs], August 2018
Andrew Zhai and Hao- Yu Wu. Classification is Strong Baseline for Deep Metric Learning. arXiv:1811.1264 [cs], August 2019.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the Knowledge in a Neural Network. arXiv:1503.0253/ [es, stat], March 2015.
John Platt. Probabilistic Outputs for Support Vector Machines and Comparisons Regularized Likelihood Methods. Adv. Large Margin Classif., 10, June 2000.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On calibration of modern neural networks. In Proceedings the 34th International Conference on Machine Learning Volume 70, ICML' 17, pages 1321-1330, Sydney, NSW, Australia, August 2017. JMLR.org.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural Tangent Kerel: Convergence and Generalization in Neural Networks. In Advances in Neural Information Processing Systems 31, pages 8571-8580. Curran Associates, Inc., 2018.
Simon S. Du, Xiyu Zhai, Barnabás Póczos, and Aarti Singh. Gradient Descent Provably Optimizes over-parameterrizd Neural Networks. In 7th International Conference Learning Representa- tions, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide Neural Networks of Any Depth Evolve Linear Models Under Gradient Descent. In Advances in Newral Information Processing Systems 32, pages 8570-8581. Curran Associates, Inc., 2019.
Roman Novak, Lechao Xiao, Jaehoon Lee, Yasaman Bahri, Greg Yang, Jiri Hron, Daniel A. Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian Deep Convolutional Networks with Many Channels Gaussian Processes. arXiv: 1810.05148 (cs, stat], February 2019a
10

---

Under review as a conference paper at ICLR 2021
Lechao Xiao, Jeffrey Pennington, and Samuel S. Schoenholz. Disentangling trainability and general- ization in deep learning. arXiv:1912.1305 [cs, stat], December 2019.
Lenaic Chizat and Francis Bach. Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss. arXiv:2002.08 04046 [cs, math, stat], March 2020.
Laurence Aitchison. Why bigger not always better: On finite and infinite neural networks. arXiv:1910.08013 [cs, stat], November 2019.
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On Lazy Training Differentiable Programming. In Advances in Neural Information Processing Systems 32, pages 2937-29477 Curran Associates, Inc., 2019.
Jachoon Lee, Samuel S. Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite Versus Infinite Neural Networks: An Empirical Study, July 2020.
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: The catapult mechanism. March 2020.
A. Krizhevsky, Learning Multiple Layers of Features from Tiny Images. Master's thesis, University of Toronto, 2009.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. In I Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pages 770-778, 2016.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large- scale hierarchical image database. 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248-255 June 2009. doi: 1L..0/CPR..0090552222.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling. arXiv:1412.3555 [cs], December 2014.
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguisties: Human Language Technologies Volume I, HLT *11, pages 142-150, USA, June 2011. Association for Computational Linguistics. ISBN 978-1-932432-87-9,
Sergey Ioffe and Christian Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. February 2015.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A Convergence Theory for Deep Learning via Over- Pree-aaattteiziiionnn. In International Conference on Machine Learning, pages 242-252. PMLR, May 2019.
Golnaz Ghiasi, Tsung- Lin, and Quoc Le. DropBlock: regularization method for convolutional networks. In Advances in Neural Information Processing Systems 31, pages 10727-10737. Curran Associates, Inc., 2018.
Niru Maheswaranathan, Alex Williams, Matthew Golub, Surya Ganguli, and David Sussillo. Reverse engineering recurrent networks for sentiment classification reveals line attractor dynamics. In Advances in Neural Information Processing Systems 32, pages 15696-15705. Curran Associates,
Inc., 2019.
Florian Wenzel, Kevin Roth, Bastiaan S. Veeling, Jakub Swiętkowski. Linh Tran, Stephan Mandt, Jasper Snoek. Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How Good is the Bayes Posterior in Deep Neural Networks Really? arXiv:2002.0244 [cs, stat], July 2020.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin. and Skye Wanderman- Mindderman-Minee. JAX: Composable transformations of Python+NumPP programs, 2018.
11

---

Under review as a conference paper at ICLR 2021
Roman Novak, Lechao Xiao, Jiri Hron, Jachoon Lee, Alexander A. Alemi, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Neural Tangents: Fast and Easy Infinite Neural Networks in Python. arXiv:1912.02803 [cs, stat], December 2019b.
12

---

Under review as a conference paper at ICLR 2021
A LINEARIZED LEARNING DYNAMICS
A. FIXED POINTS
For the linearized learning dynamics, the trajectory z(x, t) can be written in terms of the trajectories of the training set as
(xx t) z(x, 0) = θ(x, X' ()) - ()),0))
(10)
where + is the pseudo-inverse. Therefore, one can solve for z(X, tt) 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 guaranteed fixed point if the appropriate L2 regularization added to the training objective. Given a regularizer >> on the change in parameters 80 (t) 0(0), the dynamics in the linearized regime are given by
xx) (x)0xx,)))((
o(3((X))) ^goz(x)
(11)
where the last term comes from the fact that
dz i0 ōz(x) in the linearized limit. σ6
We can write down eeff-consistent equations for equilibria, which are approximately solvable in certain limits. For an arbitrary input x, the equilibrium solution x (x) is
0 px( ())() - a(32' (x))) d00z=(x (xx)
(12)
This can be rewritten in terms of the training set as
(x) =
(13)
similar to kernel learning.
remains then to solve for z (2). We have:
8z (x)
X (Dxx
(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 casiest case to analyze is when [| (1822((')] (X')ll < 1. Then
we have:
1 (1 (2) 3 x) [V(x)] - K
1
x² xx) do
(15)
which gives us
3
B
2x)x
é ()(x)) 1/K)
(16)
Kλ
Therefore the self-consistency condition for this solution is || solution to
8 @1|r << 1, which simplifies the A/
xx() Xg
(17)
This is equivalent to the solution after single step of (full-batch) SGD with appropriate learning rate. We note that unlike lincarized dynamics with Le loss and full-rank kernel, there no guarantce that the solution converges to 0 training error.
The other natural limit is [|(32((x')) of the conserved quantity of learning with cross-entropy loss. We note that the vector on the right hand side of Equation sums to for every training point. Suppose at initialization, Ô has no logit-logit interactions, as is the case for most architectures the infinite width limit with random initialization. More formally, we can write Ô = Idxx & Ô, where Ô., is M x M. Then, the have:
(2)]12 >>
We focus on the 2 class case, in order to take advantage
sum of the logits for any input x is conserved during linearized training,
1 T *xxx 1131TT [dæx @ 6.] (y) - (B2(21))
(18)
13

---

Under review as a conference paper at ICLR 2021
Multiplying the right hand side through, get
1 122(x) = n 11 (B2(2))))
(19)
(Note that Ôg has explicit dependence on the logits, there still is conserved quantity, which is more complicated to compute.)
Now we can analyze ||22 (x) > 1. With two classes, and z(X) = at initialization, we have zi -zg. Therefore, without loss of generality, focus on zj. the logit of the first class. In this limit, the leading order correction to the softmax is approximately:
(3zi) x 1270 sign(z))e 23/xil
(20)
The self-con consistencc equation then:
zi (x) = A
Ō(x, (122) - 12;00 + (zi)( -
(21)
The vector on the right hand side has entries that are (e-22lz))) for correct classifications, and ((1) for incorrect ones. If we assume that the training error is 0, then have:
B ((X, si)ssig(z))) 2
z (x)
(22)
Ao
This is still non-trivial to solve, but we see that the self consistency condition is that In(81(@1==/>a) >> 1.
Here also may be difficult train and generalize well. The individual elements of the right- hand- side 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 Le 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 't matter; rather, the ratio of 3 and Àg (compared to the appropriate norm of Ô) needs to be balanced to prevent falling into the small (3z regime (where training error might be large) the large (3 regime (where few datapoints might dominate and reduce generalization).
A.2 DYNAMICS NEAR EQUILIBRIUM
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 ²(X') = =(X') x (x) for small perturbations given by
a(x) 11 3 Id2 & 6(x, 2))a, (B2) AD} 2(x))
(23)
where og is the derivative of the softmax matrix
@z(z)
σ (z)
Oz' (o -z))
(24)
We can perform some analysis in the large and small 3 cases (once again ignoring Az). For small ||5z* (z)), we have |l - Ó|[F < which leads to:
o ==(2 (1/K - -1/K-11//4²++011 0
(25)
This matrix has K eigenvalues with value 1/K, and one zero eigenvalue (corresponding to the conservation of probability). Therefore || @ (32 & >o. and the
well-conditioned regularizer dominates the approach equilibrium.
In the large 3 case (In(3][(1/Aoo) > 1), the values of ((z((()) are exponentially close 0 (K 1 values) or 1 (the value corresponding to the largest logit). This means that az (fz((')) has exponentially small values in [|xx(X*)]l=== if any of (3xx(())) and a (32, is exponentially
14

---

Under review as a conference paper at ICLR 2021
small, the corresponding element of a (Bz(())) is as well; for the largest logit the diagonal is (32,(2))) - (()) which is also exponentially small.
From Equation 22 we have XØ < e therefore, though the o term of 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.
A.3 CONDITIONING OF DYNAMICS
Understanding the conditioning of the linearized dynamics requires understanding the spectrum of the Hessian matrix H [Idy& & 6(x; a=(32" (X)). In the limit of large model size, the first factor is dlock-diagonal with training set by training set blocks (no logit-logit interactions), and the
second term is boockdiagonal with K K blocks (no datapoint-datppimn interactions).
We will use the following lemma to get bounds on the conditioning:
Lemma: Let M AB be matrix that is the product of two matrices. The condition number ((M)= mii
has bound
B)/k(A) S R(M) S ((A)x(()
(26)
Proof: Consider the vector v that is the eigenvector of B associated with AB.
Note
that ||v///Ilv|| S AA, Analogously, for w, the eigenvector associated with AB max
A 1/1 w|| AA.min- This gives the two bounds:
AM,mi I
Aa AM.ma Ai max
(27)
This means that the condition number k(H)
AM.max is bounded by Adamii
AA Maxa Xn x(B)/n(A)
(MM)
(28)
A A,
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 a o (8z" (X)) (Xz) will lead to poor conditioning of the linearized dynamics if the NTK Ô X) is (relatively) well conditioned. This bound will be important in establishing the poor conditioning of the linearized dynamics for the large logit regime
||8z|| >> 1.
A.3.1
SMALL LOGIT CONDITIONING
For ||.2² (X()lr & 1, the Hessian H is
1 K
1 (1- 11") @ ((x,x)
(29)
Since H is the Kroenecker product of two matrices, the condition numbers multiply, and we have
H =x(@)
(30)
which is well-conditioned long the NTK is. Regardless, the well-conditioned regularization due to >g dominates the approach to equilibrium.
A.3.2 LARGE LOGIT CONDITIONING
Now consider |,8z [F > 1. Here we will show that the linearized dynamics is poorly condi- tioned, and that H: (H) is exponentially large in B.
We first try to understand o (f2z (x)) for an individual x E X. To Oth order (in an as-of-yet- undefined expansion), o, is zero at large temperature the softmax returns either 0 or which by Equation (aciioon24]gvvve gives 0 in all entries. The size the corrections end up being exponentially dependent on [| mF the entries will have broad, log-normal type distributions with magnitudes which scale
15

---

Under review as a conference paper at ICLR 2021
exp(-Blzil). There will be two scaling regimes with small number of labels in the sense VB> Vln(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 still exponential dependence on 3; we will focus on the first which is easier analyze and more realistic (e.g. for 10t labels "large" / is only ~ 15).
Let 21 be the largest of K logits, 22 the second largest, and so on. Then using Equation 24 we have:
(0 ---
(31)
for + 1,
(a=);s = bize - 2()
(32)
for # jnd and
(a:)11 (221-22)
(33)
The eigenvectors and eigenvalues can be approximately computed as:
(v1)1 V2
(v1)2
1
A1 2e -
(34)
V2
(v2)1 V2
(V2)2
A2
20-
-20(21-88
(35)
V2"
and for i 2,
(vi), = 1, (v.)1 8(22 2), A = e 9(21-2))
(36)
with all non-explicit eigenvector components 0. This expansion is valid provided that 3/K > 1 (so
that ((11-z)) >> e (1) =1)))).
Therefore the spectrum of any individual block o= (3z((x) is exponentially small in B. Using the bound in the Lemma, we have:
H ((002))) (x (x(O(X,
(37)
This is very loose bound, it assumes that the largest eigendirections of ; are aligned with the smallest eigendirections of Id2 & Ô, and vice versa. It is possible x(H) is closer in magnitude the upper bound ( < (()
Regardless, K( (H)) is exponentially large in meaning that the conditioning is exponentially poor for large || F.
B SGD AND MOMENTUM RESCALINGS
B.1 DISCRETE EQUATIONS
Consider full-batch SGD training. The update equations for the parameters are:
0 = 0 VeL
(38)
We will denote g: = VeL for ease notation.
Training with momentum, the equations of motion are given by:
Vt+1 (1 --)v= go
(39)
0 = 0 + 7/V++1
(40)
where y E [0, 1].
One key point to consider later will be the relative magnitude Ae of updates to the parameters. For SGD, the magnitude of updates 7// m//gg] For momentum with dlwll-v gradients the magnitude
n/|lg] 1/7-
16

---

Under review as a conference paper at ICLR 2021
B.2 CONTINUOUS TIME EQUATIONS
We can write down the continuous time version of the learning dynamics as follows. For SGD, for small learning have:
d0
- 7
(41)
dt
For the momentum equations we have
dv
yV g
(42)
dt
d0
= 1/v
(43)
dt
From these equations, we can see that in the continuous time limit, there are coordinate transforma- tions which can be used to cause sets of trajectories with different parameters to collapse to single trajectory. SGD is the simplest, where rescaling time to r 7tt causes learning curves to be identical for all learning rates.
single universal learning curve, there is one-parametes family of curves controlled by the ratio momm n/22. Consider rescaling time to at and bv, where and b will be chosen to put the equations in a canonical form. In our new coordinates, we have
For momentum, instead
du dr
(r/a)e (b/a)g
(44)
do
= ma/(ab)
(45)
dr
The canonical form we choose is
dv
du g
(46)
dr
d0
(47)
dr
From which we arrive at = b = /n, which gives us X = 7/ //.
Note that this not a unique canonical form; for example, we fix coefficient of -1 on V, we end up with
dv - ---v-mn/711
(48)
dr
d6 dr =
(49)
with = 7. This is different time rescaling, but still controlled by Tmom-
Working in the canonical form of Equations 46 and 47
can analyze the dynamics. One immediate
question is the difference between À < 1 and I. We note that the integral equation
((T) = (0) + [ e'e-Xa EXt X(T ''''')a 'g('')d'''
(50)
solves the differential equation for v. Therefore, for À v(t) only depends on the current value g(t) and we have ((T) R /A.) Therefore, we have, approximately:
d0 dr A
(51)
This means that for large all the curves will approximately collapse, with timescale given by VT-A - 777 (dynamics similar to SGD).
For < 1, the momentum essentially the integrated gradient across all time. If v0)) then we have
d0
R '' g(r')dr'
(52)
dr
17

---

Under review as a conference paper at ICLR 2021
In this limit, (()) is the double integral of the gradient with respect to time.
Given the form of the controlling parameter T moms we can choose to parameterize ? = Tv... Under
this parameterization, we have Tmom 5/. The dynamical equations then become:
dv dr
- g
(53)
d0
v
(54)
dr
which automatically removes explicit dependence on
One particular regime of interest is the early-time dynamics, starting from v(0) = 0. Integrating directly, we have:
1 (r) = - + 17T +
(55)
This means that T alone is the correct timescale for early learning, at least until 7ÿ ~ which in the original parameters corresponds to 1/7 (the time takes the momentum be first "fully integrated").
B.3 DETAILED ANALYSIS OF MOMENTUM TIMESCALES
One important subtlety is that 1 not the correct value to compare to. The real timescale involved is the one over which g changes significantly. We approximate this in the following way. Suppose that there is some relative change TØl of the parameters that leads to an appreciable relative change in g. Then the timescale over which & changes by that amount the to.
Ae
must compare Equation 51 approximately holds. The timescale Ta of the change is then given by:
We can compute that timescale the following way. We assume g fixed for what follows. Therefore,
Ao
g
Te C
(56)
|0]|
0
which gives
e1|10|1/|lg]
(57)
In particular, this means that the approximation is good when XT. >> 1, which gives >2/n >>>>>>>>>>>>> the former being a function of the dynamical parameters, the latter being a function of the geometry of L with respect to O.
[
One consequence of this analysis is that the |Ø| remains roughly constant, for fixed 17 and 7, late in learning when the gradients become small the dynamics shifts into the regime where À large, and we effectively have SGD.
B.4 CONNECTING DISCRETE AND CONTINUOUS TIME
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 7 and r) are scaled together. However, one typical use case for momentum is to fix the parameter 7, and sweep through different learning rates. With this setup, if g is changing slowly compared to 7 (more precisely. 7 /m > /|/|// as may be the case at later training times, the change in parameters from single step AS ~ (n/2)) and the rescaling of taking í to zyt (as for SGD) collapses the dynamics. Therefore given a collection of varying n, but fixed y curves, it is possible to get intermediate and late time dynamics the same scale.
However, early times, while the momentum still getting "up to speed" (i.e. in the first 1/7 steps), the appropriate timescale is 17 -1/2, Therefore, in order to get learning curves to collapse across
18

---

Under review as a conference paper at ICLR 2021
different 7) at early times, need to rescale y with 1 implied by Equations 46 and 47
Namely,
one must fix y and rescale y = / We note that, since y < 1, this gives us a restriction 7) < 7 for the maximum learning that can be supported by the rescaled momentum parameter.
B.5 MOMENTUM EQUATIONS WITH SOFTMAX-CROSS-ENTINN LEARNING
For cross-entropy learning with softmax inputs fz, all the scales acquire dependence on J3. we define lg
and 9s
8z 30 pm then we have, approximately. |g ~ 36.92.
Consider the goal of obtaining identical early-time learning curves for different values of B. (The curves are only globally consistent across / in the linearized regime.) In order to get learning curves to collapse, we want dE dr to be independent of 33 in the rescaled time units. We note that the change in
the loss function Ag from a single step of SGD goes as
A n/3 2g
(58)
This suggests that one way collapse learning curves is to plot them against the rescaled learning rate ijt. where ý = n7.2. While hyperparameter tuning across B, one could use = 8², sweeping over 7) in order 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 Ao for single step be small across B. We have
Ag 3B
(59)
which suggests that maximum allowable learning rates will scale as 1/3. This suggests setting " - and rescaling time as f) 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 dynamies of the loss. Starting with the rescalings from Equations
53 and 54 we have
dv
-Nv 8g
(60)
dr
d0
(61)
dr
dC
3g
(62)
dr
where g
8C Rescaling by / gets us: 832 80
dv
dßr 3 g
(63)
d0
(64)
d3T
dC
g
(65)
dBT
This rescaling causes collapse the trajectories of the at early times if 7/8 is constant for varying B. One scheme to arrive at the final canonical form, across 3, is by the following definitions of 77. 7 and
118-2 1
1 =
7= 7= BT vi
== T= B/T Vii
19

---

Under review as a conference paper at ICLR 2021
where curves with fixed will collapse. The latter two equations similar to before, except with replaced with fj. The dynamical equations are then:
dv Vv g
(66)
dr
d0
(67)
dr 3
dL
dr g V
(68)
The change in parameters from a single step (assuming constant and saturation) is
g
Ao
(69)
If we instead want the change in parameters from a single step to be invariant of 3 the continuous time approximation holds, while maintaining collapse of trajectories, we first note that
Ag
3la0z
(70)
from a single step of the momentum optimizer. To keep Ae invariant 3, we can set:
rig-1 B = 1 = 31/2211111 VTn B 3//2
Note that the relationship between y and by is the same in both schemes when measured with respect the raw learning rate 1/.
C SOFTMAX-CROSS-ENTE
GRADIENT MAGNITUDE
C.I MAGNITUDE OF GRADIENTS IN FULLY-CONNECTED NETWORKS
The value T. has nontrivial (but bounded) dependence on | F via the |ς² - (X)() 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 / correlates well with hidden units/layer) fully connected networks (Figure 7 The ratio depends details like the nonlincarities 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 CIFAR 10 when the training labels are randomly shuffled (right).
for wide (2000
Therefore in general the ratio of T. 1=/
|| large and small || depends subtly on the relation- ship between the NTK and the properties of the data distribution. A full analysis of this relationship 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.
D EXPERIMENTAL DETAILS
D.1 CORRELATED INITIALIZATION
In order to avoid confounding effects of changing and IZ'|[e with changes to Ô, we use correlated initialization strategy (similar (Chizat et al., 2019) which fixes Ô while allowing for independent variation of 8 and || Given network with final hidden layer h(xx, 0) and output weights Wo. define combined network z. (x, 0) explicitly as
(x, WO.1. Wo,2, 01, 02) Wo, 01) Wo,zh(x, 0z)
(71)
20

---

Under review as a conference paper at ICLR 2021
Normalzze Relu
f,.
y, labels)
6x:10
4x.10
(192
10 137
10
10 3x101
10
107
17 3x11 107
12"
10
||x0|
[|x
with with computed in the infinite width limit (in units of effective learning rate ij = 12),). Ratio between normalized timescales at large and small ||2]| [| depends on nonlinearity (left and middle), as well as training set statistics
Figure 7: [|[
| is highly correlated with |(²()) (000)~o(("")) -Z
(right, CIFARIO with shuffled labels).
where, at initialization, 01 = -2, and the elements of Wo,a have statistics
budd
for for *
([Wo,)ss (oo,) (c 6e0
(72)
for correlation coefficient e [-,, where dig is the Kronecker-delta which is 1 is i = and 0 otherwise. Under this approach, the initial magnitude of the training logits is given by || = BV(11 + +c)]2(()))) ()([2(||r. where ||z]||r the initial magnitude of the logits of the base model. By manipulating and c, we can independently change and || || with the caveat that || < /2,8||=|||r since e S
follows that the small B3, large ||Z|| region of the phase plane (upper left in Figure is inaccessible with most well-conditioned models where || at initialization. If only train set of weights, the Ô is independent of
Unless otherwise noted, all empirical studics in Sections and 3.2 involve training a wide resnet on CIFARIO with SGD, using GPUs, using the above correlated initialization strategy to fix Ô. All experiments used JAX (Bradbury et al. (018)) and experiments involving linearization direct computation of the NTK used Neural Tangents (Novak et al. 2019b).
21


Reviews:
[{"review_id": "w-ho4x2QZQn", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "##  Summary of the paper\nThis paper proposed the new time scaling analysis for the cross-entropy loss with temperature in the deep neural networks. The authors introduced the effective time scale, and then provided the linear and non-linear time scale based on the output of logits with temperature and stepsizes by using the NTK theory.\n\n## Strong and weak points of the paper\n### Strong points\n- Provided the novel time scaling analysis for the cross-entropy loss with many empirical validations.\n### Weak points \n- The choice of optimal $\\beta$ still remains unclear and left to the future work.  I think clarifying the tuning strategy of $\\beta$ is the most important from practical view point.\n\n## Rating\n- Clarity: The authors should clarify which part corresponds to the previous work proposed or this work proposed\n- Correctness: I did not check all the proof in detail.\n- Novelty: Since I am unfamiliar with this field, I thought that introducing the linear and non-linear time scale in Sec 2.3 and 2.4 seems intereting tools in experimetns.\n\n## Comments and Questions\n- Q) In Sec 2.2 to 2.4, what types of models and datasets are used in the numerical experiments ? Although I might overlooked, it should be explained explicitly.\n\n- Q) As far as I understood, from the ovservation of Sec 2.2 to 2.4, setting large $\\beta$ makes the times scales $\\tau_z$ and $\\tau_{nl}$ smaller, which means that  the early linear learning timescale and nonlinear timescale smaller. So, setting large $\\beta$ enhance escaping from the linear dynamics at the beggining, but the time period of the nonlinear dynamics become shorter, which results in the worse final performace. Is my understanding correct ?\n\n- Q) What is the definition of optimal $\\eta^*$ in Sec 3.1?\n\n- Following is just a comment. In Sec 3.3, the authors found that using small $\\beta\\in [10^{-2},1)$ might improve the final performance in a variety of networks.  On the other hand, when we consider Bayesian deep learning, using cold posterior that uses the temperatrue below $1$ results in better final perfomance compared to using $\\beta=1$ (For example see, Wenzel, Florian, et al. \"How good is the bayes posterior in deep neural networks really?.\" arXiv preprint arXiv:2002.02405 (2020). ). So, I just thought that there might be some connection between this small temperature phonomena.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "ipdYYw1vtR", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nStrengths:\n* The topic is novel and interesting.\n* Empirical results confirm the importance of tuning temperature.\n\nWeaknesses:\n* Theoretical contribution seems unclear.\n* Experiments are limited to two tasks.\n\nI am leaning towards rejecting the paper. My biggest concern is about the theoretical contribution. The main conclusion of the theory seems to be that smaller beta is better, because it helps neural networks exit linear regime faster. But this doesn’t explain the experiment results: optimal beta varies a lot across models, and it is sometimes quite large (beta >= 1). The paper empirically shows that small beta causes more instability as an explanation, but there is no theoretical explanation. Therefore, it is unclear how to use the current theory.\n\nIdeally, I hope the theory can be extended to explain why small beta causes instability (the conclusion section mentions this as future work), and/or how neural network architecture affects optimal beta, but these extensions do not seem obvious.\n\nThe experiments can also be expanded. Currently, there are only two tasks in the experiments. While the results are impressive, I would be more convinced if there are more tasks/models. For example, I wonder what the optimal beta is for state-of-the-art BERT-based models.  What about structured prediction tasks? Does the size of training set affects optimal beta?\n\nOn the positive side, I think the topic is novel and interesting, and the current empirical results are solid. If the theory can be extended to explain the tradeoff between small/large beta and the role of architecture, I would recommend this paper. Alternatively, the paper can also be improved by expanding experiments.\n\nOther suggestion and question:\n- Section 2 (theory) may be easier to read if there is a short summary of main claims/results.\n- For IMDb experiment, is there a reason for choosing GRU instead of more recent BERT-based models?\n- Some of the figure fonts are too small.\n\nFeedbacks after author response:\nI am maintaining my rating after reading the author response. It is a close decision. I like the topic, but I think the draft still has room for improvement to become a great paper. The updated draft is much clearer and answers some of my questions. Most importantly, the updated theory section explains why small beta can be bad: it slows training. While I appreciate the clarification, I think this argument still doesn't fully align with the experiment results. As the CIFAR-10 experiment points out, using small beta can still lead to good performance, so the main problem for small beta seems to be instability (rather than slow training), which the current theory couldn't explain. To improve the theory, I hope the paper can provide more insights into the instability caused by small beta. For example, I wonder if this is somehow connected to the slow training argument; perhaps the failed runs indeed suffer from slow training. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "SBf0u7pl_dy", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The soft-max layer of the neural network typically does not scale the outputs of the incoming layers, i.e., the scale is set to one. This paper analyses the how the dynamics of the training is effected by a non-unit scale $\\beta$. The mathematics involves elementary calculus, and it is clear and easy to follow. The analysis primarily concerns the initial/early phases of the training, or what the authors refer to as short times --- note to authors: I think there should be a better phrasing than using \"short times\" --- where $\\tau$ is small. The analysis is divided into two cases: linear dynamics where the Hessian/second-order-term is negligible , and the non-linear dynamics where it is not. The experiments are illustrating but inconclusive.\n\nAdditional comments and suggestions:\n1. Sections 2.2 to 2.3 can be restructured to more bring out that both the linear and non-linear dynamics concern the early phases of the training --- unless of course I misread these sections.\n2. Sections 3.1 and 3.2 does not use batch -normalisation. Although some readers can reasonably guess the reason, it is worthwhile to explicitly say why this is the case.\n3. I don't think section 3.3 contributes to the central message of the paper. Hence, I think it can be move to the appendix, and some from the appendix can be moved to the main paper.\n4. It is not mentioned how the first plot of Figure 2 and the two plots in Figure 3 are obtained.\n\nThe paper deserves an *accept* because it is fundamentally correct and it is one of the \"secrets-of-the-trade\", and I am glad that it is written. It is *not a stronger accept* because I feel the \"right units for comparison\" (see section 4) could already be common knowledge among those who has looked hard at soft-max.\n\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "O_TeZ-fVvqP", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nThe theory developed in this paper is rudimentary. I don't see how the theory in Section 2 is crucial for the understanding of the impact of $\\beta$. The results basically say proper normalization on the timescale is needed for a fair comparison among different $\\beta$. There are neither formal statements nor proper discussions on the significance of the results. \n\nFrom the experiment results, the performance is relatively insensitive to the parameter $\\beta$ with batch normalization. I don't see a clear motivation for dropping batch normalization or tuning $\\beta$ with batch normalization. It is apparent that tuning for the best $\\beta$ is always no worse than setting $\\beta=1$. But the impact of $\\beta$ seems to be insignificant, and the so-called optimal $\\beta$ varies wildly across different experiments. The performance with different $\\beta$ is not presented in Sections 3.2.2 and 3.2.3. \n\nFinally, the manuscript is poorly written and needs to be largely reworked to be considered for publication. None of the key concepts identified in the paper are formally defined. For example, it is unclear what is the mathematical meaning of $\\ll$ and $\\sim$ in Section 2.3, and what is the precise meaning of the phrase \"no longer be neglected\" in Section 2.4. Abbreviations like NTK are not introduced. Curves in the plots like Figure 6(c) dashed orange are not labeled.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Three reviewers are mildly positive, while one is negative. The substantive comments of the reviewers are consistent with each other; it is merely their evaluations that differ. 

One contribution of the paper is that it shows how using temperature tuning can yield similar accuracy to using batch normalization; this is useful because batch normalization is not always possible. The revised paper shows improvements, and we appreciate the engagement of the authors with the reviewer comments. However, there are remaining weaknesses such as a weak argument based on the empirical.results.

This paper can be improved based on the comments made by the reviewers. We encourage the authors to resubmit to a future venue.
Decision: Reject