技術要約: RecurrentGPT

問題提起

トランスフォーマー言語モデルのスケーリングにおいて、**表現力（expressivity）とメモリ効率（memory efficiency）**の間には根本的な緊張関係が存在する。標準的なアーキテクチャは、深さとパラメータを硬直的に結合させている。すなわち、層を追加するたびに一意の重みセットが導入されるため、固定されたハードウェア予算の下ではメモリのボトルネックが生じ、実効的な深さが制限される。逆に、標準的な重み共有戦略（例：すべての層を深さに対して共有する手法）は、進化する表現に対して一様な変換を強制してしまう。この静的なアプローチは、深さが価値を生む根源である「機能的な多様性」を崩壊させる。なぜなら、第1回目の再帰における隠れ状態と第8回目のそれは根本的に異なるにもかかわらず、両者が同じ静的な変換を強制されるからである。

本論文は、推論システムの能力はパラメータの広さではなく、反復の深さに宿るものであると断じている（チューリング [2] および現代のテストタイム・コンピューティング [3–5] を引用）。課題は、すべての層に固有の重みを持つことによるメモリコストを回避しつつ、かつ単純な重み共有によって引き起こされる表現の崩壊を避けながら、入力表現の深く反復的な洗練を実現することである。

メソドロジー: RecurrentGPT アーキテクチャ

著者らは、固定されたプレリュード（前奏）およびコーダ（終奏）ブロックでコアを挟み込むことで、モデルの深さとパラメータ数を切り離した再帰的深さトランスフォーマーである RecurrentGPT を提案する。

1. アーキテクチャの構成

モデルは $n_{pre} + n_{rec} \times R + n_{coda}$ の構成に従う：

• プレリュード ($n_{pre}$ ブロック): 入力から安定したコンテキスト依存の条件付け信号 ($h^{(pre)}$) を生成するために一度だけ適用される。
• 共有コア ($n_{rec}$ ブロック): $R$ 回再帰的に適用されるトランスフォーマーブロックの集合。標準的な重み共有とは異なり、このコアは一様に適用されるのではなく、ステップごとに変調される。
• コーダ ($n_{coda}$ ブロック): 最終的に洗練された状態を処理してロジットを生成する。

総ユニーク・パラメータ数は、再帰回数 $R$ に依存せず、$\Theta((n_{pre} + n_{rec} + n_{coda}) \cdot d^2)$ となる。

2. 適応型ゲーティングと再帰的射影

表現の崩壊を防ぎ、単一の重みテンソルが複数の特化した層として振る舞うことを可能にするために、RecurrentGPT は要素ごとの学習済みゲートと再帰的射影を導入している。

各再帰ステップ $r$ において：

1. 入力準備: 前ステップの隠れ状態 $h^{(r-1)}$ が、固定されたプレリュード出力 $h^{(pre)}$ および注入された確率的ノイズ $\epsilon_x$ と共に射影される。
   $$\tilde{h}^{(r)} = W_{proj} [h^{(r-1)} + \epsilon_x, h^{(pre)}]$$
2. 共有変換: 共有ブロック ($B_{shared}$) がこの入力を処理し、提案値 $o^{(r)}$ を生成する。
   $$o^{(r)} = B_{shared}(\tilde{h}^{(r)})$$
3. 要素ごとのゲーティング: 学習されたゲート $g^{(r)}$ が、提案値が現在の状態をどの程度置き換えるかを決定する。ゲートは、正規化された現在の状態、正規化されたプレリュード、および注入されたノイズ $\epsilon_g$ に基づいて条件付けられる。
   $$g^{(r)} = \sigma(f_{gate}([LN(h^{(r-1)}), LN(h^{(pre)})]/\tau + \epsilon_g))$$
   $$h^{(r)} = g^{(r)} \odot h^{(r-1)} + (1 - g^{(r)}) \odot o^{(r)}$$

主要な設計上の選択:

• 初期化: ゲートのバイアスは $+4$ に初期化され、これにより学習開始時には $g^{(r)} \approx 0.98$ となる。これにより、最初は残差ストリームがほとんど変化せずに通過することを保証し、モデルがどの要素を段階的に上書きすべきかを学習できるようにする。
• ノイズ: 加法的ノイズ ($\epsilon_x$) およびゲートノイズ ($\epsilon_g$) は、脆弱な完全一致パターンやゲートの崩壊を防ぐ。
• 深さサンプリング: 学習中、再帰の深さ $r$ は $\{1, \dots, R\}$ から一様にサンプリングされる。これにより、補助的な損失なしに、推論時の**創発的な早期終了（early exit）**を可能にするための各終了ポイントが暗黙的に訓練される。

主な貢献

1. アーキテクチャ: RecurrentGPT は、要素ごとのゲーティングと確率的摂動を利用して表現の崩壊を防ぎ、共有コアが再帰ステップを通じて明確に異なる挙動をとることを可能にする。
2. 統一された評価: モデルは、isoFLOPS（計算量を一致させ、パラメータを削減）および isoPARAMS（パラメータ数を一致させ、計算量を増加）の両方のレジームの下で評価される。
   • isoFLOPS: RecurrentGPT は、9つのすべての「予算別スケールセル」において、Mixture-of-Recursions (MoR) およびヘビーテイル深さサンプリングを上回る。小規模スケールでは、64%少ないパラメータで高密度な GPT-2 Small ベースラインと同等の性能を示す。中・大規模スケールでは、トークン予算が増えるにつれて高密度モデルの品質に近づく。
   • isoPARAMS: 深い再帰は、一貫した検証損失の向上をもたらす（例：パラメータ数が一致した場合、非再帰型と比較して 2.76 対 2.84）。
3. 創発的な早期終了: 補助的な損失なしに、モデルは自然に推論時の早期終了をサポートし、再帰ステップの半分のみを実行した場合でも、精度の約92%を維持する。

実験結果

• パフォーマンス: isoFLOPS 制約下において、3レイヤーの RecurrentGPT は 12レイヤーの GPT-2 Small ベースラインの精度に匹敵する。中規模スケール（24レイヤー相当）では、高密度ベースラインの 2.84 に対し、2.89 の検証損失を達成しつつ、354M ではなく 124M のパラメータのみを使用する。
• 効率性: isoFLOPS 下では、モデルは高密度ベースラインのユニーク・パラメータの 36–37% しか保持しない。
• ダウンストリームタスク: 大規模ベンチマークにおいて、isoPARAMS バリアントは 9 つのタスクのうち 8 つで高密度ベースラインを上回り、平均 +2.10 ポイントの改善を示した。isoFLOPS バリアントは、わずか 37% のパラメータで高密度ベースラインの平均性能（42.08 対 42.05）に匹敵する。
• レイテンシとメモリ: torch.compile 使用時、RecurrentGPT-Large は生成レイテンシの増加をわずか 10% に抑えつつ、ピーク GPU メモリを 59%、パラメータを 63% 削減する。
• メカニズム的洞察: 以下の二段階の計算戦略が明らかになった：
  • 初期ステップ（書き込み重視）: ゲートがより開いており、モデルは急速な再構築を行い、最初の2ステップで損失減少の約77%を解決する。
  • 後半ステップ（コピー重視）: ゲートが締まり、モデルはターゲットを絞った修正を行う。
  • トークン・ルーティング: モデルは明示的なルーティング信号なしに、困難なトークン（初期損失が高いトークン）に対して計算資源を暗黙的に多く割り当てる。

意義と主張

本論文は、適応的な深さの再利用が、パラメータを品質へと交換するための原理的な戦略であることを主張している。再帰を現在の隠れ状態、固定されたプレリュード、およびノイズに基づいて条件付けることで、モデルは固有の重みなしに機能的な多様性を達成する。

著者らは、RecurrentGPT が、単一の重みテンソルを複数の特化した層として振る舞わせることで、メモリ効率と表現力の間の緊張関係を解決することを強調している。結果は以下のことを示している：

1. パラメータ効率: 最小限のレイテンシオーバーヘッドで、大幅なパラメータ削減（最大 63%）が可能である。
2. 計算量と品質のトレードオフ: このアーキテクチャは、単一のチェックポイントが複数の動作ポイントを提供することで、推論速度と品質の間の連続的なトレードオフを可能にする。
3. スケーリング則: このアプローチは、損失対パラメータ数、および損失対トレーニング FLOPs のスケーリング則を改善し、小・中・大規模において一貫した利点を示す。

結論として、現在の実装では推論時の再帰深さは固定されているものの、学習されたゲーティング機構は、動的なトークンごとの停止（halting）や知識蒸留のための強固な基盤を提供している。