yifanzhang-pro/recurrent-looped-tranformer

解決する課題

Recurrent Looped Transformer (RLT) は、標準的な Transformer が抱える長さの汎化（学習時に見たことのない長いシーケンスでタスクを実行する能力）の制限に対処します。特に、精度を維持するためにシーケンス長に応じて計算パスを拡張する必要があるアルゴリズムタスク（パリティや状態追跡など）を対象としています。

仕組み

RLT は Transformer アーキテクチャを修正し、トークン間での再帰を導入します：

• 再帰的フィードバック (Recurrent Feedback)： 前のトークンの最終デコーダー状態が、ゲート付きマージ (gated merge) を介して現在のトークンのデコーダー入力にフィードバックされます。これにより、シーケンス長に対して線形に増加する計算パスが作成されます。
• ハイブリッドメモリ (Hybrid Memory)： グローバルなキー・バリュー (KV) メモリには並列因果エンコーダーを使用し、ローカルメモリにはスライディングウィンドウアテンション (SWA) を備えたデコーダーを使用します。
• 状態の連続性 (State Continuity)： 再帰的な状態は、プロンプトと応答の境界をシームレスにまたいで継続します。
• フィードバックのバリエーション (Feedback Variants)： 本プロジェクトでは異なるフィードバック間隔 ($B$) を調査しています：RLT-1 はトークンごとに更新、RLT-2 は $B$ トークンのチャンクごとに更新、RLT-0 はフィードバックなしです。

対象読者

シーケンスモデリング、長さの汎化、LLM におけるアルゴリズム推論に取り組んでおり、標準的な Transformer の固定深度の制限を超えたい研究者や開発者。

ハイライト

• 強力な長さの汎化： 学習データよりも大幅に長いシーケンス（例：パリティタスクで 256 ビット）において、パリティおよび S5 スワップタスクで 100% の精度を達成。
• 再帰的パス： 計算パスは $t$ トークンに対して $tL_D$ 個のデコーダーブロックを通過し、トークンあたりのコストを一定に保ちます。
• 柔軟な深度割り当て： エンコーダーとデコーダー間で層を分割することが、異なるアルゴリズムタスクのパフォーマンスにどのように影響するかを実験で示しています。
• 効率のトレードオフ： RLT-2（チャンク化フィードバック）は、一部のタスクで高い精度を維持しつつ、RLT-1 と比較して学習を高速化（最大 2.27 倍）します。