今回の論文
今回取り上げるのは、David Raposo、Sam Ritter、Blake Richards、Timothy Lillicrap、Peter Conway Humphreys、Adam Santoro による 2024 年の論文「Mixture-of-Depths: Dynamically allocating compute in transformer-based language models」です。公開元は arXiv で、研究分野は Transformer の効率化、条件付き計算、言語モデルの学習・推論最適化です。URL は https://doi.org/10.48550/arXiv.2404.02258 です。
この論文を選んだ理由は、モデルを小さくするのではなく「どのトークンにどれだけ計算を使うか」を学習させる発想が実務的だからです。推論コストを下げたい場面だけでなく、同じ学習予算でより良いモデルを作る設計にもつながるため、開発や事業のヒントを得やすい題材です。
どんな技術か
Mixture-of-Depths は、Transformer の全トークンに全層で同じだけ計算をかけるのは無駄ではないか、という問題意識から生まれた技術です。各層でルーターがトークンごとの重要度を出し、上位 k 個のトークンだけを self-attention と MLP に通し、それ以外はその層をスキップして残差経路で先へ送ります。
ポイントは、どのトークンを重く処理するかは動的に変わる一方で、各層で処理するトークン数 k 自体は固定されていることです。これにより、計算グラフやテンソルサイズは静的なまま保てます。つまり Mixture-of-Depths は、ハードウェア効率を落とさずに、トークン単位で計算の濃淡を付ける Transformer だと考えるとわかりやすいです。
課題
Transformer は通常、シーケンス中のすべてのトークンに対して、すべての層で同じ self-attention と MLP を実行します。しかし実際には、すべてのトークンが同じ難しさではありません。次の予測に強く効くトークンもあれば、比較的単純なトークンや、深い処理を毎層は必要としないトークンもあります。
難しいのは、この「必要なところだけ計算したい」という発想を、そのまま実装すると動的計算グラフになりやすいことです。動的グラフは理論上は柔軟ですが、GPU や TPU のような現在のハードウェアでは、テンソルサイズが毎回変わると実装や最適化が難しくなります。条件付き計算の研究は以前からありますが、実運用しやすい形へ落とし込むのが大きな壁でした。
既存手法にも限界があります。たとえば early exit 系は、一度抜けたトークンは後段の層に戻れません。MoE は expert を切り替えることで疎な計算を実現しますが、通常は「どの expert を使うか」の話であって、「そのトークンは今回この層の重い計算自体を受けなくてよいか」という深さ方向の配分とは少し違います。
この課題を解く必要があるのは、学習でも推論でも Transformer の計算コストが極めて大きいからです。チャット、コード生成、検索、エージェントなどの実システムでは、レイテンシ、GPU コスト、バッチ効率、KV キャッシュの肥大化がすべてボトルネックになりえます。トークンごとに計算の濃淡を付けられれば、同じ予算でより大きいモデルを学習したり、同等性能でより速く推論したりできる可能性があります。
用語解説
- 条件付き計算
- 入力に応じて使う計算量や経路を変える考え方です。この論文はまさにその一種で、各トークンが各層で重い計算を受けるか、スキップするかを学習で決めます。
- Residual connection（残差接続）
- Transformer 層で入力をそのまま次へ流す経路です。Mixture-of-Depths では「その層で重い計算をしない」トークンがこの経路を通るため、完全に情報を捨てるのではなく、深さ方向の計算だけを節約できます。
- Top-k ルーティング
- ルーターが付けたスコアに基づき、上位 k 個だけを選ぶ方式です。Mixture-of-Depths では各層で処理するトークン数を固定するために使われ、動的な選択と静的な計算量を両立する要になります。
- MoE（Mixture of Experts）
- 複数 expert のうち一部だけを各トークンへ適用する疎なモデルです。Mixture-of-Depths は MoE のルーティング発想を深さ方向へ持ち込み、「expert を選ぶ」のではなく「その層を通るかスキップするか」を決めます。
- IsoFLOP 比較
- 総学習 FLOPs を揃えてモデル同士を比較する考え方です。この論文では、同じ学習予算なら MoD がより低い loss を出せるか、あるいは同等性能で 1 ステップ当たりの計算量を減らせるかをこの枠組みで見ています。
技術の仕組み
Mixture-of-Depths の重要な点は、「どのトークンを処理するか」は動的に決めるが、「各層で何個処理するか」は固定していることです。これにより、条件付き計算の柔軟さと、ハードウェア実装のしやすさを両立しています。
基本アイデア
通常の Transformer では、各層ですべてのトークンが self-attention と MLP を通ります。Mixture-of-Depths では、各層にルーターを置き、各トークンに 1 つのスカラー重みを出します。そのスコア上位 k 個のトークンだけが、その層の重い計算に参加します。
選ばれなかったトークンは、その層では更新されず、残差経路で次の層へ進みます。ただし次の層でも再び選ばれる可能性はあります。ここが early exit と違う点です。あるトークンは中間層をいくつか飛ばしつつ、後段の層で再び重要トークンとして処理されることがあります。
モデル構造
構造としては、通常の Transformer block の前に小さな router を付けるイメージです。router はトークンごとにスコアを出し、そのスコアで top-k 選択を行います。選ばれたトークンだけが self-attention と MLP に入り、出力後に元のシーケンス位置へ戻されます。
この設計の面白い点は、attention と MLP の両方をまとめてスキップ対象にしていることです。つまり「このトークンは今回この層で深い変換を受けなくてよい」と判断したら、丸ごと重い計算を省けます。MoE のように expert 間を切り替えるのではなく、block 自体を抜ける発想です。
計算予算の決め方
論文では capacity という考え方で計算予算を定義します。各層で何トークンまで重い計算へ入れてよいかをあらかじめ決めておきます。たとえば系列長の 50% しか処理しないなら、その層の self-attention や MLP の FLOPs は通常よりかなり小さくなります。
重要なのは、capacity を事前固定することで、テンソルの最大サイズが明確になることです。実際にどのトークンが選ばれるかは毎回変わっても、ハードウェア側から見ると「各層で最大 k 個を処理する」という静的な実装にできます。これが研究上のアイデアにとどまらず、実装可能な効率化として成立している理由です。
ルーティング方式
論文は token-choice と expert-choice の考え方を整理したうえで、Mixture-of-Depths では expert-choice 的な top-k 選択を使っています。要するに「トークンが好きな経路を選ぶ」のではなく、「その層の重い計算側が上位 k トークンを選ぶ」形です。
この方式の利点は、各層で必ずちょうど k 個が選ばれるので、負荷が崩れないことです。MoE で問題になりやすい load balancing をかなり単純化できます。一方で、あるトークンが複数の層で頻繁に選ばれ、別のトークンはほとんど選ばれないことも起こります。そこを学習でうまく使い分けるのがポイントです。
学習方法
学習自体は言語モデリングの通常の目的関数で進みますが、router が極端な振る舞いをしないよう補助的な制約も入ります。論文では、各層で capacity に見合った割合のトークンが選ばれるような分布になることを狙っています。
結果として、モデルは「どのトークンにどの層で追加計算を使うと loss が下がるか」を学習します。著者らの観察では、後続トークンの予測が難しい部分、つまりエントロピーが高い予測に関連するトークンほど、より多くの層で処理される傾向が見られました。これは、モデルが難しい箇所へ計算を寄せていることを示唆します。
推論方法
学習中は top-k 選択に未来トークン情報が暗に混ざるため、そのままでは自己回帰生成に使いにくい問題があります。そこで論文では、推論時には predictor ベースのルーティングを使います。つまり、現在までの文脈だけから「このトークンが top-k に入るか」を予測する補助分類器を使います。
この predictor は学習早期から 97% 以上の精度で top-k の選択を模倣でき、自己回帰評価でも性能劣化は小さいと報告されています。これにより、学習時の柔軟な routing と、推論時の因果的な実装をつなげています。
重要な工夫
論文の技術的な肝は 3 つあります。1 つ目は、トークン数上限を固定して静的グラフを保つことです。2 つ目は、early exit のように一度抜けたら終わりではなく、後段層で再参加できることです。3 つ目は、MoE の routing 発想を深さ方向へ持ち込み、block を使うか使わないかを学習可能にしたことです。
著者らはさらに、MoE と組み合わせた MoDE も試しています。これは「どの expert を使うか」と「この層自体を通るか」を両方扱う設計で、条件付き計算をさらに細かく組み合わせられる可能性を示しています。
実験と結果
この論文の評価は、「同じ学習予算ならより良い loss を出せるか」と、「同等性能なら 1 ステップ当たりの計算を減らせるか」の 2 軸で見ると理解しやすいです。
何を検証したのか
著者らは主に次の点を検証しています。
- 同じ総学習 FLOPs で比較したとき、MoD が通常の Transformer より低い学習 loss を達成できるか
- MoD が 1 forward pass 当たりの FLOPs を減らしつつ、同等以上の性能を保てるか
- 学習時の top-k routing を、推論時の因果的 predictor routing に置き換えても性能が落ちにくいか
- MoE と組み合わせたときも利点が積み上がるか
どんな設定を使ったのか
論文では decoder-only の言語モデルで比較を行い、isoFLOP の枠組みで複数サイズを評価しています。重要なのは、単に同じパラメータ数で比べるのではなく、総学習計算量を揃えて比較している点です。これにより「MoD は小さいから速いだけなのか、それとも計算配分そのものが賢いのか」を切り分けています。
また、routing block を全層に入れるだけでなく、full-attention block と交互に入れる構成も試しています。論文中では、少なくともこの設定で 12.5% capacity を隔層に入れる構成が有望だったと報告されています。
どのような結果が出たのか
要点はかなり明快です。MoD は、同じ学習 FLOPs と同程度の学習時間で、ベースライン Transformer と同等かそれ以上の loss を達成しました。そのうえで、forward pass 当たりの FLOPs は小さくでき、自己回帰サンプリングでは 50% 以上高速に step を回せる構成もあったと報告されています。
論文の分析では、isoFLOP 最適な MoD は、同じ学習予算で見るとベースラインより大きいモデルを使えることが多く、そのぶん loss をさらに下げられます。一方で、ベースラインと同等 loss を狙うなら、より小さい MoD 構成でも達成できるため、step 速度の改善が得られます。つまり、浮いた計算を「より大きなモデルを学習する」か「より速く回す」かの両方に振り向けられる設計です。
推論時の predictor は機能したのか
ここは実用上かなり重要です。学習時の top-k routing は非因果的ですが、推論では未来トークンを見られません。論文では predictor による置き換えを行い、その予測精度は学習のかなり早い段階から 97% 以上に達し、自己回帰評価でも性能劣化は小さいと報告しています。
この結果から言えるのは、MoD の利点が学習中だけの見かけ倒しではなく、実際の生成時にもある程度持ち込めるということです。研究アイデアとして面白いだけでなく、サービング設計の候補として見られる理由はここにあります。
結果から何が言えるのか
この論文が示したのは、「Transformer の無駄は重み数だけではなく、トークンごとの計算配分にもある」という点です。全トークンへ均等に深い計算を当てるより、重要トークンへ計算を寄せたほうが、同じ予算で良い学習効率や推論効率を得られる可能性があります。
また、速度改善の源泉が単純な量子化やカーネル最適化ではなく、「そもそもどの計算をしないか」を学習で決める点も重要です。これは今後の長文処理、メモリ設計、エージェント、マルチモーダル処理にも波及しやすい発想です。
何に使える？
Mixture-of-Depths は、単なる論文上の効率化ではなく、実システムの設計方針としていくつかの応用先が考えられます。
低レイテンシな LLM 推論
まずわかりやすいのは、チャットやコード補完のように step ごとの速度が重要な推論です。全トークンを毎層深く処理しないので、forward pass 当たりの FLOPs を減らしやすく、同等品質でより速い生成を狙えます。GPU コストが効く API サービスでは特に価値があります。
長い文脈を扱うシステム
長文コンテキストや大きな履歴を持つシステムでは、すべてのトークンが常に同じ深さで処理される必要はないかもしれません。重要な命令、最新の発話、検索で取れた核心部分へ計算を寄せ、背景的なトークンは浅く流す設計は、長文推論の効率化に向いています。
エージェントやツール利用
エージェントは、ツール定義、実行履歴、中間思考、観測ログなど多様なトークンを抱えます。このとき、すべてのトークンを常に同じ重さで処理するのは非効率です。MoD 的な考え方は、「今の判断に効くトークンへだけ深い計算を当てる」設計に近く、将来的には tool use や memory lookup への routing とも相性がよさそうです。これは論文の議論を踏まえた推測です。
MoE や KV キャッシュ最適化との併用
MoD は MoE と競合するというより補完関係にあります。MoE が expert の疎化なら、MoD は深さ方向の疎化です。両方を併用できれば、アクティブなパラメータとアクティブなトークンの両面で効率化できます。また、論文でも大規模化時のメモリ節約や KV キャッシュへの好影響が示唆されており、サービング基盤と合わせて効かせる余地があります。
開発や事業へのヒント
この論文から得られるヒントは、「モデルの質を保ちつつ速くする方法は、モデル圧縮だけではない」ということです。計算そのものの配り方を学習対象にする発想は、今後の AI プロダクト設計で重要になりそうです。
推論最適化をアーキテクチャで解く
多くの開発現場では、遅いなら量子化、蒸留、キャッシュ、ハードウェア増強という順で考えがちです。MoD は、それに加えて「アーキテクチャ自体が重要度に応じて計算を節約する」選択肢を示しています。モデルを一から作る企業や独自事前学習を行うチームなら、研究投資の対象として十分ありえます。
小規模プロダクトでも思想は使える
論文そのものをすぐ実装しなくても、「入力のどこが本当に重要か」を意識して計算資源を配る発想は、小規模プロダクトでも活かせます。たとえば RAG の再ランキング、長文要約前の文選別、エージェント履歴の圧縮などで、全要素を均等に扱わない設計へつながります。
独自モデルを持つ企業の差別化ポイントになる
API 利用だけでは取り込みにくい技術ですが、逆に言えば独自推論基盤や独自モデルを持つ企業には差別化余地があります。レイテンシや GPU コストはそのまま粗利に響くので、同等品質で step 速度を上げられる設計は、プロダクト競争力に直結します。
今後注目すべき方向性
今後は、MoD のような深さ方向 routing が、長期メモリ、tool routing、マルチモーダル token selection と結びついていく流れが重要だと考えられます。論文でも、将来的には query と key/value の routing を分けたり、memory lookup のような別計算へ振り分けたりする可能性が議論されています。単なる速度改善にとどまらず、「どの種類の計算をどのトークンへ回すか」という設計へ進む余地があります。
限界
Mixture-of-Depths にも注意点はあります。まず、既存の公開済み大規模モデルへ後付けで簡単に差し込める技術ではありません。router を含めた構造で学習する必要があり、既存 dense Transformer のそのままの置き換えよりは導入ハードルが高いです。
また、実装難易度もあります。top-k routing、トークンの再配置、因果的 predictor、カーネル実装まで含めて考えると、単純な Transformer よりサービング系の複雑さは増えます。理論上の FLOPs 削減が、そのまま現場の wall-clock 改善になるかは、実装の質にかなり依存します。
データ依存性にも注意が必要です。どのトークンが重要かはタスクやコーパスで変わるため、すべての分布で同じ capacity 設定が最適とは限りません。論文でも最適 capacity は経験的に決める必要があると示されています。
さらに、skip されたトークンはその層で更新されないため、極端に攻めた設定では情報変換が不足する可能性があります。重要トークン判定が外れると、局所的には性能劣化が起こりえます。推論時 predictor の誤判定もゼロではありません。
最後に、これは主に decoder-only 言語モデルでの結果です。エンコーダ・デコーダ、マルチモーダル、長期メモリ統合などへ一般化したときに同じ利点がどこまで出るかは、今後の検証が必要です。
よくある質問
Q. Mixture-of-Depths は early exit と何が違うのですか？
A. early exit は、一度抜けたトークンが後段層へ戻らない設計が一般的です。Mixture-of-Depths では、ある層をスキップしたトークンが次の層や後段層で再び選ばれることがあります。つまり「完全に打ち切る」のではなく、「必要な層だけ通す」設計です。
Q. Mixture-of-Depths は MoE の代わりになりますか？
A. 代わりというより補完関係です。MoE はどの expert を使うかを選ぶ疎化で、MoD はその層の重い計算自体を受けるかを選ぶ疎化です。論文でも両者を組み合わせた MoDE が検討されており、併用の余地があります。
Q. どんな場面で特に効きそうですか？
A. 推論速度が重要なチャット、コード生成、長文コンテキスト処理などです。とくに、すべてのトークンが毎層同じ深さの処理を受けなくてもよいタスクでは相性がよいです。一方で、短い入力や小さいモデルでは実装コストに対して利得が小さい可能性もあります。
Q. 既存の Llama などに後付けできますか？
A. 厳密には簡単ではありません。論文の中核は routing を前提にした学習であり、既存 dense モデルへ推論時だけ雑に追加して同じ結果が出るとは限りません。研究用途の再学習や、独自モデルの設計段階で検討する技術だと考えたほうがよいです。
Q. この論文から実務で一番学べることは何ですか？
A. 「計算量を減らす」とは、モデルを小さくするだけではなく、重要な入力へ計算を集中させることでも達成できる、という点です。この発想は RAG、再ランキング、履歴圧縮、ツール選択など、周辺システム設計にも応用しやすいです。
今日の学び
この論文は、Transformer が全トークンに全層で均等な計算をかけてしまうため、学習・推論の両方で無駄が大きいという課題を扱いました。そこに対して Mixture-of-Depths は、各層で上位 k 個の重要トークンだけを重く処理し、他は残差経路でスキップさせることで、静的グラフのまま動的な計算配分を実現しました。
ここから得られるヒントは、AI システムの効率化は圧縮や量子化だけでなく、「どの入力にどれだけ計算を使うか」を設計する方向でも進むということです。今後の LLM、エージェント、長文処理では、計算の配り方そのものが差別化ポイントになりそうです。