Memory Attention
Memory Attention replaces value projections with token memory plus contextual keys, improving language modeling under matched budgets.
The paper proposes Memory Attention, which builds attention values by combining layer-specific token-indexed memory with contextual keys instead of a dedicated value projection. At inference, normalization can be folded into the memory tables so value construction becomes lookup and addition, and retrieval supports CPU offloading with prefetching. Under matched training-token budgets and additional memory parameters, experiments reported better language modeling and higher average downstream performance.
- Memory Attention builds values from token memory plus contextual keys.
- Inference folds normalization into tables, leaving lookup and addition.
- Token-indexed retrieval enables CPU offloading and prefetching.
- Matched-budget tests improved modeling and average downstream scores.
Full article115 words · extracted from arxiv.org · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2609.28399