Minimization of Stochastic First-order Oracle Complexity of Adaptive Methods for Nonconvex Optimization

Adaptive methods are powerful deep learning optimizers used to find the model parameters of deep neural networks that minimize the expected risk and empirical risk loss functions [2, Section 4]. Specific adaptive methods are, for example, adaptive gradient (AdaGrad) [5], root mean square propagation (RMSProp) [26], adaptive moment estimation (Adam) [11], and adaptive mean square gradient (AMSGrad) [19] (Table 2 in [22] lists the main deep learning optimizers).

Deep learning optimizers have been widely studied from the viewpoints of both theory and practice. While the convergence and convergence rate of deep learning optimizers have been theoretically studied for convex optimization [31, 11, 19, 13, 15], theoretical investigation of deep learning optimizers for nonconvex optimization is needed so that such optimizers can be put into practical use in deep learning [28, 1, 27].

The convergence of adaptive methods has been studied for nonconvex optimization [7, 4, 30, 10], and the convergence of stochastic gradient descent (SGD) methods has been studied for nonconvex optimization [8, 3, 21, 12]. Chen et al. showed that generalized Adam, which includes the heavy-ball method, AdaGrad, RMSProp, AMSGrad, and AdaGrad with first-order momentum (AdaFom), has an $\mathcal{O}(\log K/\sqrt{K})$ convergence rate when a decaying learning rate ($\alpha_{k}=1/\sqrt{k}$) is used [4]. AdaBelief (which adapts the step size in accordance with the belief in the observed gradients) using $\alpha_{k}=1/\sqrt{k}$ has an $\mathcal{O}(\log K/\sqrt{K})$ convergence rate [30]. A method that unifies adaptive methods such as AMSGrad and AdaBelief has been shown to have a convergence rate of $\mathcal{O}(1/\sqrt{K})$ when $\alpha_{k}=1/\sqrt{k}$ [10], which improves on the results of [4, 30].

There have been several practical studies of deep learning optimizers. Shallue et al. [23] studied how increasing the batch size affects the performances of SGD, SGD with momentum [18, 20], and Nesterov momentum [17, 25]. Zhang et al. [29] studied the relationship between batch size and performance for Adam and K-FAC (Kronecker-factored approximate curvature [14]). In both studies, it was numerically shown that increasing the batch size tends to decrease the number of steps needed for training deep neural networks, but with diminishing returns. Moreover, it was shown that SGD with momentum and Nesterov momentum can exploit larger batches than SGD [23] and that K-FAC and Adam can exploit larger batches than SGD with momentum [29]. Thus, it has been shown that momentum and adaptive methods can greatly reduce the number of steps needed for training deep neural networks [23, Figure 4], [29, Figure 5]. Smith et al. [24] numerically showed that using enormous batches leads to reductions in the number of parameter updates and model training time.

As described in Section 1.1, there have been theoretical studies of the relationship between the learning rate and convergence of deep learning optimizers. Furthermore, as also described in Section 1.1, the practical performance of a deep learning optimizer strongly depends on the batch size. Our goal in this paper is to clarify theoretically the relationship between batch size and the performance of deep learning optimizers. We focus on the relationship between the batch size and the number of steps needed for nonconvex optimization of several adaptive methods.

To explicitly show our contribution, we characterize the diminishing returns reported in [23, 29] by using the stochastic first-order oracle (SFO) complexities of deep learning optimizers. We define the SFO complexity of a deep learning optimizer as $Kb$ on the basis of the number of steps $K$ needed for training the deep neural network and on the batch size $b$ used in the optimizer. Let $b^{\star}$ be a critical batch size such that there are diminishing returns for all batch sizes beyond $b^{\star}$, as asserted in [23, 29]. Thanks to the numerical evaluations in [23, 29], we know that there is a positive number $C$ such that

where $K$ and $b$ are defined for $i,j\in\mathbb{N}$ by $K=2^{i}$ and $b=2^{j}$. For example, Figure 5(a) in [23] shows $C\approx 16$ and $b^{\star}\approx 2^{8}$ for the momentum method used to train a convolutional neural network on the MNIST dataset. This implies that, while SFO complexity $Kb$ initially is not almost changed (i.e., $K$ halves for each doubling of $b$), $Kb$ is minimized at critical batch size $b^{\star}$, and there are diminishing returns once the batch size exceeds $b^{\star}$. Therefore, we can see that there are diminishing returns beyond the critical batch size $b^{\star}$ that minimizes SFO complexity.

The goal of this paper is to develop a theory demonstrating the existence of critical batch sizes, which was shown numerically by Shallue et al. and Zhang et al. [23, 29].

To reach this goal, we first examine the relationship between the number of steps $K$ needed for nonconvex optimization (Section 2) and the batch size $b$ used for each adaptive method. We show that, under certain assumptions, the lower bound $\underline{K}$ and the upper bound $\overline{K}$ of $K$ are rational functions of batch size $b$; i.e., $\underline{K}$ and $\overline{K}$ have the following forms (see Section 3 for details):

where $A$, $B$, $C$, $D$, $E$, $F$, and $G$ are positive constants depending on the parameters (e.g., constant learning rate $\alpha$ and momentum coefficient $\beta$) used in each adaptive method, and $\epsilon,\delta>0$ are the precisions for nonconvex optimization. $\underline{K}$ and $\overline{K}$ defined by (1.2) are monotone decreasing and convex for batch size $b$. Accordingly, $\underline{K}$ and $\overline{K}$ decrease with an increase in batch size $b$. For minimizing the empirical average loss function for the given number of samples $n$ (see (2.2)), $\underline{K}$ and $\overline{K}$ are minimized at $b=n$. We can see that there are asymptotic values $\underline{K}^{\star}=A/\{\epsilon^{2}-(C+D)\}$ and $\overline{K}^{\star}=E/(\delta\epsilon^{2}+G)$ of $\underline{K}$ and $\overline{K}$.

Figure 2 visualizes the relationship between $\underline{K}(b)$ and $\overline{K}(b)$. The number of steps $K$ needed for nonconvex optimization is inversely proportional to a batch size smaller than the critical batch size $b^{\star}$, and there are diminishing returns beyond the critical batch $b^{\star}$, as shown by the numerical results of Shallue et al. and Zhang et al. [23, 29]. Moreover, the critical batch size $b^{\star}$ minimizes SFO complexity $Kb$, as seen in (1.1). We thus define the critical batch size by using the global minimizer of SFO complexity $K(b)b$.

From the forms of $\underline{K}(b)$ and $\overline{K}(b)$ in (1.2), we have the forms of $\underline{K}(b)b$ and $\overline{K}(b)b$:

We can show that $\underline{K}(b)b$ and $\overline{K}(b)b$ are convex; i.e., there are global minimizers, denoted by $b_{*}$ and $b^{*}$, of $\underline{K}(b)b$ and $\overline{K}(b)b$, as seen in Figure 2 (see Section 3 for details). Therefore, we can regard batch sizes $b_{*}$ and $b^{*}$ as the critical batch sizes in the sense of minimizing $\underline{K}(b)b$ and $\overline{K}(b)b$ defined by (1.3). Suppose that $\underline{K}(b)b$ and $\overline{K}(b)b$ defined by (1.3) approximate the actual SFO complexity $Kb$; i.e.,

which implies that

where $M_{i}$ ($i=1,2,3,4,5$) are positive constants that do not depend on a constant learning rate $\alpha$ (see Section 3 for details). Hence, if precisions $\epsilon$ and $\delta$ are small, using a sufficiently small learning rate $\alpha$ could be used to approximate $K(b)b$ (Indeed, small learning rates were practically used by Shallue et al. [23]). Then, the demonstration of the existence of critical batch sizes $b_{*}$ and $b^{*}$ leads to the existence of the actual critical batch $b^{\star}\approx b_{*}\approx b^{*}$.

Let $\mathbb{R}^{d}$ be a $d$-dimensional Euclidean space with inner product $\langle\bm{x},\bm{y}\rangle:=\bm{x}^{\top}\bm{y}$ inducing norm $\|\bm{x}\|$, and let $\mathbb{N}$ be the set of nonnegative integers. Let $\mathbb{S}_{++}^{d}$ be the set of $d\times d$ symmetric positive-definite matrices, and let $\mathbb{D}^{d}$ be the set of $d\times d$ diagonal matrices: $\mathbb{D}^{d}=\{M\in\mathbb{R}^{d\times d}\colon M=\mathsf{diag}(x_{i}),\text{ }x_{i}\in\mathbb{R}\text{ }(i\in[d]:=\{1,2,\ldots,d\})\}$.

The mathematical model used here is based on that of Shallue et al. [23]. Given a parameter $\bm{\theta}$ in a set $\mathcal{S}\subset\mathbb{R}^{d}$ and given a data point $z$ in a data domain $Z$, a machine learning model provides a prediction for which the quality is estimated using a differentiable nonconvex loss function $\ell(\bm{\theta};z)$. We minimize the expected loss defined for all $\bm{\theta}\in\mathcal{S}$ by using

where $\mathcal{D}$ is the probability distribution over $Z$, $\xi$ denotes a random variable with distribution function $\mathrm{P}$, and $\mathbb{E}[\cdot]$ denotes the expectation taken with respect to $\xi$. A particularly interesting example of (2.1) is the empirical average loss defined for all $\bm{\theta}\in\mathcal{S}$:

where $S=(z_{1},z_{2},\ldots,z_{n})$ denotes the training set, $\ell_{i}(\cdot):=\ell(\cdot;z_{i})$ denotes the loss function corresponding to the $i$-th training data instance $z_{i}$, and $[n]:=\{1,2,\ldots,n\}$. Our main objective is to find a local minimizer of $L$ over $\mathcal{S}$, i.e., a point $\bm{\theta}^{\star}\in\mathcal{S}$ that belongs to the set

where $\nabla L\colon\mathbb{R}^{d}\to\mathbb{R}^{d}$ denotes the gradient of $L$. The inequality $(\bm{\theta}^{\star}-\bm{\theta})^{\top}\nabla L(\bm{\theta}^{\star})\leq 0$ is called the variational inequality of $\nabla L$ over $\mathcal{S}$ [6].

We assume that there exists SFO such that, for a given $\bm{\theta}\in\mathcal{S}$, it returns a stochastic gradient $\mathsf{G}_{\xi}(\bm{\theta})$ of the function $L$ defined by (2.1), where a random variable $\xi$ is supported on $\Xi$ and does not depend on $\bm{\theta}$. Throughout this paper, we assume three standard conditions:

1. (S1)

   $L\colon\mathbb{R}^{d}\to\mathbb{R}$ defined by (2.1) is continuously differentiable.

2. (S2)

   Let $(\bm{\theta}_{k})_{k\in\mathbb{N}}\subset\mathcal{S}$ be the sequence generated by a deep learning optimizer. For each iteration $k$,

   (2.4)

   $\displaystyle\mathbb{E}_{\xi_{k}}[\mathsf{G}_{\xi_{k}}(\bm{\theta}_{k})]=\nabla L(\bm{\theta}_{k}),$

   where $\xi_{0},\xi_{1},\ldots$ are independent samples, and the random variable $\xi_{k}$ is independent of $(\bm{\theta}_{l})_{l=0}^{k}$. There exists a nonnegative constant $\sigma^{2}$ such that

   (2.5)

   $\displaystyle\mathbb{E}_{\xi_{k}}\left[\left\|\mathsf{G}_{\xi_{k}}(\bm{\theta}_{k})-\nabla L(\bm{\theta}_{k})\right\|^{2}\right]\leq\sigma^{2}.$

3. (S3)

   For each iteration $k$, the deep learning optimizer samples a batch $B_{k}$ of size $b$ and estimates the full gradient $\nabla L$ as

   $\displaystyle\nabla L_{B_{k}}(\bm{\theta}_{k}):=\frac{1}{b}\sum_{i\in[b]}\mathsf{G}_{\xi_{k,i}}(\bm{\theta}_{k}),$

   where $\xi_{k,i}$ is the random variable generated by the $i$-th sampling in the $k$-th iteration.

In the case where $L$ is defined by (2.2), we have that, for each $k$, $B_{k}\subset[n]$ and

We consider the following algorithm (Algorithm 1), which is a unified algorithm for most deep learning optimizers,¹¹1We define $\bm{x}\odot\bm{x}$ for $\bm{x}:=(x_{i})_{i=1}^{d}\in\mathbb{R}^{d}$ by $\bm{x}\odot\bm{x}:=(x_{i}^{2})_{i=1}^{d}\in\mathbb{R}^{d}$. $\mathrm{C}(\cdot,l,u)\colon\mathbb{R}\to\mathbb{R}$ in AMSBound ($l,u\in\mathbb{R}$ with $l\leq u$ are given) is defined for all $x\in\mathbb{R}$ by $\displaystyle\mathrm{C}(x,l,u):=\begin{cases}l&\text{ if }x<l,\\ x&\text{ if }l\leq x\leq u,\\ u&\text{ if }x>u.\end{cases}$ including Momentum [18], AMSGrad [19], AMSBound [13], Adam [11], and AdaBelief [30], which are listed in Table 1.

Let $\xi_{k}=(\xi_{k,1},\xi_{k,2},\ldots,\xi_{k,b})$ be the random samplings in the $k$-th iteration, and let $\xi_{[k]}=(\xi_{0},\xi_{1},\ldots,\xi_{k})$ be the history of process $\xi_{0},\xi_{1},\ldots$ to time step $k$. The adaptive methods listed in Table 1 all satisfy the following conditions:

1. (A1)

   $\mathsf{H}_{k}:=\mathsf{diag}(h_{k,i})$ depends on $\xi_{[k]}$, and $h_{k+1,i}\geq h_{k,i}$ for all $k\in\mathbb{N}$ and all $i\in[d]$;

2. (A2)

   For all $i\in[d]$, a positive number $H_{i}$ exists such that $\sup_{k\in\mathbb{N}}\mathbb{E}[h_{k,i}]\leq H_{i}$.

We define $H:=\max_{i\in[d]}H_{i}$. The optimizers in Table 1 obviously satisfy (A1). Previously reported results [4, p.29], [30, p.18], and [10] show that $(\mathsf{H}_{k})_{k\in\mathbb{N}}$ in Table 1 satisfies (A2). For example, AMSGrad (resp. Adam) satisfies (A2) with $H=\sqrt{M}$ (resp. $H=\sqrt{M/(1-\zeta)}$), where $M:=\sup_{k\in\mathbb{N}}\|\nabla L_{B_{k}}(\bm{\theta}_{k})\odot\nabla L_{B_{k}}(\bm{\theta}_{k})\|<+\infty$.

The theoretical performance metric used to approximate a local minimizer in $\mathrm{VI}(\mathcal{S},\nabla L)$ (see (2.3)) is an $\epsilon$-approximation of the sequence $(\bm{\theta}_{k})_{k\in\mathbb{N}}$ generated by Algorithm 1 in the sense of the mean value of $(\bm{\theta}_{k}-\bm{\theta})^{\top}\nabla L(\bm{\theta}_{k})$ ($k\in[K]$); that is,

where $\epsilon>0$ and $\delta\in(0,1]$ are precisions and $\bm{\theta}\in\mathcal{S}$. It would be sufficient to consider that the right-hand side of (2.6), $\mathrm{V}(K,\bm{\theta})\leq\epsilon^{2}$, approximates a local minimizer of $L$. Since consideration of $\mathrm{V}(K,\bm{\theta})\leq\epsilon^{2}$ leads to the lower bound of $K$ satisfying $\mathrm{V}(K,\bm{\theta})\leq\epsilon^{2}$, we also consider $\delta\epsilon^{2}\leq\mathrm{V}(K,\bm{\theta})$, which leads to the upper bound of $K$.

We provide the lower bound of $K$ satisfying $\mathrm{V}(K,\bm{\theta})\leq\epsilon^{2}$ under the following conditions:

1. (L1)

   There exists a positive number $P$ such that $\mathbb{E}[\|\nabla L(\bm{\theta}_{k})\|^{2}]\leq P^{2}$, where $\bm{\theta}_{k}=(\theta_{k,i})$ is the sequence generated by Algorithm 1.

2. (L2)

   For all $\bm{\theta}=(\theta_{i})\in\mathcal{S}$, there exists a positive number $\mathrm{Dist}(\bm{\theta})$ such that $\max_{i\in[d]}\sup\{(\theta_{k,i}-\theta_{i})^{2}\colon k\in\mathbb{N}\}\leq\mathrm{Dist}(\bm{\theta})$.

Condition (L1) was used in previous studies [4, A2] and [30, Theorem 2.2] to analyze deep learning optimizers for nonconvex optimization. Here, we use (L1) to provide the upper bounds of $\mathbb{E}[\|\bm{\mathsf{d}}_{k}\|_{\mathsf{H}_{k}}^{2}]$ and $\mathbb{E}[\|\bm{m}_{k}\|^{2}]$ (see Lemma A.3 for details); i.e.,

where (S2)(2.5) and (A1) are assumed, $\tilde{\gamma}:=1-\gamma$, and $h_{0}^{*}:=\min_{i\in[d]}h_{0,i}$. This formulation implies that $A_{K}$ in (3.1) is bounded above (see Theorems A.1 and 3.1 for details).

Condition (L2) is assumed for both convex and nonconvex optimization, e.g., [16, p.1574], [11, Theorem 4.1], [19, p.2], [13, Theorem 4], and [30, Theorem 2.1]. Here, we use (L2) and (A2) to provide the upper bounds of $\Gamma_{K}$ and $B_{K}$ in (3.1) (see Theorem A.1 for details); i.e.,

where (S2)(2.5) and (L1) are assumed and $H:=\max_{i\in[d]}H_{i}$. Therefore, we have the following theorem (see the Appendix for detailed proof):

The minimization of $\underline{K}_{\epsilon}(b)b$ is as follows (The proof of Theorem 3.2 is given in the Appendix):

We provide the upper bound of $K$ satisfying $\delta\epsilon^{2}\leq\mathrm{V}(K,\bm{\theta})$ under the following conditions:

1. (U1)

   There exists $c_{1}\in[0,1]$ such that $\mathbb{E}[\|\bm{\mathsf{d}}_{k}\|_{\mathsf{H}_{k}}^{2}]\geq c_{1}\sigma^{2}/b$.

2. (U2)

   There exist $\bm{\theta}^{*}\in\mathrm{VI}(\mathcal{S},\nabla L)$ and $c_{2}\geq 0$ such that $\mathbb{E}[(\bm{\theta}^{*}-\bm{\theta}_{k})^{\top}\bm{m}_{k-1}]\geq-c_{2}(\sigma^{2}/b+P^{2})$.

3. (U3)

   For $\bm{\theta}^{*}\in\mathrm{VI}(\mathcal{S},\nabla L)$, there exists a positive number $X(\bm{\theta}^{*})$ such that, for all $k\geq 1$, $\tilde{\gamma}\mathbb{E}[\|\bm{\theta}_{1}-\bm{\theta}^{*}\|_{\mathsf{H}_{1}}^{2}]-\mathbb{E}[\|\bm{\theta}_{k+1}-\bm{\theta}^{*}\|_{\mathsf{H}_{k}}^{2}]\geq X(\bm{\theta}^{*})$, where $\tilde{\gamma}:=1-\gamma$.

We consider the case where $\mathsf{H}_{k}$ is the identity matrix needed to justify (U1). The definition of $\bm{m}_{k}:=\beta\bm{m}_{k-1}+(1-\beta)\nabla L_{B_{k}}(\bm{\theta}_{k})$ implies that $\bm{m}_{k}$ depends on $\nabla L_{B_{k}}(\bm{\theta}_{k})$. Here, we assume the following condition, which is stronger than (S2)(2.5): there exists $c_{1}\in[0,1]$ such that

Then, (S2)(2.4) and (3.3) ensure that

From $\tilde{\gamma}_{k}=1-\gamma^{k+1}\leq 1$, we have that $\|\bm{\mathsf{d}}_{k}\|^{2}=\tilde{\gamma}_{k}^{-2}\|\bm{m}_{k}\|^{2}\geq\|\bm{m}_{k}\|^{2}$ ($k\in\mathbb{N}$). Accordingly, the lower bound of $\|\bm{\mathsf{d}}_{k}\|^{2}$ depends on $c_{1}\sigma^{2}/b$.

Algorithm 1 for nonconvex optimization satisfies that there exist positive numbers $C_{1}$ and $C_{2}$ such that, for all $\bm{\theta}\in\mathcal{S}$,

(see [10, Theorem 1]). The sequence $(\bm{\theta}_{k})_{k\in\mathbb{N}}$ generated by Algorithm 1 can thus approximate a local minimizer $\bm{\theta}^{*}\in\mathrm{VI}(\mathcal{S},\nabla L)$. Meanwhile, we have that

Accordingly, for a sufficiently large $k$, the lower bound of $(\bm{\theta}^{*}-\bm{\theta}_{k})^{\top}\bm{m}_{k-1}$ depends on $-\mathbb{E}[\|\bm{m}_{k-1}\|^{2}]\geq-(\sigma^{2}/b+P^{2})$, where (S2)(2.5) and (L1) are assumed (see (3.2)). Hence, we assume (U2) for Algorithm 1.

Condition (U3) implies that $\mathbb{E}[\|\bm{\theta}_{k+1}-\bm{\theta}^{*}\|_{\mathsf{H}_{k}}^{2}]<\mathbb{E}[\|\bm{\theta}_{1}-\bm{\theta}^{*}\|_{\mathsf{H}_{1}}^{2}]<+\infty$, which is stronger than (L2), and that $\bm{\theta}_{k+1}$ approximates more appropriately $\bm{\theta}^{*}$ than $\bm{\theta}_{1}$. Since sequence $(\bm{\theta}_{k})_{k\in\mathbb{N}}$ generated by Algorithm 1 can approximate a local minimizer $\bm{\theta}^{*}\in\mathrm{VI}(\mathcal{S},\nabla L)$ (see (3.4)), (U3) is also adequate for consideration of the upper bound of $K$.

The minimization of $\overline{K}_{\epsilon,\delta}(b)b$ is as follows (The proof of Theorem 3.4 is given in the Appendix):