LEAP: Learnable End-to-End Adaptive Pruning of Large Language Models

技術的サマリー：LEAP：大規模言語モデルの学習可能エンドツーエンド適応プルーニング

問題定義
大規模言語モデル（LLM）の展開は、メモリと計算リソースの制約によってますますボトルネックに直面しています。重みのスパース性は解決策を提供しますが、その領域は構造化、半構造化（例：2:4）、非構造化のレジームに分かれています。GPU カーネル（FlashLLM、SpInfer など）やデータフローハードウェアの最近の進歩により、ボトルネックは非構造化スパース性の「実行」から、最小限の精度損失でそれを「誘発」することへと移行しました。

非構造化プルーニングの現在の最先端手法（Wanda、SparseGPT、Thanos、ADMM など）は、Optimal Brain Surgeon（OBS）の原理から導出された層ごとの代理モデルに依存しています。これらの手法は、エンドツーエンドのモデル損失ではなく、局所的な再構成誤差を最小化します。その結果、局所的な誤差が深層ネットワークで蓄積・増幅され、特に aggressive なスパース性（50〜60%）の下で顕著な精度低下を招きます。

一方、MaskLLM や PATCH のようなエンドツーエンドの学習可能マスク手法は、言語モデル化損失に対して直接マスクを最適化することに成功していますが、半構造化パターンに限定されています。そのパラメータ化は、グループ内の有効なパターンに対するカテゴリカル分布に依存しています。小規模なグループ（例：2:4 スパース性の場合 $\binom{4}{2}=6$ のパターン）では扱い可能ですが、このアプローチは非構造化スパース性に対して組み合わせ的に扱い不可能になります。幅 $d=4096$ の行で 50% のスパース性の場合、有効なマスクの数は $\binom{4096}{2048} \approx 10^{1229}$ となり、パターン対数（logits）の保存とインデックス付けを不可能にします。

手法：LEAP
著者らは、パターンごとのカテゴリカルパラメータ化を重みごとのアプローチに置き換えることで、エンドツーエンドの非構造化マスク学習を扱い可能にするフレームワーク LEAP（Learnable End-to-End Adaptive Pruning）を導入します。

• 重みごとのベルヌーイ緩和: 各有効なパターンに対数オッズ（logit）を割り当てる代わりに、LEAP は個々の重みごとに学習可能な対数オッズを割り当てます。重み行列 $W \in \mathbb{R}^{m \times n}$ に対して、並列の対数オッズ行列 $P \in \mathbb{R}^{m \times n}$ が導入されます。確率的マスク $M$ は、Gumbel-sigmoid 緩和を介して生成されます。
  $$M = \sigma\left(\frac{\alpha P + g}{\tau}\right)$$
  ここで、$g$ は Gumbel ノイズ、$\sigma$ はシグモイド関数、$\alpha$ はスケール係数、$\tau$ は温度パラメータです。この定式化は、有効なパターンの数ではなく、重みの数（$O(mn)$）に比例して線形にスケーリングします。
• 最適化戦略:
  • 凍結された重み: 事前学習された重み $W$ は凍結されたままにされます。学習されるのはマスク対数オッズ $P$ のみです。これにより、事前学習モデルの較正が維持され、展開が簡素化されます。
  • 初期化: $P$ は、ワンショットの Wanda マスク（選択されたエントリを $+s$、他のエントリを $-s$ に設定）を使用して初期化され、妥当な出発点を提供し、コールドスタートの問題を回避します。
  • スケジューリング: スケールスケジュール（$\alpha$）はシグモイドを二値値に近づけるように増加し、温度スケジュール（$\tau$）はマスクを鋭くするために減少します。
  • 損失関数: 目的関数は、較正ストリーム上の言語モデル化損失と、2 つの正則化項を組み合わせます。
    1. グローバルスパース性正則化項: 全層にわたって目標密度をグローバルに強制し、個々の層がエンドツーエンドの重要性に基づいて密度を調整できるようにします。
    2. マグニチュード感知項: 重要な重みが除外されるような退化した極小値を防ぐため、最適化をより大きなマグニチュードを持つ重みを保持する方向にバイアスします。
• トレーニング効率: この手法は、少量の一般テキスト較正ストリーム（SlimPajama）上で約 2,000 反復でトレーニングされます。

主要な貢献

1. 組み合わせ的障害の特定: 本論文は、MaskLLM/PATCH が使用するパターンごとのカテゴリカルパラメータ化が、有効なパターンの指数関数的な増加により、非構造化スパース性に対してスケーリングできないことを形式的に特定しました。
2. 扱い可能な再定式化: 著者らは、Gumbel-sigmoid を介した重みごとのベルヌーイ再定式化を提案し、LLM スケールでの非構造化マスクに対するエンドツーエンドの微分可能学習を可能にしました。
3. LEAP フレームワーク: 凍結された重み上で動作し、少量の較正データセットと約 2,000 回のトレーニングステップのみを必要とする軽量なプルーニングフレームワーク。
4. 実証的検証: 0.5B から 8B のパラメータを持つ 5 つの LLM ファミリー（Qwen-2.5、Gemma-3、LLaMA-3.2、LLaMA-3.1）を、50% および 60% のスパース性で評価しました。

結果
LEAP は、層ごとのベースライン（Wanda、SparseGPT、Thanos、ADMM）および半構造化学習可能ベースラインである MaskLLM に対して評価されました。

• 精度向上: 10 の（モデル、スパース性）設定全体で、LEAP は最良の層ごとのベースラインである ADMM よりも、6 タスクの平均ゼロショット精度を +2.59 ポイント 改善しました。
• 具体的な改善: 最大の改善は、60% スパース性における LLaMA-3.2 1B で +5.40 ポイント（50.39 対 44.99）でした。
• パープレキシティ: LEAP は、ベースラインよりも一貫して低いパープレキシティ（PPL）を達成しました。例えば、50% スパース性における LLaMA-3.1 8B では、LEAP は 7.66 の PPL を達成し、ADMM の 9.12 と比較しました。
• スパース性割り当て: エンドツーエンド最適化とグローバル予算の下での実験により、LEAP はほぼ均一なブロックごとの密度割り当てに収束することが示されました。これは、半構造化手法で見られる不均一な変動とは対照的です。
• 計算コスト: ワンショット手法（H100 で数分）よりも高価ですが（H100 で数時間）、MaskLLM（同等の規模で数千の A100 時間）よりも著しく安価です（例：LLaMA-3.1 8B で 276 GPU 時間）。

意義と主張
本論文は、LEAP が非構造化スパース性におけるエンドツーエンドのマスク学習を妨げる根本的な障害を解決すると主張しています。マスク学習を重み更新から分離することで、事前学習された重みの較正を犠牲にすることなく、基盤モデルを圧縮するための実用的な道筋を提供します。

著者らは、LEAP が 50%〜60% の非構造化レジームでマスクを生成し、既存の高速化カーネル（FlashLLM、SpInfer、MACKO など）と直接互換性があることを強調しています。意義は新しい推論カーネルを作成することではなく、層ごとの手法では以前は達成不可能だった高品質でプルーニングされたモデル上で、これらのカーネルの「展開」を可能にする点にあります。この研究は、エンドツーエンド学習を介して誘発された非構造化スパース性が、汎用ハードウェアの展開において、圧縮と精度の間の好ましいトレードオフを提供し得ることを示唆しています。

限界と範囲

• 凍結された重み: 現在の範囲では重みが凍結されています。重みとマスクの同時最適化は自然な拡張として指摘されていますが、本実装では実装されていません。
• 均一割り当て: ブロックごとの密度の均一な割り当ての観察は、研究対象のモデルとスパース性レベルに対するレジーム固有の発見として提示されており、不均一な割り当てに対する一般的な規則ではありません。
• 安全性: 著者らは、圧縮されたモデルが長尾入力において異なる失敗モードを示す可能性があることに注意を促し、安全性が重要な展開については再評価を推奨しています。