01 · ValueCore
Memory-based construction
Eliminates the independent V projection, using a layer-specific table and contextual keys instead.
Memory Attention.
Replacing the value projection with token-indexed memory and contextual keys.
Fig. 1 Token-indexed memoryM̄ℓ · one per layer
§ 01Abstract
What if values could be retrieved?
Memory Attention (MA) investigates replacing a dedicated attention value projection with learned token-indexed memory, complemented by context-dependent keys.
Each attention layer learns a token memory table; its retrieved representation contributes reusable token-specific information, while the key representation retains contextual information. At inference, row normalization can be folded into the memory table, leaving lookup and addition to construct values. CPU placement and prefetching provide a parameter-storage trade-off. Experiments report improvements in language modeling and average downstream scores under matched training-token budgets, with additional memory parameters.
§ 02Method
Memory instead of a value projection.
Standard attention produces values with a learned dense projection. MA combines a pre-normalization content key with memory retrieved using the original token ID. Normalized rows may be precomputed for inference.
Standard V = xWV becomes
VValueBuilt from a lookup and an addition.
KcontentContextual keyPre-normalization content key, not the RoPE key.
M[ids]Token memoryLayer-specific table, rows normalized, gathered by token ID.
02 · OffloadEvaluated
CPU memory, GPU compute
The token-indexed table may reside on CPU. Prefetching overlaps required lookups and transfers with computation while a conventional KV cache remains in use.
03 · RecallProposed
Reconstruct historical values
Retain sufficient K information and token IDs, then rebuild V on demand. This extension was not included in the paper's inference benchmarks.
§ 03Interactive
Attention, token by token.
Follow the input stream and compare standard causal attention with the proposed memory-lookup interpretation.
Fig. 2 Input sequenceSelect any token
TOKEN 01 / 08Space ← →
Standard attention
Keys and values are retained in the GPU cache
Memory attention
Cached keys · CPU memory lookup · reconstructed values
Current query
Active attention
Previous positions
Cached keys & values
CPU lookup
Causal mask
Conceptual animation only. In actual MA, attention values use pre-RoPE content keys. MA-Recall requires preserving or recovering the correct key representation, with additional retrieval and positional-transform costs.
§ 04Results
Quality and efficiency, in the reported settings.
The following values are from the paper's experimental tables.
Quality · MHA + MA
41.39
Average downstream score
vs 40.68 Standard, 1024 hidden, 10B tokens.
Storage · MA-Offload
7.38%
Less GPU parameter storage
vs Standard in the reference inference setup; excludes KV cache and buffers.
Latency · MA-Offload
Nooverhead
Decode forward-pass latency
17.189 ms vs 17.636 ms Standard at batch 8, 2K context — offloading adds no extra overhead.
| Configuration | Standard avg. ↑ | MA avg. ↑ | Difference |
|---|---|---|---|
| MHA · 1024 hidden | 40.68 | 41.39 | +0.71 |
| GQA · 1024 hidden | 40.90 | 41.51 | +0.61 |
| MQA · 1024 hidden | 41.11 | 41.34 | +0.23 |
| Gated · 1024 hidden | 41.40 | 42.88 | +1.48 |
Source: paper language-modeling results table and Table 4 (inference). Average is the unweighted mean of seven accuracy-based evaluations. MA-Recall is not part of the reported inference measurements.
Ablations: where and how memory helps.
Key findings from the paper's ablation studies. Unless noted, all models use 24 layers, hidden size 1,024, 10B training tokens, and roughly 0.5M tokens per batch.
Memory Attention configuration
41.38 Best in column
AIntegration
Reusing keys as values works best.
| Configuration | Params. | LAMBADA PPL ↓ | WikiText PPL ↓ | Avg. ↑ |
|---|---|---|---|---|
| V = xW std. | 336M | 50.38 | 31.84 | 41.11 |
| V = M | 526M | 58.36 | 32.94 | 40.34 |
| V = (1−λ)·xW + λ·E[i] | 532M | 47.87 | 31.12 | 40.99 |
| V = xW + M | 532M | 45.68 | 29.89 | 41.28 |
| V = K + M | 526M | 45.93 | 29.82 | 41.38 |
Memory alone (V = M) degrades the average to 40.34 — memory complements input-dependent values rather than replacing them; key–value reuse with additive memory is the best integration.
BPlacement
Enrich values, not matching.
| Memory added to | Params. | LAMBADA PPL ↓ | WikiText PPL ↓ | Avg. ↑ |
|---|---|---|---|---|
| Hidden states · H + M | 526M | 50.22 | 30.68 | 41.14 |
| Queries · Q + M | 526M | 59.69 | 33.17 | 40.55 |
| Keys · K + M | 526M | 58.23 | 32.63 | 40.35 |
| Values · V + M | 526M | 45.93 | 29.82 | 41.38 |
| Outputs · O + M | 526M | 71.83 | 34.47 | 39.83 |
Shared key–value representations (K = V). Memory helps most when enriching what attention aggregates; output-side injection is the worst, with LAMBADA perplexity rising to 71.83.
CNormalization
Normalize the memory rows.
| Configuration | Params. | LAMBADA PPL ↓ | WikiText PPL ↓ | Avg. ↑ |
|---|---|---|---|---|
| V = K + E[i] | 526M | 54.46 | 32.60 | 40.85 |
| V = KRoPE + Norm(E[i]) | 526M | 51.38 | 30.79 | 40.61 |
| V = K + Norm(E[i]) | 526M | 45.93 | 29.82 | 41.38 |
Controlling the scale of memory embeddings balances the key and memory contributions, improving both language modeling and downstream accuracy. Sharing the RoPE-transformed key (V = KRoPE + Norm(E[i])) lowers the average to 40.61 — positional encoding must stay out of the shared value path.
DAttnRes
Complementary to attention residuals.
| Configuration | Params. | LAMBADA PPL ↓ | WikiText PPL ↓ | Avg. ↑ |
|---|---|---|---|---|
| Standard | 336M | 50.38 | 31.84 | 41.11 |
| Standard + AttnRes | 336M | 39.59 | 29.77 | 41.00 |
| MA | 526M | 45.93 | 29.82 | 41.38 |
| MA + AttnRes | 526M | 38.37 | 28.88 | 41.46 |
AttnRes improves both Standard and MA; MA + AttnRes reaches the best overall numbers, showing explicit memory and residual propagation are complementary.
Source: paper ablation studies (integration, placement, normalization, and AttnRes tables).
§ 05Implementation
Built on Flash Linear Attention.
The public implementation exposes an MA layer and a Hugging Face-style causal language model. CPU-offload comparisons are currently in standalone profiling scripts; the repository README does not describe them as integrated into the primary causal LM implementation.
Joluck/memory-attention Repository map
├─fla/layers/memory_attn.pyCore layerKey projections, layer-specific token-indexed memory and value construction.
├─fla/models/memory/ModelBackbone configuration and causal language model interface.
└─profile/ProfilingReference inference experiments and CPU-offloading traces.
§ 06Citation
Cite this work.
BibTeX
@misc{kang2026memoryattention, title={Memory Attention}, author={Jiale Kang}, year={2026}, eprint={2609.28399}, archivePrefix={arXiv}, primaryClass={cs.LG}, url={https://arxiv.org/abs/2609.28399} }