TransMem: Transforming Hidden States into Memory for Large Language Models

2026-07-31

AI Summary

This work addresses the limitation of large language models in effectively reusing task-relevant evidence from historical hidden states during extended interactions, which constrains their reasoning capabilities. To overcome this, the authors propose TransMem—a lightweight, inference-time parametric memory module that operates with a frozen backbone model. TransMem converts sparse historical hidden states into reusable memory representations and dynamically modulates the current state via a gating mechanism to avoid redundant encoding. Innovatively, it incorporates evidence-conditioned self-distillation, enabling the memory module to learn transferable memory utilization strategies rather than task-specific knowledge. Experimental results demonstrate that TransMem substantially improves performance across multiple benchmarks: it achieves up to a 29.25 F1 gain on LoCoMo and HotpotQA, and raises the average accuracy on MemoryAgentBench from 29.54% to 40.00%.

Abstract

Large language model (LLM) agents increasingly operate over long interaction histories, where effective reasoning requires identifying and exploiting task-relevant evidence distributed across past observations and actions. However, useful information encoded in previously computed representations is often underutilized during subsequent generation. We propose \textbf{TransMem}, a lightweight inference-time parametric memory module that transforms sparse historical hidden states from a frozen LLM backbone into reusable memory representations. TransMem uses a lightweight gating network to dynamically apply the latent intervention to the current hidden states, without repeatedly encoding the preceding context. To learn transferable memory utilization rather than task-specific knowledge, we introduce evidence-conditioned self-distillation. A memory-augmented student processes the full context and matches the predictive distribution of an evidence-only teacher that shares the same frozen backbone. Experiments on LoCoMo, HotpotQA, and MemoryAgentBench demonstrate consistent improvements across different model architectures and scales. TransMem yields gains of 11.58--29.25 $F_1$ on LoCoMo and 10.20--13.03 $F_1$ on HotpotQA, while improving the average MemoryAgentBench accuracy from 29.54\% to 40.00\%. These results establish sparse historical hidden states as an effective and efficient memory substrate for long-context LLM agents. Our code is available at https://github.com/Haodong-Lei-Ray/TransMem.

Problem

• large language models
• long-context reasoning
• memory utilization
• hidden states
• task-relevant evidence

Innovation

• parametric memory
• hidden state transformation
• evidence-conditioned self-distillation
• long-context reasoning
• inference-time memory