通过低维克雷洛夫子空间加速劣条件问题的优化收敛。
Streaming Krylov-Accelerated Stochastic Gradient Descent
- 将随机梯度投影到克雷洛夫子空间,用流式高斯-赛德尔法计算投影系数。
- 在条件数超10³时,收敛速度和最终误差显著优于SGD与Adam。
- 适用于大规模问题,尤其适合AMD GPU上的高效实现。
我们提出SKA-SGD(Streaming Krylov-Accelerated Stochastic Gradient Descent),一种新型优化方法,通过将随机梯度投影至低维克雷洛夫子空间,加速劣条件问题的收敛。受近期s步共轭梯度法与流式高斯-赛德尔格拉姆求解器启发,该方法结合三项创新:(1) 仅需一次流式高斯-赛德尔迭代计算投影系数,数学上等价于修正格拉姆-施密特正交化;(2) 使用切比雪夫多项式基构造克雷洛夫子空间,提升数值稳定性;(3) 基于HIP的AMD GPU高效实现。理论证明,本方法以O(s²)复杂度实现接近机器精度的反向误差,而非O(s³)。实验表明,在条件数超过10³的问题上,SKA-SGD显著优于标准SGD与Adam,收敛更快且最终误差更低。GPU性能分析显示,当处理器数约64、问题规模n ≥ 10⁶时,通信规避优势超越计算开销,出现性能拐点。
原文摘要 · Abstract (English)
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$.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。