Mémoire Associative à Porte : Une Architecture Parallèle O(N) pour une Modélisation Efficace de Séquences
Gated Associative Memory: A Parallel O(N) Architecture for Efficient Sequence Modeling
August 30, 2025
papers.authors: Rishiraj Acharya
cs.AI
papers.abstract
L'architecture Transformer, reposant sur le mécanisme d'auto-attention, est devenue la norme de facto pour les tâches de modélisation de séquences. Cependant, son opération de base présente une complexité quadratique par rapport à la longueur de la séquence (O(N^2)), ce qui constitue un goulot d'étranglement majeur pour le traitement de contextes longs. Dans cet article, nous proposons le réseau Gated Associative Memory (GAM), une architecture novatrice et entièrement parallèle pour la modélisation de séquences, qui présente une complexité linéaire (O(N)) par rapport à la longueur de la séquence. Le bloc GAM remplace la couche d'auto-attention par deux voies parallèles : une convolution causale pour capturer efficacement le contexte local dépendant de la position, et un mécanisme de récupération de mémoire associative parallèle pour modéliser les motifs globaux basés sur le contenu. Ces voies sont fusionnées dynamiquement à l'aide d'un mécanisme de gating, permettant au modèle de combiner de manière flexible les informations locales et globales pour chaque token. Nous implémentons GAM à partir de zéro et menons une analyse comparative rigoureuse contre un modèle Transformer standard et une référence moderne à temps linéaire (Mamba) sur le benchmark WikiText-2, ainsi que contre le Transformer sur le jeu de données TinyStories. Nos expériences démontrent que GAM est systématiquement plus rapide, surpassant les deux références en termes de vitesse d'entraînement, et obtient une perplexité de validation finale supérieure ou compétitive sur tous les ensembles de données, s'établissant ainsi comme une alternative prometteuse et efficace pour la modélisation de séquences.
English
The Transformer architecture, underpinned by the self-attention mechanism,
has become the de facto standard for sequence modeling tasks. However, its core
computational primitive scales quadratically with sequence length (O(N^2)),
creating a significant bottleneck for processing long contexts. In this paper,
we propose the Gated Associative Memory (GAM) network, a novel, fully parallel
architecture for sequence modeling that exhibits linear complexity (O(N)) with
respect to sequence length. The GAM block replaces the self-attention layer
with two parallel pathways: a causal convolution to efficiently capture local,
position-dependent context, and a parallel associative memory retrieval
mechanism to model global, content-based patterns. These pathways are
dynamically fused using a gating mechanism, allowing the model to flexibly
combine local and global information for each token. We implement GAM from
scratch and conduct a rigorous comparative analysis against a standard
Transformer model and a modern linear-time baseline (Mamba) on the WikiText-2
benchmark, as well as against the Transformer on the TinyStories dataset. Our
experiments demonstrate that GAM is consistently faster, outperforming both
baselines on training speed, and achieves a superior or competitive final
validation perplexity across all datasets, establishing it as a promising and
efficient alternative for sequence modeling.