Memory Attention: Teaching Transformers to Look Up Values Instead of Projecting Them

Jiale Kang submitted Memory Attention to arXiv on September 23, 2026 (arXiv:2609.28399). The paper asks a narrow, pointed question: if attention’s value matrix is supposed to supply content for aggregation, and if some of that content is token-specific rather than context-specific, why recompute it from dense matrix multiplication every time? The answer is a replacement for the value projection that uses a token-indexed lookup table instead. The result sits in an interesting position: it is not clearly better than standard attention on the same parameter budget, but it does something architecturally different enough to be worth understanding.

Why the Value Projection Is Worth Questioning

Standard self-attention computes queries, keys, and values from the current hidden state through three separate learned projection matrices: Q = XW_Q, K = XW_K, V = XW_V. The queries and keys determine which tokens attend to which; the values supply the content being aggregated. Every forward pass recomputes all three from scratch. For a given model with hidden size d and value dimension d_v, each value projection requires 2·L·S·d·d_v floating-point operations per newly processed token, where L is the number of layers and S is the batch-sequence count.

The question Kang asks is whether the value projection is doing more work than it needs to. Queries and keys need to respond to context, the same token in different positions produces different query/key vectors because the hidden state reflects surrounding text. Values carry the actual information being passed forward through the attention aggregation. Some of that information is genuinely context-dependent. But some of it is token-specific: the same token, in any context, contributes a representation that has a stable, learnable component. That stable component does not need to be recomputed from a matrix multiplication at every layer and every position.

Related work has circled this observation before. Value Embedding (KoszarskyB, 2024), DeepEmbed (BoPeng, 2025), Per-Layer Embeddings (Gemma Team, 2025), STEM (Sadhukhan et al., 2026), and Engram (Cheng et al., 2026) all add lookup-based parameters alongside existing projections. Memory Attention takes a different step: it removes the value projection entirely and replaces it with a combination of the contextual keys and a token-indexed memory table.

How Memory Attention Works

Each MA layer maintains a learnable memory table E in R^(N × d_v), where N is the vocabulary size. The token memory for a sequence is M = Norm(E[s]), where E[s] retrieves the rows indexed by the input token IDs and Norm applies RMSNorm independently within each key/value head. Values are then constructed as:

V = K + M

The memory supplies token-specific content; the keys provide context-dependent information accumulated through all preceding layers. The key projection W_K and the memory table E are learned jointly, so both can adapt to their complementary roles. The standard attention weighting and aggregation operations remain unchanged, only value construction changes.

For attention with RoPE positional embeddings, Kang applies positional transformations only for attention scoring. The content keys used in value construction (V = K + M) come from K before the RoPE rotation; the rotated versions Q_R and K_R are used for computing attention weights. This preserves the standard attention mechanism while allowing values to be constructed from token memory plus unrotated content keys.

Inference and Offloading

At inference, both the memory tables and RMSNorm parameters are fixed. Because RMSNorm acts independently on each retrieved row, it can be precomputed and folded into the tables before inference starts:

E_bar[i] = Norm(E[i]), for i = 1, …, N

Online value construction then becomes V = K + E_bar[s], requiring only a lookup and element-wise addition. The value-construction arithmetic drops from approximately 2·L·S·d·d_v FLOPs for standard attention to L·S·d_v for MA. The paper is careful to note that this arithmetic reduction does not directly translate to lower latency: memory access, kernel execution, and synchronization remain relevant.

Because lookup addresses depend only on token IDs and layer indices, not on hidden states, the required entries can be identified before the corresponding hidden states exist. MA-Offload stores the folded tables in CPU memory and prefetches the needed entries to the GPU before each layer runs, allowing retrieval and transfer to overlap with model computation. A separate extension, MA-Recall, avoids a persistent value cache entirely by retaining content keys and token IDs and reconstructing values on demand, trading cache storage for reconstruction compute.

Results From the Experiments

Experiments train models on NVIDIA H800 GPUs using the flash-linear-attention framework (Yang and Zhang, 2024), with RoPE and RMSNorm. The backbone uses attention blocks and gated MLPs. Within each configuration, Standard and MA use matched training token budgets; MA introduces additional memory parameters because the tables are sized by vocabulary (32,000 tokens) rather than by hidden dimension.

Table 2 in the paper reports language modeling and downstream results across four attention configurations. Selected numbers:

MHA, 24 layers, hidden size 1,024, 10B training tokens: Standard at 373M parameters achieves WikiText PPL 31.55 and average downstream 40.68. MA at 1,135M parameters achieves WikiText PPL 28.64 and average 41.39, a gain of 0.71 percentage points.
GQA, same configuration: Standard at 349M achieves LAMBADA PPL 50.46. MA at 729M achieves 43.68, and average improves from 40.90 to 41.51 (+0.61pp).
Gate attention, same configuration: Standard at 398M achieves average 41.40. MA at 1,160M achieves 42.88 (+1.48pp), with the strongest gains on ARC-Challenge (+1.02pp) and OBQA (+3.00pp).
MHA, 24 layers, hidden size 2,048, 20B tokens: Standard at 1,364M achieves WikiText PPL 21.52. MA at 2,836M achieves 20.79; average improves from 49.06 to 50.22 (+1.16pp), with a 3.00pp gain on OpenBookQA and 2.65pp on ARC-Challenge.

Retrieval experiments using single-needle NIAH tasks show larger gains. Within the 2,048-token training window, MA’s mean score is 97.4% at 1K tokens versus Standard’s 82.6%. At 4K tokens, twice the training context, MA scores 41.9% versus Standard’s 25.9%. Both models degrade substantially beyond the training window; MA is better at extrapolation but not robustly so.

Training token efficiency: at selected matched-loss operating points, MA achieves 1.42× token efficiency for L24-D1024 and 1.16× for L24-D2048, meaning MA reached the same loss levels with approximately 29.6% and 13.8% fewer training tokens, respectively. This measures tokens, not wall-clock time.

Inference latency (Table 4), reference workload, batch size 8, prefill length 2,048, one decoding step with cached history of 2,048 tokens, BF16 precision on a single H800:

Standard: prefill 90.970ms, decode 17.636ms, GPU parameters 2,602 MiB
MA: prefill 88.202ms (3.04% faster), decode 17.788ms (+0.86%), GPU parameters 5,410 MiB (2.08× Standard)
MA-Offload: prefill 90.686ms, decode 17.189ms (2.53% faster than Standard), GPU parameters 2,410 MiB (7.38% less than Standard), CPU parameters 3,000 MiB

MA-Offload’s total parameter count (2,836M) is 2.08× Standard’s (1,364M), but its GPU parameter storage is 192 MiB below Standard because the value projection matrices have been replaced by CPU-resident tables. The inference prototype includes prefetching and computation-transfer overlap; these timings describe model forward passes, not end-to-end serving.

Limitations and Open Questions

The paper is direct about what these results do not show. The quality improvements at matched training token budgets do not isolate the architectural contribution from the parameter count increase. MA always uses more total parameters than Standard, roughly 3× for the smaller models, 2× for the larger. Whether a Standard model trained with more parameters (to match MA’s total count) would outperform MA on the same token budget is not tested. Kang states plainly: “These results characterize the quality and efficiency of the combined design; they do not isolate the contribution of its structure from the increase in parameter capacity.”

Reduced value-construction arithmetic does not uniformly reduce latency. Memory access patterns, kernel execution, and synchronization can dominate. MA-Offload demonstrates feasibility, not a guaranteed speed advantage across workloads. The paper explicitly notes: “Moving the tables off the GPU reduces parameter storage but does not guarantee a reduction in either total GPU memory usage or inference latency.”

The implementation is described as a prototype. Experiments on additional backbones (sliding window attention, linear attention) and larger models including mixture-of-experts are listed as ongoing. The current results cover MHA, GQA, MQA, and gated attention at scales up to roughly 2.8B parameters.

MA-Recall, the extension that eliminates the persistent value cache by reconstructing values from retained keys, is analyzed but not measured in inference experiments. Its practical benefit depends heavily on attention pattern, cache representation, memory placement, and reconstruction implementation.

What This Means for Engineering Teams

The most immediately actionable observation from this paper is the CPU offloading result. MA-Offload operates with GPU parameter storage 7.38% below a comparably performing standard model while adding 2.08× total capacity through CPU-resident tables. If memory table prefetching can be effectively overlapped with computation, which the prototype demonstrates is at least plausible, teams dealing with GPU memory constraints during inference gain a new tool. The parameter capacity is not on the GPU, but it is accessible within the forward pass latency budget.

The MA-Recall extension is the more speculative but more interesting long-term signal. If historical values can be reconstructed from retained keys and token IDs instead of cached from a KV store, the size of the KV cache for long-context inference shrinks by 50% (one key representation per position rather than key+value). At very long contexts, KV cache size is frequently the binding constraint on how many sequences fit in memory simultaneously. This is not yet measured; it is an analysis. But it points toward a direction where the transformer’s memory can be decomposed into tiers with different storage costs and access patterns.

For teams building on LLM infrastructure, the attention architecture is converging toward a three-way split: contextual computation that must happen online, reusable token-indexed representations that can be cached or prefetched, and positional/structural information that may be derivable rather than stored. MA is an early, concrete implementation of the second category taking over work currently done entirely by the first. How well this generalizes to larger models and production workloads will become clearer as experiments on MoE architectures and larger scales complete. Teams watching attention architecture research should track whether the token-efficiency gains hold at scale, which is the more important question than the parameter-storage arithmetic.

Code is available at github.com/Joluck/memory-attention.

Key Takeaways

Memory Attention replaces the value projection matrix with V = K + Norm(E[s]), where E is a per-layer token-indexed table; the standard attention weighting and aggregation operations are unchanged.
At inference, normalization folds into the tables offline, reducing online value-construction arithmetic from 2·L·S·d·d_v to L·S·d_v FLOPs per token.
MA achieves 1.42× training token efficiency at L24-D1024 and 1.16× at L24-D2048 at matched-loss operating points, with consistent perplexity improvements across MHA, GQA, MQA, and gated attention configurations.
MA-Offload stores the memory tables on CPU and prefetches needed entries during the forward pass; at the reference workload this reduces GPU parameter storage by 7.38% relative to Standard while hosting 2.08× total parameters.
Quality improvements do not isolate architectural benefit from parameter count: MA always has more total parameters than the Standard baseline, and the paper does not test a parameter-matched Standard comparison.
MA-Recall, not yet measured in inference experiments, analyzes reconstructing KV values from retained keys plus token IDs, which would reduce the KV cache size by 50% at equal key/value precision.

Work With Origins AI

Origins AI builds production AI systems for engineering teams. If your team is evaluating attention architecture choices or GPU memory constraints during inference at scale, talk to our team.