PaperlayerView exact revision on arXiv

Source-linked AI summary

Cognitive Memory in Large Language Models

Lianlei Shan, Shixian Luo, Zezhou Zhu, Yu Yuan, Yong Wu

arXiv:2504.02441v2cs.CLcs.AI

TL;DR

LLMs can retrieve and summarize information but lack stable, structured memory that reliably persists over time. This paper surveys memory architectures and mechanisms, including text-based acquisition, management, and retrieval. It identifies Lamina’s higher throughput per dollar than traditional solutions as one reported result and discusses memory parameterization as a future direction.

  • Problem

    LLMs can retrieve and summarize information, but they lack stable, structured memory that reliably endures over extended periods.

    Source: §1 · p. 4

  • Method

    The paper surveys memory types and mechanisms, including text-based memory acquisition, management, and utilization.

    Source: §1 · p. 4, §3 · p. 9

  • Results

    Lamina achieves 1.48 to 12.1 times higher throughput per dollar than traditional solutions using attention offload.

    Source: §4.3.1 · p. 26

  • Takeaways & Limitations

    The paper presents memory parameterization as a promising direction for addressing the storage and retrieval limits of current memory methods.

    Source: §7 · p. 33

    Locality-sensitive hashing provides approximate retrieval results that may contain errors, though parameter choices can reduce them to some extent.

    Source: §3.3.5 · p. 16

Abstract

from arXiv · show

This paper examines memory mechanisms in Large Language Models (LLMs), emphasizing their importance for context-rich responses, reduced hallucinations, and improved efficiency. It categorizes memory into sensory, short-term, and long-term, with sensory memory corresponding to input prompts, short-term memory processing immediate context, and long-term memory implemented via external databases or structures. The text-based memory section covers acquisition (selection and summarization), management (updating, accessing, storing, and resolving conflicts), and utilization (full-text search, SQL queries, semantic search). The KV cache-based memory section discusses selection methods (regularity-based summarization, score-based approaches, special token embeddings) and compression techniques (low-rank compression, KV merging, multimodal compression), along with management strategies like offloading and shared attention mechanisms. Parameter-based memory methods (LoRA, TTT, MoE) transform memories into model parameters to enhance efficiency, while hidden-state-based memory approaches (chunk mechanisms, recurrent transformers, Mamba model) improve long-text processing by combining RNN hidden states with current methods. Overall, the paper offers a comprehensive analysis of LLM memory mechanisms, highlighting their significance and future research directions.

1 Introduction

LLMs can retrieve and summarize information, but they lack stable, structured memory that reliably persists over extended periods. This paper surveys memory types and mechanisms, then considers cognitive memory as a future direction.

1 Introduction

  • The paper distinguishes memory from profiling and knowledge: it is a dynamic repository of experiences that engages in decision-making.Profiling concerns an agent’s identity, functions, and operating context, while knowledge supplies factual information or learned representations.Source: §1 · p. 4
  • Memory can support context-rich responses, reduce hallucinations by anchoring answers in stored facts, and streamline data retrieval and processing.The paper also presents long-term memory as a potential basis for adapting to new tasks with limited data or interactions.Source: §1 · p. 4
  • The paper surveys memory types and roles in agentic workflows, how models remember information, and generative AI’s impact on memory.It proceeds from cognitive architecture to current LLM-memory research and future trends in cognitive memory.Source: §1 · p. 4

2 Cognitive Architecture

The section surveys cognitive architectures and uses sensory, short-term, and long-term memory as its organizing categories. It contrasts human memory’s adaptive, context-sensitive qualities with LLMs’ reliance on learned patterns and input context.

2 Cognitive Architecture

  • Cognitive architectures model aspects of human cognition, including memory, perception, attention, and decision-making, and have been studied through psychological experiments.Soar supports multiple learning mechanisms; ACT-R models working and long-term memory; ICARUS emphasizes autonomous learning and adaptation.Source: §2 · pp. 4–5

2.1 Categories of Memories

  • By retention duration, memory is categorized as iconic, short-term, or long-term; other classifications concern content and conscious involvement.The section also describes iconic, logical, emotional, and procedural memory, as well as explicit and implicit memory.Source: §2.1 · p. 5

2.1.1 Sensor Memory

  • In LLM systems, sensory memory corresponds to an API request or prompt, analogous to humans’ fleeting sensory input.Human sensory information is either discarded or passed to short-term memory after the first few seconds of perception.Source: §2.1.1 · p. 5

2.1.2 Short-Term Memory

  • LLM short-term memory is the text or embeddings processed within the immediate context window, and the model usually loses that context when the session ends.The section relates working memory to holding multiple inputs during planning and other multi-step tasks.Source: §2.1.2 · p. 6

2.1.3 Long-Term Memory

  • LLM long-term memory can use external databases, vector stores, or graph structures to keep relevant information available for future queries.The section relates human long-term memory to enduring knowledge, skills, and experiences, including declarative and procedural forms.Source: §2.1.3 · pp. 7–8

2.2.1 Human Memory

  • Human memory includes sensory, short-term, and long-term forms, and long-term memories can be modified or reconstructed through experience and context.The section describes sensory traces as fleeting, short-term memory as temporary storage, and long-term memory as including declarative and procedural forms.Source: §2.2.1 · p. 8

2.2.2 LLMs Information

  • LLMs encode statistical language patterns in parameters and generate responses from prompts rather than retaining explicit records of prior interactions.Adaptation or learning occurs through retraining on new data, unlike human memory’s continuous experiential change.Source: §2.2.2 · p. 8

2.2.3 Parallels and Differences

  • Humans and LLMs both use context, but human memory changes with experience while LLMs’ learned knowledge remains fixed until retraining.LLMs lack adaptive forgetting, and apparent forgetting in long conversations can result from token-length limits rather than a cognitive process.Source: §2.2.3 · p. 9

3 Text-based Memory

Text-based memory covers selecting or compressing historical information, managing stored memories, and retrieving relevant memories for current queries.

3 Text-based Memory

  • Text-based memory spans acquisition, management, and utilization: storing selected or compressed history, maintaining it, and retrieving relevant memories from a query.Retrieval methods include full-text, SQL-based, semantic, tree-based, hash-based, multi-pass, and hybrid search.Source: §3 · pp. 9–10

3.1 Memory Acquisition

  • Acquisition can retain all or selected conversation content, summarize it, or combine selection and summarization.The approaches range from preset filtering to summaries of conversation-level events and user information.Source: §3.1 · p. 10, §3.1.1 · p. 10, §3.1.2 · p. 10

3.1.2 Memory Summarization

  • Examples include COMEDY’s edited conversation summaries and Memory Bank’s daily-event and global summaries with evolving personality insights.Memory Bank aggregates daily personality insights into a global understanding used to tailor responses to user traits.Source: §3.1.2 · pp. 10–11

3.2.1 Memory Update

  • Update strategies include LRU eviction, forgetting-curve decay, recursive conversation-summary updates, and integration of new and old knowledge.Memory Bank models retention with R = e^-tS, where recall increases memory strength and resets elapsed time.Source: §3.2.1 · p. 12

The Strategy of Integrating New and Old Knowledge

  • TiM organizes ideas in hash-indexed groups through insertion, forgetting, and merging of similar ideas.Forgetting can remove unnecessary or contradictory ideas, while merging combines similar ones.Source: The Strategy of Integrating New and… · p. 12

3.2.2 Memory Access

  • A Working Memory Hub coordinates inputs, outputs, and interaction history; storage options include trees, tables, hashes, vectors, triples, and graphs.Its Interaction History Window provides short-term caching and context anchoring, while the Episodic Buffer retrieves complete scenarios beyond token limits.Source: §3.2.2 · pp. 12–13, §3.2.3 · p. 13

3.2.4 Handling Contradictory Memories

  • Conflicting memories can be merged using source context, disambiguated by retrieving context, or retained when an NLI model may have misjudged a contradiction.These strategies allow context-dependent interpretations to coexist when appropriate.Source: §3.2.4 · p. 14

3.3 Memory Utilization

  • Retrieval spans full-text, SQL, and semantic search as well as tree and hash traversal and repeated searches that refine a query.Tree search uses GPT to aggregate child text and choose traversal actions; hash search uses query embeddings and LSH to find similar memories.Source: §3.3.1 · p. 15, §3.3.2 · p. 15, §3.3.3 · p. 15

4 KV Cache-Based Memory

The KV-cache chapter organizes memory around acquisition, management, and retrieval, with acquisition divided into KV selection and compression.

Acquisition

  • KV acquisition selects key-value pairs to retain and compresses the cache through low-rank methods, merging, multimodal compression, quantization, or multi-level compression.Selection chooses some KVs to save and discards the rest.Source: §4 · pp. 16–17, Figure 6 · p. 17

Management

  • KV-cache management includes offloading to cheaper memory, OS-related optimization, storage structures, and shared attention.Source: §4 · p. 16

Utilization

  • KV-cache retrieval covers efficient attention, locality-sensitive hashing, vector retrieval, and other optimization methods.Source: §4 · p. 17

4.1 Memory Acquisition: KV Selection

KV-selection approaches use regularities in attention and token values to decide which tokens to retain, with methods ranging from cache eviction to broader attention designs.

4.1.1 Strategy of Regularity

  • D2O uses layer-specific attention-density observations and a compensation mechanism to reconsider discarded tokens and recall them when they meet a similarity threshold.The method is presented as optimizing cache size without fine-tuning while retaining key context.Source: §4.1.1 · p. 17
  • Other regularity-based approaches exploit persistent attention patterns, recent-query similarity, attention sinks, or structured long-range attention.StreamingLLM retains initial attention-sink KVs alongside sliding-window KVs; LONGNET uses dilated attention, while ETC restricts long-input tokens to global inputs or local neighbors.Source: §4.1.1 · pp. 17–18

4.1.2 Value Vector Norm and Entropy Strategy

  • VATP evaluates token importance using both attention scores and the norms of value vectors.Source: §4.1.2 · p. 18

4.1.3 Score Strategy

Score-based KV selection includes methods that adjust accumulated attention scores, estimate query-specific page importance, and retain tokens with high entropy or learned scores.

4.1.3 Score Strategy

  • A2SF discounts older accumulated attention scores with a forgetting factor to compare tokens of different ages more fairly.The approach addresses uneven accumulation in decoder models caused by masking.Source: §4.1.3 · p. 19
  • Quest uses query-dependent upper bounds on page attention to select top-K pages for normal self-attention, reducing memory movement.Its estimate uses channel-wise minima and maxima of key vectors; the selected-page token count is the token budget.Source: §4.1.3 · p. 19
  • Keyformer retains key tokens selected by a scoring function, while entropy-based selection retains tokens with higher entropy values.Keyformer’s score adds Gumbel noise to unnormalized logits and uses a temperature parameter.Source: §4.1.3 · pp. 19–20

4.1.4 Special Token Embedding Strategy

The section describes KV-cache and attention methods that select, retain, or summarize tokens to reduce memory and computation. Approaches include head-specific caching, chunk-summary tokens, learned pruning, and differentiable sparse selection.

4.1.4 Special Token Embedding Strategy

  • RazorAttention preserves full KV caches for retrieval heads and discards remote tokens in other heads, with compensation tokens to recover information.The method is training-free and motivated by observations that most heads focus locally while a small number attend broadly.Source: §4.1.4 · p. 20
  • Special marker tokens let attention control access to consecutive text chunks, while <SR> tokens aggregate information from each chunk.Mohtashami and Jaggi [2023] use markers to represent chunks; Luo et al. [2024] insert <SR> and modify the attention mask to integrate chunk information.Source: §4.1.4 · p. 20

4.1.5 Learning-based Token Selection Strategy

  • Learned Token Pruning uses per-layer thresholds on tokens’ average attention probability, with a differentiable soft mask enabling threshold learning.The method can adaptively change the pruned sequence length.Source: §4.1.5 · p. 20
  • SPARSEK Attention selects a fixed number of key-value pairs per query using learned sparse patterns; SparseK relaxes selection to a k-sum constraint for differentiable top-k.Its soft-thresholding solution performs selection in the forward pass while preserving gradient updates in the backward pass.Source: §4.1.5 · pp. 20–21
  • Parameterized pruning gates can discard context at each layer, while query-key-based signals determine whether earlier tokens influence the current token.The described gate allows dynamic context adjustment during generation, reducing unnecessary computation and memory use.Source: §4.1.5 · p. 21

4.1.6 Different layers and different heads employ distinct token selection strategies.

These methods vary token or cache selection across attention layers, heads, channels, and sequence blocks. Several use observed attention or cache structure to allocate memory selectively.

4.1.6 Different layers and different heads employ distinct token selection strategies.

  • SQUEEZEATTENTION reallocates cache budgets across attention layers, addressing optimization methods that focus on sequence or batch dimensions.Its layer-wise allocation is based on analyzing attention-layer behavior during inference.Source: §4.1.6 · p. 21
  • ThinK prunes query-dependent KV-cache channels, targeting redundancy associated with imbalanced magnitudes and low-rank attention weights.It aims to minimize attention-weight loss while selectively pruning the least important channels.Source: §4.1.6 · p. 21
  • Mixture of Attention profiles configurations across heads and layers, selecting sparse patterns and scaling rules for different input sizes.The analysis finds some heads broaden focus for longer sequences while others remain focused on fixed-length local contexts.Source: §4.1.6 · p. 22
  • Zebra alternates grouped local- and global-attention layers to reduce the computational and memory demands of full attention.The architecture balances local and global attention in grouped layers.Source: §4.1.6 · p. 22
  • LONGHEADS selects important context blocks within the pretrained length for each head, using the model’s dot-product attention to avoid processing tokens beyond that length.The method is presented as extending long-context handling without additional training.Source: §4.1.6 · p. 22
  • SnapKV uses an end-of-prompt observation window and voting across attention heads to select important KV pairs, then appends them to that window in a smaller cache.Pooling-based clustering captures local information and filters irrelevant tokens before voting.Source: §4.1.6 · p. 22
  • PyramidKV allocates more cache to lower layers and less to higher layers, following the reported progression from dispersed attention to focus on key tokens.The account describes attention as broadly dispersed at lower layers, integrating within contexts, then focusing on key tokens at higher layers.Source: §4.1.6 · p. 22
  • Adaptive budget allocation matches pruning budgets to head characteristics, addressing quality loss associated with evenly distributing budgets across heads.The method is described as optimizing a theoretical L1-loss upper bound and reducing L1 pruning loss in practice.Source: §4.1.6 · p. 22

4.1.7 Different layers and different heads use the same token selection strategy

The section covers cache compression and attention sharing approaches that apply common strategies across model components. They exploit latent representations, low-rank structure, key norms, or cross-layer similarity.

4.1.7 Different layers and different heads use the same token selection strategy

  • DeepSeek-V2 compresses the KV cache into latent vectors with Multi-head Latent Attention and uses DeepSeekMoE for sparse computation.The passage describes these components as supporting efficient inference and economical model training, respectively.Source: §4.1.7 · p. 22
  • Yu et al. [2024] compress KV heads by optimizing the transition from MHA to GQA and using a specialized key-cache strategy compatible with RoPE.The approach is motivated by low-rank characteristics of the KV cache and seeks to minimize compression error.Source: §4.1.7 · pp. 22–23
  • Devoto et al. [2024] retain keys with the lowest L2 norms and their corresponding values, based on a reported correlation between key norms and attention scores.The strategy requires no additional training or substantial model modifications and is directly applicable to Transformer-based decoder models.Source: §4.1.7 · p. 23
  • LISA reduces redundant cross-layer attention computation by sharing similar attention weights, using head alignment and difference compensation modules.Experiments reportedly find similarity between attention weights across layers in most cases.Source: §4.1.7 · p. 23
  • Layer-compressed KV caching reduces memory consumption and improves inference speed by decreasing the number of cached layers.Source: §4.1.7 · p. 23

4.1.8 Locality-Sensitive Hashing Strategy

The section describes hashing-based and dynamic sparse attention for reducing long-sequence computation, alongside backtracking that retains tokens influential to final predictions.

4.1.8 Locality-Sensitive Hashing Strategy

  • Reformer uses locality-sensitive hashing to reduce self-attention complexity from O(L²) to O(LlogL), alongside reversible residual layers and feed-forward chunking to reduce memory use.Its hashing attends among keys assigned to the same bucket.Source: §4.1.8 · p. 23
  • Sparse Flash Attention uses dynamic sparse structure and hash algorithms to reduce computation for large sequences.Source: §4.1.8 · p. 23

4.1.9 Backtracking Strategy

  • Attention backtracking tracks attention importance from output to input to retain tokens that significantly affect the final prediction.The passage notes that forward-phase token pruning can discard tokens important to subsequent layers.Source: §4.1.9 · p. 23

4.2 Memory Acquisition: KV Compression

The section surveys KV-cache compression through low-rank methods, merging similar states, and hierarchical memory updates across input segments.

4.2.1 Low-Rank Compression

  • Palu [2024] compresses hidden dimensions with low-rank projection, while LESS [2024] combines sparse KV selection with a fixed-size low-rank residual cache.Palu caches compressed intermediate states and reconstructs keys and values; LESS lets queries access information discarded by sparse selection.Source: §4.2.1 · pp. 23–24

4.2.2 KV Merging

  • KVMerger [2024a] identifies mergeable KV states and combines each set using Gaussian-kernel-weighted merging.The approach is motivated by token-level key-state similarity and reported dataset-agnostic KV-cache sparsity.Source: §4.2.2 · p. 24

4.2.3 Multi-Level KV Compression

  • Hierarchical Memory Transformer [2024] encodes each segment locally, uses attention to update global memory, then compresses that memory as new segments arrive.It recalls relevant history when producing each segment's output and concatenates segment outputs into the final representation.Source: §4.2.3 · p. 25

4.3 Memory Management

KV-cache management covers offloading, operating-system-inspired memory systems, shared attention, and the relationship between cache compression and attention.

4.3.1 Offload

  • 1.48 to 12.1 times higher throughput per dollar than traditional solutions is reported for Lamina's attention-offload strategy.The strategy assigns attention computation to memory-optimized devices while reserving high-end accelerators for other work.Source: §4.3.1 · p. 26
  • He and Zhai [2024] use CPU resources across multiple nodes for memory-bound KV-cache access, reducing transfer overhead and increasing GPU throughput on the model's other part.Their motivation is that KV caches limit batch size and CPU–GPU bandwidth bottlenecks ordinary host-memory offloading.Source: §4.3.1 · p. 26

4.3.2 Memory Management with OS System

  • PagedAttention [2023] applies paging ideas to KV-cache allocation, targeting near-zero cache-memory waste and sharing across requests in vLLM.The motivating issue is that dynamically growing and shrinking per-request caches can fragment memory and limit batch size.Source: §4.3.2 · p. 26

4.3.3 Shared Attention

  • CacheBlend [2024a] reuses non-prefix text-chunk caches by recomputing a small subset of token values, while Shared Attention [2024] shares computed attention weights.CacheBlend reports the same generation quality as full prefill; Shared Attention targets both cache size and inference computation.Source: §4.3.3 · p. 27

4.3.4 Relationship between KV Cache Compression and Attention

  • Analysis of softmax attention [2023] proves directional convergence toward a maximum-margin solution that distinguishes locally optimal from non-optimal tokens.The authors characterize token optimality using value embeddings and problem geometry.Source: §4.3.4 · p. 27

5 Parameters-Based Memory

The section treats parameterization as a distinct memory strategy: previous memories become model parameters, with LoRA, test-time training, and mixture-of-experts approaches covered.

5 Parameters-Based Memory

  • Parameter-based memory turns previous memories into model parameters, distinguishing it from text-based and KV-cache storage.The section groups the strategies as LoRA, Test Time Training, and Mixture of Experts.Source: §5 · p. 28

5.1 LoRA For Memory Parameterization

  • LoRA-based context compression [2023] compresses accumulated attention key/value pairs and reports matching full-context performance with 5× less context memory.The method integrates conditional LoRA during inference without fine-tuning the full model weights.Source: §5.1 · p. 28

5.2 Test Time Training For Memory Parameterization

  • Test Time Training adapts a model to test data by updating shared feature-extractor parameters using an auxiliary-task loss.The described update occurs for each test sample; the method is presented for settings where test and training distributions differ.Source: §5.2 · pp. 28–29

5.3 MoE For Memory Parameterization

  • BTX [2024] trains expert models on separate data subsets, merges their feedforward layers into an MoE, then fine-tunes the router on the full training set.Its expert training is parallel and asynchronous, reducing communication costs; sparse activation keeps inference FLOPs from increasing significantly despite more parameters.Source: §5.3 · p. 29

6 Hidden-State-Based Memory

Hidden-state-based memory combines RNN hidden states with other methods to process long texts, spanning chunk mechanisms, recurrent transformers, and Mamba. These approaches use chunking, recurrent state, or selective state-space processing to handle sequence information.

6.1 Chunk Mechanism

  • Chunking divides long text into semantically coherent segments, embeds each separately, and may overlap boundaries to preserve context for retrieval or recommendation.Segments can be based on character counts, sentences, paragraphs, or semantics.Source: §6.1 · p. 30
  • Cascaded sub-buffers retain important tokens evicted from earlier buffers, while exponential moving-average attention scores guide token retention.The buffers preserve the same total cache size; lower-scoring tokens are discarded.Source: §6.1 · p. 30
  • AutoCompressor summary vectors encode useful document information, support downstream tasks and retrieval or reranking, and reduce context-learning inference cost.The vectors are generated from long documents and used to improve language modeling for subsequent passages.Source: §6.1 · p. 31
  • Other chunk-based approaches include MEGABYTE’s patches with local and global models, and ChunkAttention’s prefix-based KV caching and two-stage attention.MEGABYTE is described as enabling subquadratic self-attention; ChunkAttention dynamically detects and removes redundant KV caches.Source: §6.1 · p. 31

6.2 Recurrent Transformer

  • Recurrent Transformers combine recurrent long-sequence information retention with Transformer parallel processing and global dependency capture; RWKV is described as efficiently processing long text and capturing long-term dependencies.RWKV is also described as avoiding vanishing or exploding gradients and offering parallel computing capabilities.Source: §6.2 · pp. 31–32

6.3 Mamba

  • Mamba combines a Transformer MLP for channel mixing, an SSM block for sequence modeling, and an input-dependent selective mechanism.Its workflow expands and discretizes the input, then computes state updates and outputs in recursive or convolutional forms.Source: §6.3 · pp. 32–33

7 Discussion

Hidden-state and KV-cache memories last only for a session, while text-based and parameter-based memories can persist across sessions. The paper identifies human-like summarization, integration, and selective forgetting as important challenges for LLM memory.

7 Discussion

  • KV-cache and hidden-state memories vanish after a session, whereas text-based and parameter-based memories can endure across sessions.Source: §7 · p. 33
  • LLMs lack humans’ summarization, generalization, and memory integration abilities, and without active intervention retain prior information, potentially increasing storage burdens and slowing retrieval.Source: §7 · p. 33
  • Memory parameterization is a proposed direction for enabling summarization, knowledge integration, and selective forgetting, but should avoid compromising memory or substantially increasing model parameters.LoRA, MoE, TTT, and new modules are among explored methods.Source: §7 · p. 33

Related papers

Loading 2504.02441v2…