Thank you for your detailed review. It is a pleasure to have an expert reviewer. Below we address each of your questions.

Q1 mistakes …in Table 1

The rates we report throughout are the anytime rates of each method. We do not report the fixed horizon rate. We will now include the following exert in our revision:

We say $x_t$ has an anytime convergence rate if, once parameters are fixed, there exists a real function $r:[0,+\infty) \to [0,+\infty)$ such that $\mathbb{E}[f(x_t) - f(x)] \leq r(t)$ and $r(t) \to 0$ as $t \to + \infty$.

We say an algorithm has a finite horizon rate if, for every tolerance $\epsilon>0$, there exists a choice of parameters and there exists $T \geq 1$ such that $\mathbb{E}[f(x_T) - f(x)] \leq \epsilon$.

We focus on anytime rates because these are the rates that best translate into practical performance. Where-as fixed horizon rates depend on very small step sizes, and are completely impractical when implemented.

For example, SGD with a constant stepsize does not have an anytime convergence rate, but it has a finite horizon rate. The anytime rate for SGD on convex smooth problems is $O(\log(t)/\sqrt{t})$ (Theorem 5.7 in [Garrigos & Gower, 2023]) and the fixed horizon rate is $O(1/\sqrt{t})$ (Theorem 5.5 in [Garrigos & Gower, 2023]) which is the rate you pointed out.

Q2 Rates for SPSmax and NGN … not invariant under function scaling

SPSmax and NGN do depend on the function scaling. If we turn to eq (4) for SPSmax, it is not invariant because the constant $\gamma_b$ does not scale with the function. Where as the learning rate of NGN is given by $\gamma_k = \frac{\sigma f(x_k)}{f(x_k) + \sigma \| \nabla f(x_k)\|^2}$ The $\| \nabla f(x_k)\|^2$ scales quadratically, whereas $f(x_k)$ scales linearly with $f$. Thus NGN cannot be invariant to scaling the function.

Q3 $\sigma_{pos}$ is likely equal to $\sigma_{*}^2$

These two constants are not equivalent, instead $\sigma_{*}^2 = f(x_{\*}) - \sigma_{pos}.$ In the setting where $\sigma_{pos} =0$ (regression models with no weight decay), we still could have $\sigma_*^2 \neq 0$

Q4 novelty compared to previous SPS* results

The novelty of our analysis of SPS* is in the smooth setting. Outside of the interpolation setting, we give the first convergence for SPS* in the smooth setting. Previous results in Gower et al., 2021 and Garrigos et al., 2023 only establish that SPS* converges to a fixed neighborhood of the solution.

Q5 impact of mini-batching

To see the dependency on mini-batching, we change the expectation in eq (9) to be over the mini-batching, instead of directly over $\xi$. Then $A$ and $B$ will now depend on the batch size $b$. For example, if we use sampling without replacement in the smooth $n$--finite sum setting, then according to Proposition 3.8, item (iv) [Gower et al 2019] the constants $A$ and $B$ in eq (9) have the following form

$A = \frac{n}{b}\frac{b-1}{n-1} L_{full} + \frac{n-b}{b(n-1)} L_{\max}$ $B = \frac{1}{b}\frac{n-b}{n-1} \sigma^2_*$

Where $L_{full}$ and $L_{\max}$ are the full batch and max over stochastic functions smoothness constant, respectively.

For the non-smooth setting, the right hand side in (11) would be $G^2/b$ if we used $b$--sampling with replacement.

Q6 inconsistent definitions of $\sigma_*$ throughout

Thank you and well spotted. We will no be consistent and use function noise $\sigma_f^2 =f(x_*) - \mathbb{E}{\inf f_{\xi}}$ throughout.

Q7 Table 1 uses function suboptimality and distance to the solution.

We will convert all results to function suboptimality, at the cost of an additional fixed $1/\mu$ multiplicative factor for the strongly convex problems.

Q8 Missing references

Thank you for the two additional references, we will include them both, and comments about their relative smoothness results.

Q9 Table 1 SGD has "X" in the first column? Constrain to the ball

We agree, but this would require access to the ball $B_D$ that contains $x_*$, so that the iterates $x_k$ can be projected back onto $B_D$. But this is a different oracle access. Depending on the problem, we would not have access to this ball. But we can include a footnote on this.

Q10 Why not define $\sigma_*$ in (13)

We use Corollary 2.3 to give a precise setting, and Theorem 2.1 and eq (9) for the general setting. But we do agree that Corollary 2.3 also holds when using gradient noise given by $\sigma_* = L \mathbb{E}\|g_{\xi}(x_*)\|^2$.

Q11 Why introduce Bregman div in (20)

We introduce the Bregman div because it is a positive term that provides a minor speed-up of IAM over SPS*.

Q12 What if $B=0$ in Theorem G.1

Good question, the result in Theorem G.1 does not hold for $B=0.$ Instead when $B=0$, from equation (49) we can prove that SPS* converges linearly according to $\mathbb{E}\\| x_{t+1}-x_*\\|^2 \leq (1-\mu/(2A))\mathbb{E}\\|x_{t}-x_*\\|^2$ We will edit the theorem to include this remark.

Dear Reviewer,

thank you for your thoughtful review, and doing the expert work of checking our proofs and appendix, and catching the type-O on Lines 307-308. Regarding the weakness you highlighted

It is not clear that the theory is relevant for the experiments in Figure 1 as the setting is highly non-convex. However, this is a common issue faced by papers in convex optimization theory and is not the focus of the paper.

We do of course agree. But we have written a small paragraph “Defense of convex optimization in machine learning” in response to reviewer UUoW rejection of our work because of the convexity assumption. We copy the relevant parts of our response here for your consideration, though it does not truly address your very valid point.

Defense of convex optimization in machine learning.

Stochastic non-smooth convex optimization has a surprising impact on deep learning practice.

For instance, the development of the Adagrad method [Duchi2011] was based on its convergence analysis in the non-smooth convex setting. Adagrad went on to be widely used in Deep learning, until the development of exponentiated weighted variants of Adagrad such as RMSprop and ADAM. Thus non-smooth convex analysis has a significant role to play in inspiring the most popular optimization ADAM for deep learning. Much more recently, the Schedule-free method [Defazio2024] which was developed to have fast last-iterate convergence in the non-smooth convex setting, set a new record in the AlgoPerf competition in the self-tuning category:

https://mlcommons.org/2024/08/mlc-algoperf-benchmark-competition/

The fact that non-smooth convex analysis is a reasonable model for predicting performance on Deep learning, including training large language models, is now a well-established conundrum [Schaipp2025]. Where-as, the convergence analysis of SGD in the non-convex setting has not had such an impact on Deep learning practice. As such, convex optimization has a surprisingly important and active role to play in deep learning practice.

[Duchi2011] Duchi, John, Elad Hazan, and Yoram Singer. "Adaptive Subgradient Methods for Online Learning and Stochastic Optimization." Journal of Machine Learning Research 12 (7): 2121–2159, 2011

[Defazio2024] Aaron Defazio, Xingyu Alice Yang, Ahmed Khaled, Konstantin Mishchenko, Harsh Mehta, Ashok Cutkosky, “The Road Less Scheduled”. Neurips Oral 2024

[Schaipp2024] Fabian Schaipp, Alexander Hägele, Adrien Taylor, Umut Simsekli, Francis Bach, “The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training”, arXiv:2501.18965, 2025

Dear reviewer, the main issue you have raised is regarding our assumption in our theorems that the loss is convex. We will divide our answer into two parts. In our first part, we will show how to relax global convexity to a form of local convexity. Indeed, this extension follows straightforwardly by using that SPS* and IAM have the very unusual property of being monotonic (Eq (8) and Lemma 3.1). In our second part, we write a defense for convex optimization in machine learning. After which we answer a minor point below.

We hope the reviewer will reconsider their position in light of our response.

Extending beyond convexity.
Our theorems do not require global convexity. Instead, we need only assume that for the given starting iterate $x_0 \in \mathbb{R}^d$, that the $f_\xi$ functions are almost surely convex within the closed ball

Using your notation, we need convexity in an $\delta$-neighborhood around $x^*$ where $\delta = \Vert x_0 - x^* \Vert .$

The reason we are able to relax the global convexity assumption is because SPS* and IAM have the very unusual property (amongst stochastic methods) of being monotonic. That is, for the iterates $x_t$ (or $z_t$ for IAM), we have that

The same cannot be said about SGD or any variant of SGD that we are aware of, since the iterates may escape the closed ball with some probability. We will include a remark on this in our revision.

Defense of convex optimization in machine learning.
As a mere formality, we first point out that convex optimization is one of the official topics of interest listed in ICML’s 2025 call for papers:
https://icml.cc/Conferences/2025/CallForPapers

As such, we do not believe it is a valid reason to recommend rejecting our paper because our analysis focuses on stochastic convex optimization. Moreover, stochastic non-smooth convex optimization has had, and continues to have a surprising impact on deep learning practice.

For instance:

• Adagrad [1] was developed based on its favourable non-smooth convex anytime convergence. Adagrad went on to be widely used in Deep learning, until the development of exponentiated weighted variants of Adagrad such as RMSprop and ADAM. Thus non-smooth convex analysis has a significant role to play in inspiring the most popular optimization method (ADAM) for deep learning

• The Schedule-free method [2], designed based on its fast last iterate convergence for non-smooth convex settings, just recently set a record and one first place in the AlgoPerf competition in the self-tuning category for deep learning (https://mlcommons.org/2024/08/mlc-algoperf-benchmark-competition/).

The fact that non-smooth convex analysis has some predictive power on the performance in deep learning (including LLMs) is now a well-established conundrum [3]. In contrast, non-convex SGD analysis has had far less practical impact on deep learning. Thus, convex optimization remains relevant to deep learning practice, justifying its inclusion in ICML’s topics of interest.

Q1. Do the objective functions considered in Section 4 satisfy the assumptions?

A1. The Poisson regression experiments mentioned in Section 4 (and relegated to Appendix L.1 due to space constraints), is a globally convex problem. In fact, this is one problem for which SPS* and IAM provenly converge (see Remark 2.4), and for which we are unaware of any other method that converges. This experiment was to highlight the predictive power of our theory. As for the Black-box Model Distillation in Section 4.1, the student models are decoder only transformers, as such, they are non-convex deep neural networks. Thus our theorems do not explicitly hold for these problems.

References

[1]: Duchi, J., Hazan, E., & Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12, 2121–2159, 2011

[2]: Aaron Defazio, Xingyu Alice Yang, Ahmed Khaled, Konstantin Mishchenko, Harsh Mehta, Ashok Cutkosky. The Road Less Scheduled, Oral at NeurIPS 2024.

[3]: Fabian Schaipp, Alexander Hägele, Adrien Taylor, Umut Simsekli, Francis Bach. The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training, arXiv:2501.18965, 2025

We thank the reviewer for their thoughtful review, and appreciating our work.

Q1. It would be helpful if further discussions are provided on how to approximate the optimal loss value in general settings, and how the convergence analysis would be affected by the approximation error.

A1. This is a good question. We are currently working on a follow-up on how the optimal loss value could be approximated on language models. The optimal cross-entropy loss over the training data is lower bounded by the average number of bits used to compress this data when using an optimal compression method. The final loss over a batch is then given by the average number of bits used to compress this batch. Though we do not have access to an optimal compressor, compressing is in general an easier task than predicting the next token, and thus, we do have access to very good compressors. This would suggest the following approximation: Compress the training data using a good compressor $\implies$ Use the average number of bits used to encode a batch of data as an approximation to the optimal batch loss $\implies$ Run IAM with this approximation. We can consider adding this to potential future work, if the reviewer thinks it is insightful or partially answers their question. We are also considering how to extend our convergence analysis to allow for an approximation error. This has proved challenging, and we cannot guarantee that our revision will include results on this.

We thank the reviewer for their comprehensive review. We have divided our answers in Main Questions, and Questions

MQ1 anytime convergence formally?

We say $x_t$ has an anytime convergence rate if, once parameters are fixed, there exists a real function $r:[0,+\infty) \to [0,+\infty)$ such that $\mathbb{E}[f(x_t) - f(x)] \leq r(t)$ and $r(t) \to 0$ as $t \to + \infty$.

We say an algorithm has a finite horizon convergence if, for every tolerance $\epsilon>0$, there exists a choice of parameters and there exists $T \geq 1$ such that $\mathbb{E}[f(x_T) - f(x)] \leq \epsilon$.

If an algorithm has an anytime convergence rate then it has a finite time convergence rate, but the reciprocal is not true in general. The counterexample is SGD with constant stepsize, which does not have an anytime convergence rate, but has a finite time convergence.

MQ2 main contribution is (SPS*) anytime rate?

Our main contributions are the anytime rates for SPS* in the smooth (convex+ str convex) setting, proposing the IAM method, all the last iterate analysis of IAM, and identifying a reasonable application (black box model distillation).

MQ3 Difference between bounded domain and line 270

The bounded domain assumes we know a bounded set $\mathcal{B}\subset \mathbb{R}^d$ such that $x_* \in \mathcal{B}$. If we know such a set, we can modify SGD to project $x_t$ onto $\mathcal{B}$.

Our assumption on line 270 and Eq (9), is very different and is almost for free. We require that the gradients verify some inequality on a certain ball, that we do not need to know. As we emphasize in Remark 2.4, for finite sum problems our assumption is always true, because convex functions always have bounded gradients on every bounded set.

MQ4 The teacher should be close to the minimum? Related to online imitation learning?

Yes, the teacher's loss is assumed to be close to the minimum student loss, and thus we can use the teacher loss to set the learning rate of the student using the IAM. There is some relation to the online imitation learning setting, both involve a student learning from a teacher, but there are differences. In online imitation learning the student learns a policy, and has to interact with an environment. Where-as in our setting there is no environment or policy.

MQ5 Unclear "access to $f_{\xi}(x^*)$ not the same as interpolation"

This is a good question. In short, assuming access to $f_{\xi}(x^*)$ imposes no restriction on the problem itself. This could be any stochastic optimization problem. Rather, this is an assumption as to what we know regarding the problem, and what we have access to.

The interpolation assumption does impose a constraint on what type of problems we are considering. Interpolation restricts our problem to models that can perfectly fit every single data point. This can (approximately) occur in image classification tasks, where the neural network is sufficiently overparameterized with respect to the data set. This cannot occur when training a generative language model, because it is impossible to perfectly guess the next token in a sequence due to ambiguities in languages. Because of this, the cross-entropy loss is never zero, and is always above the entropy of language (around 2.0). Thus the interpolation assumption cannot hold in generative language modelling. In contrast, assuming access to $f_{\xi}(x^*)$ imposes no restriction to the model itself, and it can occur in setting such a generative language modelling. For instance, such as out black-box model distillation problem

Q1 Comparisons to Adagrad

In the non-smooth setting, Adagrad has a guaranteed convergence of $\sqrt{2}GD/\sqrt{T}$, which is $\sqrt 2$ worse than our rates in for SPS* and IAM in Eq (12) and (20). The difference between this rate of Adagrad and our rates for IAM, is this rate holds for the last iterate of IAM and Adagrad assumes access to a bounded domain $D$ which contains the solution. We can include numerical comparisons to Adagrad, but AdamW is SOTA for LLMs, not Adagrad, which perform poorly for LLMs.

Q2 Comparisons SLS variant, SSO

Neither SLS, or the SSO variant, are applicable in our black-box model distillation setting. SLS assumes that we can use the teacher to generate a candidate sequence, where-as SSO assumes we have access to the logits of the teacher. But our black-box setting assumes we only have access to the final loss of the teacher, and not the logits.

Q2 Discussion on computational overhead

Certainly, SPS* and IAM impose a small additional computational overhead as compared to SGD and SGD-M. Both SPS* and IAM need only store one additional scalar, and do two inner products with respect to the gradient. We will add this remark to the appendix.

Q3 Random seeds

We use $42$ to set all random seeds.

Q4 ​​No failure modes

We investigate failure modes regarding misspecification of the optimal loss values in Section L2.