Memoria Associativa a Cancelli: Un'Architettura Parallela O(N) per la Modellazione Efficiente di Sequenze

Gated Associative Memory: A Parallel O(N) Architecture for Efficient Sequence Modeling

August 30, 2025

Autori: Rishiraj Acharya

cs.AI

Abstract

L'architettura Transformer, sostenuta dal meccanismo di self-attention, è diventata lo standard de facto per le attività di modellazione di sequenze. Tuttavia, la sua primitiva computazionale di base scala quadraticamente con la lunghezza della sequenza (O(N^2)), creando un significativo collo di bottiglia per l'elaborazione di contesti lunghi. In questo articolo, proponiamo la rete Gated Associative Memory (GAM), una nuova architettura completamente parallela per la modellazione di sequenze che presenta una complessità lineare (O(N)) rispetto alla lunghezza della sequenza. Il blocco GAM sostituisce il livello di self-attention con due percorsi paralleli: una convoluzione causale per catturare in modo efficiente il contesto locale dipendente dalla posizione, e un meccanismo di recupero della memoria associativa parallela per modellare schemi globali basati sul contenuto. Questi percorsi sono dinamicamente fusi utilizzando un meccanismo di gating, consentendo al modello di combinare in modo flessibile informazioni locali e globali per ogni token. Implementiamo GAM da zero e conduciamo un'analisi comparativa rigorosa rispetto a un modello Transformer standard e a una baseline moderna a tempo lineare (Mamba) sul benchmark WikiText-2, nonché rispetto al Transformer sul dataset TinyStories. I nostri esperimenti dimostrano che GAM è costantemente più veloce, superando entrambe le baseline in termini di velocità di addestramento, e raggiunge una perplexity di validazione finale superiore o competitiva su tutti i dataset, stabilendola come una promettente ed efficiente alternativa per la modellazione di sequenze.

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.