🤖 AI

Looped State-Space Language Models with Adaptive Exit-State Selection

本論文は、Mambaベースの状態空間モデルにループ構造を適用することで、計算の深さを増すことにより推論能力とパラメータ効率を向上させると同時に、予測の深さを最適化するための適応的な終了状態選択を導入していることを実証しているが、実際の推論時における節約を実現するにはさらなる状態処理メカニズムが必要であると指摘している。

原著者： Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

公開日 2026-07-14

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

CC BY 4.0

原著者： Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

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

あなたは、非常に優秀で超スマートなロボットシェフ「マンバ（Mamba）」を想像してみてください。このシェフは驚異的な速さとメモリ効率の良さで有名で、大量の異なる道具を備えた巨大なキッチンを必要とすることなく、料理を次々と作り上げることができます。しかし、一つだけ落とし穴があります。本当にトリッキーなレシピ（複雑な数学パズルや「隠された材料を見つける」ゲームなど）に直面すると、マンバは時々行き詰まってしまうことがあります。もっと深く考える必要がありますが、深く考えるための「層（レイヤー）」が足りないのです。新しい巨大なキッチンを作ることは、莫大な費用がかかるためできません。

そこで、立命館大学と東京大学の研究者たちが、シンプルで大胆な問いを投げかけました。「もし、もっと大きなキッチンを作るのではなく、マンバに同じ料理を何度も何度も作らせ、そのたびに味を洗練させていくとしたらどうだろうか？」

「ループする」キッチン：一人のシェフ、多くの工程

AIの世界では、ほとんどのモデルは工場の組立ラインのようなものです。データ（文章の中の一つの単語など）は、24個の異なるステーション（層）を通り抜け、終わります。ロボットにもっと深く考えてほしい場合、通常はステーションを追加します。

著者たちは異なる方法を試みました。彼らは「ループ型マンバ（Looped Mamba）」を構築しました。マンバが、たった一つの非常に有能なステーションであると想像してください。データを24個の異なるステーションへと流す代わりに、彼らはデータを同じステーションへと24回連続で送ります。

魔法： ロボットは、すべてのパスにおいて全く同じ脳（パラメータ）を使用します。これは、学生が4つの異なる混乱した段落を読むのではなく、理解を深めるために同じ教科書の段落を4回読み返すようなものです。
結果： 「Mano」（モジュロ演算）や「p-hop induction」（隠されたパターンの発見）と呼ばれるトリッキーな論理パズルにおいて、この「読み返し」戦略は素晴らしい成果を上げました。わずか4層のモデルを6回ループさせたもの（4 ⊗ 6）は、24個のユニークな層を持つ標準的なモデルを圧倒しました。これは、推論においては、「思考の数（ツールの数）」よりも「思考の深さ」が重要であることを示唆しています。

「出口ゲート（Exit Gate）」：最高のバージョンを選ぶ

しかし待ってください。もしロボットが永遠に段落を読み返し続けたらどうなるでしょうか？それは時間の無駄です。研究者たちは、「出口ゲート」と呼ばれる賢い機能を加えました。

このゲートは、スマートなセレクター（選択器）のようなものです。マンバがデータを処理する際、ゲートはこう問いかけます。「答えは明確になってきているか？」

もし2回のループで答えがすでに明確であれば、ゲートはこう言います。「ストップ！このパスの結果を使用します。」
もし問題が難しい場合は、ゲートはこう言います。「続けて！」と言い、3回目または4回目のループの結果を選択します。

これは「適応型出口状態選択（Adaptive Exit-State Selection）」と呼ばれます。論文では、このゲートが1億4000万パラメータ程度の小さなモデルにおいて、ゲームチェンジャーとなったことが示されています。これにより、モデルは各単語に対して「ちょうど良い」思考時間を選択できるようになり、常に最大ループ数の結果を使用するように強制する場合よりも、しばしば優れたパフォーマンスを発揮しました。

しかし、ひねりがあります。研究者たちは、これがまだコンピュータの電気代や時間を実際に節約しているわけではないということを非常に慎重に注記しています。たとえゲートがより早い段階の結果を「選択」したとしても、ロボットはバックグラウンドですべてのループを物理的に実行しています。なぜなら、マンバのメモリ状態は連続しているからです。つまり、ループ3の結果はループ2によって残された状態に依存し、それはループ1に依存しています。後続の単語のための連鎖を壊さずに、途中でロボットを「脱落」させることはできません。したがって、ゲートは「予測の深さ（どのバージョンの答えを使うか）」を選択しますが、「実時間での計算量（wall-clock computation）」は変わりません。実際に時間を節約するには、ロボットのメモリ状態を扱う新しい方法が必要であり、著者たちはこれを将来の課題としています。

現実的な検証：うまくいかなかったこと（そして、まだ可能性のあること）

この論文は、解決できなかったことについても率直に述べています。

大きいことが常に良いとは限らない（この手法において）： 彼らが、計算量（FLOPs）が同等になるように調整した標準的な非ループ型モデルと比較したところ、標準的なモデルは「モデルがどれほど混乱しているか（パープレキシティ）」という基本的なテストにおいて、依然として勝利しました。ループ型モデルは推論タスクには優れていましたが、単に次の単語を予測するという点では、標準的な深いモデルの方が依然としてわずかに優れていました。
サイズが重要： 「出口ゲート」は1億4000万パラメータのモデルでは見事に機能しました。しかし、3億7000万パラメータのより大きなモデルで試したところ、ゲートによる改善はあまり見られませんでした。大きなモデルは、とにかくすべてのループを実行することを好むようです。著者らは、大きなモデルは最終的な、深い思考のパスにより強く依存している可能性があると示唆しています。
すべてに対する魔法の杖ではない： 研究者たちは、合成パズル（数学の木構造やパターン探索など）を用いてテストを行いました。これらはコンセプトが機能することを証明していますが、この「ループ」のトリックが、小説を書いたり複雑な法的案件を解決したりするような、オープンエンドで現実世界の推論においても同様にうまく機能するかどうかは、まだ分かっていないと彼らは認めています。

結論

この論文は、同じスマートな脳を何度も再利用することが、より大きく高価な脳を構築することなく、AIを推論において賢くするための強力な方法であることを示唆しています。これは「パラメータ効率の良い」スケーリングの軸です。

証明されたこと： 特定の論理パズルにおいて、小さなマンバモデルをループさせることは、標準的なモデルと同等のサイズで、かつより大きなモデルに匹敵する性能を発揮します。
示唆されたこと： このアプローチは、特にステップ・バイ・ステップの深い思考を必要とするタスクにおいて、将来のAIをより効率的にするための鍵となる可能性があります。
未解決の課題： 「出口ゲート」を（単に結果を選択するだけでなく）実際に計算時間を節約するようにする方法や、今日の最大のAIを動かしている数十億のパラメータを持つ大規模モデルでもこれが通用するかどうかを、私たちはまだ解明する必要があります。

要約すると、著者たちは単に大きなロボットを作ったのではありません。より小さなロボットに、問題に対して二度目（あるいは三度目、四度目）の検討を行うことで、より深く考える方法を教えたのです。そして今のところ、その「二度目の検討」は非常に有望な兆しを見せています。

技術要約：適応型エグジット状態選択を備えたループ型状態空間言語モデル

問題提起

近年の大規模言語モデル（LLM）の進歩は、主にモデルサイズ、学習データ、および計算集約的な推論のスケールアップによって推進されてきた。しかし、このスケーリングのパラダイムは、Transformerベースの自己注意機構における巨大なキー・バリュー（KV）キャッシュの維持に起因する、メモリ消費量の増大や推論コストの急速な上昇といった、重大なデプロイメント上の課題をもたらしている。Mamba-2に代表される状態空間モデル（SSM）は、良好なロングコンテキストのスケーリング特性を持つ効率的な代替案を提供するが、「独立したパラメータよりも計算の深さを優先する」という原理（近年、Looped Transformerにおいて実証されたもの）がSSMにも適用可能であるかどうかは依然として不明である。さらに、既存の研究はほぼ排他的にTransformerバックボーンを用いたループ構造に焦点を当てており、MambaやハイブリッドMamba-Transformerアーキテクチャ内において、再帰的なパラメータ共有がどのように機能するかについての理解には空白が残されている。

手法

著者らは、Looped MambaおよびLooped Hybrid Mamba–Transformerアーキテクチャを提案する。コアとなるメカニズムは、単一のMamba（またはハイブリッド）ブロックのパラメータを層間で共有し、それを明示的な有限深度のループを通じて繰り返し適用することである。

アーキテクチャ

表記: モデルは $N \otimes R$ と表記される。ここで、 $N$ は独立してパラメータ化された層の数、 $R$ は再帰ステップ数である。実効深度は $NR$ となる。
Looped Mamba: 共有された $N$ $N$ 層のMambaバックボーン $M_\theta$ $M_{θ}$ を繰り返し適用する。
- $h^{(r)} = M_\theta(h^{(r-1)})$ （ただし $r = 1, \dots, R$ ）。
- 最終出力は、ステップ $R$ における隠れ状態から導出される。
Looped Hybrid: Mamba-2層と、少数の因果的自己注意層の両方を含むバックボーンにループを拡張する。このループはハイブリッドブロック全体に適用され、状態空間の更新とアテンションによる相互作用の両方を通じて、隠れ表現の反復的な洗練を可能にする。
学習プロトコル:
1. 事前学習: 最終ループ目的関数（損失はステップ $R$ でのみ計算）を用いて最適化される。
2. Exit-Gate ステージ1: すべての再帰ステップからの出力を用いて、エグジット深度に対するトークンレベルの分布を学習するために、エントロピー正則化された適応計算目的関数（Ouroから適応）を導入する。
3. Exit-Gate ステージ2: バックボーンとヘッドを凍結し、次のステップが予測損失を改善する場合にのみ継続することを促すように、エグジットゲートのみを学習する。

適応型エグジット状態選択

本論文は、Ouroの2段階のエグジットゲートをMambaに適応させている。推論時には、閾値 $q$ が累積エグジット確率に基づいて各トークンのエグジットステップ $z_i$ の選択を制御する。

意味論的注記: 著者らは、SSMにおいては状態の連続性の制約により、全 $R$ ステップの状態チェインを保持しなければならないことを明確にしている。したがって、「エグジット」はどの再帰ステップの隠れ状態を言語モデリングヘッドに渡すかを選択するものであり、追加の状態処理メカニズムなしには、本質的にウォールクロック計算量や実行ステップ数を削減するものではない。

主な貢献

制御された経験的研究: 著者らは、非ループ型のMambaおよびTransformerのベースラインに対し、等パラメータ（同一の $N$ 、異なる $R$ ）および等FLOPs（同一の有効深度 $NR$）のプロトコルの下で、Looped Mambaおよびハイブリッドモデルの厳密な比較を提示している。
推論タスクの評価: 合成推論タスク（Manoおよび $p$ -hop induction）および言語モデルの事前学習とダウンストリームベンチマークにおいて、本アプローチを検証している。
適応型エグジットメカニズム: ループ型状態空間モデルのための閾値制御された適応型エグジット状態選択メカニズムの導入により、中間深度におけるダウンストリーム性能の向上を実証している。

実験結果

合成推論タスク

Manoタスク（剰余演算）: Looped Mambaは、パラメータが一致した非ループ型のベースラインを一貫して上回る。最も困難な設定（ $L=24$ ）において、 $4 \otimes 6$ モデルは 99.35% の精度を達成したが、これは $4 \otimes 1$ ベースラインの 7.50% と比較して顕著である。等FLOPs条件（有効深度24）では、 $4 \otimes 6$ モデルは $24 \otimes 1$ の非ループ型ベースライン（99.35% vs. 37.45%）を大幅に上回っており、再帰的なパラメータ共有が、非共有の層を増やすよりも強力な帰納バイアスを提供することを示唆している。
$p$ -hop Induction（再帰的検索）: 浅い非ループ型Mambaモデルは苦戦した（例： $1 \otimes 1$ で約38-41%）。一方で、ループを追加することで性能が劇的に向上した。 $1 \otimes 12$ モデルはほぼ完全な精度（約99.89%）を達成しており、これは再帰的検索においては、個別のパラメータ数よりも再帰的な深度が重要であることを示している。

言語モデリングおよびダウンストリームベンチマーク

等パラメータ: 固定されたパラメータ予算の下で、再帰的な深度はMamba-2およびLLaMA-2バックボーンの両方において検証パープレキシティ（PPL）を改善した。
等FLOPs: 深い非ループ型モデル（ $24 \otimes 1$ ）は、ループ型モデルよりも検証パープレキシティにおいて優位性を維持した。しかし、ループ型モデルもダウンストリームベンチマークでは競争力のある性能を維持し、時にはより深い非ループ型ベースラインを上回った。
適応型エグジット:
- 140Mスケール: 適応型エグジットゲーティングはダウンストリーム性能を向上させた。140M Mamba-2モデルにおいて、ステージ2のエグジットゲートは平均スコア 38.87、平均選択エグジットステップ 3.00 を達成し、fixed-loop4 ベースライン（37.79）を上回った。
- 370Mスケール: 最良のダウンストリーム性能は、依然として fixed loop4 設定の付近に集中していた。エグジットゲート付きのバリアントは、ピーク時の固定ループのスコアを超えられなかった。これは、より大きなモデルは最終的な再帰ステップにより強く依存しているか、あるいはより優れたゲートの較正が必要であることを示唆している。
- トレードオフ: エグジットゲーティングは、選択された深度と性能の間の制御可能な境界を提供し、固定の最大深度を強制することなく、中間的な再帰状態を効果的に利用することを可能にする。

意義と主張

本論文は、再帰的な深度が、Mambaおよびハイブリッド言語モデルにおける有望な、パラメータ効率の高いスケーリング軸を提供することを主張している。主な知見は以下の通りである：

深度 vs. パラメータ: 多くの推論タスクは、追加の独立したパラメータよりも、より大きな計算深度から恩恵を受ける。Mamba層を再帰的に再利用することで、パラメータ数を増やすことなく、組成的な計算および再帰的な計算能力を大幅に向上させることができる。
ハイブリッドの生存可能性: ループされた深度の原理は、アテンション層がMamba層と共にループされるハイブリッドMamba-Transformerアーキテクチャにも効果的に適用される。
適応型選択: 現在のSSMの実装では全状態の実行が必要であるが、適応型エグジット状態選択により、最適な予測深度を選択することができ、中間深度におけるダウンストリーム性能を向上させることができる（特に140Mスケールにおいて）。

著者らは、実験が370Mパラメータまでのモデルに限定されていること、および実際の推論時の計算節約を実現するには、再帰ステップを完全にスキップするための状態処理メカニズムに関する将来の研究が必要であることを述べ、自らの主張に対して謙虚な姿勢を保っている。結果は、完全な特性評価ではなく、制御された経験的証拠として提示されている。