Stepsize modulation for cross entropy from local quadratic optimization

Introduction

The selection of the learning rate is a very consequential decision when training neural nets. Despite much research in automatically selecting step sizes, typical training with convolutional nets is done with some fixed schedule for which maximum learning rate and weight decays are tuned. In this paper, we ask whether one can modulate the standard cosine schedule using batch dependant stepsizes to improve training. More precisely, can we inspect certain properties of a minibatch and use it to decide whether a bigger or smaller stepsize should be used compared to the scheduled value of the epoch. Can this speed up training and increase test accuracy?

The Polyak stepsize

Consider a loss function $L(\theta)$. We aim to find a value of $\theta$ that makes $L$ as small as possible. The gradient descent method starts from an arbitrary $\theta_0$ and iteratively computes

Already in 1969, Polyak proposed to automatically select the stepsize using a linear approximation of the loss function [24]. Consider the linear approximation

*National Research University Higher School of Economics, Faculty of Computer Science, Moscow, Russia; This work/article is an output of a research project (HSE-BR-2025-024) implemented as part of the Basic Research Program at HSE University. This research was supported in part through computational resources of HPC facilities at HSE University.

†Luxoft Serbia.

For $w=\eta_t\nabla L(\theta_t)$, this simplifies to

Solving for $L(\theta_{t+1}) \approx0$, gives

It was directly clear that if $L$ is smooth and strongly convex, then gradient descent with this stepsize converges exponentially fast in $t$ [24]. A bit later, Polyak proved that for all convex functions (possibly non-smooth) this method has $O(1/\sqrt{t})$ convergence: $\liminf_{t}\sqrt{t}(L(\theta_t)-L(\theta^*))=0$.

In machine learning, the above gradients are estimated on small random selections of the training data, called minibatches. This means that for each step, a random approximation $L_t(\theta)$ of the loss is generated. The Polyak stepsize under this function is called the stochastic Polyak stepsize, for which $O(1/t)$ convergence was proven under star-convexity of the loss [9 theorem 4.4], and this is locally satisfied in deep neural nets [15, 33].

Unfortunately, the Stochastic Polyak stepsize is very unstable and can spike. This often happens in the beginning of training when $L(\theta_t)$ is large and the gradients are small (especially if the softmax is used). In the case of cross entropy loss, the curve does not reach zero and is locally convex for most minibatches. Hence, the Polyak stepsize overestimates, as illustrated in Figure 1 and Figure 2.

For this reason, many variants for practical use have been proposed [3, 19, 14, 20, 23, 21]. However, they lead often lead to unstable training, especially when combined with momentum. Moreover, they are not able to replace manual learning rate selection because of sensitive dependencies on newly introduced hyper parameters. (Also for LLMs there are such reports [22].) Perhaps, an exception is the *momo* scheduler from [26]`, which is designed to optimize the loss aggregated from the last few minibatches. It works well with momentum and weight decay. Moreover, it has has only a single hyper parameter that is easy to tune, but it has not been tested systematically against competitors.

Second order methods

Approaches with test points. To select the learning rate, one may consider a quadratic approximation of the loss,

Each optimization method in a typical machinelearning optimizer defines a direction, for example, for basic SGD the direction is $w=-\eta\nabla L(\theta_t)$ or with momentum, it is $w=-\eta\nabla L(\theta_t)-\mu v_{t-1}$. After substituting in $L$, one obtains an equation of the form $f(\eta)=a\eta^2+b\eta+c$ from which an optimal stepsize $\eta$ can be computed. In [5, 34], the coefficients $a,b,c$ are obtained from 3 evaluations of $f$. Since gradient descent already provides 1 evaluation, this method makes 2 extra forward evaluations, and hence, is %66 slower (assuming a backwards pass takes twice the time of a forward pass). These methods reduce the training loss very fast, but at the expense of the test loss [34]. A similar technique is applied in [5] in which the computed learning rate is smoothed using a moving average and improvements on some fine tuning tasks are reported.²

In [8] it is noted that $w \cdot\nabla L(\theta_t)$ is fast to compute, thus to find the coefficients $a,b,c$, only 1 extra forward evaluation is needed. The paper presents an algorithm that is essentially parameter free and provably converges to a stationary point for Lipschitz functions. Unfortunately, they do not report test performance.

Barzilai-Borwein stepsizes. Another approach is to calculate the solution using the available gradient information, thus one needs

In the Barzilai-Borwein method [2] one assumes $v = \nabla L(\theta_t)$ and estimates the above value by

This inspired a number of automated learning rate methods [6, 28, 30]. Again, the method suffers from instability, but [30] results come close to standard methods such as SGD with cosine annealing.

1st order estimates of the 2nd order. In [1], the stepsize (1) is estimated directly with the already computed gradient. The denominator involves the Hessian, which is expensive to compute. However, typical networks of the last decade have wide hidden layers, and this allows for a good approximation with gradients. This is because preactivations of the output neurons are locally linear functions, as explained in many papers about the neural tangent kernels, see for example [12, 17, 25].

For this, consider the loss as a function of output preactivations

The Hessians of $L$ can be written using the Hessians of $\mathcal{L}$ and the Hessians of the output preactivations $F_a$. Let $J_{a\mu} = \frac{\partial F_a}{\partial\theta_\mu}$ be the Jacobian of $F$.

For large neural nets, the spectral norm of the Hessians $F_a$ are proportional to $O(\text{depth}/\text{width})$, see [25]. In the outer product approximation of $H_L$ all these Hessian are neglected and only the 1st term remains, see for example [4 p 243, sect.8.1.6]. In [1] this first term is computed using the available gradient, and (1) is obtained with negligible computational overhead. They show in a few experiments that this reduces the train error better than SGD, but do not report test errors.

Below, we propose a method that uses a single test point to estimate this 2nd order stepsize in the same way.

Other methods

In the paper [18], a wide range of optimizers is modified in a generic way: in each minibatch, the existing learning rate schedule is randomly increased or decreased by a few percentages. Despite its simplicity, this method is reported to consistently increase test error, including on Cifar10 and Resnet18 experiments.

This easy method is similar to other regularization methods like dropout [27] or neftune [13] that increase randomness during training. The motivation is that solutions with wide minima have better generalization and it is argued that these perturbations can help to escape narrow local minima. See [31] for this theory.

Finally, the paper [7] uses a complex algorithm from the field of online convex optimization to select the learning rate. [todo: test error?]

Two proposed stepsizes

We aim to minimize to optimize parameters $\theta$ to train a neural nets. For this, a standard run of stochastic gradient descent is applied, which means that we start from a randomly generated $\theta_0$ and at each step, a minibatch loss $L_t(\theta)$ is used to update $\theta_t$ using

where $\mu$ and $\lambda$ are the parameters for momentum and weight decay.

Notation. Consider a neural net with $n_{\mathrm{out}}$ output neurons and a train set whose examples are denoted by $(x_\alpha,y_\alpha)$ with Greek indices. For each training example, the label $y_\alpha$ is given as a one-hot encoding of the correct label, thus $y_{\alpha c}=1$ if $x_\alpha$ belongs to the $c$-th class. The preactivations of the output neurons are denoted by the vector $f(x_\alpha,\theta)$, where $\theta$ is the list of all trainable weights of the net.

We now define the minibatch loss $L_t$. Let

For a list of train examples $S_t$, the $t$-th minibatch loss is

Notation. We write $\sum_\alpha$ to denote $\sum_{\alpha\in S_t}$ and sums over indices $i,j,k$ range from 1 to $n_{\mathrm{out}}$.

The second order outer approximated stepsize (so-dz)

Fix $t$ and let $L(\tilde{\eta})=L_t(\theta_t-\tilde{\eta}\tilde{v}_t)$. We select the stepsize $\eta_{\mathrm{so}}$ to minimize the quadratic approximation of $L_t(\tilde{\eta})$. As explained in subsection 1, we may assume that $z_{\alpha i}$ is linear in $\tilde{\eta}$ and we can obtain the derivative from only the gradient. Still, to increase stability, we use a test point.

Let us explain this plan in more details. We have

To compute $L'(0)$ we view $L$ as a function of the output preactivations $z_{\alpha i}$ and apply the chain rule. The 2nd derivative $L''(0)$ is computed similarly, while the 2nd derivatives of $z_{\alpha i}$ are neglected, see appendix A for details. Let $\delta_{ij}$ be the Kronecker function, which equals 1 if $i=j$ and 0 otherwise.

Finally, in this expression, $\frac{d z_{\alpha j}}{d\eta}$ is estimated by evaluating the output preactivations $z$ with a test update of $\eta$ using a fixed learning rate $\eta^{\mathrm{test}}$. More precisely, let

Let $q_\alpha^{\mathrm{test}}$ and $q_\alpha$ be the softmax values evaluated with parameters $\theta_t^{\mathrm{test}}$ and $\theta_t$. Then, estimate $\frac{d z_\alpha}{d\eta}$ by $(z_\alpha^{\mathrm{test}}-z_\alpha)/\eta^{\mathrm{test}}$. Thus,

This finishes the construction. The pseudocode is given in algorithm 1. In the experiments, we fixed $\eta^{\mathrm{test}}=5\cdot10^{-5}$.

Algorithm 1 Stepsize modulation algorithms with $\hat{\eta}_{\mathrm{so}}$ and $\hat{\eta}_{\mathrm{cs}}$

Initialize: $\theta_0$ randomly, $v_0 = 0$

for epoch $= 1, \ldots, T$ do     sample a fixed-size mini-batch $S$ and use data $\{(x_\alpha, y_\alpha)\}_{\alpha\in S}$     $\eta_{\mathrm{epoch}} \leftarrow\operatorname{schedulerValue}(\operatorname{epoch})$     $\vartheta_t \leftarrow\theta_t - (\mu v_t + \lambda\theta_t)$     $z_\alpha\leftarrow f(x_\alpha, \vartheta_t)$     $q_\alpha\leftarrow\operatorname{softmax}(z_\alpha)$     $\mathcal{L}_t \leftarrow\operatorname{crossEntropy}(z_S, y_S)$     $\theta^{\mathrm{test}} \leftarrow\vartheta_t - \eta^{\mathrm{test}}\nabla_\theta\mathcal{L}_t$     $z_\alpha^{\mathrm{test}} \leftarrow f(x_\alpha, \theta^{\mathrm{test}})$     $q_\alpha^{\mathrm{test}} \leftarrow\operatorname{softmax}(z_\alpha^{\mathrm{test}})$     $\hat{\eta} \leftarrow\text{eq (3) or (4) using } z_S, z_S^{\mathrm{test}}, q_S, q_S^{\mathrm{test}}$     $\eta_{t+1} \leftarrow\operatorname{modulation}(\hat{\eta}, \eta_{\mathrm{epoch}})$     $\theta_{t+1} \leftarrow\theta^{\mathrm{test}} - (\eta_{t+1} - \eta^{\mathrm{test}})\nabla_\theta\mathcal{L}_t$     $v_{t+1} \leftarrow\mu v_t + \lambda\theta_t + \eta_{t+1}\nabla_\theta\mathcal{L}_t$ return $\theta_{T+1}$

The closest softmax stepsize (cs-dq)

A linear approximation of the softmax values in $\eta$ is used to select a value that minimizes the average Euclidean distance to the one-hot encodings $y_\alpha$. Thus we minimize the sum of their 2-norms

Recall that sums over $\alpha$ range over the $t$-th minibatch $S_t$. Solving the derivative to $\eta$ gives

Finally, we use the same probe for $\eta^{\mathrm{test}}$ as before to estimate $\frac{\mathrm{d}q_\alpha}{\mathrm{d}\eta}$.

Both stepsizes are close

Despite being derived by different philosophies, the averages of these learning rates are close to each other. Indeed, we trained Resnet18 on Cifar10 with standard SGD (max learning rate 0.02, momentum $\mu= 0.9$, weight decay $\lambda= 0.005$) and cosine learning rate schedule. During this run we computed the average values of the Polyak stepsize and the 2 above stepsizes. In figure 2, we compare these stepsizes. We conclude that after halving the Polyak stepsize, these values are close to each other during the middle part of the training.

[To do: discuss fluctuations and compare values during minibatches.]

Practical refinements

Outlier removals in stepsize estimates

If for some training example $\alpha$ of a minibatch, the softmax values do not seem to improve, then we drop this input in the estimation of $\hat{\eta}_{\mathrm{so}}$ and $\hat{\eta}_{\mathrm{cs}}$. More precisely, we drop $\alpha$ from minibatch $S_t$ if $\cos\phi_\alpha< 1/4$ where

The effect is small, because only few points satisfy the condition $\cos(\varphi_\alpha) < 1/4$, as shown in figure 3. But it seems to help to stabilize the estimates [can we say something more convincing or precise]?

Subsampling and computational overhead

The estimation of the stepsize requires an extra forward pass. This increases the runtime typically by 33%, because the time for a backward pass is usually twice the time of a forward pass.

In our 2nd series of experiments, we took 20 random inputs from the remaining minibatch $S_t$, (which had 128 elements, before the above pruning), and estimated the learning rate from them. This reduced the overhead from 33% to about 21% (or better on GPU’s with few cores and fast memory).

Modulate an existing learning rate schedule

The effect of changing the learning rate for a given step also has an effect on future steps, because of momentum. Several papers report that this is a challenge to design automated learning rate schemes, for example in [29 footnote 1 on p7]. Moreover, the several theoretical studies suggest that during initial learning it is beneficial to keep the stepsize larger, in order to avoid parameters $\theta_t$ approaching sharp local minima. For these reasons, we focussed on perturbing existing schedules, rather than reconstructing them from scratch.

• Let $\eta_{\mathrm{ref},e}$ be a reference learning schedule evaluated in epoch $e$. In experiments we will use cosine decay.

• Let $\eta_t$ be the stepsize from 1 of the above methods, thus $\hat{\eta}_{\mathrm{so}}$ or $\hat{\eta}_{\mathrm{cs}}$ at batch $t$.

• Let $\bar{\eta}_t$ be the average of $\eta_t$ over last $N = 50$ batches.

• Let $\beta= 0.05$ be fixed. [Todo: it turns out that we have only lookahead during experiments of our paper.]

Rather than clipping, a smoothly perturbed version of $\eta_{\mathrm{ref},e}$ is obtained by multipliyin it with a the following factor in $[1-\beta\pi/2,1+\beta\pi/2]$,

The arctan function was used because the derivative in the origin equals 1.

Lookahead

In the 1st series of experiments, the lookahead technique from [32] was applied. The idea is to stabilize an existing optimizer Opt (like SGD, AdamW, etc) by letting it first make $k$ updates and then make a single update of the aggregated “gradients”. The optimizer has 2 parameters $k$ and $\alpha$. We use the default values $k = 5$ and $\alpha= 0.8$.

More precisely, the optimizer Opt is viewed as a fast optimizer and is used to create a slow optimizer as follows. Let $\theta_t$ denote the current parameters of the slow optimizer at the start of its slow step. The fast optimizer is initialized with $\phi_{t,0} = \theta_t$ and computes $k$ updates $\phi_{t,1},\ldots,\phi_{t,k}$. Finally, the slow optimizers sets

Experiments

We use the well known Resnet18 and Resnet34 models and trained them from scratch on Cifar10 and Cifar100. Data augmentation and the modification of the last hidden layer was applied as is now standard, see for example [10].

The 2 baseline experiments are as follows.

• Standard SGD with cosine learning rate decay (without warmup). The max-learning rate is 0.02, the momentum parameter is $\mu= 0.9$, the weight decay is $5e-3$ and the minibatch size is 128.

• Lookahead with the above SGD as fast optimizer. The standard values $k = 5$ and $\alpha= 0.8$ are used.

We have a default experiment which uses Resnet18 on Cifar10 with 50 epochs. All pictures where generated with this default experiment. The max-learning rate and weight decay parameters reported above where tuned for this experiment. Our 2 stepsizes were designed using this same experiment. However, below we report these baselines for different settings with 100 and 200 epochs, and we also report the experiments with Resnet34 and Cifar100. In all experiments the hyper parameters remain the same and are fixed.

Net	optimiser	Cifar10		Cifar100
		100 epochs	200 epochs	100 epochs	200 epochs
Resnet18	SGD	95.256 ± 0.152	95.476 ± 0.099	78.063 ± 0.239	78.649±0.259
	la	95.287 ± 0.12	95.603 ± 0.149	78.123 ± 0.296	78.543 ± 0.254
	la, so-dz	95.399±0.127	95.608 ± 0.102	78.157±0.229	78.593 ± 0.241
	la, cs-dq	95.379 ± 0.097	95.673±0.167	78.114 ± 0.277	78.554 ± 0.204
Resnet34	SGD	95.342 ± 0.106	95.536 ± 0.127	78.891 ± 0.249	79.357 ± 0.259
	la	95.471 ± 0.134	95.629 ± 0.153	78.968±0.265	79.537±0.197
	la, so-dz	95.550 ± 0.109	95.711±0.097	78.919 ± 0.225	79.461 ± 0.174
	la, cs-dq	95.574±0.102	95.663 ± 0.087	78.959 ± 0.204	79.519 ± 0.256

Net	Optimiser	Cifar10		Cifar100
		100 epochs	200 epochs	100 epochs	200 epochs
Resnet18	SGD	95.266 ± 0.164	95.552 ± 0.130	77.896 ± 0.262	78.610 ± 0.404
	lookahead	95.294 ± 0.148	95.520 ± 0.154	78.225±0.263	78.723±0.179
	la, so-dz	95.363 ± 0.071	95.558 ± 0.147	78.176 ± 0.153	78.595 ± 0.180
	SGD, cs-dq	95.425±0.134	95.653±0.135	78.051 ± 0.212	78.353 ± 0.300
Resnet34	SGD	95.436 ± 0.132	95.428 ± 0.198	78.833 ± 0.249	79.347 ± 0.190
	lookahead	95.438 ± 0.165	95.667 ± 0.139	78.863 ± 0.257	79.592±0.257
	SGD, so-dz	95.574±0.153	95.625 ± 0.151	78.941±0.236	79.415 ± 0.141
	SGD, cs-dq	95.433 ± 0.107	95.703±0.158	78.876 ± 0.302	79.427 ± 0.289

In all experiments, the hyper parameters of the stepsize $\hat{\eta}_{\mathrm{so}}$ and $\hat{\eta}_{\mathrm{cs}}$ are $\eta^{\mathrm{test}} = 5 \cdot10^{-5}$, the perturbation $\beta= 0.05$ and the average stepsize is computed from 50 minibatches. Experiments were performed on V100 GPU’s on the cluster of the Higher School of Economics in Moscow [16]. The code is available here https://github.com/Brbauwens/adaptive_stepsize_modular.

For both the so-dz and cs-dq stepsizes, the initial test error decreases more rapidly and the curve is more monotone. However, recall that this method requires 33% more compute (and 21% more compute with subsampling of the test probe). Especially, this is true when combined with lookahead, see figure 4 for the default experiment. However, near the end, improvement becomes slower.

Our methods often improve error on the test set. In a 1st series of experiments, lookahead was applied and no subsampling was done. All 4 methods were applied to the same minibatches. The best test set accuracies are given in tables 1. For the Cifar10 training set, both the so-dz and cs-dq stepsizes exceed both baselines. Unfortunately, for Cifar100 the results look random.

In a 2nd series of experiments, subsampling was applied and lookahead was switched off. The runs of the baselines were repeated with different random seeds. The results in table 2 suggest that this does not decrease performance. Again for Cifar10, the best results is either for so-dz or cs-dq, but in the case of Resnet34, the gains look random. Also, for Cifar100, the methods do not provide reliable improvements.

Conclusion

Two methods were proposed to automatically generate a stepsize for each minibatch. In $\hat{\eta}_{\mathrm{so}}$ stepsize optimizes a 2nd order approximation of the minibatch loss, and the $\hat{\eta}_{\mathrm{cs}}$ optimizes the softmax values. We argued theoretical and experimentally that both methods give closely related stepsizes.

These stepsizes were used to improve the cosine decay schedule. Experiments show that they speed up convergence beyond the 20%-33% especially in the middle of the training. Moreover, the decrease is more monotone. They give essentially the same accuracy, and in a few experiments, their best accuracy is improved compared to standard SGD and its combination with lookahead.

Derivation of so-dz stepsize

Fix the minibatch is fixed and we drop subscripts $t$. Below, the quantities $q$ and $z$ are considered for the updated parameter $\theta_{t+1}$, which depends on the updated learning rate through (2), thus $q=q(\theta_{t+1}(\eta))$ and $z=z(\theta_{t+1}(\eta))$.

We need to evaluate $L'(0)$ and $L''(0)$ and write it as a function of $q$ and $\frac{dz}{d\eta}$. Recall that $y_{\alpha j}$ is the one hot encoding of the label of training input $x_\alpha$. One calculates that

We have

To calculate $L''$, we neglect the 2nd derivatives of the output preactivation $z$. Thus,