Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

FlashMem: Distilling Intrinsic Latent Memory via Computation Reuse

Yubo Hou, Zhisheng Chen, Tao Wan, Zengchang Qin

TL;DR

FlashMem is a framework that distills intrinsic memory directly from transient reasoning states via computation reuse, and identifies the last hidden state as a sufficient statistic for the interaction history, effectively bridging the gap between efficiency and persistent cognition.

Abstract

The stateless architecture of Large Language Models inherently lacks the mechanism to preserve dynamic context, compelling agents to redundantly reprocess history to maintain long-horizon autonomy. While latent memory offers a solution, current approaches are hindered by architectural segregation, relying on auxiliary encoders that decouple memory from the reasoning backbone. We propose FlashMem, a framework that distills intrinsic memory directly from transient reasoning states via computation reuse. Leveraging the property that internal representations uniquely encode input trajectories, FlashMem identifies the last hidden state as a sufficient statistic for the interaction history. This enables a Shared-KV Consolidator to synthesize memory by attending directly to the backbone's frozen cache, eliminating redundant re-parameterization. Furthermore, a parameter-free Cognitive Monitor leverages attention entropy to adaptively trigger consolidation only when high epistemic uncertainty is detected. Experiments demonstrate that FlashMem matches the performance of heavy baselines while reducing inference latency by 5 times, effectively bridging the gap between efficiency and persistent cognition.

FlashMem: Distilling Intrinsic Latent Memory via Computation Reuse

TL;DR

FlashMem is a framework that distills intrinsic memory directly from transient reasoning states via computation reuse, and identifies the last hidden state as a sufficient statistic for the interaction history, effectively bridging the gap between efficiency and persistent cognition.

Abstract

The stateless architecture of Large Language Models inherently lacks the mechanism to preserve dynamic context, compelling agents to redundantly reprocess history to maintain long-horizon autonomy. While latent memory offers a solution, current approaches are hindered by architectural segregation, relying on auxiliary encoders that decouple memory from the reasoning backbone. We propose FlashMem, a framework that distills intrinsic memory directly from transient reasoning states via computation reuse. Leveraging the property that internal representations uniquely encode input trajectories, FlashMem identifies the last hidden state as a sufficient statistic for the interaction history. This enables a Shared-KV Consolidator to synthesize memory by attending directly to the backbone's frozen cache, eliminating redundant re-parameterization. Furthermore, a parameter-free Cognitive Monitor leverages attention entropy to adaptively trigger consolidation only when high epistemic uncertainty is detected. Experiments demonstrate that FlashMem matches the performance of heavy baselines while reducing inference latency by 5 times, effectively bridging the gap between efficiency and persistent cognition.
Paper Structure (66 sections, 9 equations, 4 figures, 8 tables, 1 algorithm)

This paper contains 66 sections, 9 equations, 4 figures, 8 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Token-level and Parametric Memory.
  4. Generative Latent Memory.
  5. Efficiency and Cognitive Control.
  6. Methodology
  7. Preliminaries and Problem Formulation
  8. Agent Formulation and the Efficiency Gap.
  9. Overview of FlashMem
  10. Inference Workflow.
  11. The Cognitive Monitor: Entropy-Based Uncertainty Perception
  12. Mitigating Attention Sinks.
  13. Entropy Decision and Adaptive Intervention.
  14. The Memory Consolidator: Shared-KV Generation
  15. State Projection via Injectivity.
  16. ...and 51 more sections

Figures (4)

  • Figure 1: The architecture of FlashMem. It features a parameter-free Cognitive Monitor (top right) that activates memory generation based on attention entropy uncertainty, and a lightweight Memory Consolidator (bottom right). The consolidator utilizes a Shared-KV design to efficiently generate latent memories by reusing the frozen backbone's cache, taking the hidden state $h_t$ as the initiation seed.
  • Figure 2: Efficiency Trade-offs across Layer Depths. VRAM usage scales linearly with depth, while latency increases monotonically. Combined with Table \ref{['tab:layer_acc']}, results favor a minimal depth configuration.
  • Figure 3: Entropy Dynamics during Reasoning. We visualize the attention entropy of FlashMem (Blue) vs. Vanilla (Red) across four models on challenging GSM8K samples. The vertical dashed lines indicate memory injection points triggered by the Cognitive Monitor. Post-injection, FlashMem exhibits a notable reduction in entropy compared to the baseline, reflecting increased confidence and reasoning stability.
  • Figure 4: Visualization of Attention Dynamics.Top: Consolidator's cross-attention accurately targets key information (e.g., variable definitions) in the context. Middle: FlashMem's backbone consistently attends to the injected latent tokens (marked by yellow boxes) throughout generation, as evidenced by the continuous vertical attention lines. Bottom: The Vanilla model exhibits diffuse attention, failing to capture distant dependencies.