TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

技術的サマリー：TAPIOCA – タスク認識型プルーニングが OOD モデル能力を向上させる理由

問題定義
最近の研究では、推論時に大規模言語モデル（LLM）から層を除去すること、特に TALE（Task-Aware Layer Elimination）のようなタスク認識型プルーニング手法を用いることが、再学習なしに特定タスクのパフォーマンスを向上させることが示されている。しかし、これらの向上を駆動するメカニズムは依然として不明である。核心的な問いは、これらの改善が一般的な圧縮の恩恵、正則化効果、あるいはモデル内部表現と入力分布との間の特定の相互作用に起因するものかどうかである。さらに、これらの改善が普遍的なものか、それとも特定のデータシフトのタイプに固有のものかも不明である。

手法
著者らは、制御された合成実験と、微調整済み LLM における大規模評価の組み合わせを用いて、この現象を調査する。

1. 制御された回帰設定: 本研究は、コンテキスト内線形回帰に対してゼロから訓練された 12 層トランスフォーマーから始まる。訓練分布は明示的に定義され（係数は $U(-1, 1)$ からサンプリング）、モデルは分布内（ID）データ（同一分布）と分布外（OOD）データ（係数がより広い区間 $U(-\sigma, \sigma)$、ただし $\sigma > 1$ からサンプリング）の両方で評価される。著者らは、特定の分布における検証精度を最大化するように層を除去する貪欲なプルーニングアルゴリズムである TALE を適用する。
2. 大規模 LLM 評価: 発見事項は、数学的推論とコード生成に特化した微調整済み LLM（Llama 3.1 8B および GPT-OSS 120B）に拡張される。ここで、「適応された」分布は微調整タスク（例：MATH500）であり、OOD タスクには MMLU や BoolQ などの一般的なベンチマークが含まれる。
3. 幾何学的および因果的解析: プルーニングが機能する「理由」を説明するために、著者らは隠れ状態の幾何学を解析する。層ごとの表現ノルムとペアワイズなトークン距離を測定する。また、層変換を近似する局所線形代理モデルを採用し、逆代理マップの挿入や残差更新のスケーリングといった因果的介入を行い、幾何学的歪みの効果を分離する。

主要な発見

• OOD 固有の向上: 主要な実証的発見は、タスク認識型プルーニングが分布シフト下でのみパフォーマンス向上をもたらすということである。モデルが適応された分布（ID）で評価される場合、TALE はパフォーマンスを向上させないか、最適な精度に必要な層を除去することで能動的にパフォーマンスを低下させる。対照的に、OOD データでは、TALE は一貫して精度を向上させる層を特定し、除去する。
• 表現幾何学のミスマッチ: 著者らは、OOD 入力が ID 入力と比較して歪んだ表現幾何学を誘発することを示す。具体的には、OOD 入力は、層ごとの隠れ状態ノルムとペアワイズ距離が ID データで観察されるプロファイルから系統的に逸脱させる。回帰設定では、OOD 入力は層全体で距離を膨張させる。LLM では、OOD の軌道が中層から後層において ID 軌道から分岐する。
• 増幅メカニズム: 本研究は、特定の層が「分布依存性の増幅器」として機能することを特定する。ID 入力に対しては、これらの層は有用な微細化を行う。しかし、OOD 入力に対しては、同じ層が高利得・低ランクの増幅を示し、表現をタスク適応された幾何学からさらに遠ざける。TALE はこれらの特定の層を除去することでパフォーマンスを向上させ、幾何学的歪みの増幅を防ぐ。
• 収縮だけでなく整列: プルーニングは単に活性化ノルムを抑制するものではない。分布シフトの方向を反転させた実験（広い分布で訓練し、狭い分布でテスト）において、プルーニングは ID プロファイルと整合するように表現距離を「拡大」する。核心的なメカニズムは、収縮が必要か拡大が必要かにかかわらず、適応中に学習された幾何学に OOD 表現を再整列させることである。
• 因果的証拠: 因果的介入は、幾何学的歪みがパフォーマンス低下の駆動因子であることを確認する。プルーニングされた層の後に逆代理マップを挿入（層の幾何学的効果を元に戻す）することで、OOD パフォーマンスはプルーニング済みモデルのレベルにほぼ回復する。同様に、選択された層の残差更新をスケーリングし、層の寄与を減衰させると、精度が継続的に向上し、トポロジカルなアーティファクトや単純な冗長性が唯一の原因ではないことが排除される。

意義と主張
本論文は、推論時プルーニングの成功に対する新たな幾何学的説明を提供すると主張する。層の重要性は固定的ではなく、タスクと入力分布の両方に条件付けられると論じる。著者らは、事前学習モデルがそれぞれ局所的な表現幾何学に関連付けられた複数のタスク特化機能エンコードしていると仮定する。ID 入力はこの幾何学に適合するが、OOD 入力はこれを侵害する。

この研究の意義は以下の点にある：

1. プルーニングの限界の明確化: タスク認識型プルーニングは万能のパフォーマンス向上策ではなく、分布シフトに対する修正メカニズムであることを確立する。
2. 幾何学的解釈: 層の冗長性の理解を、静的なパラメータ数から、動的で分布依存性の幾何学的歪みにシフトさせる。
3. メカニズム的洞察: パフォーマンス向上が、モデル容量の増加やノイズの削減ではなく、内部表現の軌道をモデルの適応された幾何学に一致させるように修正することから生じるという因果的証拠を提供する。

著者らは、プルーニングが、入力分布とモデルの学習済みタスク幾何学との間のミスマッチを増幅する層を減衰させることで OOD パフォーマンスを向上させ、結果としてモデルの内部状態を最適化された条件と再整列させると結論づける。