Attention's value projection can be a vocabulary lookup: 2.08x parameters in less GPU memory

Memory Attention

Jiale Kang

cs.LG

2026-09-24

Memory Attention builds values as contextual keys plus a per-layer vocabulary lookup, reaching matched loss with 13.8-29.6% fewer tokens and less GPU-resident memory.

What problem this solves

Standard attention derives Q, K and V from the same hidden states through three dense projections. V is the content being aggregated, yet a large share of what it carries is token identity, which the paper describes as reusable across contexts. Paying a dense matmul plus GPU-resident weights for it at every layer is a fixed cost. Prior work added token-specific components on top of the value projection: value embeddings mix word vectors into V, PLE injects per-layer embeddings, Engram retrieves n-gram-indexed memory. This paper asks the more radical question: can WV be removed outright, with a token-indexed lookup supplying identity and the existing keys supplying context?

Method

The construction fits in one line: V = K + Norm(E[s]).

• Each layer keeps a learnable table E of shape vocabulary size x total KV dimension. With a 32k vocab and dv = 1,024 for MHA, one layer's table holds about 32.8M parameters, roughly 0.79B across 24 layers.
• The row indexed by the current token ID is retrieved, RMSNorm-ed within each KV head, and added to K. Q and K are projected as usual; RoPE applies only to the scoring copies; the attention operator itself is untouched.

The division of labor matters: memory says what the token is; K says what it means here. K comes from hidden states that have already absorbed preceding context and layers, so context dependence survives. The same word in different sentences fetches the same row but adds a different K. WK and the tables are trained jointly to support value construction.

At inference, the normalization folds into the tables, so online value construction is one lookup plus one addition, no matmul. Because lookup addresses depend only on token IDs and layer index, the tables can live in CPU memory behind a prefetch pipeline that overlaps transfer with compute. An analyzed extension, MA-Recall, would cache K and token IDs and rebuild historical values as V = K + normalized lookup, halving the KV cache in theory; it is never measured. Relative to MLA: both aggregate a shared representation; MLA applies a learned value mapping after aggregation, while MA aggregates keys plus additive memory.

Results

Matched training token budgets, same seeds, FineWeb-10BT pretraining, average over seven zero-shot tasks:

WikiText perplexity drops from 31.55 to 28.64 for MHA. The larger model (24 layers, d=2,048, 20B tokens) moves from 49.06 to 50.22 on average, with +3.00 on OpenBookQA and +2.65 on ARC-Challenge, and slight drops on ARC-Easy and PIQA. Single-needle retrieval shows the widest gap: within the 2K training window MA averages 96.8% versus 82.1% for standard attention; at 4K extrapolation it is 41.9% versus 25.9%; the largest single gain is 62.4% to 98.2%.

Matched-loss token efficiency is 1.42x for the small model and 1.16x for the large one, meaning 29.6% and 13.8% fewer training tokens.

Inference (single H800, batch 8, 2,048-token prefill plus one decode step with 2,048 tokens of history):

MA-Offload holds 2.08x the total parameters in 7.38% less GPU parameter storage than the baseline, at comparable latency.

Why it matters

If the substitution holds at scale, attention loses one dense matmul and gains a new home for parameters: vocab-indexed tables that can sit in CPU RAM behind prefetching, decoupling model capacity from GPU memory. That axis is orthogonal to quantization and KV-cache compression. The retrieval gains suggest a second effect worth studying: token identity flowing directly into V may keep rare tokens from being washed out across layers. Honestly, this reads as a feasibility study rather than a production recipe: experiments stop at 2.8B parameters and 20B tokens, it is a single-author preprint, and it has not been peer reviewed.

Limitations

• The authors' own headline caveat: gains cannot be separated from added capacity. MHA goes from 373M to 1,135M parameters for +0.71 average points, hardly a striking exchange rate.
• Training compute goes up, not down. Measured training FLOPs are 1.78x baseline (934 versus 525 in Figure 2); combined with 1.16x token efficiency, matched-loss total compute works out to roughly 1.5x baseline. What is saved is GPU memory and inference arithmetic, not training cost.
• Table size is bound to the vocabulary; larger vocabs scale every layer's table, and out-of-vocabulary tokens are not discussed.
• 4K extrapolation degrades sharply for both models, so long-context robustness is unproven; MA-Recall is analyzed but never measured; latency numbers come from a prototype at a single operating point, and the paper says the small deltas should not be read as consistent speedups.
• The memory accounting covers parameter storage only, excluding KV caches, activations and transfer buffers, so total GPU memory is not measured.
• ARC-Easy and PIQA dip at the larger scale; the wins are not uniform.

Terms

• value projection: the dense matrix WV mapping hidden states to attention values; the component this paper deletes
• GQA/MQA: attention variants that shrink the number of KV heads, sharing them across query heads
• MLA: Multi-head Latent Attention from DeepSeek, which compresses K and V into a shared latent vector
• RMSNorm: normalization that rescales a vector by its own root-mean-square
• RoPE: Rotary Position Embedding, applied to Q and K for scoring
• NIAH: needle-in-a-haystack, a retrieval test that hides a key fact inside a long document

Source

• arXiv 2609.28399v1
• PDF
• Hugging Face

What people are saying

• Memory Attention: Paper Replaces Value Projection With Lookup, Cuts GPU Storage 7.38% serrjoa

Related papers

• Stopgrad flow maps converge exponentially even though the update is not a gradient
• Orthogonal 18-harness RL lifts Qwen3-4B in-space Pass@1 by 9.8 over one harness
• Rank-1 Activation Subspace Lifts Qwen3-0.6B Zeroth-Order Tuning Past MeZO
• Prompting a budget isn't enough: Apple compares 11 rewards that teach agents to spend within it
• Per-cell channel Shapley beats MC-Owen on HLS burn scars at 6% of the budget
• Mila derives exact bigram regret: sublinear width leaves a Zipf floor

All paper explainers