Armijo Line-search Can Make (Stochastic) Gradient Descent Provably Faster

We thank the reviewer for their helpful feedback and address their concerns below.

(1) My major concern with this paper is that it doesn't have a related work section, and seems to only cite very old papers concerning line searches, as well as a nearly 20 year old text book...Please include this in your rebuttal and add it to the final camera ready.

Even though we do not have an explicit related work section, we believe that we have cited the most relevant work. This includes recent work on Armijo line-search (Vaswani et al., 2019a; Galli et al., 2024; Hubler et al., 2024), Polyak step-sizes (Loizou et al., 2021) and variants of normalized gradient descent (Mei et al., 2021; Axiotis & Sviridenko, 2023; Taheri & Thrampoulidis, 2023). Moreover, we have added the most relevant and recent references for each example – logistic regression (Axiotis & Sviridenko, 2023; Freund et al., 2018; Wu et al., 2024), generalized linear models (Mei et al., 2020; Hazan et al., 2015), softmax policy gradient (Mei et al., 2020; Asad et al., 2024; Lu et al., 2024), two layer neural networks (Taheri & Thrampoulidis, 2023). In the final version of the paper, we will add the following extended comparison and the corresponding references.

Comparison to GD with adaptive step-sizes on uniform smooth functions: There has been substantial work on designing adaptive step-sizes for GD. However, most of this literature (for example (Orabona & Pal, 2016; Malitsky & Mishchenko, 2019; Carmon & Hinder, 2022; Khaled et al., 2023)) considers either the uniform-smooth, convex or non-smooth but Lipschitz, convex settings. In the uniform-smooth, convex case, the focus of these papers is to design more efficient ways to adapt to the smoothness constant, and the resulting methods achieve the same rate as GD with constant step-size. In contrast, we consider Armijo line-search which is the classic way to set the step-size for GD on smooth functions. The focus of our work is to identify when $`GD-LS`$ can be provably faster than $`GD(1/L)`$. To that end, we have identified a class of non-uniform smooth functions and shown that numerous examples in machine learning satisfy these properties. We believe that we are the first to identify such a broad range of examples and demonstrate the provable improvement of $`GD-LS`$ over $`GD(1/L)`$.

Comparison to GD on non-uniform smooth functions: There has been recent work on analyzing GD (with/without clipping) on functions satisfying (L0, L1) non-uniform smoothness (Zhang et al., 2019; 2020; Koloskova et al., 2023). This work requires the knowledge of L0 and L1. In contrast, as explained in Contribution 1 and in lines 133-142 (after Proposition 1) of our paper, the proposed non-uniform smoothness condition is different and the $`GD-LS`$ algorithm does not require the knowledge of L0, L1 or other problem-dependent constants.

Comparison to $`GD-LS`$ on non-uniform smooth functions: The most relevant paper to our setup is (Hubler et al., 2024) and we have compared to this in Lines 58-64. In particular, this paper analyzes the convergence of $`GD-LS`$ on functions satisfying a different notion of (L0, L1) non-uniform smoothness. However, their algorithm requires the knowledge of L1 and the resulting rate is the same as that of $`GD(1/L)`$

We hope that this addresses the reviewer’s concerns and helps them better contextualize our contributions. If the reviewer is aware of a key reference we have missed, we would be happy to include and compare to it.

(2) I object to the use of stochastic to refer to analysis under interpolation condition... The title of the paper is misleading and would suggest to readers that there is some novelty in the paper that addresses the full stochastic setting.

As is evidenced by the numerous papers analyzing SGD under interpolation, we first note that the interpolation setting is an important special case. Moreover, this setting can serve as a building block towards analyzing the full stochastic setting. For example, the stochastic line-search (Vaswani et al., 2019b) and stochastic Polyak step-sizes (Loizou et al., 2021) were originally designed for the interpolation setting, and have been adapted to the full stochastic case by combining these techniques with appropriately decreasing step-sizes (Orvieto et al., 2022; Vaswani et al., 2022; Jiang & Stich, 2023). Moreover, the proof for $`SGD-SLS`$ in Section 6 is more involved compared to the standard SGD proofs. Given this, we think it is important to highlight that our paper also considers the stochastic setting, and hence we added "stochastic” to the title. Note that we clarify that we are only considering the interpolation setting in the abstract

We do understand the reviewer's point, and will change the title to Armijo Line-search Can Make (Stochastic) Gradient Descent Go Fast to suggest that the speedup is for some special settings.

We thank the reviewer for their helpful feedback and address their concerns below.

(1) The targeted optimization problem in this paper is limited in terms of application perspective. The authors primarily focus on logistic regression and generalized linear models. Although the authors showed that object for softmax policy gradient also satisfies the assumptions in appendix, the application sides of this paper to real-world use cases are not significant enough for a venue like ICML.

It is not required for papers in ICML to have "application sides...to real-world use cases''. In fact, the call for papers lists "Theory of Machine Learning (statistical learning theory, bandits, game theory, decision theory, etc.)'' explicitly as a topic of interest.

Having said that, we believe that our paper does have practical implications. In particular, our paper analyzes when $`GD-LS`$ (a standard optimization method widely used in practice) can be provably faster than $`GD(1/L)`$. Besides achieving adaptivity to the smoothness in a principled manner and hence requiring minimal hyper-parameter tuning (an important consideration from a practical standpoint), $`GD-LS`$ can result in faster convergence on numerous convex (e.g. logistic regression) and non-convex problems (e.g. generalized linear models, 2 layer neural networks, softmax policy gradient) important in machine learning. We have argued that $`GD-LS`$ can thus replace more specialized algorithms tailored for these specific problem settings. Having a single algorithm that can work well across different applications is important from a practical standpoint.

Furthermore, as evidenced by the experimental results in [Vaswani et al, 2019], $`SGD-SLS`$ (analyzed in Section 6 of our paper) results in good empirical performance on convex problems that satisfy our assumptions. Under the interpolation assumption, Vaswani et al, 2019 can only prove an $O(1/\epsilon)$ rate for $`SGD-SLS`$ on smooth convex losses such as logistic regression. However, the empirically, the convergence seems to be faster than $O(1/\epsilon)$ (for example, see the linear convergence on the $`mushrooms`$ dataset in Figure 3 of [Vaswani et al, 2019]. Our paper (specifically Corollary 5) can explain this fast convergence. We believe that it is important to understand the behaviour of optimization algorithms, and that these insights can often lead to developing better methods in practice.

Finally, it is important to note that this same method - $`SGD-SLS`$ and its non-monotone variants [Galli et al, 2024] result in strong empirical performance on non-convex losses with deep neural networks [Vaswani et al, 2019, Galli et al 2024] where our assumptions are not necessarily satisfied. This demonstrates that developing methods that are theoretically principled in simple settings can result in empirical gains for real-world use cases.

(2) One potential improvement for this is that the author may list one or two informal version of their main theorem and provide a clear and succinct description of their contributions.

Thank you for the good suggestion. We will consider making this change in the final version of the paper.

(3) Moving the explanation on taking $\eta_{\max} = \infty$ before Theorem 1 will minimize the confusion.

We will make the change.

We thank the reviewer for their helpful feedback and address their concerns below. For all unaddressed comments, we agree with the reviewer's suggestion and will make the corresponding change.

(1) Lower-bounds for $`GD(1/L)`$

For logistic regression, Theorem 3 in Wu et al, 2024 shows that $`GD(1/L)`$ (or more generally, GD with any constant step-size that guarantees monotonic descent) cannot have a convergence rate faster than $\Omega(1/\epsilon)$. Hence, both $`GD-LS`$ and $`SGD-SLS`$ on logistic regression are provably faster than their constant step-size analogs.

For the softmax policy gradient in Section 5.1, Theorems 9, 10 in Mei et al, 2020 show that GD with any constant step-size cannot achieve a rate faster than $\Omega(1/\epsilon)$ on both bandit and MDP problems. We mention this after Corollary 4 in the paper. For the GLM problem in Section 5.2, we are not aware of a lower bound for $`GD(1/L)`$ or normalized GD. We will include these comparisons.

(2) Softmax PG - Setup

We note that the linear convergence result in Corollary 3 also applies to the general MDP problem. See Proposition 6 for the problem setting. $`GD-LS`$ in Section 5.1 is equivalent to softmax policy gradient in the "deterministic'' or "exact'' setting. This setting is commonly used as a testbed for analyzing policy gradient algorithms [Mei et al 2020; Section 4, Lu et al 2024]. For general MDPs, the deterministic setting corresponds to knowing the rewards and transition probabilities, and is the same setting under which classic RL algorithms such as value iteration or policy iteration are analyzed. This was our motivation to consider the problem setup in Section 5.1.

(3) Softmax PG - Handling stochasticity

For the general MDP problem, in cases where the rewards/transition probabilities are unknown, the policy gradient is estimated by interacting with the environment. Standard policy gradient results (e.g [Theorem 29, Agarwal et al, 2020] can then prove convergence up to an error incurred because of the estimation.

For the specific case of stochastic multi-armed bandits with unknown rewards, recent work [Mei et al, 2023] has proven stronger global convergence results. In particular, Mei et al, 2023 use importance sampling to construct the stochastic softmax policy gradient and use it with constant step-size SGD. The resulting algorithm can be proven to converge to the optimal arm. From Proposition 3, we know that the objective function satisfies Assumptions 2, 3, and 4. It also satisfies the interpolation condition [Theorem 4.2 in Mei et al, 2023], required for the $`SGD-SLS`$ analysis in Section 6. Hence, we believe that a variant of $`SGD-SLS`$ could be a good algorithm for the stochastic multi-armed bandit problem. However, since the softmax policy objective is non-convex, we cannot directly use the results in Section 6 and leave this interesting direction to future work.

(4) ....not clear if gains from GD-LS would extend to more typical problems...

We have tried to argue that losses that satisfy our non-uniform smoothness assumptions are in fact, not as rare. To illustrate this, we have given numerous examples ranging from classification in supervised learning to policy gradient in reinforcement learning. In general, applications where we use losses with an exponential tail (e.g. exponential, logistic loss) or those that use a softmax function to parameterize probabilities can benefit from using a line-search. Another example that we did not study in detail is noise contrastive estimation [Lu et al 2021], where GD is slow because of the flatness of the landscape, and $`GD-LS`$ can potentially improve the convergence.

We agree that large gains from $`GD-LS`$ are not always possible (e.g. for quadratics) and if the function is only uniformly smooth, $`GD-LS`$ will converge at the same rate as $`GD(1/L)`$. However, even in this case, $`GD-LS`$ enables setting the step-size without estimating or conservatively bounding the smoothness constant and thus has a practical benefit.

(5) Comparison to Polyak step-size

The proof of Theorem 1 (and non-convex results) relies on the monotonic decrease in the function values guaranteed by $`GD-LS`$. GD with the Polyak step-size is not a monotonic descent method, and it is unclear how to extend the proofs.

(6) Experimental Evaluation

See Point (4) in our response to Rev. jUfU

(7) Comparison to recent work on adaptive step-sizes

See Point (1) in our response to Rev. 8uuj

(8) Cor. 2 vs Cor. 6:

Cor. 2 is for $`GD-LS`$, Cor. 6 is for GD with the Polyak step-size.

(9) Proposition 3: Def. of reward gap

It is the difference between the rewards of the best (optimal) and second-best arm [see Lemma 17 in Mei et al, 2020].

(10) Assumption 2 (a) and (b)

These are separate (but related) conditions that need to both be satisfied.

We thank the reviewer for their helpful feedback and address their concerns below.

(1) Independence of Assumptions

Assumptions 3 and 5 do indeed imply Assumption 2(b). The reason we did not use Assumption 5 as our main assumption is because Assumption 5 cannot be satisfied for logistic regression, multi-class classification with $L_c = 0$. We explained this in Lines 137-142, and note that having $L_c = 0$ is essential to derive a linear convergence rate for logistic regression in Corollary 2. To see this from a technical perspective, consider $f$ to be a finite sum $f(\theta) = \frac{1}{n}\sum_i f_i(\theta)$ (for example, the logistic regression loss in Proposition 1). If $f_i$ satisfies Assumption 5 with constants $L_c, L_g$, it does necessarily imply that $f$ also satisfies Assumption 5 with the same constants (see Proposition 7 for a counter-example).

In contrast, using Assumption 2 as our main assumption enables us to prove that the losses corresponding to logistic regression, multi-class classification satisfy this assumption with $L_0 = 0$, and this enables us to prove the linear convergence rate in Corollary 2. Again, from a technical perspective, we can do this because if $f_i$ satisfies Assumption 2 with constants $L_0$ and $L_1$, then $f$ also satisfies this assumption with the same constants (please refer to the proof of Lemma 5).

In summary, we only use Assumption 5 for $f_i$ (and not $f$) corresponding to common finite-sum losses as a way (using Lemma 5) to show that $f$ satisfies Assumption 2. For example, see the proof of Proposition 1, where we instantiate this logic for logistic regression.

(2) Positivity Assumption (Assumption 1)

From an algorithmic perspective, Assumption 1 is quite benign. As the reviewer suggests, it is possible to shift $f$ by subtracting $f^*$ and ensure that the resulting function $h(\theta) := f(\theta) - f^\ast$ is always non-negative. Since $f^\ast$ is a constant, $\nabla h(\theta) = \nabla f(\theta)$, meaning that we do not require the knowledge of $f^\ast$ to implement GD. Moreover, the Armijo condition on $h(\theta)$ has $f^\ast$ on both sides, implying that implementing $`GD-LS`$ on the non-negative function $h(\theta)$ does not require the knowledge of $f^*$. We use this property in Section 5.1 (see Proposition 3 and the corresponding discussion). Finally, we note that most common losses in machine learning (cross-entropy, exponential, and squared) are non-negative.

(3) Missing Reference

Thank you for directing us to this paper. The assumption in their paper is more restrictive than ours, as they assume strong convexity, whereas we only rely on gradient domination. This enables us to consider non-convex losses such as generalized linear models. Additionally, they focus on exact line search, while we consider Armijo line search. We will cite this paper in the final version.

(4) Experimental Validation

In the deterministic setting (with access to full gradients), $`GD-LS`$ is a standard algorithm for (non)-convex minimization and is widely used in practice. In the stochastic setting, Vaswani et al. (2019b); Galli et al. (2024) have empirically evaluated the performance of $`SGD-SLS`$ and its non-monotone variant for convex losses corresponding to linear and logistic regression and non-convex losses with deep neural networks. These supervised learning experiments have demonstrated the efficacy of $`SGD-SLS`$. As we explain in Section 6, under the interpolation assumption, Vaswani et al. (2019a) can only prove an O(1/ϵ) rate for SGD-SLS on smooth convex losses such as logistic regression. However, empirically, the convergence seems to be faster than O(1/ϵ) (for example, see the linear convergence on the mushrooms dataset in Vaswani et al. (2019b, Figure 3). Our paper (specifically Corollary 5) can explain this fast convergence.

For the softmax policy optimization problem in RL, we note that Lu et al. (2024, Figure 1) have shown the linear convergence of $`GD-LS`$ on tabular Markov decision processes. However, they could only prove an O(1/ϵ) convergence rate for the resulting algorithm (Lu et al., 2024, Theorem 1). In order to attain a linear convergence rate in theory, they propose a non-standard line-search scheme that requires the knowledge of the optimal value function. On the other hand, Corollary 4 in our paper proves a linear convergence rate for $`GD-LS`$ (the standard line-search) on bandit problems, and the same proof extends to the MDP setting (using Proposition 6 in our paper).

We will include these comparisons and explanations in the final version of the paper.