Batch-wise Adaptive Pruning: Periodic Neuron Activation-Aware Weight Pruning for Language Reasoning Model

技術要約：大規模推論モデルのためのバッチ単位適応型プルーニング

問題提起
大規模推論モデル（LRM）は、長い思考連鎖（Chain-of-Thought）の生成を通じて複雑なタスクで高い性能を達成しますが、膨大な計算コストを伴います。プロダクション環境では、高スループットを実現するためにバッチ推論（複数のサンプルを同時に処理すること）が必要です。しかし、TEALやCATSといった既存の学習不要な適応型プルーニング手法は、このレジームにおいて著しく失敗します。

その核心的な失敗メカニズムは二重の要因によるものです：

1. 分布シフト（Distribution Shift）： 既存の手法は、バッチ全体に対して単一のプルーニングマスクを使用します。これらは、個々のサンプルの活性化（例：平均値）を集計し、単一サンプルの分布に基づいてオフラインで較正された閾値を適用します。この集計は活性化分布をシフトさせ、閾値がデータに適合しなくなるため、実現されるスパース性比率にドリフトが生じ、推論精度の崩壊を招きます。
2. 静的対動的トレードオフ（Static vs. Dynamic Trade-off）： 静的なプルーニング手法（固定パターン）は、長い推論過程における進化する活性化パターンを追跡できません。一方で、毎トークンのステップごとに新しいマスクを選択することは計算コストが極めて高く、高速化のメリットを打ち消してしまいます。

手法
著者らは、LRMのバッチ推論向けに特別に設計された、学習不要のバッチ単位適応型プルーニングを提案します。このアプローチは、主に2つのコンポーネントに基づいています：

1. 周期的なTop-k選択：

   • 閾値ベースの選択（分布シフトに敏感）の代わりに、集計された重要度スコアに対してTop-k選択を採用します。Top-k選択はスコアの相対的な順序のみに依存するため、バッチ集計によって引き起こされる分布シフトに対して堅牢です。
   • 全てのデコーディングステップでマスクを更新する計算オーバーヘッドを避けるため、プルーニングマスクは各デコーディングステップではなく、周期的に（$T_{trans}$ ステップごと）更新されます。
   • バッチ集計： $B$ 個のサンプルからなるバッチに対して、個々のサンプルの重要度スコアは**要素ごとの最大値（element-wise maximum）**を用いて単一のベクトルに集計されます（平均ではなく）。この保守的なアプローチにより、バッチ内の他のサンプルで活性化が低かったとしても、少なくとも一つのサンプルにとって重要なニューロンが保持され、プルーニングされるのを防ぎます。

2. 活性化メモリ（Activation Memory）：

   • LLMの活性化パターンの観察により、重要なニューロンは長い推論生成中に**周期的な再発火（periodic re-firing）**を示すことが明らかになっています（例：約22トークンごとに再発火し、再発火率は約76%）。
   • これを捉えるために、本手法は重要度スコアを蓄積する活性化メモリを保持します。メモリは要素ごとの最大値（$m_n \leftarrow \max(m_{n-1}, s_n)$）によって更新されます。
   • これにより、以前のフェーズで重要であると特定されたニューロンが、一時的に活性化が低下したとしても、保持候補として残り続け、重要なコンポーネントを失うことなく進化するパターンに適応できるようになります。

アルゴリズムの流れ
本手法は、自己回帰生成中に3つのフェーズで作動します：

1. 入力フェーズ： 初期プロンプトを処理し、初期の重要度スコアを計算します。
2. 高密度探索フェーズ（Dense Exploration Phase）： 推論の方向性を確立し、コンテキスト固有の活性化統計を収集するために、一定数の高密度ステップ（$T_{init}$）を実行します。ここでメモリが更新され、最初のプルーニングマスクが生成されます。
3. 周期的な適応型プルーニングフェーズ： スパースな計算（$T_p$ ステップ）と高密度な探索（$T_E$ ステップ）を交互に繰り返します。各高密度探索ステージの終了時に、活性化メモリが更新され、蓄積されたメモリに基づいて新しいTop-kマスクが生成されます。

主な貢献

1. バッチによる崩壊の経験的特定： 本論文は、既存の学習不要なプルーニング手法（静的および適応型の両方）が、バッチ推論下での推論タスクにおいて著しく精度を低下させることを初めて経験的に示しました。具体的には、適応型手法であるTEALは、バッチサイズ4において、高密度モデルと比較して平均精度が58.4〜67.7ポイント低下します。
2. 堅牢なバッチ・プルーニング・フレームワーク： 著者らは、閾値ベースの選択を周期的なTop-k選択に置き換え、再発火するニューロンの重要性を追跡するための活性化メモリを組み込むことで、バッチ設定下でも精度を維持する手法を導入しました。
3. 較正不要（Calibration-Free）な動作： 外部コーパスを用いたオフライン較正（WandaやTEALなど）を必要とする手法とは異なり、本手法は実行時の活性化のみに依存するため、較正が不要です。

実験結果
実験は、DeepSeek-R1-Distill-Qwen-7BおよびDeepSeek-R1-Distill-Llama-8Bを用い、5つの推論ベンチマーク（GSM8K, MATH500, MINERVA, AMC23, GPQA-DIAMOND）に対して行われました。

• 精度： ターゲットスパース性50%、バッチサイズ4において、提案手法はQwen-7Bモデルで従来の最先端（TEAL）を平均精度で39.7ポイント（54.0% vs. 14.3%）上回り、Llama-8Bモデルでは24.1ポイント上回りました。
• 安定性： バッチサイズ1では高い性能を示すもののバッチサイズが増加すると崩壊するTEALとは異なり、提案手法はすべてのバッチサイズ（1, 4, 8, 16）において一貫した性能を維持します。
• 高速化： Qwen-7Bモデルにおいて、実際のスパース性50%で高密度推論に対して1.40倍の高速化を達成しました。この高速化はバッチ設定においてより顕著です。なぜなら、本手法は重み行列の次元を構造的に削減して計算量を直接減らすのに対し、ゼロの活性化をスキップする手法は、ワークロードがメモリ帯域制限から計算量制限へと移行するにつれて効果が低下するためです。
• 高スパース性： 提案手法は、ターゲットスパース性が高いレベル（例：ターゲットスパース性50%で71.0%の精度）でも強力な性能を維持しますが、ベースラインは40%のスパース性を超えると急激に劣化します。

意義と主張
本論文は、その手法が大規模推論モデルの展開における決定的なボトルネック、すなわち既存の適応型プルーニング技術とバッチ推論との不適合性を解決することを主張しています。バッチ集計に内在する分布シフトの問題を解決し、ニューロンの再発火の周期性を利用することで、本手法は、従来の先行手法で見られた性能崩壊を起こすことなく、効率的で高スループットな推論を可能にします。著者らは、これらの利得が、集計演算子の具体的な選択（max vs. mean）ではなく、周期的なTop-k選択と活性化メモリによってもたらされていることを強調しています。