系列全体を通る再帰経路

Recurrent Looped Transformer(RLT)と呼ばれる新しい言語モデルのアーキテクチャー案は、プロンプトから応答のすべてのトークンに至るまで、1つの連続した計算を引き継ぐ方法を示している。この設計は因果的エンコーダーと再帰型デコーダーを組み合わせ、生成が進む間、デコーダーの最終隠れ状態と、層ごとのスライディングウィンドウ注意機構のキャッシュを保持する。

プロジェクトの著者らは、このシステムを、時間方向の深さを拡張できる潜在空間での推論を実現する手段として提示している。この表現は、各トークンに対して無制限の計算を行うという意味ではない。そうではなく、新しいトークンが加わるたびに系列全体を通る再帰経路が長くなる一方、個々のトークンに割り当てるデコーダーの計算量は一定に保たれる。プロジェクトが示す具体的な構成では、エンコーダーとデコーダーはそれぞれ48層を持つ。系列がtトークンに達すると、その再帰経路は48×t個のデコーダーブロックを通過したことになる。

エンコーダーは、既知のトークンから全体を対象とするキー・バリューメモリーを作る。次にデコーダーが、そのメモリーを参照し、各スライディングウィンドウ注意層で直近のキャッシュを保持しながら、トークンを順番に処理する。直前のデコーダー出力は次の更新に入力される。この構造は、プロンプト処理と応答生成の間にある構造的な断絶をなくし、両方を同じ状態遷移に参加させることを意図している。

学習と実行の要件

この提案は、アーキテクチャーを実用化するための実行手法も概説している。既知のプロンプトトークンは因果関係を保ったバッチで処理でき、独立したエンコーダーの処理は並列実行できる。再帰処理の中核の周囲には、重みとメモリーの再利用、系列のバッチ化、活性化チェックポイントが配置される。ただし、各更新が次の更新に必要な状態を構築するため、デコーダーの更新は引き続きトークンの順番どおりに行う必要がある。

学習の厳密性を保つには、さらに制約が加わる。時間方向の完全な誤差逆伝播では、再帰出力、デコーダーのキャッシュ、エンコーダーメモリーを通る勾配経路を保持しなければならず、これらの経路のいずれかを切り離すと勾配が変わる。パラメーター更新後は、以前に計算したキャッシュは現在の方策を厳密に表すものではなくなる。そのため著者らは、現在の方策による再実行では、プロンプトの状態とデコーダーの注意キャッシュを含む履歴全体を再構築しなければならず、一方で行動確率は、実際にトークンを生成したサンプラーに対応していなければならないとしている。

今回の公開内容はアーキテクチャーの提案であり、この手法がすでに推論能力の向上やハードウェア費用の削減を実現したことを示す証拠ではない。著者らは、実際の推論性能の向上、ハードウェア効率、強化学習のスケーリングについては今後の検証に委ねると明示している。現時点での貢献は、新しいトークンごとに実行する論理ブロック数を増やすことなく、生成される系列に合わせてTransformerの計算履歴を伸ばすための、具体的かつ検証可能な設計を示した点にある。