DarwinLM: Evolutionary Structured Pruning of Large Language Models

技術要約: DarwinLM

問題提起
大規模言語モデル（LLM）は大きな成功を収めているが、その膨大な計算コストが、特にリアルタイムアプリケーションにおけるデプロイの妨げとなっている。構造化プルーニング（Structured Pruning）は、主流のハードウェア上でエンドツーエンドの速度向上を実現するためにモデルを圧縮する解決策となるが、既存の手法は以下の2つの重要な要素を考慮できていないことが多い：

1. 非一様な感度（Non-uniform Sensitivity）: モデルの各コンポーネント（レイヤー、アテンションヘッド、MLP次元）は、プルーニングに対する感度が異なる。一様な圧縮は、しばしば性能の崩壊を招く。
2. 圧縮後の学習（Post-Compression Training）: プルーニング手法は、単に有能な部分構造を特定するだけでなく、その構造が圧縮後のファインチューニングを通じて回復可能であることを保証しなければならない。ZipLMなどの従来の手法は、局所的なレイヤーごとの誤差やワンショットの性能を最適化していることが多いが、これはダウンストリームタスクの性能やファインチューニングによる回復度合いと必ずしも相関しない。

手法
著者らは、最適な非一様スパース性割り当てを見出すために進化論的探索プロセスを利用した、学習を考慮した構造化プルーニング手法であるDarwinLMを提案している。このパイプラインは、主に2つのステージで構成される：

1. 学習を考慮した選択を伴う進化論的探索:

   • 初期化: 探索は、二次の情報（Optimal Brain Surgeon定式化）を用いたワンショット構造化プルーニングによって生成された「親」モデルから開始される。
   • 突然変異（Mutation）: 各世代において、アルゴリズムは親モデルをコピーし、「レベルスイッチ」突然変異を適用することで「子」モデルを生成する。これには、グローバルな目標スパース性またはスピードアップの制約を維持するために、サブブロック間（例：あるレイヤーのスパース性を高める一方で、別のレイヤーのスパース性を下げる）でスパース性のレベルをシフトさせることが含まれる。
   • 学習を考慮した選択: コアとなる革新は、軽量なファインチューニングを選択プロセスに統合したことである。オフスプリング（子モデル）を、単なるワンショット指標（小さなキャリブレーションセットに対するKLダイバージェンスなど）のみに基づいて選択するのではなく、本手法は多段階の選択プロセスを採用している。オフスプリングは、段階的に拡大されるトークン予算（例：10K $\to$ 50K $\to$ 200K トークン）で学習される。各段階で、性能の低いモデルは排除される。これにより、選択された最も「適応度の高い」モデルが、単に静的な状態で正確であるだけでなく、継続的な学習後に堅牢であり、かつ回復可能であることを保証する。
   • 終了: 探索が収束すると、最終的な候補は、最終評価の前に、より広範なファインチューニング段階（例：10B トークン）を経る。

2. プルーニング済みレイヤー・データベース:

   • 効率的な探索を促進するため、本手法は様々な離散化されたスパース性レベルにおけるスパースなサブブロック（MLPおよびアテンションモジュール）のデータベースを事前計算する。
   • アテンションモジュールのプルーニングは、ヘッド単位で行われる。MLPについては、次元は32の倍数でプルーニングされる。
   • Group-Query Attention (GQA) を持つモデルの場合、KおよびV行列を直接プルーニングするのではなく、計算の互換性を維持するために、対応するヘッドをフォワードパス中に削除する。

3. MoEアーキテクチャへの拡張:

   • 本手法は、Mixture-of-Experts (MoE) モデル（例：Qwen3-30B-A3B）にも拡張されている。この設定では、探索はエキスパートMLP内のスパース性を最適化する一方で、MoEブロック間では一様なスパース性を維持する。エキスパートがパラメータの大部分を占めるため、アテンションモジュールは一般にプルーニング対象外となる。

主な貢献

• 学習を考慮した進化論的探索: DarwinLMは、ワンショットの精度のみを最適化するのではなく、ファインチューニング中の性能回復能力を明示的に考慮した、新しい進化論的フレームワークを導入する。
• 非一様構造化プルーニング: 本手法は、精度の損失なしに高い圧縮率を達成するために、異なるレイヤーの多様な感度を活用した、LLMのきめ細かな非一様圧縮を可能にする。
• MoEへの適応: 本研究は、MoEアーキテクチャに特化した非一様構造化プルーニングの最初の探求であり、構造化スパース性が複雑なエキスパートベースのモデルにおいても有効であることを示している。
• 効率性: 探索プロセスは非常に効率的であり、4基のコンシューマー向けGPUで約8時間で完了し、最終的なファインチューニング段階は標準的なクラスターで約半日を要する。

実験結果
著者らは、Llama-2-7B、Llama-3.1-8B、Qwen-2.5-14B-Instruct、および Qwen3-30B-A3B (MoE) に関して DarwinLM を検証している。

• 性能: DarwinLM は、ワンショット構造化プルーニングにおいて最先端（SOTA）の性能を達成している。例えば、Llama-3.1-8B を半分にプルーニングした場合、従来の最良の手法である ZipLM と比較して、平均ゼロショット精度が 5.9% 高くなっている。
• データ効率: 本手法は、回復に必要な学習予算を大幅に削減する。DarwinLM は、50B トークンのファインチューニングを必要とする ShearedLlama を上回る性能を、わずか 10B トークンで達成している。同じ 10B トークンの設定において、DarwinLM は ShearedLlama を凌駕する。
• MoE の結果: Qwen3-30B-A3B モデルにおいて、DarwinLM は 10B トークンのファインチューニング後に元のモデルの精度の $\ge$90% を保持する 16B-A2B バリアントを生成する。
• ハードウェア効率: プルーニングされたモデルは、直接的な速度向上とメモリ削減を示す。例えば、2.7B の DarwinLM モデルは、L40s GPU 上で、デンスな 7B ベースラインと比較して 1.98$\times$ のスループットを達成し、2.43$\times$ 少ないメモリを使用する。
• 比較: 本手法は、様々なスパース性レベルにおいて、一様プルーニングのベースラインや他の構造化プルーニング手法（ShortGPT、Shortened-Llama、EvoPress を含む）を凌駕しており、特に他の手法が急速に劣化する高圧縮率において顕著である。

意義
論文は、DarwinLM が構造的探索と事後学習による回復の間の溝を埋めることにより、LLM圧縮における重要な進歩を遂げたと主張している。非一様な構造化プルーニングが、膨大な再学習予算を必要とせずに、デンスなアーキテクチャおよび MoE アーキテクチャの両方に効果的に適用できることを示すことで、本研究は「高性能な圧縮には大規模な再学習予算が必要である」という概念に異を唱えている。著者らは、DarwinLM を、精度を犠牲にすることなく効率的な LLM のデプロイを可能にする、実用的かつハードウェアに依存しないソリューションとして位置づけている。