Recurrent Action Transformer with Memory

この論文は、部分観測環境における長期記憶の保持を可能にする再帰的メモリ機構を組み込んだ新しいトランスフォーマーベースのアーキテクチャ「RATE」を提案し、記憶依存タスクおよび標準的なベンチマークにおける優れた性能を実証しています。

要約

この論文は、人工知能（AI）が「長い間、重要なことを忘れずに」判断を下すための新しい仕組み「RATE（レイト）」という名前のお話です。

まるで**「記憶力に優れた探偵」**のような存在が、AI の世界に登場したのです。

🕵️‍♂️ 従来の AI の悩み：「忘れっぽい探偵」

まず、これまでの AI（特に「トランスフォーマー」と呼ばれるタイプ）は、非常に優秀な「直感の探偵」でした。
しかし、ある弱点がありました。それは**「直前の出来事しか覚えていない」**ことです。

• 例え話：
  探偵が事件現場（ゲームの画面）を見て、「さっき赤いボタンを押したから、今は青い箱を開けよう」と判断します。
  でも、もしその「赤いボタン」が10 分前に押されたもので、その間の出来事が大量に挟まっていたら？
  従来の AI は、その間の情報が多すぎて「あ、赤いボタンって何だっけ？」と忘れちゃいます。
  これを専門用語では「文脈の長さの制限」と呼びますが、要は**「記憶容量が足りなくて、昔のヒントを忘れる」**のです。

🧠 新しい解決策：「RATE（レイト）」の登場

そこで登場したのが、この論文で提案された**RATE（Recurrent Action Transformer with Memory）です。
これは、「メモ帳と、賢いメモの整理係」**を AI に搭載したようなものです。

RATE は、長い物語（ゲームの履歴）をすべて一度に読もうとするのではなく、**「章ごと」に分けて読み進めます。そして、各章が終わるたびに、「次の章で必要な重要な情報だけ」**をメモ帳に書き留めて、次の章へ持ち越します。

3 つの魔法の仕組み

RATE がなぜこんなに優秀なのか、3 つの魔法の道具で説明します。

1. 📝 魔法のメモ帳（Memory Embeddings）

   • 従来の AI は、過去の情報をすべて画面に並べて見ていましたが、RATE は「重要な情報だけ」を小さなメモに書き留めます。
   • 例え： 長い旅行で、すべての景色を写真に撮るのではなく、「目的地への道しるべ」だけをメモ帳に書き込んで、次の街へ持ち運ぶイメージです。

2. 🔄 情報の受け渡し（Recurrent Caching）

   • 前の章で計算した「隠れた知識」を、次の章でそのまま使えるようにします。
   • 例え： 前の章で「鍵の場所」を調べた結果を、次の章で「鍵を探す」ためにそのまま引き継ぐ感じです。最初から全部やり直す必要がありません。

3. 🚪 情報の選別ゲート（Memory Retention Valve / MRV）

   • これが RATE の一番のすごいところです。メモ帳に新しい情報を書き込むとき、**「本当に必要な情報か？」**を厳しくチェックする「選別係（バルブ）」が働きます。
   • 例え： 毎日届く手紙（新しい情報）の中から、**「本当に重要な手紙」**だけを選んで、古い重要な手紙（過去の記憶）を捨てないように守る係です。
   • これがないと、新しい情報が入ってくるたびに、昔の重要なヒントが上書きされて消えてしまいます。RATE はこのゲートのおかげで、「100 歩前（あるいはもっと前）のヒント」を忘れずに保持し続けることができます。

🎮 実際のテスト結果：どんなに長くても勝つ！

研究者たちは、この RATE をいろいろなゲームで試しました。

• 迷路ゲーム（T-Maze）：
  入り口で「左に行け」というヒントをもらい、迷路の奥まで進んでからそのヒントを思い出して左に曲がる必要があります。

  • 従来の AI： 迷路が長くなるとヒントを忘れ、右に行ってしまう（成功率 50%）。
  • RATE： 迷路が 100 倍長くても、入り口のヒントを忘れないで正解します（成功率 100% 近く）。

• 色を覚えるゲーム（ViZDoom）：
  最初に見た「赤い柱」の色を覚えて、後で赤いアイテムだけを集める必要があります。

  • RATE： 柱が見えなくなってからも、色を覚えて正しくアイテムを集め続けます。

• 普通のゲーム（Atari やロボット制御）：
  「記憶力」があまり必要ない普通のゲームでも、RATE は他の AI と同じくらい、あるいはそれ以上によく動きました。つまり、**「記憶がなくても大丈夫な場面でも邪魔をせず、記憶が必要な場面では最強」**という万能選手なのです。

🌟 まとめ

この論文が伝えたかったことは、**「AI に『忘れずに考える力』を持たせれば、複雑で長い問題も解決できる」**ということです。

RATE は、**「過去の重要なヒントを、必要な時まで安全に保管し、必要な時に引き出せる」という仕組みを作りました。これにより、AI はまるで「経験豊富なベテラン探偵」**のように、長い時間をかけても正しい判断を下せるようになったのです。

これは、ロボットが複雑な作業を覚えたり、長期的な計画を立てたりする未来にとって、非常に大きな一歩だと言えます。

技術概要

論文「Recurrent Action Transformer with Memory (RATE)」の技術的サマリー

本論文は、ICLR 2026 にて発表された「Recurrent Action Transformer with Memory (RATE)」という、オフライン強化学習（Offline RL）における新しいトランスフォーマーアーキテクチャを提案する研究です。

1. 背景と課題 (Problem)

強化学習（RL）において、トランスフォーマーはエージェントの軌跡をシーケンスとして扱い、方策学習をシーケンスモデリング問題として再構築することで、オフライン RL において高い性能を発揮しています。しかし、以下の根本的な課題が存在します。

• 部分観測マルコフ決定過程 (POMDP) における記憶の欠如: 現実の多くの環境は部分観測的であり、過去の情報を保持することが意思決定に不可欠です。
• 自己注意 (Self-Attention) の制約: 標準的なトランスフォーマーは、自己注意の計算コストが二次関数的 ($O(N^2)$) であるため、コンテキスト長（入力シーケンスの長さ）に制限があります。
• 長期依存関係の破綻: 報酬が疎（スパース）で、重要な手がかりが非常に過去のステップに存在するタスク（例：迷路の入り口で見た色を覚えて、出口で正しい方向を選ぶ）において、標準的なトランスフォーマー（Decision Transformer など）はコンテキストウィンドウを超えると情報を失い、性能が急激に低下します。
• 既存の解決策の限界: コンテキストウィンドウの拡張技術はトレーニングの不安定性やタスク固有のスパース性への依存などの問題を抱えており、NLP 以外の分野での汎用性が限定的です。

2. 提案手法 (Methodology)

著者らは、Recurrent Action Transformer with Memory (RATE) を提案しました。これは、トランスフォーマーに再帰的な記憶メカニズムを組み合わせたアーキテクチャです。RATE は以下の 3 つの相補的なメカニズムを統合しています。

2.1 アーキテクチャの概要

RATE は、長い軌跡を $N$ 個のセグメント（長さ $K$）に分割して処理します。各セグメント $S_n$ に対して、以下の処理を行います。

1. メモリ埋め込み (Memory Embeddings):
   各セグメントの前後に、共有されるメモリトークン $M_n$ を付加します。

   • プレフィックス（読み取り）: セグメント内のトークンが過去のメモリにアクセス（注意）できるようにします。
   • サフィックス（書き込み）: 現在のセグメントの処理結果がメモリを更新できるようにします。
     これにより、セグメント間での情報の受け渡しが可能になります。

2. 隠れ状態のキャッシュ (Recurrent Caching):
   Transformer-XL の仕組みを応用し、直前のセグメントで計算された隠れ状態をキャッシュし、次のセグメントのキー・バリューコンテキストとして再利用します。これにより、セグメント境界を超えた情報の流れを確保します。

3. メモリ保持弁 (Memory Retention Valve, MRV):
   RATE の中核となる新規メカニズムです。単にメモリを渡すだけでは、重要な情報が上書きされたり、ノイズが蓄積したりするリスクがあります。MRV は、クロス・アテンションを用いて、新しいメモリ状態 ($M_{n+1}$) を更新する際に、過去のメモリ状態 ($M_n$) を参照し、「何を保持し、何を上書きするか」を制御します。

   • 具体的には、$M_n$ を Query、$M_{n+1}$ を Key/Value としてアテンションを計算し、重要な情報を保持したまま更新を行うフィルタリング機構です。
   • 理論的解析により、MRV がメモリ情報の損失に下限を設け、カタストロフィック・フォージング（重要な情報の破壊的消失）を防ぐことが証明されています。

2.2 アルゴリズムの流れ

1. 軌跡をセグメントに分割。
2. 各セグメントにメモリトークンを付加し、トランスフォーマーに入力。
3. 出力された新しいメモリ状態を MRV を通してフィルタリングし、次のセグメントへ渡す。
4. このプロセスを再帰的に繰り返すことで、非常に長い軌跡（コンテキストウィンドウを超えた長さ）を処理可能にします。

3. 主要な貢献 (Key Contributions)

1. RATE の提案: オフライン RL 向けに、メモリ埋め込み、隠れ状態キャッシュ、MRV を組み合わせた新しいトランスフォーマーアーキテクチャを提案。
2. 記憶依存タスクでの SOTA 性能: ViZDoom-Two-Colors, T-Maze, Minigrid-Memory, Memory Maze, POPGym などの記憶が不可欠なタスクにおいて、強力なベースライン（Decision Transformer, RMT, Transformer-XL, LSTM 系など）を凌駕する性能を示しました。
3. 汎用性の証明: 記憶タスクだけでなく、Atari や MuJoCo の標準的な MDP ベンチマークにおいても、既存のオフライン RL 手法と同等かそれ以上の性能を発揮し、幅広いタスクタイプに対応できる汎用モデルであることを実証しました。
4. 理論的保証: MRV によるメモリ保持の数学的保証（定理 1）を提供し、長期にわたる情報保持の安定性を理論的に裏付けました。

4. 実験結果 (Results)

実験は、記憶負荷の高い環境と標準的な制御タスクの両方で行われました。

• 記憶依存タスク:
  • T-Maze: 訓練データ（最大 900 ステップ）を超えた 9600 ステップの推論において、RATE は 100% の成功率を維持しました。一方、Decision Transformer (DT) はコンテキストを超えると 50% 程度に低下し、RNN 系モデルも長期化に伴い性能が劣化しました。
  • ViZDoom-Two-Colors: 45 ステップで消える柱の色を記憶し、その後の 2000 ステップ以上で正しい色のアイテムを集めるタスクにおいて、RATE は高いリターンとバランスの取れた行動を示しました。
  • POPGym: 48 種類の部分観測タスクのスイートにおいて、RATE は「記憶パズル」タスクで他のすべてのモデルを大きく上回り、正の平均スコアを達成しました（他のモデルは負のスコア）。
• 標準ベンチマーク (Atari, MuJoCo):
  • Atari ゲームや MuJoCo 制御タスクにおいても、RATE は Decision Transformer や CQL、Mamba 系モデルと同等かそれ以上の性能を達成し、記憶メカニズムが単純なタスクの性能を阻害しないことを示しました。
• アブレーション研究:
  • MRV を除去すると、長期タスクでの性能が急激に低下することを確認。
  • MRV のクロス・アテンション構成（MRV-CA-2）が最も効果的であることを示しました。
  • 隠れ状態のキャッシュとメモリ埋め込みの役割分担（密な報酬ではキャッシュが重要、疎な報酬ではメモリ埋め込みが重要）を明らかにしました。

5. 意義と結論 (Significance)

本論文の RATE は、オフライン強化学習における「長期の意思決定」と「部分観測性への対応」という二つの大きな課題を解決する統合的なアーキテクチャとして確立されました。

• スケーラビリティ: 再帰的なセグメント処理と MRV により、トランスフォーマーのコンテキスト制限を超えて、数千〜数万ステップにわたる長期依存関係を効率的に学習・推論できます。
• 実用性: 複雑な記憶タスクだけでなく、一般的な制御タスクでも高性能を発揮するため、単一のモデルで多様な RL 課題を解決できる「ユニバーサルなオフライン RL モデル」としてのポテンシャルを示しました。
• 理論的深み: 単なる実験的な工夫ではなく、メモリ保持の理論的保証を含めて提案されており、記憶メカニズムの設計指針を提供しています。

結論として、RATE は、トランスフォーマーの注意機構と再帰的な記憶制御を融合させることで、従来のモデルが苦手としていた「長期・疎な報酬・部分観測」環境における意思決定を可能にする画期的なアプローチです。