Ассоциативная память с управляемыми воротами: параллельная архитектура O(N) для эффективного моделирования последовательностей

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

August 30, 2025

Авторы: Rishiraj Acharya

cs.AI

Аннотация

Архитектура Transformer, основанная на механизме самовнимания, стала де-факто стандартом для задач моделирования последовательностей. Однако её основная вычислительная операция масштабируется квадратично относительно длины последовательности (O(N^2)), что создаёт значительное узкое место при обработке длинных контекстов. В данной статье мы предлагаем сеть Gated Associative Memory (GAM) — новую, полностью параллельную архитектуру для моделирования последовательностей, которая демонстрирует линейную сложность (O(N)) относительно длины последовательности. Блок GAM заменяет слой самовнимания двумя параллельными путями: причинной сверткой для эффективного захвата локального, зависящего от позиции контекста и параллельным механизмом извлечения ассоциативной памяти для моделирования глобальных, основанных на содержании паттернов. Эти пути динамически объединяются с использованием механизма гейтирования, что позволяет модели гибко комбинировать локальную и глобальную информацию для каждого токена. Мы реализуем GAM с нуля и проводим строгий сравнительный анализ с базовой моделью Transformer и современным линейным базовым методом (Mamba) на бенчмарке WikiText-2, а также с Transformer на наборе данных TinyStories. Наши эксперименты показывают, что GAM стабильно быстрее, превосходя оба базовых метода по скорости обучения, и достигает превосходного или конкурентоспособного итогового значения perplexity на всех наборах данных, что делает её перспективной и эффективной альтернативой для моделирования последовательностей.

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.