ATLAS: Learning to Optimally Memorize the Context at Test Time

要約

主にコンテキスト内検索タスクでの有効性と大規模な学習能力により、トランスはシーケンスモデリングで最も人気のあるバックボーンとして確立されています。
しかし、それらの二次記憶と時間の複雑さは、適用性をより長いシーケンスに縛り付けたため、研究者に最新の再発性ニューラルネットワーク(長期再発メモリモジュール)などの効果的な代替アーキテクチャを探求する動機付けをしました。
多様なダウンストリームタスクでの最近の成功にもかかわらず、彼らは長いコンテキストの理解と長いシーケンスへの外挿を必要とするタスクで苦労しています。
これらの欠点は、デザインの3つのばらばらの側面から来ていることがわかります。(1)メモリのアーキテクチャと入力の特徴マッピングによって制限されている限られたメモリ容量。
(2)更新のオンライン性、つまり、最後の入力に対してのみメモリを最適化する。
(3)固定サイズのメモリの表現力の低い管理。
これら3つの側面すべてを強化するために、現在および過去のトークンに基づいてメモリを最適化し、長期メモリモデルのオンライン性を克服することにより、コンテキストを記憶することを学ぶ大容量の長期メモリモジュールであるAtlasを紹介します。
この洞察に基づいて、私たちは、元の変圧器アーキテクチャの厳格な一般化であるディープトランスフォーマーと呼ばれる変圧器のようなアーキテクチャの新しいファミリーを紹介します。
言語モデリング、常識的推論、リコール集約的、および長いコンテストの理解タスクに関する実験結果は、Atlasがトランスと最近の線形再発モデルのパフォーマンスを上回ることを示しています。
Atlasはさらに、タイタンの長いコンテキストパフォーマンスを改善し、Babilongベンチマークの10mのコンテキスト長で+80 \%精度を達成します。

要約(オリジナル)

Transformers have been established as the most popular backbones in sequence modeling, mainly due to their effectiveness in in-context retrieval tasks and the ability to learn at scale. Their quadratic memory and time complexity, however, bound their applicability in longer sequences and so has motivated researchers to explore effective alternative architectures such as modern recurrent neural networks (a.k.a long-term recurrent memory module). Despite their recent success in diverse downstream tasks, they struggle in tasks that requires long context understanding and extrapolation to longer sequences. We observe that these shortcomings come from three disjoint aspects in their design: (1) limited memory capacity that is bounded by the architecture of memory and feature mapping of the input; (2) online nature of update, i.e., optimizing the memory only with respect to the last input; and (3) less expressive management of their fixed-size memory. To enhance all these three aspects, we present ATLAS, a long-term memory module with high capacity that learns to memorize the context by optimizing the memory based on the current and past tokens, overcoming the online nature of long-term memory models. Building on this insight, we present a new family of Transformer-like architectures, called DeepTransformers, that are strict generalizations of the original Transformer architecture. Our experimental results on language modeling, common-sense reasoning, recall-intensive, and long-context understanding tasks show that ATLAS surpasses the performance of Transformers and recent linear recurrent models. ATLAS further improves the long context performance of Titans, achieving +80\% accuracy in 10M context length of BABILong benchmark.

arxiv情報

著者 Ali Behrouz,Zeman Li,Praneeth Kacham,Majid Daliri,Yuan Deng,Peilin Zhong,Meisam Razaviyayn,Vahab Mirrokni
発行日 2025-05-29 17:57:16+00:00
arxivサイト arxiv_id(pdf)

提供元, 利用サービス

arxiv.jp, Google

カテゴリー: cs.AI, cs.CL パーマリンク