Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation

基本情報

• arXiv ID: 2507.10524 (https://arxiv.org/abs/2507.10524)
• 著者: 未記載
• 所属: 未記載
• 投稿日: 2025年07月15日
• カテゴリ: cs.LG, cs.AI

簡単に説明すると

この論文は、大規模言語モデル（LLM）の計算効率を根本的に改善する新しいアーキテクチャ「Mixture-of-Recursions（MoR）」を提案しています。
従来のLLMは、すべての単語（トークン）に対して同じ量の計算をしていましたが、実際には簡単な単語と複雑な単語が混在しています。
MoRは、各トークンの難易度に応じて動的に計算量を調整する仕組みです。
具体的には、同じレイヤーセットを何度も繰り返し使用（再帰）し、軽量なルーターが各トークンに必要な再帰回数を決定します。
これにより、パラメータ数を約50％削減しながら、同等以上の性能を達成し、推論速度も最大2.18倍に向上します。

1. 研究概要

1.1 背景と動機

Transformerネットワークを数千億パラメータにスケールしたことで、印象的な推論能力が実現されました。
しかし、それに伴うメモリや計算の要求は、データセンター以外での学習や展開を困難にしています。
この問題に対して、計算量とメモリ使用量を削減する設計が模索されてきました。

モデルの効率化の軸として、「パラメータ効率」（重みの削減や共有）と「適応的計算」（必要なときだけ計算を増やす）が有望です。
レイヤータイイング（重みの再利用）や早期終了（シンプルなトークンは早く出力）などの手法が提案されてきましたが、両方を効果的に統合したアーキテクチャはありませんでした。

1.2 主要な貢献

本研究の主要な貢献は以下の通りです。

• 言語モデリングの計算効率を向上させる統一フレームワーク「Mixture-of-Recursions（MoR）」の提案
• トークンごとに動的に再帰深度を割り当てるルーターの導入
• 135Mから1.7Bパラメータのモデルでの幅広い実証検証
• 再帰別KVキャッシングによるメモリ効率の向上

2. 提案手法

2.1 手法の概要

Mixture-of-Recursions（MoR）は、Recursive Transformersの可能性を最大限に活用する統一フレームワークです。
MoRは、軽量なルーターをエンドツーエンドで学習し、トークン固有の再帰深度を割り当てます。
各トークンに必要な「思考」の深さに応じて、共有パラメータブロックが何回適用されるかを決定します。

この動的なトークンレベルの再帰は、再帰別のkey-value（KV）キャッシングを可能にします。
各トークンの再帰深度で必要なKVペアのみを選択的に保存・取得します。
これにより、メモリトラフィックが削減され、スループットが向上します。

2.2 技術的詳細

パラメータ共有戦略として、Cycle、Sequence、Middle-Cycle、Middle-Sequenceの4つを検討しました。
Cycle共有では、再帰ブロックがサイクリックに再利用されます。
Sequence共有では、各再帰ブロックが同じレイヤーを連続して再利用します。

ルーティング戦略には、Expert-choiceとToken-choiceの2つがあります。
Expert-choiceでは、各再帰深度がエキスパートとなり、上位K個のトークンを選択します。
Token-choiceでは、各トークンが最初から完全な再帰ブロックシーケンスにコミットします。

KVキャッシング戦略として、再帰別キャッシングと再帰共有を提案しています。
再帰別キャッシングでは、特定の再帰ステップにルーティングされたトークンのみがKVエントリを保存します。
再帰共有では、最初の再帰ステップでのみKVペアをキャッシュし、後続のすべての再帰で再利用します。

2.3 新規性

既存の手法はパラメータ効率と適応的計算のどちらか一方に焦点を当てていましたが、MoRはこれらを統一しました。
従来の再帰モデルは固定深度でしたが、MoRは事前学習中に動的な再帰深度を学習します。
さらに、再帰別KVキャッシングにより、推論時のメモリアクセスを約50%削減します。

3. 実験結果

3.1 実験設定

LlamaベースのTransformerアーキテクチャを使用し、SmolLMオープンソースモデルの設定を参考にしました。
FineWeb-Eduデータセットの重複除去サブセットで事前学習をしました。
評価はFineWeb-eduの検証セットと6つのfew-shotベンチマークで行いました。

3.2 主要な結果

同じ学習計算量（16.5e18 FLOPs）のもとで、MoRモデルはバニラベースラインを上回りました。
Expert-choiceルーターと2回の再帰を使用したMoRモデルは、検証損失が低く、few-shot精度が43.1%（バニラは42.3%）でした。
これはパラメータ数が約50%少ないにもかかわらず達成されました。

同じトークン数（20B）での比較では、MoRは25%少ない学習FLOPsで優れた性能を示しました。
実用的には、学習時間が19%短縮された上、ピークメモリ使用量も25%削減されました。

3.3 既存手法との比較

IsoFLOP分析で、135Mから1.7BのモデルサイズでMoRを評価しました。
135Mではバニラに劣りましたが、360M以上ではMoRがバニラを上回りました。
推論スループットでは、連続深度方向バッチングにより、最大2.06倍の高速化を達成しました。

4. 実用性評価

4.1 実装の容易性

MoRは既存のTransformerアーキテクチャに基づいており、実装が比較的容易です。
ルーターは軽量な線形層やMLPで構成でき、追加のパラメータ数は最小限です。
パラメータ共有戦略はMiddle-Cycleが最も効果的だと実験で示されました。

4.2 計算効率

分散学習の観点から、FSDPを使用すると高効率化が可能です。
通常の1回のall-gather操作で1イテレーションをサポートするのに対し、再帰モデルではN_rイテレーションをサポートします。
再帰別キャッシングにより、KVメモリとIOを約(N_r+1)/2N_r倍削減します。

4.3 応用可能性

推論タスクへの適用が有望です。
最近の研究では、推論チェーン内の冗長性が指摘されており、MoRの適応的計算が効果的です。
マルチモーダルや非テキスト領域への拡張も可能です。
再帰ブロックは本質的にモダリティに依存しないため、視覚、音声、統一マルチモーダルアーキテクチャに容易に統合できます。

5. まとめと所感

5.1 論文の意義

この論文は、LLMの効率化における重要なブレークスルーを提供しています。
パラメータ効率と適応的計算を統一した初めてのアーキテクチャであり、実用的な意義が大きいです。
特に、事前学習段階から動的な再帰深度を学習するアプローチは新しく、従来の事後的な早期終了手法の問題を回避しています。

実験結果から、大規模モデルの品質を維持しつつ、約50%の効率化を達成できることが確認されました。
これにより、計算リソースが限られた環境でも大規模モデルの能力を活用できる可能性が広がります。

5.2 今後の展望

推論MoRモデルの開発が重要な課題です。
ルーターがChain-of-Thoughtの必要性を判断して動的調整し、推論精度と効率の両方を向上できます。

より大規模なモデル（3Bパラメータ以上）での検証も必要です。
計算制約のため、本研究では1.7Bまでに留まりましたが、より大規模なスケールでの性能検証が期待されます。

スパースアルゴリズムとの互換性も有望です。
構造化スパーシティを統合することで、トークンとレイヤーの両レベルで不要な計算を動的に削减できます。
プルーニングや量子化などの技術もMoRと高い補完性を持ちます。