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.
