ElasticMem: Latent Memory as a Learnable Resource for LLM Agents

기술 요약: ElasticMem

문제 정의

장기 기억은 대규모 언어 모델(LLM) 에이전트가 확장된 상호작용 전반에 걸쳐 일관되게 추론하고, 응답을 개인화하며, 과거의 경험을 재사용할 수 있도록 하는 데 매우 중요합니다. 그러나 기존의 메모리 증강 시스템은 메모리를 고정된 자원으로 취급하며, 이는 메모리의 쿼리 의존적 유용성과 정적인 메모리 용량 할당 사이의 불일치를 야기합니다.

현재의 접근 방식은 크게 두 가지 범주로 나뉘며, 두 방식 모두 상당한 한계를 가지고 있습니다:

1. 텍스트 공간 방식(Text-Space Methods): 자연어 기록을 검색하여 컨텍스트 윈도우에 결합합니다. 해석 가능성은 높지만, 토큰 오버헤드가 크고 노이즈나 중복된 증거에 민감합니다. 즉, 무관한 메모리가 유용성과 상관없이 컨텍스트 토큰을 소비하게 됩니다.
2. 잠재 공간 방식(Latent-Space Methods): 텍스트 비용을 줄이기 위해 메모리를 연속적인 벡터나 소프트 토큰으로 압축합니다. 그러나 이들은 일반적으로 경직된 검색 메커니즘(예: 고정된 코사인 유사도)과 고정된 용량 인터페이스에 의존하며, 모든 검색된 청크에 동일한 잠재 예산을 할당합니다. 이는 결정적인 증거를 포함하는 청크와 단순히 어휘적으로 유사하지만 도움이 되지 않는 청크를 구분하지 못합니다.

핵심 과제는 메모리의 유용성이 쿼리와 태스크에 따라 매우 가변적임에도 불구하고, 검색 및 용량 할당이 학습된 결과가 아닌 고정된 설계 선택에 의해 결정된다는 점입니다.

방법론: ElasticMem

저자들은 ElasticMem을 제안하며, 이는 메모리를 **학습 가능한 탄력적 잠재 자원(learnable, elastic latent resource)**으로 취급하는 프레임워크입니다. ElasticMem은 고정된 검색과 균등한 할당 대신, 추론기의 내부 상태에 따라 메모리를 적응적으로 검색하고, 하위 유용성에 따라 각 메모리에 가변적인 잠재 예산을 할당하도록 학습합니다.

이 프레임워크는 네 가지 주요 단계로 구성됩니다:

1. 잠재 메모리 뱅크 구축 (Latent Memory Bank Construction)

오프라인 잠재 메모리 뱅크는 한 번 구축되면 훈련 중에 동결된 상태로 유지됩니다. 각 메모리 청크 $m_i$는 동결된 베이스 LLM 인코더에 의해 다음을 생성하도록 인코딩됩니다:

• 검색 키($k_i$): 청크의 최종 은닉 상태(final hidden state).
• 콘텐츠 캐시($C_i$): 청크의 마지막 $N_c$개 은닉 상태.
  이를 통해 추론 시 전체 코퍼스를 다시 인코딩하지 않고도 효율적인 검색이 가능합니다.

2. 쿼리 조건부 검색 (Query-Conditioned Retrieval)

별도의 고정된 리트리버를 사용하는 표준 시스템과 달리, ElasticMem은 추론기 자신의 은닉 상태로부터 검색 쿼리를 도출합니다.

• LoRA가 적용된 추론기는 쿼리 $q$가 주어지면 검색 제어 토큰(retrieval-control token) $a$를 샘플링합니다.
• 추론기는 $(q, a)$를 처리하여 쿼리 상태 $h_q$를 생성합니다.
• 메모리는 $h_q$와 캐싱된 검색 키 $k_i$ 사이의 코사인 유사도를 기반으로 검색됩니다. 이를 통해 검색이 단순한 표면적 의미 매칭을 넘어, 모델의 내부적 해석에 적응할 수 있게 합니다.

3. 탄력적 잠재 예산 할당 (Elastic Latent Budget Allocation)

상위 $Z$개의 후보 청크를 검색한 후, 경량 **트랜스포머 예산 정책($\pi_\phi$)**이 각 청크에 가변적인 잠재 토큰 수($b_j$)를 할당합니다.

• 예산 $b_j$는 이산 분포 $\{0, 1, \dots, B_{max}\}$에서 샘플링됩니다.
• $b_j=0$인 예산은 유사하지만 도움이 되지 않는 검색된 청크를 억제할 수 있게 합니다.
• 이 정책은 쿼리 상태, 메모리 키, 검색 점수, 학습된 랭크 임베딩을 특징으로 포함하며, 예산을 할당하기 전에 셀프 어텐션을 사용하여 메모리 간의 상호작용을 비교합니다.
• 유용한 청크는 더 큰 예산을 할 받는 반면, 중복되거나 오도하는 청크는 압축되거나 억제됩니다.

4. 잠재 메모리 주입 및 최적화 (Latent Memory Injection and Optimization)

콘텐츠 캐시에서 선택된 은닉 상태들은 학습된 **잠재 프로젝터($P_\psi$)**를 통해 소프트 메모리 토큰으로 투영되어 추론기의 입력 시퀀스에 주입됩니다.

• 훈련: 전체 과정(검색 제어, 예산 할당, 투영, 생성)은 **그룹 상대 정책 최적화(GRPO)**를 사용하여 공동으로 최적화됩니다.
• 보상 신호(예: 정확도, 태스크 성공 여부)는 단순히 최종 답변뿐만 아니라 검색 결정과 예산 할당까지 감독합니다. 이를 통해 모델은 실제 태스크 결과에 기반하여 어떤 메모리를 검색하고 각 메모리에 어느 정도의 표현 용량을 할당해야 하는지를 학습할 수 있습니다.

주요 기여

1. 탄력적 잠재 메모리 (Elastic Latent Memory): 메모리를 고정된 텍스트 또는 잠재 인터페이스에서 유용성에 따라 용량이 동적으로 할당되는 학습 가능한 탄력적 자원으로 전환하는 새로운 패러다임을 제시했습니다.
2. 추론기 상태 적응형 검색 (Reasoner-State Adaptive Retrieval): 경직된 유사도 매칭을 넘어, 추론기의 내부 상태와 샘플링된 제어 토큰에 의해 조건화되는 검색 메커니즘을 구현했습니다.
3. 유용성 인식 예산 정책 (Utility-Aware Budget Policy): 검색된 청크에 가변적인 잠재 예산을 할당하여, 모델이 유용한 증거를 우선시하고 노이즈를 억제할 수 있도록 하는 학습된 정책을 제공합니다.
4. 엔드 투 엔드 최적화 (End-to-End Optimization): 다운스트림 태스크 보상을 통해 검색, 할당 및 생성을 공동으로 최적화하는 훈련 프레임워크를 구축하여 메모리 관리를 태스크 유용성과 정렬시켰습니다.

실험 결과

ElasticMem은 Qwen2.5-3B 및 Qwen2.5-7B 백본을 사용하여 메모리 집약적 QA(PersonaMem, LoCoMo, LongMemEval)와 체화된 에이전트 의사결정(ALFWorld)을 포함하는 MemorySuite에서 평가되었습니다.

• 성능 향상:
  • QA 태스크: ElasticMem은 가장 강력한 베이스라인 대비 가중 평균 정확도를 Qwen2.5-3B에서 26.2%, Qwen2.5-7B에서 24.6% 향상시켰습니다.
  • 에이전트 태스크 (ALFWorld): ElasticMem은 가중 평균 성공률을 Qwen2.5-3B에서 66.3%, Qwen2.5-7B에서 27.2% 향상시켰습니다.
• 효율성: ElasticMem은 ALFWorld에서 비교된 모든 방법 중 가장 낮은 토큰 비용을 달 기록하며, 성능 향상이 긴 상호작용 궤적이나 과도한 메모리 사용 때문이 아님을 입증했습니다.
• 절제 연구 (Ablation Studies):
  • 예산 정책: 트랜스포머 기반 예산 정책은 랜덤, 균등, MLP 기반 할당 전략보다 훨씬 뛰어난 성능을 보여, 메모리 간 상호작용을 학습하는 것의 가치를 확인했습니다.
  • 검색: 제어 토큰을 사용한 쿼리 상태 기반 검색이 의미적 및 동결 상태 검색보다 우수한 성능을 보여, 검색을 추론기의 내부 상태에 적응시키는 것의 중요성을 강조했습니다.
  • 용량 제한: 적절한 청크당 예산 제한($B_{max}=20$)이 최상의 결과를 냈으며, 이는 과도한 예산 할당이 오히려 유용하지 않은 청크에 용량을 할당하여 성능을 저하시킬 수 있음을 시사합니다.

의의 및 주장

본 논문은 ElasticMem이 현재 메모리 시스템의 근본적인 불일치, 즉 메모리 할당의 정적인 특성과 메모리 유용성의 동적인 특성 사이의 문제를 해결한다고 주장합니다. 이 프레임워크를 통해 에이전트는 다음과 같은 능력을 갖출 수 있습니다:

• 가장 어휘적으로 유사한 쿼리가 아니더라도 유용한 증거와 전이 가능한 계획을 우선순위에 둡니다.
• 컨텍스트나 잠재 용량을 소비할 수 있는 노이즈 및 중복 메모리를 억제합니다.
• 높은 토큰 효율성을 유지하면서 추론 및 의사결정 태스크에서 탁월한 성능을 달립니다.

저자들은 탄력적 잠재 메모리가 고정된 컨텍스트 및 경직된 검색 패러다임의 한계를 넘어, 더 효율적이고 적응적이며 장기적인 추론이 가능한 LLM 에이전트를 구축하는 유망한 방향이라고 결론짓습니다.