Stochastic Steepest Descent with Acceleration for $\ell_p$-Smooth Non-Convex Optimization

Thank you for your replies. First, I would like to reply your comment: As we noted in Theorem 1 ...

The authors set $\eta = \frac{1}{L^{1/2} G^{1/(p-1)} T^{1/2}}$. Here, we must notice that $T$ is set before implementing the algorithm. For example, let $n = 50000$ be the number of training data (e.g., CIFAR) and let $b = 32$ be the batch size. Then, let us consider training DNN using the algorithm during $E =200$ epochs. Then, we have that $T = \lceil \frac{n}{b} \rceil E = 312600$. Although you use $T \to \infty$, the $T = 312600$ never diverges.

I think that the authors should address the following concern that is the third concern in my first review:

• Theorem 1 (Cont.): Although Theorem 1 does not guarantee convergence of Algorithm 1, the authors fixed the number of steps $T$ and proposed the setting of the batch size $n_t = T$, the number of gradient call $N = T^2$, and the learning rate $\eta = \frac{1}{L^{1/2} G^{1/(p-1)} T^{1/2}}$. Then, the authors claimed that the upper bound is $O(N^{- 1/4})$. The setting and the upper bound $O(N^{- 1/4})$ seem to be useful in theory. However, we must verify whether or not the real upper bound

$

U := \frac{1}{N^{1/4}} \left( L^{1/2} G^{1/(p-1)} ( f(\theta_0) - f(\theta^*) + 1/2) + \frac{2p -1}{p-1} G^{\frac{1}{p-1}} \Vert \sigma \Vert_1 \right)

$

becomes sufficiently small (e.g., $U < 10^{-4}$). I think that there is a possibility such that $U$ becomes large (e.g., $U > 10^{2}$). Therefore, I would also like to suggest that the authors could provide sufficient conditions (e.g., practical setting of $\eta$, $n_t$, and $T$) such that $U$ can be small in practical cases such as DNN training.

We thank the reviewer for their comments and suggestions.

However, I do not think that the theorem leads to convergence of Algorithm 1 ... The conditions would be appropriate in theory. However, it would not be good in practice ...

As we noted in Theorem 1 that line 228 that when $\eta = \frac{1}{L^\frac{1}{2}G^\frac{1}{p-1}T^\frac{1}{2}}$, the last term coming from the bias of the coordinate-wise rescaled stochastic gradient also decays with the number of iterations $T$, allowing the bound to converge to $0$ as $T \rightarrow \infty$, achieving the $O(\epsilon^{-4})$ rate as further explained in lines 229-233. The setting of learning rate in the order of $O(\frac{1}{\sqrt{T}})$ is also observed in the convergence analysis of SignSGD by Bernstein et al. (2018) as well as in standard SGD analysis, e.g., Theorem 5.5 in (Garrigos and Gower, 2023).

Reference:

Guillaume Garrigos, Robert M. Gower, Handbook of Convergence Theorems for (Stochastic) Gradient Methods, 2023.

However, we must verify whether or not the real upper bound $U$ becomes sufficiently small (e.g., $U < 10^{-4}$). I think that there is a possibility such that $U$ becomes large ... The above major concerns for Theorem 1 are based on unknown parameters $f(\theta_0) - f(\theta^\ast)$, $G$, $\Vert\sigma\Vert_1$ and $L$.

We would note that in practice, the number of gradient calls $N$ is usually considerably large, as it is the product of batch size, number of epochs, and number of iterations. One can always increase the number of iterations to reach the desired accuracy. Furthermore, there is intrinsically a gap between theory and practice, as the convergence rate and relevant constants are derived in the setting of worst-case analysis, whereas in practice, it's common to observe better empirical performance than that derived in theory, and the parameters need not be set as pessimistic as the worst-case theoretical derivation. In addition, the constants mentioned by the reviewer are standard assumptions in optimization analysis, and the way we set the batch size and learning rate follows from that of SignSGD (Bernstein et al. (2018)).

The proof of Theorem 1: The proof is well-written. However, I have one concern for the proof of Lemma 3 (Appendix A.1 and L269-L284). Given $t$, the authors consider two cases. The cases seem to be incorrect. If (1) is assumed, then (2) should be $\neg (1)$.

We would note that case (2) is indeed $\neg (1)$ in the context of the proof. Note that these two cases in Lemma 3 fall under the condition of sign$\left(\nabla f(\theta_t)^{(i)}\right) =$ sign$\left(g(\theta_t)^{(i)}\right)$, as indicated by the indication function in Lemma 3 (line 917). Note that by definition of $\zeta^{(i)}$ as the variable in the Lagrangian remainder from the Taylor expansion, it takes the value between $\nabla f(\theta_t)^{(i)}$ and $g(\theta_t)^{(i)}$, as we stated in line 940. Since sign$\left(\nabla f(\theta_t)^{(i)}\right) =$ sign$\left(g(\theta_t)^{(i)}\right)$, when $\nabla f(\theta_t)^{(i)} \geq 0$ and $g(\theta_t)^{(i)} \geq 0$, we must have $\zeta^{(i)} \geq 0$ by its definition from the remainder of Taylor expansion. Then we either have $g(\theta_t)^{(i)} \geq \zeta^{(i)} \geq \nabla f(\theta_t)^{(i)}$ which belongs to case (1) or $\nabla f(\theta_t)^{(i)} \geq \zeta^{(i)} \geq g(\theta_t)^{(i)}$ which belongs to case 2. Similarly, when both $\nabla f(\theta_t)^{(i)}$ and $g(\theta_t)^{(i)}$ are negative, we must have $\zeta^{(i)} < 0$, and since $\zeta^{(i)}$ is within the range of $g(\theta_t)^{(i)}$ and $\nabla f(\theta_t)^{(i)}$, their relation belongs to either case (1) or case (2).

Can you explain, how do you choose parameters for the baselines? Why do you choose a smaller lr for Adam against Stacey? Why can not we use the fact that $\ell_p$ norms are equivalent, choose $p$ for the problem based on STACEY experiments (e.g., $p=3$), adjust the parameters for Adam using the relations between $\ell_2$ and $\ell_3$ norms.

We acknowledge the concerns regarding the fairness of comparisons. To address this, we present the ablation results for AdamW on CIFAR-10 in the accompanying table. These results demonstrate that the best performance is not achieved with the largest learning rate or the highest weight decay parameter ($\lambda$). Based on these findings, we selected a learning rate of $0.01$ for AdamW, and a similar process guided the parameter selection for Adam.

STACEY is designed to operate explicitly within the $\ell_p$ framework, leveraging both primal and dual norms to achieve acceleration in non-Euclidean geometries. The choice of $p$ directly aligns with STACEY’s optimization strategy. Adam, while effective, operates in a Euclidean-like adaptive gradient framework. Adjusting Adam’s parameters to match the dynamics of an ℓp-norm optimization method like STACEY would require significant reformulation and tuning, as Adam does not inherently account for non-Euclidean geometries.

We thank the reviewer for their comments and suggestions.

The analysis in Theorem 1 is done under a restrictive bounded gradients assumption. Usually it significantly simplifies the analysis. There are many convergence results for gradient methods without it, e.g., see [1, 2] and references therein. Can the theory be improved by removing the bounded gradients assumption? What is the challenge?

We are aware of the references the reviewer mentioned, and would like to note that they focus on simpler cases, e.g., Euclidean and unbiased settings, and we would like to further justify our assumption of bounded gradient. The $\ell_p$ steepest descent is a generalization of SignSGD and takes a more complicated form as the reviewer acknowledges. The additional challenge for theoretical analysis is twofold. (1) First, for the norm of the stochastic gradient term, that is, term $C$ in line 836 in the proof of Theorem 1, conventional analysis breaks it into mean-squared error and the norm of the true gradient, in which the mean- squared-error coincides with the variance when $p=2$. In our setting, the norm is in $\ell_p$, and the stochastic gradient is coordinate-wise rescaled, leading to a $p$-dependent power on the norm. As a result, we believe it's non-trivial to bound this term by variance. (2) Second, our analysis, due to the extra coordinate-wise rescaling of the stochastic gradient, involves the multiplication of two stochastic gradient terms, e.g., line 957, both depending on the same noise, thus not independent. The expectation of the product of two dependent terms is also hard to characterize based on simple unbiased gradient estimation or other conventional assumptions. Therefore, we introduce the bounded gradient constant to reduce such a product to just one noise-dependent term.

Why the authors analyze only the stochastic descent but not STACEY?

The convergence properties of the method depend on the choice of its parameters, for which our existing results do indeed establish convergence. As we have noted, there exist lower bounds that preclude the possibility of attaining accelerated rates without introducing additional assumptions.

How can one choose an appropriate $p$ for the deep learning problem at hand? Do practitioners need to follow your recommendation of or they need to set it themselves for each particular problem?

We appreciate the reviewer’s insightful question about choosing an appropriate $p$ for deep learning problems and whether practitioners need to follow our recommendations or tune $p$ themselves. Our experiments demonstrate that $p$ can significantly impact performance. For instance:

• On CIFAR, $p \approx 2$ led to superior accuracy.
• For LLM pretraining, $p \approx 3$ offered better convergence.

Could you explain, what is written in Appendix? You provide tables for different methods, but you test only STACEY.

We kindly suggest that there has been a misunderstanding on the part of the reviewer, as we do not test only STACEY.

We thank the reviewer for their comments and suggestions.

The authors propose two variants of their algorithm Stacey, but do not provide any theoretical analysis of the algorithm even for any simplified/toy examples. The algorithm is primarily motivated by the idea of accelerating descent but since there are no convergence or stability results, it is difficult to say whether it actually achieves that goal.

The convergence properties of the method depend on the choice of its parameters, for which our existing results do indeed establish convergence. As we have noted, there exist lower bounds that preclude the possibility of attaining accelerated rates without introducing additional assumptions.

For the ImageNet experiment, Stacey has been compared only with SGD without momentum. Without momentum, SGD is known to converge extremely slow, so it is no surprise that it performs as badly as it does, and there is no other baseline in that experiment.

We thank the reviewer for their comment regarding the baseline comparison in the ImageNet experiment. We performed additional experiments, and the results are presented below.

The table presents the test accuracy (Top-1) of SGD with momentum $0.9$ (SGDM) on ImageNet at different epochs:

We see that STACEY (p,p) also outperforms SGDM consistently across epochs.

The learning rate used for Adam and AdamW ($10^{-4}$) is 100 and 500 times smaller than the learning rate used for the different versions of Stacey ($10^{-2}$ and $5\times 10^{-4}$).

For hyperparameters selection, the learning rates for Adam ($10^{-4}$) and AdamW ($5 \times 10^{-4}$) were selected based on widely-used defaults and prior literature, which ensure stability and optimal performance in these adaptive optimizers. For STACEY, learning rates ($10^{-2}$ and $5 \times 10^{-4}$) were chosen empirically to achieve the best performance in alignment with its novel dynamics in the $\ell_p$ -norm space. Larger learning rates were suitable for STACEY because of its primal-dual interpolation mechanism, which stabilizes updates.

For the CIFAR-10 experiment, some versions of Stacey are used with learning rate 0.1 whereas Adam and AdamW used with learning rates 0.001 and 0.01 respectively. Again, it does not appear to be a fair comparison. Moreover, the value of λ (presumably, weight decay parameter) for Adam and AdamW is set to 2×10−4, for Stacey it is 0.01.

We acknowledge these concerns, which we address thusly. We present the ablation results for AdamW on CIFAR-10 in the accompanying table. These results demonstrate that the best performance is not achieved with the largest learning rate or the highest weight decay parameter ($\lambda$). Based on these findings, we selected a learning rate of $0.01$ for AdamW, and a similar process guided the parameter selection for Adam.

In Theorem 1, what is the justification for choosing the batch size same as the number of iterations (T)? Typically, the batch size would be constant. In practice, the number of iterations might even be much larger than the total number of training samples available. What would a batch size of T mean in such situations?

Since for $\ell_p$ steepest descent, the stochastic gradient update is coordinate-wise rescaled and thus no longer an unbiased estimate of the coordinate-wise rescaled true gradient. An increasing batch size with respect to the number of iterations $T$ is essentially to make such bias, which would otherwise remain as a non-vanishing constant, decay with the number of iterations. We would note that the same batch size setting is applied for SignSGD in Bernstein et al. (2018), without which the $\mc{O}(\epsilon^{-4})$ rate cannot be achieved.

For the CIFAR-10 experiment, Stacey(p,p) has been used with p=2, which means that both of the descent steps reduce to just the usual gradient descent. In this case, in what ways is Stacey(2,2) different from the existing momentum-based gradient descent algorithms?

We acknowledge that, indeed, STACEY($p$,$p$) with $p=2$ reduces both descent steps to standard gradient descent in the Euclidean geometry. However, the various momentum mechanisms (which generalize Lion) mean that it does not reduce to SGD with momentum, and furthermore, STACEY is specifically designed to accommodate non-Euclidean geometries ( $p>2$ ).

In the footnote on page 6 as well lines 370-371, the authors make a distinction between "Nesterov's acceleration" and "momentum" but it is not entirely clear what exactly the distinction is. Can the authors clarify how they define these terms and what do they mean when they say they coincide in the deterministic convex setting?

The distinction was meant to contrast the workings of deterministic vs. stochastic variants of acceleration (see, e.g., [1]). We will clarify this in the updated version of the manuscript.

Can the authors clarify how the hyperparameters for various algorithms were chosen, in particular for baseline algorithms like SGD, Adam, and AdamW? Was a grid search performed for these or were they chosen based on some default values suggested in prior work? How were the values of determined for Stacey in various experiments?

A grid search was performed for the hyperparameters, based on that used in prior work.

[1] Lan, Guanghui. "An optimal method for stochastic composite optimization." Mathematical Programming 133, no. 1 (2012): 365-397.

Why is there no comparison to Lion and Adam in the ImageNet experiment?

We appreciate the reviewer’s feedback, for which we clarify thusly:

• The full ImageNet dataset is exceptionally large, making it computationally challenging to perform extensive hyperparameter tuning within our available resources.
• Following the Lion paper’s hyperparameter tuning guidelines, we have noted that their batch size (1024) is significantly larger than ours (256), due to time and resource limitations.
• The results presented in our paper are based on reasonable hyperparameter settings informed by prior work and our empirical observations.

Was cosine annealing or other schedulers used in the LLM experiment? If not, why?

Yes, we used cosine annealing, and we apologize for the omission of mentioning it in the LLM section.

We thank the reviewer for their comments and suggestions.

The results in Theorem 1, while new, are clearly inspired by the proof of Bernstein et al. (2018) for signSGD. They do require a lot more work since the update is not exactly the sign of stochastic gradient, but unlike the result of Bernstein et al. (2018), the new one also requires the gradients to be bounded.

We thank the reviewer for acknowledging our new results and would like to justify our assumption of bounded gradient. The $\ell_p$ steepest descent is a generalization of SignSGD and takes a more complicated form as the reviewer acknowledges. The additional challenge for theoretical analysis is twofold. (1) First, for the norm of the stochastic gradient term, that is, term $C$ in line 836 in the proof of Theorem 1, conventional analysis breaks it into mean-squared error and the norm of the true gradient, in which the mean-squared error coincides with the variance when $p=2$. In our setting, the norm is in $\ell_p$, and the stochastic gradient is coordinate-wise rescaled, leading to a $p$-dependent power on the norm. As a result, we believe it's non-trivial to bound this term by variance. (2) Second, our analysis, due to the extra coordinate-wise rescaling of the stochastic gradient, involves the multiplication of two stochastic gradient terms, e.g., line 957, both depending on the same noise, thus not independent. The expectation of the product of two dependent terms is also hard to characterize based on simple unbiased gradient estimation or other conventional assumptions. Therefore, we introduce the bounded gradient constant to reduce such a product to just one noise-dependent term.

We would also kindly note that when going from SGD to signSGD, the proof of Bernstein et al. (2018) also involved non-trivial assumptions like their Assumption 2 and Assumption 3, which are stronger than the conventional smoothness assumption and variance assumption. Therefore, we believe it's reasonable to introduce new assumptions when tackling more general and complicated settings.

In addition, we believe that the extra complication in the proof for handling the more general and complicated form of $\ell_p$ steepest descent, e.g., the proof of Lemma 3, highlights the technical novelty of this paper.

There is no convergence theory for the two accelerated methods presented in the paper. Even though the authors note that we shouldn't expect acceleration when the variance dominates the convergence rates, it would still be useful to know what convergence properties the methods have.

The convergence properties of the method depend on the choice of its parameters, for which our existing results do indeed establish convergence. As we have noted (and as acknowledged by the reviewer), there exist lower bounds that preclude the possibility of attaining accelerated rates without introducing additional assumptions.

What we are left with in terms of contributions is the empirical comparisons, but they are too limited to be the main contribution. In some experiments, the authors do not compare to standard baselines, and the LLM experiment is done only for the first 5000 iterations.

We emphasize that our empirical results are designed to complement our theoretical contributions, particularly the novel convergence guarantees for $\ell_p$-norm-based optimization in stochastic non-convex settings and the accelerated STACEY algorithm.

While we agree that empirical results alone may not be sufficient as the primary contribution, we believe they strongly validate the practical implications of our theoretical advancements.

Some experimental details appear to be missing, for instance, I didn't find a description of the scheduler used in the LLM experiment. I also wish the experiments were done and reported with multiple random seeds. I think it would be great if the authors added a section in the appendix with a full list of hyperparameters used in the experiments.

We appreciate the reviewers’ detailed feedback and constructive suggestions regarding the experimental setup. Below, we address these specific concerns:

• In the LLM experiments, we employed a cosine learning rate schedule with warm-up steps proportional to 10% of the total iterations. This detail, along with any additional hyperparameter settings, will be explicitly included in a new appendix section in the revised manuscript.
• We will provide a comprehensive table in the appendix summarizing all hyperparameters for each experiment, including learning rates, batch sizes, momentum terms, and optimizer-specific parameters.