技術要約: TraceNAS

問題定義

大規模言語モデル（LLM）のデプロイは、膨大なパラメータ数と計算要件によって制約を受けている。構造化プルーニングは、アーキテクチャのユニット（例：アテンションヘッド、トランスフォーマーブロック）を削除することで効率的なデプロイへの道筋を提供するが、既存の手法には以下のような重大な限界がある：

1. 局所的依存関係 vs 全局的依存関係: ほとんどのプルーニング手法は、局所的なヒューリスティック（例：活性化重み付きマグニチュードやヘシアンに基づく再構成）を用いて、層、ヘッド、またはチャネルの重要性を個別に評価する。これらのアプローチは、複雑なグローバルな構造的依存関係や、プルーニングによって誘発されるエラーのモデル全体への伝播を無視している。
2. 学習を伴う探索の計算コスト: グローバルな依存関係を捉えようとする手法は、多くの場合、学習を伴うニューラルアーキテクチャ探索（NAS）に依存している。これらのアプローチは、オーバーフィッティングを防ぐために反復的な勾配更新と未学習のトークンの継続的なストリームを必要とし、その探索プロセスは、プルーニング後のリカバリ学習自体と同じくらい計算コストがかかる。
3. 既存のゼロショット・プロキシの不適切さ: 現在の学習フリーのプロキシは、学習可能性（例：勾配の安定性、損失曲面の滑らかさ）に焦点を当て、学習可能なランダム初期化アーキテクチャを特定するように設計されている。これらは「機能的継承（functional inheritance）」、すなわち、事前学習済みの損失曲面を保持することによってプルーニングされたモデルが性能を回復できる能力を評価できていない。

メソドロジー

TraceNASは、探索時の学習を伴わずに、LLMの深さと幅を共同で最適化するように設計された、学習フリーのNASフレームワークとして提案されている。

1. 探索空間とエンコーディング

本フレームワークは、事前学習済みのLLMをスーパーネットワークとして扱う。探索空間 $G$ は、以下に関する結合分布として定義される：

• 深さ ($d$): $L$ 個のトランスフォーマーブロックに対する離散的なマスクであり、残差ストリームの連続性を維持しながら、非一様な深さの削減を可能にする。
• 幅 ($\kappa$): アテンションおよびMLPサブブロックのパラメータ保持率であり、モデル全体にわたるヘテロジニアス（不均一）なスパース性を可能にする。

2. インプレースによるアーキテクチャの実装

多数の候補モデルをインスタンス化するメモリオーバーヘッドを避けるため、TraceNASはインプレース・マスキング戦略を採用している：

• 幅のマスキング: 重みチャネルに対してバイナリマスクを生成するために、活性化重み付き積（Wandaに類似）を使用する。この $O(d^2)$ のヒューリスティックは、アウトライヤー特徴量を保持しながら、機能的に重要なチャネルを特定する。
• 深さのマスキング: ブロックが無効化された場合は、そのブロック全体をバイパスし、残差接続を通じて活性化をルーティングする。
• 実行: フレームワークは、最小限のキャリブレーションセットに対して勾配を計算するためにベースモデルの重み行列を一時的に変更し、その後元の重みを復元することで、後続のイテレーションのためにベースモデルを保持する。

3. ゼロショット・プロキシ: 勾配トレース相関

コアとなる革新は、プルーニングされた候補 ($M_{sub}$) と事前学習済みのベースモデル ($M_{base}$) の間の機能的整合性を測定する、スケール不変なプロキシ $\Phi$ である。

• 勾配トレース: キャリブレーションセットに対して勾配が計算される。スケーラビリティを確保するため、これらの勾配は低ランク（LoRA）部分空間に投影され、LLMの固有の低次元性を活用することで、メモリオーバーヘッドを数桁削減する。
• 整合性指標: プロキシは、各層におけるプルーブルドモデルとベースモデルの標準化された勾配トレース間のピアソン相関係数を計算する。これによりスケール不変性が確保され、マグニチュードの変化ではなく方向的な整合性に焦点が当てられる。
• スパース性重み付き集計: 最終的なスコア $\Phi$ は、層ごとの相関の加重和であり、ここで重みはパラメータ保持率に対応する。これにより、グローバルなスコアが高容量の領域に固定され、極端にプルーニングされたサブブロックからのノイズがランキングを歪めるのを防ぐ。

4. 進化的探索

TraceNASは、離散的な探索空間をナビゲートするために進化的戦略を採用している：

• 初期化: グローバルな重要度の事前分布に基づいて集団が生成される。
• 進化: 集団は交叉と突然変異を行う。深さの進化にはビット反転突然変異と連続範囲の交叉を用い、幅の進化には補間とガウスジッターを用いる。
• 選択: 集団はプロキシ $\Phi$ のみによってランク付けされ、プルーニング後の継続的事前学習（CPT）における性能回復の可能性が最も高いモデルを選択する。

主な貢献

1. スケーラブルな学習フリーNAS: TraceNASは、学習時のトレーニングを伴わずに、非一様なアーキテクチャの発見を可能にする、深さと幅の共同プルーニングのための統一されたフレームワークを提供する。
2. 機能的継承プロキシ: プルーニングされたモデルと事前学習済みモデルの間の勾配トレースの整合性を評価する、新しいスケール不変な勾配ベースのプロキシを導入した。これは、ダウンストリームの精度に対して高い相関（Spearman $\rho = 0.94$, Kendall $\tau = 0.82$）を達成している。
3. 効率性: 低ランク勾配部分空間とインプレース・マスキングを利用することで、単一のGPU上で高精度なモデル発見を可能にし、探索オーバーヘッドを削減する。

実験結果

TraceNASは、Llama-2-7B、Llama-3.1-8B、および Qwen-2.5-14B に対して評価され、それぞれ 2.7B、4.6B、8.4B パラメータにプルーニングされた。

• 探索効率: TraceNASは、A100において8.5 GPU時間で最適なアーキテクチャを特定した。これは、DarwinLMのような学習を伴うベースラインと比較して10倍の削減となる。また、DarwinLMと比較して10倍少ない98Mの探索トークンを必要とした。
• 性能:
  • Llama-2-7B: TraceNASは、8つの推論ベンチマークにおいて最高の平均精度（62.81%）を達成し、DarwinLM（62.57%）およびShearedLLaMA（62.63%）を上回った。特筆すべきは、ShearedLLaMAはTraceNASと同じトークン予算に制限された場合、性能が大幅に低下（60.86%）したことである。
  • Llama-3.1-8B & Qwen-2.5-14B: TraceNASは、複雑な推論タスク（例：ARC-Challenge, BoolQ）において、一様プルーニングのベースラインおよび学習を伴う手法（DarwinLM, E3-Pruner）を上回った。同時に、リカバリに必要なトークン数も大幅に少なかった（Minitronの20Bに対し、94B）。
• プロキシの検証: プロキシは、様々なハイパーパラメータ（サンプルサイズ、コンテキスト長、LoRAランク）に対して高い安定性を示し、勾配トレースの整合性が、制約された設定下でもモデルの継承を効果的に捉えていることを確認した。

重要性と主張

論文は、TraceNASが構造化プルーニングにおける計算コストの決定的なボトルネックに対処すると主張している。学習時（ゼロからの学習）から「継承」（事前学習済みの損失曲面の保持）へと焦点を移すことで、TraceNASは、最小限のリソース消費で高性能な非一様プルーニングモデルの発見を可能にする。

著者らは、同手法が以下の点において優れていることを強調している：

• 学習時のトレーニングに伴う法外なコストを排除し、大規模なモデル圧縮をアクセシブルにする。
• 局所的なヒューリスティックが見落とすグローバルな構造的依存関係を捕捉し、複雑な推論能力の保持を確実にする。
• アーキテクチャ探索に必要なGPU時間を劇的に減らすことで、モデル圧縮のカーボンフットプリントを削減する。

本研究は、TraceNASを、高密度な事前学習済みモデルの機能的完全性を維持しつつ、リソース制約のある環境で効率的なLLMをデプロイするためのスケーラブルな基盤として位置づけている。