Memory-Augmented Self-Attention

• Memory-augmented self-attention is a technique that extends standard attention by incorporating persistent or dynamic memory to capture longer-range dependencies.
• It employs methods such as recurrent-output memory, learnable key–value vectors, and episodic chunk summaries to blend past representations with current queries.
• Empirical studies in NLP, speech, and vision demonstrate that these memory mechanisms can improve model efficiency, performance, and robustness in long-context scenarios.

Searching arXiv for relevant papers on memory-augmented self-attention and closely related architectures. Memory-augmented self-attention denotes a family of attention mechanisms in which the query at a position reads not only from the immediately available sequence states but also from an explicitly designated memory. That memory may be a sliding window of prior recurrent outputs, a bank of persistent learnable key–value vectors, segment summaries, memory tokens, chunk-level episodic representations, or a parametric fast-weight state updated online. In each case, the central operation remains attention, but the key/value pool or the attention weights are altered so that current computation is conditioned on stored representations beyond the default token sequence. Early formulations used a differentiable memory over recent recurrent states in language modeling (Daniluk et al., 2017); later Transformer variants introduced persistent memory vectors and memory tokens as first-class attention participants (Sukhbaatar et al., 2019, Burtsev et al., 2020); recent long-context systems couple local attention with episodic or parametric memory to improve scaling and robustness at 16k–256k tokens (Chaudhury et al., 20 Feb 2025, Wang et al., 14 Feb 2026).

1. Formal structure and conceptual scope

At its most general, self-attention computes

$$\mathrm{Attn}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V.$$

Memory augmentation modifies this template in one of two ways. The first is to enlarge the key/value set by concatenating memory entries to the sequence-derived keys and values. The second is to compute a separate memory-derived signal and use it to reweight token-level self-attention. Both retain differentiable, content-based addressing, but they differ in where the memory enters the computation.

A useful technical distinction is between memory as additional attention targets and memory as an additional control signal. Persistent key–value memory in attention-only Transformers exemplifies the first pattern: token queries attend jointly over contextual states and a bank of learnable memory vectors (Sukhbaatar et al., 2019). Episodic attention in long-context decoders exemplifies the second: chunk-level relevance scores are broadcast to tokens and used to reweight attention into the stored KV cache (Chaudhury et al., 20 Feb 2025).

Family	Memory carrier	Representative form
Recurrent-output memory	Past hidden states	$Y_t=[h_{t-L}\cdots h_{t-1}]$ (Daniluk et al., 2017)
Persistent memory	Learnable key/value vectors or [mem] tokens	Concatenated to keys/values or prepended to the sequence (Sukhbaatar et al., 2019, Burtsev et al., 2020)
Segment or episodic memory	Chunk summaries or chunk embeddings	Attention over a memory bank of prior segments or chunks (Wu et al., 2020, Chaudhury et al., 20 Feb 2025)
Parametric memory	Fast weights updated online	Local attention plus a non-linear TTT memory network (Wang et al., 14 Feb 2026)

This scope is broader than a single architecture family. Memory-augmented self-attention appears in recurrent LLMs, encoder-only and decoder-only Transformers, streaming speech recognizers, video segmentation networks, reinforcement-learning agents, and hybrid convolution-attention systems.

2. Memory carriers, read–write paths, and memory policies

A precise way to analyze these models is to decompose them into encoding, consolidation/storage, and retrieval. That decomposition is stated explicitly for memory-augmented Transformers in work linking model design to human memory, which also distinguishes static memory from dynamic memory (Raccah et al., 2022). Static memory is fixed content used only via retrieval; dynamic memory is updated as new input is processed and therefore requires write and forgetting policies.

Recurrent-output memory is the simplest dynamic case. In a LLM built on an LSTM, the memory at time $t$ can be a sliding window of previous outputs,

$Y_t=[h_{t-L}\cdots h_{t-1}],$

with attention performed over that window. The key–value–predict factorization introduced in this setting separates address vectors $k_t$, stored content $v_t$, and a prediction vector $p_t$, so that lookup, stored content, and next-word prediction no longer share a single overloaded representation (Daniluk et al., 2017).

Persistent-memory Transformers replace dynamic writes with learned global parameters. One formulation augments each head with persistent key and value matrices $\mathbf M_k,\mathbf M_v$, so every query can attend jointly over token-derived keys/values and these memory vectors (Sukhbaatar et al., 2019). A closely related design prepends trainable [mem] tokens to the sequence, turning memory into a low-dimensional global workspace that is updated layer by layer along with ordinary token states; a stricter bottleneck variant forces non-local information to pass through those memory tokens (Burtsev et al., 2020).

Segment-level memory introduces explicit write operations at chunk boundaries. In streaming acoustic modeling, the sequence is divided into short segments with left and right context, and each processed segment writes a summary vector into a memory bank $(\mathbf m_1,\dots,\mathbf m_{n-1})$. Queries for the current segment attend jointly over the current segment and that bank of memories, making all processed segments directly accessible in one hop (Wu et al., 2020). A different streaming design augments restricted self-attention with an LSTM memory state $h_t$, so that local self-attention over $Y_t=[h_{t-L}\cdots h_{t-1}]$0 is combined with a recurrent summary of the entire past; here memory is not a bank of slots but a causal state trajectory (Luo et al., 2021).

Recent long-context LLM systems have pushed the separation between token memory and higher-level memory further. EpMAN stores chunk embeddings as an episodic memory and uses query-to-chunk similarity to produce chunk-level attention, which is then broadcast to tokens and used to reweight decoder self-attention into the KV cache (Chaudhury et al., 20 Feb 2025). AllMem replaces full global attention with a hybrid token mixer,

$Y_t=[h_{t-L}\cdots h_{t-1}]$1

where the memory path is a non-linear test-time-trainable fast-weight network updated online from the key/value stream (Wang et al., 14 Feb 2026).

3. Representative architectural realizations

In language modeling, memory augmentation originally arose as an attempt to extend prediction beyond immediate recurrent state. Attention over recent outputs, key–value attention, and key–value–predict attention all instantiate the same principle: past states are treated as a differentiable memory queried by the present state (Daniluk et al., 2017). Persistent-memory Transformers later reframed the feed-forward layer itself as a kind of attention over global basis vectors and made those vectors explicit as persistent key–value memory, yielding an all-attention architecture without a separate FFN (Sukhbaatar et al., 2019). Memory tokens generalized the same idea by turning memory into extra sequence positions that can be read from, written to, or even used as a bottleneck for global information flow (Burtsev et al., 2020).

Speech recognition developed several distinct memory-augmented self-attention variants. One line augments self-attention layers with persistent memory vectors that encode information “beyond the whole utterance level,” including both key–value memory and input-embedding memory structures inside DFSMN-SAN encoders (You et al., 2019). A second line targets streaming ASR by segmenting the acoustic sequence and attending over a bank of segment memories in addition to a short local segment, thereby preserving streamability while recovering direct access to distant processed segments (Wu et al., 2020). A third line combines restricted self-attention with recurrent memory inside an RNN-T-style transducer, using local self-attention for nearby structure and an LSTM memory state for long-range causal context (Luo et al., 2021).

Vision and video models instantiate memory differently but preserve the same addressing logic. TMANet stores encoded features from several past video frames in an external temporal memory and lets current-frame queries attend over all pixels in all memory frames, replacing optical flow with content-based memory reads (Wang et al., 2021). EAANet does not introduce a large external memory bank across images, but it adopts memory-based efficient attention mechanisms in which Longformer-style global tokens act as static memory slots and permit distant regions to communicate at near-linear cost (Zhang et al., 2022).

In reinforcement learning, memory-augmented self-attention can be realized as a sliding window of past observations rather than a separate external matrix. A multi-head dot-product attention module operating over semantic feature vectors across time functions as a transient working memory: the same Q/K/V mechanism both selects relevant features and retrieves from a temporal buffer of prior observations (Bramlage et al., 2020). This makes memory use directly inspectable through attention matrices rather than latent recurrent states.

4. Empirical behavior and effective memory use

A recurrent finding is that adding memory does not guarantee that long-range memory will actually be used. In neural language modeling, attention over differentiable memory was found to place almost all mass on the five most recent positions, and increasing the memory window beyond small values produced no significant perplexity improvement. On the Wikipedia corpus, the Key–Value–Predict model obtained test perplexities of 76.1 with window size 1, 75.8 with size 5, 76.0 with size 10, and 75.8 with size 15; a simple 4-gram RNN reached 75.9, essentially matching the attentive model’s 75.8 (Daniluk et al., 2017). This result established that representational separation can help performance without materially extending the effective attention span.

In video segmentation, memory augmentation is more visibly non-local. TMANet constructs memory from several past frames and uses current-frame queries to attend over all memory pixels. With a ResNet-50 backbone it reports 80.3% mIoU on Cityscapes and 76.5% mIoU on CamVid, and its ablations show that increasing the memory from 4 to 6 frames yields almost no additional gain, indicating saturation rather than monotonic benefit from longer memory (Wang et al., 2021).

In reinforcement learning, attention over a temporal buffer acts as an explicit working memory rather than an implicit one. A multi-head attention-over-time agent solves a partially observable T-maze and converges slightly faster than a two-layer LSTM baseline, whereas a single-head temporal attention model performs poorly on the same task. Because inputs are semantic features rather than pixels, the attended stimuli are directly interpretable and reveal head specialization over both features and time (Bramlage et al., 2020).

Speech results show that persistent or segmental memory can materially improve recognition. In large-vocabulary continuous speech recognition, adding persistent memory to DFSMN-SAN yields a further 5% to 11% relative improvement in CER over the already stronger DFSMN-SAN baseline (You et al., 2019). For online ASR, adding recurrent memory to restricted self-attention improves WER relative to restricted-self-attention baselines by 13.5 on WSJ and 7.1 on SWBD, without much computation-cost increase (Luo et al., 2021). In streaming Transformer acoustic models, augmented memory self-attention outperforms existing streamable Transformer methods and achieves over 15% relative error reduction compared with the LC-BLSTM baseline on LibriSpeech (Wu et al., 2020).

Long-context LLMs show a different empirical pattern: memory is valuable when it changes how attention is distributed across vast context windows. EpMAN, which reweights decoder self-attention using chunk-level episodic attention, is reported to be stronger and more robust from 16k to 256k tokens than baseline decoders trained with self-attention and popular retrieval-augmented generation frameworks (Chaudhury et al., 20 Feb 2025). AllMem, which replaces global attention with sliding-window attention plus non-linear TTT memory, reports a 0.83 average drop relative to full attention on 37k LongBench with a 4k window, and at 128k context on InfiniteBench its 8k-window variant outperforms full attention (Wang et al., 14 Feb 2026).

5. Long-context scaling, efficiency, and the distinction between model memory and device memory

Memory augmentation is often motivated by the quadratic compute and storage cost of full self-attention, but the term “memory” covers two different issues. One is representational memory: persistent vectors, episodic banks, segment summaries, or fast weights that store information for later retrieval. The other is device-memory efficiency: restructuring exact attention so that longer sequences fit in accelerator memory.

Representationally, the modern long-context trend is to combine precise local attention with a compact global memory. AllMem exemplifies this design: sliding-window attention supplies accurate local token interactions, while a fixed-size parametric memory supplies long-range information independent of sequence length. At 128k context, the paper reports that AllMem’s memory path uses about 1/9 the FLOPs and KV cache of full attention, while the cache size of the memory path does not grow with context length (Wang et al., 14 Feb 2026). EpMAN takes a different route: it retains the full KV cache but overlays it with a chunk-level episodic relevance distribution, improving which parts of the cache the model actually uses (Chaudhury et al., 20 Feb 2025).

Vision work often frames the same trade-off as memory-efficient attention rather than explicit long-term storage. EAANet replaces full self-attention in an attention-augmented ConvNet with Longformer- or Linformer-style efficient attention. In that setting, Longformer global tokens are interpretable as static memory slots, and the empirical point is that memory scales much better with input size than normal self-attention, allowing higher-resolution images than AA-Net with full attention (Zhang et al., 2022).

Implementation-level work makes the representational/device distinction explicit. Exact self-attention need not require $Y_t=[h_{t-L}\cdots h_{t-1}]$2 device memory: one can compute attention by streaming over keys and values and keeping only running summaries, obtaining $Y_t=[h_{t-L}\cdots h_{t-1}]$3 memory for a single query, $Y_t=[h_{t-L}\cdots h_{t-1}]$4 memory for self-attention in the theoretical sequential formulation, and a practical accelerator implementation with $Y_t=[h_{t-L}\cdots h_{t-1}]$5 memory. At sequence length 16384, this reduces attention memory overhead by 59X for inference and 32X for differentiation (Rabe et al., 2021). This line of work changes the scheduling of attention computation, not the representational role of memory inside the model.

6. Limitations, misconceptions, and open directions

A recurring misconception is that larger memory automatically yields longer effective context. The strongest counterexample is the early language-modeling evidence that differentiable memory was used almost entirely as a five-step context buffer, and that a fixed-window concatenation baseline could match the attentive model (Daniluk et al., 2017). That result does not imply that memory augmentation is unhelpful; it implies that usefulness depends on the memory representation, the optimization problem, and the evaluation target.

Another challenge is that the memory interface itself can degrade. Work on attention factorization through shared memory reports that direct input↔memory attention is suboptimal, identifies a phenomenon termed memory degradation, and shows that performance may be substantially improved by filtering the input signal before it communicates with memory (Yorsh et al., 2024). This is an architectural warning: bottleneck memories are not automatically efficient or expressive merely because they reduce attention cost.

The choice of memory signal is also contested. In cognitively inspired work on memory-augmented Transformers, raw attention weights had essentially no predictive power for human memory except a very weak effect in one GPT-2 layer, whereas surprisal correlated positively with human memory effects and is presented as a more plausible encoding-policy variable (Raccah et al., 2022). This complicates any simplistic use of attention magnitude as a proxy for what should be written to memory.

Results on pretrained models are similarly mixed. Augmenting a pretrained masked LLM with memory tokens in GLUE fine-tuning shows mixed results, even though the same paper reports positive correlations between memory and performance for machine translation and language modeling (Burtsev et al., 2020). Memory augmentation therefore appears to be highly task-dependent: it is consistently beneficial in some settings, only weakly helpful in others, and sometimes beneficial only when the training objective explicitly teaches the model how to use it.

Current research directions follow from these limitations. Proposed remedies include forcing or encouraging attention over longer history rather than allowing purely local fallbacks, using event boundaries or surprisal as encoding signals, refining taxonomies of static and dynamic memory, integrating short-term parametric memory with external persistent memory systems, and analyzing what exactly fast-weight or episodic memory encodes during long-context generation (Daniluk et al., 2017, Raccah et al., 2022, Wang et al., 14 Feb 2026). Together, these lines of work suggest that memory-augmented self-attention is not a single mechanism but a design space: the central problem is no longer whether attention can access memory, but how memory should be represented, updated, constrained, and made genuinely useful.