PATCH: Learnable Tile-level Hybrid Sparsity for LLMs

以下は、論文「PATCH: LEARNABLE TILE-LEVEL HYBRID SPARSITY FOR LLMs」の詳細な技術的サマリーです。

1. 問題定義

大規模言語モデル（LLM）は優れた性能を発揮しますが、展開時に莫大なメモリおよび計算コストという課題を抱えています。モデルプルーニングは主要な解決策の一つですが、既存のアプローチは根本的なトレードオフに直面しています。

• 非構造化スパース性: 任意の位置に非ゼロ要素を許容するため高精度を維持できますが、不規則なメモリアクセスパターンを生み出し、効率的な GPU 加速（例：Tensor Cores）を阻害します。
• 半構造化スパース性（例：2:4）: 4 要素あたり 2 つの非ゼロ要素という厳格なパターンを強制するため、ハードウェアに親和性が高く推論を加速します。しかし、固定された 50% のスパース性比率と全層への均一な適用は、特にワンショットプルーニング手法を使用する場合、顕著な精度低下を招くことがよくあります。

現在の手法は、非構造化プルーニングの柔軟性と半構造化プルーニングのハードウェア効率性の間のギャップを埋めることに苦慮しており、特にスパース性比率を連続的に適応させたり、層間で非均一にスパース性を割り当てたりする能力の欠如が問題となっています。

2. 手法：PATCH

著者らは、ハードウェア互換性を維持しつつ 0% から 50% の連続的なスパース性比率を実現するフレームワーク、PATCH（Pruning with a Learnable Tile-level Configuration for Hybrid Sparsity）を提案します。

中核メカニズム

PATCH は重み行列をタイル（例：$b_1 \times b_2$）に分割します。各タイルに対して、モデルは「密（0% スパース性）」のままにするか「2:4 疎（50% スパース性）」にするかを学習して決定します。

• ハイブリッドマスク生成: 最終的なマスクは、「密」オプションと「2:4 疎」オプションの重み付き組み合わせです。
• 学習可能な分布:
  1. タイル選択: 学習可能なログイト分布が、タイルが密か疎かの確率を決定します。これは微分可能な最適化を可能にするため、Gumbel-Softmaxを用いてサンプリングされます。
  2. パターン選択: 疎と指定されたタイル内では、別の学習可能な分布が、6 つの可能な置換の中から特定の 2:4 パターンを選択します。
• 最適化目的: 訓練目的関数は以下の組み合わせです。
  • 標準的なモデリング損失（例：次のトークン予測）。
  • スパース性正則化: 目標となるグローバルなスパース性比率（$\rho$）からの逸脱を罰する項であり、0–50% の範囲を精密に制御可能にします。
  • 重み正則化: 勾配伝播を支援するため、より大きな重みの大きさを促進します。

変種

• PATCHJoint: タイルレベルの選択（密 vs 疎）と、疎タイル内の微細な 2:4 パターンの両方を共同で最適化します。10 億パラメータ未満の小型モデルに使用されます。
• PATCHTile: メモリ効率に優れた変種です。ここでは 2:4 パターンは高品質なワンショット手法（MaskLLM など）から初期化して固定され、タイルレベルの密/疎の決定のみが最適化されます。これにより、限られた GPU メモリ環境でも 130 億パラメータまでの大規模モデルへのスケーリングが可能になります。

推論と展開

• STOICC 統合: 標準的な GPU ライブラリ（cuBLAS、cuSPARSELt）は混合された密/疎タイルをサポートしていないため、PATCH は Triton ベースのコンパイラであるSTOICCを利用します。STOICC はカーネル設定（タイルサイズ、例：$128 \times 128$）を自動調整し、ハイブリッド実行を効率的に処理します。
• ハードウェア互換性: 疎タイルが厳密に 2:4 パターンに準拠することを保証することで、PATCH は NVIDIA/AMD の Tensor Cores を利用して加速を実現し、密タイルは標準的な密カーネルを利用します。

3. 主要な貢献

1. ハイブリッドスパース性フレームワーク: 同じ重み行列内で密タイルと 2:4 疎タイルを動的に混合することで、**連続的なスパース性比率（0%–50%）**を実現する最初の手法を導入しました。
2. 学習可能な非均一割り当て: 固定された 2:4 手法とは異なり、PATCH は適応的にスパース性を割り当てることを学習します。敏感な層（例：Transformer ブロックの初期/後期、Attention メカニズム）では密度を維持し、冗長な層（例：中間ブロック、MLP のアップ/ゲート/ダウン行列）にはスパース性を適用します。
3. エンドツーエンドの微分可能な訓練: Gumbel-Softmax を使用してマスク選択と重みプルーニングを共同最適化し、ワンショットプルーニングに伴う精度低下を回避します。
4. 実用的な展開: ハイブリッドタイルレベルスパース性とコンパイラ（STOICC）の統合を初めて成功させ、一般消費者向け GPU 上で実世界の速度向上を実現しました。

4. 実験結果

著者らは、Qwen-2.5、LLaMA-2/3、Gemma-3 などの 0.5B から 13B パラメータのモデルに対し、8 つのゼロショット下流タイルと WikiText2 のパープレキシティを用いて PATCH を評価しました。

• 精度:

  • PATCH は、最先端の 2:4 プルーニング手法（MaskLLM、SparseGPT、Wanda、ProxSparse）を一貫して上回ります。
  • LLaMA-2 7Bにおいて、PATCH（25% スパース性）は平均精度**51.58%**を達成し、50% スパース性の MaskLLM（48.62%）や ProxSparse（45.92%）を上回りました。
  • PATCH は密モデルとの差を大幅に縮小します。25% スパース性において、密モデルの精度から約 2 パーセントポイント以内に回復します。
  • パープレキシティ: PATCH はすべてのベースラインよりも低いパープレキシティを達成します。LLaMA-2 7B（25% スパース性）の場合、PPL は5.86であり、MaskLLM の6.78より優れています。

• 速度向上と効率性:

  • NVIDIA A6000 GPU（LLaMA-2 7B）において、PATCH は密モデルベースラインに対して1.18 倍から 1.38 倍のエンドツーエンドの速度向上を実現しました。
  • メモリフットプリントは密モデルの0.59 倍–0.76 倍に削減されました。
  • PATCH は、速度向上を提供しない非構造化プルーニングや、速度向上を正当化するほど精度低下が激しいことが多い硬直的な 2:4 プルーニングを上回ります。

• アブレーション研究:

  • グローバル vs レイヤーワイド: 層ごとの可変スパース性を許容するグローバルなスパース性目標は、均一なレイヤーワイド目標よりも大幅に優れています。
  • タイルサイズ: 性能はタイルサイズに対して頑健ですが、$4 \times 4$が最も細かい制御を提供します。ハードウェア効率性から、より大きな実行タイル（$128 \times 128$）を使用することが推奨されます。
  • 初期化: この手法は初期化の事前分布に対して頑健であり、グローバルなスパース性目標は訓練中に動的な再割り当てを可能にします。

5. 意義

PATCH は、「精度対加速」のジレンマを解決することで、効率的な LLM 推論における重要な前進を表しています。

• ギャップの埋め合わせ: ハードウェア加速の恩恵をもたらす半構造化スパース性と、柔軟な非構造化割り当てを成功裡に組み合わせました。
• 実用性: 理論的なプルーニングを超え、既存のコンパイラ基盤（STOICC）を用いた実ハードウェア上での実際の速度向上を実証しました。
• スケーラビリティ: メモリ効率に優れた変種（PATCHTile）は、マスクの訓練に莫大な計算予算を必要とすることなく、大規模モデル（13B 以上）にもハイブリッドスパース性を適用可能であることを証明しました。
• 将来の方向性: このフレームワークはパターンに依存せず、新しいハードウェアがサポートするにつれて理論的には他の N:M パターン（例：4:8）にも拡張可能であり、将来の LLM 圧縮のための多用途なソリューションとなります。

要約すると、PATCH は、モデルの品質と推論速度のトレードオフを前例のない粒度で調整することを可能にする、柔軟で学習可能かつハードウェア効率的なプルーニング戦略を提供します。