MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs

技術概要：MaskPro – LLM における厳密な (N:M) 疎性を対象とした線形空間確率学習

問題定義
大規模言語モデル（LLM）の急速なスケーリングは、推論効率において重大なボトルネックを生み出し、効果的な圧縮技術の必要性を高めています。半構造化疎性、特に $M$ 個の連続する重みから $N$ 個の要素を保持する (N:M) パターンは、加速とメモリ削減のためのハードウェアフレンドリーな解決策を提供します。しかし、最適な (N:M) マスクを特定することは、強力な組み合わせ最適化問題です。既存のアプローチは、明確な限界を持つ 2 つのカテゴリーに分類されます。

1. ルールベース手法: これらは、手作りの指標（例：層ごとの再構成誤差の最小化）を用いた貪欲な層別探索を採用します。計算コストは低いものの、本質的なバイアスに悩まされ、エンドツーエンドのタスク性能との間に大きな乖離が生じます。
2. 学習ベース手法（例：MaskLLM）: これらは生成損失に対する勾配降下を通じて直接マスクを最適化します。効果的である一方、 prohibitive なトレーニングコストを伴います。具体的には、MaskLLM は $M$ 個の重みのグループ内のすべての可能なマスクに対するロジットを保存する必要があり、メモリ複雑度は $O(\binom{M}{N} \frac{d}{M})$ となります。$M$ が大きい場合、これは指数関数的にスケーリングし、大規模モデルでは非現実的となります。さらに、広大な組み合わせ空間は方策勾配に高い分散を生じさせ、トレーニングの不安定さを引き起こします。

手法：MaskPro
著者らは、(N:M) 疎性を効率的かつ安定的に学習するために設計された、新しい線形空間確率フレームワークであるMaskProを提案します。

1. 線形空間確率緩和:
   $M$ 個の重みのグループに対する $\binom{M}{N}$ 個のすべての可能なマスクの離散分布を定義する代わりに、MaskPro は問題を再定式化します。各グループ内の $M$ 個の個々の要素に対する離散分布を確立します。(N:M) マスクは、この分布から非復元 N 回サンプリングを実行することで生成されます。

   • 理論的基盤: 著者らは、任意の (N:M) マスクが $N$ 個の異なる基底ベクトルの確率的和として表現できることを証明しています。
   • メモリ効率: この再定式化により、学習可能なパラメータ（ロジット）の数を $O(\binom{M}{N} \frac{d}{M})$ から線形の $O(d)$ に削減します（ここで $d$ は重みの総数です）。これにより、以前の学習ベース手法に関連するメモリの爆発的増加が排除されます。

2. 洗練された方策勾配推定量（PGE）:
   ロジットを最適化するために、MaskPro は方策勾配を利用します。著者らは、ミニバッチの違いによる分散（データサンプリングからの損失変動が、マスクサンプリングからの損失変動をしばしば凌駕する）により、LLM において生の損失値を用いた標準的な方策勾配が不安定になることを特定しました。

   • 損失残差: 生の損失 $f(m \odot w, \xi)$ を使用する代わりに、MaskPro は初期基準マスク $m_0$ に対する損失残差に基づいてロジットを更新します：$f(m \odot w, \xi) - f(m_0 \odot w, \xi)$。これにより、マスク変更の影響を分離します。
   • 移動平均トラッカー: トレーニングをさらに安定させ、数値的な爆発を防ぐために、基準を平滑化する移動平均トラッカー $\delta$ が導入されます。更新則は以下のようになります：
     $$\pi_{t+1} = \pi_t - \eta (f(m_t \odot w, \xi) - f(m_0 \odot w, \xi) - \delta) \nabla \log p(m_t|\pi_t)$$
   • 理論的保証: 本論文は、これらの洗練された推定量が真の方策勾配の不偏推定量であることを示す証明を提供しています。さらに、サンプリングされたマスクの損失が基準損失の 2 倍未満であるという条件（実際には容易に満たされる）の下では、洗練された推定量の分散は、標準的な方策勾配の分散よりも厳密に低いことが示されています。

主要な貢献

1. 線形空間フレームワーク: MaskPro は、$M$ 個の要素に対する離散分布からの非復元 N 回サンプリングとして選択プロセスをモデル化することにより、(N:M) 疎性の学習に必要なメモリ要件を、指数関数的な $O(\binom{M}{N} \frac{d}{M})$ から線形の $O(d)$ に削減します。
2. 強化された方策勾配: 著者らは、損失残差と移動平均トラッカーを用いた分散低減更新メカニズムを導入し、大規模な組み合わせ空間における方策勾配の高い分散によって引き起こされるトレーニングの不安定さに対処しました。
3. 理論的解析: 推定量の不偏性と、提案された更新則の分散低減特性に関する包括的な証明が提供されています。

実験結果
Gemma-7B、Vicuna-7B、LLaMA-2-7B、DeepSeek-7B といった複数の LLM において、HellaSwag、RACE、PIQA などのさまざまな下流タスクおよび (2:4)、(4:8)、(8:16) などの疎性パターンに対して、広範な実験が行われました。

• 性能: MaskPro は、SparseGPT、Wanda、Pruner-Zero などの最先端のルールベース手法を一貫して上回り、メモリ集約的な MaskLLM と同等の性能を達成します。多くの場合、精度で平均 2% 以上の改善が見られます。
• メモリスケーラビリティ: MaskLLM は 7B モデルの (4:8) 疎性において 8x A100 GPU で Out-of-Memory (OOM) エラーに遭遇するのに対し、MaskPro は疎性パターンに関わらず安定したメモリ使用量（約 36-37 GB）を維持します。
• データ効率: MaskPro はデータ量に対する顕著なロバスト性を示します。MaskLLM が収束するために数十万のサンプルを必要とするのに対し、MaskPro はわずか 1 つのデータサンプルでトレーニングしても安定した性能を達成します。
• トレーニング効率: この手法は、MaskLLM と比較してはるかに少ないトレーニングサンプルとトレーニング時間を必要としながら、ルールベース手法と同等のメモリ使用量を維持します。

意義と主張
本論文は、MaskPro が学習ベースの半構造化疎性のメモリおよび安定性のボトルネックを解決することで、LLM の実用的でスケーラブルなローカル展開への viable な道筋を提供すると主張しています。著者らは、このアプローチが学習ベース手法の高い性能と、ルールベース手法の効率性の間のギャップを埋めていることを強調しています。メモリフットプリントを線形空間に削減し、洗練された方策勾配を通じてトレーニングプロセスを安定化させることで、MaskPro は、これまでそのようなタスクに関連していた prohibitive な計算オーバーヘッドなしに、大規模モデル上で最適な (N:M) マスクを学習することを可能にします。この手法は、データ不足やモデルの規模に対してロバストな、確率的マスク学習のための汎用フレームワークとして提示されています。