GradMem: 테스트 시간 경사 하강법을 통한 메모리 컨텍스트 기록 학습

긴 문서를 처리할 때 발생하는 막대한 메모리 비용 문제를 해결하기 위해, 모델의 가중치는 고정한 채 메모리 토큰만 실시간으로 최적화하여 정보를 압축 저장하는 새로운 패러다임을 구축했다. 이는 기존의 단순 압축 방식보다 훨씬 높은 정보 보존력을 보여주며 추론 효율성을 극대화했다.

Transformer의 Self-Attention은 시퀀스 내 모든 토큰 쌍의 관계를 계산하는데, 이 과정에서 발생하는 KV-cache는 문맥이 길어질수록 메모리 사용량이 선형적으로 증가하여 하드웨어 한계에 부딪혔다. 기존의 압축 방식은 정보를 한 번 훑으며 요약하려 했으나, 이 과정에서 중요한 세부 정보를 놓치는 경우가 많았다. GradMem은 이 문제를 해결하기 위해 모델의 가중치는 그대로 둔 채, 입력된 문맥을 가장 잘 표현할 수 있도록 소수의 '메모리 토큰' 임베딩만을 테스트 시점에 직접 최적화했다. 이는 마치 시험 공부를 할 때 중요한 내용을 골라 요약 노트를 반복해서 다듬는 과정과 같으며, 재구성 손실(Reconstruction Loss)을 통해 메모리 토큰만으로도 원래 문장을 복원할 수 있을 만큼 정보를 밀도 있게 압축했다. 결과적으로 단 5번의 경사 하강 업데이트만으로도 방대한 정보를 정확하게 저장하며, 한 번 생성된 메모리는 이후 수많은 쿼리에 대해 원본 문맥 없이도 즉각적인 답변을 제공할 수 있게 됐다.

GradMem은 WRITE와 READ라는 두 단계의 명시적 프로세스로 작동한다. WRITE 단계에서는 입력 컨텍스트 C를 고정된 크기의 메모리 M으로 압축하며, 이때 자기지도 재구성 손실 L_write를 최소화하는 방향으로 메모리 토큰의 임베딩을 갱신한다. [메모리 토큰 M과 이전 토큰 시퀀스 t_<i를 입력으로] -> [Transformer 모델 f_theta를 통해 다음 토큰 t_i가 나타날 확률의 로그값을 계산하고 이를 모두 더해 음수를 취하면] -> [재구성 손실이라는 하나의 숫자가 나오고] -> [이 값이 작을수록 메모리가 원본 문맥을 더 정확하게 요약하고 있음을 의미한다]. 이 최적화 과정은 메타 학습된 초기값 M_0 덕분에 단 1~5회의 경사 하강 단계만으로도 충분한 성능을 냈다. READ 단계에서는 최적화된 메모리 M_hat과 새로운 쿼리 Q만을 결합하여 최종 정답 Y를 생성하며, 이때 원본 컨텍스트 C에는 전혀 접근하지 않는다. 이러한 구조는 긴 문맥을 한 번만 읽고 압축하여 재사용할 수 있게 함으로써 추론 시의 연산 효율을 극대화했다.

연상 키-값 검색(Associative KV-retrieval) 실험에서 GradMem은 동일한 메모리 크기를 가진 기존의 순방향 쓰기 방식인 RMT를 크게 앞질렀다. 8개의 메모리 토큰을 사용할 때, GradMem은 96쌍의 키-값을 88%의 정확도로 검색해낸 반면, 순방향 방식은 8쌍을 넘어서면 정확도가 급격히 하락했다. 또한 경사 하강 단계 K를 늘릴수록 메모리 용량이 효과적으로 확장됨이 확인됐다. 자연어 처리 태스크인 bAbI와 SQuAD에서도 GradMem은 경쟁력 있는 성능을 보였다. 특히 SQuAD(short) 데이터셋에서 GradMem(K=5)은 54.9%의 정확도를 기록하며, 압축을 수행하지 않은 Pythia-160m 모델의 48.9%를 상회하는 결과를 나타냈다. 추론 지연 시간 분석에서는 동일 컨텍스트에 대해 약 64회 이상의 쿼리가 발생할 경우, KV-cache를 유지하는 표준 Transformer보다 전체 연산 비용이 낮아지는 손익분기점을 달성했다.

GradMem 아키텍처는 테스트 시간 훈련(TTT)의 원리를 입력 프리픽스 메모리 토큰에 적용한 형태이다. 모델의 파라미터 theta는 고정된 상태에서 메모리 토큰 M만이 최적화 대상이 된다. 학습 시에는 MAML 스타일의 2차 미분 구조를 사용하여, 내부 루프의 K단계 최적화 이후 외부 루프에서 태스크 손실을 최소화하도록 초기값 M_0와 모델 theta를 동시에 학습시켰다. 기술적으로 가장 큰 도전 과제인 Attention 메커니즘에 대한 double-backward 연산 비용을 줄이기 위해, Flash HVP와 유사한 효율적인 헤시안-벡터 곱(Hessian-Vector Product) 구현체를 도입했다. 이를 통해 1024 토큰 시퀀스 기준 역전파 연산 시간을 약 40% 단축하고 GPU 메모리 점유율을 60GB에서 30GB 수준으로 절반가량 낮추는 데 성공했다. 이러한 최적화는 GradMem이 더 긴 문맥과 더 많은 최적화 단계를 수용할 수 있게 하는 핵심적인 기반이 됐다.

훈련 과정에서 unrolled optimization 단계를 거쳐야 하므로 일반적인 학습보다 더 많은 GPU 메모리와 연산 시간이 소요된다. 또한 테스트 시점에 경사 하강법을 수행해야 하므로 첫 번째 쿼리에 대한 응답 지연 시간(Latency)이 순방향 모델보다 길어질 수 있다.