Recent years have witnessed the rapid development of deep neural networks in a wide range of applications. The success of neural networks should be partially attributed to optimization algorithms such as stochastic gradient descent (SGD) and its variants [Kingma and Ba, 2014]. These popular optimization methods, widely adopted in practice, have been proved to be extremely efficient and effective for finding the local minima of the objective function [Bottou, 2010;Bottou, 2012;Le et al., 2011]. SGD exploits the finite sum structure of the objective function, avoids the expensive computation of exact gradient, and thus provides a feasible and efficient optimization solution in practice [Bottou, 2012]. The stochastic nature also helps escaping saddle points, when Newton method cannot [Ge et al., 2015;Jin et al., 2017]. It has also been shown both empirically and * equal contribution theoretically that SGD and its variants improve the generalization performance of the trained models [Hardt et al., 2016;Chaudhari et al., 2017].
Despite the above advantages, one limitation with SGD type methods has manifested in recent years. Calculated using partial data, the stochastic gradient tends to fluctuate [Ruder, 2016] and deviate from the full gradient descent direction, causing deteriorated optimization efficiency. This issue becomes more serious when researchers continuously develop powerful yet large size networks to address big data challenges. With large networks, we are often forced to use small batch size due to the memory constraints. For example, in the high resolution image synthesis task [Wang et al., 2018], it requires a GPU of 24GB memory to train a generative adversarial network (GAN) using full resolution images with batch size 1. Such scenario demands small batch sizes, causing large variance in stochastic gradients, and thus inefficient optimization.
Variance reduction techniques have been developed to alleviate such issue [Johnson and Zhang, 2013;Reddi et al., 2016;Lei and Jordan, 2017]. These techniques reduce the noise of stochastic gradient by progressively estimating the difference between noisy and noiseless gradient. However, these techniques require the access to gradient computed by full batch or large batch size, and thus is impractical in the limited memory/small batch size scenario.
In this paper, we propose a new optimization method to reduce the variance of stochastic gradients with small batch size. Our method, called Trajectory Regularized Algorithm (TRAlgo), computes each stochastic gradient by locally adding a proximal type regularizer. The locally regularized loss function becomes strongly convex and can be solved efficiently [Rakhlin et al., 2012]. The local regularizer stabilizes the stochastic gradient, and thus reduces its variance. We prove that with small batch size, our method can achieve a much better stochastic gradient variance compared with SGD, and thus a better convergence rate. Our algorithm is fairly general and in principle can be combined with many existing SGD type methods. It opens the door toward novel optimization techniques that uniquely suit the big data world. In summary, our main contributions include:
• We propose a new algorithm using proximal regularizer to address the issue of noisy stochastic gradients, due to small batch size.
• We prove that our algorithm converges no slower than SGD. Furthermore, under the constraint of small batch size, our method reduces the variance of stochastic gradients, and thus has a better convergence rate than SGD. • We combine our algorithm with various state-of-the-art optimization algorithms (vanilla SGD and Adam). We empirically show that the proposed method improves the optimization performance when training deep neural networks with small batch sizes.
Note that the power of a proximal type regularizer has been exploited [Allen-Zhu, 2018;Lin et al., 2015]. However, existing works only focus on the optimization efficiency. To the best of our knowledge, TRAlgo is the first to use such technique to address noisy gradients, under the small batch size constraint.
The remainder of this paper is organized as follows. We introduce our notation and TRAlgo in Section 2. In Section 3, we combine TRAlgo with SGD (called TRSGD) and analyze the gradient complexity of both TRSGD and SGD. Experimental results are presented in Section 4. Finally, we conclude our work in Section 5.
We introduce the notations, assumptions, as well as our proposed TRAlgo framework. Throughout this paper, we denote by f (w) = 1 n n i=1 f i (w) the loss function, where w ∈ R d represents parameters in the optimization problem. We use
• to denote the L 2 norm. In this paper, we only consider convergence to a first order stationary point : ∇f (x) 2 ≤ .
The following assumptions are widely adopted in the analysis of convex and non-convex optimization algorithms. They regulate the behavior of zero-th order, first order and second order derivatives of f (x). Assumption 1. We assume following conditions hold:
In following assumption we apply δ to describe how 'non-
In general, δ is uniformly bounded by Hessian upper bound and could be equivalent to in the worst case. We can view δ as the minimum value of λ so that f (w)+ λ 2 w 2 is convex. Assumption 3 (Variance bounded SGD). We assume function f (x) is σ variance bounded:
The last assumption regulates the variance of stochastic gradient which is common in analyzing the convergence behavior of Stochastic Gradient Descent.
Algorithm 1 presents our Trajectory Regularized framework. It run for T iterations. At iteration t, we will construct Algorithm 1 TRAlgo (T, λ, w 0 , S)
Run stochastic gradient type method with decreasing step size on G t (w) for S iterations to compute w t+1 5: Option I: 6: Pick ŵ = w T 7: return ŵ
In practice 8: Option II: 9: Pick ŵ uniformly randomly from w 1 , .., w T 10: return ŵ
For analysis an auxiliary function G t (w) by adding a regularization term to f (x), and run stochastic gradient type optimization algorithm on G t (w) for S iterations. The result w after S steps is the next gradient step, w t+1 . After T iterations, depending on the context, we may return w T as the final output (in practice), or return a random sample from all previous w * (for analysis).
The key idea of TRAlgo is to apply a proximal type regularizer λ 2 w -w t 2 with λ ≥ 2δ to make the loss function f (x) strongly convex. The regularizer stablizes the gradient and thus reduces its variance, as we will prove in the next section. To avoid introducing bias as a fixed regularizer does [Hoerl and Kennard, 1970], the regularizer is locally adaptive (centered at current w t ). Quadratic proximal regularizer has been used before [Parikh et al., 2014], but only for better optimization of non smooth L 1 loss function. Instead, TRAlgo uses the it for a trajectory regularization purpose. In TRAlgo, the choice of the stochastic gradient method in each iteration is flexible. One can apply a vanilla SGD or its variants: Adam [Kingma and Ba, 2014] and AMSGrad [Reddi et al., 2018]. In next section, we use SGD as an example to analyze the benefits of TRAlgo, while in Section 4, we empirically apply TRAlgo to both SGD and Adam.
In this section we incorporate SGD into our framework TRAlgo (called TRSGD) and compare it with Vanilla SGD. In Section 3.1 we review SGD and restate some known results of SGD for a direct comparison with analysis for TRSGD. In Section 3.3 we formally define TRSGD and analyze its convergence rate using batch size 1. While TRSGD achieves comparable gradient complexity with Vanilla SGD in general it attains a better convergence rate when batch size is constrained. The analysis focuses on quantifying how number of iterations S run on auxiliary function G w (w) affects the convergence. Our analysis demonstrates even when batch size is limited, TRSGD could still control variance.
In Algorithm 2 we describe the Stochastic Gradient Descent method. We use I k ⊂ [n] to denote the random sampled
For practice 7: Option II: 8: Pick ŵ uniformly randomly from w 1 , .., w S 9: return ŵ For analysis small batch and ∇f
) is known as noise in stochastic gradient. Next we introduce some well known results about the convergence behavior of SGD and demonstrate how the noise affects the convergence. One can find similar statements in [Allen-Zhu, 2018;Ge et al., 2015;Reddi et al., 2016].
Lemma 1. Under assumptions 1 and 3, the iterative update of SGD with step size η ≤ 1 2 and batch size M satisfies the following inequality:
Recall that the SGD update rule is w t+1 = w t -η(∇f (w t ) + σ t ) where σ t represents noise in the Stochastic Gradient with
The left hand side of Eq. ( 1) is the improvement made in each iteration in expectation and the right hand side gives a guarantee on the progress for iteration t. One can observe that large variance will slow down the improvement brought by gradient. To alleviate this issue, one may increase batch size M . However, in our setting, when M is restricted to small, we cannot mitigate the negative impact of variance. Such phenomenon is shown in Figure 1. By telescoping Lemma 1 and using the fact that f (x) -f (y) ≤ 2B, we have the convergence rate of SGD. Theorem 1. Under assumptions 1 and 3, for SGD with step size η = 1/(2 ), batch size M , and any T , it holds that: In particular, denote by N = T M the total number of gradient calls applied in the algorithm, picking T = 4 √ BN /σ, the gradient complexity to reach E ∇f (w) 2 ≤ is:
Above Theorem suggests that if picking the output of SGD uniformly randomly over w 1:T , one can achieve
. One could increase training iteration T and batch size M so that the expected gradient squared norm meets the first order stationary point criterion. An optimal choice balancing ( B) /T and σ 2 T /N would be T = 4 √ BN /σ. Although the above statements reveal the power of SGD in efficiency, it is not realistic to set M = Ω( √ N ) in training large size deep neural network. Indeed, in the practical setting we are addressing, M may be a small constant (say 1). The right hand side of Eq. ( 2) can be arbitrarily bad with noisy gradients. In next section, we will show how TRAlgo addresses the issue and still keep a tight bound even with M = 1.
Our method TRAlgo requires the choice of a base optimization algorithm. In this section, we instantiate our method using the Vanilla SGD. We analyize the convergence performance of this instantiation, called TRSGD. Definition 1. TRSGD is defined as applying SGD in TRAlgo:
While the advantage of decreasing stepsize (usually set to be 1/k) for SGD type methods has been well studied in strongly convex case [Rakhlin et al., 2012;Shamir and Zhang, 2013], it remains unclear how to properly apply such strategy in non-convex stochastic optimization. In the rest of this section we show one can benefit from such strategy due to the convexity of auxiliary function G t (w).
First we state a technical lemma which follows [Rakhlin et al., 2012]. We use the following lemma to quantify the progress made by running SGD for S iterations with decreasing step size. Lemma 2. If one runs SGD on G t (w) = f (w)+ size 1, one have
where w * t+1 = arg min w {G t (w)}. Now we are ready to prove the key lemma which describes one round progress of TRSGD. One can make a direct comparison between the following lemma and Lemma 1 by setting η = 1/(2 ) in Eq. ( 1). Lemma 3. Under assumptions 1, 2 and 3, the iterative update of TRSGD with λ = 2 and batch size 1 satisfies the following inequality:
2 we can derive (3).
While in Eq. ( 1) the issue of variance could be alleviated by increasing the batch size M for SGD, one can alleviate such issue by increasing the number of iterations S in Eq. (3) for TRSGD. One can set S = Θ(M ) so both algorithms have a comparable one round convergence progress. However, in practice the choice of M is limited due to the memory restriction but the choice of S is not. We choose batch size to be 1 for simplicity and the analysis can be generalized to any batch size.
We finish this section by telescoping Eq. ( 3) to derive the following theorem. The inequality bounds the expected squared norm of gradient of TRSGD, similar to Theorem 1. Theorem 2. Under assumptions 1, 2 and 3, the TRSGD with batch size 1 , λ = 2 achieves the following convergence rate:
In particular, let N = T S be the total number of gradient call applied in the algorithm, by picking
By setting the batch size to be 1, the gradient complexity of TRSGD is comparable to SGD in a rate of N = O 1/ 2 to meet the criterion E ∇f (w t ) 2 ≤ .
Interpretation of theorems. We can now explicitly compare Eq. ( 4) and Eq. ( 2). While the function loss term B/T appear in both inequality the σ 2 term appears in different fashion due to their strategies in monitoring variance. When gradient is associated with big variance, SGD increases batch size M and TRSGD runs more iteration on auxiliary function G t (w). The strategy of TRSGD is not restricted by memory which is adorable in training huge network. The choice of S depends on the batch size restriction. In the case where one only affords to set M = 1, one can take S = Θ( √ N ) to optimally balance B/T and ( 2 (σ 2 + g 2 )T )/N in the gradient complexity. On another hand, when it is possible to set M = Θ( √ N ), the TRAlgo framework will not be the first choice as there is no need for variance control.
In this section we apply the proposed framework to momentum SGD and Adam algorithms, called TRSGD and TRAdam, respectively, and empirically illustrate the effectiveness of our methods when training deep neural networks with small batch sizes. We first compare the performance of different algorithms on image classification tasks on CI-FAR10 and CIFAR100 datasets, and then show the application of our framework in the training of GANs using an image to image translation task on the Cityscapes dataset [Cordts et al., 2016]. We also perform ablation study on the batch size and iteration number S on TRSGD. [Huang et al., 2017]. Four algorithms (TRSGD, SGD, TRAdam, Adam) are used to train these networks on CI-FAR10 and CIFAR100. In all experiments, the number of epoch is 200 and the batch size is 8. For TRSGD and SGD, the momentum is 0.9, the learning rate is initially set to 0.1 and decayed to 0.01, 0.001 at epoch 100 and 150, respectively. For TRAdam and Adam, β 1 = 0.9, β 2 = 0.999, the learning rate is initially set to 0.001 and decayed to 0.0001, 0.00001 at epoch 100 and 150, respectively. For both TRSGD and TRAdam, λ = 0.1 and S = 10.
Image to image translation using GAN. In this task we follow the network structure and parameter settings of [Isola et al., 2017] to train a model that generates images from semantic labels on Cityscapes dataset. The only difference is that we replace the Adam optimizer with our TRAdam algorithm. In TRAdam, λ = 0.1, S = 10 and the other parameters are the same as Adam.
Image classification. The train and test curves of SGD vs. TRSGD, and Adam vs. TRAdam using a batch size of 8 are shown in Fig. 2 and Fig. 3, respectively. Only results of ResNet101 are presented due to the space limitation. We can observe that our TRSGD and TRAdam achieve faster convergence and better accuracies compared to SGD and Adam on both datasets. The test accuracies of all models are reported in Table 1. It is evident that our TRSGD and TRAdam outperform SGD and Adam in all cases. Because of the large gradient variance when the batch size is small, SGD and Adam are not able to obtain better performance with deeper networks. However, our algorithms can handle the noisy gradients by the regularizer, thus achieving higher accuracies using ResNet101 and DenseNet121, especially on CIFAR100 dataset.
Image to image translation using GAN. To evaluate the image quality generated by the model, we use the "FCNscore" as in [Isola et al., 2017]. It evaluates how realistic the generated images are based on the semantic segmentation results using a fully convolutional neural network. The pre-trained model predicts labels from the synthesized images. Then the labels are compared with the ground-truth labels those images were synthesized from. We adopt two pre-trained semantic segmentation models to com- pute "FCN-score": FCN-8s [Long et al., 2015] and DPC-xception71 [Isola et al., 2017]. FCN-8s is used in Isola et al. [2017] while DPC-xception71 has the state-of-the-art performance on the Cityscapes dataset thus is more reliable than FCN-8s to compute the score. The evaluation results using two models are reported in Table 2 andTable 3. For the results evaluated by FCN-8, TRAdam is only slightly better in the per-pixel accuracy compared to Adam. However, TRAdam outperforms Adam in all three metrics when evaluated by DPC-xception71. The ground-truth metrics (GT in the tables) using original images for segmentation are also provided for reference. When computing the ground-truth values, we use the resized 256×256 version of original images because the resolution of synthesized images is 256×256. Some synthesized images are presented in Fig. 4.
In this section we explore the effect of batch size and inner iteration number S on our TRSGD algorithm. ResNet18 and CIFAR10 dataset are used in all experiments of this section. the variance of stochastic gradients. In Table 5 we test the performance of TRSGD under different values of S. The performance of TRSGD is not sensitive with regard to the value of S, One can also observe that when S = 50, the accuracy becomes slightly worse than other cases. This is because when gradient complexity is fixed, large S will results in small T (number of total iterations of TRSGD) thus a under fitting phenomenon may happen, as shown in Theorem 2.
In this paper we present a Trajectory Regularization framework to train deep neural network with noisy gradients. We combine our framework with SGD and show it relaxes the batch size constraint of SGD, while achieving a comparable gradient call complexity. Our theoretical analysis is supported by the empirical evidence. Besides, we illustrate that our framework can be applied to improve the performance of a popular SGD type algorithm Adam in the small batch size setting, on both image classification tasks and image-toimage translation task using GAN. The proposed framework may be applied to other GAN-related tasks because of the large memory requirements of these tasks. We leave this for the future work.