技術要約：再帰を用いない再帰型ネットワークの事前学習

問題提起

再帰型ニューラルネットワーク（RNN）の学習には、本質的に、初期の入力が将来の出力にどのように影響するかを決定する、長いシーケンスにわたる**クレジット割り当て問題（credit assignment problem）を解決する必要があります。標準的な手法であるBackpropagation Through Time (BPTT)**は、RNNを時間軸方向に展開し、計算グラフ全体を通じて勾配を逆方向に伝播させます。このアプローチには、以下の2つの決定的な制限があります。

1. 逐次的なボトルネック: BPTTは本質的に時間に対して逐次的であるため、並列化が妨げられ、現代のハードウェアにおけるスケーラビリティが制限されます。
2. 勾配の不安定性: 勾配の経路長は $O(T)$ （$T$ はシーケンス長）となります。これにより、勾配消失または勾配爆発が生じ、長期的な依存関係の学習が困難になり、モデルが初期の情報を忘れてしまう「近接バイアス（recency bias）」を引き起こします。

Transformerは並列化の問題を解決しましたが、過去の履歴として全トークンを保持するため、圧縮された固定サイズのメモリを持ちません。一方、Linear RNNは固定メモリと並列性を提供しますが、線形遷移関数に制約されているため、状態追跡のような複雑な逐次タスクに対する表現力が限定的です。

手法：教師ありメモリ学習 (Supervised Memory Training: SMT)

本論文では、再帰的な情報の伝播を完全に回避して非線形RNNを学習させる手法である**Supervised Memory Training (SMT)**を提案しています。SMTは、RNNの学習を、一ステップのメモリ遷移ラベルに対する教師あり学習へと還元します。

コアメカニズム

SMTは、「何を記憶するか」（メモリ表現）と**「メモリをいかに更新するか」**（メモリダイナミクス）を分離します。

1. 教師モデル（予測状態学習）: 時間並列なTransformerベースのエンコーダ・デコーダ・ペアを訓練し、過去のコンテキストを**予測状態（predictive state）**であるメモリ状態 $m_t$ へと圧縮します。この状態は、将来の出力を予測するために必要な情報のみを保持し、不要な詳細は破棄します。エンコーダはコンテキスト $x_{\le t}$ を $m_t$ に写像し、デコーダは $m_t$ と将来の入力を用いて将来の出力 $y_{\ge t}$ を予測します。
2. 生徒モデル（RNNダイナミクス）: RNN $f_\theta$ は、教師のメモリ遷移を模倣するように訓練されます。RNNを展開する代わりに、一ステップの更新関数 $(m_t, x_{t+1}) \to \hat{m}_{t+1}$ を学習します。
3. 損失関数:
   • デコーディング損失 ($L_{dec}$): 圧縮されたメモリから未来を予測するように教師を訓練します。
   • ダイナミクス損失 ($L_{dyn}$): 現在のメモリ $m_t$ と入力 $x_{t+1}$ から次のメモリ状態 $\hat{m}_{t+1}$ を予測するようにRNNを訓練します。この際、教師による正解の $m_{t+1}$ を用いて教師あり学習を行います。
   • 一様性損失 ($L_{unif}$): メモリ空間が単一の点に崩壊するのを防ぎます。

学習パイプライン

1. 事前学習 (SMT): 教師エンコーダ・デコーダとRNNを、上記の損失関数を用いて結合的（または逐次的）に訓練します。極めて重要な点は、このフェーズにおいてRNNは一度も展開されないことです。任意の2つのトークン間の勾配経路は $O(1)$ です。なぜなら、クレジットはRNNの再帰ステップを通じたバックプロパゲーションではなく、教師の予測状態目的関数によって割り当てられるからです。
2. ファインチューニング (DMT - DAgger Memory Training): RNNが教師の正解メモリではなく、自身の予測したメモリ（教師の真値からドリフトする可能性がある）を使用することによる「訓練とテストの不一致」に対処するため、軽量なDAggerフェーズを導入します。RNNを自身の予測状態を用いて展開し、オンポリシーの模倣学習を通じて、このドリフトを修正するためにダイナミクス損失を計算します。

主な貢献

• 非線形RNNの時間並列学習: SMTは、展開を伴わずに表現力豊かな非線形RNNの訓練を可能にし、$O(1)$ の勾配経路長を実現します。
• 安定した長期的なクレジット割り当て: 問題を予測状態に対する教師あり学習に還元することで、SMTは $O(T)$ の経路に伴う勾配消失・爆発の問題を排除し、長期的な依存関係の安定した学習を可能にします。
• 固定メモリ推論: Transformerとは異なり、SMTで訓練されたRNNは推論時にも固定のメモリサイズ ($O(M)$) を維持するため、無制限のシーケンス（例：生涯学習）に適しています。
• 表現とダイナミクスの分離: この手法は、メモリ圧縮（並列な教師が担当）とメモリダイナミクス（RNNが担当）を分離し、RNNの最適化ランドスケープを単純化します。

実験結果

著者らは、合成タスクおよび実世界のデータセットを用いて、SMTをBPTTベースラインおよびTransformerと比較評価しました。

• 合成タスク: 勾配の安定性（リトリーバル）、メモリ容量（文字列コピー）、状態追跡（スタック操作）、連合想起（associative recall）、およびインコンテキスト学習をテストするタスクにおいて、SMT→DMTは一貫してBPTTを上回りました。BPTTはシーケンス長が増加するにつれて著しく苦戦しましたが、SMTの性能は安定していました。
• 画素シーケンスモデリング (Attneaveのタスク): MNISTおよびSketchyデータセット（1D画素シーケンスとしてモデル化）において、BPTTで訓練されたRNNは長期的な構造的依存関係（例：ストロークの連続性）を捉えることに失敗しました。一方、SMT→DMTで訓練されたRNNは一貫性のある画像を生成することに成功し、数百ステップにわたる情報の統合における優れた能力を示しました。
• 効率性: SMTおよびSMT→DMTは、同等またはそれ以上の性能を達成するために、BPTTよりも大幅に少ない逐次FLOPs（シリアル計算の尺度）を必要としました。特に長期メモリを必要とするタスクにおいて顕著でした。言語モデリングにおけるデータ効率はBPTTと同等でしたが、画素モデリングにおいてはより優れていました。
• スケーリング則: SMT→DMTは、コンテキスト長、メモリサイズ、およびモデルパラメータの増加に伴い、スムーズな性能向上を示しました。注目すべきは、大規模化においてRNNの性能が教師となるTransformerの性能に接近したことです。
• 汎化: SMTで訓練されたRNNは、訓練時に見られたシーケンス長よりも長いシーケンスに対して、Transformer教師よりも優れた汎化性能を示しました。これは、時間的抽象化に対するより堅牢な帰納バイアスを示唆しています。

意義と主張

本論文は、SMTが過去の経験の時間的抽象化を構築するモデルのスケーリングを解き放つ経路を提供すると主張しています。非線形RNNが安定した勾配と固定メモリを用いて長期的な依存関係を学習できるようにすることで、SMTはBPTT（不安定性と逐次性）とTransformer（メモリ増大の制限）の両方の限界に対処します。

著者らは、SMTを主に事前学習アルゴリズムとして位置付けています。SMTはメモリ表現を効果的に学習しますが、ドリフトを軽減し、特定のダウンストリームタスクに適応させるためには、軽量なポストトレーニングフェーズ（DMTなど）が必要であることも認めています。これにより、モデルが教師の表現力の限界を超えることも可能になります。この研究は、時間並列学習と固定メモリ推論、そして最大級の表現力を組み合わせることが、将来のシーケンスモデルへの有望な方向性であり、エージェントが人間の一生のような無制限の地平から学習することを可能にする可能性があることを示唆しています。