StateX：ポストトレーニングによる状態の拡張でRNNの想起（リコール）性能を向上

要旨: 線形アテンションや状態空間モデルのような再帰型ニューラルネットワーク（RNN）は、長いコンテキストを処理する際の計算がトークンごとに一定であることから人気を集めてきました。しかし、これらの再帰型モデルは、長いコンテキストから文脈情報を正確に想起することを必要とするタスクでは苦戦します。なぜなら、文脈情報のすべてが固定サイズの再帰状態に圧縮されてしまうからです。先行研究では、想起能力は再帰状態のサイズと正の相関があることが示されていますが、大きな再帰状態をもつRNNを直接学習させると学習コストが高くなります。本論文では、事前学習済みのRNNの状態を効率的に拡張するポストトレーニングの枠組みであるStateXを提案します。人気のある2つのRNNクラス、線形アテンションと状態空間モデルに対して、StateX内でポストトレーニングのためのアーキテクチャ修正を設計し、モデルパラメータの増加なし、または増加が無視できる程度で、状態サイズをスケールアップできるようにします。最大1.3Bパラメータのモデルでの実験により、StateXは高いポストトレーニングコストを負うことなく、また他の能力を損なうこともなく、RNNの想起およびインコンテキスト学習の性能を効率的に向上させることが示されました。