SUMO: Subspace-Aware Moment-Orthogonalization for Accelerating Memory-Efficient LLM Training

要約

低ランクの勾配ベースの最適化方法により、大規模な言語モデル（LLM）のトレーニング中にメモリ効率が大幅に改善され、パフォーマンスを犠牲にすることなく制約されたハードウェア内の操作が可能になります。
しかし、これらの方法は主に記憶の節約を強調し、多くの場合、標準的な等方性が最も急な降下技術に依存しているため、収束の潜在的な加速を見落とします。
この論文では、動的に適応した低次元サブスペース内での正確な直交化のために正確な単数値分解（SVD）を使用するオプティマイザーであるSUMO（サブスペースを意識したモーメント正統派）を提案し、標準誘導の急な降下最適化ステップを促進します。
最適化ステップを損失状況のスペクトル特性と明示的に整列させることにより、SUMOは、Newton-Schulz Orthogonalization近似のような一般的に使用される方法に関連する近似エラーを効果的に軽減します。
理論的には、これらの近似誤差の上限を確立し、モーメントの条件数に依存していることを証明し、LLMトレーニング中に分析的に実証する条件を証明します。
さらに、私たちは理論的にも経験的にも、SVDを介した正確な直交化が収束率を大幅に改善し、全体的な複雑さを減らすことを示しています。
経験的評価は、SUMOが収束を加速し、安定性を向上させ、パフォーマンスを向上させ、最新の方法と比較してメモリ要件を最大20％削減することを確認します。

要約(オリジナル)

Low-rank gradient-based optimization methods have significantly improved memory efficiency during the training of large language models (LLMs), enabling operations within constrained hardware without sacrificing performance. However, these methods primarily emphasize memory savings, often overlooking potential acceleration in convergence due to their reliance on standard isotropic steepest descent techniques, which can perform suboptimally in the highly anisotropic landscapes typical of deep networks, particularly LLMs. In this paper, we propose SUMO (Subspace-Aware Moment-Orthogonalization), an optimizer that employs exact singular value decomposition (SVD) for moment orthogonalization within a dynamically adapted low-dimensional subspace, enabling norm-inducing steepest descent optimization steps. By explicitly aligning optimization steps with the spectral characteristics of the loss landscape, SUMO effectively mitigates approximation errors associated with commonly used methods like Newton-Schulz orthogonalization approximation. We theoretically establish an upper bound on these approximation errors, proving their dependence on the condition numbers of moments, conditions we analytically demonstrate are encountered during LLM training. Furthermore, we both theoretically and empirically illustrate that exact orthogonalization via SVD substantially improves convergence rates while reducing overall complexity. Empirical evaluations confirm that SUMO accelerates convergence, enhances stability, improves performance, and reduces memory requirements by up to 20% compared to state-of-the-art methods.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google