技術要約: T2MLR: 時間的中間層再帰（Temporal Middle-Layer Recurrence）を備えたTransformer

問題提起

著者らは、標準的なTransformerアーキテクチャにおける推論能力に関する根本的な限界を指摘している。Transformerはインコンテキストでの検索には優れているが、その自己回帰的なデコーディングメカニズムは情報のボトルネックを生み出す。各デコーディングステップにおいて、豊かで高次元な潜在表現は、次のステップへの入力として供給される前に、逐次、疎なトークンスペース（ワンホットベクトル）へと圧縮されてしまう。このプロセスは、中間的な抽象推論状態が時間を超えて持続することを妨げ、モデルに注意（Attention）や高度に圧縮された入力埋め込みを通じて間接的に必要な情報を再構築することを強いる。

既存の潜在推論へのアプローチは、推論を完全に連続空間内で行う（多くの場合、最終層で行う）、あるいは線形結合を通じて不確実性を伝播させることで、これを軽減しようと試みている。しかし、これらの手法は通常、再帰的な信号を入力埋め込みへとフィードバックするため、抽象的な推論が行われることが知られている中間層をバイパスしてしまう。あるいは、「ループ型」Transformerアーキテクチャは、全レイヤーを反復することで実効的な深さを増大させるが、これはループの数に応じて推論コストが線形に増加するため、デプロイメントにおいては非現実的である。

手法

本論文では、推論時の深さを増大させることなく、抽象的な中間表現がデコーディングステップを越えて持続し、進化することを可能にする新しいアーキテクチャ、T2MLR (Transformers with Temporal Middle-Layer Recurrence) を導入する。

アーキテクチャ設計

T2MLRは、標準的な$L$層のデコーダーのみのTransformerを修正し、前のトークンのより深い層から現在のトークンのより初期の層へと接続される、ゲート付き再帰パスを導入する。

1. 再帰境界 (Recurrence Boundaries): アーキテクチャは、$\ell_{start}$ と $\ell_{end}$ という2つのレイヤーインデックスによってパラメータ化され、「中間」ブロックのレイヤーを定義する。
2. 再帰キャッシュ ($R_t$): 各デコーディングステップ $t$ に対して、定数サイズのキャッシュベクトルが保持される。これは、前のステップにおけるレイヤー $\ell_{end}$ 後の表現を格納する。
3. 融合モジュール ($\Phi$): 現在のトークンの表現（$h^{(\ell_{start}-1)}_t$）を、前のトークンからの再帰キャッシュ（$R_{t-1}$）と融合させる処理が、現在のトークンのレイヤー $\ell_{start}$ で行われる。この融合は、学習可能なゲート機構によって制御される：
   $$\Phi(h, R) = h + \tanh(\gamma_{cur}) \sigma(f_{cur}([h, R])) \odot h + \tanh(\gamma_{rec}) \sigma(f_{rec}([h, R])) \odot W_{rec}R$$
   この設計により、モデルは現在のトークンのストリームと、時間的に蓄積された潜在状態のどちらに依存するかを学習することができる。
4. キャッシュ更新: レイヤー $\ell_{end}$ を通過した後、再帰キャッシュは、残差接続とRMSNormを用いて、新しい表現と以前のキャッシュを用いて更新される。

学習戦略: 時間的並列性 (Temporal Parallelism)

潜在推論アーキテクチャにおける主要な課題は、標準的なシーケンス並列学習（トークン並列）が、ステップ間で依存関係が存在する場合に崩れることである。これに対処するため、著者らはヤコビ固定点反復を用いた近似的な時間的並列学習スキームを提案している：

• 逐次的なトークン処理の代わりに、モデルは並列に $d_{forward}$ 回の精緻化反復を行う。
• 標準的なフォワードパスから派生した初期キャッシュから開始し、このキャッシュを現在のトークンの表現と融合させ、中間ブロックを精緻化し、キャッシュを更新する作業を反復する。
• 勾配は、Truncated Backpropagation Through Time (TBPTT) に類似した、切り捨てられた深さ ($d_{backward}$) を伴ってこれらの反復を通じて逆伝播される。
• このアプローチにより、モデルは高密度な教師強制（Teacher Forcing）とシーケンス並列性を備えた学習が可能となり、これは著者らが連続的な連鎖思考（Chain-of-Thought）のバリアントの中で唯一であると主張している能力である。

主な貢献

1. アーキテクチャの革新: T2MLRは、入力埋め込みや最終出力ではなく、Transformerの中間層に特化した再帰パスを導入する。これにより、標準的な自己回帰インターフェースを維持しつつ、潜在状態の持続を可能にする。
2. スケーラブルな学習: 著者らは、近似的な時間的並列性を用いてこれらの再帰的依存関係を学習する方法を示し、スケーラビリティを損なうことなく大規模なデータセットでの事前学習およびファインチューニングを可能にすることを実証した。
3. レトロフィット（後付け）能力: T2MLRはゼロから学習する必要がないことを本論文は示している。再帰パスは、既存の学習済みモデル（例：SmolLM2-1.7B）に挿入し、大幅な推論能力の向上を実現するためにファインチューニングすることができる。
4. 効率性: ループ型Transformerとは異なり、T2MLRは（定数サイズのキャッシュと推論時の単一フォワードパスにより）無視できる程度の推論オーバーヘッド（最大で1トークンあたり約8%）しか加えない。

実験結果

著者らは、合成ベンチマーク、自然言語事前学習、およびダウンストリームの推論タスクを用いてT2MLRを評価している。

• 合成推論 (S5-Retrieval): T2MLRは、状態追跡とインコンテキスト検索を同時に必要とするタスクにおいて、標準的なTransformerおよびLSTMの両方を大幅に上回る。Transformerは深さの制約により長いシーケンス長で失敗し、LSTMは固定された状態圧縮により検索において失敗するが、T2MLRは高い精度を維持する。
• 事前学習 (FineWeb-Edu): 135Mパラメータのモデルにおいて、T2MLRのバリアントは、ゼロショットNLPベンチマーク（ARC、HellaSwagなど）において、パラメータ数が一致するTransformerのベースラインと同等またはそれを上回る性能を一貫して示した。特に、中間層の再帰（例：レイヤー13–18の約20%をループする）は、全レイヤーの再帰よりも優れた性能を示すことが多い。
• ダウンストリーム推論: 数学（GSM-Aug）、マルチホップQA（HotPotQA）、変数割り当てなどのタスクにおいて、T2MLRはベースラインに対して一貫した利点を示している。ここでも、中間層の再帰バリアントは全レイヤーの再帰を上回っており、ターゲットを絞った再帰がグローバルなループよりも効果的であることを示唆している。
• スケーリング: パフォーマンスの向上は、361Mおよび1Bパラメータへのスケールアップ、および50Bトークンへの事前学習の拡張においても持続し、成長する。
• レトロフィット: 学習済みのSmolLM2-1.7B-InstructモデルにT2MLRを適用し、数学データでファインチューニングしたところ、GSM8Kの精度が35.8%から39.9%へ、MATH500の精度が12.8%から18.0%へと向上した。これは、フルな事前学習なしでの実用的な適用可能性を証明している。

意義と主張

本論文は、効果的な潜在推論を行うためには、すべてのレイヤーをループさせる（推論コストを増大させる）必要も、再帰を入力埋め込みに限定する（抽象的な計算をバイパスしてしまう）必要もないと断じている。代わりに、著者らは、ターゲットを絞った中間層の再帰が、有意な推論能力を解き放つのに十分であると主張している。

T2MLRの意義は、以下の点にある：

1. 推論の深さとコストの分離: ループ型アーキテクチャに伴う線形な推論ペナルティを課すことなく、時間ステップを越えて持続する時間的潜在パスを導入する。
2. 導入障壁の低下: 既存の学習済みモデルにレトロフィット可能であることを示すことで、高度な推論能力を、ゼロからの事前学習コストのわずかな割合で現在のモデルに追加できることを示している。
3. 帰納バイアスの検証: 結果は、Transformerにおける再帰の「スイートスポット」は、入力や出力のインターフェースではなく、抽象的な計算が行われる中間層であることを示唆している。

著者らは、Jacobi近似ステップにより、T2MLRが学習時のオーバーヘッド（再帰ブロックのサイズに応じて2〜4倍のウォークタイム）を伴うことを認めている。彼らはこれをトレードオフとして捉えている。すなわち、T2MLRは学習計算を犠牲にすることで、推論コストが低く、高性能な推論タスクのための推論アーキテクチャを獲得しているのである。