Restart techniques are common in gradient-free optimization to deal with multimodal functions. Partial warm restarts are also gaining popularity in gradientbased optimization to improve the rate of convergence in accelerated gradient schemes to deal with ill-conditioned functions. In this paper, we propose a simple warm restart technique for stochastic gradient descent to improve its anytime performance when training deep neural networks. We empirically study its performance on the CIFAR-10 and CIFAR-100 datasets, where we demonstrate new state-of-the-art results at 3.14% and 16.21%, respectively. We also demonstrate its advantages on a dataset of EEG recordings and on a downsampled version of the ImageNet dataset. Our source code is available at https://github.com/loshchil/SGDR
1 INTRODUCTION
Deep neural networks (DNNs) are currently the best-performing method for many classification problems, such as object recognition from images (Krizhevsky et al., 2012a; Donahue et al., 2014) or speech recognition from audio data (Deng et al., 2013). Their training on large datasets (where DNNs perform particularly well) is the main computational bottleneck: it often requires several days, even on high-performance GPUs, and any speedups would be of substantial value.
The training of a DNN with n free parameters can be formulated as the problem of minimizing a function f : IRn → IR. The commonly used procedure to optimize f is to iteratively adjust xt ∈ IRn (the parameter vector at time step t) using gradient information ∇ft(xt) obtained on a relatively small t-th batch of b datapoints. The Stochastic Gradient Descent (SGD) procedure then becomes an extension of the Gradient Descent (GD) to stochastic optimization of f as follows:
xt+1 = xt − ηt∇ft(xt), (1)
where ηt is a learning rate. One would like to consider second-order information
xt+1 = xt − ηtH−1t ∇ft(xt), (2)
but this is often infeasible since the computation and storage of the inverse Hessian H−1t is intractable for large n. The usual way to deal with this problem by using limited-memory quasiNewton methods such as L-BFGS (Liu & Nocedal, 1989) is not currently in favor in deep learning, not the least due to (i) the stochasticity of ∇ft(xt), (ii) ill-conditioning of f and (iii) the presence of saddle points as a result of the hierarchical geometric structure of the parameter space (Fukumizu & Amari, 2000). Despite some recent progress in understanding and addressing the latter problems (Bordes et al., 2009; Dauphin et al., 2014; Choromanska et al., 2014; Dauphin et al., 2015), state-ofthe-art optimization techniques attempt to approximate the inverse Hessian in a reduced way, e.g., by considering only its diagonal to achieve adaptive learning rates. AdaDelta (Zeiler, 2012) and Adam (Kingma & Ba, 2014) are notable examples of such methods.
Intriguingly enough, the current state-of-the-art results on CIFAR-10, CIFAR-100, SVHN, ImageNet, PASCAL VOC and MS COCO datasets were obtained by Residual Neural Networks (He et al., 2015; Huang et al., 2016c; He et al., 2016; Zagoruyko & Komodakis, 2016) trained without the use of advanced methods such as AdaDelta and Adam. Instead, they simply use SGD with momentum 1:
vt+1 = µtvt − ηt∇ft(xt), (3) xt+1 = xt + vt+1, (4)
where vt is a velocity vector initially set to 0, ηt is a decreasing learning rate and µt is a momentum rate which defines the trade-off between the current and past observations of ∇ft(xt). The main difficulty in training a DNN is then associated with the scheduling of the learning rate and the amount of L2 weight decay regularization employed. A common learning rate schedule is to use a constant learning rate and divide it by a fixed constant in (approximately) regular intervals. The blue line in Figure 1 shows an example of such a schedule, as used by Zagoruyko & Komodakis (2016) to obtain the state-of-the-art results on CIFAR-10, CIFAR-100 and SVHN datasets.
In this paper, we propose to periodically simulate warm restarts of SGD, where in each restart the learning rate is initialized to some value and is scheduled to decrease. Four different instantiations of this new learning rate schedule are visualized in Figure 1. Our empirical results suggest that SGD with warm restarts requires 2× to 4× fewer epochs than the currently-used learning rate schedule schemes to achieve comparable or even better results. Furthermore, combining the networks obtained right before restarts in an ensemble following the approach proposed by Huang et al. (2016a) improves our results further to 3.14% for CIFAR-10 and 16.21% for CIFAR-100. We also demonstrate its advantages on a dataset of EEG recordings and on a downsampled version of the ImageNet dataset.
1More specifically, they employ Nesterov’s momentum (Nesterov, 1983; 2013)
2 RELATED WORK
2.1 RESTARTS IN GRADIENT-FREE OPTIMIZATION
When optimizing multimodal functions one may want to find all global and local optima. The tractability of this task depends on the landscape of the function at hand and the budget of function evaluations. Gradient-free optimization approaches based on niching methods (Preuss, 2015) usually can deal with this task by covering the search space with dynamically allocated niches of local optimizers. However, these methods usually work only for relatively small search spaces, e.g., n < 10, and do not scale up due to the curse of dimensionality (Preuss, 2010). Instead, the current state-of-the-art gradient-free optimizers employ various restart mechanisms (Hansen, 2009; Loshchilov et al., 2012). One way to deal with multimodal functions is to iteratively sample a large number λ of candidate solutions, make a step towards better solutions and slowly shape the sampling distribution to maximize the likelihood of successful steps to appear again (Hansen & Kern, 2004). The larger the λ, the more global search is performed requiring more function evaluations. In order to achieve good anytime performance, it is common to start with a small λ and increase it (e.g., by doubling) after each restart. This approach works best on multimodal functions with a global funnel structure and also improves the results on ill-conditioned problems where numerical issues might lead to premature convergence when λ is small (Hansen, 2009).
2.2 RESTARTS IN GRADIENT-BASED OPTIMIZATION
Gradient-based optimization algorithms such as BFGS can also perform restarts to deal with multimodal functions (Ros, 2009). In large-scale settings when the usual number of variables n is on the order of 103 − 109, the availability of gradient information provides a speedup of a factor of n w.r.t. gradient-free approaches. Warm restarts are usually employed to improve the convergence rate rather than to deal with multimodality: often it is sufficient to approach any local optimum to a given precision and in many cases the problem at hand is unimodal. Fletcher & Reeves (1964) proposed to flesh the history of conjugate gradient method every n or (n + 1) iterations. Powell (1977) proposed to check whether enough orthogonality between ∇f(xt−1) and ∇f(xt) has been lost to warrant another warm restart. Recently, O’Donoghue & Candes (2012) noted that the iterates of accelerated gradient schemes proposed by Nesterov (1983; 2013) exhibit a periodic behavior if momentum is overused. The period of the oscillations is proportional to the square root of the local condition number of the (smooth convex) objective function. The authors showed that fixed warm restarts of the algorithm with a period proportional to the conditional number achieves the optimal linear convergence rate of the original accelerated gradient scheme. Since the condition number is an unknown parameter and its value may vary during the search, they proposed two adaptive warm restart techniques (O’Donoghue & Candes, 2012):
• The function scheme restarts whenever the objective function increases.
• The gradient scheme restarts whenever the angle between the momentum term and the negative gradient is obtuse, i.e, when the momentum seems to be taking us in a bad direction, as measured by the negative gradient at that point. This scheme resembles the one of Powell (1977) for the conjugate gradient method.
O’Donoghue & Candes (2012) showed (and it was confirmed in a set of follow-up works) that these simple schemes provide an acceleration on smooth functions and can be adjusted to accelerate stateof-the-art methods such as FISTA on nonsmooth functions.
Smith (2015; 2016) recently introduced cyclical learning rates for deep learning, his approach is closely-related to our approach in its spirit and formulation but does not focus on restarts.
Yang & Lin (2015) showed that Stochastic subGradient Descent with restarts can achieve a linear convergence rate for a class of non-smooth and non-strongly convex optimization problems where the epigraph of the objective function is a polyhedron. In contrast to our work, they never increase the learning rate to perform restarts but decrease it geometrically at each epoch. To perform restarts, they periodically reset the current solution to the averaged solution from the previous epoch.
3 STOCHASTIC GRADIENT DESCENT WITH WARM RESTARTS (SGDR)
The existing restart techniques can also be used for stochastic gradient descent if the stochasticity is taken into account. Since gradients and loss values can vary widely from one batch of the data to another, one should denoise the incoming information: by considering averaged gradients and losses, e.g., once per epoch, the above-mentioned restart techniques can be used again.
In this work, we consider one of the simplest warm restart approaches. We simulate a new warmstarted run / restart of SGD once Ti epochs are performed, where i is the index of the run. Importantly, the restarts are not performed from scratch but emulated by increasing the learning rate ηt while the old value of xt is used as an initial solution. The amount of this increase controls to which extent the previously acquired information (e.g., momentum) is used.
Within the i-th run, we decay the learning rate with a cosine annealing for each batch as follows:
ηt = η i min +
1 2 (ηimax − ηimin)(1 + cos( Tcur Ti π)), (5)
where ηimin and η i max are ranges for the learning rate, and Tcur accounts for how many epochs have been performed since the last restart. Since Tcur is updated at each batch iteration t, it can take discredited values such as 0.1, 0.2, etc. Thus, ηt = ηimax when t = 0 and Tcur = 0. Once Tcur = Ti, the cos function will output −1 and thus ηt = ηimin. The decrease of the learning rate is shown in Figure 1 for fixed Ti = 50, Ti = 100 and Ti = 200; note that the logarithmic axis obfuscates the typical shape of the cosine function.
In order to improve anytime performance, we suggest an option to start with an initially small Ti and increase it by a factor of Tmult at every restart (see, e.g., Figure 1 for T0 = 1, Tmult = 2 and T0 = 10, Tmult = 2). It might be of great interest to decrease ηimax and η i min at every new restart. However, for the sake of simplicity, here, we keep ηimax and η i min the same for every i to reduce the number of hyperparameters involved.
Since our simulated warm restarts (the increase of the learning rate) often temporarily worsen performance, we do not always use the last xt as our recommendation for the best solution (also called the incumbent solution). While our recommendation during the first run (before the first restart) is indeed the last xt, our recommendation after this is a solution obtained at the end of the last performed run at ηt = ηimin. We emphasize that with the help of this strategy, our method does not require a separate validation data set to determine a recommendation.
4 EXPERIMENTAL RESULTS
4.1 EXPERIMENTAL SETTINGS
We consider the problem of training Wide Residual Neural Networks (WRNs; see Zagoruyko & Komodakis (2016) for details) on the CIFAR-10 and CIFAR-100 datasets (Krizhevsky, 2009). We will use the abbreviation WRN-d-k to denote a WRN with depth d and width k. Zagoruyko & Komodakis (2016) obtained the best results with a WRN-28-10 architecture, i.e., a Residual Neural Network with d = 28 layers and k = 10 times more filters per layer than used in the original Residual Neural Networks (He et al., 2015; 2016).
The CIFAR-10 and CIFAR-100 datasets (Krizhevsky, 2009) consist of 32×32 color images drawn from 10 and 100 classes, respectively, split into 50,000 train and 10,000 test images. For image preprocessing Zagoruyko & Komodakis (2016) performed global contrast normalization and ZCA whitening. For data augmentation they performed horizontal flips and random crops from the image padded by 4 pixels on each side, filling missing pixels with reflections of the original image.
For training, Zagoruyko & Komodakis (2016) used SGD with Nesterov’s momentum with initial learning rate set to η0 = 0.1, weight decay to 0.0005, dampening to 0, momentum to 0.9 and minibatch size to 128. The learning rate is dropped by a factor of 0.2 at 60, 120 and 160 epochs, with a total budget of 200 epochs. We reproduce the results of Zagoruyko & Komodakis (2016) with the same settings except that i) we subtract per-pixel mean only and do not use ZCA whitening; ii) we use SGD with momentum as described by eq. (3-4) and not Nesterov’s momentum.
The schedule of ηt used by Zagoruyko & Komodakis (2016) is depicted by the blue line in Figure 1. The same schedule but with η0 = 0.05 is depicted by the red line. The schedule of ηt used in SGDR is also shown in Figure 1, with two initial learning rates T0 and two restart doubling periods.
4.2 SINGLE-MODEL RESULTS
Table 1 shows that our experiments reproduce the results given by Zagoruyko & Komodakis (2016) for WRN-28-10 both on CIFAR-10 and CIFAR-100. These “default” experiments with η0 = 0.1 and η0 = 0.05 correspond to the blue and red lines in Figure 2. The results for η0 = 0.05 show better performance, and therefore we use η0 = 0.05 in our later experiments.
SGDR with T0 = 50, T0 = 100 and T0 = 200 for Tmult = 1 perform warm restarts every 50, 100 and 200 epochs, respectively. A single run of SGD with the schedule given by eq. (5) for T0 = 200 shows the best results suggesting that the original schedule of WRNs might be suboptimal w.r.t. the test error in these settings. However, the same setting with T0 = 200 leads to the worst anytime performance except for the very last epochs.
SGDR with T0 = 1, Tmult = 2 and T0 = 10, Tmult = 2 performs its first restart after 1 and 10 epochs, respectively. Then, it doubles the maximum number of epochs for every new restart. The main purpose of this doubling is to reach good test error as soon as possible, i.e., achieve good anytime performance. Figure 2 shows that this is achieved and test errors around 4% on CIFAR-10 and around 20% on CIFAR-100 can be obtained about 2-4 times faster than with the default schedule used by Zagoruyko & Komodakis (2016).
Since SGDR achieves good performance faster, it may allow us to train larger networks. We therefore investigated whether results on CIFAR-10 and CIFAR-100 can be further improved by making WRNs two times wider, i.e., by training WRN-28-20 instead of WRN-28-10. Table 1 shows that the results indeed improved, by about 0.25% on CIFAR-10 and by about 0.5-1.0% on CIFAR-100. While network architecture WRN-28-20 requires roughly three-four times more computation than WRN-28-10, the aggressive learning rate reduction of SGDR nevertheless allowed us to achieve a better error rate in the same time on WRN-28-20 as we spent on 200 epochs of training on WRN28-10. Specifically, Figure 2 (right middle and right bottom) show that after only 50 epochs, SGDR (even without restarts, using T0 = 50, Tmult = 1) achieved an error rate below 19% (whereas none of the other learning methods performed better than 19.5% on WRN-28-10). We therefore have hope that – by enabling researchers to test new architectures faster – SGDR’s good anytime performance may also lead to improvements of the state of the art.
In a final experiment for SGDR by itself, Figure 7 in the appendix compares SGDR and the default schedule with respect to training and test performance. As the figure shows, SGDR optimizes training loss faster than the standard default schedule until about epoch 120. After this, the default schedule overfits, as can be seen by an increase of the test error both on CIFAR-10 and CIFAR-100 (see, e.g., the right middle plot of Figure 7). In contrast, we only witnessed very mild overfitting for SGDR.
4.3 ENSEMBLE RESULTS
Our initial arXiv report on SGDR (Loshchilov & Hutter, 2016) inspired a follow-up study by Huang et al. (2016a) in which the authors suggest to takeM snapshots of the models obtained by SGDR (in their paper referred to as cyclical learning rate schedule and cosine annealing cycles) right before M last restarts and to use those to build an ensemble, thereby obtaining ensembles “for free” (in contrast to having to perform multiple independent runs). The authors demonstrated new state-of-
the-art results on CIFAR datasets by making ensembles of DenseNet models (Huang et al., 2016b). Here, we investigate whether their conclusions hold for WRNs used in our study. We used WRN28-10 trained by SGDR with T0 = 10, Tmult = 2 as our baseline model.
Figure 3 and Table 2 aggregate the results of our study. The original test error of 4.03% on CIFAR-10 and 19.57% on CIFAR-100 (median of 16 runs) can be improved to 3.51% on CIFAR-10 and 17.75% on CIFAR-100 when M = 3 snapshots are taken at epochs 30, 70 and 150: when the learning rate of SGDR with T0 = 10, Tmult = 2 is scheduled to achieve 0 (see Figure 1) and the models are used with uniform weights to build an ensemble. To achieve the same result, one would have to aggregate N = 3 models obtained at epoch 150 of N = 3 independent runs (see N = 3,M = 1 in Figure 3). Thus, the aggregation from snapshots provides a 3-fold speedup in these settings because additional (M > 1-th) snapshots from a single SGDR run are computationally free. Interestingly, aggregation of models from independent runs (when N > 1 and M = 1) does not scale up as well as from M > 1 snapshots of independent runs when the same number of models is considered: the case of N = 3 and M = 3 provides better performance than the cases of M = 1 with N = 18 and N = 21. Not only the number of snapshots M per run but also their origin is crucial. Thus, naively building ensembles from models obtained at last epochs only (i.e., M = 3 snapshots at epochs 148, 149, 150) did not improve the results (i.e., the baseline of M = 1 snapshot at 150) thereby confirming the conclusion of Huang et al. (2016a) that snapshots of SGDR provide a useful diversity of predictions for ensembles.
Three runs (N = 3) of SGDR with M = 3 snapshots per run are sufficient to greatly improve the results to 3.25% on CIFAR-10 and 16.64% on CIFAR-100 outperforming the results of Huang et al. (2016a). By increasing N to 16 one can achieve 3.14% on CIFAR-10 and 16.21% on CIFAR-100. We believe that these results could be further improved by considering better baseline models than WRN-28-10 (e.g., WRN-28-20).
4.4 EXPERIMENTS ON A DATASET OF EEG RECORDINGS
To demonstrate the generality of SGDR, we also considered a very different domain: a dataset of electroencephalographic (EEG) recordings of brain activity for classification of actual right and left hand and foot movements of 14 subjects with roughly 1000 trials per subject. The best classification results obtained with the original pipeline based on convolutional neural networks [R. Schirrmeister et al. Convolutional neural networks for EEG analysis: Design choices, training strategies, and feature visualization., under review at Neuroimage] were used as our reference. First, we compared the baseline learning rate schedule with different settings of the total number of epochs and initial learning rates (see Figure 4). When 30 epochs were considered, we dropped the learning rate by a factor of 10 at epoch indexes 10, 15 and 20. As expected, with more epochs used and a similar (budget proportional) schedule better results can be achieved. Alternatively, one can consider SGDR and get a similar final performance while having a better anytime performance without defining the total budget of epochs beforehand.
Similarly to our results on the CIFAR datasets, our experiments with the EEG data confirm that snapshots are useful and the median reference error (about 9%) can be improved i) by 1-2% when model snapshots of a single run are considered, and ii) by 2-3% when model snapshots from both hyperparameter settings are considered. The latter would correspond to N = 2 in Section (4.3).
4.5 PRELIMINARY EXPERIMENTS ON A DOWNSAMPLED IMAGENET DATASET
In order to additionally validate our SGDR on a larger dataset, we constructed a downsampled version of the ImageNet dataset [P. Chrabaszcz, I. Loshchilov and F. Hutter. A Downsampled Variant of ImageNet as an Alternative to the CIFAR datasets., in preparation]. In contrast to earlier attempts (Pouransari & Ghili, 2015), our downsampled ImageNet contains exactly the same images from 1000 classes as the original ImageNet but resized with box downsampling to 32 × 32 pixels. Thus, this dataset is substantially harder than the original ImageNet dataset because the average number of pixels per image is now two orders of magnitude smaller. The new dataset is also more difficult than the CIFAR datasets because more classes are used and the relevant objects to be classified often cover only a tiny subspace of the image and not most of the image as in the CIFAR datasets.
We benchmarked SGD with momentum with the default learning rate schedule, SGDR with T0 = 1, Tmult = 2 and SGDR with T0 = 10, Tmult = 2 on WRN-28-10, all trained with 4 settings of the initial learning rate ηimax: 0.050, 0.025, 0.01 and 0.005. We used the same data augmentation procedure as for the CIFAR datasets. Similarly to the results on the CIFAR datasets, Figure 5 shows that SGDR demonstrates better anytime performance. SGDR with T0 = 10, Tmult = 2, ηimax = 0.01 achieves top-1 error of 39.24% and top-5 error of 17.17% matching the original results by AlexNets (40.7% and 18.2%, respectively) obtained on the original ImageNet with full-size images of ca. 50 times more pixels per image (Krizhevsky et al., 2012b). Interestingly, when the dataset is permuted only within 10 subgroups each formed from 100 classes, SGDR also demonstrates better results (see Figure 8 in the Supplementary Material). An interpretation of this might be that while the initial learning rate seems to be very important, SGDR reduces the problem of improper selection of the latter by scanning / annealing from the initial learning rate to 0.
Clearly, longer runs (more than 40 epochs considered in this preliminary experiment) and hyperparameter tuning of learning rates, regularization and other hyperparameters shall further improve the results.
5 DISCUSSION
Our results suggest that even without any restarts the proposed aggressive learning rate schedule given by eq. (5) is competitive w.r.t. the default schedule when training WRNs on the CIFAR10 (e.g., for T0 = 200, Tmult = 1) and CIFAR-100 datasets. In practice, the proposed schedule requires only two hyper-parameters to be defined: the initial learning rate and the total number of epochs.
We found that the anytime performance of SGDR remain similar when shorter epochs are considered (see section 8.1 in the Supplemenary Material).
One should not suppose that the parameter values used in this study and many other works with (Residual) Neural Networks are selected to demonstrate the fastest decrease of the training error. Instead, the best validation or / and test errors are in focus. Notably, the validation error is rarely used when training Residual Neural Networks because the recommendation is defined by the final solution (in our approach, the final solution of each run). One could use the validation error to determine the optimal initial learning rate and then run on the whole dataset; this could further improve results.
The main purpose of our proposed warm restart scheme for SGD is to improve its anytime performance. While we mentioned that restarts can be useful to deal with multi-modal functions, we do not claim that we observe any effect related to multi-modality.
As we noted earlier, one could decrease ηimax and η i min at every new warm restart to control the amount of divergence. If new restarts are worse than the old ones w.r.t. validation error, then one might also consider going back to the last best solution and perform a new restart with adjusted hyperparameters.
Our results reproduce the finding by Huang et al. (2016a) that intermediate models generated by SGDR can be used to build efficient ensembles at no cost. This finding makes SGDR especially attractive for scenarios when ensemble building is considered.
6 CONCLUSION
In this paper, we investigated a simple warm restart mechanism for SGD to accelerate the training of DNNs. Our SGDR simulates warm restarts by scheduling the learning rate to achieve competitive results on CIFAR-10 and CIFAR-100 roughly two to four times faster. We also achieved new stateof-the-art results with SGDR, mainly by using even wider WRNs and ensembles of snapshots from
SGDR’s trajectory. Future empirical studies should also consider the SVHN, ImageNet and MS COCO datasets, for which Residual Neural Networks showed the best results so far. Our preliminary results on a dataset of EEG recordings suggest that SGDR delivers better and better results as we carry out more restarts and use more model snapshots. The results on our downsampled ImageNet dataset suggest that SGDR might also reduce the problem of learning rate selection because the annealing and restarts of SGDR scan / consider a range of learning rate values. Future work should consider warm restarts for other popular training algorithms such as AdaDelta (Zeiler, 2012) and Adam (Kingma & Ba, 2014).
Alternative network structures should be also considered; e.g., soon after our initial arXiv report (Loshchilov & Hutter, 2016), Zhang et al. (2016); Huang et al. (2016b); Han et al. (2016) reported that WRNs models can be replaced by more memory-efficient models. Thus, it should be tested whether our results for individual models and ensembles can be further improved by using their networks instead of WRNs. Deep compression methods (Han et al., 2015) can be used to reduce the time and memory costs of DNNs and their ensembles.
7 ACKNOWLEDGMENTS
This work was supported by the German Research Foundation (DFG), under the BrainLinksBrainTools Cluster of Excellence (grant number EXC 1086). We thank Gao Huang, Kilian Quirin Weinberger, Jost Tobias Springenberg, Mark Schmidt and three anonymous reviewers for their helpful comments and suggestions.
8 SUPPLEMENTARY MATERIAL
8.1 50K VS 100K EXAMPLES PER EPOCH
Our data augmentation procedure code is inherited from the Lasagne Recipe code for ResNets where flipped images are added to the training set. This doubles the number of training examples per epoch and thus might impact the results because hyperparameter values defined as a function of epoch index have a different meaning. While our experimental results given in Table 1 reproduced the results obtained by Zagoruyko & Komodakis (2016), here we test whether SGDR still makes sense for WRN-28-1 (i.e., ResNet with 28 layers) where one epoch corresponds to 50k training examples. We investigate different learning rate values for the default learning rate schedule (4 values out of [0.01, 0.025, 0.05, 0.1]) and SGDR (3 values out of [0.025, 0.05, 0.1]). In line with the results given in the main paper, Figure 6 suggests that SGDR is competitive in terms of anytime performance.

This heuristic to improve gradient descent in image classification is simple and effective, but this looks to me more like a workshop track paper. Demonstration of the algorithm is limited to one task (CIFAR) and there is no theory to support it, so we do not know how it will generalize on other tasks
Working on DNNs for NLP, I find some observations in the paper opposite to my own experience. In particular, with architectures that combine a wide variety of layer types (embedding, RNN, CNN, gating), I found that ADAM-type techniques far outperform simple SGD with momentum, as they save searching for the right learning rate for each type of layer. But ADAM only works well combined with Poliak averaging, as it fluctuates a lot from one batch to another.
Revision:
-  the authors substantially improved the contents of the paper, including experiments on another set than Cifar
-  the workshop track has been modified to breakthrough work, so my recommendation for it is not longer appropriate
I have therefore improved my rating


Stochastic Gradient Descent (SGD) and its variants are the most important optimization algorithms used in large scale machine learning. Mini-batch version of stochastic gradient is often used in practice for taking advantage of hardware parallelism. In this work, we analyze the effect of mini-batch size over SGD convergence for the case of general non-convex objective functions. Building on the past analyses, we justify mathematically that there can often be a large difference between the convergence guarantees provided by small and large mini-batches (given each instance processes equal number of training samples), while providing experimental evidence for the same. Going further to distributed settings, we show that an analogous effect holds with popular Asynchronous Gradient Descent (ASGD): there can be a large difference between convergence guarantees with increasing number of learners given that the cumulative number of training samples processed remains the same. Thus there is an inherent (and similar) inefficiency introduced in the convergence behavior when we attempt to take advantage of parallelism, either by increasing mini-batch size or by increase the number of learners.
1 INTRODUCTION
Stochastic gradient descent (SGD) and its parallel variants form the backbone of most popular deep learning applications. Consequently, there has been a significant interest in investigating their convergence properties. SGD has been shown to satisfy an asymptotic convergence rate of O(1/S) for convex objective functions [Nemirovski et al. (2009)] and an asymptotic convergence rate of O(1/ √ S) for general non-convex objective functions with mini-batch size 1 in [Ghadimi & Lan (2013)] or with arbitrary mini-batch sizes in [Dekel et al. (2012)].
Although SGD converges asymptotically with the same rate irrespective of mini-batch size, it has been reported that for large mini-batch sizes, often it is slower to converge - for example, see Wilson & Martinez (2003) for the detailed graphical illustrations therein for the effect of increasing batchsize or see Bottou (2010) for the comments on the tradeoffs of mini-batching. In this work, we are interested in using theoretical analysis for justifying such practical observations or comments. In particular, we show the following:
• We consider general non-convex objective functions and show that prior to reaching asymptotic regime, SGD convergence could get much slower (inferred from the difference in the convergence rate guarantees from Theorem 2) with using higher mini-batch size, assuming a constant learning rate. Here, to evaluate the convergence rate guarantee we use the measure of average gradient norm since we are considering general non-convex objectives. As a consequence of slower convergence, the number of training samples required to attain a certain convergence guarantee (in terms of average gradient norm) increases as the mini-batch size increases. We build the analysis based on the framework in Ghadimi & Lan (2013).
• Further, we investigate Asynchronous Stochastic Gradient Descent (ASGD) which is one of the most popular asynchronous variants of SGD [Dean et al. (2012); Li et al. (2014a);
Chilimbi et al. (2014)]. Recently, Lian et al. (2015) extended the results of SGD convergence to ASGD and showed that it converges asymptotically with a convergence rate of O(1/ √ S). In our analysis we show that prior to the asymptotic regime, with using higher
number of learners ASGD convergence could get much slower in terms of average gradient norm attained after cumulatively processing a fixed number of training samples (this slow-down is inferred from the difference in the convergence guarantees from Theorem 4). • This suggests that there is an inherent limit on harnessing parallelism with SGD either
by increasing mini-batch size or increasing the number of learners (even when we do not take into account overheads such as communication cost). The difference in convergence behavior caused by increasing mini-batch size for SGD and by increasing the number of learners with ASGD was found to be similar (See Theorem 2, Theorem 4 and the discussion at the end of Section 4).
For rest of the paper, we use the following notation: Let F (x, zi) denote a non-convex function of a parameter vector x and a training sample zi selected from the training set {z1, z2, . . . , zn}. Our aim is to find a parameter vector that minimizes the objective function f(x) = EzF (x, z). Towards this, we use mini-batch SGD, where in kth iteration, we select a random mini-batch zk = {z1k, z2k, . . . , zMk } of size M and perform the following update:
xk+1 = xk − γ ·G(x, zk) = xk − γ · M∑ i=1 G(x, zik) (1)
In the above equation, γ denotes the learning rate and we have G(x, zk) = ∑M i=1G(x, z i k) where G(x, zik) denotes the gradient of the objective function f(x) with respect to a training sample z i k ∈ zk. We define Df := f(x1) − f(x∗) where x1 is the initial parameter vector and x∗ is a local optima towards which SGD proceeds. We also denote by S the total number of training samples to be processed.
Additionally, we make the following standard assumptions, see e.g., [Lian et al. (2015), Ghadimi & Lan (2013)]:
A.1 Unbiased estimator: We assume that the expectation of G(x, z) equals to the true value of the gradient, i.e., EzG(x, z) = ∇f(x) ∀x.
A.2 Bounded variance: We assume that there exists a constant σ2 such that Ez(‖G(x, z)−∇f(x)‖2) ≤ σ2 ∀x.
A.3 Lipschitzian Gradient: We assume f(x) to satisfy Lipschitzian smoothness, i.e., there exists a constant L such that ‖∇f(x)−∇f(y)‖ ≤ L · ‖x− y‖ ∀x, y.
The paper is organized as follows: Section 2 discussed some related work. We follow it up by analyzing impact of mini-batch size on SGD in Section 3. Later we extend the analysis to ASGD in Section 4. We provide experimental evidence regarding our analysis in Section 5 and conclude by discussing future directions in Section 6.
2 RELATED WORK
In recent years, there have been several works analyzing convergence properties of SGD and its variants. In Nemirovski et al. (2009), SGD has been shown to have a convergence rate of O(1/S) for convex objective functions, where S is the number of samples seen, and this rate is in terms of distance of objective function from the optimal value. When the objective cost functions are non-convex, as is the case with most deep learning applications, the rate of convergence of SGD in terms of average gradient norm has been shown to be O(1/ √ S) asymptotically by Ghadimi & Lan (2013). The results in Dekel et al. (2012) can be interpreted as showing the applicability of this convergence rate also for the mini-batches of size in M , where now S takes form of MK with K being the number of mini-batches processed.
Among the distributed variants of SGD, ASGD has been the most popular variant Dean et al. (2012); Li et al. (2014a); Chilimbi et al. (2014). Practically it has been observed that ASGD is often slower to converge with increasing number of learners [Seide et al. (2014); Chan & Lane (2014); Dean et al.
(2012)]. Although these works did not ignore communication overhead, in Section 4 we investigate the inherent inefficiency in ASGD without communication overhead costs. In Lian et al. (2015), it was proved that in the asymptotic regime the convergence rate of O(1/ √ S) can be extended to ASGD when the “age” of updates was bounded by the number of learners.
There exist several other sequential and distributed variants of SGD. For example, SGD with variance reduction techniques to mitigate the effects of gradient variance have been discussed in [Johnson & Zhang (2013); Xiao & Zhang (2014)]. SGD with co-ordinate descent/ascent and its distributed variants has been studied in [Hsieh et al. (2008); Richtárik & Takác (2013); Fercoq et al. (2014); Konečnỳ et al. (2014); Qu & Richtárik (2014); Liu et al. (2015); Jaggi et al. (2014); Nesterov (2012)]. The convergence properties of asynchronous stochastic coordinate descent have been analyzed in Liu & Wright (2015). More recently, the authors in Meng et al. (2016) have studied combining variance reduction techniques, randomized block co-ordinate descent [Richtárik & Takáč (2014)] and Nesterov acceleration methods [Nesterov (2013)] and analyzed its theoretical properties.
There have been several recent works which attempt to mitigate the effect of degrading convergence for large mini-batches (e.g., see [Li et al. (2014b)] where in each mini-batch a regularized objective function is optimized to compute updated parameter vector), or there are works which attempt to select mini-batch dynamically for better performance, for example see [Byrd et al. (2012); Tan et al. (2016); De et al. (2016)], or there have been works which attempt to improve SGD performance by intelligent selection of training samples, e.g., [Needell et al. (2014); Bouchard et al. (2015)].
3 THE IMPACT OF MINIBATCH SIZE ON SGD
In this section, we build on the SGD convergence analysis in Ghadimi & Lan (2013). In particular, we consider the convergence guarantees from Theorem 2.1 in Ghadimi & Lan (2013) restricted to constant learning rate. However, we first modify their analysis to allow mini-batches of arbitrary sizes. Building on, we show that SGD with smaller mini-batches can have better convergence guarantees than with using larger mini-batches. As a consequence, we observe that for larger mini-batches, a larger number of samples is needed for the convergence rate to fall below a certain threshold.
Lemma 1. With mini-batches of size M and a constant learning rate γ, after K iterations of SGD, we have ∑K
k=1 E(‖∇f(xk)‖ 2 )
K ≤
1 γ − LMγ22 · ( Df S + Lσ2γ2 2 ) (2)
Proof outline. Using the property of Lipschitzian gradient, we get
f(xk+1) ≤ f(xk) + 〈∇f(xk), xk+1 − xk〉+ L
2 ‖xk+1 − xk‖2
= f(xk)− γk · 〈 ∇f(xk),
M∑ i=1 G(xk, z i k)
〉 + γ2kL
2 ∥∥∥∥∥ M∑ i=1 G(xk, z i k) ∥∥∥∥∥ 2
(3)
Let us define δik = G(xk, z i k) − ∇f(xk) and δk = ∑M i=1 δ i k. Using it, the above equation can be rewritten as
f(xk+1)− f(xk) ≤ −γk · 〈∇f(xk), δk +M · ∇f(xk)〉+ γ2kL
2 ‖δk +M · ∇f(xk)‖2
⇒ f(xk+1)− f(xk) ≤ −γk · 〈∇f(xk), δk〉 −Mγk · ‖∇f(xk)‖2
+ γ2kL
2
( ‖δk‖2 + 2M〈δk,∇f(xk)〉+M2‖∇f(xk)‖2 ) Rest of the proof involves adding such inequalities over first K − 1 updates and bounding the ‖δk‖2 using assumption A.2 and 〈δk,∇f(xk)〉 using assumption A.1 from Section 1. Finally, we use f(xK)− f(x∗) ≤ Df and rearrange the terms to get the desired bound.
In the following theorem, we justify that given a fixed number of samples S to be processed, SGD with smaller mini-batches SGD can have better convergence guarantees than with using larger minibatches. Note that α used in the theorem statement is a measure of (the square root of) the number of training samples processed.
Theorem 2. Let α := √ Sσ2
LDf . Let 4Ml ≤ α ≤ Mh/4. Then the convergence guarantee for SGD
for mini-batch size Mh after processing S training samples can be worse than the convergence guarantee for mini-batch size Ml by a factor of 2Mh/(α √ 2 +Ml).
Proof outline. For a fixed number S of training samples to be processed, we now minimize the right hand side of Equation 2 to find the best convergence guarantee supported by Lemma 1. Let γ = c · √ Df/(SLσ2), where c is a scalar multiplier. Substituting it in Equation 2 and after some algebraic manipulations, we get∑K k=1 E(‖∇f(xk)‖ 2 )
K ≤
( 1 c + c 2
1− cM2α
) · √ DfLσ 2
S (4)
By applying simple calculus, it can be shown that the value c∗ of c which minimizes right hand side of the above equation is given by
c∗ = M
α ·
( −1 + √ 1 + 2α2
M2
) (5)
In appendix, we show that with M = Ml ≤ α/4, we have c∗ ≈ √ 2 −Ml/α and consequently the coefficient of √ DfLσ2/S from Equation 4 evaluates to approximately √ 2 +Ml/α. Whereas
for M = Mh ≥ 4α, we show that in the appendix using that c∗ = α/Mh and the coefficient of√ DfLσ2/S from Equation 4 evaluates to approximately 2Mh/α. Combining these observations, we get that the convergence guarantees for Ml and Mh can differ by a factor of 2Mh/(α √ 2 +Ml) after processing S training samples. See the appendix for complete proof.
Note that while it could be possible to show that in general smaller mini-batches converge faster than larger mini-batches in theory, we used 4Ml ≤ α ≤ Mh/4 in Theorem 2 for the sake of simplicity. Also, although we theoretically justified faster convergence smaller mini-batches in Theorem 2, the exact factors by which bigger mini-batches can be worse can vary in practice. Here our purpose is to give theoretical support to the practical observations of smaller mini-batches being faster. Theorem 3. The number of samples needs to be processed in order for SGD achieve the same convergence guarantee increases as the mini-batch size increases.
Proof. For the same values of γ and S, the value of the bound from Equation 2 becomes worse (i.e., increases) as M increases. This is because for a fixed γ the quantity γ − LMγ2/2 must decrease as M increases. Consequently for given S, the best value of convergence guarantee (i.e., smallest average gradient norm) attained by varying γ must become worse (i.e., higher) asM increases. Thus in order to reach the same convergence guarantee, SGD must process more training samples with increasing mini-batch size.
Now we will proceed to the next section, where we will show that with ASGD (which is one of the most popular distributed variants of SGD) increasing number of learners can lead to slower convergence given a fixed total number of samples to be processed, and the effect is similar to that of increasing mini-batch size for SGD.
4 ASYNCHRONOUS STOCHASTIC GRADIENT DESCENT
ASGD typically has a parameter server maintaining the parameters (i.e., the weights in the neural network) and multiple learners. Each learner asynchronously repeats the following:
• Pull: Get the parameters from the server.
• Compute: Compute the gradient with respect to randomly selected mini-batch (i.e., a certain number of samples from the dataset). • Push and update: Communicate the gradient to the server. Server then updates the pa-
rameters by subtracting this newly communicated gradient multiplied by the learning rate.
We assume that the update performed by the server is atomic, i.e., the server does not send or receive parameters while it updates the parameters. Now we express kth update step of the ASGD algorithm in terms of our notation. Note that for kth update, the partial gradient computed by a learner can be with respect to an older parameter vector. This is because while computing the partial gradient, the parameter vector could have been updated because of the partial gradients sent in by other learners. Let xk−τ be the parameter vector used by a learner to compute the partial gradient to be used in kth update. Then the equation for kth update of ASGD becomes:
xk+1 = xk − γ ·G(xk−τ , z) (6)
Lian et al. (2015) showed that when the age of the updates τ is bounded by the number of learners N , then ASGD asymptotically converges with a rate ofO(1/ √ S) where S is the cumulative number of training samples processed. From Theorem 1 in Lian et al. (2015), the convergence rate guarantee (expressed in the terms of average gradient norm) for ASGD with N learners after processing K updates becomes ∑K
k=1 E(‖∇f(xk)‖ 2 )
K ≤ 2Df MKγ + σ2Lγ + 2σ2L2MNγ2 (7)
s.t. LMγ + 2L2M2N2γ2 ≤ 1 (8)
The terms independent of the number of updates K in Equation 7 indicate that with a constant learning rate, there is a limit on how close the algorithm can reach to the optimum without lowering the learning rate. Although asymptotically, it can be shown that Equation 7-8 lead to O(1/ √ S) convergence (see Lian et al. (2015)), we now investigate the convergence behavior prior to such a regime. We have the following theorem about the effect of increasing the number of learners on ASGD convergence guarantee:
Theorem 4. Let N > 1 be the number of learners and let α = √ Kσ2
MLDf ≤ N , then the optimal
ASGD convergence rate guarantee for 1 learner and N learners can differ by a factor of approximately Nα .
The proof of above theorem is in the same spirit as that of Theorem 2 and can be found in the appendix. Note that without asynchronous nature of ASGD, the analysis for synchronous distributed SGD would be the same as the analysis for SGD from Section 3. This is because synchronous SGD, where each of the N learners compute the gradient for a random mini-batch of size M , equivalently represents SGD with mini-batch size MN . The asynchronous nature of ASGD introduces extra factors to be taken into account such as the “age” of the updates (i.e., the situation where the gradient returned by a learner may have been computed by an older parameter vector). Theorem 5. For a constant mini-batch size M , the number of samples needs to be processed in order to achieve the same convergence guarantee increases as the number of learners increases.
Proof outline. The range of γ permissible by Equation 8 becomes smaller as N increases. Rest of the proof combines the observations that the minimum attained by the convergence guarantee by Equation 7 must become worse if the range of γ decreases and N increases. For complete proof, please see the appendix.
Discussion: From Theorem 2, for sequential SGD, there could be a difference of 2Mh/(α √ 2+Ml) between the convergence guarantee of mini-batch sizesMl, Mh with 4Ml ≤ α ≤Mh/4. Assuming Ml to be far smaller than α, this factor becomes approximately √ 2Mh/α. This is comparable with the difference N/α between ASGD convergence guarantee of 1 learner and N learners. Although exact numerical multipliers may differ depending on the tightness of the original convergence bound, it points to the similarities between slow-down caused by bigger mini-batch sizes with SGD and larger number of learners with ASGD. At a high level, ASGD with higher number of learners (with
bounded age/staleness of updates) can be thought of SGD with some effective mini-batch size. This effective mini-batch size could be dependent on the number of learners as well as the age/staleness of updates.
5 EXPERIMENTS
Experiment setup: We carry our experiments with CIFAR-10 dataset [cif (accessed January 11, 2016a)] which contains 50, 000 training samples and 10, 000 test samples, each associated with 1 out of 10 possible labels For the CIFAR-10 dataset, our aim is to predict the correct label the input images. For our experiments, we train convolutional neural network shown in Table 1, which is taken from [cif (accessed January 11, 2016b)]. It is a fairly standard convolutional network design consisting of a series of convolutional layers interspersed by max-pooling layers. The convolutional layers outputs are filtered with rectified linear unit before applying max-pooling. Additionally, it also uses Dropout layers which act as regularization [Srivastava et al. (2014)]. At the end, it has a fully connected layer with 10 outputs (equal to the number of labels). The number of parameters to be learned for CIFAR-10 network is ≈ 0.5 million. We use cross-entropy between the input labels and the predicted labels in the final output layer, i.e., F (x, z) (see Section 1 for the notation) is the cross-entropy error and f(x) is the average crossentropy error over all training samples. The implementation of neural network was done using Torch. The target platform for our experiments is a Magma system with 16 GPUs connected to an IBM Power8 host. In our ASGD Downpour implementation, the learners are run on the GPUs, while the parameter server is run on the CPU.
We carried our experiments for 100 epochs, here by an epoch we mean a complete pass of the training data. We chose the learning rate to be 0.01 as it was seen to be performing well at the end
of our experiments with respect to test accuracy (i.e., classification accuracy on the test data). For ASGD part of experiments, we randomly partitioned the training data between all N learners in the beginning of each epoch. At the end of each epoch we measured the test accuracy. See Figure 1 for the results of our SGD experiments. We can observe that as the mini-batch size increases, the test error converges slower with respect to the number of epochs. See Figure 2 for the result of our experiments with ASGD experiments. Again, we can observe that as the number of learners increases, the convergence of test error becomes slower.
These observations agree with our justifications from Section 3 and 4. Moreover, they show that there are similarities between the slow-down caused by increasing mini-batch size with SGD and increasing the number of learners with ASGD. Thus, exploiting parallelism, either by increasing mini-batch size or by increasing the number of learners introduces an inherent inefficiency in the convergence behavior, even after disregarding other overheads such as communication time when we increase the number of learners.
6 CONCLUSION AND FUTURE DIRECTIONS
In this paper, we theoretically justified faster convergence (in terms of average gradient norm attained after processing a fixed number of samples) of SGD with small mini-batches or that of ASGD with smaller number of learners. This indicates that there is an inherent inefficiency in the speed-up obtained with parallelizing gradient descent methods by taking advantage of hardware. It would be interesting to see if such a conclusion holds for more advanced update methods than vanilla SGD, for example methods using momentum and its variants.
A APPENDIX
Lemma 1. With mini-batches of size M and a constant learning rate γ, after K iterations of SGD, we have ∑K
k=1 E(‖∇f(xk)‖ 2 )
K ≤
1 γ − LMγ22 · ( Df S + Lσ2γ2 2 ) (2)
Proof. Using the property of Lipschitzian gradient, we get
f(xk+1) ≤ f(xk) + 〈∇f(xk), xk+1 − xk〉+ L
2 ‖xk+1 − xk‖2
= f(xk)− γk · 〈 ∇f(xk),
M∑ i=1 G(xk, z i k)
〉 + γ2kL
2 ∥∥∥∥∥ M∑ i=1 G(xk, z i k) ∥∥∥∥∥ 2
(9)
Let us define δik = G(xk, z i k) − ∇f(xk) and δk = ∑M i=1 δ i k. Using it, the above equation can be rewritten as
f(xk+1)− f(xk) ≤ −γk · 〈∇f(xk), δk +M · ∇f(xk)〉+ γ2kL
2 ‖δk +M · ∇f(xk)‖2
⇒ f(xk+1)− f(xk) ≤ −γk · 〈∇f(xk), δk〉 −Mγk · ‖∇f(xk)‖2
+ γ2kL
2
( ‖δk‖2 + 2M〈δk,∇f(xk)〉+M2‖∇f(xk)‖2 ) Generating such inequalities over K mini-batches and adding them, we get
−Df ≤ f(xK+1)− f(x1)
≤ K∑ k=1 ( −γk · 〈∇f(xk), δk〉 −Mγk‖∇f(xk)‖2
+ γ2kL
2
( ‖δk‖2 + 2M〈δk,∇f(xk)〉+M2‖∇f(xk)‖2 )) Simple rearrangement of terms gives us
K∑ k=1 ( Mγk − γ2kLM 2 2 ) · ‖∇f(xk)‖2
≤ Df + K∑ k=1 ( −γk · 〈∇f(xk), δk〉+ γ2kL 2 ( ‖δk‖2 + 2M · 〈δk,∇f(xk)〉 )) (10)
Now we observe that E(‖δk‖2) = E(‖ ∑M i=1 δ i k‖ 2 ) ≤ M · ∑M i=1 E(‖δik‖ 2 ) ≤ Mσ2, using Assumption A.2. From the assumption of the stochastic gradient being unbiased estimator of the true gradient (Assumption A.1), we also know that E(〈δk,∇f(xk)〉) = ∑M i=1 E( 〈 δik,∇f(xk)
〉 ) =〈∑M
i=1 E(δik),∇f(xk) 〉 = 0. Taking the expectation of both sides in Equation 10 with respect to
randomly selected samples and using these observations we get the following equation: K∑ k=1 ( Mγk − LM2γ2k 2 ) · E(‖∇f(xk)‖2) ≤ Df + LMσ2 2 K∑ k=1 γ2k
The above equation is equivalent to Equation 2.11 from Ghadimi & Lan (2013) modified to allow arbitrary mini-batch sizes. Restricting ourselves to allow constant learning rate and after simplifying the above equation, we get∑K
k=1 E(‖∇f(xk)‖ 2 )
K ≤
1 γ − LMγ22 · ( Df S + Lσ2γ2 2 ) (11)
Theorem 2. Let α := √ Sσ2
LDf . Let 4Ml ≤ α ≤ Mh/4. Then the convergence guarantee for SGD
for mini-batch size Mh after processing S training samples can be worse than the convergence guarantee for mini-batch size Ml by a factor of 2Mh/(α √ 2 +Ml).
Proof outline. For a fixed number S of training samples to be processed, we now minimize the right hand side of Equation 2 to find the best convergence guarantee supported by Lemma 1. Let γ = c · √ Df/(SLσ2), where c is a scalar multiplier. Substituting it in Equation 2 and after some algebraic manipulations, we get∑K k=1 E(‖∇f(xk)‖ 2 )
K ≤
( 1 c + c 2
1− cM2α
) · √ DfLσ 2
S (12)
By applying simple calculus, it can be shown that the value c∗ of c which minimizes right hand side of the above equation is given by
c∗ = M
α ·
( −1 + √ 1 + 2α2
M2
) (13)
• Consider the case ofM =Ml ≤ α/4. Denote c∗ by c∗l for this case. WithM =Ml ≤ α/4, we have √ 1 + 2α2/M2l ≈ √ 2α/Ml. Thus we get c∗l ≈ √ 2−Ml/α. Further, we can write
the following using Ml is little compared to α and other simple known approximations: 1
1− c ∗ lMl 2α
= 1
1− ( √ 2− Mlα ) Ml 2α ≈ 1 1− Ml√
2α
≈ 1 + Ml√ 2α
(14)
1 c∗l = 1√ 2− Mlα = 1√ 2 · 1 1− Ml√
2α
≈ 1√ 2 · ( 1 + Ml√ 2α ) ≈ 1√ 2 + Ml 2α (15)
c∗l 2 = 1√ 2 − Ml 2α (16)
Using Equation 14, 15, 16, we get that the coefficient of √ DfLσ2/S from Equation 12 evaluates to approximately √ 2 +Ml/α. • Consider Mh ≥ 4α. Denote c∗ by c∗h for this case. With M = Ml ≤ α/4, we have√ 1 + 2α2/M2h ≈ 1+α2/M2h , using the approximation √ 1 + ≈ 1+ /2 for small . Thus
c∗h ≈ α/Mh. Since α is much smaller thanMh, we can approximate 1/c∗h+c∗h/2 ≈ 1/c∗h ≈ Mh/α and we also have 1− c∗hM/(2α) = 1/2. Thus the coefficient of √ DfLσ2/S from Equation 12 evaluates to approximately 2Mh/α.
Combining the above observations, we get that the convergence guarantees forMl andMh can differ by a factor of 2Mh/(α √ 2 +Ml) after processing S training samples.
Theorem 4. Let N > 1 be the number of learners and let α = √ Kσ2
MLDf ≤ N , then the optimal
ASGD convergence rate guarantee for 1 learner and N learners can differ by a factor of approximately Nα .
Proof. We have γ = c · √ Df/(MKLσ2) = c αML from the definition of α. Substituting this in Equation 7, we get∑K k=1 E(‖∇f(xk)‖ 2 )
K ≤
2Df MKc · √
Df MKLσ2
+ σ2L · c · √
Df MKLσ2
+ 2σ2L2MN · c αML
· c · √
Df MKLσ2
=
( 2
c + c+
2Nc2
α
) · √ DfLσ2
MK (17)
From the definition of α, we have K = α2MLDf/σ2. Using it in the above equation, we get∑K k=1 E(‖∇f(xk)‖ 2 )
K ≤
{( 2
c + c+
2Nc2
α ) · 1 α } · σ 2 M (18)
Similarly, given γ = c · √
Df MKLσ2 = c αML , the condition in Equation 8 can also be expressed as:
c α +
2N2c2
α2 ≤ 1 ⇒ 2N2c2 + αc− α2 ≤ 0
Since learning rate (and hence c) is always positive, the above equation gives us
0 ≤ c ≤ α 4N2
· (−1 + √ 1 + 8N2)
Thus finding the optimal learning rate (within the regime of Equation 7 and 8) is equivalent to solving the following:
minimize ( 2
c + c+
2Nc2
α ) · 1 α · σ 2 M (19)
s.t. 0 ≤ c ≤ α 4N2
· (−1 + √ 1 + 8N2) (20)
Now, by means of solving the above optimization, we will investigate how much the convergence guarantee can differ as the number of learners increase. In particular, we will look at the difference in the guarantee for 1 learner and N ′ learners where N ′ ≥ 16. Taking the derivative of Equation 19 with respect to c and setting it to 0, we get the following:
4Nc3 + αc2 − 2α = 0 (21)
Let c∗1 and c ∗ N ′ denote the solutions to the above equation for 1 and N ′ learners respectively. Notice that for N = 1 and for α ≥ 16, the square term dominates in Equation 21 and c∗1 ≈ √ 2 (which also satisfies the constraint in Equation 20). Thus we get
For N = 1: ∑K k=1 E(‖∇f(xk)‖ 2 )
K . 2 √ 2 · σ 2 αM (22)
However, for N = N ′ and 16 ≤ α ≤ N ′, the cubic term dominates in Equation 21. Thus the value of c which satisfies Equation 21 is approximately 3 √ α/(2N ′). However, the upper bound in Equation 20 for large N ′ becomes c . α/( √ 2N ′). For the range of α under consideration, this
is smaller than 3 √ α/(2N ′), thus we get c∗N ′ = α/( √ 2N ′). Thus for 16 ≤ α ≤ N ′, Equation 18 becomes
For N = N ′: ∑K k=1 E(‖∇f(xk)‖ 2 )
K .
2 √ 2N ′
α · σ
2
αM (23)
Thus comparing Equation 22 and 23, we see that the ASGD convergence guarantee for N = 1 and N = N ′ learners can differ by a factor of N ′
α for 16 ≤ α ≤ N ′.

This paper shows that when a larger mini-batch is used (in the serial setting), the number of samples needed to be processed for the same convergence guarantee is larger. A similar behavior is discussed for using multiple learners in asynchronous SGD. This behavior has been known in convex optimization (e.g., "Better Mini-Batch Algorithms via Accelerated Gradient Methods", NIPS 2011). There, the convergence is of the form O(1/\sqrt{bT}+1/T), and so using bT samples only lead to \sqrt{b} time improvement. This paper extends a similar result to the nonconvex case (but the underlying mathematics is mainly from [Ghadimi & Lan, 2013]). However, this behavior is also known and indeed has been summarized in the deep learning textbook (chapter 8). Hence, the novelty is limited.
The theoretical results in this paper suggest that it is best not to use mini-batch. However, using mini-batch is often beneficial in practice. As discussed in the deep learning textbook, using mini-batch allows using a larger learning rate (note that this paper assumes the same learning rate for both mini-batch sizes). Moreover, multicore architectures allows parallel execution of mini-batch almost for free. Hence, the practical significance of the results is also limited.
Other comments:
- Equation (4): since the same number of samples (S) is processed and S=MK, where M is the mini-batch size and K is the number of mini-batches processed (as mentioned in the first paragraph of section 2), when two different mini-batch sizes are considered (M_l and M_h), their K's should differ. However, the same K is used on the LHS of (4).
- Figures 1 and 2: As convergence speed is of main interest, why not show the training objective instead?


Though deep neural networks have shown great success in the large data domain, they generally perform poorly on few-shot learning tasks, where a classifier has to quickly generalize after seeing very few examples from each class. The general belief is that gradient-based optimization in high capacity classifiers requires many iterative steps over many examples to perform well. Here, we propose an LSTMbased meta-learner model to learn the exact optimization algorithm used to train another learner neural network classifier in the few-shot regime. The parametrization of our model allows it to learn appropriate parameter updates specifically for the scenario where a set amount of updates will be made, while also learning a general initialization of the learner (classifier) network that allows for quick convergence of training. We demonstrate that this meta-learning model is competitive with deep metric-learning techniques for few-shot learning.
1 INTRODUCTION
Deep learning has shown great success in a variety of tasks with large amounts of labeled data in image classification (He et al., 2015), machine translation (Wu et al., 2016), and speech modeling (Oord et al., 2016). These achievements have relied on the fact that optimization of these deep, high-capacity models requires many iterative updates across many labeled examples. This type of optimization breaks down in the small data regime where we want to learn from very few labeled examples. In this setting, rather than have one large dataset, we have a set of datasets, each with few annotated examples per class. The motivation for this task lies not only in the fact that humans, even children, can usually generalize after just one example of a given object, but also because models excelling at this task would have many useful applications. Firstly, they would help alleviate data collection as we would not require millions of labeled examples to attain reasonable performance. Furthermore, in many fields, data exhibits the characteristic of having many different classes but few examples per class. Models that are able to generalize from few examples would be able to capture this type of data effectively.
There seem to be two main reasons why gradient-based optimization fails in the face of few labeled examples. Firstly, the variants of gradient-based optimization algorithms, such as momentum (Nesterov, 1983), Adagrad (Duchi et al., 2011), Adadelta (Zeiler, 2012), and ADAM (Kingma & Ba, 2014), weren’t designed specifically to perform well under the constraint of a set number of updates. Specifically when applied to non-convex optimization problems, with a reasonable choice of hyperparameters these algorithms don’t have very strong guarantees of speed of convergence, beyond that they will eventually converge to a good solution after what could be many millions of iterations. Secondly, for each separate dataset considered, the network would have to start from a random initialization of its parameters, which considerably hurts its ability to converge to a good solution after a few updates. Transfer learning (Caruana, 1995; Bengio et al., 2012; Donahue et al., 2013) can be applied to alleviate this problem by fine-tuning a pre-trained network from another task which has more labelled data; however, it has been observed that the benefit of a pre-trained network greatly decreases as the task the network was trained on diverges from the target task (Yosinski et al., 2014). What is needed is a systematic way to learn a beneficial common initialization that would
∗Work done as an intern at Twitter. Sachin is a PhD student at Princeton University and can be reached at sachinr@princeton.edu.
serve as a good point to start training for the set of datasets being considered. This would provide the same benefits as transfer learning, but with the guarantee that the initialization is an optimal starting point for fine-tuning.
Previous work has suggested one manner in which to acquire quick knowledge from few examples, through the idea of meta-learning (Thrun, 1998; Schmidhuber et al., 1997). Meta-learning suggests framing the learning problem at two levels. The first is quick acquisition of knowledge within each separate task presented. This process is guided by the second, which involves slower extraction of information learned across all the tasks.
We present a method here that addresses the weakness of neutral networks trained with gradientbased optimization on the few-shot learning problem by framing the problem within a meta-learning setting. We propose an LSTM-based meta-learner optimizer that is trained to optimize a learner neural network classifier. The meta-learner captures both short-term knowledge within a task and long-term knowledge common among all the tasks. By using an objective that directly captures an optimization algorithm’s ability to have good generalization performance given only a set number of updates, the meta-learner model is trained to converge a learner classifier to a good solution quickly on each task. Additionally, the formulation of our meta-learner model allows it to learn a task-common initialization for the learner classifier, which captures fundamental knowledge shared among all the tasks.
2 TASK DESCRIPTION
We first begin by detailing the meta-learning formulation we use. In the typical machine learning setting, we are interested in a dataset D and usually split D so that we optimize parameters θ on a training set Dtrain and evaluate its generalization on the test set Dtest. In meta-learning, however, we are dealing with meta-sets D containing multiple regular datasets, where each D ∈ D has a split of Dtrain and Dtest.
We consider the k-shot, N -class classification task, where for each dataset D, the training set consists of k labelled examples for each of N classes, meaning that Dtrain consists of k ·N examples, and Dtest has a set number of examples for evaluation. We note that previous work (Vinyals et al., 2016) has used the term episode to describe each dataset consisting of a training and test set.
In meta-learning, we thus have different meta-sets for meta-training, meta-validation, and metatesting (Dmeta−train, Dmeta−validation, and Dmeta−test, respectively). On Dmeta−train, we are interested in training a learning procedure (the meta-learner) that can take as input one of its training setsDtrain and produce a classifier (the learner) that achieves high average classification performance on its corresponding test setDtest. Using Dmeta−validation we can perform hyper-parameter selection of the meta-learner and evaluate its generalization performance on Dmeta−test.
For this formulation to correspond to the few-shot learning setting, each training set in datasets D ∈ D will contain few labeled examples (we consider k = 1 or k = 5), that must be used to generalize to good performance on the corresponding test set. An example of this formulation is given in Figure 1.
3 MODEL
We now move to the description of our proposed model for meta-learning.
3.1 MODEL DESCRIPTION
Consider a single dataset, or episode, D ∈ Dmeta−train. Suppose we have a learner neural net classifier with parameters θ that we want to train on Dtrain. The standard optimization algorithms used to train deep neural networks are some variant of gradient descent, which uses updates of the form
θt = θt−1 − αt∇θt−1Lt, (1)
where θt−1 are the parameters of the learner after t − 1 updates, αt is the learning rate at time t, Lt is the loss optimized by the learner for its tth update, ∇θt−1Lt is the gradient of that loss with respect to parameters θt−1, and θt is the updated parameters of the learner.
Our key observation that we leverage here is that this update resembles the update for the cell state in an LSTM (Hochreiter & Schmidhuber, 1997)
ct = ft ct−1 + it c̃t, (2) if ft = 1, ct−1 = θt−1, it = αt, and c̃t = −∇θt−1Lt. Thus, we propose training a meta-learner LSTM to learn an update rule for training a neural network. We set the cell state of the LSTM to be the parameters of the learner, or ct = θt, and the candidate cell state c̃t = ∇θt−1Lt, given how valuable information about the gradient is for optimization. We define parametric forms for it and ft so that the meta-learner can determine optimal values through the course of the updates.
Let us start with it, which corresponds to the learning rate for the updates. We let it = σ ( WI · [ ∇θt−1Lt,Lt, θt−1, it−1 ] + bI ) ,
meaning that the learning rate is a function of the current parameter value θt−1, the current gradient ∇θt−1Lt, the current loss Lt, and the previous learning rate it−1. With this information, the metalearner should be able to finely control the learning rate so as to train the learner quickly while avoiding divergence.
As for ft, it seems possible that the optimal choice isn’t the constant 1. Intuitively, what would justify shrinking the parameters of the learner and forgetting part of its previous value would be if the learner is currently in a bad local optima and needs a large change to escape. This would correspond to a situation where the loss is high but the gradient is close to zero. Thus, one proposal for the forget gate is to have it be a function of that information, as well as the previous value of the forget gate:
ft = σ ( WF · [ ∇θt−1Lt,Lt, θt−1, ft−1 ] + bF ) .
Additionally, notice that we can also learn the initial value of the cell state c0 for the LSTM, treating it as a parameter of the meta-learner. This corresponds to the initial weights of the classifier (that
the meta-learner is training). Learning this initial value lets the meta-learner determine the optimal initial weights of the learner so that training begins from a beneficial starting point that allows optimization to proceed rapidly. Lastly, note that though the meta-learner’s update rule matches the cell state update of the LSTM, the meta-learner also bears similarity to the GRU (Cho et al., 2014) hidden state update, with the exception that the forget and input gates aren’t tied to sum to one.
3.2 PARAMETER SHARING & PREPROCESSING
Because we want our meta-learner to produce updates for deep neural networks, which consist of tens of thousands of parameters, to prevent an explosion of meta-learner parameters we need to employ some sort of parameter sharing. Thus as in Andrychowicz et al. (2016), we share parameters across the coordinates of the learner gradient. This means each coordinate has its own hidden and cell state values but the LSTM parameters are the same across all coordinates. This allows us to use a compact LSTM model and additionally has the nice property that the same update rule is used for each coordinate, but one that is dependent on the respective history of each coordinate during optimization. We can easily implement parameter sharing by having the input be a batch of gradient coordinates and loss inputs (∇θt,iLt,Lt) for each dimension i. Because the different coordinates of the gradients and the losses can be of very different magnitudes, we need to be careful in normalizing the values so that the meta-learner is able to use them properly during training. Thus, we also found that the preprocessing method of Andrychowicz et al. (2016) worked well when applied to both the dimensions of the gradients and the losses at each time step:
x→
{( log(|x|) p , sgn(x) ) if |x| ≥ e−p
(−1, epx) otherwise
This preprocessing adjusts the scaling of gradients and losses, while also separating the information about their magnitude and their sign (the latter being mostly useful for gradients). We found that the suggested value of p = 10 in the above formula worked well in our experiments.
3.3 TRAINING
The question now is how do we train the LSTM meta-learner model to be effective at few-shot learning tasks? As observed in Vinyals et al. (2016), in order to perform well at this task, it is key to have training conditions match those of test time. During evaluation of the meta-learning, for each dataset (episode), D = (Dtrain, Dtest) ∈ Dmeta−test, a good meta-learner model will, given a series of learner gradients and losses on the training set Dtrain, suggest a series of updates for the classifier that pushes it towards good performance on the test set Dtest.
Thus to match test time conditions, when considering each dataset D ∈ Dmeta−train, the training objective we use is the loss Ltest of the produced classifier on D’s test set Dtest. While iterating over the examples in D’s training set Dtrain, at each time step t the LSTM meta-learner receives (∇θt−1Lt,Lt) from the learner (the classifier) and proposes the new set of parameters θt. The process repeats for T steps, after which the classifier and its final parameters are evaluated on the test set to produce the loss that is then used to train the meta-learner. The training algorithm is described in Algorithm 1 and the corresponding computational graph is shown in Figure 2.
3.3.1 GRADIENT INDEPENDENCE ASSUMPTION
Notice that our formulation would imply that the losses Lt and gradients∇θt−1Lt of the learner are dependent on the parameters of the meta-learner. Gradients on the meta-learner’s parameters should normally take this dependency into account. However, as discussed by Andrychowicz et al. (2016), this complicates the computation of the meta-learner’s gradients. Thus, following Andrychowicz et al. (2016), we make the simplifying assumption that these contributions to the gradients aren’t important and can be ignored, which allows us to avoid taking second derivatives, a considerably expensive operation. We were still able to train the meta-learner effectively in spite of this simplifying assumption.
3.3.2 INITIALIZATION OF META-LEARNER LSTM
When training LSTMs, it is advised to initialize the LSTM with small random weights and to set the forget gate bias to a large value so that the forget gate is initialized to be close to 1, thus enabling gradient flow (Zaremba, 2015). In addition to the forget gate bias setting, we found that we needed to initialize the input gate bias to be small so that the input gate value (and thus the learning rate) used by the meta-learner LSTM starts out being small. With this combined initialization, the metalearner starts close to normal gradient descent with a small learning rate, which helps initial stability of training.
3.4 BATCH NORMALIZATION
Batch Normalization (Ioffe & Szegedy, 2015) is a recently proposed method to stabilize and thus speed up learning of deep neural networks by reducing internal covariate shift within the learner’s hidden layers. This reduction is achieved by normalizing each layer’s pre-activation, by subtracting by the mean and dividing by the standard deviation. During training, the mean and standard deviation are estimated using the current batch being trained on, whereas during evaluation a running average of both statistics calculated on the training set is used. We need to be careful with batch normalization for the learner network in the meta-learning setting, because we do not want to collect mean and standard deviation statistics during meta-testing in a way that allows information to leak between different datasets (episodes), being considered. One easy way to prevent this issue is to not collect statistics at all during the meta-testing phase, but just use our running averages from metatraining. This, however, has a bad impact on performance, because we have changed meta-training and meta-testing conditions, causing the meta-learner to learn a method of optimization that relies on batch statistics which it now does not have at meta-testing time. In order to keep the two phases as similar as possible, we found that a better strategy was to collect statistics for each datasetD ∈ D during Dmeta−test, but then erase the running statistics when we consider the next dataset. Thus, during meta-training, we use batch statistics for both the training and testing set whereas during meta-testing, we use batch statistics for the training set (and to compute our running averages) but then use the running averages during testing. This does not cause any information to leak between different datasets, but also allows the meta-learner to be trained on conditions that are matched between training and testing. Lastly, because we are doing very few training steps, we computed the running averages so that higher preference is given to the later values.
Algorithm 1 Train Meta-Learner Input: Meta-training set Dmeta−train, Learner M with parameters θ, Meta-Learner R with parameters Θ.
1: Θ0 ← random initialization 2: 3: for d = 1, n do 4: Dtrain, Dtest ← random dataset from Dmeta−train 5: θ0 ← c0 . Intialize learner parameters 6: 7: for t = 1, T do 8: Xt,Yt ← random batch from Dtrain 9: Lt ← L(M(Xt; θt−1),Yt) . Get loss of learner on train batch 10: ct ← R((∇θt−1Lt,Lt); Θd−1) . Get output of meta-learner using Equation 2 11: θt ← ct . Update learner parameters 12: end for 13: 14: X,Y ← Dtest 15: Ltest ← L(M(X; θT ),Y) . Get loss of learner on test batch 16: Update Θd using∇Θd−1Ltest . Update meta-learner parameters 17: 18: end for
4 RELATED WORK
While this work falls within the broad literature of transfer learning in general, we focus here on positioning it relative to previous work on meta-learning and few-shot learning.
4.1 META-LEARNING
Meta-learning has a long history, but has grown to prominence recently as many have advocated for it as a key to achieving human-level intelligence in the future (Lake et al., 2016). The ability to learn at two levels (learning within each task presented, while accumulating knowledge about the similarities and differences between tasks) is seen as being crucial to improving AI. Previous work has used a variety of techniques in the meta-learning setting.
Schmidhuber (1992; 1993) explored using networks that learn how to modify their own weights over a number of computations steps on the input. The updating of the weights is defined in a parametric form that allows the prediction and weight-change process to be differentiable end-to-end. The work of Bengio et al. (1990; 1995) and Bengio (1993) considered learning update rules for neural networks that are biologically plausible. This property is enforced by allowing the parametric form of the update to only have as input local information at each hidden unit to determine the weight change. Different optimization methods, such as genetic programming or simulated annealing, are used to train the learning rule.
In Santoro et al. (2016), a memory-augmented neural network is trained to learn how to store and retrieve memories to use for each classification task. The work of Andrychowicz et al. (2016) uses an LSTM to train a neural network; however, they are interested in learning a general optimization algorithm to train neural networks for large-scale classification, whereas we are interested in the few-shot learning problem. This work also builds upon Hochreiter et al. (2001) and Bosc, both of which used LSTMs to train multi-layer perceptrons to learn on binary classification and timeseries prediction tasks. Another related method is the work of Bertinetto et al. (2016), who train a meta-learner to map a training example to the weights of a neural network that is then used to classify future examples from this class; however, unlike our method the classifier network is directly produced rather than being fine-tuned after multiple training steps. Our work also bears similarity to Maclaurin et al. (2015), who tune the hyperparameters of gradient descent with momentum by backpropagating through the chain of gradient steps to optimize the validation performance.
4.2 FEW-SHOT LEARNING
The best performing methods for few-shot learning have been mainly metric learning methods. Deep siamese networks (Koch, 2015) train a convolutional network to embed examples so that items in the same class are close while items in different classes are far away, according to some distance metric. Matching networks (Vinyals et al., 2016) refine this idea so that training and testing conditions match, by defining a differentiable nearest neighbor loss involving the cosine similarities of embeddings produced by a convolutional network.
5 EVALUATION
In this section, we describe the results of experiments, examining the properties of our model and comparing our method’s performance against different approaches∗. Following Vinyals et al. (2016), we consider the k-shot, N -class classification setting where a meta-learner trains on many related but small training sets of k examples for each of N classes. We first split the list of all classes in the data into disjoint sets and assign them to each meta-set of meta-training, meta-validation, and meta-testing. To generate each instance of a k-shot, N -class task dataset D = (Dtrain, Dtest) ∈ D , we do the following: we first sampleN classes from the list of classes corresponding to the meta-set we consider. We then sample k examples from each of those classes. These k examples together compose the training set Dtrain. Then, an additional fixed amount of the rest of the examples are sampled to yield a test set Dtest. We generally have 15 examples per class in the test sets. When training the meta-learner, we iterate by sampling these datasets (episodes) repeatedly. For metavalidation and meta-testing, however, we produce a fixed number of these datasets to evaluate each method. We produce enough datasets to ensure that the confidence interval of the mean accuracy is small.
For the learner, we use a simple CNN containing 4 convolutional layers, each of which is a 3 × 3 convolution with 32 filters, followed by batch normalization, a ReLU non-linearity, and lastly a 2× 2 max-pooling. The network then has a final linear layer followed by a softmax for the number of classes being considered. The loss function L is the average negative log-probability assigned by the learner to the correct class. For the meta-learner, we use a 2-layer LSTM, where the first layer is a normal LSTM and the second layer is our modified LSTM meta-learner. The gradients and losses are preprocessed and fed into the first layer LSTM, and the regular gradient coordinates are also used by the second layer LSTM to implement the state update rule shown in (1). At each time step, the learner’s loss and gradient is computed on a batch consisting of the entire training set Dtrain, because we consider training sets with only a total of 5 or 25 examples. We train our LSTM with ADAM using a learning rate of 0.001 and with gradient clipping using a value of 0.25.
5.1 EXPERIMENT RESULTS
The Mini-ImageNet dataset was proposed by Vinyals et al. (2016) as a benchmark offering the challenges of the complexity of ImageNet images, without requiring the resources and infrastructure necessary to run on the full ImageNet dataset. Because the exact splits used in Vinyals et al. (2016) were not released, we create our own version of the Mini-Imagenet dataset by selecting a random 100 classes from ImageNet and picking 600 examples of each class. We use 64, 16, and 20 classes for training, validation and testing, respectively. We consider 1-shot and 5-shot classification for 5 classes. We use 15 examples per class for evaluation in each test set. We compare against two baselines and a recent metric-learning technique, Matching Networks (Vinyals et al., 2016), which has achieved state-of-the-art results in few-shot learning. The results are shown in Table 1.
The first baseline we use is a nearest-neighbor baseline (Baseline-nearest-neighbor), where we first train a network to classify between all the classes jointly in the original meta-training set. At metatest time, for each dataset D, we embed all the items in the training set using our trained network and then use nearest-neighbor matching among the embedded training examples to classify each test example. The second baseline we use (Baseline-finetune) represents a coarser version of our metalearner model. As in the first baseline, we start by training a network to classify jointly between all classes in the meta-training set. We then use the meta-validation set to search over SGD hyperparameters, where each training set is used to fine-tune the pre-trained network before evaluating on
∗Code can be found at https://github.com/twitter/meta-learning-lstm.
the test set. We use a fixed number of updates for fine tuning and search over the learning rate and learning rate decay used during the course of these updates.
For Matching Networks, we implemented our own version of both the basic and the fully-conditional embedding (FCE) versions. In the basic version, a convolutional network is trained to learn independent embeddings for examples in the training and test set. In the FCE version, a bidirectional-LSTM is used to learn an embedding for the training set such that each training example’s embedding is also a function of all the other training examples. Additionally, an attention-LSTM is used so that a test example embedding is also a function of all the embeddings of the training set. We do not consider fine-tuning the network using the train set during meta-testing to improve performance as mentioned in Vinyals et al. (2016), but do note that our meta-learner could also be fine-tuned using this data. Note that to remain consistent with Vinyals et al. (2016), our baseline and matching net convolutional networks have 4 layers each with 64 filters. We also added dropout to each convolutional block in matching nets to prevent overfitting.
For our meta-learner, we train different models for the 1-shot and 5-shot tasks, that make 12 and 5 updates, respectively. We noticed that better performance for each task was attained if the metalearner is explicitly trained to do the set number of updates during meta-training that will be used during meta-testing.
We attain results that are much better than the baselines discussed and competitive with Matching Networks. For 5-shot, we are able to do much better than Matching Networks, whereas for 1-shot, the confidence interval for our performance intersects the interval for Matching Networks. Again, we note that the numbers do not match the ones provided by Vinyals et al. (2016) simply because we created our version of the dataset and implemented our own versions of their model. It is interesting to note that the fine-tuned baseline is worse than the nearest-neighbor baseline. Because we are not regularizing the classifier, with very few updates the fine-tuning model overfits, especially in the 1-shot case. This propensity to overfit speaks to the benefit of meta-training the initialization of the classifier end-to-end as is done in the meta-learning LSTM.
5.2 VISUALIZATION OF META-LEARNER
We also visualize the optimization strategy learned by the meta-learner, in Figure 3. We can look at the it and ft gate values in Equation 2 at each update step, to try to get an understanding of how the meta-learner updates the learner during training. We visualize the gate values while training on different datasets Dtrain, to observe whether there are variations between training sets. We consider both 1-shot and 5-shot classification settings, where the meta-learner is making 10 and 5 updates, respectively. For the forget gate values for both tasks, the meta-learner seems to adopt a simple weight decay strategy that seems consistent across different layers. The input gate values are harder to interpret to glean the meta-learner’s strategy. However, there seems to a be a lot of variability between different datasets, indicating that the meta-learner isn’t simply learning a fixed optimization strategy. Additionally, there seem to be differences between the two tasks, suggesting that the meta-learner has adopted different methods to deal with the different conditions of each setting.
6 CONCLUSION
We described an LSTM-based model for meta-learning, which is inspired from the parameter updates suggested by gradient descent optimization algorithms. Our LSTM meta-learner uses its state to represent the learning updates of the parameters of a classifier. It is trained to discover both a good initialization for the learner’s parameters, as well as a successful mechanism for updating the learner’s parameters to a given small training set for some new classification task. Our experiments demonstrate that our approach outperforms natural baselines and is competitive to the state-of-theart in metric learning for few-shot learning.
In this work, we focused our study to the few-shot and few-classes setting. However, it would be more valuable to train meta-learners that can perform well across a full spectrum of settings, i.e. for few or lots of training examples and for few or lots of possible classes. Our future work will thus consider moving towards this more challenging scenario.
ACKNOWLEDGMENTS
We thank Jake Snell, Kevin Swersky, and Oriol Vinyals for helpful discussions of this work.

In light of the authors' responsiveness and the updates to the manuscript -- in particular to clarify the meta-learning task -- I am updating my score to an 8.
-----
This manuscript proposes to tackle few-shot learning with neural networks by leveraging meta-learning, a classic idea that has seen a renaissance in the last 12 months. The authors formulate few-shot learning as a sequential meta-learning problem: each "example" includes a sequence of batches of "training" pairs, followed by a final "test" batch. The inputs at each "step" include the outputs of a "base learner" (e.g., training loss and gradients), as well as the base learner's current state (parameters). The paper applies an LSTM to this meta-learning problem, using the inner memory cells in the *second* layer to directly model the updated parameters of the base learner. In doing this, they note similarities between the respective update rules of LSTM memory cells and gradient descent. Updates to the LSTM meta-learner are computed based on the base learner's prediction loss for the final "test" batch. The authors make several simplifying assumptions, such as sharing weights across all second layer cells (analogous to using the same learning rate for all parameters). The paper recreates the Mini-ImageNet data set proposed in Vinyals et al 2016, and shows that the meta-learner LSTM is competitive with the current state-of-the-art (Matchin Networks, Vinyals 2016) on 1- and 5-shot learning.
Strengths:
- It is intriguing -- and in hindsight, natural -- to cast the few-shot learning problem as a sequential (meta-)learning problem. While the authors did not originate the general idea of persisting learning across a series of learning problems, I think it is fair to say that they have advanced the state of the art, though I cannot confidently assert its novelty as I am not deeply familiar with recent work on meta-learning.
- The proposed approach is competitive with and outperforms Vinyals 2016 in 1-shot and 5-shot Mini-ImageNet experiments.
- The base learner in this setting (simple ConvNet classifier) is quite different from the nearest-neighbor-on-top-of-learned-embedding approach used in Vinyals 2016. It is always exciting when state-of-the-art results can be reported using very different approaches, rather than incremental follow-up work.
- As far as I know, the insight about the relationship between the memory cell and gradient descent updates is novel here. It is interesting regardless.
- The paper offers several practical insights about how to design and train an LSTM meta-learner, which should make it easier for others to replicate this work and apply these ideas to new problems. These include proper initialization, weight sharing across coordinates, and the importance of normalizing/rescaling the loss, gradient, and parameter inputs. Some of the insights have been previously described (the importance of simulating test conditions during meta-training; assuming independence between meta-learner and base learner parameters when taking gradients with respect to the meta-learner parameters), but the discussion here is useful nonetheless.
Weaknesses:
- The writing is at times quite opaque. While it describes very interesting work, I would not call the paper an enjoyable read. It took me multiple passes (as well as consulting related work) to understand the general learning problem. The task description in Section 2 (Page 2) is very abstract and uses notation and language that is not common outside of this sub-area. The paper could benefit from a brief concrete example (based on MNIST is fine), perhaps paired with a diagram illustrating a sequence of few-shot learning tasks. This would definitely make it accessible to a wider audience.
- Following up on that note, the precise nature of the N-class, few-shot learning problem here is unclear to me. Specifically, the Mini-ImageNet data set has 100 labels, of which 64/16/20 are used during meta-training/validation/testing. Does this mean that only 64/100 classes are observed through meta-training? Or does it mean that only 64/100 are observed in each batch, but on average all 100 are observed during meta-training? If it's the former, how many outputs does the softmax layer of the ConvNet base learner have during meta-training? 64 (only those observed in training) or 100 (of which 36 are never observed)? Many other details like these are unclear (see question).
- The plots in Figure 2 are pretty uninformative in and of themselves, and the discussion section offers very little insight around them.
This is an interesting paper with convincing results. It seems like a fairly clear accept, but the presentation of the ideas and work therein could be improved. I will definitely raise my score if the writing is improved.