🤖 machine learning

Adaptive Computation Depth via Learned Token Routing in Transformers

本論文は、文脈の難易度に基づいて個々のトークンに対して冗長な層計算を動的にスキップすることを可能にする軽量かつエンドツーエンドで微分可能なゲーティング機構であるトークン選択アテンション（TSA）を導入し、明示的な深さ正則化を必要とすることなく、最小限の品質低下で大幅な効率向上を実現するものである。

原著者： Ahmed Abdelmuniem Abdalla Mohammed

公開日 2026-05-08

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

CC BY 4.0

原著者： Ahmed Abdelmuniem Abdalla Mohammed

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

工場で、コンベアベルトを流れてくるすべての製品が、全く同じ処理を受けると想像してください。単純で完璧なウィジェットであれ、複雑で壊れたウィジェットであれ、すべての製品はすべての工程で全く同じ時間を費やします。これが現在の標準的な AI モデル（トランスフォーマー）の仕組みです：どの単語が理解しやすいか難しいかに関係なく、文章内のすべての単語が、全く同じ数の「思考層」を通過して処理されます。

この論文は、**トークン選択的アテンション（TSA）**と呼ばれる新しいシステムを紹介しています。TSA は、工場のすべての工程に、賢く自動的な門番を設置したと考えることができます。

問題点：「すべてに同じサイズ」の工場

標準的な AI モデルでは、「To be, or not to be」と書かせると、モデルは「To」という単語と「question」という単語に、同じだけの知能を費やします。

簡単な単語（「the」や「is」など）は、単純なウィジェットのようなものです。これらはあまり処理を必要としません。
難しい単語（複雑な概念や珍しい名前など）は、壊れたウィジェットのようなものです。これらを修正するには、追加の時間と注意が必要です。

現在、工場は、難しいものと同じ強度で簡単なウィジェットを処理することで、エネルギーを無駄にしています。

解決策：賢い門番（TSA）

著者たちは、AI の各層の間に、小さく軽量な「門番」（小さなニューラルネットワーク）を追加しました。この門番は、各単語（トークン）を個別に観察し、「この単語は次の処理工程に進む必要があるか、それともスキップできるか？」と問います。

判断: 門番は単に「はい」か「いいえ」と言うだけではありません。それは確率スコア（調光スイッチのようなもの）を与えます。単語が簡単であれば、門番は「次のステップをスキップできる」「作業を半分だけ行えばよい」と言うかもしれません。単語が難しい場合は、「最後まで通過する」と言います。
魔法: 最も素晴らしい点は、門番が自ら学習することです。どの単語が難しいかを人間が教える必要はありません。誤りを最小化しようとする試行を通じて純粋に学習します。特定の単語のステップをスキップすることでエラーが発生すれば、門番は次回はその単語をアクティブに保つように学習します。スキップしても問題なければ、エネルギーを節約するように学習します。

実務での仕組み

この論文は、主に以下の 2 つの点でこれをテストしました：

単純な論理パズル: AI が数字のリストをコピーする必要がある場合、門番は「これは簡単だ」と判断し、作業の約 65% をスキップしました。数字をソートする必要がある場合（これはより難しい）、スキップしたのは約 27% でした。これは、より難しいタスクにはより多くのステップが必要であることを自然に理解しました。
物語の執筆: シェイクスピアのように書くよう求められたり、ウィキペディアの記事を要約させたりした場合、AI は計算能力の**14% から 23%**を節約しました。
- 句読点、スペース、一般的な単語は簡単であり、素早く処理できることを学習しました。
- 固有の内容語は、完全な「工場」処理を必要とすることを学習しました。

結果：より速く、より賢く

品質の低下なし: AI は標準モデルと同じように書きましたが、使用するエネルギー（計算能力）は大幅に少なくなりました。
「早期終了」より優れている: 他の手法は、AI が「自信」を持っている場合に、文全体を早期に停止させようとします。TSA はより賢明です。個々の単語が不要なステップを通過するのを止めつつ、難しい単語は進めさせます。論文によると、同じ量のエネルギーを節約するように設定された場合、TSA はこれらの古い手法よりも良い結果を生み出しました。
実際の速度向上: 研究者が実際にコードをコンピュータで実行したとき、スキップされた単語の計算をコンピュータが実際に行わなかったため、実際の速度向上（約 2.3% 高速）が見られました。

結論

この論文は、AI モデルを「最も賢い方法で怠けさせる」方法を示しています。すべての単語に同じ量の作業を強いるのではなく、TSA は AI に、どの単語が簡単で、どの単語が追加の助けを必要とするかを、その場で判断させます。簡単な製品はラインを素通りし、難しい製品は完全な VIP 待遇を受け、品質を犠牲にすることなく時間とエネルギーを節約する工場を持っているようなものです。

技術的概要：トランスフォーマーにおける学習済みトークンルーティングによる適応的計算深度

問題定義

標準的なトランスフォーマーアーキテクチャは、トークンの文脈的難易度や情報密度に関わらず、シーケンス内のすべてのトークンに固定された層数を適用します。この設計は、トークンごとの適応性よりもアーキテクチャの単純さを優先しており、結果として計算効率の重大な非効率性を生んでいます。実際には、予測可能な文脈における一般的なトークンは、新規の構成における稀なトークンに比べてはるかに少ない処理を必要としますが、両者はすべてのフォワードパスにおいて同一の計算リソースを受け取ります。この非効率性は、推論規模において特にコストがかかります。なぜなら、支配的なコストは、すべてのトークンに対する全層を通じた完全なフォワードパスだからです。既存の解決策、例えば RNN 向けの適応的計算時間（ACT）、ユニバーサルトランスフォーマー、深さの混合（MoD）、および早期終了メカニズムは、この問題に対処しようと試みていますが、しばしばハード選択、累積予算、または事後の信頼度閾値に依存しており、これらは必ずしも微分可能または最適であるとは限りません。

手法：トークン選択的アテンション（TSA）

本論文は、連続するトランスフォーマーブロック間の残差更新に対して、学習済みかつトークンごとのゲートを導入するメカニズムである**トークン選択的アテンション（Token-Selective Attention, TSA）**を提案します。

アーキテクチャとルーティング

TSA は、プレノームのデコーダー専用トランスフォーマーにおいて、各ブロック $f_l$ （ $l = 0, \dots, L-2$ ）の後に軽量なルーター $r_l$ を挿入します。ステムブロック $f_0$ は常に無条件に実行されます。

ルーター設計: 各ルーターは、シグモイド出力を持つ 2 層のマルチレイヤーパーセプトロン（MLP）であり、各トークンに対して連続的な停止確率 $p_l \in (0, 1)$ を生成します。隠れ次元は $d/4$ （16 で切り捨て）であり、ベースモデルに対して約 1.7% のパラメータオーバーヘッドを追加します。
ゲート付き更新: 後続のブロック $f_{l+1}$ $f_{l + 1}$ に対する残差更新は、 $(1 - p_l)$ $(1 - p_{l})$ によってスケーリングされます。
- $p_l = 0$ の場合、更新は標準的です。
- $p_l = 1$ の場合、状態は変更されず、実質的にそのブロックをスキップします。
- 補間は滑らかであり、このメカニズムはストライトスルー推定子、ガンベルサンプリング、または強化学習を必要とせずに完全に微分可能であることを保証します。

訓練と正則化

訓練中は、すべての層が完全に実行されます。ゲートは残差更新をスケーリングしますが、計算をスキップすることはありません。総損失は $L = L_{task} + L_{depth}$ と定義されます。

深度正則化: 早期停止を穏やかに促す項 $L_{depth} = \lambda \cdot \frac{1}{L-1} \sum (1 - p_l)$ が追加されます。
創発的挙動: 注目すべきは、論文が示すところによれば、ルーティングメカニズムは $\lambda = 0$ （明示的な深度圧力なし）であっても、タスク損失勾配のみによって駆動され、操作をスキップすることを学習するということです。

推論

推論時には、メカニズムは「スパース」モードに切り替わります。 $p_l > 0.5$ のトークンは、gather/scatter 操作を介してフィードフォワードネットワーク（FFN）から除外されますが、正確なキー・バリューの意味を保持するためにアテンションは密のままです。これにより、実際の FLOPs の節約とウォールクロック時間の高速化が実現されます。

主要な貢献

微分可能なトークンルーティング: ベーストランスフォーマーのアーキテクチャ変更なしに、トークンごと・層ごとのルーティングを可能にする、残差更新に対する単純な連続的なソフトゲート。
勾配駆動による創発: 難易度に比例するルーティングが、タスク損失勾配のみから自然に創発するという証拠。深度正則化（ $\lambda=0$ ）がなくても、ルーターはトークン - 層操作の約 20% をスキップすることを学習します。
効率性と品質のトレードオフ: 文字レベル言語モデリングにおける検証により、トークン - 層操作（TLOps）を 14–23% 節約しつつ、品質の低下は 0.5% 未満であることを示しました。
早期終了との優位性: 同等の効率性において、TSA は早期終了のベースラインよりも低い検証損失を達成し、エンドツーエンドで学習されたルーティング決定が、推論時の事後の信頼度閾値設定よりも優れていることを示唆しています。
実世界の高速化: 汎用ハードウェア（Apple M1 Pro）におけるスパース推論による、実質的なウォールクロック時間の高速化（最大 2.3%）の実証。

実験結果

合成タスク

コピーおよびソートタスク（長さ 10 のシーケンス）において：

コピータスク: ルーターはタスクが恒等写像であることを学習し、精度の低下なしに TLOps の 54.9% をスキップしました（ $\alpha = 0.341$ ）。
ソートタスク: ルーターはより多くの計算を割り当て、TLOps の 22.5% をスキップしました（ $\alpha = 0.730$ ）。これはより高い難易度を反映しています。

文字レベル言語モデリング

Tiny-Shakespeare: TSA は、検証損失の増加がわずか 0.006 ナット（+0.4%）であるのに対し、TLOps の 22.8% を節約しました（ $\alpha = 0.726$ ）。
enwik8: より多様なコーパス（Wikipedia）において、TSA は品質の劣化なし（わずかに良い検証損失）に TLOps の 13.9% を節約しました（ $\alpha = 0.833$ ）。
ロバスト性: このメカニズムは、2 桁の範囲（ $[0, 0.1]$ ）にわたる正則化ハイパーパラメータ $\lambda$ に対してロバストです。 $\lambda=0.05$ において、品質コストを無視できるレベルで TLOps の 50.4% を節約しました。

早期終了との比較

一致したアクティブ割合（ $\alpha \approx 0.726$ ）において、TSA は早期終了のベースラインよりも0.71% 低い検証損失を達成しました。これは、エンドツーエンドでルーティング決定を学習することが、推論時に信頼度閾値を適用するよりも優れた決定を生み出すことを示しています。

ウォールクロックスループット

Soft-TSA（訓練）: $\alpha$ に関わらず約 1% のオーバーヘッドを追加しました。
Sparse-TSA（推論）: Tiny-Shakespeare（ $\alpha=0.726$ ）で2.3% の高速化を達成し、enwik8（ $\alpha=0.833$ ）ではブレイクイーブン性能を示しました。高速化には、CPU-GPU 同期オーバーヘッドを克服するためにバッチサイズ $\ge 64$ が必要でした。

意義と主張

本論文は、TSA が固定深度トランスフォーマーの非効率性に対する最小限かつアーキテクチャ的に単純な解決策を提供すると主張しています。その主な意義は、難易度に関する明示的な監督や複雑な強化学習シグナルなしに、適応的計算深度をタスク目的から直接学習できることを実証した点にあります。

著者らは、ルーターが「学習済みノイズゲート」として機能し、残差貢献が冗長な場合に更新を減衰させると強調しています。これにより、モデルはトークンの難易度に基づいて計算を動的に割り当てることができ、合成タスクでは最大 55%、言語モデリングでは 14–23% の大幅な効率向上を、最小限の品質トレードオフで達成します。この研究は、離散選択や事後の閾値設定よりも、エンドツーエンドの微分可能ルーティングが適応的深度のためのより効果的な戦略であることを示唆し、汎用ハードウェアにおけるより効率的な推論への道を開いています。

論文は規模に関しては控えめな立場をとり、すべての実験が約 500 万〜600 万パラメータのモデルで行われたことに言及し、より大規模なモデル（1000 万〜1 億パラメータ以上）へのスケーリング挙動は未解決の課題であると述べています。さらに、現在の高速化はバッチサイズに依存しており、アテンションのスパース化はまだ実装されていない点にも触れています。