Paper: Iteration and Stochastic First-order Oracle Complexities of Stochastic Gradient Descent using Constant and Decaying Learning Rates (ICLR-2024)

Reviewer comment (Theory):
In the abstract, the authors made a claim that "SGD using a step-decay learning rate and a small batch size reduces the SFO (Stochastic First-Order) complexity to find a local minimizer of a loss function." However, upon reviewing the paper, it becomes apparent that the study primarily focuses on demonstrating the convergence of SGD to a stationary point rather than specifically proving convergence to a local minimizer.

Author response:
Thank you for pointing out. I will revise our claim as "SGD using a step-decay learning rate and a small batch size reduces the SFO complexity to find a stationary point of a loss function."

Reviewer comment (Evaluation):
I do not think SFO is a reasonable measure. In practice, different individual gradients are calculated parallelly and the corresponding time does not accumulate across samples.

Author response:
Thank you for your valuable comment. It is important to evaluate the implementation time of SGD. In this paper, in both theory and practice, we study the relationship between the batch size and the iteration complexity/the SFO complexity without computing $b$ gradients parallelly.

Reviewer comment (Theory):
The definition of K_{\epsilon} and N_{epsilon} seems to be weird, since the learning rate does not appear in any side of the equation.

Author response:
Thank you again for your comment. In the revision, we would like to replace (2) with $K_\epsilon (n,b,\alpha_k,\Delta,L,\sigma^2)$ and  $N_\epsilon (n,b,\alpha_k,\Delta,L,\sigma^2)$.

Reviewer comment (Novelty):
I wonder what is the novelty of Theorem 3.1. Is not it a very basic analysis of SGD?

Author response:
Theorem 3.1 is used to provide Theorems 3.2 and 3.3. Our main results are Theorems 3.2 and 3.3. Hence, we may replace Theorem 3.1 by e.g., Lemma 3.1/Proposition 3.1.

Reviewer comment (Theory):
I find the result of Decay 4 problematic. Specifically, in Theorem 3.2, isn't $\underline{\alpha}$ itself depends on $K(b)$? How can $K(b)$ be further calculated by $\underline{\alpha}$? That being said, when $T$ is independent of $\epsilon$, T is in the same order as $\varepsilon$ as P. Therefore, $\underline{\alpha}$ depends exponentially on $K$. Applying this to Theorem 3.2, it indicates $K(b)$ is also exponentially dependent over $\varepsilon$ and contradicts Theorem 3.4.

Author response:
We apologize for the insufficient analysis of Decay 4. Based on [1], we would like to revise our analysis of Decay 4.

First, we modify the definition of Decay 4 as follows. Let $\alpha_0 > 0$, $\beta \geq 2$, $T, P \geq 1$, and $K = TP$.
A step-decay learning rate is
$$(\alpha_k)_{k=0}^{K-1}
= \left(\underbrace{\alpha_0, \alpha_0, \cdots, \alpha_0}_T,
\frac{\alpha_0}{\beta}, \frac{\alpha_0}{\beta}, \cdots, \frac{\alpha_0}{\beta}, \cdots,
\frac{\alpha_0}{\beta^{P-1}}, \frac{\alpha_0}{\beta^{P-1}}, \cdots, \frac{\alpha_0}{\beta^{P-1}} \right),$$
which is monotone decreasing for $k$.
We assume that $\alpha_0 \leq \frac{1}{L}$, which implies that, for all $k \in [0:K-1]$, $\alpha_k \leq \frac{1}{L}$.

Next, we can show that
\begin{align*}
\alpha_k \left(1 - \frac{L \alpha_k}{2} \right) \mathbb{E} \left[ \Vert  \nabla f({\theta}_{k}) \Vert^2 \right] 
\leq \mathbb{E} [ f (\theta_k)-f (\theta_k+1)]+\frac{L \sigma^2 \alpha_k^2}{2b}
\end{align*}
based on Lemma A.1 (our paper).

We can thus modify the result of Decay 4 in Theorem 3.1 as follows:

Let $P = \log_{\alpha_0} K/2$. Then, the sequence $(\theta_k)$ generated by Algorithm 1 under (C1)-(C3) satisfies that, for all $K \geq 1$,
\begin{align*}
\min_{k \in [0:K-1]} \mathbb{E}\left[\Vert \nabla f({\theta}_k) \Vert^2 \right]
\leq
\frac{D_3 (\sqrt{K}-2 )}{\sqrt{K}-1}+\frac{D_4}{(\sqrt{K}-1)b},
\end{align*}
where $D_3 = \frac{(\beta -1) \Delta}{\alpha_0}$ and $D_4 = \alpha_0 L \sigma^2 (\beta -1)$.

Then, we have that $K$ and $N$ of SGD using Decay 4 needed to an $\epsilon$-approximation are
$$K(b) = \frac{1}{(D_3 - \epsilon^2)^2} \left( (2D_3 - \epsilon^2)^2  - \frac{2 D_4 (2 D_3 -  \epsilon^2)}{b} + \frac{D_4^2}{b^2} \right).$$
$$N(b) = \frac{b}{(D_3 - \epsilon^2)^2} \left( (2D_3 - \epsilon^2)^2  - \frac{2 D_4 (2 D_3 -  \epsilon^2)}{b} + \frac{D_4^2}{b^2} \right).$$
As a result, we have the following: 
- The iterations $K(b)$ is monotone decreasing for $b < \frac{D_4}{2D_3 - \epsilon^2}$ and convex for $b < \frac{3 D_4}{2(2 D_3 - \epsilon^2)}$ .
- The SFO complexity $N(b)$ is convex for $b > 0$ and $N'(b) > 0$ holds for all $b > 0$.
- We have that $K_\epsilon = O(1/\epsilon^4)$ and $N_\epsilon = O(1/\epsilon^4)$. 

We will revise the manuscript based on the above discussion.  

[1] Wang, Xiaoyu, Sindri Magnússon, and Mikael Johansson. “On the convergence of step decay step-size for stochastic optimization.” Advances in Neural Information Processing Systems 34 (2021): 14226-14238.

Reviewer comment (Theory):
On page 3, Is $N_{\epsilon}$ just $bK_{\epsilon}$? If yes, why not use the simpler one?

Author response:
Thank you. We will revise it accordingly.

Reviewer comment (Theory):
The convergence rate given for SGD with the step-decay schedule is misleading and leads to incorrect iteration and oracle complexities for this method.

Author response:
We apologize for the insufficient analysis of Decay 4. Based on [1], we would like to revise our analysis of Decay 4.

First, we modify the definition of Decay 4 as follows. Let $\alpha_0 > 0$, $\beta \geq 2$, $T, P \geq 1$, and $K = TP$.
A step-decay learning rate is
$$(\alpha_k)_{k=0}^{K-1}
= \left(\underbrace{\alpha_0, \alpha_0, \cdots, \alpha_0}_T,
\frac{\alpha_0}{\beta}, \frac{\alpha_0}{\beta}, \cdots, \frac{\alpha_0}{\beta^{P-1}}, \frac{\alpha_0}{\beta^{P-1}}, \cdots, \frac{\alpha_0}{\beta^{P-1}} \right),$$
which is monotone decreasing for $k$.
We assume that $\alpha_0 \leq \frac{1}{L}$, which implies that, for all $k \in [0:K-1]$, $\alpha_k \leq \frac{1}{L}$.

Next, we can show that
\begin{align*}
\alpha_k \left(1 - \frac{L \alpha_k}{2} \right) \mathbb{E} \left[ \Vert  \nabla f({\theta}_{k}) \Vert^2 \right]
\leq \mathbb{E} [ f (\theta_k)-f (\theta_k+1)]+\frac{L \sigma^2 \alpha_k^2}{2b}
\end{align*}
based on Lemma A.1 (our paper).

We can thus modify the result of Decay 4 in Theorem 3.1 as follows:

Let $P = \log_{\alpha_0} K/2$. Then, the sequence $(\theta_k)$ generated by Algorithm 1 under (C1)-(C3) satisfies that, for all $K \geq 1$,
\begin{align*}
\min_{k \in [0:K-1]} \mathbb{E}\left[\Vert \nabla f({\theta}_k) \Vert^2 \right]
\leq
\frac{D_3 (\sqrt{K}-2 )}{\sqrt{K}-1}+\frac{D_4}{(\sqrt{K}-1)b},
\end{align*}
where $D_3 = \frac{(\beta -1) \Delta}{\alpha_0}$ and $D_4 = \alpha_0 L \sigma^2 (\beta -1)$.

Then, we have that $K$ and $N$ of SGD using Decay 4 needed to an $\epsilon$-approximation are
$$K(b) = \frac{1}{(D_3 - \epsilon^2)^2} \left( (2D_3 - \epsilon^2)^2  - \frac{2 D_4 (2 D_3 -  \epsilon^2)}{b} + \frac{D_4^2}{b^2} \right).$$
$$N(b) = \frac{b}{(D_3 - \epsilon^2)^2} \left( (2D_3 - \epsilon^2)^2  - \frac{2 D_4 (2 D_3 -  \epsilon^2)}{b} + \frac{D_4^2}{b^2} \right).$$
As a result, we have the following:
- The iterations $K(b)$ is monotone decreasing for $b < \frac{D_4}{2D_3 - \epsilon^2}$ and convex for $b < \frac{3 D_4}{2(2 D_3 - \epsilon^2)}$ .
- The SFO complexity $N(b)$ is convex for $b > 0$ and $N'(b) > 0$ holds for all $b > 0$.
- We have that $K_\epsilon = O(1/\epsilon^4)$ and $N_\epsilon = O(1/\epsilon^4)$.

We will revise the manuscript based on the above discussion.

Reviewer comment (Writing):
The manuscript is unnecessarily "mathy" and many equations could be omitted while maintaining the same results. See, for example, Equations 1 and 2. While this makes the writing seem superficially impressive, it is difficult to read and detracts from the flow of the paper.

Author response:
We kindly agree with your comment. We will revise the manuscript accordingly.

Reviewer comment (Experiments):
None of the experiments serve to verify the theoretical derivations. I would liked to see at least one synthetic experiment for which the problem constants are known (e.g. a simple quadratic) and $b^*$ can be computed. Plots similar to that in Figures 1/2 could then show that $b^*$ does, in fact, obtain the optimal oracle complexity as claimed.

Author response:
Thank you for your valuable comments. In the revision, we clarify that $n$ can be taken arbitrarily large and $\epsilon$ is a sufficiently small positive constant. 

We now prepare a synthetic experiment. Hence, we plan to reply further comments to Comment 4 as soon as possible.

Reviewer comment (Theory):
"SGD using a decaying learning rate...": some additional comment on the type of learning rate decay is needed here. SGD with step-size $\alpha_k = 1 / \log(k)$ does not converge as $O(1/\sqrt{K})$ despite $\alpha_k \rightarrow 0$.

Author response:
Thank you for your valuable comment. In the revision, we would like to discuss SGD with $\alpha_k = \frac{1}{\log k}$.

Reviewer comment (Theory):
First math display in Section 1.3.2: this bound should mention somewhere that $b > n$ isn't possible and $b = n$ reduces to full-batch gradient descent. As a result, it is not always feasible to select the batch-size to minimize the oracle complexity.

Author response:
Thank you for pointing out. We will add the condition $b \leq n$ to (C3). Also, we will revise Section 1.3.2 accordingly.

Reviewer comment (Theory):
"Accordingly, small batch sizes are appropriate for a decaying learning rate or a step-decay learning rate": why is this true? You have said that SFO complexity has no positive stationary point, but that doesn't imply it is increasing in $b$ or that a small batch-size minimizes the complexity over the positive integers. Can you please address this fact?

Author response:
Thank you for your valuable comment. In (Decay 2), $N(b)$ is convex for $b > 0$ (see Theorem 3.3). Moreover, in (Decay 2), $b^\star < 0$ holds, and hence, $N(b)$ is monotonically increasing for $b>0$. In (Decay 4), we have that $N(b)$ is convex for $b > 0$ and $N'(b) > 0$ holds for all $b > 0$. Therefore, small batch sizes are appropriate for (Decay 2) and (Decay 4) in the sense of minimizing the SFO complexities.

Reviewer comment (Theory):
Theorem 3.1: I am concerned by the presentation of the convergence rates for SGD with step-decay. Firstly, $\underline{\alpha}$ does not appear to be defined anywhere. From the proof in the appendix, it seems $\underline{\alpha} = \alpha_{K-1} = \alpha \eta^{K/T-1}$. This quantity depends on $K$ — it is exponentially decreasing every $K/T$ iterations — so that it is incorrect to write it as a constant factor. Similarly, $D_3$ depends on $T$, which may or may not have a relationship with $K$ depending on algorithm parameters.

Author response:
We apologize for the insufficient analysis of Decay 4. Based on [1], we would like to revise our analysis of Decay 4.

First, we modify the definition of Decay 4 as follows. Let $\alpha_0 > 0$, $\beta \geq 2$, $T, P \geq 1$, and $K = TP$.
A step-decay learning rate is
$$(\alpha_k)_{k=0}^{K-1}
= \left(\underbrace{\alpha_0, \alpha_0, \cdots, \alpha_0}_T,
\frac{\alpha_0}{\beta}, \frac{\alpha_0}{\beta}, \cdots, \frac{\alpha_0}{\beta^{P-1}}, \frac{\alpha_0}{\beta^{P-1}}, \cdots, \frac{\alpha_0}{\beta^{P-1}} \right),$$
which is monotone decreasing for $k$.
We assume that $\alpha_0 \leq \frac{1}{L}$, which implies that, for all $k \in [0:K-1]$, $\alpha_k \leq \frac{1}{L}$.

Next, we can show that
\begin{align*}
\alpha_k \left(1 - \frac{L \alpha_k}{2} \right) \mathbb{E} \left[ \Vert  \nabla f({\theta}_{k}) \Vert^2 \right]
\leq \mathbb{E} [ f (\theta_k)-f (\theta_k+1)]+\frac{L \sigma^2 \alpha_k^2}{2b}
\end{align*}
based on Lemma A.1 (our paper).

We can thus modify the result of Decay 4 in Theorem 3.1 as follows:

Let $P = \log_{\alpha_0} K/2$. Then, the sequence $(\theta_k)$ generated by Algorithm 1 under (C1)-(C3) satisfies that, for all $K \geq 1$,
\begin{align*}
\min_{k \in [0:K-1]} \mathbb{E}\left[\Vert \nabla f({\theta}_k) \Vert^2 \right]
\leq
\frac{D_3 (\sqrt{K}-2 )}{\sqrt{K}-1}+\frac{D_4}{(\sqrt{K}-1)b},
\end{align*}
where $D_3 = \frac{(\beta -1) \Delta}{\alpha_0}$ and $D_4 = \alpha_0 L \sigma^2 (\beta -1)$.

Then, we have that $K$ and $N$ of SGD using Decay 4 needed to an $\epsilon$-approximation are
$$K(b) = \frac{1}{(D_3 - \epsilon^2)^2} \left( (2D_3 - \epsilon^2)^2  - \frac{2 D_4 (2 D_3 -  \epsilon^2)}{b} + \frac{D_4^2}{b^2} \right).$$
$$N(b) = \frac{b}{(D_3 - \epsilon^2)^2} \left( (2D_3 - \epsilon^2)^2  - \frac{2 D_4 (2 D_3 -  \epsilon^2)}{b} + \frac{D_4^2}{b^2} \right).$$
As a result, we have the following:
- The iterations $K(b)$ is monotone decreasing for $b < \frac{D_4}{2D_3 - \epsilon^2}$ and convex for $b < \frac{3 D_4}{2(2 D_3 - \epsilon^2)}$ .
- The SFO complexity $N(b)$ is convex for $b > 0$ and $N'(b) > 0$ holds for all $b > 0$.
- We have that $K_\epsilon = O(1/\epsilon^4)$ and $N_\epsilon = O(1/\epsilon^4)$.

We will revise the manuscript based on the above discussion.

Reviewer comment (Theory):
Theorem 3.4: In addition to the issue with the complexity of step-decay raised previously, this theorem assumes that $b \leq n$ can be chosen arbitrarily large in order to obtain the desired complexity. For example, SGD with a constant step-size requires $b \geq 2 C_2 \epsilon$, which diverges to infinity as $\epsilon \rightarrow 0$. But this is not sensible because $n$ is assumed to be a fixed, finite number of training examples. If this is not the case, then the authors must specify somewhere that they assume a setting where $n$ can be taken arbitrarily large.

Author response:
Thank you for your valuable comments. In the revision, we clarify that $n$ can be taken arbitrarily large and $\epsilon$ is a sufficiently small positive constant. 

We now prepare a synthetic experiment. Hence, we plan to reply further comments to Comment 4 as soon as possible.

Reviewer comment (Novelty):
The paper's theoretical results are known in the literature (e.g., [r1], [r2]).

Author response:
We sincerely apologize for our claim without checking carefully the previous results.

Corollaries 6 and 7 in [r2] are the same as our results for SGD with a constant learning rate in Theorems 3.1 and 3.3, since the randomized stochastic projected gradient free algorithm in [r2] that is one of stochastic zeroth-order (SZO) methods and coincides with SGD can be applied to the situations where only noisy function values are available. In particular, Corollary 6 shows the convergence rate of the SZO method using a fixed batch size and Corollary 7 indicates the SZO complexity that is the same as the SFO complexity (our paper).  Hence, Corollaries 6 and 7 in [r2] lead to the finding that the iteration complexity of SGD is $O(1/\epsilon^2)$ and the SFO complexity of SGD is $O(1/\epsilon^4)$.

Therefore, we deleted our claim such that “To the best of our knowledge, this is the first paper to provide that the iteration complexity of SGD using a constant learning rate and the critical batch size is $O(1/\epsilon^2)$.”

Reviewer comment (Theory):
The hardness result [r3] says that: whatever batch size the SGD algorithm can choose, the SFO cannot be better than \(O(1/\epsilon^4)\).

Author response:
Reference [r3] showed that the iteration complexity of the general SGD does not perform better than the order of $\frac{1}{\epsilon^4}$. However, we would like to emphasize that there is a room such that, if specific learning rate and batch size are used, then SGD will perform better than the order of $\frac{1}{\epsilon^4}$. Our paper shows that, if we can use a constant learning rate $\alpha < 2/L$ and the critical batch size $b^* = 2 C_2/\epsilon^2$, then the iteration complexity breaks through the order of $\frac{1}{\epsilon^4}$ and becomes $O(\frac{1}{\epsilon^2})$. We can also check that using specific batch size performs well in practice. For example, Figure 2 (our paper) indicates that the SFO complexity of SGD using $b = 2^4$ is better than the ones using other batch sizes.

Reviewer comment (Novelty):
Can you describe how your approach is better than the references I gave above (e.g., [r1, r2])?

Author response:
Thank you for your valuable comments.

Assumptions (C1)--(C3) (our paper) are general and well used in previous papers. From the experimental results in the previous paper [1] and our experimental results (see Figures 2 and 6, and so on), it is clear that the SFO complexity of SGD depends on both the learning rate and the batch size. Therefore, we study the relationship between batch size and the SFO complexity of SGD using specific learning rates (Constant and Decay 1--4) from the viewpoints of both theory and practice. Our results show that

- The number of iterations $K$ needed for nonconvex optimization of SGD is a monotone decreasing and convex function with respect to the batch size $b$ (Theorem 3.2);
- The SFO complexity $N$ is convex with respect to the batch size $b$ (Theorem 3.3);
- The existence of the critical batch size that is a global minimizer of $N$ depends on learning rates (Theorem 3.3).

The above results (Theorems 3.2 and 3.3) are our contributions compared with the previous results in [r1] and [r2].