TransMem: 대규모 언어 모델을 위한 숨겨진 상태를 메모리로 변환

TransMem: Transforming Hidden States into Memory for Large Language Models

대규모 언어 모델(LLM) 에이전트는 점점 더 긴 상호 작용 기록에서 작동하며, 효과적인 추론은 과거의 관찰 및 행동에 분산된 작업 관련 증거를 식별하고 활용하는 것을 필요로 합니다. 그러나 이전에 계산된 표현에 인코딩된 유용한 정보는 종종 후속 생성 과정에서 충분히 활용되지 않습니다. 본 논문에서는 가볍고 추론 시간에 적용 가능한 파라미터 메모리 모듈인 **TransMem**을 제안합니다. TransMem은 동결된 LLM의 희소한 과거 숨겨진 상태를 재사용 가능한 메모리 표현으로 변환합니다. TransMem은 경량 게이팅 네트워크를 사용하여 현재 숨겨진 상태에 잠재적인 개입을 동적으로 적용하며, 이전 컨텍스트를 반복적으로 인코딩하지 않습니다. 작업별 지식 대신 전송 가능한 메모리 활용 방법을 학습하기 위해 증거 기반 자기 훈련(self-distillation)을 도입했습니다. 메모리가 강화된 학생 모델은 전체 컨텍스트를 처리하고 동일한 동결된 백본을 공유하는 증거만 사용하는 교사 모델의 예측 분포와 일치시킵니다. LoCoMo, HotpotQA 및 MemoryAgentBench에서의 실험 결과는 다양한 모델 아키텍처 및 규모에서 지속적인 성능 향상을 보여줍니다. TransMem은 LoCoMo에서 11.58–29.25의 $F_1$ 점수 향상과 HotpotQA에서 10.20–13.03의 $F_1$ 점수 향상을 가져왔으며, MemoryAgentBench의 평균 정확도를 29.54%에서 40.00%로 향상시켰습니다. 이러한 결과는 희소한 과거 숨겨진 상태가 긴 컨텍스트를 처리하는 LLM 에이전트를 위한 효과적이고 효율적인 메모리 기반임을 입증합니다. 저희 코드는 https://github.com/Haodong-Lei-Ray/TransMem에서 확인할 수 있습니다.

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.