Learned Memory Tokens in Neural Models

Updated 2 July 2026

Learned Memory Tokens are trainable, context-dependent vectors designed to store, retrieve, and adapt information within deep neural models.
They integrate as prepended tokens, latent memory banks, or recurrent carriers, enabling efficient fine-tuning, long-context inference, and scalable memory.
These tokens enhance model performance by reducing computational complexity while supporting applications in language, vision, and sequential decision-making.

Learned memory tokens are trainable, context-dependent vector representations designed for explicit information retention, retrieval, or task adaptation within deep neural models. Unlike standard token embeddings, which correspond to single vocabulary items and are fixed post-training, memory tokens encode variable-length sub-sequences or latent knowledge, can be dynamically generated or updated, and serve as dedicated channels for storing and transmitting information either across tasks, time steps, layers, or even entire datasets. Recent research demonstrates that learned memory tokens are central for scalable memory, efficient fine-tuning, long-context inference, continual learning, reversible data storage, and architectural interpretability, with implementations spanning vision transformers, LLMs, retrieval-augmented architectures, and sequence-to-sequence systems.

1. Mathematical Definitions and Architectural Integration

Formally, given a set of memory tokens $M = \{m_1, ..., m_{N_m}\}$ with $m_i \in \mathbb{R}^d$ , these vectors can be introduced either as (a) augmentations to the model input, (b) persistent latent vectors queried via cross-attention, (c) sequence-state carriers updated recurrently, or (d) compressed representations of historical context.

Several key paradigms exist:

Prepending tokens (autoregressive LLMs): A unique memory token (e.g., <MEMORY>) with a trainable embedding $M$ is added to the input sequence and optimized via autoregressive cross-entropy, yielding a reversible mapping from the token to an entire stored text. At inference, the model consumes only the learned $M$ , generating the memorized sequence exactly (Sastre et al., 17 Jun 2025).
Latent memory banks (transformers): A persistent memory bank $M \in \mathbb{R}^{N_m \times d}$ is queried at one or more layers by the model’s hidden state $H \in \mathbb{R}^{B \times L \times d}$ , using cross-attention:

$\mathrm{MemRead}(H, M) = \mathrm{softmax}\left(\frac{QK^\top}{\sqrt{d_h}}\right) V$

where $Q = \mathrm{Norm}(H)W_Q$ , $K = MW_K$ , and $V = MW_V$ . A routing mechanism may select a sparse relevant subset from a large memory bank for computational tractability (Tibrewal et al., 22 Mar 2026).

Recurrent memory (sequential decision models): Memory tokens are included within the input at each step $m_i \in \mathbb{R}^d$ 0, passed recurrently via self-attention, and updated either by direct overwrite or exponential moving average:

$m_i \in \mathbb{R}^d$ 1

with $m_i \in \mathbb{R}^d$ 2 as the new candidate state (Cherepanov et al., 10 Jun 2026).

Chunk-wise compression (scalable memory): A sequence is divided into fixed-size chunks. Each chunk is encoded to a memory token using an encoder $m_i \in \mathbb{R}^d$ 3, and these tokens are fed to a separate decoder (possibly in parallel), drastically reducing memory and compute complexity (Badger, 13 Feb 2026, Chen et al., 6 Mar 2026).

2. Training Objectives and Information Retention

The precise learning objective depends on the architectural context:

Autoregressive memory training: The memory token is optimized for a fixed target sequence $m_i \in \mathbb{R}^d$ 4 under the cross-entropy loss

$m_i \in \mathbb{R}^d$ 5

while all model weights $m_i \in \mathbb{R}^d$ 6 are frozen except $m_i \in \mathbb{R}^d$ 7 (Sastre et al., 17 Jun 2025).

Autoencoding and information-retention: To maximize recoverable information, an encoder produces memory tokens for input chunks, with a decoder trained to reconstruct the original tokens from just these summaries. Information retention is directly measured as token-level accuracy in reconstructing the input, often exceeding 90% in jointly-trained autoencoders, but remaining much lower in vanilla LLMs trained solely on next-token prediction (Badger, 13 Feb 2026).
Combined objectives: Hybrid training with both next-token and copy or blank-copy losses encourages the formation of memory tokens that are both causally useful for language modeling and information-rich for downstream decoding (Badger, 13 Feb 2026).
Sparse memory routing: For massive latent memory banks, a discrete router or softmax-based sparse selection ensures only a relevant subset of tokens is attended per sequence, enabling practical scaling to hundreds of thousands of memory vectors (Tibrewal et al., 22 Mar 2026).
Recurrent memory: In memory-augmented sequence models, memory tokens are updated across temporal steps using truncated backpropagation through time (TBPTT) or gradient-detached moving averages, with the sole loss computed on downstream prediction (e.g., action loss in VLA models) (Cherepanov et al., 10 Jun 2026).

3. Computation, Efficiency, and Scalability

The core motivation for learned memory tokens is to circumvent the $m_i \in \mathbb{R}^d$ 8 complexity of dense self-attention and to support memory beyond the intrinsic context-window of transformers.

Memory tokenization/compression: By dividing a length- $m_i \in \mathbb{R}^d$ 9 sequence into $M$ 0 chunks and encoding each to a single token, total attention cost is reduced from $M$ 1 to $M$ 2, balancing encoder and decoder costs for $M$ 3 and enabling sub-quadratic inference (Badger, 13 Feb 2026).
Memory sparse attention (MSA): Each document or sequence chunk is projected to a small set of mean-pooled keys/values (the memory tokens), which are then distributed across hardware for scalable retrieval. A lightweight router selects relevant subsets for any given query, facilitating inference over 100M tokens with $M$ 4 accuracy degradation (Chen et al., 6 Mar 2026).
Chapter-based memory (MoC): A global memory bank is partitioned into chapters. For each input, a learned router selects $M$ 5 chapters (e.g., $M$ 6 out of $M$ 7), greatly reducing both memory-access FLOPs and VRAM requirements, and enabling scaling to $M$ 8 tokens (Tibrewal et al., 22 Mar 2026).
Learned eviction and latent compression: In streaming LLM inference, a compact online-updated latent memory augments the bounded attention cache, allowing accurate recovery of evicted information while bounding GPU memory and compute (Yang et al., 25 May 2026).

4. Empirical Performance and Application Domains

Learned memory tokens exhibit distinct empirical benefits across multiple domains:

Domain	Memory Token Role	Key Result
LLMs (reversible memory)	Store/replay sequences via token emb.	Llama 3.1 8B achieves 100% seq. recall (up to ~240 tokens) (Sastre et al., 17 Jun 2025)
Transformers (MoC)	Persistent, learnable memory bank	Perplexity/QA accuracy exceeds iso-FLOP baselines, reduces catastrophic forgetting (Tibrewal et al., 22 Mar 2026)
Vision Transformers	Per-layer adaptation tokens	Parameter-efficient fine-tuning, preserves old-task performance (Sandler et al., 2022)
RL & VLA models	Recurrent state across steps	μVLA: success rate 0.84 vs. 0.42 for baseline; improves partial-observability (Cherepanov et al., 10 Jun 2026)
Scalable long-context	Chunkwise/latent tokenization	MSA: 8.8% degradation from 16K to 100M tokens, robust QA at ultra-long scale (Chen et al., 6 Mar 2026)
Latent memory in streaming	KV eviction, online compression	IndexMem: +30 points on RULER-16K over heuristics; near lossless retrieval (Yang et al., 25 May 2026)

These mechanisms underpin notable advances in model efficiency, continual learning, and robust performance as context length and knowledge scale increase.

5. Theoretical and Mechanistic Insights

Function tokens as emergent memory keys: In mainstream LLMs, a small number of high-frequency “function tokens” (e.g., punctuation, articles) act as implicit memory indices, activating the majority of latent features both during memory retrieval (inference) and consolidation (training). These tokens steer model computation, facilitate knowledge organization, and serve as natural control points for model behavior—an empirical finding termed the “function token hypothesis” (Zhang et al., 9 Oct 2025).
Non-invertibility of next-token objectives: Standard next-token prediction is inherently non-invertible for information retention, motivating hybrid or autoencoding objectives when high-fidelity memory formation is required (Badger, 13 Feb 2026).
Memory as addressable, interpretable parameter space: Persistent memory banks permit inspection and intervention at the token or chapter level, and explicit architectural separation between backbone and memory supports superior knowledge retention and reduced catastrophic forgetting (Tibrewal et al., 22 Mar 2026).
Curriculum and freezing for memory models: Staged training—autoencoding encoder pretraining, encoder freezing, then decoder adaptation—yields higher information recall and faster convergence in memory-based architectures (Badger, 13 Feb 2026).

6. Limitations, Open Problems, and Future Directions

Despite their effectiveness, learned memory tokens present several open challenges:

Computational intensity: Some training pipelines (e.g., reversible memory token fitting) require thousands of full-model backward passes per memory token, far from the efficiency of BERT-style pooling (Sastre et al., 17 Jun 2025).
Task specificity and generalization: Many mechanisms serve verbatim replay or adaptive recall but are not directly suited for question answering or compositional generalization without further model adaptation (Sastre et al., 17 Jun 2025, Badger, 13 Feb 2026).
Update and write mechanisms: Most persistent memories are read-only at inference; dynamic or episodic write capabilities remain an active area for model personalization and continual adaptation (Tibrewal et al., 22 Mar 2026).
Interpretability and security risks: Memory tokens can encode arbitrary content, including adversarial or harmful information, raising new questions about control, auditing, and interpretability (Sastre et al., 17 Jun 2025).
Scalability versus latency: While chunkwise compression and routing yield dramatic memory savings, per-chunk processing and router lookup can become a new bottleneck at extreme scales (Chen et al., 6 Mar 2026).

Key directions include token-level and hierarchical routing, dynamic memory writes, principled scaling laws for memory–performance trade-offs, combining latent memory with retrieval-augmented and external databases, and leveraging function-token insights for model steering and controllability.

Learned memory tokens complement and, in some cases, unify existing memory paradigms:

Latent persistent memory: Explicit associative memory banks in transformers generalize Mixture-of-Experts, sophisticated external memory (e.g., DNC), and chapter-based systems (Tibrewal et al., 22 Mar 2026).
Chunkwise autoencoding: Encoder–decoder memory models leverage neural summarization, parallelism, and compositionality beyond simple sliding-window or pooling (Badger, 13 Feb 2026, Chen et al., 6 Mar 2026).
Recurrent state and continual learning: Self-attention-based recurrence with learnable tokens outperforms both naive state passing and complex hierarchical memory in partially observable settings (Cherepanov et al., 10 Jun 2026).
Implicit memory routing: Function tokens act analogously to keys in a distributed key–value store, dynamically steering computation and memory retrieval (Zhang et al., 9 Oct 2025).
KV cache eviction and streaming memory: Efficient streaming architectures with learned eviction and compression reconstruct essential context under compute constraints (Yang et al., 25 May 2026).

This expanding repertoire establishes learned memory tokens as a core primitive for memory-intensive, scalable, and adaptive deep learning systems across vision, language, and sequential decision domains.