オンラインサブスペース降下法を用いたメモリ効率の良いLLMトレーニング

Memory-Efficient LLM Training with Online Subspace Descent

August 23, 2024

著者: Kaizhao Liang, Bo Liu, Lizhang Chen, Qiang Liu

cs.AI

要旨

最近、メモリ効率の高いLLMトレーニングアルゴリズムが広く注目されています。これらの手法は、勾配の低ランク構造を活用して、最適化器の状態を特異値分解（SVD）によって得られる射影行列を使用して部分空間に射影します。ただし、これらのアルゴリズムの収束性は、射影行列の更新ルールに高度に依存しています。本研究では、射影行列の任意の更新ルールに対する初めての収束保証を提供します。この保証は、Hamiltonian Descentで分析できる最も一般的な最適化器（LION、Adamなど）に適用できます。理論的理解に基づいて、我々はSVDを使用せずに新しいサブスペース降下最適化器ファミリーであるOnline Subspace Descentを提案します。Online Subspace Descentは、射影行列を固有ベクトルで更新するのではなく、オンラインPCAで更新します。Online Subspace Descentは柔軟であり、トレーニングに最小限のオーバーヘッドしか導入しません。我々は、C4データセット上の60Mから7BパラメータのLLaMAモデルの事前トレーニングタスクにおいて、Online Subspace Descentが、異なる設定で最先端の低ランクトレーニング手法よりも低いperplexityと優れた下流タスクのパフォーマンスを達成し、完全ランクのベースラインとの差を縮めることを示します。

English

Recently, a wide range of memory-efficient LLM training algorithms have gained substantial popularity. These methods leverage the low-rank structure of gradients to project optimizer states into a subspace using projection matrix found by singular value decomposition (SVD). However, convergence of these algorithms is highly dependent on the update rules of their projection matrix. In this work, we provide the first convergence guarantee for arbitrary update rules of projection matrix. This guarantee is generally applicable to optimizers that can be analyzed with Hamiltonian Descent, including most common ones, such as LION, Adam. Inspired by our theoretical understanding, we propose Online Subspace Descent, a new family of subspace descent optimizer without SVD. Instead of updating the projection matrix with eigenvectors, Online Subspace Descent updates the projection matrix with online PCA. Online Subspace Descent is flexible and introduces only minimum overhead to training. We show that for the task of pretraining LLaMA models ranging from 60M to 7B parameters on the C4 dataset, Online Subspace Descent achieves lower perplexity and better downstream tasks performance than state-of-the-art low-rank training methods across different settings and narrows the gap with full-rank baselines.