A Resizable Mini-batch Gradient Descent based on a Multi-Armed Bandit

1 INTRODUCTION . Gradient descent ( GD ) is a common optimization algorithm for finding the minimum of the expected loss . It takes iterative steps proportional to the negative gradient of the loss function at each iteration . It is based on the observation that if the multi-variable loss functions f ( w ) is differentiable at point w , then f ( w ) decreases fastest in the direction of the negative gradient of f at w , i.e. , −∇f ( w ) . The model parameters are updated iteratively in GD as follows : wt+1 = wt − ηtgt , gt = ∇wf ( wt ) where wt , gt , and ηt are the model parameters , gradients of f with respect to w , and learning rate at time t respectively . For small enough ηt , f ( wt ) ≥ f ( wt+1 ) and ultimately the sequence of wt will move down toward a local minimum . For a convex loss function , GD is guaranteed to converge to a global minimum with an appropriate learning rate . There are various issues to consider in gradient-based optimization . First , GD can be extremely slow and impractical for large dataset : gradients of all the data have to be evaluated for each iteration . With larger data size , the convergence rate , the computational cost and memory become critical , and special care is required to minimize these factors . Second , for non-convex function which is often encountered in deep learning , GD can get stuck in a local minimum without the hope of escaping . Third , stochastic gradient descent ( SGD ) , which is based on the gradient of a single training sample , has large gradient variance , and it requires a large number of iterations . This ultimately translates to slow convergence . Mini-batch gradient descent ( MGD ) , which is based on the gradient over a small batch of training data , trades off between the robustness of SGD and the stability of GD . There are three advantages for using MGD over GD and SGD : 1 ) The batching allows both the efficiency of memory usage and implementations ; 2 ) The model update frequency is higher than GD which allows for a more robust convergence avoiding local minimum ; 3 ) MGD requires less iteration per epoch and provides a more stable update than SGD . For these reasons , MGD has been a popular algorithm for machine learning . However , selecting an appropriate batch size is difficult . Various studies suggest that there is a close link between performance and batch size used in MGD Breuel ( 2015 ) ; Keskar et al . ( 2016 ) ; Wilson & Martinez ( 2003 ) . There are various guidelines for selecting a batch size but have not been completely practical Bengio ( 2012 ) . Grid search is a popular method but it comes at the expense of search time . There are a small number of adaptive MGD algorithms to replace grid search Byrd et al . ( 2012 ) ; De et al . ( 2016 ) ; Friedlander & Schmidt ( 2012 ) . These algorithms increase the batch size gradually according to their own criterion . However , these algorithms are based on convex loss function and hard to be applied to deep learning . For non-convex optimization , it is difficult to determine the optimal batch size for best performance . This paper considers a resizable mini-batch gradient descent ( RMGD ) algorithm based on a multi-armed bandit for achieving best performance in grid search by selecting an appropriate batch size at each epoch with a probability defined as a function of its previous success/failure . At each epoch , RMGD samples a batch size from its probability distribution , then uses the selected batch size for mini-batch gradient descent . After obtaining the validation loss at each epoch , the probability distribution is updated to incorporate the effectiveness of the sampled batch size . The benefit of RMGD is that it avoids the need for cumbersome grid search to achieve best performance and that it is simple enough to apply to any optimization algorithm using MGD . The detailed algorithm of RMGD are described in Section 4 , and experimental results are presented in Section 5 . 2 RELATED WORKS . There are only a few published results on the topic of batch size . It was empirically shown that SGD converged faster than GD on a large speech recognition database Wilson & Martinez ( 2003 ) . It was determined that the range of learning rate resulting in low test errors was considerably getting smaller as the batch size increased on convolutional neural networks and that small batch size yielded the best test error , while large batch size could not yield comparable low error rate Breuel ( 2015 ) . It was observed that larger batch size are more liable to converge to a sharp local minimum thus leading to poor generalization Keskar et al . ( 2016 ) . It was found that the learning rate and the batch size controlled the trade-off between the depth and width of the minima in MGD Jastrzkebski et al . ( 2017 ) . A small number of adaptive MGD algorithms have been proposed . Byrd et al . ( 2012 ) introduced a methodology for using varying sample size in MGD . A relatively small batch size is chosen at the start , then the algorithm chooses a larger batch size when the optimization step does not produce improvement in the target objective function . They assumed that using a small batch size allowed rapid progress in the early stages , while a larger batch size yielded high accuracy . However , this assumption did not corresponded with later researches that reported the degradation of performance with large batch size Breuel ( 2015 ) ; Keskar et al . ( 2016 ) ; Mishkin et al . ( 2017 ) . Another similar adaptive algorithm , which increases the batch size gradually as the iteration proceeded , was done by Friedlander & Schmidt ( 2012 ) . The algorithm uses relatively few samples to approximate the gradient , and gradually increase the number of samples with a constant learning rate . It was observed that increasing the batch size is more effective than decaying the learning rate for reducing the number of iterations Smith et al . ( 2017 ) . However , these increasing batch size algorithms lack flexibility since it is unidirectional . Balles et al . ( 2017 ) proposed a dynamic batch size adaptation algorithm . It estimates the variance of the stochastic gradients and adapts the batch size to decrease the variance . However , this algorithm needs to find the gradient variance and its computation depends on the number of model parameters . Batch size can also be considered as a hyperparameter , and there have been some proposals based on bandit-based hyperparameter ( but not batch size ) optimization which maybe applicable for determining the best fixed batch size . Jamieson & Talwalkar ( 2016 ) introduced a successive halving algorithm . This algorithm uniformly allocates a budget to a set of hyperparameter configurations , evaluates the performance of all configurations , and throws out the worst half until one configuration remains . Li et al . ( 2017 ) introduced a novel bandit-based hyperparameter optimization algorithm referred as HYPERBAND . This algorithm considers the optimization problem as a resource allocation problem . The two algorithms mentioned above are not adaptive , and for searching a small hyperparameter space , the two algorithms will not be very effective . The experimental results in this paper show that adaptive MGD tends to perform better than fixed MGD . 3 SETUP . Let B = { bk } Kk=1 be the set of possible batch size and π = { πk } Kk=1 be the probability distribution of batch size where bk , πk , and K are the kth batch size , the probability of bk to be selected , and number of batch sizes respectively . This paper considers algorithm for multi-armed bandit over B according to Algorithm 1 . Letwτ ∈ W be the model parameters at epoch τ , and w̃t be the temporal parameters at sub iteration t. Let J : W → R be the training loss function and let g = ∇J ( w ) be the gradients of training loss function with respect to the model parameters . ητ is the learning rate at epoch τ . Let ` : W → R be the validation loss function , and yk ∈ { 0 , 1 } be the cost of choosing the batch size bk . In here , yk = 0 if the validation loss decreases by the selected batch size bk ( well-updating ) and yk = 1 otherwise ( misupdating ) . The aim of the algorithm is to have low misupdating . For the cost function yk , graduated losses such as hinge loss and percentage of nonnegative changes in validation loss can be variations of 0-1 loss . However , there are no differences in regret bound among them in this setting and it is experimentally confirmed that there are little performance gaps among them . Therefore , this paper introduces the 0-1 loss , which is simple and basic . 4 RESIZABLE MINI-BATCH GRADIENT DESCENT . The resizable mini-batch gradient descent ( RMGD ) sets the batch sizes as multi arms , and at each epoch it samples one of the batch sizes from probability distribution . Then , it suffers a cost of selecting this batch size . Using the cost , probability distribution is updated . 4.1 ALGORITHMS . The overall framework of the RMGD algorithm is shown in Figure 1 . The RMGD consists of two components : batch size selector and parameter optimizer . The selector samples a batch size from probability distribution and updates the distribution . The optimizer is usual mini-batch gradient . Selector samples a batch size bkτ ∈ B from the probability distribution πτ at each epoch τ where kτ is selected index . Here bk is associated with probability πk . The selected batch size bkτ is applied to optimizer for MGD at each epoch , and the selector gets cost ykτ from optimizer . Then , the selector Algorithm 1 Resizable Mini-batch Gradient Descent Input : B = { bk } Kk=1 : Set of batch sizes π0 = { 1/K , . . . , 1/K } : Prior probability distribution Procedure : 1 : Initialize model parameters w0 2 : for epoch τ = 0 , 1 , 2 , . . . 3 : Select batch size bkτ ∈ B from πτ 4 : Set temporal parameters w̃0 = wτ 5 : for t = 0 , 1 , . . . , T − 1 where T = dm/bkτ e 6 : Compute gradient gt = ∇J ( w̃t ) 7 : Update w̃t+1 = w̃t − ητgt 8 : end for 9 : Update wτ+1 = w̃T 10 : Observe validation loss ` ( wτ+1 ) 11 : if ` ( wτ+1 ) < ` ( wτ ) 12 : Get cost ykτ = 0 13 : else 14 : Get cost ykτ = 1 15 : end if 16 : for i = 1 , 2 , . . . , K 17 : if i = kτ 18 : Set temporal probability π̃i = πiτe −βykτ /πiτ 19 : else 20 : Set temporal probability π̃i = πiτ 21 : end if 22 : end for 23 : Update ∀i ∈ [ K ] , πiτ+1 = π̃i/ ∑ j π̃ j 24 : end for updates probabilities by randomized weighted majority , for i = kτ , π̃i = πiτe −βykτ /πiτ for i 6= kτ , π̃i = πiτ ∀i , πiτ+1 = π̃i/ ∑ j π̃j where β ∈ ( 0 , 1 ) is positive hyperparameter . When τ = 0 , πτ = { 1/K , . . . , 1/K } . Optimizer updates the model parameters w. For each epoch , temporal parameters w̃0 is set to wτ , and MGD iterates T = dm/bkτ e1 times using the selected batch size bkτ where m is the total number of training samples : w̃t+1 = w̃t − ητgt , gt = ∇J ( w̃t ) . After T iterations at epoch τ , the model parameters is updated aswτ+1 = w̃T . Then , the optimizer obtains validation loss ` , and outputs cost as follows : ykτ = { 0 if ` ( wτ+1 ) < ` ( wτ ) 1 otherwise . The RMGD samples an appropriate batch size from a probability distribution at each epoch . This probability distribution encourages exploration of different batch size and then later exploits batch size with history of success , which means decreasing validation loss . Figure 2 shows an example of training progress of RMGD . The figure represents the probability distribution with respect to epoch . The white dot represents the selected batch size at each epoch . In the early stage of training , 1dxe is the least integer that is greater than or equal to x commonly , all batch sizes tend to decrease validation loss : π is uniform . Thus , all batch size have equal probability of being sampled ( exploration ) . In the later stages of training , the probability distribution varies based on success and failure . Thus , better performing batch size gets higher probability to be sampled ( exploitation ) . In this case , 256 is the best performing batch size .

Summary: The authors consider the problem of determining the minibatch size for SGD by first fixing a set of candidate sizes, and then learning a distribution over those sizes using a MAB algorithm. A minibatch size is first sampled from the distribution, then one training epoch is performed. A validation error is then computed, and if it is lower than that of the last epoch, the cost of the minibatch is taken to be zero (otherwise one), and the distribution is updated. This is Algorithm 1.