GradMem: 테스트 타임 그래디언트 하강을 통해 컨텍스트를 메모리에 기록하는 학습 방법

초록

많은 대규모 언어 모델 애플리케이션은 긴 문맥을 조건으로 사용해야 합니다. 트랜스포머는 일반적으로 과거 활성화 값을 각 계층별로 대규모 KV 캐시에 저장하여 이를 지원하지만, 이는 상당한 메모리 오버헤드를 초래합니다. 바람직한 대안은 압축 메모리입니다. 즉, 문맥을 한 번 읽고 간결한 상태로 저장한 후, 그 상태에서 여러 쿼리에 답하는 방식입니다. 우리는 추론 시점에 원본 문맥에 접근하지 않고 답을 생성해야 하는 문맥 제거 설정에서 이를 연구합니다. 우리는 샘플별 추론 시 최적화를 통해 문맥을 메모리에 기록하는 GradMem을 소개합니다. GradMem은 주어진 문맥에 대해 모델 가중치는 고정한 채 소규모 접두사 메모리 토큰 집합에 대해 몇 단계의 경사 하강법을 수행합니다. GradMem은 모델 수준의 자기 지도 문맥 재구성 손실을 명시적으로 최적화하여, 순전파 전용 방법과 달리 반복적 오류 수정이 가능한 손실 기반 기록 연산을 구현합니다. 연관 키-값 검색 과제에서 GradMem은 동일한 메모리 크기를 가진 순전파 전용 메모리 기록 방법을 능가하며, 추가적인 경사 하강 단계는 반복적인 순전파 기록보다 용량 확장에 훨씬 더 효과적입니다. 우리는 더 나아가 GradMem이 합성 벤치마크를 넘어서서 전이됨을 보여줍니다: 사전 학습된 언어 모델을 사용하여 bAbI 및 SQuAD 변형을 포함한 자연어 처리 과제에서 메모리에 인코딩된 정보만을 의존하여 경쟁력 있는 결과를 달성합니다.

English

Many large language model applications require conditioning on long contexts. Transformers typically support this by storing a large per-layer KV-cache of past activations, which incurs substantial memory overhead. A desirable alternative is ompressive memory: read a context once, store it in a compact state, and answer many queries from that state. We study this in a context removal setting, where the model must generate an answer without access to the original context at inference time. We introduce GradMem, which writes context into memory via per-sample test-time optimization. Given a context, GradMem performs a few steps of gradient descent on a small set of prefix memory tokens while keeping model weights frozen. GradMem explicitly optimizes a model-level self-supervised context reconstruction loss, resulting in a loss-driven write operation with iterative error correction, unlike forward-only methods. On associative key--value retrieval, GradMem outperforms forward-only memory writers with the same memory size, and additional gradient steps scale capacity much more effectively than repeated forward writes. We further show that GradMem transfers beyond synthetic benchmarks: with pretrained language models, it attains competitive results on natural language tasks including bAbI and SQuAD variants, relying only on information encoded in memory.

GradMem: 테스트 타임 그래디언트 하강을 통해 컨텍스트를 메모리에 기록하는 학습 방법

GradMem: Learning to Write Context into Memory with Test-Time Gradient Descent

초록

Support