From Layers to Submodules: Rethinking Granularity in Replacement-Based LLM Compression

技術要約：レイヤーからサブモジュールへ (SUBFIT)

問題提起
大規模言語モデル（LLM）のポストトレーニング圧縮は、アーキテクチャ構成要素を削除することでデプロイコストを削減することを目的としている。層全体や連続したブロックを削除する手法は効果的であるが、残差ストリーム（residual stream）を断絶させ、しばしば急激な品質低下を引き起こす。近年の「置換ベース」のアプローチは、削除されたコンポーネントを軽量な適合済みモジュールで置き換えることで、この問題に対処している。しかし、既存の手法には2つの制約的な設計上の制約がある：

1. フルレイヤー粒度： トランスフォーマー層全体を削除の原子単位として扱う。
2. 連続的な選択： 連続した深度範囲内でコンポーネントを削除する。

著者らは、これらの制約は過度に制限的であると主張している。学習済みトランスフォーマーにおける冗長性は、連続した領域に限定されているわけではなく、また、AttentionとFeedForward（FFN）の出力間で均等に分布しているわけでもない。したがって、異なるサブモジュール・タイプを近似するためには異なる戦略が必要であり、削除可能なコンポーネクションは連続した深度範囲内に集まる必要はない。

手法：SUBFIT
本論文では、個々のAttentionおよびFFNサブモジュールの粒度で動作するポストトレーニング圧縮手法である SUBFIT (Submodule-level Fitted residual replacement) を導入する。プロセスは主に2つのステージで構成される：

1. 非連続的な選択 (Non-Contiguous Selection):

   • SUBFITは、スパース性予算 $s$ に基づき、モデルの深さ全体にわたって独立して、対象となる数のAttentionおよびFFNサブモジュールを選択する。
   • 選択は非連続的であり、モジュールは連続したブロックを形成することなく、スコアによって独立して選ばれる。
   • スコアリング戦略:
     • Attention: モジュールは、キャリブレーション用の活性化値に基づいて計算された「インパクト・スコア」によってランク付けされる。これは、サブモジュールが残差ストリームをどれだけ回転させるか、およびその相対的な大きさ（magnitude）を測定するものである。
     • FFN: ベースモデルについては、削除の影響とサロゲート誤差を組み合わせた「置換を考慮したスコア」が使用される。インストラクション・チューニング済みモデルについては、不適切なキャリブレーションによる残差強度のノイズを避けるため、コサイン類似度のみを用いたスコアが使用される。
     • 順序付け: Attentionサブモジュールが先に選択・置換される。その後、Attentionの置換によって誘発される残差ストリームのシフトを考慮するために、FFNサブモジュールがスコアリングおよび選択される。

2. 適合済み残差置換 (Fitted Residual Replacement):

   • 選択された各サブモジュールは、内部の計算ではなく、元のサブモジュールの残差寄与 ($y = F_\ell(x)$) を近似する軽量なパラメータ化されたサロゲート $\hat{F}_\ell$ によって置換される。
   • サロゲート構造: サロゲートは、アフィン補正（要素ごとのゲート、バイアス、およびセンタリング）と、低ランク演算子 $U_\ell = V_\ell^\top W_\ell$ によって定義される構造化された相互作用項を使用する。
   • 非対称なパラメータ化:
     • Attention: 各削除されたAttentionサブモジュールは、深さによる異質な役割を反映して、それぞれ独立した低ランク基底 ($V_\ell$) とパラメータを受け取る。
     • FFN: すべての選択されたFFNレイヤーに対して、共有された入力基底 $V$ が使用される。出力因子 ($W_\ell$) とアフィン項のみがレイヤー固有のまま残る。このクロスレイヤーのパラメータ共有により、デプロイ時のコストが大幅に削減される。
   • フィッティング: パラメータは、キャリブレーション統計量（リッジ正則化最小二乗法および固有値分解）を用いた閉形式のフィッティングを通じて解析的に回収されるため、勾配ベースの最適化やフルモデルのファインチューニングを必要としない。

主な貢献

1. 制約の特定: 著者らは、フルレイヤーの粒度および連続的な選択が、既存の置換ベースの圧縮における制限要因であることを特定した。
2. SUBFITフレームワーク: 非連続的な選択を伴う、サブモジュールレベル（AttentionおよびFFN）での結合された選択と置換を行う新しい手法。
3. 非対称設計: Attentionのバイパスは独立した基底を使用する一方で、FFNのバイパスは共通の入力基界を共有してパラメータ・オーバーヘッドを最小限に抑えるという、特定のアーキテクチャ。
4. 包括的な評価: 10のLLM（ベースモデル5つ、インストラクション・チューニング済みモデル5つ）および5つのスパース性レベル（12.5%から37.5%）にわたる広範なテストを行い、4つのベースライン（Streamline, ReplaceMe）に対して改善されたトレードオフを示した。

実験結果

• パープレキシティ (ベースモデル): SUBFITは、10のモデル-スパース性設定のうち9つにおいて、最も低いパープレキシティの劣化を達成した。最強のベースラインとの差は、アグレッシブな圧縮下で顕著に拡大する（例：25%のスパース性において、SUBFITのパープレキシティ劣化は2.42倍であるのに対し、最強のベースラインは4.34倍である）。
• ダウンストリーム精度 (インストラクション・チューニング済み): 25%のスパース性において、SUBFITは平均して**84.6%**の密なダウンストリーム精度を維持し、最強のベースライン（81.6%）を上回った。SUBFITは、25%のスパース性で80%以上の精度を、37.5%のスパース性で73%以上の精度を維持できる唯一の手法である。
• 安定性: SUBFITは、高スパース性において不安定性を示す（例：標準偏差が数桁増大する）ベースラインと比較して、パープレキシティの劣化における分散が著しく低い。
• 効率性: 本手法は、スパース・カーネルやカスタムハードウェアを必要とせずに、測定可能な推論速度向上（最初のトークン生成までの時間で1.18倍〜1.40倍）およびKVキャッシュの節約（スパース性に比例）をもたらす。
• アブレーション研究:
  • 残差補償: FFNの置換が安定性の主要因であり、AttentionとFFNの両方を置換することが最良の結果をもたらす。
  • 共有基底: 共有FFN基底は、パラメータ数を大幅に削減しながら、レイヤーごとの基底と同等の性能を達成する。
  • ランクの感度: FFNの品質はランク4096まで向上する一方、Attentionはランク256で飽和し、この非対称なランク設定を正当化している。

意義と主張
本論文は、粒度がポストトレーニング圧縮における未探索の軸であると主張している。レイヤーレベルの置換からサブモジュールレベルの置換へと移行することで、SUBFITは同じパラメータ予算の下で、連続的なフルレイヤー置換よりも密なモデルの挙動をより多く回復する。結果は、非連続的かつサブモジュール固有の回復が、特にアグレッシブな圧縮において、より強力な総合的な品質・効率のトレードオフを提供することを示唆している。本手法は、デプロイが容易（閉形式のフィッティング）であり、既存の推論スタックと互換性があるように設計されている。著者らは、このアプローチが他の除去ユニット（例えば、Attentionヘッドや個々のニューロン）にも一般化する可能性があることを控えめに示唆しているが、本研究においてそれらの特定のユニットへの即時的な一般化を主張するものではない。