💻 computer science

SHIFT-LLM: Distribution Shift Correction in Depth-Pruned LLMs

SHIFT-LLMは、Transformerブロックの削除によって引き起こされる分布シフトを補正するために、閉形式の回帰を通じて較正された軽量な線形残差アダプタを挿入することで、深さ方向のプルーニングが行われた大規模言語モデルの精度を回復させる、学習不要のポストプルーニング・フレームワークである。

原著者： Ali Bahri, Hang Li, Hongliang Li, Zhitang Chen

公開日 2026-08-27

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Ali Bahri, Hang Li, Hongliang Li, Zhitang Chen

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

大規模言語モデルは、物語を書いたり、問題を解決したり、複雑な質問に答えたりすることができる、今日の最も高度な人工知能ツールの背後にあるエンジンです。これらのシステムは、情報の層が下の層から受け取った情報を洗練させながら積み重なっていく、多層階のビルのように、デジタル処理ユニットが何層にも重なって構築されています。ビルに階層が多いほど、最終的な出力はより詳細で洗練されたものになりますが、これは構造を非常に重く、運用コストを高くすることにもなります。多くの実用的なアプリケーション、特に電力やメモリが限られているデバイスにおいては、これらの巨大なモデルはあまりに大きすぎて使用することができません。これらを小さくするために、研究者たちは建物からフロア全体を取り除く手法を開発しました。これは「デプス・プルーニング（深さの枝刈り）」として知られるプロセスです。この手法は負荷を軽減しモデルを高速化することには成功していますが、多くの場合、残されたフロアが躓いてしまいます。なぜなら、取り除かれたフロンドは上の層へ特定の種類の情報を渡すように設計されていたため、それらを取り去ると、上の層は学習時に期待していたものとは一致しない信号を受け取るようになり、混乱してしまうからです。この不一致は「分布シフト」と呼ばれ、モデルの精度を低下させ、ミスを引き起こします。その結果、開発者はエラーを修正するために、膨大な時間と計算能力を投じてモデルを再学習させることを余儀なくされます。

Huawei Noah's Ark Labの研究チームは、再学習を必要とせずにこの問題を解決する、「SHIFT-LLM」と呼ばれる新しいアプローチを導入しました。取り除かれたフロアを再構築したり、残されたフロアを再学習させたりする代わりに、彼らの手法は、レイヤーが削除されたすべての箇所に、小さく軽量な補正モジュールを挿入します。もし、あるビルからフロアを取り除いた場合、そこに、隙間を完璧に埋めるカスタムメイドのシンプルなスロープを設置し、上の階の人々が、もし欠けているフロアがまだ存在していた場合と同じ指示を正確に受け取れるようにすることを想像してみてください。デジタル世界において、このスロープは、情報の直接的な経路を維持しながら、微小で計算された調整を加える「線形アダプター」です。この調整は、取り除かれたレイヤーが果たしたはずの具体的な貢献を模倣するように設計されており、実質的に、欠けているフロアがまだ存在しているかのように他のモデルを欺くのです。

この手法の素晴らしさは、その調整をどのように算出するかという点にあります。欠けているレイヤーが何をすべきだったかを理解するために数千時間のトレーニングを行うのではなく、研究者たちは少量のサンプルデータを使用して、失われた情報が正確にどのようなものであるかを測定します。そして、失われた情報を復元するために必要な正確な補正を、単純な数学的計算を用いて決定します。このプロセスは完全に自動化されており、通常、壊れたモデルを修正するために必要とされる複雑で反復的な学習プロセスである「勾配ベースの最適化」を必要としません。この閉形式（クローズドフォーム）の計算を用いることで、チームは数日間の計算時間ではなく、わずか数百の例を用いて数分でモデルの内部状態を修正することができます。その結果、プルーニングされたモデルは、より少ないレイヤーを持つことによるスピードと効率性を維持しながら、元の巨大なモデルとほぼ同等の精度を発揮します。

研究者たちは実験において、15億パラメータの小型モデルから140億パラメータの大型モデルまで、5つの異なる大規模言語モデルのファミリーに対してこの技術をテストしました。彼らはどのレイヤーを取り除くかを決定するために6つの異なる手法を適用し、一般的な知識と推論をテストするために設計された7つのベンチマークで結果を評価しました。結果は一貫していました。補正モジュールは、ほとんどすべての構成において、プルーニング中に失われた精度を正常に回復させたのです。ある80億パラメータのモデルに関する顕著なケースでは、この手法は標準的なベンチマークにおいて、驚異的な15.7ポイントの精度を回復しました。これは、通常であれば広範で高価な再学習を必要とする利得です。モデルがプルーニング後にすでにファインチューニングされている場合でも、この補正モジュールを追加することでさらなるブーストが得られたことは、これら二つのアプローチが性能をさらに高めるためにうまく機能することを示唆しています。

研究者たちはまた、これらの補正モジュールをさらに圧縮してスペースを節約したり、複数のレイヤーが連続して削除された場合にそれらを統合したりできることも実証しました。これにより、修正によるオーバーヘッドによってプルーニングによる効率向上のメリットが失われないようにしています。この研究は、大規模言語モデルを小型化する際の主な障害は、単に部品を取り除くことではなく、それらの間の情報の流れを乱すことであることを示唆しています。単純で的を絞った調整によってその流れを修復することに焦点を当てることで、SHIFT-LLMは、強力なAIモデルを日常的な用途に実用化するための、一般的かつ効率的な方法を提供します。この研究は、適切な補正さえあれば、デジタルな脳の重くて冗長な部分を取り除いても、明確に思考する能力を失うことなく、より速く、より安価で、より身近な人工知能への扉を開けることを裏付けています。

技術要約: SHIFT-LLM

問題提起

大規模言語モデル（LLM）における深さプルーニング（Depth pruning）は、Transformerブロック全体を削除することで推論効率を向上させることを目的としています。しかし、このプロセスは**分布シフト（distribution shift）**という重大な問題を引き起こします。中間層のブロックが削除されると、そのブロックが生成していたはずの残差更新（residual update）が失われます。その結果、ダウンストリームの層は、事前学習時に遭遇したものとは系統的に異なる隠れ状態を受け取ることになります。このシフトはネットワークを通じて伝播し、大幅な精度の低下を招きます。プルーニング後のファインチューニングによって性能を回復させることは可能ですが、それには膨大な計算リソース、最適化時間、および勾配ベースの学習が必要であり、リソース制約のあるシナリオにおける適用性を制限しています。

手法: SHIFT-LLM

著者らは、勾配計算を伴わずに分布シフトを軽減するために設計された、トレーニングフリーのポストプルーニング補正フレームワークであるSHIFT-LLMを提案しています。その中核となるメカニズムは、各プルーニング箇所に挿入される**線形残差アダプター（Linear Residual Adapter: LRA）**です。

1. 線形残差アダプター（LRA）の定式化

LRAは、元の出力隠れ状態をゼロから再構築するのではなく、元の**アイデンティティ・パスウェイ（identity pathway）を保持した上で、軽量なアフィン残差補正（affine residual correction）**を加えます。

構造: プルーニングされた層 $\ell$ に対して、元の出力 $h^{(\ell)} = h^{(\ell-1)} + f_\ell(h^{(\ell-1)})$ は、 $\hat{h}^{(\ell)} = h^{(\ell-1)} + g_\ell(h^{(\ell-1)})$ によって近似されます。
補正関数: アダプター $g_\ell(\cdot)$ は、アフィン変換としてモデル化されます： $g_\ell(h) = hA^{(\ell)} + \mathbf{1}(b^{(\ell)})^\top$ 。ここで、 $A^{(\ell)}$ は重み行列、 $b^{(\ell)}$ はバイアスベクトルです。
目的: LRAは、フル隠れ状態ではなく、欠落した残差更新（ $\delta^{(\ell)} = f_\ell(h^{(\ell-1)})$ ）を具体的に近似することを目的としています。これにより、モデルが事前学習した残差構造を活用することで、近似問題を簡略化しています。

2. 閉形式による推定（トレーニングフリー）

LRAのパラメータ（ $A^{(\ell)}$ および $b^{(\ell)}$ ）は、バックプロパゲーションや反復的な最適化なしに推定されます。

データ収集: 元の未プルーニングモデルを用いて、少量のホールドアウト・キャリブレーション・セットに対してフォワードパスを実行します。入力隠れ状態（ $H^{(\ell)}$ ）と、対応するグランドトゥルースの残差更新（ $\Delta^{(\ell)}$ ）を収集します。
リッジ回帰: パラメータは、予測された残差更新と実際の残差更新の間の誤差を最小化するように、閉形式の最小二乗回帰（リッジ回帰）を通じて解かれます：
$X^{(\ell)*} = \left((H^{(\ell)}_{aug})^\top H^{(\ell)}_{aug} + \lambda I\right)^{-1} (H^{(\ell)}_{aug})^\top \Delta^{(\ell)}$
ここで、 $H^{(\ell)}_{aug}$ はバイアス列を含み、 $\lambda$ は正則化係数です。
効率性: このプロセスは、プルーニングされた層ごとに、単一のフォワードパスと線形解法を必要とするため、勾配計算やファインチューニングを必要としません。

3. 実用的な拡張

低ランク分解: フィットされた重み行列 $A^{(\ell)}$ は、切断特異値分解（SVD）によって圧縮でき、パラメータ数を $d^2$ から $2dr + d$ へと削減可能です。
正確なマージ: LRAはアフィン写像を定義するため、隣接するプルーニング箇所にある連続したLRAは、近似誤差なしに単一の変換へと数学的にマージすることができ、さらなる圧縮を実現します。

主な貢献

SHIFT-LLM フレームワーク: 深さプルーニングによって引き起こされる隠れ状態の分布シフトを明示的にターゲットとした、トレーニングフリーのポストプルーニング補正手法。
線形残差アダプター（LRA）: 元の残差構造を保持しながら、閉形式の回帰によって校正された軽量なアフィン補正を通じて、欠落した残差更新を近似する新しいコンポーネント。
モジュール性と汎用性: 本フレームワークは特定の層選択戦略（例：Block Influence、Taylor sensitivity、Magnitude）に依存せず、多様なモデルファミリーに適用可能です。
相補性: 本手法は既存のパラメータ効率の良いファインチューニング（例：LoRA）と併用できるように設計されており、ファインチューニング単独による回復を超えた追加の性能向上を提供します。

実験結果

著者らは、5つのモデルファミリー（Qwen2-1.5B, Qwen1.5-7B, Llama-3.1-8B-Instruct, Vicuna-7B, Qwen2.5-14B）、6つの層選択基準、および7つのゼロショット・ベンチマークを用いてSHIFT-LLMを評価しました。

トレーニングフリーの回復: ファインチューニングを行うことなく、SHIFT-LLMは深さプルーニングによって失われた精度を一貫して回復させます。
- Llama-3.1-8B-Instructにおいて、逆順プルーニング（25%の深さ）を用いた場合、本手法は平均ゼロショット精度で +15.74 ポイントの向上を達成しました。
- Qwen1.5-7Bでは、Magnitude- $\ell_1$ プルーニングの下で +6.88 ポイントの利得に達しました。
- WikiText-2におけるパープレキシティは、ほとんどの構成において大幅に減少しました（例：Llama-3.1-8Bの逆順プルーニング下では、1000ポイント以上の減少）。
スケーラビリティ: 本手法は Qwen2.5-14B にスケールアップした場合でも効果的であり、ゼロショット精度で最大 +5.36 ポイントの利得を示しました。
ファインチューニングとの互換性: LoRAまたは部分層ファインチューニングと組み合わせた場合、SHIFT-LLMは相補的な回復を提供し、評価されたほとんどの設定でさらなる性能向上を実現しました（例：Llama-3.1-8B-InstructにおいてLoRAと組み合わせた際に +1.14 ポイント向上）。
SOTAとの比較: Navigation LLM パイプラインと統合することで、Llama-3.1-8B-Instructの平均精度を59.08から59.54へと向上させ、強力なベースラインを劣化させることなく価値を付加できることを実証しました。

重要性と主張

本論文は、SHIFT-LLMが、深さプルーニングに固有の分布シフト問題に対する、シンプルで汎用的かつ効果的な解決策を提供すると主張しています。元のモデルの残差構造を保持しながら、欠落した残差更新を明示的に近似することで、再学習の計算コストをかけることなく分布シフトを緩和します。

著者らは、このアプローチが、ポストプルーニングの学習が望ましくない、あるいは実行不可能な設定において特に重要であることを強調しています。数百のキャリブレーション・サンプルのみを使用し、勾配計算を一切行わずに、大幅な精度回復（最大 +15.7 ポイント）を実現できる能力を持つ本フレームワークは、広範なファインチューニングに代わる実用的な選択肢となります。さらに、連続するアダプターのマージや低ランク分解による重みの圧縮が可能であることは、深さプルーニングによる効率性の利点を維持、あるいはさらに強化することを保証します。