今回の論文
今回取り上げるのは、Sangmin Bae、Yujin Kim、Reza Bayat らによる論文「Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation」です。2025年7月14日に arXiv で公開された論文で、公開元は arXiv、研究分野は LLM アーキテクチャ設計、推論最適化、学習効率化です。URL は https://arxiv.org/abs/2507.10524 です。
この論文を選んだ理由は、LLM を軽くする方法としてありがちな「量子化」や「蒸留」だけでなく、モデルの通し方そのものを変えているからです。同じブロックを使い回しつつ、必要なトークンにだけ追加計算を回すという考え方は、推論系プロダクトでも学習系の設計でも応用しやすいです。
どんな技術か
Mixture-of-Recursions は、Transformer の中間ブロックを再帰的に何度も適用しながら、すべてのトークンを同じ深さまで処理しないようにする技術です。簡単にいうと、モデル全体を深く大きくする代わりに、同じ計算ブロックを繰り返し使い、難しいトークンだけ多く考えさせます。
通常の Transformer では、各トークンが全レイヤーを必ず通ります。しかし実際には、句読点や定型表現のような簡単なトークンと、文脈依存が強い難しいトークンでは必要な計算量が違うはずです。Mixture-of-Recursions はそこに着目し、ルーターが各トークンに必要な再帰回数を決めます。
その結果として、重み共有によるパラメータ削減と、トークン単位の適応的な計算配分を同時に実現しようとしています。論文では、これによって同じ FLOPs 予算でより多くの学習トークンを処理でき、推論時の KV キャッシュも減らせると示しています。
課題
この技術が解決しようとしているのは、LLM を大きくすると性能は上がりやすい一方で、学習コスト、メモリ使用量、推論レイテンシが急激に悪化するという課題です。
何が難しいのかというと、既存の効率化手法は改善軸が分かれがちだからです。重み共有はパラメータ数を減らせますが、すべてのトークンに同じ計算をする限り、無駄な FLOPs は残ります。逆に early exit のような適応計算は計算量を減らせますが、追加の学習設計や推論時の実装が複雑になりやすいです。
既存の方法ではどこに限界があるのかというと、まず固定深さの再帰モデルは「同じ重みを使い回す」点では効率的でも、簡単なトークンと難しいトークンを区別できません。また、通常の early exit は層の途中で抜ける設計が多く、Transformer の重み共有ときれいに統合しにくいという問題があります。
なぜこの課題を解く必要があるのかというと、実際の AI システムでは計算資源の制約が常に効くからです。長文コンテキストを扱うチャット、エージェント、コード補完、バッチ推論などでは、少しのメモリ削減やスループット改善がそのままコスト差になります。特に「モデル品質は落としたくないが、GPU 台数もメモリも限られる」という場面では、モデル構造レベルの工夫が効きます。
実際のシステムでは、すべてのトークンを均一に深く処理する必要がないケースが多いです。たとえば定型文の多い業務文書、ログ解析、RAG での長いコンテキスト処理では、重要トークンだけを深く見る設計にできれば、品質を保ちながらコストを下げられる可能性があります。
用語解説
- 再帰型 Transformer
- 同じレイヤーブロックを複数回適用して、実効的な深さを稼ぐ Transformer です。この論文では普通の深い Transformer を作る代わりに、共有ブロックを繰り返し通す土台として使われます。
- 重み共有
- 複数の層で同じパラメータを使い回す設計です。Mixture-of-Recursions のパラメータ削減効果の中心であり、「大きなモデルらしい深さ」を少ない固有パラメータで表現するために重要です。
- ルーティング
- 各トークンをどの経路や深さに送るかを決める仕組みです。この論文では各トークンに何回再帰ブロックを通すかを決める役割を持ち、計算量の自動配分を支えています。
- KV キャッシュ
- 自己注意の key と value を保存して、デコード時に再計算を減らす仕組みです。長文推論では便利ですがメモリを多く使うため、MoR では「深く計算するトークンだけ KV を厚く持つ」設計が効きます。
- isoFLOP 比較
- 総 FLOPs をそろえてモデルを比較する評価方法です。単純な精度比較ではなく、「同じ計算予算ならどれだけ良いモデルになるか」を見るため、この論文の価値を理解する上で重要です。
技術の仕組み
Mixture-of-Recursions のポイントは、1つの共有ブロックを何度も回すこと自体ではなく、その回数をトークンごとに変えるところにあります。これによって、重み共有、適応計算、KV キャッシュ削減が1つの枠組みの中でつながります。
基本アイデア
通常の Transformer は、全トークンが全レイヤーを順番に通ります。MoR では、最初と最後の一部レイヤーは通常通り持ちながら、中間部分を共有された再帰ブロックとして扱います。そして各トークンがその共有ブロックを何回通るかを、軽量なルーターが決めます。
イメージとしては、「全員に同じ授業を最後まで受けさせる」のではなく、「理解が早い人は早く抜け、難しい部分だけ補習を重ねる」に近いです。これが token-level adaptive computation の本質です。
モデル構造
論文では複数の共有方法を比較していますが、最も良かったのは Middle-Cycle です。これは最初と最後のレイヤーは固有のまま残し、中間レイヤー群だけを共有して再帰させる構成です。
この設計が効く理由は、入力に近い層と出力に近い層には固有の役割を残しつつ、計算量の大きい中間部分で共有の恩恵を最大化できるからです。論文でも 10B トークン学習の比較で Middle-Cycle が他の共有方式より一貫して良い NLL を示しています。
ルーティングの仕組み
MoR では各トークンに対して、再帰を何回回すかをルーターが決めます。論文では主に expert-choice と token-choice の2種類を比較しています。
expert-choice は各再帰ステップごとに「このトークンを次の再帰に進めるか」を選ぶ方式です。難しいトークンだけが次の深さへ進み、簡単なトークンは早めに計算を終えます。論文の主結果でも、この方式が最も安定して高性能でした。
一方の token-choice は、最初に各トークンの再帰回数をまとめて決める方式です。実装は分かりやすいですが、論文では N_r=3 の比較で few-shot 平均が 40.0% と、expert-choice の 42.6% を下回りました。つまり「最初に深さを固定する」より、「途中で残すかどうかを段階的に決める」ほうがうまくいったわけです。
KV キャッシュの扱い
MoR のもう1つの実務的なポイントは、再帰段数ごとに必要なトークンだけ KV キャッシュを持つことです。再帰が深くなるほど、生き残るトークン数は減ります。すると、深い段で保持する KV も減らせます。
論文では、この recursion-wise caching によって KV メモリと IO を圧縮し、推論スループットを改善できると説明しています。別案として、最初の再帰で作った KV を後段でも再利用する KV sharing も試していますが、これはメモリには有利でも性能はやや落ちました。そのため論文の主結果では、共有しすぎない independent caching が主に使われています。
学習方法
モデルは Llama ベースの Transformer を土台にし、SmolLM 系の設定を参照してゼロから事前学習しています。学習データは FineWeb-Edu の重複除去済みサブセットで、評価には FineWeb-Edu の validation NLL と、LAMBADA、HellaSwag、PIQA、WinoGrande、ARC、MMLU の few-shot ベンチマークが使われています。
重要なのは、MoR が後付けの推論トリックではなく、最初からこの構造で学習されていることです。つまり「推論時だけ token pruning をする」のではなく、「深さを可変にする前提でモデルを育てる」設計です。これにより、学習時と推論時の挙動がずれにくくなります。
処理の流れ
1つのトークンが処理される流れを単純化すると、まず入力が通常の初段レイヤーを通ります。次に共有された再帰ブロックに入り、ルーターが「このトークンはもう十分か、それとももう1段深く処理するか」を判断します。残ったトークンだけが次の再帰へ進み、そのたびに注意計算の対象トークン数と KV 保持量が減っていきます。最後に終端側のレイヤーで出力が作られます。
この流れは、Sequence Length を横方向に削る token pruning と少し似ていますが、MoR は主に深さ方向の計算量制御をしています。つまり「どの token を何層分考えるか」を制御する設計です。
実験と結果
論文では、MoR が本当に効率と性能の両立に寄与するのかを、学習時の isoFLOP 比較、固定トークン数比較、推論スループット比較、共有方式やルーティング方式のアブレーションで検証しています。
何を検証したのか
主な検証点は3つあります。1つ目は、同じ学習 FLOPs なら Vanilla Transformer より良い性能を出せるかです。2つ目は、同じ学習トークン数ならより少ない FLOPs とメモリで同程度以上の性能が出るかです。3つ目は、推論時に KV キャッシュ削減が実際のスループット改善につながるかです。
データセットと評価指標
事前学習には FineWeb-Edu の重複除去済みデータを使い、主実験では 20B トークンまたは同等 FLOPs 条件で比較しています。評価指標は、言語モデルとしての負の対数尤度である NLL と、6つの few-shot ベンチマーク平均精度です。推論面では throughput を比較しています。
同じ FLOPs なら、より少ないパラメータで Vanilla を上回った
16.5e18 FLOPs をそろえた比較では、Vanilla 315M モデルの few-shot 平均が 42.3% だったのに対し、MoR の expert-choice、再帰2回、167M パラメータ版は 43.1% でした。Validation NLL も 2.7824 から 2.7511 へ改善しています。
ここで重要なのは、MoR が単に「小さいけれど少し負けるモデル」ではなかった点です。固有パラメータはほぼ半分まで減っているのに、同じ計算予算ではむしろ性能が上がっています。論文では、効率が良いぶん同じ FLOPs で 27B トークンまで学習できたことが効いていると説明しています。
同じ 20B トークンでも学習コストを削れた
20B トークン固定の比較でも、MoR の expert-choice、再帰2回モデルは 12.3e18 FLOPs で済みながら、Vanilla より良い NLL と few-shot 平均 42.9% を出しました。論文本文では、Vanilla 比で学習時間を 19% 短縮し、ピークメモリを 25% 削減したと述べています。
これは実務的にかなり大きいです。同じデータ量で同じくらいの品質を出すだけでも意味がありますが、MoR はそこからさらに少し勝っています。学習予算が厳しいチームほど、このタイプの改善は効きます。
推論スループットも改善した
論文では continuous depth-wise batching と組み合わせた推論評価も行っています。その結果、MoR-4 は最大バッチ条件で Vanilla 比 2.06 倍、固定バッチでも 1.60 倍のスループット改善を示しました。MoR-2 や MoR-3 でも一貫して高速化しています。
ここでのポイントは、MoR が「理論上 FLOPs が減る」だけでなく、KV キャッシュ量の削減によって実際の推論処理でも有利になったことです。長文推論や多数同時リクエストをさばく環境では、この差が直接インフラコストに跳ねます。
ただし大規模化ではまだ課題が残る
一方で 1.7B スケールの比較では、MoR の expert-choice 版は平均 48.4% と強い結果を出しつつも、Vanilla の 48.9% をわずかに下回りました。論文自身も、現在の MoR 設計は大規模化で最適化余地が残ると述べています。
つまりこの論文の主張は「常に Vanilla を完全に置き換える」ではなく、「中小規模から中規模でかなり有望で、さらに推論効率の改善も見込める」という読み方が正確です。
何に使える？
Mixture-of-Recursions の考え方は、単に論文用モデルを作る話ではなく、計算資源の使い方を変える設計として複数の用途に応用しやすいです。
長文チャットや社内アシスタント
長い会話履歴や社内ドキュメントを読むアシスタントでは、すべてのトークンに同じ深さの計算をかける必要がない場面が多いです。MoR 的な発想を使うと、重要文脈や曖昧なトークンにだけ計算を厚く配分する設計を検討できます。長文になるほど KV キャッシュ効率の差も効きやすいです。
RAG の生成側最適化
RAG は検索側ばかり最適化されがちですが、取得した長いコンテキストを読む生成モデル側のコストも大きいです。MoR のように「読み込みの深さをトークンごとに変える」設計は、引用が必要な箇所や質問に近い箇所だけ深く処理する発想につながります。これは論文の直接検証ではなく応用上の推測ですが、かなり自然な方向性です。
エージェントのツール実行前後の判断
エージェントでは、ユーザー入力の一部やツール返り値の一部だけが本当に重要なことがあります。たとえば API レスポンスのうち結論に関わる部分だけを深く処理したい場面です。MoR の「トークンごとに thinking depth を変える」考え方は、将来的な agent backbone 設計にもつながりそうです。
エッジ寄り・小規模 GPU 環境
固有パラメータ数を抑えつつ性能を稼げるため、モデルサイズの制約が厳しい環境でも相性があります。量子化だけに頼らず、構造側で軽くする選択肢として有効です。特に少数 GPU での継続学習や専用モデル構築では、学習メモリと時間の削減がそのまま実行可能性に関わります。
深さ方向の適応計算が効く他モダリティ
この論文自体は言語モデル中心ですが、「難しい入力にだけ深い計算を回す」という考え方は画像や音声にも拡張しやすいです。すでに token merging や adaptive computation の流れは他分野にもあるので、MoR 的な再帰構造は今後マルチモーダル側でも出てきやすいと考えられます。ここは将来予想を含みます。
開発や事業へのヒント
この論文から得られるヒントは、「モデルを軽くする」と「品質を保つ」を別々に考えないほうがよい、ということです。構造設計の段階で計算配分を変えれば、品質と効率を同時に触れます。
全入力を均一に処理しない設計を検討する
多くの AI アプリでは、入力長や候補数を削る方向の最適化が先に来ます。しかし MoR は、入力を削らなくても「重要な部分だけ深く考える」ことでコストを抑えられる可能性を示しています。RAG、要約、コード解析などでも、深さ方向の可変計算は設計候補になります。
共有ブロックを前提にした専用モデル設計は現実的
LoRA や量子化は既存モデルの後処理として便利ですが、サービスが継続運用段階に入り、特定用途向けに自前モデルを設計するなら、最初から共有ブロック前提で作る価値があります。MoR は「最初からこの構造で学習する」重要性を示しています。
推論コスト改善は KV キャッシュの設計まで含めて考える
推論最適化というとデコードアルゴリズムや量子化に目が向きやすいですが、論文では KV キャッシュの持ち方が throughput に効いています。長文 LLM サービスを作るなら、モデル本体だけでなく KV メモリの削減策を設計の中心に置くべきだという示唆があります。
小さなプロダクトでも部分導入の発想は使える
自前で MoR を実装しなくても、「難しいケースだけ追加計算を回す」という発想は使えます。たとえば軽い一次判定モデルと重い二次判定モデルの段階実行、重要文だけ再ランクして深く読む RAG、曖昧な問い合わせだけ追加ツールを呼ぶエージェントなどは、MoR と同じ思想です。
今後注目すべき方向性
今後は、token pruning、early exit、recursive model、MoE、KV cache compression が別々に進むのではなく、統合的に設計される流れが強まりそうです。MoR はその先行例として見ると面白いです。特に「どのトークンにどれだけ計算を割くか」を学習させる方向は、推論コストが事業性に直結するプロダクトで注目しておく価値があります。
限界
まず、実装は通常の Transformer より複雑です。再帰ブロック、ルーター、深さごとのバッチング、KV キャッシュ管理まで含めると、学習も推論も素直な既存実装では済みません。論文の効果を実運用で再現するには、カーネルや serving 側の対応が重要になります。
次に、性能は routing 設計にかなり依存します。論文でも token-choice は expert-choice より明確に弱く、ルーターの設計次第で性能を落としやすいことが分かります。つまり「再帰にすれば勝てる」のではなく、どのトークンを深く計算するかの学習が核心です。
また、1.7B スケールでは Vanilla をわずかに下回っており、大規模化で常に優位とはまだ言えません。論文でも、非共有ブロックを増やすことや depth-specific LoRA / MoE の導入など、追加改良の余地を挙げています。
データ依存性にも注意が必要です。評価は FineWeb-Edu と一般的な few-shot ベンチマークが中心で、コード特化、マルチモーダル、実運用ログ、RAG 特化条件で同じ傾向が出るかは未検証です。
さらに、MoR の利点は「可変深さを前提に最初から学習する」ときに出やすいので、既存巨大モデルへ後付けで簡単に適用できるわけではありません。導入障壁は低くありません。
よくある質問
Q. Mixture-of-Recursions は Mixture-of-Experts と同じものですか？
A. 似ているのは「ルーティングで計算を振り分ける」という発想だけです。MoR は複数の expert を持つのではなく、共有された同じブロックを何回通すかをトークンごとに変えます。つまり expert を選ぶというより、thinking depth を選ぶ設計です。
Q. この技術は推論専用ですか、それとも学習にも効きますか？
A. 論文の主張は両方です。推論では KV キャッシュ削減と depth-wise batching による高速化があり、学習では同じ FLOPs でより多くのトークンを処理できるため、結果として品質改善にもつながっています。
Q. 既存の LLM に後付けで適用できますか？
A. 部分的な発想は流用できますが、論文と同じ形での適用は簡単ではありません。MoR は再帰構造とルーティングを前提に最初から学習しているため、既存モデルへそのまま差し込むより、自前アーキテクチャとして設計する寄りの技術です。
Q. 量子化や speculative decoding と比べると何が違いますか？
A. 量子化や speculative decoding は既存モデルの実行方法を改善する色が強いです。一方 MoR は、モデル内部でどのトークンに何回計算を回すかという構造そのものを変えます。したがって競合というより、組み合わせ対象として見るほうが自然です。
Q. 実務でまず真似するなら、どこから取り入れるのがよいですか？
A. いきなり再帰 Transformer を実装するより、「難しい入力だけ追加計算を回す」二段階設計から入るのが現実的です。これは RAG の再読込、ツール選択の追加判定、長文要約の再精査などに落とし込みやすく、MoR の発想を小さく試せます。
今日の学び
この論文は、LLM を大きくすると性能は上がるが、計算量とメモリが重くなりすぎるという課題を扱っています。そこで、共有レイヤーを再帰的に使い回しつつ、難しいトークンだけ深く処理する Mixture-of-Recursions という構造で解こうとしました。
そこから得られるヒントは、効率化を単なる後処理にせず、モデルがどこに計算を使うべきかを構造として学ばせることです。AI アプリや基盤モデル開発でも、「全入力を均一に処理しない」という発想は、今後かなり重要になりそうです。