Table of Contents
Fetching ...

Convergence Analysis of Randomized Subspace Normalized SGD under Heavy-Tailed Noise

Gaku Omiya, Pierre-Louis Poirion, Akiko Takeda

TL;DR

This paper tackles nonconvex stochastic optimization in high dimensions by introducing randomized subspace updates to cut per-iteration cost. It develops two methods: RS-SGD, with high-probability convergence under sub-Gaussian noise, and RS-NSGD, which normalizes projected gradients to robustly handle heavy-tailed noise with bounded $p$-th moments; both come with in-expectation and high-probability guarantees. The results show that RS-SGD achieves the same $\epsilon$-stationary oracle-complexity order as prior in-expectation analyses, while RS-NSGD can offer improved complexity over full-dimensional NSGD in certain regimes and remains provably reliable under heavy-tailed noise. Empirical studies on synthetic heavy-tailed problems and character-level language modeling corroborate the theoretical findings, demonstrating practical gains in memory/communication-constrained settings and motivating broader use of Haar-based randomized subspaces in large-scale learning.

Abstract

Randomized subspace methods reduce per-iteration cost; however, in nonconvex optimization, most analyses are expectation-based, and high-probability bounds remain scarce even under sub-Gaussian noise. We first prove that randomized subspace SGD (RS-SGD) admits a high-probability convergence bound under sub-Gaussian noise, achieving the same order of oracle complexity as prior in-expectation results. Motivated by the prevalence of heavy-tailed gradients in modern machine learning, we then propose randomized subspace normalized SGD (RS-NSGD), which integrates direction normalization into subspace updates. Assuming the noise has bounded $p$-th moments, we establish both in-expectation and high-probability convergence guarantees, and show that RS-NSGD can achieve better oracle complexity than full-dimensional normalized SGD.

Convergence Analysis of Randomized Subspace Normalized SGD under Heavy-Tailed Noise

TL;DR

This paper tackles nonconvex stochastic optimization in high dimensions by introducing randomized subspace updates to cut per-iteration cost. It develops two methods: RS-SGD, with high-probability convergence under sub-Gaussian noise, and RS-NSGD, which normalizes projected gradients to robustly handle heavy-tailed noise with bounded -th moments; both come with in-expectation and high-probability guarantees. The results show that RS-SGD achieves the same -stationary oracle-complexity order as prior in-expectation analyses, while RS-NSGD can offer improved complexity over full-dimensional NSGD in certain regimes and remains provably reliable under heavy-tailed noise. Empirical studies on synthetic heavy-tailed problems and character-level language modeling corroborate the theoretical findings, demonstrating practical gains in memory/communication-constrained settings and motivating broader use of Haar-based randomized subspaces in large-scale learning.

Abstract

Randomized subspace methods reduce per-iteration cost; however, in nonconvex optimization, most analyses are expectation-based, and high-probability bounds remain scarce even under sub-Gaussian noise. We first prove that randomized subspace SGD (RS-SGD) admits a high-probability convergence bound under sub-Gaussian noise, achieving the same order of oracle complexity as prior in-expectation results. Motivated by the prevalence of heavy-tailed gradients in modern machine learning, we then propose randomized subspace normalized SGD (RS-NSGD), which integrates direction normalization into subspace updates. Assuming the noise has bounded -th moments, we establish both in-expectation and high-probability convergence guarantees, and show that RS-NSGD can achieve better oracle complexity than full-dimensional normalized SGD.
Paper Structure (43 sections, 18 theorems, 279 equations, 4 figures, 3 tables, 1 algorithm)

This paper contains 43 sections, 18 theorems, 279 equations, 4 figures, 3 tables, 1 algorithm.

Key Result

Lemma 1

Assume that the stochastic gradient $\nabla f(x,\xi)$ satisfies ass:unbiasedass:subgauss. For a fixed $x \in \mathbb{R}^d$, define the mini-batch gradient estimator and set Then for all $\lambda$ satisfying $|\lambda| \le 1/\bar{\sigma}$, it holds that

Figures (4)

  • Figure 1: The value of $\tau$ as a function of $r$ for $d = 1000$.
  • Figure 2: Synthetic quadratic with heavy-tailed noise ($d=100$, $\bar{B}=4$): oracle calls vs. $\|\nabla F(x)\|=\|\Lambda x\|$ (mean $\pm$ 1 std., 5 seeds).
  • Figure 3: Character-level language modeling on PTB ($d=5079,\bar{B}=128$) and WikiText-2 ($d=13700,\bar{B}=128$): oracle calls vs. training loss (mean $\pm$ 1 std., 5 seeds).
  • Figure 4: Character-level language modeling on PTB ($d=5079,\bar{B}=128$) and WikiText-2 ($d=13700,\bar{B}=128$): oracle calls vs. validation/test loss (5 seeds; mean $\pm$ 1 std.).

Theorems & Definitions (32)

  • Lemma 1
  • Theorem 4.1
  • Proposition 1
  • Proposition 2
  • Definition 1
  • Theorem 5.1
  • Theorem 5.2
  • Definition 2
  • Lemma 2
  • proof
  • ...and 22 more