Titans: Learning to Memorize at Test Time
https://arxiv.org/pdf/2501.00663
概要
Over more than a decade there has been an extensive research effort of how effectively utilize recurrent models and attentions. While recurrent models aim to compress the data into a fixed-size memory (called hidden state), attention allows attending to the entire context window, capturing the direct dependencies of all tokens. This more accurate modeling of dependencies, however, comes with a quadratic cost, limiting the model to a fixed-length context. We present a new neural long-term memory module that learns to memorize historical context and helps an attention to attend to the current context while utilizing long past information. We show that this neural memory has the advantage of a fast parallelizable training while maintaining a fast inference. From a memory perspective, we argue that attention due to its limited context but accurate dependency modeling performs as a short-term memory, while neural memory due to its ability to memorize the data, acts as a long-term, more persistent, memory. Based on these two modules, we introduce a new family of architectures, called Titans, and present three variants to address how one can effectively incorporate memory into this architecture. Our experimental results on language modeling, common-sense reasoning, genomics, and time series tasks show that Titans are more effective than Transformers and recent modern linear recurrent models. They further can effectively scale to larger than 2M context window size with higher accuracy in needle-in-haystack tasks compared to baselines.
ChatGPT訳：
過去10年以上にわたり、再帰モデルとアテンションをいかに効果的に活用するかについて広範な研究が行われてきた。再帰モデルはデータを固定サイズのメモリ（隠れ状態と呼ばれる）に圧縮することを目的とする一方、アテンションはコンテキスト全体に注目し、すべてのトークン間の直接的な依存関係を捉える。しかし、この依存関係のより正確なモデル化には二次的な計算コストが伴い、モデルを固定長のコンテキストに制限してしまう。
本研究では、新しいニューラル長期記憶モジュールを提案する。これは、過去のコンテキストを記憶し、アテンションが現在のコンテキストに注目しながら、長期の過去情報を活用できるようにするものである。我々は、このニューラルメモリが並列化可能な高速学習を実現しつつ、推論も高速に維持できることを示す。メモリの観点から見ると、アテンションはその限定的なコンテキストによって短期記憶として機能するが、データを記憶する能力を持つニューラルメモリは、より長期的で持続性のある記憶として機能すると考えられる。
これら2つのモジュールに基づき、新しいアーキテクチャのファミリー「Titans」を提案し、メモリを効果的に統合するための3つのバリアントを提示する。我々の実験結果によると、言語モデリング、常識推論、ゲノム解析、時系列タスクにおいて、TitansはTransformersや最新の線形再帰モデルよりも高い効果を示した。さらに、ベースラインと比較して、「干し草の中の針」タスクにおいて、200万以上のコンテキストウィンドウサイズまでスケール可能であり、高い精度を維持できることが確認された。
Transformer 2.0 と呼ばれてる?
多分｡
