すべてがつながっている：テスト時の記憶化、注意バイアス、保持、そしてオンライン最適化を巡る旅

It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization

April 17, 2025

著者: Ali Behrouz, Meisam Razaviyayn, Peilin Zhong, Vahab Mirrokni

cs.AI

要旨

効率的で効果的なアーキテクチャのバックボーンを設計することは、基盤モデルの能力を向上させるための研究努力の核心に位置づけられてきた。人間の認知現象である注意バイアス（特定の事象や刺激を優先する自然な傾向）に着想を得て、我々はTransformer、Titans、そして現代的な線形リカレントニューラルネットワークを含むニューラルアーキテクチャを、内部目的（注意バイアスと呼ばれる）を用いてキーと値のマッピングを学習する連想記憶モジュールとして再概念化した。驚くべきことに、既存のシーケンスモデルのほとんどが、(1) ドット積類似度、または (2) L2回帰目的のいずれかを注意バイアスとして利用していることが観察された。これらの目的を超えて、我々は代替的な注意バイアス設定とその効果的な近似法を提示し、それらの訓練手順を安定化させた。さらに、現代の深層学習アーキテクチャにおける忘却メカニズムを保持正則化の一形態として再解釈し、シーケンスモデルのための新しい忘却ゲートを提案した。これらの洞察に基づいて、我々はMirasという深層学習アーキテクチャを設計するための一般的なフレームワークを提示する。Mirasは、(i) 連想記憶アーキテクチャ、(ii) 注意バイアス目的、(iii) 保持ゲート、(iv) 記憶学習アルゴリズムの4つの選択肢に基づいて構築される。我々は、既存の線形RNNの能力を超えながら、高速で並列化可能な訓練プロセスを維持する3つの新しいシーケンスモデル（Moneta、Yaad、Memora）を提示する。実験結果は、Mirasにおける異なる設計選択が、様々な強みを持つモデルを生み出すことを示している。例えば、Mirasの特定のインスタンスは、言語モデリング、常識推論、リコール集約型タスクなどの特殊なタスクにおいて卓越した性能を発揮し、Transformerや他の現代的な線形リカレントモデルを凌駕することさえある。

English

Designing efficient and effective architectural backbones has been in the core of research efforts to enhance the capability of foundation models. Inspired by the human cognitive phenomenon of attentional bias-the natural tendency to prioritize certain events or stimuli-we reconceptualize neural architectures, including Transformers, Titans, and modern linear recurrent neural networks as associative memory modules that learn a mapping of keys and values using an internal objective, referred to as attentional bias. Surprisingly, we observed that most existing sequence models leverage either (1) dot-product similarity, or (2) L2 regression objectives as their attentional bias. Going beyond these objectives, we present a set of alternative attentional bias configurations along with their effective approximations to stabilize their training procedure. We then reinterpret forgetting mechanisms in modern deep learning architectures as a form of retention regularization, providing a novel set of forget gates for sequence models. Building upon these insights, we present Miras, a general framework to design deep learning architectures based on four choices of: (i) associative memory architecture, (ii) attentional bias objective, (iii) retention gate, and (iv) memory learning algorithm. We present three novel sequence models-Moneta, Yaad, and Memora-that go beyond the power of existing linear RNNs while maintaining a fast parallelizable training process. Our experiments show different design choices in Miras yield models with varying strengths. For example, certain instances of Miras achieve exceptional performance in special tasks such as language modeling, commonsense reasoning, and recall intensive tasks, even outperforming Transformers and other modern linear recurrent models.