ATLAS: Leren om de context tijdens testtijd optimaal te onthouden

Samenvatting

Transformers zijn gevestigd als de meest populaire backbones in sequentiemodellering, voornamelijk vanwege hun effectiviteit in in-context retrieval taken en het vermogen om op grote schaal te leren. Hun kwadratische geheugen- en tijdscomplexiteit beperkt echter hun toepasbaarheid in langere sequenties, wat onderzoekers heeft gemotiveerd om effectieve alternatieve architecturen te verkennen, zoals moderne recurrent neural networks (ook wel long-term recurrent memory modules genoemd). Ondanks hun recente succes in diverse downstream taken, worstelen ze met taken die een langetermijncontextbegrip en extrapolatie naar langere sequenties vereisen. We observeren dat deze tekortkomingen voortkomen uit drie losstaande aspecten in hun ontwerp: (1) beperkte geheugencapaciteit die wordt begrensd door de architectuur van het geheugen en de feature mapping van de input; (2) de online aard van updates, waarbij het geheugen alleen wordt geoptimaliseerd met betrekking tot de laatste input; en (3) minder expressief beheer van hun vaste geheugengrootte. Om deze drie aspecten te verbeteren, presenteren we ATLAS, een long-term memory module met hoge capaciteit die leert om de context te onthouden door het geheugen te optimaliseren op basis van de huidige en vorige tokens, waardoor de online aard van long-term memory modellen wordt overwonnen. Op basis van dit inzicht presenteren we een nieuwe familie van Transformer-achtige architecturen, genaamd DeepTransformers, die strikte generalisaties zijn van de originele Transformer-architectuur. Onze experimentele resultaten op het gebied van taalmodellering, gezond verstand redeneren, recall-intensieve en langetermijncontextbegrip taken laten zien dat ATLAS de prestaties van Transformers en recente lineaire recurrent modellen overtreft. ATLAS verbetert verder de langetermijncontextprestaties van Titans, met een nauwkeurigheid van +80\% in een contextlengte van 10M van de BABILong benchmark.

English

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.

ATLAS: Leren om de context tijdens testtijd optimaal te onthouden

ATLAS: Learning to Optimally Memorize the Context at Test Time

Samenvatting

Support