LOST: Low-rank and Sparse Pre-training for Large Language Models
LOST: Low-rank and Sparse Pre-training for Large Language Models
基本情報
- arXiv ID: 2508.02668v1 (https://arxiv.org/abs/2508.02668)
- 著者: Jiaxi Li (University of Surrey), Lu Yin (University of Surrey), Li Shen (Sun Yat-sen University), Jinjin Xu (Bytedance), Liwu Xu (Alibaba Group), Tianjin Huang (University of Exeter), Wenwu Wang (University of Surrey), Shiwei Liu (University of Oxford), Xilu Wang (University of Surrey)
- 所属: University of Surrey, University of Oxford, Sun Yat-sen University, Bytedance, Alibaba Group, University of Exeter
- 投稿日: 2025年8月6日
- カテゴリ: cs.LG
簡単に説明すると
この論文は、大規模言語モデル（LLM）を効率的に事前学習するための新しい手法「LOST（Low-rank and Sparse Training）」を提案しています。
LOSTは、重み行列を低ランク成分とスパース成分に分解することで、メモリ使用量と計算コストを約50%削減しながら、フルランクモデルと同等またはそれ以上の性能を達成します。具体的には、特異値分解（SVD）を使用して、大きな特異値に関連する主要な部分空間を低ランク成分として保持し、残りの部分空間から重要なチャネルを選択してスパース成分として補完する設計となっています。
GitHubリポジトリが公開されており（https://github.com/JiaxiLi1/LOST-Low-rank-and-Sparse-Training-for-Large-Language-Models）、実装の詳細やコードを確認できます。この手法は、限られた計算リソースでLLMを訓練したい研究者や開発者にとって、実用的な価値が高い研究成果です。
1. 研究概要
1.1 背景と動機
大規模言語モデル（LLM）は様々な分野で卓越した性能を示していますが、数十億のパラメータを持つモデルの事前学習には膨大な計算リソースとメモリが必要となります。この問題に対処するため、ファインチューニングの分野では、LoRA（Low-Rank Adaptation）をはじめとする低ランク近似手法が人気を集めています。LoRAは、事前学習済みの重みを固定したまま、低ランクアダプタのみを更新することで、メモリ使用量と計算コストを約70-90%削減します。
しかし、これらの効率的なファインチューニング手法の成功にもかかわらず、より計算集約的な事前学習段階における有効性は十分に探求されていませんでした。既存の低ランク構造を用いたニューラルネットワークの学習に関する研究には、次のような制限がありました。1つ目は主に小規模モデルに限定されていたこと、2つ目はフルランクでのウォームアップ訓練に依存していたこと、3つ目はフィードフォワードネットワーク（FFN）レイヤーのみに限定されていたことです。
さらに、低ランク重みによるLLMの事前学習は、一貫してフルランク訓練と比較して性能が劣るという問題がありました。これは、低ランク構造が重み行列間に複雑な依存関係を導入することが原因です。その訓練ダイナミクスへの影響が不明確であることも1つの要因として考えられます。
1.2 主要な貢献
本研究では、以下の3つの主要な貢献をしています。
- 低ランクとスパース成分を組み合わせた新しい訓練手法LOSTを提案し、フルランクの事前学習を必要とせずにメモリ効率を約50%向上させたLLM事前学習を可能にしました
- 低ランクとスパース成分の相補的な設計により、両成分が互いに補完し合う新しいアプローチを開発しました
- 60Mから7Bまでの様々なサイズのLLaMAモデルでC4データセットを用いた事前学習を実施し、性能と効率性の向上を実証しました
2. 提案手法
2.1 手法の概要
LOSTは、標準的な初期化手法（Kaiming初期化など）を用いて初期化されたフルランク重み行列W ∈ R^(m×n)から開始します。この行列を低ランク成分とスパース成分に分解しますが、従来の手法のように両成分を単純に組み合わせるのではなく、特異値分解（SVD）を活用して、低ランクとスパース成分が直交するランク部分空間で互いに補完し合うように設計されています。
具体的な処理フローは以下の通りです。まず、初期化されたフルランク重み行列WにSVDを適用し、特異値と特異ベクトルを取得します。次に、上位r個の特異値とそれに対応する特異ベクトルを用いて低ランク成分W_l = AB^Tを構築します。そして、残りの特異値から構成される補完行列W_compから、チャネルごとの重要度に基づいてスパース成分W_sを生成します。最後に、低ランク成分とスパース成分を結合係数γで組み合わせて最終的な出力を生成します。
2.2 技術的詳細
低ランクモデリングでは、初期化されたフルランク行列WにSVDを適用します：
W = UΣV^T = Σ_{i=1}^{rank(W)} σ_i u_i v_i^T
ここで、Uは左特異ベクトル行列、Vは右特異ベクトル行列、Σは特異値を対角要素とする行列です。目標ランクrが与えられたとき、上位r個の特異値とそれに対応する特異ベクトルを選択して、低ランク近似W_l = AB^Tを構築します：
- A = U_r Σ_r^{1/2} ∈ R^(m×r)
- B = V_r Σ_r^{1/2} ∈ R^(n×r)
さらに、表現力を向上させるため、AとBの間にSiLU非線形活性化関数を追加しています。
チャネルワイズスパースモデリングでは、残りの部分空間W_comp = Σ_{i=r+1}^{rank(W)} σ_i u_i v_i^Tから重要なチャネルを選択します。目標スパース率ρが与えられたとき、k = ⌈ρ·n⌉個のチャネルを保持します。チャネルの重要度は、L2ノルムに基づいて計算されます：
CI_j = ||W_comp[:, j]||_2, j = 1, ..., n
最も重要度の高いk個のチャネルが選択され、対応する重みがスパース行列W_s ∈ R^(m×k)として保存されます。
訓練時の順伝播では、以下の式で計算が行われます：
o = γ · σ(xA)B^T + (1-γ) · x_{[:, I]} W_s^T
ここで、γ ∈ [0,1]は両成分の相対的な重要度を制御する係数です。
2.3 新規性
LOSTの主要な新規性は、低ランクとスパース成分の相補的な設計にあります。従来の手法では、低ランクとスパース成分は独立に初期化され、単純に組み合わされていました。これに対し、LOSTは以下の点で革新的です。
第一に、SVDに基づく初期化により、低ランク成分は最大の特異値に関連する主要な部分空間を捉え、スパース成分は残りの部分空間から重要な情報を保持するように設計されています。これにより、両成分が直交する部分空間で動作し、情報の重複を避けながら相補的に機能します。
第二に、チャネルワイズの構造化スパース性を採用することで、メモリ効率を大幅に改善しています。従来の要素ワイズスパース性では、バイナリマスクやインデックスの保存に元のスパース成分の2倍のメモリが必要でしたが、チャネルワイズスパース性では、選択されたチャネルのインデックスのみを保存すればよく、メモリオーバーヘッドが大幅に削減されます。
第三に、低ランク行列間に非線形活性化関数を挿入することで、パラメータを追加することなくモデルの表現力を向上させています。
3. 実験結果
3.1 実験設定
実験はColossal Clean Crawled Corpus（C4）データセットを使用して、LLaMAベースのアーキテクチャで60Mから7Bまでの様々なモデルサイズで実施されました。実装には、Pre-normalization、RMSnorm、SwiGLU活性化関数が含まれています。メモリ効率を高めるためBF16フォーマットを使用し、既存の文献に従ってオプティマイザ設定、コサイン学習率減衰、ウォームアップ戦略を採用しています。
すべてのモデルサイズで、係数γは0.7に設定され、スパース成分の構築にはSVDベースの初期化でランク256を使用しています。スパース性レベルは0.01に設定され、低ランク成分のランクは他のベースラインと同等のパラメータ数を維持するように調整されています。
3.2 主要な結果
LOSTは、すべてのモデルサイズにおいて、追加のメモリオーバーヘッドなしにベースライン手法を上回る性能を示しました。フルランクモデルとの比較では、350Mパラメータで同等の性能を達成し、他のスケールではフルランクモデルを上回りました。具体的には、1Bで3.5%、130Mで1.2%、60Mで5.3%低いパープレキシティを達成しています。
最も近いベースラインであるSLTrainとの比較では、LOSTはすべてのモデルスケールで性能とメモリ効率の両方で大幅な改善を示しました。例えば、1Bモデルでは、LOSTがパープレキシティ15.02を達成したのに対し、SLTrainは16.14でした。
メモリ使用量の詳細な分析では、LOSTを含むすべての低ランク手法が、フルランクモデルと比較してメモリ使用量をほぼ半分に削減することを示しています。1Bモデルでの実際のメモリ消費量の内訳を見ると、重み、活性化、勾配、オプティマイザ状態のすべてのコンポーネントでメモリ削減が達成されています。
3.3 既存手法との比較
LOSTは、Full-Rank、LoRA、ReLoRA、GaLore、LORO、CoLA、SLTrainといった既存手法と比較されました。すべてのモデルサイズにおいて、LOSTは最も低いパープレキシティを達成しています。
特筆すべきは、7Bモデルでのスケーリング性能です。計算リソースの制約により40K ステップのみの訓練でしたが、LOSTは8-bit Adamや8-bit GaLoreを上回る性能を示しました。標準のLOSTは40Kステップでパープレキシティ16.48を達成し、8-bit LOSTでも17.59という競争力のある性能を維持しました。
ファインチューニングタスクでも、LOSTの汎用性が実証されました。RoBERTa-baseモデルをGLUEベンチマークでファインチューニングした結果、LOSTはすべてのタスクで競争力のある、または優れた性能を達成しました。
4. 実用性評価
4.1 実装の容易性
LOSTの実装は比較的簡単で、既存の深層学習フレームワークに容易に統合できます。主要なコンポーネントは、SVDベースの初期化、チャネルワイズスパース選択、そして訓練時の効率的な順伝播計算です。GitHubで公開されているコードは、これらの実装の詳細を提供しており、研究者や開発者が自身のプロジェクトに適用することを容易にしています。
特に、SVDは初期化時に一度だけ実行されるため、訓練時のオーバーヘッドはありません。また、チャネルワイズスパース性により、推論時のメモリアクセスパターンも効率的に保たれます。
4.2 計算効率
LOSTは、パラメータ数をmnからr(m+n)+mkに削減することで、大幅な計算効率の改善を実現しています。ここで、k = ⌈ρ·n⌉であり、k << min(m,n)です。低ランク成分とチャネルワイズスパース成分の組み合わせにより、メモリ効率が向上し、特に要素ワイズインデックスの保存に関連する大幅なオーバーヘッドが排除されます。
実験結果は、LOSTがフルランクモデルと比較してメモリ使用量をほぼ半分に削減しながら、同等またはそれ以上の性能を達成することを示しています。これは、限られた計算リソースでLLMを訓練する必要がある環境において、特に価値があります。
4.3 応用可能性
LOSTは、様々な規模のLLMに適用可能であることが実証されています。60Mから7Bまでの幅広いモデルサイズでの成功は、この手法のスケーラビリティを示しています。また、事前学習だけでなくファインチューニングタスクでも有効であることが示されており、幅広い応用が期待できます。
特に、計算リソースが限られている研究機関や、エッジデバイスでのLLM展開を検討している企業にとって、LOSTは実用的な解決策を提供します。また、環境への配慮から計算効率を重視する研究開発においても、重要な技術となる可能性があります。
5. まとめと所感
5.1 論文の意義
この研究は、LLMの効率的な事前学習という重要な課題に対して、実用的かつ理論的に健全な解決策を提供しています。低ランクとスパース成分の相補的な設計という新しいアプローチにより、従来の手法の限界を克服し、フルランクモデルと同等またはそれ以上の性能を達成しています。
特に注目すべきは、SVDベースの初期化により、両成分が直交する部分空間で動作するという理論的な裏付けがある点です。これにより、情報の重複を避けながら、モデルの表現力を最大限に活用します。また、チャネルワイズスパース性の採用により、実用的なメモリ効率の改善も実現しています。
アブレーション研究の充実も本論文の強みです。低ランク初期化方法、スパース成分の構築方法、結合係数γの影響など、様々な設計選択の影響を詳細に分析しており、提案手法の各要素の重要性を明確に示しています。
5.2 今後の展望
今後の研究の方向性として、いくつかの興味深い可能性が考えられます。まず、より大規模なモデル（70B以上）への適用や、マルチモーダルモデルへの拡張が期待されます。また、動的にランクやスパース性を調整する適応的な手法の開発も、さらなる効率改善につながる可能性があります。
理論的な観点からは、低ランクとスパース成分の最適な組み合わせに関する理論的な分析や、訓練ダイナミクスのより深い理解が求められます。また、他の効率化技術（量子化など）との組み合わせによる、さらなる圧縮の可能性も探求する価値があります。
実用的な観点では、特定のドメインやタスクに特化した最適化や、分散学習環境でのLOSTの適用など、実世界での展開に向けた研究が重要になるでしょう。