Understanding Transformer from the Perspective of Associative Memory

2025/05/26 by Shu Zhong, Mingyu Xu, Zhong, Shu +6 · 1 voice · 13 citations
Computer Science · Psychology · #Artificial Intelligence (cs.AI) #Educational and Psychological Assessments #FOS: Computer and information sciences #Machine Learning (cs.LG) #cs.AI #cs.LG

paper · pdf · doi:10.48550/arxiv.2505.19488

openalex publication_date 2025/05/26 · arxiv published 2025/05/26 · arxiv updated 2025/05/26 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

In this paper, we share our reflections and insights on understanding Transformer architectures through the lens of associative memory--a classic psychological concept inspired by human cognition. We start with the basics of associative memory (think simple linear attention) and then dive into two dimensions: Memory Capacity: How much can a Transformer really remember, and how well? We introduce retrieval SNR to measure this and use a kernel perspective to mathematically reveal why Softmax Attention is so effective. We also show how FFNs can be seen as a type of associative memory, leading to insights on their design and potential improvements. Memory Update: How do these memories learn and evolve? We present a unified framework for understanding how different Transformer variants (like DeltaNet and Softmax Attention) update their "knowledge base". This leads us to tackle two provocative questions: 1. Are Transformers fundamentally limited in what they can express, and can we break these barriers? 2. If a Transformer had infinite context, would it become infinitely intelligent? We want to demystify Transformer architecture, offering a clearer understanding of existing designs. This exploration aims to provide fresh insights and spark new avenues for Transformer innovation.

Citations

• It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization
• RWKV-7 "Goose" with Expressive Dynamic State Evolution
• Forgetting Transformer: Softmax Attention with a Forget Gate
• MoBA: Mixture of Block Attention for Long-Context LLMs
• DeltaProduct: Improving State-Tracking in Linear RNNs via Householder Products
• Test-time regression: a unifying framework for designing sequence models with associative memory
• MiniMax-01: Scaling Foundation Models with Lightning Attention
• Titans: Learning to Memorize at Test Time
• KV Shifting Attention Enhances Language Modeling
• Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues
• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
• Parallelizing Linear Transformers with the Delta Rule over Sequence Length
• Understanding Transformer Reasoning Capabilities via Graph Algorithms
• Length Generalization of Causal Transformers without Position Encoding
• The Illusion of State in State-Space Models
• Mechanistic Design and Scaling of Hybrid Architectures
• Gated Linear Attention Transformers with Hardware-Efficient Training
• YaRN: Efficient Context Window Extension of Large Language Models
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• Retentive Network: A Successor to Transformer for Large Language Models
• The Impact of Positional Encoding on Length Generalization in Transformers
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
• Robust Speech Recognition via Large-Scale Weak Supervision
• In-context Learning and Induction Heads
• Toy Models of Superposition
• What Can Transformers Learn In-Context? A Case Study of Simple Function Classes
• The Parallelism Tradeoff: Limitations of Log-Precision Transformers
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• NormFormer: Improved Transformer Pretraining with Extra Normalization
• Highly accurate protein structure prediction with AlphaFold
• RoFormer: Enhanced Transformer with Rotary Position Embedding
• Learning Transferable Visual Models From Natural Language Supervision
• Linear Transformers Are Secretly Fast Weight Programmers
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Hopfield Networks is All You Need
• Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
• Language Models are Few-Shot Learners
• GLU Variants Improve Transformer
• Root Mean Square Layer Normalization
• Attention Is All You Need
• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• Correlation Matrix Memories

Cited by

• Attention Residuals
• Raven: High-Recall Sequence Modeling with Sparse Memory Routing
• Distributed Dynamic Associative Memory via Online Convex Optimization
• Higher-order Linear Attention
• Local Linear Attention: An Optimal Interpolation of Linear and Softmax Attention For Test-Time Regression
• Understanding the Mixture-of-Experts with Nadaraya-Watson Kernel
• Short window attention enables long-term memorization
• Distributed Associative Memory via Online Convex Optimization
• UniPool: A Globally Shared Expert Pool for Mixture-of-Experts
• Bridging Compositional and Distributional Semantics: A Survey on Latent Semantic Geometry via AutoEncoder
• Learning to Disentangle Latent Reasoning Rules with Language VAEs: A Systematic Study
• T-SHRED: Symbolic Regression for Regularization and Model Discovery with Transformer Shallow Recurrent Decoders
• Comba: Improving Bilinear RNNs with Closed-loop Control

Discussions

• この論文はマジ重要。こいつだけはしっかり読んで理解しきらないと今後まずい。 今日はざっと眺めて、明日がっつり行くとしますか。 Understanding Transformer from the Perspective of Associative Memory arxiv.org/abs/2505.19488 [bsky, 1 points, 0 comments]

Related

• Beyond Scaling Laws: Understanding Transformer Performance with Associative Memory
• Understanding Factual Recall in Transformers via Associative Memories
• Do Transformers Need Deep Long-Range Memory
• TransformerFAM: Feedback attention is working memory
• An Introduction to Transformers