Temperature check: theory and practice for training models with softmax-cross-entropy losses
Authors: Atish Agarwala
Samuel Stern Schoenholz
Jeffrey Pennington
Yann Dauphin
Venue: Accepted by TMLR ()

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\textbf{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, experimental 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 sensitivity to the random seed.

Published in Transactions on Machine Learning Research (03/2023)
Temperature check: theory and practice for training models with softmax-cross-entrop losses
Atish Agarwala
thetish eectlll@@oogle.ooo
Google Research, Brain Team
Jeffrey Pennington
spenninn
Google Research, Brain Team
Yann Dauphin
ynd com
Google Research, Brain Team
Sam Schoenholz OpenAl"
schsam@openai..omm
Reviewed on OpenReview: /ttps://opentss
// openreveew....e..mom net/ forum? id- LBA2Jj5Gqm
Abstract
The softmax function combined with cross-entropy loss is a principled approach to modeling probability distributions that has become ubiquitous in deep learning. The softmax function is defined by lone hyperparameter, the temperature, that is commonly set to or regarded 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 coftmax-cross-elropp loss and show that the learning dynamics depend crucially on the inverse-temperature 3 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 variety of model architectures trained on CIFAR10, ImageNet, and IMDB sentiment analysis. We find that generalization performance depends strongly the temperature, but only weakly the initial logit magnitude. We provide evidence that the dependence of generalization on B is not due to changes in model confidence, but is dynamical phenomenon. It follows that the addition of 3 as tunable hyperparameter is key to maximizing model performance. Although we find the optimal 33 to be sensitive to the architecture, experimental results suggest that luning 3 over the range 10-2 to 10t improves performance over all architectures studied. We find that smaller may lead to better peak performance the cost of sensitivity the random seed.
1
Introduction
Deep learning has led to breakthroughs across slew of classification tasks (eCunn 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- blass-probbbilttie 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 saltmax-crros-entropy learning. A particularly successful approach has been to improve overfitting by roducing model confidence; this has been done by regularizing outputs using confidence regularization (Pereyra 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.
* Work done while Google Brain.

---

Published in Transactions on Machine Learning Research (03/2023)
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 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-wei wide networks trained with softmax-cross-entroe 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-performing models operate best away from
loss have been
the linearized regime (Aitchison 2019 Novak et al.) 2019a). 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, investigate the training of neural networks with softmax-cross 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 3 = 1/T have identical learning curves if the learning rate is scaled as n n 33... We show that networks begin to learn over a timescale T3 ~ /2/ where Z are the inputs to the softmax. We also study the deviation from the linearized dynamics and show that there is another independent timescale Tut B/nj- Combined our analytical results suggest that at early times networks with smaller 3 become more non-linear over an equal change in loss. Previous results suggest that neural networks will perform best in this regime where they
losses. In general this
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. (016)) on ImageNet (Deng et al. [200)], and GRUs (Chung 2014) the IMDB sentiment analysis task (Maas et al.] 2011). In the case of residual networks, we
ct al. 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 | >> but otherwise largely independent of ||2'|22 Moreover, learning becomes slower and less stable at very small 3; indeed, the optimal learning rate scales like 17" ~ 1/3 and the resulting early learning timescale can be written as T ~ / For all models studied, our experiments show strong performance for 8 E [10 1,002,,000 10'J although the specific optimal / is architecture dependent. Emphatically, the optimal / is often far from 1. For models without batch normalization, smaller 3 can give stronger results 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, 10").
2 Theory
We begin with precise description of the problem setting before discussing a theory of learning at short times. We will show the following:
The inverse temperature 3 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- linear learning regime. We will see empirically that earlier access to the non-linear rogime 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 B will train more slowly.
2

---

Published in Transactions on Machine Learning Research (03/2023)
All numerical results in this section are using Wide Resnet Zagoruyko and Komodakis 2017 trained CIFAR10.
2.1 Basic model and notation
We consider a classification task with K classes. For an N dimensional input x, let z(x, 0) be the pre-softmax output of a classification model parameterized by 0 E RP, such that the classifier predicts the class corresponding to the largest output value 2, We will mainly consider 0 trained by SGD on a training set (X.)) of M input-label pairs. We focus on models trained with cross-entropy loss with a non-trivial inverse temperature B. The sottmax-cross-cntropy loss can be written as
K
K Σ y (Z, 0)))
((0,X,3) = Σ } (2
(1)
where we define the rescaled logils Z = 82 and oZ); = - e is the softmax function. Here Z(2', 0) is the M x K dimensional matrix of rescaled logits on the training set.
=
As we will see later, the statistics of individual a (Z); will have a strong inffuence on the learning dynamics. While the statistics of o (Z); intractable for intermediate magnitudes, || they can be understood in the limits of large and small || For fixed model z(x, 0), controls the certainty of the predicted probabilities. Values of / such that 3 < 1/||z//2 will give small values of || < 1, and the outputs of the softmax will be close to 1/K independent of i (the maximum-entrops distribution on K classes). Larger values of 3 such that B > 1/||z||2 will lead to large values of || /2 > 1; the resulting distribution has probability close to 1 on one label, and (exponentially) close to on the others.
The continuous time learning dynamics (exact in the limit of small learning rate) are given by:
K
(zz (2, (2,)(()
Σ θ = nB 2=
- (Z,,,()((X,0()))
(2)
00
for learning rate n- We will drop the explicit dependence of Z; from here onward, and will denote time dependence Z, (X)t) explicitly where needed.
In function space, the dynamics of the model outputs on an input x are given by
dx,(x)
K
n3 j=1
Σ
(x, -o(z(
(3)
where we define the M x M x K x K dimensional tensor Op. the empirical neural tangent kernel (NTK), as
(00),, ix,(x)
(22((A)
00
00
for class indices and j, which is block-diagonal in the infinite-wider limit.
From Equation 3 we see that the early-time dynamics in function space depend on B, the initial softmax input 7(X. 0) on the training set, and the initial Ôp. 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 3 and ||Z ||= = IZ(X' O)]] (the norm of the M x dimensional matrix of training logits), independent of Ôø. We note that though our theoretical quantities will be defined for full-batch gradient descent, because of self-averaging they will be relevant for SGD as well.
2.2 Linearized dynamics
For small changes 0, the tangent kernel approximately constant throughout learning (Jacot et al.) 2018). In this section, we study this linearized regime and drop the explicit dependence of the kernel. We note that networks reach this regime width N increases, long lim B(N)N (Chizat al. 2019
3

---

Published in Transactions on Machine Learning Research (03/2023)
The linearized dynamics of z(x, only depend the initial value of θ and the 8-scaled logit values ( t). This suggests that there is a universal timescale across 3 and which can be used to compare linearized trajectories with different parameter values. Indeed, if we define an effective learning rate ij = nn3, we have
dz,(x) dt
Σ
-z(2(22)))
which removes explicit 3 dependence of the dynamics. We note that similar rescaling exists for the continuous time versions of other optimizers like momentum (Appendix B
20-
20
3-10
3= 10
3-10"
Loss
1.5
Loss
1.5
3-10"
1.0
1.0
10
Linearized
0.5
0.5
10
10
10
10
109
10
10a
107
10'
10a
7j0
Figure 1: For fixed random model initialization, training while varying 7 and B only (left), learning curves collapse when plotted against fjt ntt (right). Learning curves to collapse to the learning curve of the lincarized model at early times, in contrast to un-scaled curves. Models with large / follow linearized dynamics the longest.
The effective learning rate is useful for understanding the nonlinear dynamics, plotting learning curves versus 7f causes carly-time collapse for fixed Z for single model varying only / and т) (Figure D We see that there is strong, monotonic, dependence of the time at which the nonlinear model begins to deviate from its linearization on 3. We will return to and explain this phenomenon in Section 2.4
Unless otherwise noted, we will analyze all timescales in units of ij instead of 7,, as will allow for the appropriate carly-time comparisons between models with different
2.3 Early learning timescale
We now define and compute the early learning timescale, Tz, that measures the time takes for the logits change significantly from their initial value. Specifically, we define Tx such that for < Ts we expect [Z(x, Z(x, 0) |r < (x(x0))le and for > Tz, [Zx, t) O)l, ~ |Z(x, (or larger). This is synonymous with the timescale over which the model begins to learn. As will show below, x | Therefore in units of îj, тı only depends on || and not B.
To see this, note that at very short times it follows from Equation 5 that
(x,xt) Z. (x,
Σ (0); ( (Z)(())))
(6)
It follows that can define timescale over which the logits (on the training set) change appreciably from their initial value
1r
(7)
7 OxX, X) - o (Z)))
where the norms are once again taken across all classes as well as training points. This definition has the desired properties for < T₂ and t > Tg.
In units of i T, depends only on |Z in two ways. The first is a linear scaling in || the second comes from the contribution from the gradient |iθ(X, x) a (Z(( 0)) As previously discussed, since (Z²)
4

---

Published in Transactions on Machine Learning Research (03/2023)
saturates at small and large values of |/ it follows that the gradient term will also saturate for large and small [ and the ratio of saturating values is some ((1) constant. independent of || and B.
Test accuracy vs. time, 1.00s
Test accuracy
10
0.18
12 I 10-1
0.18
I 10
F- 1.00
Accuracy
0.16
12"1|,-100
I 10
Accuracy
12
10
10
R 1.00x
1.00
|12*|
12
0.14
1.00x
B- 1.00s
02
10
10
10a
10"
105
16 10
10P
102
104
10P
||
f/7:
Figure 2: The timescale Tz depends only on || in units 811 (left, (. Ta depends linearly on F, up to an O(1) coefficient which saturates at large and small || (left, main). Accuracy increases more quickly for small initial |[, though late time dynamics are similar (center). Rescaling time to t/t, causes early accuracy curves to collapse (right).
The quantitative and conceptual nature of T, can both be confirmed numerically. We compute the empirical fz explicitly [x; 1) Zx(x, 0)]/4|| for short times (Figure 2 left). When plotted over a wide range of || and 3, the ratio f fz/l|z"|| || (in rescaled time units) undergoes saturating, O(1) variation from small to large [| The quantitative dependence of the transition on the NTK is confirmed in Appendix C Additionally, for lixed /3 and varying || rescaling time by 1/fz causes accuracy curves to collapse at early times (Figure 2 middle), even they are very different at early times without the rescaling (right). We note here that the late time accuracy curves seem similar across || without rescaling, a point which we will return to in Section 3.2
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, Tals corresponding to the time over which the network deviates appreciably from the linearized equations. We will show that Tat & B/tj. Therefore, in terms of 3 and [Z"] networks with small will access the nonlinear regime early learning, while networks with large 3 will be effectively lincarized throughout training. We note that similar point was raised in Chizat al. (2019) primarily in the context of MSE loss.
We define Fad to be the timescale over which the change in Ôg (which contributes to the second order term in Equation can no longer be neglected. Examining the second time derivative of Z, we have
dzz
K
-(0o),) x) o(((())) d [00))(x (2) oZZ
Σ
(8)
dt2
dt
linearized dynamics
nonlincarized dynamics i
The first term is the second derivative under fixed kernel, while the second term is due the change in the kernel (neglected in the lincarized limit). A direct calculation shows that the second term, which we denote Znt; can be written as
K
(zz)(X) (Qru6) 00
~m))=B-
( (2,(2))o(Zz,()) oZZ)(X))
(9)
00
This gives us a nonlinear timescale Tud defined, at initialization, by Tud || // /1. interpret Taa the time it takes for changes in the kernel to contribute to learning.
We can
5

---

Published in Transactions on Machine Learning Research (03/2023)
Though computing [Z O)lF exactly is analytically intractable, its basic scaling in terms of 3 and |Z 12 (and therefore, that of Tnt) is computable. We first note the explicit 3- dependence. The remaining terms are independent of 3 and vary by at most (1)) with || indeed as described above, - 0)) saturates for large and small [| Morevoer, the derivative, 0z(xx.0) is the square rool of the NTK and, at 88
initialization, it is independent of || Together with our analysis of Ta we have that, up to some O(1) dependence on ||e,, Tad x 8/n, Therefore, the degree of nonlincarity early in learning is controlled via B alone.
105
10"
10
10a
B= 10
Tm
102
Tall
40
0=
== 10
3x100
== 10
101
== 10a
2x 10'
10
- 10a
B= 10a
10
105
10
10
101
105
1
10
101
|
||
Figure 3: The time to deviation from linearized dynamics, Tat, has large deviation over 8 and || (left), which can be largely explained by linear dependence on 3 (right), in units of fj = B²1 There an 0(1) dependence on || which is consistent across varying / for fixed ||
Once again we can confirm the quantitative and conceptual understanding of Tnd mumerically. Qualitatively, we see that for fixed [ models with smaller 3 deviate sooner from the linearized dynamics when learning curves are plotted against ijt (Figure D We compute the empirical Fut by taking [l Zun (f)]/{] at small
times, where Zin (t) is the solution Equation predicted by the theory.
T has an (1)) dependence on | [| only (Figure)
2.5 Learning rates and learning speed
The timescales, Ty and Tat, along with the rescaled learning rate îj, can be combined gain information about training with SGD. We know that changes in the loss proportional to f early times. The ratio Tz/Tnt tells us something about how much the network deviates from the linearized dynamics by the time the softmax inputs (and therefore the loss) have changed by an appreciable amount. For fixed Z we have Tz/TaI a/3 for a 3-independent constant a. This suggests that as we decrease 8, the network becomes more non-linear for an equivalent change in the loss.
Another key point is that largest allowable lcarning rate is controlled by the curvature of the loss function Du al. (2019); Allen-Zhu et al. (2019). necessary condition for training is that the curvature be small compared with the step size, which happens when Tnd S where c is architecture-depnnae O(1) which has not yet been calculated theoretically (see Lewkowyez et al. (2020) for empirical calculations of c). This predicts a maximum effective learning rate vj of 0(3) which in turn implies a raw learning rate of O(8-1) (as 1) 32)
constant of
If ij is O then Te O(3-1). This means that, for carly learning, networks with smaller 3 will take more SGD steps to reach an appreciable change in logits. Therefore, we predict that networks with small / take longer to train (which in practice as well).
6

---

Published in Transactions on Machine Learning Research (03/2023)
Optimal learning rate vs. /
10
0
1
10
10a
104
10
102
10
10
10'
10
10a
Figure 4: Optimal learning rate 17" for WRN on CIFAR10 scales as 1/3.
3
Experimental Results
3.1 Optimal Learning Rate
We begin our empirical investigation by training wide resnets (Szegody et al. 2016) without batch normalization on CIFAR10, as the plots from Figure I suggest that this model can interpolate between the lincarized and non-linear regimes. In order understand the effects of the different timescales on learning, we control B and || independently by using 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 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 B. The optimal learning rate '7" has a clear 1/3 dependence (Figure 4 for WRV with BatchNorm... 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 stiill converges. Again, predicted, networks with smaller 8 also learn more slowly in terms of number of SGD steps.
Thus, at small 3 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 a significant amount before the dynamics become appreciably nonlinear.
3.2
Phase plane
In the precoding discussion two quantitics emerged that control the behavior of carly-time dynamics: the svvrse-temperature, B, and the rescaled logits |Z In attempting to understand the behavior of real neural networks trained using coftmax-cross-entrqqy loss, it therefore makes sense to try to reason about this behavior by considering neural networks that span the 3 1 phase plane, the space of allowable pairs (8. [l))) . By construction, the phase plane is characterized by the timescales involved in early learning. To summarize, Tz ~ 1 / sets the timescale for early learning, with larger values of || leading to longer time before significant accuracy gains are made (Section 2.3 Mcanwhile, Tat ~ B/n controls the timescale for learning dynamics to leave the linearized regime with small 3 leading immediate departures from linearity, while models with large 3 may stay linearized throughout their learning trajectories (Section 2.4
In Figure 5 (a), we show schematic of the phase planc. The colormap shows the test performance of a wide residual welw (Zagoruyko and Komodakis) 2017), without batch normalization, trained on CIFAR10 in different parts the phase plane. The value of makes a large difference in generalization, with optimal performance achieved at 10-2, In general, larger 3 performed worse than small as expected. Morcover, we observe similar generalization for all sufficiently large B: this is be expected since models in this regime close to their linearization throughout. training (see Figure] 1 and we expect the linearized models to
7

---

Published in Transactions on Machine Learning Research (03/2023)
(b)
Seoe (20 (3909
(a)
1.00
|| 0:
Slow
10a
Poor conditioning
10a
3011
-0.95
101
Early learning
Accuracy
181
107
127
8.09
10-0
0.90
(c)
x
224
10
4
10
Fast
10-1
10-7
10-1
10
10
16
10
0.40
Tal Xc 8/9
Less
More
Linearity
Figure 5: Properties of early leariing dynamics, which affect generalization, can be determined by location in the B- 8-|| phase plane (a). At optimal learning rate 17,, small / and larger | leads to slower early learning (b), and larger 3 increases time before nonlinear dynamics contributes to lcarning. Large || has poorly conditioned linearized dynamics. Gemeralization for a wide resnet trained on CIFARIO is highly sensitive to B, and relatively insensitive to ||Z|| outside poor conditioning regime. Final logit variance is relatively insensitive to paramcters (c).
have (-inddpendent performance. Generalization was largely insensitive to || lr so long as the network was sulficiently well-conditioned to be trainable. This suggests that long term learning is insensitive to
In Figure (b), we plot the accuracy after steps of optimization (with the optimal learning rate). For fixed |L,,,, the training speed was slow for the smallest B and then bocame faster with incrcasing B. For fixed 3 the training speed was fastest for small || F and slowed || increased. Both these phenomena were predicted by our theory and shows that both parameters are important in determining the early-time dynamics. However, note that the relative accuracy across the phase plane early times did not correlate with the performance at late times.
This highlights that differences generalization are a dynamical phenomenon. Another indication of this fact is that at the end of training, at time fy. the final training set logit values | F = [| (y) tend towards independent of the initial I and || (Figure (c)). With the exception of the poorly-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 cerlainly (a typical motivation for luning B).
3.3 Architecture Dependence of the Optimal
Having demonstrated that 8 controls the generalization performance of neural networks with softmax- cross-entropy loss, we now discuss the question of choosing the optimal B. We note that without the special correlated initialization, || is generally proportional to 3. This means that large leads poor conditioning and training in practice. However, small 3 have small initial || and they avoid the conditioning issues early on. We also saw that ||Z*|| eventually reaches some O(1) value over wide range of B. Additionally, is relatively insensitive to |. Our results in the previous section suggest that Tut is more important for generalization; this leads to the conjecture that tuning / may have interesting effects in practical models.
We investigate this hypothesis over number of different architectures. We find the optimal choice of B be strongly architecture dependent. Whether or not the optimal 3 can be predicted analytically an open question that we leave for future work. Nonetheless, show that all architectures considered display optimal
8

---

Published in Transactions on Machine Learning Research (03/2023)
3 between approximately 10 and 10'. We observe that by taking the time to tune 8 it is often the case that performance can be improved over the naive setting of =
3.3.1
Wide Resnet on CIFAR10
(a)
Accuracy average
(b)
Accuracy vs. å
(c)
1.00
0.88
Accuracy
0.90
Accuracy
Accuracy
0.88
0.85
0.85
0.84
1:00
0.80
1.11
0.80
1.11
0.82
1.23
0.80
Max over soeds
0.75
1.57
0.75
1.37
Averago seeds
0.70
0.70
DN
DN
0.65
0.65
/
Figure 6: Dependence of test accuracy for various architectures with 3 tuning. (a) For WRN with batchnorm, trained CIFAR10, the optimal B z 10. Without batchnorm, the performance of the network can be nearly recovered with 8-scaling alone with B = 10-2 Even poorly conditioned networks (achieved by increasing weight scale ow)) recover performance. (b) For 3 10-2, learning is more variable, 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 3 = but peak performance is small B.
In Figure 6 (a) we show the accuracy against for several wide residual networks whose weights are drawn from normal distributions of different variances, o trained without batchnorm, well as network with of trained with batchnorm (averaged over 10 seeds). The best average performance attained for 8 < 1, Tw without batchnorm, and in particular networks with large du dramatically improved with B tuning. The network with batchnorm is better at all 8, with optimal B ~ 10. However, see that the best performing seed is often at lower 8 (Figure 6 (b)), with larger networks competitive with ow = 1, and even with batchnorm at fixed B (though batchnorm with B = 10 still performs the best). This suggests that small 3 can improve best case performance, at the cost of introducing variability. Our results emphasize the importance of tuning B, especially for models that have not otherwise been optimized.
3.3.2 Resnet50 on ImageNet
Table 1: Accuracy on Imagenet dataset for ResNet-50. Tuning 3 significantly improves accuracy.
Method
Accuracy (%)
ResNet-50
GGhiasi et al. 2018
76.51 + 0.07
ResNet-50
Dropout Ghiasi et al. 2018
76.80 + 0.04
ResNet-50
Label Smoothing Ghiasi et al. 2018
77.17 + 0.05
Re
Temperature check (p = 0.3)
77.37 + 0.02
Motivated by our results CIFAR10, experimentally explored the effects of 3 tunable hyperparameter for ResNet-50 trained 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 shows that scaling 3 improves accuracy for ResNet-50 with BatchNorm. However, we did not find that using < was optimal for ResNett-50 without normalization. This further emphasizes the subtle architecture dependence that warrants further study.

---

Published in Transactions on Machine Learning Research (03/2023)
3.3.3 GRUs on IMDB Sentiment Analysis
To further explore the architecture dependence of optimal 3, 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 peak performance near B ~ I averaged over ensemble of works,,s, but we observe that smaller 3 can give better optimal performance at the expense of variability over initialization.
3.4 Proposed tuning procedure
Our results suggest the following tuning procedure for networks trained with
Train/tune model as normal. Note the optimal learning rate yo-
For the best parameter set, sweep over B E [10-2, 104), scaling learning rate nn/B.
If best performing 3 is at an endpoint of 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 B > 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.
4 Discussion
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 CIFAR10 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 be made by more carefully tuning B jointly with other hyperparameters, in particular the learning rate schedule and batch size. Our theory does not predict the optimal 3.
One key lesson of our theoretical work is that properties of learning dynamics must be compared using the right units. Other work has shown that learning curves with the same learning rate may not be comparable as the width of network is varied (Yang el al., 2022). Without the correct normalization, it is easy to draw incorrect conclusions. For example, in our work, we see that Tat X 1/3n, which at first glance suggests that models with smaller 3 will become nonlinear more slowly than their large & counterparts. However, analyzing Tht with respect to the effective learning rate f 33n yields Tna x B/nj. (Note that analogous rescaling strategy can align the learning curves in (Ynng et. al. 2022 there the rescaled learning rate is = 11/ n/TT,
where n is the width of the network.)
In fact, networks with smaller 3 tend to deviate more from the lincarized regime even when learning curves are rescaled to maximize overlap at early times, contrast. to ner with large B which with the right. initialization can remain the lincarized regime throughout training (Figure 1 right).. There is other work which applies a particular 3 scaling to get more non-linear behavior; in particular, it has been shown that the width of network is taken to infinity, choosing 3 3= 1/n gives a well-defined infinite-width limit
which is distinct from the lincarized NTK regime (Yang 2021 Yang et al. 2022 Our work shows that
manipulating 3 is useful and interesting in the finite-width regime as well.
Our analysis suggests a range of good B, but does not predict the optimal value. Indeed the optimal magnitude may be less than or greater than 1 depending on architecture and choice of other hyperparameters. Architettrr-ddppadeent, non-scaling multiplicative factors to key learning parameters have been observed in other contexts (Lewkowyzz et al. 2020), and their numerical estimation is in general difficult. Extending
10

---

Published in Transactions on Machine Learning Research (03/2023)
the theoretical results to 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.
References
Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L.. D. Jackel. Backpropa- gation Applied to Handwritten Zip Code Recognition. Neural Computation, 1(4):541 551, December 1989. ISSN 0899-7667. doi: n0.116/neeo..198.111. wtt..o........... 1989.1.4.541.
Alex Krizhevsky, Πya Sutskever, and GeofTrey E Hinton. ImageNet Classification with Deep Convolutional Neural Networks. In Advances in Neurol Information Processing Systems 25, pages 1097-1105. Curran Associates, Inc., 2012.
Sergey Zagoruyko and Nikos Komodakis. Wide Residual Networks. arXiv: 1005.01100 [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-31 -18, December 2005. ISSN 1433-3058. doi: 10.
Pavel Golik, Patrick Doetsch, and Hermann Ney. Cross-Entropy vs. 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 Geolfrey Hinton. Regularizing Neural Networks by Penalizing Confident Output Distributions. arXiv:17010065/4 [cs]. January 2017.
Rafael Müller, Simon Kornblith, and Geoffrey E Hinton. When does label smoothing help? In Advances in 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 Pallern Recognition, pages 2818 2826, 2016.
Zhirong Wu, Alexei A. Efros, and Stella X. Yu. Improving Generalization via Scalable Neighborhood Component Analysis. ar 1xXx:1808.04699 [es), August 2018.
Andrew Zhai and Hao-Yu Yu Wu. Classification is a Strong Baseline for Deep Metric Learning. arXi: [es], August 2019.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the Knowledge in a Neural Network. arXio:1503.025: [es, statf, 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 modern neural networks. In Proceedings of the 34th International Conference on Machine Learning Volume 70, ICML/17, pages 1321-1330, Sydncy, NSW, Australia, August 2017. JMLR.org.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural Tangent Kernel: Convergence and Goneralization 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 P6ezos, and Aarti Singh. Gradient Descent. Provably Optimizes Over- parameterized Neural Networks. In 7th International Conference Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
Jachoon Lee, Lechao Xiao, Samuel Schocnholz, 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 Neural Information Processing Systems 32, pages 8570 8570-8581. Curran Associates, Inc., 2019.
11

---

Published in Transactions on Machine Learning Research (03/2023)
Roman Novak, Lechao Xiao, Jachoon Lee, Yasaman Bahri, Greg Yang, Jiri Hron, Daniel A. Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes. arXic:1810.0511 (cs, statj, February 2019a.
Lechao Xiao, Jeffrey Pennington, and Samuel S. Schoenholz. Disentangling trainability and generalization in deep learning. arXix:1113 13053 fcs, 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.04: 04486 jes, math, stat), March 2020.
Laurence Aitchison. Why bigger is not always better: On finite and infinite neural networks. 191:::1910.08013 (es, statj, November 2019.
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On Lazy Training in Differentiable Programming. Tn Advances in Neural Information Processing Systems 32, pages 2937-2947. Curran Associates, Inc., 2019.
Jaehoon 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 Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 770-778, 2016.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hicrarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248 255, June 2009, doi: 111109/CVP.2000.52222
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling. arxiv:14 WYv::412.338 112.3555 [es], 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 1, HLT '11, pages 142-150, USA, June 2011. Association for Computational Linguistics. ISBN 978-1-932432-87-9.
Sergey Toffe and Christian Szegody. 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- Paramoterization. In International Conference on Machine Learning, pages 242-252. PMLR, May 2019.
Golnaz Ghiasi, Tsung-Yi Lin, and Quoc V Le. DropBlock: A rogularization mothod for convolutional networks. 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. Tn 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 Snock, Tim Salimans, Rodolphe Jenatton, and Scbastian Nowozin. How Good is the Bayes Postcrior Deep Neural Networks Really? 2Xx::2002.02405 [es, statj, July 2020.
Greg Yang. Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, March 2022.
12

---

Published in Transactions on Machine Learning Research (03/2023)
Greg Yang. Tensor Programs I: Wide Feedforward Recurrent Neural Networks of Any Architecture are Gaussian Processes. Xrxi::1100 12478 (cond-mat, physicssmath-ph) May 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. arXic:1912.02810 fes, stat/, December 2019b.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman-Milne, JAX: Composable transformations of Python+NummP programs, 2018.
13

---

Published in Transactions on Machine Learning Research (03/2023)
A Linearized learning dynamics
A.1 Fixed points
For the linearized learning dynamics, the trajectory z(x, 6) can be written in terms of the trajectories of the training set as
zxx,t) z(x, 0) = é(x, X) (x, (z( z(),)))
(10)
where + is the pseudo-inverse. Therefore, if one can solve for z(X,t), (), then in principle properties of generalization computable.
However, in general Equation 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 the training objective. Given a regularizer _\el16011 on the change parameters 50 = (t) 0(0), the dynamics in the lincarized regime are given by
3 does not admit analytic solution even for fixed 6, in contrast to the case
2(x) - Br0 x) o(3x(x))) \odz(x)
(11)
where the last term comes from the fact that
8z ag 60 6z(x) in the linearized limit.
We can write down self-consistent equations for equilibria, which are approximately solvable in certain limits. For an arbitrary imput x, the equilibrium solution z (x) is
0 = (,0 x) - (B2 (x))) - Aadz" (x)
(12)
This can be rewritten in terms of the training set as
8z" = (Xx,
(13)
similar to kernel learning.
It remains then to solve for z" (X). We have:
8z' 6x' 8
-0(22
(14)
We immediately note that the solution depends on the initialization. We assume z(x, 0) = 0, so oz z order simplify the analysis. The easiest case analyze when [|]3 < 1. Then we have:
8 x(xx Xg
z' xx)x)
(1 Bz (x))
(15)
K
which gives us
3
3
xx(x)
1
6(x
()xx) - 1/K)
(16)
X
Therefore the self-consistency condition for this solution is || h1 Ô11 || 1, which simplifies the solution
z' z'(x) = Ao (1x1))
(17)
This is equivalent to the solution after single step of (full-batch) SGD with appropriate learning rate. We note that unlike linearized dynamics with L2 loss and full-rank kernel, there is no guarantce that the solution converges to 0 training error.
The other natural limit is || (X) > 1. We focus on the class case, in order to take advantage of the conserved quantity of learning with cross-cntropy loss. We note that the vector on the right hand side of Equation 3 sums to for every training point. Suppose at initialization, Ô, has no logit-logit interactions, is the case for most architectures in the infinite width limit with random initialization. More formally, we
14

---

Published in Transactions on Machine Learning Research (03/2023)
can write θ. Id & θ₂ where θ₂ is M x M. Then, the sum of the logits for any input x is conserved during linearized training, as we have:
1 (xx() n11TT [ @ - (Bz(21))
(18)
Multiplying the right hand side through, we gel
1 1'((x) = n, ( =(z((X))) 0
(19)
(Note that if θg has explicit dependence on the logits, there still is a conserved quantity, which more complicated to compute.)
Now we analyze (x)]( >> 1. With two classes, and x(X) 0 at initialization, we have z = = Therefore, without loss of generality, we focus zi, the logit of the first class. In this limit, the leading order correction to the softmax is approximately:
oB22)) se 12;20 - sign(z])) 23/2,1
(20)
The self-consistency equation is then:
- 0(x (1x)) 12320 + sign(zi)e sign(zije-221eii) 201zil
zi(x) = Ao
(21)
The vector on the right hand side has entries that are ((e-2/121) for correct classifications, and ((1) for incorrect ones, If we assume that the training error is 0, then we have:
3 (sign)))sgnn 280111
(x)
(22)
Xa
This is still non-trivial to solve, but we see that the self consistency condition is that ln(31]@]l=/>o) > 1.
Here also may be difficult to train and generalize well. The individual elements of the hight-hand-sing vector are broadly distributed due to the exponential so the outputs of the model are sensitive to ro/my only depend a small number of datapoints. Even if the equilibrium solution has training loss, generalization error may be high for the same reasons.
This suggests that even for NTK learning (with L2 regularization), the scale of ||5z|| 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 docsn't matter; rather, the ratio 3 and Ao (compared to the appropriate norm of Ô) needs to be balanced lo prevent falling into the small fz regime (where training error might be large) or the large Iz 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) z(x) z (X) for small perturbations is given by
a(x)= = [n [I @ (3), + So} 2(x)
(23)
where a is the derivative of the softmax matrix
Or Oz( Oz
diag(a(2)) (z))
(24)
We can perform some analysis in the large and small 3 cases (once again ignoring A²), For small ||,3z* (xzz)(()) (2²) [[F,
we have |l % Q1|æ & 1 which leads to:
x -11K 0
(25)
15

---

Published in Transactions on Machine Learning Research (03/2023)
This matrix has K - eigenvalues with value 1/K, and one zero cigenvalue (corresponding to the conservation of probability). Therefore || @ , 0 (82" (X)) & λo, and the well-conditioned regularizer dominates the approach to equilibrium.
In the large case (In($1(611/2o) > 1), the values of ((3((z)) are exponentially close to 0 (K - values) or (the value corresponding the largest logit). This means that (=((2()) has exponentially small values in [| if any one of (X), and ((2z(( is exponentially small, the corresponding element of (3z(()) is well; for the largest logit the diagonal is (f2z((X)) a (32,()))) which is also exponentially small.
From Equation 22 we have Ag < f2g28lii,; therefore, though the o; 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.
A.3 Conditioning of dynamics
Understanding the conditioning the linearized dynamics requires understanding the spectrum of the Hessian matrix H Id2 @ ((x o= (32' (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 x 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 matrices. The condition number s(M)
And
AM
has bound
<< x(B)/k(A) S ((M) S x(A)x(B)
(26)
Proof: Consider the vector v that is the eigenvector of B associated with AB, Note that ||Av|//|||| <
A Analogously, for w, the cigenvector associated with AB,max |wwl//||/| 2 A This gives us the two bounds:
Min S B,in AM,max AA AB.
(27)
This means that the condition number k(H)
max is bounded by Add
AA max B 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 (fzz (X)) will lead to poor conditioning 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 [| > 1.
A.3.1 Small logit conditioning
For [| & 1, the Hessian H is
(1 1 = K K 117) & ((x,x))
1
(29)
Since H is the Kroenecker product of two matrices, the condition numbers multiply, and we have
R HH) = x(A)
(30)
which is well-conditioned so long as the NTK is. Regardless, the well-conditioned regularization due to X dominates the approach to equilibrium.
16

---

Published in Transactions on Machine Learning Research (03/2023)
A.3.2 Large logit conditioning
Now consider (X)]l= >> 1. Here we will show that the linearized dynamics is poorly conditioned, and that s(H) is exponentially large B.
We first try to understand a (32"(x)) for individual x X. To Oth order (in an as-of-yet-udeefned expansion), o₂ 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 |,22 [p; the entries will have broad, log-normal type distributions with magnitudes which scale as exp(-3lzil). There will be two scaling regimes one with a small number of labels the sense vB > Vln(K, where the largest logit dominates the statistics, and one where the number of labels large (and the central limit theorem applies to the partition function). In both cases, however, there still exponential dependence on B3 we will focus on the first which is easier analyze and more realistic (e.g. for 10 labels "large" B is only 15).
Let 21 be the largest of K logits, 22 the second largest, and so on. Then using Equation 24 have:
(az)a = -(te----
(31)
for # 1,
(0)) = -yee-1as--1 D(e)-aa) - -0-
(32)
for + j and
(0=)1== = -(arr
(33)
The eigenvectors and eigenvalues can be approximately computed as:
1
(v1)1 = V2
(v1)2 =
d -2-- 9(21-22)
(34)
V2
(v2)1 V
1
(v2)2
1
\2
1
-23(21 22)
(35)
V2¹
2
and for i >>> 2,
(v.), 1, (v,)1 #(2-22) a ((21=))
(36)
with all non-explicit eigenvector components This expansion is valid provided that 3/K > (2)) > (1)))
(so that
Therefore the spectrum of any individual block a= (B((())) is exponentially small in B. Using the bound in the Lemma, we have:
H 2 (10 /
(37)
This is a very loose bound, as assumes that the largest eigendirections of a₂ are aligned with the smallest. eigendirections of Id, & O, and vice versa. is possible s(H) is closer in magnitude to the upper bound 222 =x) =x)==xx)xx( x)),
Regardless, n (H) is exponentially large in 3 meaning that the conditioning is exponentially poor for large [|z" |F.
B SGD and momentum rescalings
B.1 Discrete equations
Consider full-batch SGD training. The update equations for the parameters 0 are:
o O, VoL
(38)
We will denote g: = VoL for ease of notation.
Training with momentum, the equations of motion are given by:
Ve1 = (1 - -)v, gr
(39)
17

---

Published in Transactions on Machine Learning Research (03/2023)
0 02 + 7/4++1
(40)
where y 1].
One key point to consider later will be the relative magnitude Ao of updates to the parameters. For SGD, the magnitude of updates is n nllg|l For momentum with slowly-varying gradients the magnitude is nlgll//>
B.2 Continuous time equations
We can write down the continuous time version of the learning dynamics follows. For SGD, for small learning rates we have:
d0 7g
(41)
dt
For the momentum equations we have
dv
VV g d0 r/v
(42)
dt
(43)
dt
From these equations, we can see that in the contimuous 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 r = nyt causes learning curves to be identical for all learning rates.
For momentum, instead of a single universal learning curve, there is one-parameter family of curves controlled by the ratio Tiom n/?2. Consider rescaling time to at and = bv, where and b will be chosen to put the equations in canonical form. In our new coordinates, we have
du
dr = (7/a)y (b/a)g
(44)
d0
= m/bbbb)
(45)
dr
The canonical form we choose is
dv
-u -
(46)
dr
d0
(47)
dr
From which we arrive at a = b VT. which gives us X = 7/.
Note that this not a unique canonical form; for example, we a coefficient of -1 on v, we end up with
dv - -
(48)
d0
(49)
dr
with a y. This different time rescaling, but still controlled by Tmome
Working in the canonical form of Equations 46 and 47
analyze the dynamics. One immodiate question is the difference between À & 1 and X > 1. We note that the integral equation
((7)= = v(0) + [ c-Mrr'')( g(r))dr'
(50)
18

---

Published in Transactions on Machine Learning Research (03/2023)
solves the differential equation for Therefore, for À > 1, v(t) only depends the current value g(t) and we have ((7) g(T)/X. Therefore, we have, approximately:
d0
1 A
(51)
dr
This means that for large À all the curves will approximately collapse, with timescale given by vnA y (dynamics similar to SGD).
For Ả & 1, the momentum is essentially the integrated gradient across all time. If v() = 0, then we have
dr J g(T')dr'
de
(52)
In this limit, ( is the double integral of the gradient with respect to time.
Given the form the controlling parameter Toom; we can choose to parameterize 7 = 5v. Under this parameterization, have T mom = 572 The dynamical equations then become:
dv 5v
(53)
dr
d0
(54)
dr
which automatically removes explicit dependence 17- One particular regime of interest is the early-time dynamics, starting from =(0) = 0. Integrating directly, we have:
ig
(55)
This means that t alone is the correct timescale for early learning, at least until v5 ~ 1 which in the original parameters corresponds to 1/y (the time it takes for the momentum to be first "fully integrated").
B.3 Detailed analysis of momentum timescales
One important subtlety is that not the correct value to compare to. The real timescale involved the one over which g changes significantly. We can approximate this the following way. Suppose that there is some relative change [[0] timescale over which 0 changes by that amount is the one we must compare À to.
Ae of the parameters that leads to an appreciable relative change in g. Then the
We can compute that timescale in the following way. We assume g fixed for what follows. Therefore, Equation 51 approximately holds. The timescale T, of the change is then given by:
A
g
0
To
(56)
which gives
To ~ |
(57)
In particular, this means that the approximation is good when Xre >> 1, which gives 72/n > being a function of the dynamical parameters, the latter being function of the geometry of L with respect to 0.
0 the former
One consequence of this analysis is that if the ||]| remains roughly constant, for fixed n and y, late learning when the gradients become small the dynamics shifts into the regime where À is large, and we effectively have SGD.
19

---

Published in Transactions on Machine Learning Research (03/2023)
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 contimuous expectation can be informative.
With momentum, we only have perfect collapse when y and rj 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 y (more precisely, 7 > |gll/||[|]),
may be the case at later training times, the change in parameters from a single step is Ag ~ (n/r)) ln)r)))lll and the rescaling of taking to npt (as for SGD) collapses the dynamics. Therefore given collection of varying 7, but fixed y curves, it is possible to get intermediate and late time dynamics on the same scale.
However, at early times, while the momentum still getting "up speed' (i.e. the first 1/7 steps), the appropriate timescale is n 11/2, Therefore, in order get learning curves collapse across different 1y at early times, we need to rescale 7 with 1/ as implied by Equations 46 and 47 Namely, one must fix and rescale y タ. We note that, since 7 < 1, this gives us restriction 1) < 7-2 for the maximum learning that can be supported by the rescaled momentum parameter.
B.5 Momentum equations with softmax-cross-entropy learning
For cross-entropy learning with softmax inputs fz, all the scales acquire dependence on B. If define then we have, approximately, ||g|| R Blag..
4
=
ac J32
and I=
= 8z
D0
Consider the goal of obtaining identical early-time learning curves for different values 3. (The curves are only globally consistent across in the linearized regime.) In order to get learning curves collapse, we want dc to be independent of in the rescaled time units. We note that the change in the loss function Ac from single step of SGD goes as
Dc n3
(58)
This suggests that one way to collapse learning curves is to plot them against the rescaled learning rate 7j, where î = = While hyperparameter tuning across 3, one could use 17 = 3², sweeping over ń in order to easily obtain comparable learning curves.
However, better goal for a learning rate rescaling is try and stay within the continuous time limit that is, to control the change in parameters Ae for single step to be small across B. We have
Ao nBL.9=
(59)
which suggests that maximum allowable learning rates will scale as 1/8. This suggests setting 1) hh3 and rescaling time as f/3 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
dv -ÿiv 0gg
(60)
dr
d0
(61)
dr
dL
3g 1
(62)
dr
where g
BC,
Rescaling 7 by 3 gets us:
832 50
dv
3
(63)
d3r
B
20

---

Published in Transactions on Machine Learning Research (03/2023)
d0
(64)
d3T 3
dC
dßT g
(65)
This rescaling causes a collapse of the trajectories of the at early times if 5/3 constant for varying 3
One scheme to arrive at the final canonical form, across 3, by the following definitions of n 7, and
7 = n
= B7 = ViA
B Vit
where curves with fixed 7 will collapse. The latter two equations similar to before, except with y) replaced with ij. The dynamical equations are then:
dv
.v g
(66)
dr
d0
(67)
dt 3
dC
g
(68)
dr
The change in parameters from single step (assuming constant and saturation) is
Ag
(69)
7,
If we instead want the change in parameters from a single step to be invariant of B so the continuous time approximation holds, while maintaining collapse of trajectories, we first note that
V Blag=
Do
(70)
from a single step of the momentum optimizer. To keep Ag invariant of B, we can set:
-13-1 - = 3\12
B Viñi 32/2 Vini
Note that the relationship between and 7 is the same in both schemes when measured with respect to the raw learning rale v).
C Softmaxx-coss-entroop gradient magnitude
C.1 Magnitude of gradients in fully-connected networks
The value T2 has nontrivial (but bounded) dependence on ||Z||e via the |ς² - (Z))) term Equation 7 We can confirm the dependence for highly overparameterized models by using the theoretical é. In particular, for wide neural networks, the tangent kernel block-diagonal in the logits, and easily computable.
21

---

Published in Transactions on Machine Learning Research (03/2023)
The numerically computed 7z/ /Z|l/ correlates well with | (Z for wide (2000 hidden units/layen fully connected networks (Figure 7 The rat 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)..
Nommliized 7.. Relu
Normalized T,. Erf
Normalized Erf (shuftled labels)
2x 107
6x10
10R
107
6x100
12
4x10
3x
13 184"
16+
"0"""
for
16
10P
16 3x10"
164
10"
105
|
||
||2" 11,
Figure 7: 1l is highly correlated with 10(0' (2 with ê computed in the infinite width limit (in units of effective learning rale 7 B2)) Ratio between normalized limescales at large and small ||0|| depends on noulinearity (left and middle), as well as training set statistics (right, CIFAR10 with shuffled labels).
Therefore in general the ratio of 12/1|/0/l at large and small |Z 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.
Experimental details
D.1 Correlated initialization
In order avoid confounding effects of changing and || with changes é, we use a corvelated initialization strategy (similar (Chizat al. 2019 which fixes é while allowing for independent variation of 3 and |/ Given network with final hidden layer h(x, 0) and output weights W oo, define combined network (x, 0) explicitly
ze( Wo,1, 0o.2,0,,00 = Wo, (x, + Wo.h((x 02)
(71)
where, at initialization, 0, 02 and the elements of W o,a have statistics
deddz for a = b
El El(Wo.))
(72)
c. Gudz
for a + b
for correlation coofficient e [-1, where d, is the Kronockor-detaa which is 1 i = j and 0 otherwise. Under this approach, the initial magnitude of the training set logits is given by |/ /1 + c)]l,"|1r, where |F is the initial magnitude of the logits of the base model. By manipulating 3 and we can independently change @ and || with the caveat that || S /2 since c 1. It follows that the small B, large | r region of the phase plane (upper left in Figure 5 is inaccessible with most vell-conditioned models at initialization. If only train one set of weights, the Ô is independent of
where ||z||]
D.2 Wide ResNet (WRN) experiments
All empirical studies Sections 2 and 3.2 involve training Wide ResNet (Zagoruyko and Komodakis 2017 with no BatchNorm on CIFAR10. The main hyperparameters of the networks can be found Tables and respectively. In addition we used the correlated initialization strategy above to fix &.
22

---

Published in Transactions on Machine Learning Research (03/2023)
trained GPUs with learning rates 10-5/3 for 8 E [10-3, 11103,110"") 10100"] and 10 10 E [10", 10"). This is smaller than the optimal learning rate, and was chosen to cleanly compare the trajectories to the linearized learning trajectory. The linearized trajectory was generated using the Neural Tangents library Novak et al. (019bb)
For Figure
For Figures and we trained on GPU using float64 precision with learning rate 10-2/3. For the phase plane in Figure we trained on TPUs. We swept over a learning rate range [10-3.5/8, 10-2/3] for each point in the grid and report the result for the optimal learning rate in that range. We found that without BatchNorm, learning rate of 10 1052.5/3 was most commonly optimal. In the few cases where the optimal value during the grid search was 10 10-2/3, we individually validated that a larger learning rate gave worse performance.
Figure was obtained by training WRN with BatchNorm, using the geometry from Table B and performing learning rate sweep centered at 10-2/// 10-2/B.
All networks were implemented in JAX (Bradbury et. al. 2018 and experiments involving linearization
direct computation of the NTK used Neural Tangents NNvaa et al. 2019b
Parameter
Value
Batch size
4096
Block type
B22
Blocks per group (1)
Widening factor (k)
Number of groups (d) e regularization
5 5010
Table 2: Fixed parameters for wide ResNet experiments in Section
Paramcter
Value
Batch size
4096
Block type
B22 2)
Blocks per group (1)
4
Widening factor (k)
10
Number of groups (d) l regularization
5 10-4
Table Fixed parameters for wide ResNet experiments in Section 3.2
23


Reviews:
[{"review_id": "kWvp93EFu6", "reviewer": "Reviewer_d6eE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The strengths of the paper include its novel and interesting topic and the confirmation of the importance of temperature tuning through empirical results. The paper's contribution can be widely adopted by the deep learning community. However, the connections between the theoretical contribution and experiment results may need to be clarified, and the experiments are limited to only three tasks.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a new time scaling analysis for the cross-entropy loss with temperature in deep neural networks. The authors introduce the concept of effective time scale and provide a theoretical analysis of linear and non-linear time scales based on the output of logits with temperature and step sizes using the Neural Tangent Kernel (NTK) theory. The analysis shows that smaller inverse temperatures (beta) in the softmax layer of the neural network result in a faster exit from the linear regime, which is associated with better performance. The paper includes experimental results on image classification and sentiment analysis tasks, which confirm the benefits of temperature tuning for neural network generalization.\n\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Q1: Figure 4 and Section 3.1 mention that the optimal learning rate is the learning rate with the best generalization rate with beta. This is unclear to me. Can you explicitly define the optimal learning rate and explain how it is derived?\n\nQ2: The beta range, [10^-2, 10^1] is not derived from the theoretical part from my understanding. One suggestion is that the paper should have more discussion towards the empirical settings, specifically the conclusion of the beta range, since it is highlighted in the abstract and might confuse readers.\n\nQ3: In section 3.4, the tuning procedure is time-consuming for training on large datasets, and this claim is self-evident. (More like a suggestion) One good point is, can the method do early stopping? Moreover, do we have any guidelines for the early stopping when we observe a specific beta that should be good for training?\n\n", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xauToJ01TZ", "reviewer": "Reviewer_12gG", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strength\n\nThe paper provides insights into how the inverse-temperature operator affects the early stage of supervised learning. The conclusion suggests that too large $\\beta$ prevents the training from accessing the non-linear regime, which harms the generalization, while too small $\\beta$ leads to slow training. I have not seen any other paper conducting such an in-depth analytic study on the parameter $\\beta$. The topic is interesting, the conclusions are informative, and the analysis is easy to follow.\n\n## Weakness\n\nDespite the appealing conclusions, there is still a non-negligible gap between the theory/experiments and practical deep learning scenarios:\n\n1. The theory is developed based on the continuous time learning dynamics which is exact only in the limit of small learning rate. The analysis focuses on the \"timescale\" by which the model leaves the linearized regime and enters the nonlinear regime, under different scales of $\\beta$, the stepsize, and the magnitude of logits. As the analysis manipulates these quantities, especially the stepsize, it is not clear whether the results still reveal the dynamic of the discretized case with practical stepsizes.\n  \n2. Many numerical experiments are carried out for the wide ResNet models, and the authors claim the architecture to be “well within the regime where the theory applies“. The reason why this architecture is within the desired regime is not clear. More justification is needed.\n  \n3. The correlated initialization used in the experiments, where $\\beta$ and $||\\beta \\mathbf{z}||$ can be adjusted independently, is very special. Practically-used initialization schemes do not allow tuning $\\beta$ while keeping $||\\beta \\mathbf{z}||$ constant. Instead, $||\\beta \\mathbf{z}||$ is always linear in the size of $\\beta$. In this case, it is not clear how to interpret the theoretical results of this paper.\n  \n4. Even though the theory tells a lot about how $\\beta$ affects the early-stage training dynamics, the experiments for ResNet-50 on ImageNet seem to imply that \"$\\beta$ depends on the actual architecture, but we need fine-grained tuning to find it\". The only information from the theory that can be directly utilized in the actual training seems to be \"tuning $\\beta$ may improve the performance\". If the current $\\beta$ is bad, the theory seems not even able to tell us whether we should increase or decrease $\\beta$.\n  \n\nI understand that for many theoretical papers, additional efforts are needed for the theories to play a part in guiding practice. However, a thorough discussion on the limitation of the theory can clearly shed light on the topic studied. Therefore, I strongly encourage the authors to clarify and discuss the above issues.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the training of supervised deep learning models with softmax-cross-entropy loss. Specifically, the paper considers the early stage of the continuous time learning dynamics (i.e., the gradient flow), where the training will experience a linearized regime. By deriving the time scale to access the non-linear regime, the authors develop a theory on how the training dynamics depend on the inverse-temperature parameter $\\beta$ in the softmax operator, as well as the magnitude of the logits $||\\beta \\mathbf{z}||_2$.\n\nThe theoretical results make two major claims:\n1) Small \f$\\beta$ causes training to access the non-linear learning regime. (Empirical evidence further shows that  We will see empirically\nthat earlier access to the non-linear regime can improve generalization).\n2) The theory suggests the best learning rate $\\eta$ to be the timescale to leave the linearized learning regime. This implies that too small $beta$ will make the model learn slowly.\n\nNumerical results on wide ResNet are provided to support the theoretical claims. This paper also presents experiments on other deep learning models, such as ResNet-50 on ImageNet, to show that fine-tuning $\\beta$ can further improve the test accuracy. Although the best $\\beta$ depends on the network architecture,  the results suggest tuning $\\beta$ over the range $10^{-2}$ to $10^1$ for all architectures.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "For major concerns, please refer to ``weakness'' in the last section. Below are a few additional questions\n\n1. In the introduction, a sentence reads \"increasing access to the non-linear regime can improve generalization.\" What does \"increasing access\" mean? Should it be \"earlier access\"?\n\n2. In Eq. (1), both $\\mathcal{Y}_i$ and the ln() term seem to be $M$-dimensional vectors. Does $\\cdot$ represent the inner product? Should it be $\\mathcal{Y}_i^\\top$ instead of $\\mathcal{Y}_i$?\n\n3. It seems that the theory is not limited to deep learning settings, but all supervised learning settings for classification tasks. If so, why does the paper restrict the discussion under deep learning settings?", "broader_impact_concerns": "Not applicable.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "XZnQVUP9sv", "reviewer": "Reviewer_fn91", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n1. **[Interesting Topic]** Softmax-cross-entropy is prevalently used in the entire machine learning community. Its temperature coefficient is almost always set to 1 by default, and most researchers & practitioners do not even notice that it's a tunable hyper-parameter. Even if some works tuned this hyper-parameter, they usually just treated it as a black-box hyper-parameter. This work studies this problem deeply, and I think the topic is of interest to a wide range of audiences in the machine learning community.\n\n2. **[Solid Theory]** The authors leverage some theoretical tools of NTK to study the soft-cross-entropy loss. Notice that NTK was proposed for MSE loss and most NTK works focus on MSE. In my perspective (as a researcher who also used NTK as a tool), it's non-trivial to develop an NTK-based theoretical analysis for cross-entropy. Surely the authors use lots of approximations, but they have good reasonings -- also, I believe certain assumptions & approximations are necessary to carry out an NTK-based analysis for cross-entropy. Overall, I think the theoretical analysis part is a solid contribution -- and the theory could be used as a theoretical tool for other researchers studying neural nets under cross-entropy loss, of independent interest.\n\n3. **[Thorough Experiments]** The experiment part is also very good. First, the authors used multiple datasets (image: CIFAR10, ImageNet, text: IMDB) and various neural network structures (ResNet, Wide ResNet, GRU), and the experiment results support the theoretical results.\nFurthermore, the authors observe some patterns of the $\\beta$ coefficient, and offer useful tips to practitioners for the choice/tuning of this hyper-parameter.\n\n## Weaknesses\n\n**[The effect of Width]** The authors adopt the NTK tool, which usually requires the neural net wide to be sufficiently large. However, in this paper, the authors do not quantitively define the magnitude of the width. I can see that the authors surely assume the width to be large, e.g., in Sec. 2.2, the authors state, \"For small changes in θ, the tangent kernel is approximately constant throughout learning,\" which is a result of large width. From my perspective, it's crucial to know how width affects the theory -- i.e., when the theory breaks/holds as a result of small/large width? \n\nIn the discussion section, the authors state, \"*it has been shown that as the width n of a network is taken to infinity, choosing β = 1/n gives a well-defined infinite-width limit which is distinct from the linearized NTK regime (Yang et al., 2022; Yang, 2021). Our work shows that manipulating β is useful and interesting in the finite-width regime as well.*\" In my opinion, it's not 100% clear that the current theory is in the finite-width regime, since the authors have not defined the width quantitively -- we cannot rule out the possibility that the theory of this work only holds for (almost) infinitely wide networks. The authors need to quantitatively analyze the effect of width theoretically to convince readers that the theoretical results apply to moderately wide neural nets.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the temperature coefficient in the commonly used softmax-cross-entropy loss. Theoretically, the authors consider the NTK regime (i.e., neural nets are wide enough & learning rate is small), and study the effect of the inverse temperature, $\\beta$, on the learning dynamics of neural nets. Empirically, the authors conducted experiments on multiple datasets (vision & text) with various neural net structures to corroborate the theoretical results. Based on the theoretical & empirical analyses, the authors offer a suggestion to practitioners: use $\\beta\\in[10^{-2},10^1]$, and usually a smaller $\\beta$ is better for neural net generalization.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I quite appreciate the paper and enjoy reading it. There are some minor suggestions that I think can improve the paper -- it's totally fine not adopting my suggestions.\n\n1. **[Discuss Related Works in a Separate Section]** The authors discussed related prior works across this paper, but the discussions are distributed over the entire paper -- for example, [1] is a prior work that is closely related to this work, and some observations of the current paper are overlapped with this prior one; the authors mentioned this prior work loosely in several places, but I think it's better to discuss this prior work in an individual \"Related Works\" section, which can help readers to understand better all relations & differences between this past work and the current one. The same suggestion applies to other related works as well.\n\n2. **[Theoretical Analysis in the Large-Depth Regime]** Clearly, this paper considers the large-width scenario since it uses the NTK technique. It would also be interesting to study how large depth affects the soft-cross-entropy & learning dynamics. For instance, [2] studies the NTK for several wide neural net structures (fully-connected, resnet) and provides insights on how depth affects neural net trainability & generalization. I think the techniques of [2] can be used to further study the problem in the current paper, and maybe the authors will also find some interesting relations between the network depth & the temperature coefficient, etc.\n\n\n[1] Lewkowycz et al. The large learning rate phase of deep learning: the catapult mechanism. arXiv 2020.\n\n[2] Xiao et al. Disentangling Trainability and Generalization in Deep Neural Networks. ICML 2020.", "broader_impact_concerns": "No broader impact concern.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "aUkZokwxTR", "reviewer": "Reviewer_TjA9", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\n1.This paper investigates an important problem, that how the inverse-temperature of softmax-cross-entropy loss, as well as the magnitude the logits at initialization, affects the dynamics of logits during training. In fact, the inverse-temperature/temperature of softmax-cross-entropy loss is widely used in for the self-supervised learning, not only for the supervised classification tasks. I think the audiences of TMLR will have interest in this paper. \n\n2.This paper will provide significant contributions, if the findings are well supported by the experimental evidences with details that can be reproducible. (I have concerns on the findings/claims)\n\n3.This paper is well organized, and well compare the previous related work, as to me. \n\n**Weaknesses:**\n\n1.My main concerns are that the details of the experiments are not clear, it is difficult for me to understand how the experiments are conducted (e.g, the configurations, hyper-parameters), and further recognize the experimental evidences for supporting the claims. (I believe the main claims concluded in this paper are mainly based on their experimental observations.)\n\n(1)\tIn Figure 1, what is the details to change the x-axis $\\eta t$ (e.g., the learning rate used during training) ? Does the setup is all the same for different $\\beta$?. It is important to know these information. E.g., In Eqn.3, the X is defined over the all dataset, How about the batch size used in the experiments (full batch gradient or mini-batch gradient) ?\n\n(2)\tIn Figure 2, “We compute $\\tau_z$ explicitly as…. for short times (Figure 2, left).” Does this paper mean that the $\\tau_z$ is calculated not based on the definition in Eqn.7? If yes, why not use the definition in Eqn.7  and how this approximation holds ? if not, how does this paper explicitly calculate the NTK? \n\n(3)\tWhat is the experimental details in Figure 5 (learning rate, weight decay, batch size and so on)? What is the optimal learning rate $\\eta^*$ in the experiments and how to set/get the optimal learning rate here? I noted that Figure 5 uses the wide Resnet without batch normalization, how about the results on the original wide ResNet (with batch normalization)? Do the observations still hold using the original wide ResNet?\n\n(4)\tI think it is difficult to reproduce the experiments shown in the figures, if this paper does not provide the code (including the experimental setups, hyper-parameters). It should provide the details for their experimental setups (at least in the supplementary materials) for reproduction, and it is better to provide the codes for the experiments in section 2.  \n\n2.I also have concerns on some claims/description of this papers. \n\n(1)The claims/conclusions of this paper obtained in Section 2 is just only performed on the the Wide Resnet (with batch normalization, I think) trained on CIFAR10, but based on the description in the last paragraph of introduction section, it seems that this paper claims the obtained main claims/conclusions are experimentally supported over multiple architectures (with batch normalization or not) and datasets (tasks). By the way, I did wonder whether the claims/conclusions of this paper still hold over multiple architectures (with batch normalization or not) and datasets (tasks) described in the introduction. \n\n(2)In the abstract, “We find that smaller $\\beta$ may lead to better peak performance at the cost of learning stability.” This observation seems go against my experience in training neural networks. It seems large logits (large beta) usually shows the training instability in the initial stage. But when I check the details in section 3, I find this instability means the results over different random seeds has high variance, not the instability during training. I think this paper should clearly illustrate the differences is the abstract (if this sentence is preserved) \n\n\n(3) I have concern on this statement “ Here, networks begin to learn over a timescale $\\tau_z \\sim \\|Z^0 \\|_2/ \\tilde{\\eta}$, where $Z^0$ are the initial logits of the network after being \nmultiplied by $\\beta$. This implies that we expect learning to begin faster for networks with smaller logits.”. Based on $\\tilde{\\eta}=\\beta^2 \\eta$, it seems that smaller logits ($\\beta$) will make the learning begin slower. Indeed, in section 5, this paper also states that “Therefore, we predict that networks with small $\\beta$ take longer to train (which we see in practice as well).”, which seems to contradict to the statement (I concerned). Similar concerns also hold for the nonlinear learning dynamics \n\n(4)For this statement, “Thus we see that, in fact, networks with smaller $\\beta$ tend to become more non-linearized before much learning has occurred, compared to networks with large $\\beta$ which can remain in the linearized regime throughout training. Our numerical results confirm this intuition developed using the theoretical analysis.” I donot recognize how the numerical results confirm the claims, I expect this paper can provide clear logic in how the numerical results confirm the claims. \n\n\n\n\n \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper focuses on investigating how the (inverse) temperature of softmax-cross-entropy loss impacts the training dynamics or generalization performance of a neural network. It develops a theory of early learning for models trained with softmax-cross-entropy loss, borrowing the tool of Neural Tangent Kernel (NTK), and shows how the inverse-temperature, as well as the magnitude the logits at initialization, affects the dynamics of logits during training. These observations are supported by experiments conducted on wide residual network for CIFAR-10 classifications. Furthermore, this paper conducts experiments to investigate how the inverse-temperature affects the generalization performance, and obtain several interesting observations empirically. ", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Provide the experimental details, and make it possible to reproduce the data in the figures based on the experimental details.\n\n2. For this statement, “Thus we see that, in fact, networks with smaller $\\beta$ tend to become more non-linearized before much learning has occurred, compared to networks with large $\\beta$ which can remain in the linearized regime throughout training. Our numerical results confirm this intuition developed using the theoretical analysis.” I donot recognize how the numerical results confirm the claims, I expect this paper can provide clear logic in how the numerical results confirm the claims.", "broader_impact_concerns": "No impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a finding that tuning the inverse temperature of the softmax-cross-entropy loss can improve the generalization performance of the neural network. The paper provides a solid theoretical analysis based on NTK to support the finding. This finding may be of interest to researchers and practitioners working on deep learning, as it provides a new method for improving the performance of neural networks. Besides,  the paper is well-organized, clearly written, and compares well with related works. There were some concerns but addressed well as noted by the reviewers. All reviewers recommend acceptance. 
Decision: Accept (poster)