Summary
This paper studies non-convex stochastic optimization under the assumption of mean-square smoothness. It introduces, Ada-STORM, a variant of the STORM algorithm, which achieves the optimal rate $O(T^{-1/3})$. Unlike vanilla STORM, Ada-STORM eliminates the $O(\log T)$ factor and does not require the Lipschitz assumption. The algorithm's success hinges on a carefully designed adaptive learning rate schedule. Instead of the standard AdaGrad learning rate $\eta_t \approx (\sum_{i=1}^t \\|g_i\\|^2)^{-1/3}$ used in STORM, Ada-STORM employs $\eta_t = \min\\{T^{-1/3}, T^{-(1-\alpha)/3}(\sum_{i=1}^t \\|v_i\\|^2)^{-\alpha}\\}$ where $\alpha < 1/3$ and $v_t$ is the STORM update.
Strengths
The proposed algorithm, Ada-STORM, achieves the optimal rate $O(T^{-1/3})$ without the log factor. Unlike STORM+ which also eliminates the log factor, this algorithm does not assume $f$ to be Lipschitz and bounded above. This improvement represents a significant contribution. Moreover, the required modification is straightforward: it involves merely adjusting the learning rate scheduler $\eta_t$ and the momentum factor $\beta_t$. I particularly enjoy the constant momentum constant $\beta_t\equiv T^{-2/3}$ rather than time-varying and dependent on $\eta_t$.
In addition, the proposed learning rate scheduler is of interest in its own right. The paper extends the application of this learning rate technique to related problems, achieving optimal rates in composite optimization and finite-sum optimization (ERM). Finally, empirical experiments also validate Ada-STORM's performance, which matches or outperforms existing STORM variants and other state-of-the-art optimizers.
Weaknesses
One major concern pertains to the technical correctness of the proof of Theorem 1, in particular from line 440-444. I don't immediately see how $\mathbb{E}[\sum_{t=1}^{s-1} \\|v_t\\|^2] \le C_0T^{1/3}$ implies $\mathbb{E}[\frac{1}{T}\sum_{t=1}^{s-1}\\|v_t\\|] \le \sqrt{C_0}T^{-1/3}$, and similarly how the equation in line 441 implies line 442. My concern arise from the fact that $s$ is a random variable dependent on the entire history $\mathcal{H}_T$. As a result, applying Jensen's inequality like $\mathbb{E}[\frac{1}{T}\sum\_{t=1}^s\\|v_t\\|] \le \sqrt{\mathbb{E}[\frac{1}{T}\sum\_{t=1}^s\\|v_t\\|^2]}$ seems to be inappropriate when $s$ is a random variable dependent $v_t$'s. It's likely that I miss something, and I encourage the authors to elaborate this part of the proof to make it clear.
Another concern is about the doubling trick. In each stage $k$, the guarantee is
$$\frac{1}{2^{k-1}} \sum_{t=2^{k-1}}^{2^k} \mathbb{E}\\|\nabla f(x_t)\\| \lesssim \frac{\Delta_f^{3/4}+O(1)}{(2^{k-1})^{1/3}}.$$
Since the algorithm is not resetting $x_t=x_0$ in the new stage, the parameter $\Delta_f$ is not constant. In particular, in stage $k$, $\Delta_f = \mathbb{E}[f(x_{2^{k-1}}) - \inf f(x)]$. Consequently, the doubling trick implicitly assumes $f$ has bounded function value in order to bound $\Delta_f$ in all stages.
If the authors can address my major concerns, I am inclined to revise my score upwards, given the otherwise strong results of the paper.
Questions
- Is there a specific reason why the authors prefer $\alpha$ as large as possible in the trade-off of $\Delta^{1/2(1-\alpha)} + \sigma^{1/(1-\alpha)} + L^{1/2\alpha}$. Is it due to the assumption (or some empirical observation) that the smoothness constant $L$ is usually the dominating parameter?
- In the experiments, did the authors apply any learning rate scheduler (e.g., linear decay or cosine decay) on top of the adaptive learning rate? Although somewhat tangential, it would be interesting to compare the new adaptive scheduler with other popular schedulers.