Automated Learning Rate Scheduler for Large-batch Training
Abstract
Large-batch training has been essential in leveraging large-scale datasets and models in deep learning.
While it is computationally beneﬁcial to use large batch sizes, it often requires a specially designed learning rate (LR) schedule to achieve a comparable level of performance as in smaller batch training.
Especially, when the number of training epochs is constrained, the use of a large LR and a warmup strategy is critical in the ﬁnal performance of large-batch training due to the reduced number of updating steps.
In this work, we propose an automated LR scheduling algorithm which is eﬀective for neural network training with a large batch size under the given epoch budget.
In speciﬁc, the whole schedule consists of two phases: adaptive warmup and predeﬁned decay, where the LR is increased until the training loss no longer decreases and decreased to zero until the end of training.
Here, whether the training loss has reached the minimum value is robustly checked with Gaussian process smoothing in an online manner with a low computational burden.
Coupled with adaptive stochastic optimizers such as AdamP and LAMB, the proposed scheduler successfully adjusts the LRs without cumbersome hyperparameter tuning and achieves comparable or better performances than tuned baselines on various image classiﬁcation benchmarks and architectures with a wide range of batch sizes.
1. Introduction
In modern deep learning tasks, scaling up both the sizes of dataset and model has shown promising improvements.
However, a large number of samples and slow gradient computations lead to considerably longer training time, and therefore the use of large batch size in stochastic optimization with multiple computational nodes has gain popularity to speed up for such a large-scale training.
It is desirable for the optimization algorithm with the increased batch sizes to maintain the performances without increasing the amount of processed samples, namely the number of training epochs.
Recently, several works have been proposed for successful large-batch training with stochastic gradient descent (SGD), and most of them achieve it by specially designed learning rate (LR) schedules and especially LR scaling (Krizhevsky, 2014; Goyal et al., 2017;
Hoﬀer et al., 2017).
For instance, linear (Krizhevsky, 2014; Goyal et al., 2017) or square- root (Hoﬀer et al., 2017)
LR scaling rules according to the batch sizes eﬀectively alleviate the performance degradation in large-batch training by compensating the reduced number of optimization steps.
In addition, gradual LR warmup heuristic (Goyal et al., 2017) reduces the instability caused by large LRs and has become a standard for large-batch training.
More recently, layer-wise adaptive LR scalings have been proposed to further increase the batch size and are applied to diverse tasks (You et al., 2017, 2020; Ginsburg et al., 2019).
However, these LR scaling and warmup heuristics are sensitive to its hyperparameters including the LR and the warmup schedule, hence require intensive tuning eﬀort.
There are a number of studies attempting to automate the LR scheduling to reduce such cumbersome tuning and to enhance the performance.
Online or oﬄine LR search algorithms using Bayesian optimization (Picheny et al., 2020) or meta-optimization (Schraudolph, 1999;
Chen et al., 2017; Baydin et al., 2018; Donini et al., 2020) have been proposed.
However, they have been limited to small (surrogate) tasks due to complexity and the feasibility of them in large-batch training is unknown.
We propose an automated online LR scheduler for large-batch training.
The schedule consists of the warmup phase, where the LR is increased until the training loss no longer decreases, and the decay phase that decreases the LR.
To robustly decide the phase transition on the ﬂy, we employ Gaussian process (GP) smoothing (Rasmussen et al., 2006).
This
GP-based online detection is robust to small bumps and has a low computational burden.
In addition, to cover a wide range of possible LR values while ensuring stability, the LR is exponentially increased from a very small value (10 − 5 ) to a very large value (1.0) up to the half of the given epoch budget.
As a result, the proposed LR scheduler can eﬃciently and automatically ﬁgure out not only the initial and peak LRs but also the warmup length in a data-driven way as the training progresses.
From the perspective of having the warmup and the decay phase, SALSA (Zhang et al., 2020) is most similar to this work.
However, SALSA relies on the backtracking line search for the warmup phase which can be burdensome, and it was not evaluated on large batch sizes.
We empirically demonstrate that the proposed LR scheduler together with an adaptive optimizer such as AdamP (Heo et al., 2021) or LAMB (You et al., 2020) achieves comparable or better performance compared to a tuned baseline in case of large batch size on various image classiﬁcation benchmarks and architectures.
2. Methods
We propose an automated LR scheduler, called AutoWU (stands for automated warmup), which consists of two phases warmup and decay .
Figure 1 describes how AutoWU works.
2.1 Warmup Phase
During the warmup phase, the LR starts from a very small value and is increased following an exponential schedule.
It is tested at the end of each epoch whether the loss has hit the minimum value as training progresses.
If the minimum has been detected, then the LR follows the predeﬁned decay schedule.
This scheme combines the warmup strategy which enhances the stability in the early phase of training with the automated LR selection.
2.1.1 Warmup by exponential schedule
The proposed exponential schedule in the warmup phase has three hyperparameters: the initial LR η min , the maximum possible LR η max , and the maximum possible fraction ρ w ∈
[Figure]
(0 , 1) of warmup steps.
Given the total number of steps T , the warmup schedule is deﬁned as η 0 = η min and
[Equation]
for
[Equation]
where
[Equation]
Here, if the training loss keeps decreasing, then the warmup phase can last for (cid:98) ρ w T (cid:99) steps, and η t can grow up to η max .
In comparison to the other growth rates such as the linear growth rate, this exponential growth rate enables stable and ﬁne-grained LR exploration, especially in the early stage for large-batch training.
We set η min = 10 − 5 and η max = 1 to ensure that LRs can sweep a wide enough range of values.
Moreover, we set ρ w = 0 .
5 to ensure that the LR does not grow too fast.
Results regarding the sensitivity of AutoWU with respect to ρ w and the comparison with the linear growth are found in Appendix C.2.
2.1.2 Online minimum detection via Gaussian process
A loss trajectory may exhibit local and global ﬂuctuations due to the stochastic nature of a mini-batch computation and a highly non-convex loss landscape.
This makes it diﬃcult to robustly detect whether a loss is no longer decreasing or not in an online manner.
As a remedy, we propose to use GP regression to smooth the loss curve and then conduct a decision test based on the predictive distribution.
Suppose that we have been observed (noisy) loss values L 0 , · · · , L t − 1 at step t .
We model the loss curve s (cid:55)→
L s (deﬁned for s ∈ [0 , t ]) by a GP with a homoskedastic noise:
L s = f ( s/t )
+ (cid:15), where f ∼ GP ( θ, K (cid:96),σ f ( · , · )) and (cid:15) ∼ N (0 , σ 2 n ) , (2) where θ , K (cid:96),σ f and σ 2 n denote the mean, the covariance kernel, and the noise level, respectively.
When conditioned on the observed sequence L 0 , · · · , L t − 1 , the predictive distribution of f is again a GP and can be computed straightforwardly.
We are interested in whether f is no longer decreasing, i.e. , there exists x ∗ < 1 such that f ( x ∗ ) <
f
Hence, it is natural to compute the probability
[Equation]
and decide to end the warmup phase if (3) is large.
However, since (3) is not easy to compute, we compute the following lower bound instead:
[Equation]
which holds for any S ⊆
Here, S is chosen to be 500 equally-spaced points in [0 , 1].
GP model details
We use the GP model with the constant mean θ and the squaredexponential kernel K (cid:96),σ f ( x, x (cid:48) )
=
σ 2 f · e −| x − x (cid:48) | 2 / (2 (cid:96) 2 ) , where the length-scale (cid:96) is ﬁxed to 0 .to
prevent model overﬁtting to local variations, and all other parameters θ , σ f , and σ n are ﬁtted by Adam (Kingma and Ba, 2014) with LR 0 . 01 for 100 steps.
Test
details.
There are three main hyperparameters involved in the test: n test , conﬁdence c , and patience p . We used n test = 5, c = 0 . 95 and p = 3 in all experiments. Suppose that the loss trajectory C
= { (0 , L 0 ) , · · · , ( t − 1 , L t − 1 ) } has been observed at step t .
To conduct the test, we ﬁrst choose a random subset C 0 ⊆ C of maximumsamples and ﬁt the GP parameters with respect to the marginal log-likelihood of C 0 .
Then, we further sample n test random subsets C 1 , · · · , C n test of C , each of maximum 500 samples, and infer f i conditioned on C i for i = 1 , · · · , n test .
Finally, we compute P min ( f i ) in Eqn.
(4) and if the majority of P min ( f i )’s exceed the conﬁdence c , then it is regarded as that the minimum has been detected.
We start the decay phase only if the minimum has been detected in p consecutive tests.
This is to prevent premature ending of warmup due to random spikes in the loss trajectory, which was observed more often in large-batch setting in our preliminary experiments.
We remark that larger p implies the switch happening later; this may result in a too large LR.
To remove the impact of the choice of p , we set the starting LR of the decay phase to be η t ∗ =
γ t ∗ η 0 where
[Equation]
The overall algorithm is summarized in Appendix A (Algorithm 1).
Computation time.
Time complexity of exact GP inference increases cubically with the number of samples.
Nonetheless, those computations are not very burdensome since they are only conducted once per epoch in this work, and we use subsampling so that the number of samples is kept constant regardless of the task (particularly the batch size).
Moreover, we use GPyTorch (Gardner et al., 2018) for eﬃcient GP computations on GPUs.
Time overhead per test (ﬁtting and 5 inferences) was less than a second in average, which is typically much smaller than the overall gradient computation time per epoch.
2.2 Decay Phase
In the decay phase, AutoWU follows an LR schedule with the predeﬁned shape, but whose starting LR is adaptively determined in the warmup phase.
To remove any sophistication in evaluation, we only consider two simple types of schedule in the decay phase: cosine (Loshchilov and Hutter, 2017) or constant-then-cosine .
In case of cosine decay, the LR starts with the value determined by the warmup phase and is annealed toward zero with cosine schedule.
This schedule is particularly appealing since it does not introduce any additional hyperparameters to consider.
On the other hand, constant-then-cosine decay maintains LR constant until a predetermined fraction of epochs is left and then follows [Table] cosine schedule.
We set the fraction 0.2 ( i
[Table]
e. cosine decay in the last 20% of epochs) in all experiments.
We have empirically observed that an automated LR decay based on the convergence test such as SASA+ (Zhang et al., 2020) was not decisively superior than these schedules.
This probably implies that a warmup strategy including the peak LR is more critical for the ﬁnal performance, especially in large-batch training setting.
3. Experiments
We evaluate our algorithm on image classiﬁcation tasks with three benchmark datasets:
CIFAR-10, CIFAR-100 (Krizhevsky et al., 2009) and ImageNet (Deng et al., 2009).
We mainly consider convolutional networks: ResNet-18, Wide-ResNet-28-10 (Zagoruyko and
Komodakis, 2016) for CIFAR-10 and CIFAR-100 respectively, and ResNet-50 (He et al., 2016) for ImageNet.
We also conduct evaluations on ImageNet with EﬃcientNet-B0 (Tan and Le, 2019) and a vision transformer ViT-S/16 (Dosovitskiy et al., 2020), and the results by these other architectures are included in Appendix B.
AutoWU only requires to know how the loss values are changing over the course of training, hence it is capable to be used with any stochastic descent algorithm.
In our evaluations, the state-of-the-art Adam-based optimizer, AdamP (Heo et al., 2021) is used as the base optimizer due to their better stability than SGD and its variants with large batch sizes.
The results when coupled with the layer-wise adaptive optimizer, LAMB (You et al., 2020), are also presented in Appendix B.
The proposed AutoWU scheduler is compared with the conventional baseline LR schedule on each task with multiple batch sizes: { 256 , 1024 , 8192 , 16384 } for CIFAR (200 epochs) and { 1024 , 4096 , 16384 , 32768 } for ImageNet (120 epochs).
The baseline LR schedule consists of 5 epochs of a linear warmup from 0 to a predetermined peak LR followed by a cosine decay to 0.
Here, following a common LR scaling practice for large-batch training, the peak LR is scaled with the batch-size according to the square-root scaling rule, speciﬁcally, η base (cid:112) B/ 256 for batch size B .
We set η base = 0 .
001 in all experiments by empirically tuning it as the common base LR that is well performed across tasks and architectures.
We verify that the base LR as well as the warmup length strongly aﬀect the performances of the baseline LR scheduler for large-batch training (see Appendix C.1), however the proposed
AutoWU scheduler removes these tuning eﬀorts.
[Figure]
As shown in Table 1, the proposed AutoWU performs on par with or better than the baseline LR schedule across diverse batch sizes and tasks.
Especially, with large batch sizes of 8K and 16K for CIFAR-100 and 16K and 32K for ImageNet, AutoWU with constantthen-cosine decay signiﬁcantly reduces the performance drops in comparison to the conventional baseline.
In addition, overall, the constant-then-cosine decay is slightly better than the cosine decay when combined with AutoWU which means that as we ﬁnd the safe peak
LR automatically for a long warmup time, retaining the found LR longer would be better.
Figure 2 shows the LR schedule that AutoWU found and the training loss curve in each task.
For any batch size and task, AutoWU robustly detects the minimum loss in an online manner by GP.
Also, interestingly, the starting LR of the decay phase by AutoWU grows as the batch size increases, and the rate is similar to the baseline scaling of the peak
LR even though the warmup epochs are quite diﬀerent.
Actually, the found combination of the peak LR and the warmup length is similar to the search results on the baseline schedule as shown in Appendix C.1.
More ablations on AutoWU are presented in Appendix C.
4. Conclusion
In this work, an automated LR scheduler especially for fast training using large batch sizes is proposed.
The proposed LR schedule has the warmup phase followed by the decay phase, where the LR is exponentially increased until the training loss no longer decreases and then decreased to zero until the end of training.
The online detection of the minimum loss has been eﬃciently and robustly realized by GP regression.
Empirical evaluation demonstrates that our automated LR scheduler appropriately adapts the whole warmup procedure as well as LRs for any batch size and task, and consequently results in comparable or better performances in comparison to the ﬁne-tuned LR schedulers.
Our implementation will be available at https://github.com/kakaobrain/autowu .
We would like to acknowledge and thank Woonhyuk Baek and Brain Cloud Team at Kakao Brain for their support.
Appendix A. More Details
The overall algorithm of AutoWU is depicted in Algorithm 1.
All implementations are based on PyTorch 1.7.0 and GPyTorch 1.3.0.
Experiments are conducted on NVIDIA Tesla V100 GPUs with 32GB memory.
We use the weight decay of 0 .
1, β 1 = 0 .
9, β 2 = 0 .
999, (cid:15) =
10 − 8 and δ = 0 .
1 for AdamP and the weight decay of 0 . 1, β 1 = 0 .
9, β 2 = 0 .
999 and (cid:15) =
10 − 6 for LAMB.
In all experiments involving CIFAR-10 or CIFAR-100, the corresponding network is trained with AutoAugment (Cubuk et al., 2018), Cutout (DeVries and Taylor, 2017), and label smoothing (M¨uller et al., 2019) with factor 0.1 for 200 epochs.
We report the mean and the standard deviation from three independent runs with diﬀerent random seeds.
In case of ResNet-50 training on ImageNet, the standard set of augmentations as in He et al.
(2016) is used for ImageNet and the network is trained for 120 epochs.
We report the result of a single run with a ﬁxed random seed in all ImageNet experiments due to resource constraints.
Appendix B. Further Evaluations
In this section, further evaluation results which were not discussed in Section 3 are presented. [Table] [Figure]
Firstly, the comparison between AutoWU and the baseline schedule when coupled with
LAMB is made.
The peak LR of the baseline is scaled as η base (cid:112) B/ 256 for batch size B .
The base LR η base is chosen 0 .
01 rather than 0 .
001, since the latter has resulted in worse performances.
The results are summarized in Table [Table] and the plot of LR schedules and training loss curves can be found in Figure 3.
We observe that general tendency is similar to that of AutoWU coupled with AdamP, but the resulting performances are slightly worse than those obtained by the baseline schedule with an exception of batch size 8192 on
CIFAR-100.
However, we emphasize that AutoWU does not require tuning, in contrast to the fact that LAMB can be more sensitive to η base (Chen et al., 2021) and adjusting η base to 0.01 in baselines was necessary to achieve the reported performance.
We also evaluate the performance of AutoWU on ImageNet classiﬁcation task with two other architectures: EﬃcientNet-B0 and ViT-S/16.
The training conﬁguration of
EﬃcientNet-B0 is identical to that of ResNet-50, except that the label smoothing of factor 0.1 is used.
For ViT-S/16, it is trained for 300 epochs with augmentations as described in Touvron et al.
The peak LR of the baseline is scaled as η base (cid:112) B/ 256 with η base = 0 .
001 for EﬃcientNet-B0 and η base = 0 .
0005 for ViT-S/16.
Results are found in
[Table] [Figure] Table 3, and the plots of LR schedules and training loss curves are found in Figure 4.
Even though the performances are a little bit lower compared to the tuned baseline, AutoWU stably works on other architectures without any speciﬁc hyperparameter tuning
[Table]
Appendix C. Ablation Studies C.1 Sensitivity of the baseline with respect to the warmup schedule
We compare the performances of the baseline LR scheduling according to various numbers of warmup epochs and peak LRs.
Namely, the LR is linearly increased from 0 to the peak LR η peak for n warmup epochs then decayed to 0 via cosine schedule, and the experiment is carried out for η peak ∈ { 0 . 002 , 0 . 004 , 0 . 008 , 0 . 016 , 0 . 032 } and n warmup ∈ { 5 , 20 , 40 , 60 } .
Here, we use
ResNet-50 trained by AdamP with a batch size of 16384 on ImageNet, and therefore the conﬁguration ( η peak , n warmup ) = (0 . 008 , 5) corresponds to the baseline schedule (presented in Table 1).
Figure 5 demonstrates how the performance of ResNet-50 on ImageNet changes with respect to η peak and n warmup .
The value of optimal LR is either 0.008 or 0.016 when n warmup is ﬁxed, which includes our baseline.
On the other hand, we ﬁnd that the validation accuracy is improved as n warmup increases and the best performance is 75.84% obtained by ( η peak , n warmup ) =
We remark that not only the performances of AutoWU with cosine and constant-then-cosine decay (75.22% and 75.62%, respectively) are fairly close to the best performance but also the peak LR and the warmup epochs found by
AutoWU , (0 . 0215 , 40) as shown in Figure 2(c), are similar to the best conﬁguration on the baseline schedule.
This validates the eﬀectiveness of the proposed algorithm.
C.2 Dependency of AutoWU on the warmup schedule
We have argued that the exponential growth in the warmup phase stabilizes the training and enables a ﬁne-grained LR exploration.
To demonstrate this, we compare the linear growth and the exponential growth in case of ResNet-50 training on ImageNet with AdamP and batch size 16384.
Speciﬁcally, we consider the linear schedule in the warmup phase deﬁned as
[Equation]
for
[Equation]
(6) where η min = 10 − 5 and η max ∈ { 0 . 1 , 1 . 0 } .
In both choices of η max , the cosine schedule is used in the decay phase and all hyperparameters are set identical as in case of AutoWU with exponential warmup schedule.
Figure 6 shows the LR schedule, the training loss and the validation accuracy of the baseline, AutoWU with exponential growth, and AutoWU with linear growth as described above.
In terms of the validation accuracy, the baseline and AutoWU with exponential growth achieves 75.02% and 75.22% respectively, while AutoWU with linear growth [Table] achieves 29.65% for η max = 1 . 0 and 73.56% for η max = 0 .
This implies that the linear warmup makes AutoWU very sensitive to the choice of η max , hence a signiﬁcant amount of eﬀort to tune η max must be made, obliterating the purpose of automated LR scheduling.
The growth factor γ in Eqn. 1 in AutoWU with the exponential growth is completely determined by ρ w when η min and η max are set to the suﬃciently small and large values; γ is increased if ρ w is decreased.
Therefore, we also compare the performances of AutoWU with the cosine decay for diﬀerent values of ρ w .
Speciﬁcally, ResNet-50 is trained on ImageNet with a batch size of 16384, AdamP, and
AutoWU for ρ w ∈ { 0 . 125 , 0 . 25 , 0 . 5 } , and the results are summarized in Table 4.
Here, ρ w = 0 .
5 corresponds to the default conﬁguration which is also reported in Table 1.
The best performances are attained by ρ w = 0 .
25 or ρ w = 0 . 5, and when ρ w = 0 .
125, the performances are degraded.
We plot the relation between the starting LR η t ∗ of the decay phase and the choice of ρ w in Figure 7.
Larger batch size or faster growth implies the bigger starting LR in the decay phase in general, but this breaks down when the batch size becomes very large.
We have observed that this “critical batch size” is larger with larger ρ w ( i.e. smaller γ ), and supports the intuition that slow growth of LR in the warmup phase stabilizes the overall training dynamics.
Such an observation is also consistent with the previous works (Smith and Topin, 2019; Cohen et al., 2021).
Appendix D. Further discussions
We did not present the study investigating the impact of other hyperparameters including the conﬁdence c , patience p , and the choice of decay schedule on the performance of the proposed algorithm AutoWU .
While we leave a thorough study as a future work, we brieﬂy discuss below on a few observations from our preliminary experiments.
• It seems that the choice of conﬁdence c does not matter much, as P min ( f ) in Eqn.
(4) become very close to 1 when the loss starts to increase.
We also note that conducting n test > 1 tests on sub-sampled loss curves made the test more robust.
• The choice of p was more critical.
Note that AutoWU is an algorithm which detects whether the loss minimum lies in the past or not; this implies that AutoWU may be susceptible to a local ﬂuctuation of loss, especially when the number of steps is small (or equivalently, the batch size is large).
To reduce such a sensitivity, the patience p to wait whether the observed minimum is local or not is introduced.
Utilizing a large value of p makes the test more robust, but simultaneously lets LR to grow until a large value hence training may diverge.
We found that setting p = 3 was working well for all experiments.
• We argued that SALSA (Zhang et al., 2020) is the most similar work to AutoWU .
In terms of warmup scheduling, SALSA employs a line search based method to warmup the initial training dynamics and ends the warmup phase when the loss stops to decrease , in the similar spirit as this work.
However, the extra computation required for the line search was less appealing than a simple exponential schedule, and we have observed that the line search based method in general is sensitive to the Armijo constant especially for large batch sizes.