Recurrent Looped Transformer：トークン間の再帰で状態追跡を長い入力へ外挿する

入力が長くなっても、途中の変化を正しく積み重ねられるか。RLTが扱うのは、この状態追跡の問題だ。たとえば五つの要素の並びに交換操作を順番に適用し、最後の並びを答える課題では、各操作の結果が次の操作の出発点になる。入力を参照できるだけでなく、変化を順序どおりに反映する必要がある。

通常の固定層数のTransformerでは、各位置に適用する計算の深さは、入力の長さに応じて増えない。著者はこれを、学習時より長い系列で状態追跡が難しくなる要因として捉える。ここでいう長さ外挿とは、短い系列で学習した処理を、より長い同種の系列にも適用することだ。異なる種類の仕事を解けるかという問いとは分けて読む必要がある。

再帰とは、前の処理結果を次の処理へ戻して使うことだ。関連研究にも、時間方向のフィードバック、前のトークンの潜在状態の注入、ブロック単位の再帰、深さ方向での重み共有がある。RLTの中心は、入力を読む部分と、状態を引き継いで更新する部分を分ける点にある。

入力側の因果エンコーダは、各位置から未来を見ない制約の下で表現を作る。再帰デコーダは、その表現に直前のトークンの最終状態を加えて処理する。五つの要素の交換課題なら、直前までの並びを表す情報を、次の交換を処理するために渡す設計と考えられる。ただし、内部状態が人間に読める並びそのものになると保証されているわけではない。

この分離によって、各トークンが通るブロック数を固定したまま、状態更新をつなぐ計算経路は系列の長さに応じて伸びる。交換操作が増えれば、それに沿って前の結果を受け取る更新も増える。

ただし、「トークンあたり固定」を一定の実行時間と読むと誤る。RLTは現在位置までのエンコーダメモリにも注意を向けるため、その注意の演算量は文脈長とともに増える。また、前の最終状態を受け取る部分には順番待ちが生じる。長さ外挿を助ける可能性のある構造が、既知の入力を一度に処理する並列性には制約を加える。この兼ね合いが評価のもう一つの軸だ。

処理は、因果エンコーダが既知の入力を位置ごとに並列処理するところから始まる。ここで各トークンの表現と、デコーダが参照できる接頭辞のメモリを作る。接頭辞は、現在位置までの入力を指す。

次に、直前のトークンの最終デコーダ状態を正規化し、現在のエンコーダ表現に統合する。統合にはゲートを使う。ゲートは、現在の入力と過去の状態から、戻す情報の各成分をどの程度通すかを決める仕組みだ。主比較では、フィードバック全体の強さを定める係数を0.1にしている。

デコーダは、近い過去への注意、現在位置までのエンコーダメモリへの注意、変換処理を順に適用する。保持する状態には最終出力だけでなく、各層の近い過去の注意用履歴も含まれる。最終出力から次のトークンを予測し、その状態を次の位置へ渡す。プロンプトを読むときも応答を生成するときも、同じ状態更新を使う。

毎トークン状態を更新する構成がRLT-1だ。対してRLT-2は、複数トークンをまとめたチャンクの境界で更新する。たとえば4トークン単位なら、チャンク内では直前の境界状態を共有し、最後の出力だけを次のチャンクへ渡す。未来を見ない制約を保ちつつ、既知のトークンをチャンク内で並列処理できる。

これは更新のつなぎ方を変える仕組みであり、評価するデコーダブロックの総数を減らすものではない。また、答えを自己回帰的に生成するときは、いずれも1トークンずつ進む。既知入力の並列処理と生成速度は分けて考える必要がある。

著者はフィードバックを除いたRLT-0も比較した。こうして構成の一部を取り除いて働きを調べる実験をアブレーションと呼ぶ。ただしRLT-0では、初期状態、正規化、ゲート、状態の射影も一緒に除去する。精度が下がれば再帰統合の重要性を示す材料になるが、フィードバックだけの効果を完全に切り分ける比較ではない。

主比較は、8層のdecoder-only Transformerと、合計8層をエンコーダとデコーダに配分したRLT-1だ。4+4、5+3、6+2、7+1、8+0の五つを調べる。表記の前半がエンコーダ、後半がデコーダの層数で、8+0でもゲート付きの再帰統合は残る。共通するモデル幅は512、注意のヘッド数は4で、三つの初期化を使う。

課題はパリティ、加算、標準型とスワップ型のS₅追跡、括弧なしと括弧付きのmod-5の六つだ。mod-5は、式を計算した結果を5で割った余りを答える課題である。訓練範囲はパリティが3〜40ビット、S₅が32操作、加算が1〜8桁。mod-5は括弧なしが奇数長3〜39、括弧付きが長さ3〜40だ。

パリティ、加算、S₅は2,000ステップ、主要なmod-5比較は5,000ステップ学習する。したがって、mod-5の主要結果を他の課題と同じ学習予算で得たものとして扱うことはできない。

形式課題は各長さで共通の1,024例を使い、最後の答えまたは状態が正しいかを採点する。加算は各桁数で共通の256組を使い、正しい過去の答えを与えるteacher forcingの条件で、答えトークンの正解率を測る。これは答え全体を自力で生成できた割合ではない。また、S₅の最終状態が正しいことから、途中のすべての状態も正しかったとは言えない。

評価するモデルは、学習範囲内の検証損失が最小になったチェックポイントから選ぶ。テスト成績は選択に使わず、同値なら最も早い時点を採用する。報告の「平均±標本標準偏差」は、固定した共通データに対する三つの初期化の差を表す。以下の実験値はいずれも著者報告であり、このばらつきは実務で遭遇するデータ変動全体を表すものではない。

256ビットのパリティでは、RLT-1の5+3と7+1がともに100.00±0.00%だった。最大40ビットでの学習から、三つの初期化すべてで正答した。一方、Transformerは50.07±1.63%で、偶然水準付近にとどまる。ただし、成功はすべての層配分に共通するわけではない。

スワップ型S₅では、32操作で学習した4+4が、8倍長の256操作で97.30±2.76%を得た。Transformerは0.85±0.30%だった。さらに512操作へ伸ばすと、4+4は55.70±25.78%となる。比較対象への優位は残るが、精度も初期化間の安定性も低下する。この課題の256操作では、デコーダ層が多いほど改善し、1層と0層は偶然水準付近だった。

5,000ステップのmod-5比較でも、学習範囲外の中間的な長さで差が出た。括弧なしの長さ63では6+2が93.36±5.69%、Transformerが33.20±2.33%。括弧付きの長さ64では5+3が67.97±2.91%、Transformerが46.71±1.21%だった。ただし、さらに長い式では精度が下がり、初期化差も大きくなる。

4+4の64ビット・パリティでは、毎トークン更新するRLT-1が100.00±0.00%、4トークンごとに更新するRLT-2が98.99±1.66%だった。フィードバックを除くRLT-0は50.23±3.80%、Transformerは48.47±1.18%である。更新間隔を広げても、この条件では高精度を保った。

同じ4+4でも、64操作のスワップ型S₅では結果が大きく違う。毎トークン更新は100.00±0.00%だったが、4トークン単位では19.60±6.54%まで下がった。平均の差は80.40パーセントポイントで、相対的な減少率を表す数字ではない。RLT-0とTransformerは、120通りを一様に予測する水準付近にとどまる。

著者は、フィードバック除去後には全分割で、長さ64以降のパリティとスワップ型S₅が偶然水準付近になるとも報告する。再帰統合が成功に関わるという解釈を支える一方、チャンク化の結果は、必要な更新頻度を課題横断で決められないことを示す。

加算の32桁では、RLT-1の各分割の答えトークン正解率が14.89〜15.74%、Transformerが16.84±1.45%だった。標準S₅の512操作でも、RLT-1は0.72〜1.43%、Transformerは0.81±0.06%である。スワップ型での成功を、標準型へそのまま移すことはできない。

加算ではフィードバック係数を0.1から0.03または0.01へ下げる追加実験も行った。63比較のうち37件は改善し、26件は悪化した。差は−3.86〜+5.60パーセントポイントで、比較する両側の標本標準偏差を超えた差は3件だった。どの分割にも、全試験幅を改善する設定はなかった。再帰を入れる、あるいは戻す情報を弱めるだけで、長い加算を解決できたとは言えない。

学習過程の安定性にも注意が必要だ。パリティ5+3の一つの初期化では、検証正解率がステップ500の100%から600の48.05%へ下がり、700で100%へ戻った。最良の検証時点を選んだ最終的な報告値だけでは、学習途中の後退は見えにくい。

チャンク化の速度測定は限定的だ。4+4のmod-5二課題を、一つのシードで共有クラスタのCPU上で測ったところ、4トークン単位のRLT-2はRLT-1の2.27倍の訓練ステップ速度だった。RLT-0は4.17〜4.30倍である。これはステップ1,001〜2,000の平均で、検証、チェックポイント保存、ログを除く。GPU性能、既知入力を読み込むprefillの遅延、生成時間、金銭的コストの測定ではない。

層数を揃えても、モデル規模と計算量は一致しない。RLT-1は2,610万〜2,873万パラメータ、Transformerは2,531万で、注意構造も異なる。RLT-0は各分割で787,968パラメータ少ない。このため、主比較の差を再帰だけに、除去実験の差をフィードバックだけに帰すには留保がある。

評価の中心は小規模な教師あり合成課題であり、自然言語の大規模学習、実務課題、強化学習の性能は確立していない。大きなチャンクで学習して後から小さくする案も未検証で、実験では各サイズを最初から別々に学習している。次に見るべき証拠は、対象課題での長さ外挿、更新頻度による精度差、そして実際の環境での処理時間を併せた評価だ。