PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding Transformer from the Perspective of Associative Memory

Shu Zhong, Mingyu Xu, Tenglong Ao, Guang Shi

arXiv:2505.19488v1cs.LGcs.AI

TL;DR

The paper addresses how Transformer mechanisms work, how their memory capacity and updates shape behavior, and whether expressivity or unlimited context presents theoretical challenges. It analyzes Transformers as associative memories and concludes that retrieval, update rules, and context growth have distinct theoretical implications.

  • Problem

    Transformer internals remain complex and not fully understood, motivating analysis of memory capacity, update mechanisms, expressivity, and unlimited-context learning.

    Source: §1 · p. 1

  • Method

    The paper analyzes Transformer components as associative memories, using retrieval SNR and a kernel perspective for capacity and a generalized recurrent framework for memory updates.

    Source: §1 · pp. 1–2, §2.1 · p. 4

  • Results

    Softmax Attention’s exponential kernel outperforms Linear Attention in retrieval accuracy and memory capacity; memory may converge as context grows, degrading in-context learning.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Associative-memory analysis offers a unified framework for understanding Transformer memory capacity and update dynamics and identifying avenues for architectural innovation.

    Source: §4 · p. 21

    Without optimizations such as memory decay, vanilla linear attention has an unbounded objective, an exploding memory norm, and cannot perform very long in-context learning.

    Source: §2.2.2 · p. 16

Abstract

from arXiv · show

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.

1 Introduction

The paper interprets Transformer components as associative memories and analyzes their capacity and update strategies. This perspective frames questions about retrieval, model expressivity, and learning with unlimited context.

1 Introduction

  • Attention forms short-term contextual associations, whereas FFNs encode persistent long-term associative memories distilled during training.Source: §1 · pp. 1–2
  • Retrieval SNR measures associative-memory retrieval quality, and a kernel analysis explains Softmax Attention’s retrieval and capacity advantage over Linear Attention.Source: §1 · p. 2
  • The paper compares Attention’s exponential kernel with FFNs, discusses polysemanticity, and proposes kernel design, multihead mechanisms, sparsity, and gating as improvement avenues.Source: §1 · p. 2
  • A generalized recurrent update framework recasts Linear Attention, DeltaNet, Gated Attention, and Softmax Attention as distinct associative-memory update strategies.Source: §1 · p. 2
  • The proposed associative-memory model combines Softmax Attention’s retrieval accuracy with DeltaNet’s delta-rule updates and is proved more expressive than standard Transformers.The paper relates memory-update strategies to expressivity measured by circuit complexity.Source: §1 · p. 2
  • As context grows, memory may converge and degrade in-context learning, while the paper analyzes factors that could mitigate convergence.Source: §1 · p. 2

2 Associative Memory

This section examines Transformer attention and FFNs as associative memories, focusing on retrieval capacity, kernel and architectural choices, and memory-update rules. It also considers how updates behave over long contexts, including DeltaFormer’s expressivity, differences between DeltaNet and softmax attention, and chunking for DeltaNet training.

2 Associative Memory

  • Attention is an associative memory with S_t = Σ_i v_iϕ(k_i)⊤ and recall map f_S(q) = Sϕ(q); FFNs also implement associative memory.Attention stores short-term contextual information in the KV cache, discarded after inference, while FFNs encode persistent knowledge in weights learned during training.Source: §2 · pp. 3–4

2.1 Memory Capacity

  • Reliable linear associative recall requires the number of stored pairs to be on the order of or less than key dimensionality, dk ≳ N; beyond this, noise degrades retrieval.This capacity limit helps explain why linear Transformers can underperform on precise long-context retrieval tasks [20] [28]. Kernel mapping and sparsity are proposed as capacity-improvement strategies with structural similarities.Source: §2.1 · p. 5

2.1.1 Kernel Mapping

  • With τ = √dk, the exponential kernel corresponds to Softmax Attention and reduces required feature dimensionality from O(N) to O(log2 N), while ReLU-kernel FFNs retrieve less precisely.The standard two-layer, bias-free ReLU FFN can be interpreted as an associative memory; ReLU suppresses negative-similarity pairs but remains less precise than the exponential kernel.Source: §2.1.1 · pp. 5–7

2.1.2 Vignette 1: Rethinking Attention and FFN from the Perspective of Associative Memory

  • The authors hypothesize that noisier ReLU recall enables polysemantic superposition and denser knowledge, while precise attention favors token-specific recall; SoLU increased monosemanticity but hurt TriviaQA accuracy.Their analysis indicates that SoLU offers higher retrieval precision than ReLU, supporting a tradeoff between knowledge density and retrieval precision.Source: §2.1.2 · pp. 7–8
  • The authors view normalization mainly as a training-stability mechanism rather than a fundamental design requirement; NormFormer [37] adds LayerNorm after the FFN activation.They note that linear-attention variants [20, 40, 48] show normalization can be replaced by a simple per-head norm.Source: §2.1.2 · p. 8
  • Multihead attention computes distinct similarity patterns and can trade retrieval precision for expressivity or superposition, motivating renewed consideration of multihead mechanisms in FFNs.The authors also propose sparse routing for attention, where keys and values are dynamic, contrasting with static FFN expert parameters; MoBA [22] is given as an attention-routing example.Source: §2.1.2 · pp. 8–9
  • Attention gating can dynamically amplify or suppress query responses to specific keys, while FFN-style gating motivates a simpler non-cumulative alternative left for future work.Gated Linear Attention [47] and FoX [21] are cited as attention-gating examples; the proposed simpler gating hypothesis remains untested.Source: §2.1.2 · pp. 9–10

2.2 Memory Update

  • Classical associative memory adds each new value-key outer product, an update interpretable as gradient descent on an objective that can drive memory-response norms without bound.Decay mechanisms such as those in [47] can mitigate the resulting numerical instability.Source: §2.2 · p. 10
  • Linear attention's unbounded objective risks norm explosion; DeltaNet uses L2 regression and a key-dependent decay to attenuate overlapping historical values before writing, with momentum also available.The cited momentum finding is a reference rather than a strict equivalence to DeltaNet [2].Source: §2.2 · p. 11
  • Softmax attention's exponential kernel improves recall, and normalization regularizes memory updates, but its objective can vanish with long contexts and suppress updates.The passage also notes that adding memory decay can improve softmax-attention performance [21].Source: §2.2 · pp. 11–12, §2.2.1 · p. 12

2.2.1 Vignette 2: Softmax Attention with Delta Rule

  • DeltaFormer integrates the kernel trick with the delta rule by processing each value to remove information overlapping with historical memory before writing, while retaining softmax-style recall.Its general form permits independent kernels, gates, and a retrieval vector; Transformer and DeltaNet are special cases.Source: §2.2.1 · pp. 12–13
  • DeltaFormer can track n elements for n ≥5, which the authors use to show expressivity beyond TC0 and theoretically reaching NC1.The cited analysis reports that constant-depth Transformers without chain-of-thought are limited to TC0 [23].Source: §2.2.1 · p. 15
  • Chunking the sequence enables fully parallel computation within each chunk, reducing DeltaNet’s training time complexity from O(T) to O(T/C), while softmax attention has O(1) time complexity.C denotes the chunk size.Source: §2.2.1 · pp. 15–16

2.2.2 Vignette 3: What Will Happen in Infinitely Long In-Context Learning?

  • As context length approaches infinity, single-head DeltaNet’s memory converges to a fixed value, making it a bi-gram model that loses its in-context learning capability.The analysis assumes Wk is typically invertible and the input second-moment matrix is generally positive definite.Source: §2.2.2 · p. 16
  • Multi-head DeltaNet’s low-rank projection matrices make its optimal memory depend on the data distribution rather than being fixed, which may alleviate in-context learning degradation.The data-distribution dependence follows because the input second-moment term cannot be eliminated when the key projection is not the identity.Source: §2.2.2 · p. 16
  • Vanilla linear attention has an unbounded optimization objective and an exploding memory norm, whereas DeltaNet’s bounded objective avoids these issues and offers a more reasonable basis for analysis.The passage notes that linear attention without optimizations such as memory decay cannot perform very long in-context learning.Source: §2.2.2 · p. 16
  • As context length grows, softmax attention’s gradient approaches zero and its memory updates gradually cease, diminishing in-context learning as relevant information becomes harder to retrieve from the long context.The passage describes the long KV cache as having a low signal-to-noise ratio.Source: §2.2.2 · p. 17
  • Removing softmax attention’s normalization prevents its gradient from vanishing but leaves an unbounded optimization objective that would eventually cause model collapse.Source: §2.2.2 · p. 17
  • Although memory updates may stagnate in theory, multi-head attention, memory decay, short key-value convolutions, and multilayer interactions can prevent smooth convergence to the theoretical fixed solution.Fluctuations in single-layer optimization, including those introduced by the input second-moment matrix, may accumulate and amplify across layers.Source: §2.2.2 · p. 17

3 Experiments

Across swapping and graph-reachability experiments, single-layer DeltaFormer performs better than Transformer, while the multi-head comparison shows different head-count effects for Linear and Softmax Attention. Kernel choice, curriculum learning, and positional embeddings also affect swapping performance.

3.1 Multi-head Efficiency Trade-off

  • Reducing head count to one benefits Linear Attention but degrades Softmax Attention, supporting different capacity-versus-expressivity hypotheses for the two mechanisms.The comparison used GPT-2 small with 12 layers and 768 hidden dimensions; linear attention replaced Softmax with a linear kernel and both settings used per-head RMS normalization.Source: §3.1 · pp. 17–18

3.2 Track the Exchange of Elements

  • DeltaFormer outperforms Transformer across simple κ1 designs on swapping, and a single layer tracks exchanges of five elements.Adding Transformer layers did not improve performance; the authors speculate optimization may be involved.Source: Figure 4 · p. 18, §3.2 · p. 19
  • Similarity-function choice strongly affects swapping: κ1 closer to the construction performs better, while unsuitable choices can accumulate update errors.The authors link this sensitivity to perturbations of a potentially ill-conditioned inverse; the calculation for u is less stable than the calculation for o.Source: §3.2 · p. 19
  • At fixed d, linear-kernel performance degrades as n≥d increases; nonlinear functions can store information through superposition.The stress test used d=128 and n∈{128, 256, 512}.Source: §3.2 · p. 19, Figure 5 · p. 19
  • Gradually extending the training window from 32 to 256 yields better performance with fewer computations and samples than starting at length 256.The curriculum doubles training length when accuracy reaches 0.99.Source: §3.2 · pp. 19–20, Figure 6 · p. 20
  • Without RoPE, direct training at length 256 produces random scores, whereas the 32-to-256 curriculum reaches 100% accuracy and shows better length generalization at extension points.The authors speculate RoPE may hinder expression and extrapolation while helping optimization.Source: §3.2 · p. 20, Figure 6 · p. 20

3.3 Reachability of Directed Acyclic Graphs

  • On 32-node reachability, multilayer Transformer struggles to reach 100% accuracy, while a single-layer DeltaFormer performs very well.The authors speculate optimization also constrains Transformer effectiveness on this task.Source: §3.3 · p. 20, Figure 7 · p. 21

4 Conclusion

The paper presents associative memory as a unifying lens on Transformer behavior while emphasizing that understanding these architectures remains an ongoing effort. It raises open questions about architecture, memory consolidation, and computational limits.

4 Conclusion

  • The conclusion leaves open whether associative memory has an optimal architecture, whether attention and FFNs could consolidate contextual experience into persistent knowledge, and whether higher parallel complexity or new hardware could change what models solve.The proposed consolidation could support continuous, adaptive test-time learning and potentially mitigate challenges from infinitely long contexts.Source: §4 · p. 21

A Derivation of the Inverse SNR Formula

  • The inverse-SNR derivation starts from a general kernel-based recall equation and specializes the derivations for linear and exponential kernels.The linear associative-memory form is treated as a special case of Eq. 9.Source: A Derivation of the Inverse SNR Formula · p. 26

B.1 General Form

  • A scalar optimization objective for the general recurrent memory update exists under the stated construction only when the parameter matrices At and Bt are symmetric.The passage says nonsymmetric matrices generally do not yield a scalar potential whose gradient is the desired vector field.Source: §B.1 · pp. 28–29

B.2 Representative Cases

  • With momentum coefficient η in (0, 1), the recurrent formulation adds a constant-coefficient term to Ct and updates the optimization objective.Source: §B.2 · p. 30

C Delta Rule with Kernel Function

D Re-examining DeltaNet from the Perspective of Information Aggregation

D.1 Information Aggregation

Vanilla attention writes each timestep’s value into memory independently, whereas the delta-rule formulation recovers DeltaNet and writes only incremental information. This difference gives the delta rule both writing and deletion capabilities.

D.1 Information Aggregation

  • The information-aggregation formulation recovers DeltaNet from an alternative perspective.Source: §D.1 · p. 32
  • Unlike vanilla attention’s independent value writes, the delta rule removes information redundant with historical values based on key similarity, enabling both memory writing and deletion.Source: §D.1 · p. 32

D.2 Enhancing Retrieval Capability of DeltaNet

The paper proposes replacing DeltaNet’s key-similarity computation with an exponential kernel and integrating the delta rule with softmax attention.

D.2 Enhancing Retrieval Capability of DeltaNet

  • The proposed retrieval enhancement replaces DeltaNet’s key-similarity computation with an exponential kernel and introduces a normalization term, integrating the delta rule with softmax attention.Source: §D.2 · pp. 32–33

D.3 Efficient Chunk-wise Implementation

The recurrent computation for the proposed sequence update is costly and sequential, so the paper describes parallel matrix operations and chunking as runtime trade-offs. It also presents an iterative method for block matrix inversion on GPUs.

D.3 Efficient Chunk-wise Implementation

  • For T ≫ d, the recurrent update costs O(T^2d) and has O(T) training-time complexity because each step depends on previous states, limiting GPU parallelism.Source: §D.3 · p. 33
  • Rewriting the update as a matrix system permits parallel matrix operations, but the stated inversion approach raises complexity to O(T^3).Source: §D.3 · p. 33
  • Chunking reduces recurrent steps from O(T) to O(N) while changing complexity from O(T^2d) to O(T^2d + TCd + TC^2), trading computation for reduced runtime.Source: §D.3 · p. 33

D.3.1 Efficient Matrix Inversion on GPUs.

  • The GPU inversion procedure assumes a strict lower-triangular C × C matrix A with AC = 0, then iterates to compute (I + A)^-1.Source: §D.3.1 · pp. 34–35

E State Tracking of DeltaFormer

The section establishes a state-exchange construction for attention under assumptions on the keys and a function f. It states that DeltaFormer can express this attention and notes spatial-complexity advantages for state tracking.

E State Tracking of DeltaFormer

  • DeltaFormer can express the attention mechanism described by the state-exchange theorem.The section also analyzes DeltaFormer’s spatial-complexity advantages in state tracking.Source: E State Tracking of DeltaFormer · pp. 35–36
  • The construction assumes unit-norm keys with pairwise inner products bounded by ϵ(d, n), and a function f satisfying Assumption 2.The keys are initialized with corresponding values as n key-value pairs.Source: E State Tracking of DeltaFormer · p. 36
  • For t>n, querying either of two selected keys retrieves the other key’s stored value, while querying any other key retrieves its own value.This implements an exchange of the values associated with the selected keys.Source: E State Tracking of DeltaFormer · p. 36

E.1 Proof of Theorem 1

The proof of Theorem 1 introduces an auxiliary lemma and specifies the conditions under which it yields an identity used in the proof.

E.1 Proof of Theorem 1

  • An auxiliary lemma is introduced to facilitate the proof of Theorem 1.The theorem concerns state exchange.Source: §E.1 · pp. 36–37
  • The lemma assumes the keys satisfy Assumption 1 and that k>n is the difference between two keys from the set.It also requires f to satisfy Assumption 2.Source: §E.1 · p. 37
  • Under those conditions, the lemma establishes an identity used in the proof.Source: §E.1 · p. 37

E.1.1 Proof of Lemma 1

The lemma is proved by dividing the index cases according to whether l is within the original key range or represents a difference of keys, then checking possible index equalities.

E.1.1 Proof of Lemma 1

  • The proof separates l≤n from l>n, representing the latter as the difference k_l1−k_l2 for l2<l1.Source: §E.1.1 · p. 37
  • When exactly one index pair matches, the function evaluates to 1 or −1 according to which of i or j matches l1 or l2.The cases i=l1 and j=l2 yield 1; i=l2 and j=l1 yield −1.Source: §E.1.1 · p. 37
  • When both pairs match, i=l1 and j=l2 yields 2, while the reversed matching contradicts the ordering condition.Source: §E.1.1 · p. 38
  • The proof combines the case analyses to complete the lemma.Source: §E.1.1 · p. 38

E.1.2 Formally Prove Theorem 1

Theorem 1 is proved by induction over exchange steps. The result states that the model can trace exchanges of n states regardless of how large the exchange step becomes.

E.1.2 Formally Prove Theorem 1

  • At t=n+1, the rules make it possible to trace the states exchanged between t1 and t2.Source: §E.1.2 · p. 38
  • For t>n+1, assuming the proposition holds at t−1, the proof shows that the retrieved states for {k1,…,kn} are correct at step t.Source: §E.1.2 · pp. 38–39
  • By induction, the model can trace the exchange of n states regardless of how large the exchange step t is.The base case establishes tracing at t=n+1, and the induction step retrieves the correct states for the original keys.Source: §E.1.2 · pp. 38–39

E.2 Re-exam Assumptions

The section revisits assumptions behind state tracking, deriving a dimension condition and a GQA-like kernel construction. Its toy experiment reports stronger S5 tracking as query count increases with nonlinear kernels.

E.2 Re-exam Assumptions

  • Assumption 1 requires d = O(log n) under the stated α = 0.01 and ε = 0.1 setting.The passage reports ϵ(d, n) ≤ 0.11 < 1/8 for these parameter choices.Source: §E.2 · p. 39
  • Four shared query heads can theoretically express an f(·) satisfying Assumption 2, while combining multiple kernel terms addresses rounding’s optimization difficulty.The construction uses weights a_j; the passage says the linear combination can track at least S_d.Source: §E.2 · pp. 39–40
  • As query count increases, nonlinear-kernel DeltaFormer becomes stronger at tracking S5, suggesting GQA-like grouping can increase kernel complexity and model capacity.The experiment uses head dim=3 to track exchanges of five elements across 16 exchanges.Source: Figure 8 · p. 40, §E.2 · p. 40
  • State exchange between t1 and t2 is achievable when d = O(log n) and learned projections satisfy k_t = k_t1 − k_t2 and v_t = 0.This is the condition stated in the summary of Theorem 1.Source: §E.2 · p. 40

E.3 Stronger Compression Capability

The section compares the space needed to track exchanges with and without periodic rewriting. It argues that nonlinear kernels expand the delta rule’s tracking capability, while softmax retrieves richer information than linear attention.

E.3 Stronger Compression Capability

  • Tracking T exchanges of n states requires O(T log n) space with d = O(log n) and a KU cache of length T.The stated storage requirement combines the dimension and cache length.Source: §E.3 · p. 40
  • Rewriting the recalled states into a KU cache of length O(n) yields total space O(n log n), compared with O(n2) for an identity mapping.The cited comparison notes that Peng et al. [26] use a 5 × 5 matrix to track five-element exchange.Source: §E.3 · p. 40
  • Nonlinear kernels let the delta rule track exponentially many states rather than only the previous n = O(d) states, while softmax retrieves richer information from longer contexts than linear attention.Source: §E.3 · p. 40

E.4 Toy Experiment

The toy experiment code generates five-element swap sequences and trains models on the resulting tracking task. The analytical section formulates an L2 memory objective and states its optimal solution condition.

E.4 Toy Experiment

  • The toy code generates random swap sequences over five elements and tracks the first element through the resulting permutations.The reference code is provided for the DeltaFormer state-tracking toy experiment.Source: §E.4 · pp. 40, 43
  • Training doubles sequence length after accuracy exceeds 0.999, with the initial sequence length set to k=16.Source: §E.4 · p. 45

F Analytical Solution of Associative Memory

  • For fixed test-time weights, the analytical section formulates L2 memory optimization over inputs x, with k = W_kx and v = W_vx, and gives the optimum by setting the gradient to zero.The text states that this yields the optimal memory matrix S∗.Source: F Analytical Solution of Associative… · p. 46

Related papers

Loading 2505.19488v1…