A Closer Look at Deep Learning Heuristics: Learning rate restarts, Warmup and Distillation

Introduction

The introduction of heuristics such as normalization layers (Ioffe & Szegedy, 2015; Ba et al., 2016), residual connections (He et al., 2016), and learning rate strategies (Loshchilov & Hutter, 2016; Goyal et al., 2017; Smith, 2017) have greatly accelerated progress in Deep Learning. Many of these ingredients are now commonplace in modern architectures, and some of them have also been buttressed with theoretical guarantees (Balduzzi et al., 2017; Poggio & Liao, 2017; Hardt & Ma, 2016). However, despite their simplicity and efficacy, why some of these heuristics work is still relatively unknown. Existing attempts at explaining these strategies empirically have been limited to intuitive explanations and the use of tools such as linear interpolation between two models and low-dimensional visualizations (Li et al., 2017) of the loss surface. In our work, we instead use recent tools built specifically for analyzing deep networks, viz., mode connectivity (Garipov et al., 2018) and singular value canonical correlation analysis (SVCCA) (Raghu et al., 2017). We investigate three strategies in detail: (a) cosine learning rate decay, (b) learning rate warmup, and (c) knowledge distillation, and list the summary of our contributions at the end of this section.

Cosine annealing (Loshchilov & Hutter, 2016), also known as stochastic gradient descent with restarts (SGDR), and more generally cyclical learning rate strategies (Smith, 2017), have been recently proposed to accelerate training of deep networks (Coleman et al., 2018). The strategy involves reductions and restarts of learning rates over the course of training, and was motivated as means to escape spurious local minima. Experimental results have shown that SGDR often improves convergence both from the standpoint of iterations needed for convergence and the final objective.

Learning rate warmup (Goyal et al., 2017) also constitutes an important ingredient in training deep networks, especially in the presence of large or dynamic batch sizes. It involves increasing the learning rate to a large value over a certain number of training iterations followed by decreasing the learning rate, which can be performed using step-decay, exponential decay or other such schemes. The strategy was proposed out of the need to induce stability in the initial phase of training with large learning rates (due to large batch sizes). It has been employed in training of several architectures at scale including ResNets and Transformer networks (Vaswani et al., 2017).

Further, we investigate knowledge distillation (KD) (Hinton et al., 2015). This strategy involves first training a (teacher) model on a typical loss function on the available data. Next, a different (student) model (typically much smaller than the teacher model) is trained, but instead of optimizing the loss function defined using hard data labels, this student model is trained to mimic the teacher model. It has been empirically found that a student network trained in this fashion significantly outperforms an identical network trained with the hard data labels. We defer a detailed discussion of the three heuristics, and existing explanations for their efficacy to sections 3, 4 and 5 respectively.

Finally, we briefly describe the tools we employ for analyzing the aforementioned heuristics. Mode connectivity (MC) is a recent observation that shows that, under circumstances, it is possible to connect any two local minima of deep networks via a piecewise-linear curve (Garipov et al., 2018; Draxler et al., 2018). This shows that local optima obtained through different means, and exhibiting different local and generalization properties, are connected. The authors propose an algorithm that locates such a curve. While not proposed as such, we employ this framework to better understand loss surfaces but begin our analysis in Section 2 by first establishing its robustness as a framework.

Deep network analyses focusing on the weights of a network are inherently limited since there are several invariances in this, such as permutation and scaling. Recently, Raghu et al. (2017) propose using CCA along with some pre-processing steps to analyze the activations of networks, such that the resulting comparison is not dependent on permutations and scaling of neurons. They also prove the computational gains of using CCA over alternatives ((Li et al., 2015)) for representational analysis and employ it to better understand many phenomenon in deep learning.

We use mode connectivity and CCA to improve understanding of cosine annealing, learning rate warmup and knowledge distillation. For mode connectivity, we also establish the robustness of the approach across changes in training choices for obtaining the modes.

We demonstrate that the reasons often quoted for the success of cosine annealing are not substantiated by our experiments, and that the iterates move over barriers after restarts but the explanation of escaping local minima might be an oversimplification.

We show that learning rate warmup primarily limits weight changes in the deeper layers and that freezing them achieves similar outcomes as warmup.

We show that the latent knowledge shared by the teacher in knowledge distillation is primarily disbursed in the deeper layers.

Empirical tools

Garipov et al. (2018) introduce a framework, called mode connectivity, to obtain a low loss (or high accuracy, in the case of classification) curve of simple form, such as a piecewise linear curve, that connects optima (modes of the loss function) found independently. This observation suggests that points at the same loss function depth are connected, somewhat contrary to several empirical results claiming that minima are isolated or have barriers between themDraxler et al. (2018) independently report the same observation for neural network loss landscapes, and claim that this is suggestive of the resilience of neural networks to perturbations in model parameters..

To demonstrate that the curve-finding approach works in practice, Garipov et al. (2018) use two optima found using different initializations but a common training scheme which we detail below. We explore the limits of this procedure by connecting optima obtained from different training strategies. Our goal of this investigation is to first establish the robustness of the framework in order to seamlessly use it as a tool for analysis. In particular, we experiment with different initializations, optimizers, data augmentation choices, and hyperparameter settings including regularization, training batch sizes, and learning rate schemes. We note in passing that while the framework was proposed to connect two points in the parameter space that are at equal depth in the loss landscape, it is well-defined to also connect points at different depths; in this case, the path corresponds to one that minimizes the average loss along the curve.

Conventional wisdom suggests that the different training schemes mentioned above will converge to regions in the parameter space that are vastly different from each other. Examples of this include size of minibatches used during training (Keskar et al., 2016), choice of optimizer (Heusel et al., 2017; Wilson et al., 2017), initialization (Goodfellow et al., 2016) and choice of regularizer. Having a high accuracy connection between these pairs would seem counterintuitive.

For obtaining the reference model (named mode $G$), we train the VGG-16 model architecture (Simonyan & Zisserman, 2014) using CIFAR-10 training data (Krizhevsky et al., 2014) for $200$ epochs with SGD. We then build 6 variants of the reference mode $G$ as follows: we obtain mode $A$ using a training batch size of $4000$, mode $B$ by using the Adam optimizer instead of SGD, mode $C$ with a linearly decaying learning rate instead of the step decay used in mode $G$, mode $D$ using a smaller weight decay of $5\times 10^{-6}$, mode $E$ by increasing the variance of the initialization distribution to $3\times\sqrt[]{2/n}$ and mode $F$ using no data augmentation. Note that for the set of modes $\{A,B,C,D,E,F\}$, all the other hyper-parameters and settings except the ones mentioned above are kept same as that for mode $G$. We use the mode connectivity algorithm on each of the $6$ pairs of modes including $G$ and another mode, resulting in curves $GA$, $GB$, $GC$, $GD$, $GE$, and $GF$.

Figure 1 shows the validation accuracy for models on each of the $6$ connecting curves during the $20$th, $40$th, $60$th and $80$th epochs of the mode connectivity training procedure. As described in Section 2.1, for a polychain curve $GX$ (connecting modes $G$ and $X$ using the curve described by $\theta$), model parameters $\phi_{\theta}(t)$ on the curve are given by $p_{\phi_{\theta}(t)}=2(tp_{\theta}+(0.5-t)p_{G})\;\text{if}\;0\leq t\leq 0.5$ and $p_{\phi_{\theta}(t)}=2((t-0.5)p_{X}+(1-t)p_{\theta})\;\text{if}\;0.5<t\leq 1$ where $p_{G}$, $p_{\theta}$ and $p_{X}$ are parameters of the models $G$, $\theta$, and $X$ respectively. Thus $\phi_{\theta}(0)=G$ and $\phi_{\theta}(1)=X$.

In a few epochs of the curve training, for all 6 pairs, we can find a curve such that each point on it generalizes almost as well as models from the pair that is being connected. Note that by virtue of existence of these $6$ curves, there exists a high accuracy connecting curve (albeit with multiple bends) for each of the ${7\choose 2}$ pairs of modes. We refer the reader to Appendix 7 for a t-SNE plot of the modes and their connections, and also for additional plots and details. Having established the high likelihood of the existence of these curves, we use this procedure along with interpolation of the loss surface between parameters at different epochs as tools to analyze the dynamics of SGD and SGDR.

2 CCA for measuring representational similarity

Canonical correlation analysis (CCA) is a classical tool from multivariate statistics (Hotelling, 1936) that investigates the relationships between two sets of random variables. Raghu et al. (2017) have proposed coupling CCA with pre-processing steps like Singular Value Decomposition (SVD) or Discrete Fourier Transform (DFT) to design a similarity metric for two neural net layers that we want to compare. These layers do not have to be of the same size or belong to the same network.

Stochastic Gradient Descent with Restarts (SGDR)

Loshchilov & Hutter (2016) introduced SGDR as a modification to the common linear or step-wise decay of learning rates. The strategy decays learning rates along a cosine curve and then, at the end of the decay, restarts them to its initial value. The learning rate at the $t$-th epoch in SGDR is given by the following expression in (1) where $\eta_{min}$ and $\eta_{max}$ are the lower and upper bounds respectively for the learning rate. $T_{cur}$ represents how many epochs have been performed since the last restart and a warm restart is simulated once $T_{i}$ epochs are performed. Also $T_{i}=T_{mult}\times T_{i-1}$, meaning the period $T_{i}$ for the learning rate variation is increased by a factor of $T_{mult}$ after each restart. Figure 2(b) shows an instance of this learning rate schedule.

While the strategy has been claimed to outperform other learning rate schedulers, little is known why this has been the case. One explanation that has been given in support of SGDR is that it can be useful to deal with multi-modal functions, where the iterates could get stuck in a local optimum and a restart will help them get out of it and explore another region; however, Loshchilov & Hutter (2016) do not claim to observe any effect related to multi-modality. Huang et al. (2017) propose an ensembling strategy using the set of iterates before restarts and claim that, when using the learning rate annealing cycles, the optimization path converges to and escapes from several local minima. We empirically investigate if this is actually the case by interpolating the loss surface between parameters at different epochs and studying the training and validation loss for parameters on the hyperplane passing throughThis hyperplane is the set of all affine combinations of $w_{a},w_{b}\text{ and their connection }\theta=w_{a-b}$. the two modes found by SGDR and their connectivity. Further, by employing the CCA framework as described in Section 2.2, we investigate the progression of training, and the effect of restarts on the model activations.

We train a VGG-16 network (Simonyan & Zisserman, 2014) on the CIFAR-10 dataset using SGDR. For our experiments, we choose $T_{0}=10$ epochs and $T_{mult}=2$ (warm restarts simulated every 10 epochs and the period $T_{i}$ doubled at every new warm restart), $\eta_{max}=0.05$ and $\eta_{min}=10^{-6}$. We also perform VGG training using SGD (with momentum of $0.9$) and a step decay learning rate scheme (initial learning rate of $\eta_{0}=0.05$, scaled by $5$ at epochs $60$ and $150$). Figure 2(b) shows the learning rate variation for these two schemes on a logarithmic scale and Figure 2(a) shows the validation accuracy over training epochs with these two learning rate schemes.

In order to understand the loss landscape on the optimization path of SGDR, the pairs of iterates obtained just before the restarts $\{w_{30},w_{70}\},\{w_{70},w_{150}\}$ and $\{w_{30},w_{150}\}$ are given as inputs to the mode connectivity algorithm, where $w_{n}$ is the model corresponding to parameters at the $n$-th epoch of training. Figure 2(c) shows the training loss for models along the line segment joining these pairs and those on the curve found through mode connectivity. For the baseline case of SGD training, we connect the iterates around the epochs when we decrease our learning rate in the step decay learning rate scheme. Thus, we chose $\{w_{55},w_{65}\},\{w_{145},w_{165}\}$ and $\{w_{55},w_{165}\}$ as input pairs to the mode connectivity algorithm. Figure 2(d) shows the training loss for models along the line segments joining these pairs and the curves found through mode connectivity.

From Figure 2(c), it is clear that for the pairs $\{w_{30},w_{150}\}$ and $\{w_{70},w_{150}\}$ the training loss for points on segment is much higher than the endpoints suggesting that SGDR indeed finds paths that move over a barriera path is said to have moved over or crossed a barrier between epoch $m$ and $n$ ($n>m$) if $\exists$ $w_{t}\in\{\lambda w_{m}+(1-\lambda)w_{n}|\lambda\in\}$ such that $\mathcal{L}(w_{t})>\max\{\mathcal{L}(w_{m}),\mathcal{L}(w_{n})\}$ in the training loss landscape. In contrast, for SGD (without restarts) in Figure 2(d) none of the three pairs show evidence of having a training loss barrier on the line segment joining them. Instead there seems to be an almost linear decrease of training loss along the direction of these line segments, suggesting that SGD’s trajectory is quite different from SGDR’s. We present additional experiments, including results for other metrics, in Appendix 8.

To further understand the SGDR trajectory, we evaluate the intermediate points on the hyperplane in the $D$-dimensional space defined by the three points: $w_{70}$, $w_{150}$ and $w_{70-150}$, where $w_{70-150}$ is the bend point that defines the high accuracy connection for the pair $\{w_{70},w_{150}\}$. Figures 3(a) and 3(b) show the training and validation loss surface for points in this subspace, respectively. Note that the intermediate iterates do not necessarily lie in this plane, and thus are projected. We refer the reader to Appendix 8 for additional details on the projection, and analogous results with $w_{30}$ and $w_{70}$.

Figure 3(a) suggests that SGDR helps the iterates converge to a different region although neither of $w_{70}$ or $w_{150}$ are technically a local minimum, nor do they appear to be lying in different basins, hinting that Huang et al. (2017)’s claims about SGDR converging to and escaping from local minima might be an oversimplification.We note in passing that during our experiments, a strategy that consistently performed well is one of a cosine (or linear) decay over the entire budget. We hypothesize that this decay, and less so the restarts, plays a major role in the success of such strategies. Another insight we can draw from Figure 3(a) is that the path found by mode connectivity corresponds to lower training loss than the loss at the iterates that SGDR converges to ($\mathcal{L}(w_{150})>\mathcal{L}(w_{70-150})$). However, Figure 3(b) shows that models on this curve seem to overfit and not generalize as well as the iterates $w_{70}$ and $w_{150}$. Thus, although gathering models from this connecting curve might seem as a novel and computationally cheap way of creating ensembles, this generalization gap alludes to one limitation in doing so; Garipov et al. (2018) point to other shortcomings of curve ensembling in their original work. In Figure 3, the region of the plane between the iterates $w_{70}$ and $w_{150}$ corresponds to higher training loss but lower validation loss than the two iterates. This hints at a reason why averaging iterates to improve generalization using cyclic or constant learning rates (Izmailov et al., 2018) has been found to work well.

Finally, in Figure 14 in Appendix 9, we present the CCA similarity plots for two pairs of models: epochs 10 and 150 (model at the beginning and end of training), and epochs 150 and 155 (model just before and just after a restart). For standard SGD training, Raghu et al. (2017) observe that the activations of the shallower layers bear closer resemblance than the deeper layers between a partially and fully trained network from a given training run. For SGDR training, we witness similar results (discussed in Appendix 9), meaning that the representational similarities between the network layers at the beginning and end of training are alike for SGDR and SGD, even though restarts lead to a trajectory that tends to cross over barriers.

Warmup learning rate scheme

Learning rate warmup is a common heuristic used by many practitioners for training deep neural nets for computer vision (Goyal et al., 2017) and natural language processing (Bogoychev et al., 2018; Vaswani et al., 2017) tasks. Theoretically, it can be shown that the learning dynamics of SGD rely on the ratio of the batch size and learning rate (Smith et al., 2017; Jastrzebski et al., 2017; Hoffer et al., 2017). And hence, an increase in batch size over a baseline requires an accompanying increase in learning rate for comparable training. However, in cases when the batch size is increased significantly, the curvature of the loss function typically does not support a proportional increase in the learning rate. Warmup is hence motivated as a means to use large learning rates without causing training instability. We particularly focus on the importance of the learning rate schedule’s warmup phase in the large batch (LB) training of deep convolutional neural networks as discussed in Goyal et al. (2017). Their work adopts a linear scaling rule for adjusting the learning rate as a function of the minibatch size, to enable large-batch training. The question we aim to investigate here is: How does learning rate warmup impact different layers of the network?

Using CCA as a tool to study the learning dynamics of neural networks through training iterations, we investigate the differences and similarities for the following $3$ training configurations - (a) large batch training with warmup (LB + warmup), (b) large batch training without warmup (LB no warmup) and (c) small batch training without warmup (SB no warmup). We train a VGG-11 architecture on the CIFAR-10 (Krizhevsky et al., 2014) dataset using SGD with momentum of $0.9$. Learning rate for the small batch case (batch-size of $100$) is set to $0.05$, and for the large batch cases (batch-size of $5000$) is set to $2.5$ as per the scaling rule. For the warmup, we increase the learning rate from to $2.5$ over the first $200$ iterations. Subsequently, we decrease the learning rate as per the step decay schedule for all runs, scaling it down by a factor of $10$ at epochs $60$, $120$ and $150$. We plot the learning rate and validation accuracy for these $3$ cases in Figure 4(b) and (a).

Using CCA and denoting the model at the $j$-th iteration of a training setup by $\textit{iter}_{j}$, we compare activation layers from $\textit{iter}_{0}$ (init.) and $\textit{iter}_{200}$ (end of warmup) for each of the three runs, presented in Figures 5(a), (b) and (c), and also layers from $\textit{iter}_{200}$ (end of warmup) and $\textit{iter}_{2990}$ (end of training) for the LB + warmup case, presented in Figure 5(d). Figure 4(c) plots the similarity for layer $i$ of $\textit{iter}_{a}$ with the same layer of $\textit{iter}_{b}$ (this corresponds to diagonal elements of the matrices in Figure 5) for these three setups.

An evident pattern in Figures 5(a), (b) and (c) is the increase in similarity for the last few layers (stack of fully-connected layers) for the LB + warmup and SB cases, which is absent in the LB without warmup case. This suggests that when used with the large batch size and learning rate, warmup tends to avoid unstably large changes in the fully-connected (FC) stack for this network configuration. To validate this proposition, we train using the LB without warmup setup, but freezing the fully-connected stack for the first $200$ iterations (LB no warmup + FC freeze). Figure 4(d) shows the validation accuracy for this training run in comparison to the three training setups discussed before. The performance is comparable at the end of warmup by freezing the FC stack, suggesting the validity our proposition in this case. We refer the reader to Appendix 10 for analogous results for ResNet-18 and ResNet-32 (He et al., 2016); thus also demonstrating the generality of our claim. Finally, note from Figure 4(d) that no qualitative difference exists in the trajectory beyond the warmup when compared to the standard training approach (Raghu et al., 2017).

Knowledge Distillation

We study knowledge distillation as proposed by Hinton et al. (2015) using CCA to measure representational similarity between layers of the teacher and student model. Distillation involves training a “student” model using the output probability distribution of a “teacher” model. This has been widely known to help the student model perform better than it would, if it were trained using hard labels due to knowledge transfer from the teacher model. The reason often quoted for the success of distillation is the transfer of dark knowledge from the teacher to the student (Hinton et al., 2015), and more recently, as an interpretation of importance weighing (Furlanello et al., 2018). We investigate if this knowledge transfer is limited to certain parts of the network, and if representational similarity between layers of the student and teacher model and a student can help answer this question.

To construct an example of distillation that can be used for our analysis, we use a VGG-16 model (Simonyan & Zisserman, 2014) as our teacher network and a shallow convolutional network ([conv, maxpool, relu]x2, fc, relu, fc, fc, softmax) as the student network. We train the shallow network for CIFAR-10 using the teacher’s predicted probability distribution (softened using a temperature of $5$), ($S_{\text{distilled}}$), and for the baseline, train another instance of the same model in a standard way using hard labels, ($S_{\text{indep.}}$). Over 5 runs for each of the two setups, we find the distillation training attains the best validation accuracy at $85.18\%$ while standard training attains its best at $83.01\%$. We compare their layer-wise representations with those of the teacher network ($T$).

Figure 6 shows the CCA plots and the absolute value of their difference. The scores of these two pairs are quite similar for the shallow layers of the student network relative to the deeper layers, suggesting that the difference that knowledge distillation brings to the training of smaller networks is restricted to the deeper layers (fc stack). Similar results are obtained through different configurations for the student and teacher when the student benefits from the teacher’s knowledge. We hypothesize that the dark knowledge transferred by the teacher is localized majorly in the deeper (discriminative) layers, and less so in the feature extraction layers. We also note that this is not dissimilar to the hypothesis of Furlanello et al. (2018), and also relates ot the results from the literature on fine-tuning or transfer learning (Goodfellow et al., 2016; Yosinski et al., 2014; Howard & Ruder, 2018) which suggest training of only higher layers.

Conclusion

Heuristics have played an important role in accelerating progress of deep learning. Founded in empirical experience, intuition and observations, many of these strategies are now commonplace in architectures. In the absence of strong theoretical guarantees, controlled experiments aimed at explaining the the efficacy of these strategies can aid our understanding of deep learning and the training dynamics. We investigate three such heuristics: cosine annealing, learning rate warmup, and knowledge distillation. For this purpose, we employ recently proposed tools of mode connectivity and CCA. Our empirical analysis sheds light on these heuristics and suggests that: (a) the reasons often quoted for the success of cosine annealing are not evidenced in practice; (b) that the effect of learning rate warmup is to prevent the deeper layers from creating training instability; and (c) that the latent knowledge shared by the teacher is primarily disbursed in the deeper layers. Our empirical experiments and hypotheses open new questions, and encourage a deeper exploration into improving and better understanding these heuristics.

References

Appendix

Additional Results on Robustness of MC

The learning rate is initialized to $0.05$ and scaled down by a factor of $5$ at epochs $\{60,120,160\}$ (step decay). We use a training batch size of 100, momentum of 0.9, and a weight decay of $0.0005$. Elements of the weight vector corresponding to a neuron are initialized randomly from the normal distribution $\mathcal{N}(0,\,\sqrt[]{2/n})$ where $n$ is the number of inputs to the neuron. We also use data augmentation by random cropping of input images.

2 Plots

Figures 7, 8 and 9 show the Validation Loss, Training Accuracy and Training Loss respectively for the curves joining the $6$ pairs discussed in Section 2.1.1. These results too, confirm the overfitting or poor generalization tendency of models on the curve.

3 t-SNE visualization for the 7 modes

We use t-SNE (Maaten & Hinton, 2008) to visualize these $7$ modes and the $\theta$ points that define the connectivity for the $6$ pairs presented in Section 2.1.1, in a $2$-dimensional plot in Figure 10. Since t-SNE is known to map only local information correctly and not preserve global distances, we caution the reader about the limited interpretability of this visualization, it is presented simply to establish the notion of connected modes.

Additional SGDR Results

For completeness, in Figure 11, we present the Validation loss, Validation accuracy and Training accuracy results for the curves and line segments joining iterates from SGDR and SGDR discussed in Figure 2(c) and (d).

2 Projecting iterates

meaning it is the point on the plane (linear combination of $w_{70},w_{150}\;\text{and}\;\theta$) with the least $l$-2 distance from the original point (iterate in this case).

In Section 3, we present some experiments and make observations on the trajectory of SGDR by using the plane defined by the points $w_{70}$, $w_{150}$ and $w_{70-150}$. Here we plot the training loss surface in Figure 12 and the validation loss surface in Figure 13 for another plane defined by SGDR’s iterates $w_{30},w_{70}$ and their connection $w_{30-70}$ to ensure the reader that the observations made are general enough.

SGDR CCA Heatmaps

In Figure 14, we present the CCA similarity plots comparing two pairs of models: epochs 10 and 150, and epochs 150 and 155. The $(i,j)^{th}$ block of the matrix denotes the correlation between the $i^{th}$ layer of the first model and the $j^{th}$ layer of the other. A high correlation implies that the layers learn similar representations and vice versa. We present the former to compare against the typical stepwise or linear decay of SGD, and the latter to demonstrate the immediate effect of restarting on the model. Raghu et al. (2017) showed in their work that for typical SGD training, a CCA similarity plot between a partially and completed trained network reveals that the activations of the shallower layers bears closer resemblance in the two models than the deeper layers. We note that, despite the restart, a similar tendency is seen in SGDR training as well. This again suggests that the restart does not greatly impact the model, both in weights and representations, and especially so in the shallower layers. A comparison of epochs 150 and 155, i.e., before and after a restart also stands as evidence for this hypothesis.

Warmup Experiments on ResNet-18 and ResNet-32

In Figure 4(d), we show that the stability induced by warmup when training with large batches and learning rates can also be obtained by holding the FC stack frozen. This experiment was conducted on the VGG-11 network (Simonyan & Zisserman, 2014). To demonstrate the generality of our claim, we present additional experiments on two ResNet architectures: 18 and 32. The setup for this experiment is identical to the VGG-11 one with one change: instead of the learning rate being set to $2.5$, which is the learning rate for SB ($0.05$) times the batch size increase ($50\times$), we set it to $5.0$ since SB training is better with $0.1$. For the warmup case, we linearly increase the learning rate from to $5$ again for $20$ epochs. Experiments on other configurations yielded similar results. Whether these results remain true also for training larger datasets, such as ImageNet, remains to be shown and is a topic of future research.