💬 NLP

WIDE: Boosting Adaptive LLM Inference via Token-level Dynamic Width Pruning

本論文は、アテンションヘッドおよびFFNチャネルのきめ細かな選択を可能にすることでLLMのトークンレベルの動的な幅のプルーニングを実現する、初のエンドツーエンドで微分可能なフレームワークであるWIDEを紹介しており、専用のカーネル・コデザインと組み合わせることで、高い精度を維持しながら大幅なエンドツーエンドの推論加速を実現している。

原著者： Haozhe Hu, Hao Wu, Peiran Yin, Chao Han, Yunpu Ma, Xiaoyu Shen

公開日 2026-07-31

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

CC BY 4.0

原著者： Haozhe Hu, Hao Wu, Peiran Yin, Chao Han, Yunpu Ma, Xiaoyu Shen

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

あなたは、巨大で非常に賢いロボットの脳を、小さなバッテリー駆動のデバイスで動かそうとしていると想像してください。このロボットの脳は、大規模言語モデル（LLM）と呼ばれる人工知能の一種です。物語を書いたり、数学の問題を解いたり、人間のようにチャットしたりすることができます。しかし、問題は、これらの脳は非常に巨大であるということです。これらはあまりに大きく、電力も大量に消費するため、通常はスーパーコンピュータを使って動かす必要があります。そのため、スマートフォンやノートパソコンで使うには、動作が遅く、コストも高くついてしまいます。

これを解決するために、科学者たちはこれらの脳を「プルーニング（枝刈り）」しようとしてきました。プルーニングとは、巨大で茂りすぎた生け垣を整えるようなものです。植物がほとんど機能していない枝を切り落として、小さく、速くします。長い間、この最も良い方法は、植物がその日に何をしようとしているかにかかわらず、生け垣のセクション全体を永久に切り落とすことでした。しかし、これは少し不器用な方法です。時として、特定のタスクに必要だった枝まで切り落としてしまい、ロボットを物忘れさせてしまうことがあるからです。最近では、ロボットが思考している最中に、どの部分を使うかを判断できる、よりスマートな方法が発明されました。しかし、これらのスマートな手法でもまだ少し粗削りでした。彼らは、あるセクション全体を使うか、あるいは完全にスキップするかという決定を下しており、それはまるで、いくつかの特定の明かりだけをつけるのではなく、部屋全体を使うか、あるいは家全体を空っぽにするかを決めるようなものでした。

この論文では、WIDE（Width-based Inference with Dynamic Execution：幅に基づいた動的実行による推論）と呼ばれる新しいシステムを紹介しています。これは、ロボットに、脳内のあらゆる電球に対して超精密な調光器を与えるようなものです。部屋全体をオンにするかオフにするのではなく、WIDEは、単語（トークン）を処理するごとに、どの小さなニューロンのグループ（脳の細胞）を点灯させ、どれを消したままにするかを、動的に決定することができます。研究者たちは、この瞬時の判断を下すための特別な「交通管制官」（ルーター）を構築しました。また、コンピュータのハードウェアが混乱したり速度が低下したりすることなく、これらの決定を処理できる新しい方法も設計しました。その結果、ロボットは以前と同じくらい賢いまま、脳の容量の半分を削り取ったとしても、より速く、より少ないエネルギーで動作するようになりました。

問題点：「全か無か」の罠

あなたが、大規模なキッチンを運営しているシェフだと想像してください。あなたには、料理を作るために協力し合う何百人ものシェフ（ニューロン）がいます。昔は、時間を節約したい場合、キッチンのスタッフの半分を永久に解雇していました。これは、常に同じ簡単な料理を作っている場合にはうまく機能しますが、もし突然、複雑なケーキを焼く必要が出てきたら、必要だった唯一のパン職人を解雇してしまったかもしれません。これが**静的プルーニング（static pruning）**で起こることです。モデルは、どのような質問に答えているかにかかわらず、一度切り取られ、固定されます。

その後、**動的プルーニング（dynamic pruning）**が登場しました。これは、マネージャーを雇って、「おい、今回の注文にはグリル・ステーションだけでいいぞ。ベーカリーは休憩していい」と伝えるようなものでした。これは改善されましたが、マネージャーはまだ少し不器用でした。彼らは、「ベーカリー全体をスキップしろ」と言ってしまうのです。たとえ、そのベーカリーに10個のオーブンのうち2個だけが必要だったとしてもです。これは、キッチンの大きな塊に対する「全か無か」の決定でした。

問題は、現代のAIモデルは非常に複雑であるため、塊全体をスキップすると、有用な情報が失われ、回答の質を損なってしまうことが多いということです。さらに、もしマネージャーが毎秒ごとにスケジュールを変更し続けると、キッチンのスタッフは混乱し、切り替えのオーバーヘッドのせいで、実際の調理速度はあまり速くなりません。

解決策：WIDEの「調光器」

Ningbo Institute of Digital TwinとLMU Munichで働くこの論文の著者たちは、WIDEを生み出しました。部門全体を解雇したり、部屋全体をスキップしたりする代わりに、WIDEは、モデルが単語（トークン）を処理するごとに、どの小さなニューロンのグループを活性化させるかを決定することを可能にします。

モデルの脳を、巨大な照明スイッチのグリッドだと考えてください。

従来の動的方法： 「家の左翼全体を消灯せよ」
WIDE： 「この特定の単語に対しては、キッチン、しかもコンロと冷蔵庫の上にある照明だけを点けて。それ以外のキッチンは暗いままにしておけ」

WIDEは、モデルに付随する軽量な「ルーター」（小さな意思決定者）を使用することでこれを行います。このルーターは現在の単語を見て、「このアテンション・ヘッド（他の単語を見る部分）のグループが必要か？」そして「このFFNチャンネル（意味を処理する部分）のグループが必要か？」と問いかけます。そして、そのグループを使うか、スキップするかというバイナリ（二値）の選択を行います。

決定的なことに、WIDEは単に決定を下すだけでなく、ハードウェアの悩みを解決します。通常、コンピュータに計算のランダムな部分をスキップするように指示すると、コンピュータはデータの所在を探すことで停滞してしまいます。WIDEは、**マスク再配置（mask reordering）**と呼ばれる巧妙なトリックを使用します。郵便物の乱雑な山があり、いくつかは配達用で、いくつかはゴミだと想像してください。ゴミを捨ててから郵便物を仕分けようとする代わりに、WIDEはまず、すべての「保持すべき」手紙を整然としたブロックとして上に、すべての「ゴミ」の手紙を下に、山をシャッフルします。これにより、コンピュータはすべての郵便物を一つずつチェックすることなく、非常に効率的に「保持すべき」ブロックを処理できるようになります。

実験結果：スピードと賢さ

研究者たちは、Llama 3.1（80億パラメータ）やLlama 3.2（30億パラメータ）といった普及しているAIモデルを用いてWIDEをテストしました。彼らは、静的および動的プルーニングにおける既存の最高の方法と比較しました。

数値が示唆する内容は以下の通りです：

よりスマートなプルーニング： モデルの容量を50%削減（非常に積極的なトリミング）した際、WIDEは他のどの手法よりもモデルの知能をよく維持しました。例えば、Llama 3.1モデルにおいて、WIDEはプルーニング後に元の精度の**86.42%を維持しましたが、次に優れた動的方法であるSkipGPTは59.42%**しか維持できませんでした。追加のファインチューニングなしでも、WIDEはすでに最高の静的方法を上回っていました。
現実世界のスピード： 論文では、モデルが実際にどれほど速くなったかを測定しています。「プリフィル（prefill）」段階（長いプロンプトを読み込む段階）では、WIDEは元のモデルよりも1.68倍速くなりました。「デコード（decode）」段階（答えを単語ごとに生成する段階）では、1.55倍速くなりました。
カーネルレベルのマジック： 数学的な計算のみに注目すると（他のオーバーヘッドを無視すると）、スピードアップはさらに劇的で、プリフィルで最大1.98倍、デコードで最大4.95倍に達しました。これは、従来のメソッドが「スキップ」を扱う際の非効率な方法がボトルネックであったことを示唆しており、WIDEの新しいハードウェア設計がそれを解決したことを意味しています。

また、著者らはモデルが非常に戦略的に学習したことも発見しました。モデルは単にランダムに部分をスキップしたのではなく、「the」や「a」のような「退屈な」単語をスキップし、「running」や「track」のような「重要な」単語を保持することを学習しました。あるテストでは、モデルはアテンションの作業の66%をスキップしましたが、プロセッシング（処理）の作業は28%しかスキップしておらず、どこでエネルギーを節約すべきかを正確に把握していたことが示されました。

結論

WIDEは、効率的なAIの未来が、単にモデルを小さくすることではなく、残された部分を「どのように」使うかについてモデルをより賢くすることにあることを示唆しています。「部屋全体をスキップする」から「特定の明かりを調光する」へと移行し、さらにこれらの調光器を効率的に扱うためにキッチンを再設計することで、著者らは強力なAIモデルを、記憶力を失うことなく大幅に高速化し、効率化するフレームワークを作り上げました。

結果は有望ですが、論文では、これはGPUハードウェア（コンピュータ内のチップ）向けの特定的なソリューションであり、モデルがまず意思決定を学び、次に失われた精度を取り戻すために素早い調整を受けるという、2段階のトレーニングプロセスに依存していることも記されています。これは、強力なAIを日常的なデバイスで利用可能にするための重要な一歩であり、賢くなるために巨大な脳は必要ではなく、ただ適切な時に脳のどの部分を使うかを知っていればよいということを証明しています。

技術要約: WIDE - トークンレベルの動的幅プルーニングによる適応型LLM推論の強化

1. 問題提起

大規模言語モデル（LLM）は、スケールアップに伴い重大な効率性の課題に直面しています。**静的な構造化プルーニング（static structured pruning）**は、ハードウェアフレンドリーでありスループットの向上をもたらしますが、入力に依存しない性質（input-agnostic）を持つため、アグレッシブなスパース性下では固定された除去パターンをすべての入力に適用してしまい、大幅な精度低下を招くことがよくあります。対照的に、**動的プルーニング（dynamic pruning）**手法は個々の入力に合わせて計算を適応させますが、現在は以下の点に制限されています：

粒度の粗さ（Coarse Granularity）: 既存の動的方法は、主に層レベル（層全体またはサブモジュールを実行するかスキップするかを決定する）で動作します。この「全か無か」のアプローチは、レイヤー内で部分的な計算のみを必要とするトークンを捉えることができず、不必要な容量損失につながります。
システム的な非効率性（Systemic Inefficiency）: 動的プルーニングをより細かい粒度（例：トークンごとの幅の割り当て）へと押し進めようとすると、不規則な実行パターンが生じます。明示的なgather-scatter操作に依存するナイーブな実装は、過度なデバイス・ホスト間の同期、過剰なメモリトラフィック、およびストレージオーバーヘッドを引き起こし、FLOPsの理論的な削減が現代のGPUにおける実質的なエンドツーエンドのスループット向上に結びつくことを妨げます。

2. メソドロジー: WIDEフレームワーク

著者らは、プリフィル（prefill）とデコード（decode）の両方のシナリオ向けに設計された、最初のエンドツーエンドで微分可能なトークンレベルの動的幅プルーニングフレームワークであるWIDEを提案しています。このメソドロジーは、以下の3つのコアコンポーネントで構成されています。

A. 細粒度トークンレベル・ルーティング

層レベルでルーティングを行う従来のアプローチとは異なり、WIDEは各トークンが、すべてのTransformerブロック内において、**アテンションヘッド（attention heads）およびFFNチャネル（FFN channels）**の細粒度なグループを独立して選択することを可能にします。

プルーニング単位（Pruning Units）:
- アテンション（Attention）: 既存のデコーディング最適化（例：XQA）との互換性を維持するため、Grouped Query Attention (GQA) 構造に沿ったアテンションヘッドのグループをプルーニングします。
- FFN: 中間次元 ( $D'$ ) における連続したチャネルグループをプルーニングし、標準的なGEMMタイリング戦略との互換性を確保します。
ルーター・アーキテクチャ（Router Architecture）: 軽量なボトルネックルーターが、入力活性化を各グループに対するバイナリ決定（実行 vs スキップ）へとマッピングします。学習中、Gumbel-Softmax緩和を用いることで微分可能な最適化が可能になります。
学習パイプライン（Training Pipeline）: 2段階のプロセスが採用されます：
1. ルーター学習（Router Training）: ベースモデルの全パラメータを凍結し、ルーターのみをキャリブレーションデータ上で学習させて、スパースな実行パターンを学習させます。
2. LoRAリカバリ（LoRA Recovery - オプション）: モデルの品質を回復するために、線形層に低ランク適応（LoRA）モジュールを追加し、その後にファインチューニングを行います。

B. 統合されたマスク再順序付け前処理（Unified Mask-Reordering Preprocessing）

不規則なトークンごとのルーティングによる非効率性を解決するため、WIDEは、明示的なgather-scatter操作なしに、任意のルーティングマスクを構造化された実行レイアウトに変換する前処理ステップを導入しています。

メカニズム: 各ルーティンググループに対して、バイナリマスクをトークン次元 ( $M$ ) に沿ってソートします。これにより、アクティブなトークンが連続したプレフィックスとして集約されます。
利点: これにより、不規則なスパース性がCTAレベルの規則性へと変換されます。タイル内のアクティブな行は連続的になるため、カーネルはデータを既知のアクティブなプレフィックスを持つ密なブロックとして扱うことができ、トークン固有のスパーステンソルを実体化する必要がなくなります。

C. プルーニング・カーネルの協調設計（Pruning-Kernel Co-Design）

WIDEは、再順序付けられたマスクを利用するために、カスタムCUDAカーネル（CuTe、Triton、TileLangを使用）内でマルチステージのスキッピング戦略を実装しています。

CTAレベルの早期終了 ( $\ell=0$ ): 特定のグループに対してトークンのブロック全体が非アクティブである場合、CTA（Compute Thread Array）全体の実行をスキップします。
イントラブロック・スキッピング ( $\ell=1, 2$ ):
- ロード ( $\ell=1$ ): 非アクティブなメモリパケットのロードをスキップします（例：cp.async や TMA gather を使用）。
- MMA実行 ( $\ell=2$ ): 非アクティブなTensor Coreの計算フラグメントをスキップします。
  この設計により、細粒度な動的スパース性が、効率的なGPU実行パターン（例：合体されたメモリアクセス）を維持しつつ、不要な計算を排除することを保証します。

3. 主な貢献

最初のエンドツーエンド・トークンレベル動的幅プルーニング: WIDEは、動的プルーニングを粗い層/サブレイヤーの決定から、ニューロン・ブロックレベルの粒度へと押し進め、トークンがアテンションヘッドおよびFFNチャネルのグループを動的に選択できるようにしました。
統合加速フレームワーク: 本論文は、同期されたマスク・インデックス変換と不規則なgather-scatterが主要なボトルネックであることを特定しました。WIDEは、スタンドアロンのマスク再順序付けとマルチステージのイントラブロック予測によってこれを解決し、カーネルがプルーブルな操作の理想的なスピードアップに近づくことを可能にしました。
プルーニング・カーネルの協調設計: WIDEは、トークンごとの動的幅プルーニングと、プリフィルおよびデコーディングの両方に対応するGPUカーネルの協調設計を組み合わせた最初のフレームワークです。これは、細粒度な動的プルーニングがハードウェアの実行パスと共に設計されなければならないことを示しています。

4. 実験結果

実験は、キャリブレーションおよびLoRAリカバリのためにRedPajama-1Tデータセットを用い、Llama3.1-8BおよびLlama3.2-3Bに対して行われました。

精度の保持:
- 50% スパース性（キャリブレーションのみ）において、WIDEはLlama3.1-8Bにおける最強の動的深度ベースライン（SkipGPT）に対し、平均ゼロショット精度を8.80ポイント改善しました（61.84 vs 53.04）。
- LoRAリカバリを使用した場合、WIDEはLlama3.1-8Bにおいて**90.59%**の精度保持を達成し、SkipGPT（86.09%）やDDP（76.36%）のような静的な手法を大幅に上回りました。
- WIDEは、様々なグループサイズ ( $G$ ) において堅牢なパフォーマンスを維持しており、 $G=32$ または $G=64$ が最適なトレードオフをもたらすことが多いことが示されました。
スピードアップと効率性:
- カーネルレベル: WIDEは、ナイーブなgather-scatterベースラインと比較して、最大200倍のスピードアップと1,000倍のピークメモリ削減を達成しました。
- エンドツーエンド: 50%のスパース性条件下で、WIDEはプリフィルで1.68倍、デコーディングで1.55倍のスピードアップを実現しました。
- 理論的限界: 本フレームワークは、プリフィルで1.98倍、デコーディングで4.95倍という、理論的なカーネルレベルのスピードアップに近い値を達成しています。
- オーバーヘッド: スパース性が0%の場合、最良の構成は密なスループットの**98.60%**を保持しており、ルーティングおよびカーネル初期化によるオーバーヘッドが極めて小さいことを示しています。

5. 重要性と主張

本論文は、細粒度な動的プルーニングは、単なるモデル圧縮の問題としてではなく、ハードウェアの実行パスと協調設計されなければならないことを示すことで、WIDEがトークンごとの動的構造化プルーニングの新たな境地を切り開いたと主張しています。

品質と効率のトレードオフ: WIDEはパレート図を押し上げることに成功しており、既存の動的深度プルーニング手法よりも大幅に優れた品質保持を提供しながら、実用的なリアルワールドの加速を実現しています。
意味論的認識（Semantic Awareness）: ルーティングパターンの分析により、WIDEは意味内容（例：「the」のような冠詞をスキップし、「running」のような内容語を保持する）に基づいて計算予算を割り当てることを学習していることが明らかになり、トークンレベルの適応性の有効性が検証されました。
実用的な生存能力: マスク再順序付けとカーネル協調設計を通じて不規則な実行のシステムレベルの課題を解決することで、WIDEは、細粒度な動的スパース性が現代のGPUアーキテクチャ（具体的にはNVIDIA RTX 5090およびA100）において、真のエンドツーエンドのスループット向上をもたらし得ることを証明し、理論的なスパース性と実用的な推論加速の間の溝を埋めています。