Streaming Krylov-Accelerated Stochastic Gradient Descent

TL;DR

SKA-SGD addresses the slow convergence of stochastic gradient methods on ill-conditioned problems by projecting stochastic gradients onto a low-dimensional Krylov subspace constructed via a streaming Gauss-Seidel procedure. The method leverages a Chebyshev basis to maintain numerical stability and implements the projection in a memory- and communication-efficient way suitable for AMD GPUs via HIP. The authors prove variance reduction through Krylov projection, provide convergence guarantees for smooth strongly convex objectives, and show that a single forward sweep yields near-machine-precision backward error on the Gram-like system. Empirically, SKA-SGD achieves significantly faster convergence and lower final error than SGD and Adam across a range of condition numbers (up to κ ≥ 10^4), with notable variance reduction and robust late-stage performance, making it practical for large-scale ill-conditioned optimization tasks.

Abstract

We present SKA-SGD (Streaming Krylov-Accelerated Stochastic Gradient Descent), a novel optimization approach that accelerates convergence for ill-conditioned problems by projecting stochastic gradients onto a low-dimensional Krylov subspace. Directly inspired by recent advances in s-step Conjugate Gradient methods with streaming Gauss-Seidel Gram solvers \cite{dambra2025sstep}, our method extends these techniques to the stochastic optimization domain. Our approach combines three key innovations: (1) projection coefficients computed via a single streaming Gauss-Seidel iteration, which is mathematically equivalent to Modified Gram-Schmidt orthogonalization; (2) a Chebyshev polynomial basis for constructing the Krylov subspace, providing superior numerical stability; and (3) efficient implementation for AMD GPUs using HIP. We prove that our streaming approach achieves a backward error near machine precision with $O(s^2)$ complexity rather than $O(s^3)$, where $s$ is the Krylov subspace dimension. Experimental results demonstrate that SKA-SGD significantly outperforms standard SGD and Adam in convergence rate and final error, particularly for problems with condition numbers exceeding $10^3$. GPU performance analysis reveals a crossover point where communication-avoiding benefits outweigh computational overhead, typically occurring at moderate scale ($p \approx 64$ processors) for problem sizes $n \geq 10^6$.

Figures (8)

• Figure 1: Relative error $\|x_k - x^\star\| / \|x^\star\|$ for quadratic minimization with condition number $\kappa = 10^4$. SKA-SGD, using Chebyshev basis and Gauss--Seidel projection, achieves faster and more stable convergence than SGD and Adam.
• Figure 1: Convergence in $|f(x_k) - f^*|$ for a quadratic problem with condition number $\kappa=10^8$. The ULTIMATE method combining all techniques exhibits superior convergence compared to other variants and baseline methods.
• Figure 2: Error norm $\|x_k - x^*\|$ convergence on an extremely ill-conditioned problem ($\kappa=10^8$). The ULTIMATE method demonstrates significantly faster convergence in the early and middle phases of optimization.
• Figure 3: Convergence comparison for different Krylov dimensions on a quadratic problem with $\kappa=10^5$. At this moderate condition number, the benefits of increased Krylov dimension are less pronounced.
• Figure 4: Multiple trial convergence analysis: Mean ± 1 standard deviation over 10 independent trials for iterations 1800-2000. SKA-SGD (blue line with blue shaded region) demonstrates consistently lower objective values and significantly reduced variance compared to standard SGD (black dashed line with gray shaded region). The statistical advantage of SKA-SGD manifests as both improved convergence and enhanced stability across different random seeds.

Theorems & Definitions (4)

• Proposition 1
• Theorem 1
• Theorem 2
• Theorem 3