OPTIMA: Optimal One-shot Pruning for LLMs via Quadratic Programming Reconstruction

技術要約: OPTIMA: 二次計画法による再構成を用いたLLMのための最適なワンショット・プルーニング

問題提起

ポストトレーニングにおけるワンショット・プルーニングは、大規模言語モデル（LLM）の推論コストを削減することを目的としており、フルファインチューニングのコストを回避しながら、少量のキャリブレーションデータセットのみを用いてパラメータを削除します。しかし、この分野は根本的なトレードオフに直面しています。

• ヒューリスティック手法（例：Wanda）は計算コストは低いものの、残りの重みを更新せずに単に重みをゼロにするため、精度が大幅に低下することがよくあります。
• 原理的な結合最適化手法（例：Optimal Brain Surgeon）は、最適な重み更新を解くことで精度を回復できますが、巨大なヘシアン（Hessian）の逆行列計算の複雑さにより、現代のLLMの規模では計算量的に実行不可能です。
• 既存の近似手法（例：SparseGPT、Thanos）は、このギャップを埋めようと試みていますが、再構成の最適性を犠牲にしており、回復可能な精度の「余地」を残しています。

核心となる課題は、精度（再構成問題を最適に解くことによる）とスケーラビリティ（現代のアクセラレータ上で計算可能なことによる）のバランスをとる手法を開発することです。

手法

本論文は、マスク選択後のレイヤーごとの重み再構成ステップを、独立した列ごとの二次計画問題（QP）の集合として定式化するフレームワークであるOPTIMAを導入します。

1. 数学的定式化:
重み行列 $W_l$ に対してバイナリマスク $M_l$ が選択された後、目標は出力再構成誤差 $\|X_l W_l - X_l (M_l \odot \hat{W}_l)\|_F^2$ を最小化する更新された重み $\hat{W}_l$ を見つけることです。

• OPTIMAはこの問題を列ごとに分解します。各列 $j$ について、更新 $\Delta w_j$ は以下のQPとして定式化されます：
  $$\min_{\Delta w_j} \Delta w_j^T H \Delta w_j \quad \text{subject to} \quad (\Delta w_j)_i = -(w_j)_i, \forall i \in S_j$$
  ここで、$H = X_l^T X_l$ はヘシアン（入力アクティベーションの共分散）であり、$S_j$ はプルーニングされた重みのインデックスを表します。
• 重要な洞察: 同じ重み行列のすべての列は、同じヘシアン $H$ を共有しています。異なるのは線形制約（マスクによって決定される）のみです。この共有構造により、効率的なバッチ処理が可能になります。

2. ソルバーと実装:

• アルゴリズム: OPTIMAは、制約付き問題に適しており、共有ヘシアンとの行列ベクトル積に依存する第一次の原始・双対ソルバー、具体的には rAPDHG（再起動加速原始・双対ハイブリッド勾配法）を採用しています。
• アクセラレータの最適化:
  • バッチ処理: 数千の小さなQPを、単一のアクセラレータ（GPU/TPU）上で並列に解きます。
  • メモリ効率: 密なアクティベーション行列全体を保持する代わりに、ヘシアンをキャリブレーションシーケンスから逐次的に累積します。
  • 制約の処理: プルーニングされたエントリは、明示的な等式行列ではなく、タイトな変数境界を通じて強制されるため、大きな制約行列のオーバーヘッドを回避できます。
  • フォールバック・メカニズム: 列のQPがイテレーション予算内に収束しない場合（収束証明書によってチェック）、その更新は破棄され、元の重み（プルーニングされたエントリをゼロにしたもの）が保持されます。これにより、この手法がベースラインのマスクセレクターの更新よりも悪化しないことが保証されます。

主な貢献

1. QPへの定式化: 重みの再構成を、単一のヘシアンを共有するバッチ化された列ごとのQPとして定式化することで、再構成目的関数に対する列ごとのグローバルな最適性を可能にしました。
2. スケーラブルなソルバー・パイプライン: ヘシアンを逐次的に累積し、境界によって制約を強制し、rAPDHGを使用して数千のQPを並列に解く、アクセラレータフレンドリーなパイプラインを設計しました。
3. モジュール性: OPTIMAが、既存のマスクセレクター（Wanda、SparseGPT、Thanos、ProxSparse、ADMM）と互換性のある、「ドロップイン」可能な重み更新ステップとして機能することを実証しました。
4. 実証的検証: LLaMA、Gemma、Qwenファミリーにわたる広範な結果を提供し、ファインチューニングなしでの一貫した精度向上を示しました。

結果

著者らは、非構造化（50%、60%）および半構造化（2:4）スパース性レジームにおいて、LLaMA 3.1/3.2、Gemma 2/3、Qwen 2.5モデルを用いてOPTIMAを評価しました。

• 精度向上: OPTIMAは、49個の（モデル、マスクセレクター、スパース性）の設定のうち40個において、平均ゼロショット精度を向上させました。
  • 平均利得：50%非構造化で +0.61 ポイント、2:4で +0.56 ポイント、60%非構造化スパース性で +1.01 ポイント。
  • 最大利得：LLaMA 3.2 3B（60%スパース性）で +3.97 ポイント。
  • Qwen 2.5では、利得は +4.19 ポイントに達しました。
• ベースラインとの比較:
  • OPTIMAは、「更新なし」のベースライン（Wanda）を一貫して上回りました。
  • 重み更新を行うベースライン（SparseGPTやThanos）に対しても、ほとんどの設定で改善が見られましたが、ベースラインが既に重み更新を行っている場合には、利得は小さくなりました。
  • 独自のマスクを使用した場合、OPTIMAはADMMベースの厳密な再構成（Boža, 2024）と同等またはそれをわずかに上回りました。
• パープレキシティ vs 精度: 本論文では、パープレキシティとダウンストリームの精度が必ずしも相関しないことを指摘しています。OPTIMAはダウンストリームの精度を優先しており、これはこの文脈におけるプルーニングの品質の主要な指標です。
• 効率性:
  • 8Bパラメータのモデルを単一のNVIDIA H100でプルーニングするには、40時間（ピーク時60 GBメモリ）かかります。
  • 70Bモデルのプルーニングには、約350時間かかります。
  • レイヤーあたりの計算時間はヒューリスティックな手法（Wanda/SparseGPT）よりも遅いですが、これは一度限りのオフラインコストです。

重要性と主張

本論文は、原理的な列ごとのQP再構成が、ワンショット・ポストトレーニング・プルーニングにおけるヒューリスティックな重み更新に代わる実用的な選択肢であると主張しています。

• ギャップの解消: OPTIMAは、共有ヘシアン構造を利用してアクセラレータ上での並列かつ保証された解決を可能にすることで、高速だが不正確なヒューリスティックと、低速で実行不可能な厳密な手法の間のギャップを埋めます。
• ファインチューニング不要: この手法はファインチューニングなしでこれらの利得を実現しており、純粋な推論コスト削減技術となっています。
• スケーラビリティ: 分散トレーニングインフラを必要とせず、単一のアクセラレータ上で問題を解くことにより、OPTIMAは厳密な再構成が大規模モデルに対しても実現可能であることを示しています。
• モジュール性: このアプローチはマスク選択とは直交しており、固定されたマスクに依存する既存のあらゆるプルーニング戦略を強化できることを意味します。

著者らは、OPTIMAが閉形式の更新よりも遅いこと、およびレイヤーごとの再構成誤差の最小化が必ずしもエンドツーエンドのタスク性能と完全に一致しないことを認め、限界についても触れています。しかし、彼らはこの手法が、プルーニングされたモデルの有用性を最大化するための、堅牢で理論的根拠に基づいたアプローチを提供すると主張しています。