PaperlayerView exact revision on arXiv

Source-linked AI summary

Memory Attention

Jiale Kang

arXiv:2609.28399v2cs.LG

TL;DR

Language models usually construct attention values from contextual hidden states, raising the question of whether token-indexed memory can replace the dedicated value projection when paired with contextual information. Memory Attention combines layer-specific token memory with contextual keys. Under matched training token budgets, it improves language-modeling perplexity metrics and average downstream accuracy across evaluated attention configurations, while using additional memory parameters.

  • Problem

    The paper asks whether token-indexed memory can replace part of the dedicated value computation while contextual information is retained.

    Source: §1 · pp. 1–2

  • Method

    Memory Attention constructs values by adding layer-specific token memory to contextual keys, removing the independent value projection.

    Source: §1 · p. 2

  • Results

    Across evaluated attention configurations, MA improves language-modeling perplexity metrics and average downstream accuracy under matched training token budgets, with additional memory parameters.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Token-indexed tables can be placed in CPU memory and prefetched, lowering GPU parameter storage at a latency cost from retrieval and transfer.

    Source: §1 · p. 2

    Because MA includes additional memory parameters, the reported results do not isolate its structure's contribution from increased parameter capacity.

    Source: §4.2 · p. 7

Abstract

from arXiv · show

Language models typically construct attention values from contextual hidden states, even when some of their content may be reusable across contexts. We investigate whether token-indexed memory can replace the dedicated value projection when complemented by contextual information. We propose Memory Attention (MA), which forms values by combining layer-specific token memory with contextual keys. The memory supplies token-specific representations, while the keys preserve context dependence. At inference, normalization can be folded into the memory tables, reducing value construction to lookup and addition. Token-indexed retrieval also enables CPU offloading with prefetching, reducing GPU parameter storage. Under matched training token budgets and with additional memory parameters, experiments across attention configurations show improved language modeling and average downstream performance.

1 Introduction

Memory Attention asks whether token-indexed memory can replace part of the value computation while contextual keys preserve context dependence. The paper combines layer-specific token memory with keys, then reports improved language-modeling and average downstream performance under matched training token budgets, with additional memory parameters.

  • The paper asks whether explicit memory can replace part of an existing computation, rather than only supplementing model capacity with learned representations.The motivation is improving language-model quality under limited computation and GPU memory.Source: §1 · pp. 1–2
  • Memory Attention adds learned, layer-specific token memory to contextual keys, replacing the independent value projection while preserving standard attention weighting and aggregation.The figure expresses the construction as V = K + Norm(M), with M retrieved by token ID.Source: Figure 1 · p. 1, §1 · p. 2
  • At inference, normalization can be folded into memory tables, leaving lookup and addition; MA-Offload prefetches token-indexed entries from CPU memory.The tables can reside outside the GPU, with latency determined by retrieval, transfer, and overlap with computation.Source: §1 · p. 2
  • MA-Recall analyzes reconstructing values from retained keys and token memory in exchange for repeated retrieval and reconstruction, reducing persistent value-cache storage.Source: §1 · p. 2
  • Across evaluated attention configurations, MA improves language-modeling perplexity metrics and average downstream accuracy under matched training token budgets, while adding memory parameters.CPU offloading accommodates the larger parameterization with lower GPU parameter storage than the standard baseline, but adds decoding overhead.Source: §1 · p. 2
  • The evaluation examines model quality and inference behavior alongside trade-offs in parameter capacity, value-construction arithmetic, GPU storage, memory traffic, and persistent cache storage.Source: §1 · p. 2

2 Background

The background presents self-attention as contextual projections whose query–key interactions determine aggregation weights and whose values supply aggregated content. It then reviews token-indexed memory mechanisms that supplement or replace contextual representations in different ways.

  • Self-attention derives queries, keys, and values from current hidden states; query–key interactions set aggregation weights, while values provide the content aggregated.Source: §2 · pp. 2–3
  • Value embeddings retrieve token-specific vectors and interpolate them with contextual projected values, retaining the original value projection.Source: §2 · p. 3
  • Per-Layer Embeddings use a learnable embedding table at each layer, with token-specific information modulated by a hidden-state gate.Source: §2 · p. 3
  • Engram retrieves learned memory from local n-gram patterns; token-ID-based lookup allows retrieval to be prefetched before corresponding hidden states are available.Source: §2 · p. 3
  • DeepEmbedAttention combines token-specific embeddings with contextual key and value projections, using up-projections and embedding tables to support a compact KV cache.Source: §2 · p. 3
  • Bank of Values [7] replaces contextual value projections in the last third of layers with token-indexed tables, whereas Memory Attention combines token memory with contextual keys.Source: §2 · p. 3

3 Memory Attention

Memory Attention replaces the dedicated value projection with layer-specific token memory added to contextual keys. The section develops this value construction, its computation and storage trade-offs, and two extensions for offloading or reconstructing values.

3.1 Memory-Based Value Construction

  • Each MA layer retrieves and normalizes a learned memory vector per input token, then adds it to the contextual key to form the value.The same token retrieves the same vector across positions and contexts within a layer, while separate tables allow layers to learn distinct representations.Source: §3.1 · p. 4

3.2 Contextual and Token-Specific Value Decomposition

  • Attention aggregates contextual keys and token-specific memory under the same attention weights, so a token's memory contribution varies with the query and context.The decomposition replaces the independent value projection with a jointly optimized combination of shared contextual representation and token-indexed memory.Source: §3.2 · pp. 4–5

3.3 Computation and Parameter Storage

  • MA replaces Ld d_v value-projection parameters with L(Nd_v + p_norm) memory and normalization parameters, and changes training value-construction arithmetic from approximately 6LSd_v FLOPs to O(LSd_v).The added capacity is accessed through token lookup rather than a dense projection; the arithmetic comparison excludes optimizer updates, gradient-buffer initialization, and memory traffic.Source: §3.3 · p. 5
  • At inference, normalization can be folded into memory tables, leaving lookup and element-wise addition for online value construction; the cost table excludes lookup overhead and memory traffic.This arithmetic reduction does not by itself determine latency, which also depends on memory access, kernel execution, and synchronization.Source: §3.3 · pp. 5–6, Table 1 · p. 5

3.4 MA-Offload

  • MA-Offload stores folded memory tables in CPU memory and prefetches token-indexed entries to overlap retrieval and transfer with model computation.Offloading does not change total parameter count or value-construction arithmetic, and its latency depends on transfer and execution factors; the historical key and value cache remains.Source: §3.4 · p. 6

3.5 MA-Recall: An Extension for Value Reconstruction

  • MA-Recall explores reconstructing historical values from content keys and token IDs instead of retaining a persistent value cache, but it is not included in the reported inference measurements.The storage reduction trades against repeated memory access and reconstruction work, with practical benefit depending on the attention pattern, cache representation, memory placement, and implementation.Source: §3.5 · p. 6

4 Experiments

The experiments evaluate Memory Attention across language modeling, downstream tasks, retrieval, training-token efficiency, and inference trade-offs. Across matched training-token budgets, the reported evaluations show quality gains alongside additional memory parameters and trade-offs in parameter placement and latency.

4.2 Language Modeling

  • MA improves both perplexity metrics and average accuracy across all three configurations, though gains vary by task and do not extend to every benchmark.Average accuracy gains are 0.71 percentage points for smaller MHA, 0.61 for GQA, and 1.16 for larger MHA.Source: §4.2 · p. 7
  • At 4K, MA scores 41.9% mean retrieval versus Standard’s 25.9%, while outperforming it across all nine task–length combinations.Both models degrade substantially beyond the 2K training window; the results do not establish robust retrieval at longer lengths.Source: Table 3 · p. 8, §4.2 · p. 8

4.3 Ablation Studies

  • V = K + M achieves the best average downstream accuracy, 41.38, and lowest WikiText perplexity, 29.82, among the tested integration strategies.Memory alone reduces average accuracy from 41.11 to 40.34; key–value reuse is presented as an effective complementary integration strategy.Source: §4.3 · p. 8, Table 4 · p. 8
  • Adding memory to values achieves the best overall performance, while modifying queries, keys, or attention outputs generally produces higher perplexity and lower average accuracy.Output-side injection has the highest LAMBADA perplexity, 71.83.Source: §4.3 · pp. 8–9, Table 5 · p. 9
  • Normalized memory reduces LAMBADA perplexity from 54.46 to 45.93 and WikiText perplexity from 32.60 to 29.82; pre-RoPE key sharing also outperforms post-RoPE sharing.With pre-RoPE sharing, LAMBADA perplexity falls from 51.38 to 45.93 and average accuracy rises from 40.61 to 41.38.Source: §4.3 · p. 9, Table 6 · p. 9
  • With AttnRes, MA reaches LAMBADA and WikiText perplexities of 38.37 and 28.88, versus 45.93 and 29.82 for MA alone, and average accuracy of 41.46.AttnRes improves language-modeling perplexity for both Standard and MA.Source: §4.3 · p. 9, Table 7 · p. 9

4.4 Training and Inference Efficiency

  • At the shown loss operating points, MA uses 1.42× and 1.16× token efficiency for L24-D1024 and L24-D2048, respectively; these are not wall-clock speed measurements.These comparisons use additional memory parameters.Source: §4.2 · p. 8, §4.4 · p. 10, Figure 2 · p. 10
  • MA-Offload stores 1,572.864M parameters in CPU memory, cutting GPU parameter storage 55.45% versus MA; at the reference workload, prefill and decode latency are 0.31% and 2.53% lower than Standard.The prototype measurements are an operating-point result, not evidence of consistent speed advantages across workloads; storage estimates exclude KV caches, activations, and transfer buffers.Source: §4.4 · pp. 10–11, Table 8 · p. 11

A Pretraining Settings

Pretraining comparisons use matched settings within each backbone configuration, while the MA memory tables vary with the attention configuration. The setup specifies the shared training recipe, backbone variants, and layer-specific memory dimensions.

A Pretraining Settings

  • Pretraining uses FineWeb-10BT for 20,480 steps, with sequence length 2,048 and global batch size 256.The learning rate peaks at 3 × 10−4 with warmup and cosine decay; Standard and MA share training settings and random seed within each backbone configuration.Source: A Pretraining Settings · p. 13
  • The 24-layer, hidden-dimension-1,024 backbone uses MHA, GQA, or MQA with matching Standard and MA head settings within each comparison.The configurations use 16/16, 16/8, and 4/1 query/KV heads, respectively; comparisons across configurations also differ in head count and dimension.Source: A Pretraining Settings · p. 13, Table 10 · p. 14
  • Each MA layer has an independent learned token-memory table, with widths of 1,024, 512, and 256 for MHA, GQA, and MQA, respectively.Retrieved vectors are normalized per key/value head and added to corresponding keys before RoPE; query heads sharing a key/value head share its memory contribution.Source: A Pretraining Settings · p. 13

Related papers

Loading 2609.28399v2…