Test-time regression: a unifying framework for designing sequence models with associative memory

要約

シーケンスモデルは、現代の深い学習の中心にあります。
しかし、急速な進歩により、変圧器や再発の代替品など、一見無関係な建築の多様性が生まれました。
この論文では、統一されたフレームワークを紹介して、これらのシーケンスモデルを理解し、導き出し、関連するリコールの経験的重要性、コンテキストに関連するトークンを取得する能力に触発されます。
連想リコールは、2段階のプロセス、暗記、取得、回帰問題として暗記をキャストするものとして形式化します。
これらの2つのステップを組み合わせたレイヤーは、入力トークンを介して「テスト時間回帰」を介して連想リコールを実行します。
線形注意、状態空間モデル、高速プログラマー、オンライン学習者、ソフトマックスの注意を含む顕著な層は、回帰重み、リグレッサー関数クラス、テスト時間最適化アルゴリズムの3つの設計選択で定義された特別なケースとして生じます。
私たちのアプローチは、線形の注意がどのようにトークン間相関をキャプチャできないかを明確にし、ソフトマックスの注意におけるクエリキー正規化の経験的効果のための数学的正当性を提供します。
さらに、設計スペース内の未開拓の地域を照らし、ソフトマックスの注意の新しい高次一般化を導き出すために使用します。
統一を超えて、私たちの作業橋は、古典的な回帰法を備えたシーケンスモデリング、広範な文献を持つフィールドであり、より強力で理論的に原則的なアーキテクチャを開発する方法を開いています。

要約(オリジナル)

Sequence models lie at the heart of modern deep learning. However, rapid advancements have produced a diversity of seemingly unrelated architectures, such as Transformers and recurrent alternatives. In this paper, we introduce a unifying framework to understand and derive these sequence models, inspired by the empirical importance of associative recall, the capability to retrieve contextually relevant tokens. We formalize associative recall as a two-step process, memorization and retrieval, casting memorization as a regression problem. Layers that combine these two steps perform associative recall via “test-time regression” over its input tokens. Prominent layers, including linear attention, state-space models, fast-weight programmers, online learners, and softmax attention, arise as special cases defined by three design choices: the regression weights, the regressor function class, and the test-time optimization algorithm. Our approach clarifies how linear attention fails to capture inter-token correlations and offers a mathematical justification for the empirical effectiveness of query-key normalization in softmax attention. Further, it illuminates unexplored regions within the design space, which we use to derive novel higher-order generalizations of softmax attention. Beyond unification, our work bridges sequence modeling with classic regression methods, a field with extensive literature, paving the way for developing more powerful and theoretically principled architectures.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google