Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

$μ^2$-SGD: Stable Stochastic Optimization via a Double Momentum Mechanism

Tehila Dahan, Kfir Y. Levy

TL;DR

The paper tackles stochastic convex optimization with smooth losses (SCO-EOS) by introducing a novel gradient estimator that blends two momentum-based mechanisms. The μ^2-SGD algorithm uses a double-momentum approach to shrink gradient-estimation error with a single sample per iteration, achieving GD-like stability and fixed-learning-rate guarantees in both noiseless and noisy settings, and extending to an accelerated variant μ^2-ExtraSGD. Theoretical results show optimal convergence rates with broad learning-rate robustness, and empirical results on convex and non-convex tasks confirm improved stability and performance. This work offers a practical pathway to robust SGD-like optimization without heavy hyperparameter tuning, with potential for large-scale and non-convex applications.

Abstract

We consider stochastic convex optimization problems where the objective is an expectation over smooth functions. For this setting we suggest a novel gradient estimate that combines two recent mechanism that are related to notion of momentum. Then, we design an SGD-style algorithm as well as an accelerated version that make use of this new estimator, and demonstrate the robustness of these new approaches to the choice of the learning rate. Concretely, we show that these approaches obtain the optimal convergence rates for both noiseless and noisy case with the same choice of fixed learning rate. Moreover, for the noisy case we show that these approaches achieve the same optimal bound for a very wide range of learning rates.

$μ^2$-SGD: Stable Stochastic Optimization via a Double Momentum Mechanism

TL;DR

The paper tackles stochastic convex optimization with smooth losses (SCO-EOS) by introducing a novel gradient estimator that blends two momentum-based mechanisms. The μ^2-SGD algorithm uses a double-momentum approach to shrink gradient-estimation error with a single sample per iteration, achieving GD-like stability and fixed-learning-rate guarantees in both noiseless and noisy settings, and extending to an accelerated variant μ^2-ExtraSGD. Theoretical results show optimal convergence rates with broad learning-rate robustness, and empirical results on convex and non-convex tasks confirm improved stability and performance. This work offers a practical pathway to robust SGD-like optimization without heavy hyperparameter tuning, with potential for large-scale and non-convex applications.

Abstract

We consider stochastic convex optimization problems where the objective is an expectation over smooth functions. For this setting we suggest a novel gradient estimate that combines two recent mechanism that are related to notion of momentum. Then, we design an SGD-style algorithm as well as an accelerated version that make use of this new estimator, and demonstrate the robustness of these new approaches to the choice of the learning rate. Concretely, we show that these approaches obtain the optimal convergence rates for both noiseless and noisy case with the same choice of fixed learning rate. Moreover, for the noisy case we show that these approaches achieve the same optimal bound for a very wide range of learning rates.
Paper Structure (44 sections, 11 theorems, 64 equations, 41 figures, 2 algorithms)

This paper contains 44 sections, 11 theorems, 64 equations, 41 figures, 2 algorithms.

Table of Contents

  1. Introduction
  2. Setting
  3. Momentum Mechanisms
  4. Mechanism I: Anytime-GD
  5. Mechanism II: Recursive Corrected Averaging
  6. Our Approach: Double Momentum Mechanism
  7. Stability of $\mu^2$-SGD.
  8. Extensions
  9. Necessity of both Mechanisms
  10. Proof Sketch of Thm. \ref{['thm:Main']}
  11. Accelerated Version: $\mu^2-\texttt{Extra}\text{SGD}$
  12. Our Approach.
  13. Stability of $\mu^2-\texttt{Extra}\text{SGD}$ .
  14. Experiments
  15. Conclusion
  16. ...and 29 more sections

Key Result

Theorem 3.1

Let $f:\mathcal{K}\mapsto {\mathbb R}$ be a convex function with a minimum $w^*\in\mathop{\mathrm{arg\,min}}\limits_{w\in\mathcal{K}}f(w)$. Also let $\{\alpha_t\geq 0\}_t$, and $\{w_t\in\mathcal{K}\}_{t},\{x_t\in\mathcal{K}\}_{t}$, such that $\{x_t\}_{t}$ is an $\{\alpha_t\}_t$ weighted average of $

Figures (41)

  • Figure 1: MNIST: Test Accuracy Over Different Learning Rates in a Convex Setting ($\uparrow$ is better).
  • Figure 2: Test Accuracy Over a Range of Learning Rates in Non-Convex Setups ($\uparrow$ is better).
  • Figure 3: Test Accuracy at Learning Rate = 10 ($\uparrow$ is better).
  • Figure 4: Test Accuracy at Learning Rate = 1 ($\uparrow$ is better).
  • Figure 5: Test Accuracy at Learning Rate = 0.1 ($\uparrow$ is better).
  • ...and 36 more figures

Theorems & Definitions (20)

  • Theorem 3.1: Rephrased from Theorem 1 in cutkosky2019anytime
  • Theorem 4.1
  • Theorem 4.2: $\mu^2$-SGD Guarantees
  • proof
  • Theorem 5.1: See e.g. the proof of Thm. 1 in kavis2019unixgrad
  • Theorem 5.2: $\mu^2-\texttt{Extra}\text{SGD}$
  • proof : Proof of Thm. \ref{['thm:Main']}
  • Lemma C.1
  • proof : Proof of Lemma \ref{['lem:SumMart']}
  • proof : Proof of Thm. \ref{['thm:muSGD']}
  • ...and 10 more