[논문 리뷰] Neural Garbage Collection: Learning to Forget while Learning to Reason

TL;DR

최근 발표된 논문 "Neural Garbage Collection: Learning to Forget while Learning to Reason"은 언어 모델의 고질적인 메모리 문제를 해결하기 위한 혁신적인 방법론을 제시합니다. 이 연구는 언어 모델이 추론 과정에서 불필요한 정보를 담은 KV 캐시를 스스로 정리하도록 학습시키는 Neural Garbage Collection (NGC) 기법을 도입했습니다. NGC는 강화학습을 통해 최종 과제 성공 여부라는 보상 신호만으로 모델이 무엇을 기억하고 무엇을 잊을지 동적으로 결정하게 합니다. 실험 결과, NGC는 KV 캐시의 50%를 제거하여 2배 이상의 메모리 압축 효과를 보이면서도 원본 모델과 거의 동등한, 혹은 더 높은 정확도를 유지했습니다. 이는 기존의 고정 규칙 기반 메모리 관리 기법들을 능가하는 성능으로, 언어 모델의 추론 능력과 효율성을 동시에 향상시키는 새로운 가능성을 열어줍니다.

연구 배경 및 동기

거대 언어 모델(LLM)의 추론 능력은 비약적으로 발전했지만, 그 이면에는 막대한 메모리 사용량이라는 그림자가 있습니다. 특히, 복잡한 문제를 풀기 위해 사용되는 연쇄적 사고(Chain-of-Thought, CoT) 기법은 모델이 중간 추론 과정을 단계별로 생성하면서 메모리 사용량을 폭발적으로 증가시킵니다.

이 문제의 핵심에는 **KV 캐시(Key-Value Cache)**가 있습니다. 트랜스포머 아키텍처에서 어텐션 메커니즘은 이전에 생성된 모든 토큰의 Key와 Value 값을 저장해두고(캐싱), 다음 토큰을 생성할 때 이를 참조합니다. 이 때문에 추론 길이가 길어질수록 KV 캐시의 크기는 선형적으로 증가하여($O(L^2)$의 어텐션 계산을 $O(L)$로 줄여주지만, 메모리 자체는 $O(L)$로 증가) 메모리 병목 현상을 일으키고 모델의 확장성을 저해합니다.

기존 메모리 관리 기법은 '최근 N개 토큰만 남기기(Sliding Window)'나 '처음 N개 토큰은 항상 유지하기(Attention Sinks)'와 같은 고정된 규칙에 의존했습니다. 이런 방식은 간단하지만, 문맥의 중요도를 고려하지 못해 핵심 정보를 잃어버릴 위험이 큽니다. 문제의 종류나 추론 단계에 따라 중요한 정보의 위치는 계속 바뀌기 때문입니다. 따라서, 언어 모델이 스스로 문맥을 파악하고 불필요한 정보를 '잊을' 수 있는 능력을 학습하는 것이 절실합니다.

관련 연구

본 연구는 기존의 고정 규칙 기반 메모리 관리 기법의 한계를 극복하고자 합니다.

NGC는 위 기법들과 달리, 어떤 정보를 버리고 남길지를 학습 가능한 정책(policy)으로 간주합니다. 이를 통해 과제에 따라 유연하게 메모리를 관리하는 능력을 모델에 부여하는 것이 핵심 차별점입니다.

핵심 기여

1. Neural Garbage Collection (NGC) 제안: 언어 모델이 강화학습을 통해 메모리를 동적으로 관리하는 새로운 프레임워크를 제시했습니다. 모델이 스스로 무엇을 기억하고 무엇을 잊을지 학습합니다.
2. 종단간(End-to-End) 학습 프레임워크: 최종 결과(정답 여부)라는 희소한(sparse) 보상 신호만으로 메모리 관리 정책을 학습할 수 있는 종단간 학습 프레임워크를 개발했습니다.
3. 성능과 효율성 동시 달성: NGC를 통해 언어 모델의 메모리 사용량을 2배 이상 압축하면서도 원본 모델과 동등하거나 더 높은 정확도를 유지함을 실험적으로 입증했습니다.
4. 강화학습의 새로운 적용: 메모리 관리라는 결정 과정을 강화학습 문제로 공식화하여, 모델이 다양한 과제에 적응적으로 대응할 수 있는 능력을 부여했습니다.

제안 방법론

NGC는 언어 모델의 메모리 관리를 강화학습(Reinforcement Learning, RL) 문제로 정의하고 해결합니다. 모델은 '에이전트(agent)'가 되어, 과제를 해결하는 동시에 메모리를 효율적으로 관리하는 방법을 학습합니다.

강화학습 기반 메모리 관리

NGC의 학습 목표는 다음과 같은 강화학습 목표 함수를 최대화하는 것입니다.

각 요소의 의미는 다음과 같습니다.

• $J(\theta)$: 최적화하려는 목표 함수. 즉, 보상의 기댓값입니다.
• $\theta$: 학습 대상인 모델의 파라미터입니다. 여기에는 언어 생성 능력과 메모리 관리 정책이 모두 포함됩니다.
• $\pi_\theta$: 파라미터 $\theta$로 정의되는 정책(policy)입니다. 주어진 상태에서 어떤 행동(토큰 생성, 캐시 제거)을 할지 결정합니다.
• $\tau$: 하나의 에피소드(trajectory)로, 토큰 생성과 캐시 제거로 이루어진 전체 행동 시퀀스 $(s_0, a_0, s_1, a_1, ...)$를 의미합니다.
• $R(\tau)$: 에피소드 $\tau$가 끝났을 때 받는 최종 보상(reward)입니다. 예를 들어, 수학 문제를 맞혔으면 +1, 틀렸으면 0을 받습니다.

모델은 이 보상 $R(\tau)$를 최대로 만드는 방향으로 파라미터 $\theta$를 업데이트합니다. 즉, 정답을 맞히는 데 도움이 된 메모리 관리 전략은 강화하고, 그렇지 않은 전략은 억제하는 방식으로 학습이 진행됩니다.

이산적인 캐시 제거 결정(버린다/남긴다)을 미분 가능한 방식으로 학습시키기 위해 Gumbel-Top-k 트릭이 사용됩니다. 이는 각 캐시 항목에 대한 '유지 확률'을 계산한 뒤, 확률적 샘플링을 통해 상위 k개를 선택하는 과정을 경사 하강법으로 최적화할 수 있게 해주는 기법입니다.

'생성-수집(Generate-and-Collect)' 알고리즘

NGC는 다음과 같은 '생성-수집' 주기를 반복하며 추론을 진행합니다.

1. 생성 (Generate): 모델이 일정 길이(예: 64 토큰)의 추론 과정을 생성합니다. 이때 생성된 토큰들의 KV 캐시가 메모리에 쌓입니다.
2. 수집 (Collect): 생성을 잠시 멈추고, 현재까지 쌓인 KV 캐시 중에서 어떤 항목을 버릴지(garbage collection) 결정합니다. NGC 정책 네트워크가 각 캐시의 중요도를 예측합니다.
3. 압축 (Compact): 결정에 따라 불필요한 캐시를 제거하고 메모리를 압축합니다. 예를 들어, 50%를 제거하기로 했다면 중요도가 낮은 절반을 버립니다.
4. 재개 (Resume): 정리된 캐시를 바탕으로 다음 추론 생성을 이어갑니다.

이 모든 과정은 오직 최종 결과(정답 여부)에 대한 보상 신호 하나만으로 최적화됩니다. 어떤 캐시를 버리는 것이 정답을 맞히는 데 도움이 되었는지를 모델이 수많은 시도를 통해 스스로 학습하는 것입니다.

실험 설정

NGC의 성능을 검증하기 위해 복잡한 추론 능력을 요구하는 다양한 수학 벤치마크에서 실험이 진행되었습니다.

• 기반 모델: deepseek-ai/deepseek-coder-1.3b-instruct 모델을 Qwen-1.5B 스타일로 추가 파인튜닝하여 사용
• 데이터셋:
  • Countdown: 간단한 산술 연산 문제
  • AMC, AIME: 미국 수학 경시대회 문제로, 복잡한 다단계 추론 필요
• 평가 지표: 정확도(Accuracy), 메모리 사용량(KV Cache Compression Ratio)
• 베이스라인:
  • No Eviction (원본 모델)
  • Sliding Window (최신 토큰 유지)
  • H2O / StreamingLLM (초기 + 최신 토큰 유지)

실험 결과 분석

실험 결과, NGC는 메모리 효율성과 추론 성능 두 마리 토끼를 모두 잡는 데 성공했습니다.

주요 결과

위 표에서 볼 수 있듯, NGC는 KV 캐시의 50%를 제거하여 메모리 사용량을 절반으로 줄이면서도 원본 모델과 거의 동등한 정확도를 보였습니다. Countdown 데이터셋에서는 0.9%p의 미미한 하락에 그쳤고, 더 어려운 AMC 데이터셋에서는 오히려 0.5%p 성능이 향상되었습니다. 이는 불필요한 노이즈를 제거하는 것이 오히려 모델의 추론에 긍정적인 영향을 줄 수 있음을 시사합니다.

반면, Sliding Window나 H2O와 같은 고정 규칙 기반 방식들은 중요한 문맥 정보를 잃어버리면서 상당한 성능 하락을 겪었습니다.

Ablation Study (구성 요소 분석)

NGC의 핵심 설계 요소가 왜 중요한지를 확인하기 위해 일부 요소를 제거한 모델과 성능을 비교했습니다.

• NGC (Full Model): 49.6% 정확도
• Token log-probs only: 35.7% 정확도 (캐시 제거 정책을 직접 학습하지 않고, 토큰 생성 확률만을 이용해 간접적으로 학습)
• Targeted KV dropout: 2.5% 정확도 (강화학습의 off-policy 문제를 제대로 처리하지 않고 단순히 드롭아웃 적용)

결과는 NGC의 두 가지 핵심 설계, 즉 제거 정책을 직접 학습하는 것과 강화학습의 off-policy 특성을 올바르게 처리하는 것이 성능에 필수적임을 명확히 보여줍니다. 특히 Targeted KV dropout은 훈련 중 경사 폭주(exploding gradients)로 훈련이 붕괴되어, 정교한 강화학습 프레임워크의 중요성을 입증했습니다.

비판적 평가

강점

1. 높은 효율성: 상당한 메모리 압축률에도 불구하고 성능 저하가 거의 없거나 오히려 향상되는 놀라운 효율성을 보여줍니다.
2. 뛰어난 적응성: 강화학습을 통해 문제의 종류나 추론 단계에 따라 동적으로 메모리를 관리하므로, 다양한 태스크에 유연하게 적용될 수 있습니다.
3. 새로운 패러다임: 메모리 관리를 학습 가능한 정책으로 취급함으로써, LLM 최적화에 대한 새로운 연구 방향을 제시했습니다.

한계점 및 개선 방향

• 훈련 복잡성 및 불안정성: 강화학습은 보상 설계, 하이퍼파라미터 튜닝에 민감하며 훈련이 불안정한 경향이 있습니다. 안정적인 학습을 위한 추가적인 기법 연구가 필요합니다.
• 추론 오버헤드: 추론 시 '생성-수집' 주기마다 정책 네트워크를 실행해야 하므로 약간의 지연 시간(latency)이 추가될 수 있습니다. 이 오버헤드가 메모리 절약의 이점보다 크지 않은지 고려해야 합니다.
• 일반화 성능: 실험이 주로 수학 추론 문제에 집중되어 있어, 코드 생성, 스토리텔링, 요약 등 다른 종류의 긴 문맥 태스크에서도 동일한 효과를 보일지는 추가 검증이 필요합니다.

실무 적용 가이드

• 긴 문맥 추론 서비스: CoT를 사용하거나 매우 긴 문서를 처리해야 하는 RAG(검색 증강 생성) 시스템 등에서 메모리 병목을 해결하기 위한 효과적인 솔루션이 될 수 있습니다.
• 온디바이스 AI: 스마트폰이나 엣지 디바이스와 같이 메모리가 제한된 환경에서 LLM을 구동할 때, NGC는 모델의 가용성을 크게 높일 수 있습니다.
• 구현 시 고려사항: NGC를 직접 구현하거나 파인튜닝할 때는 강화학습에 대한 깊은 이해가 필요합니다. PPO(Proximal Policy Optimization)와 같은 안정적인 RL 알고리즘을 기반으로 시작하고, 작은 규모의 문제에서 충분한 실험을 거치는 것이 좋습니다.

결론

Neural Garbage Collection은 언어 모델이 단순히 지식을 저장하는 수동적인 존재를 넘어, 자신의 '기억'을 능동적으로 관리하는 주체가 될 수 있음을 보여준 획기적인 연구입니다. 이는 제한된 컴퓨팅 자원 내에서 더 길고 복잡한 추론을 수행할 수 있는 길을 열어주며, 언어 모델의 효율성과 능력을 동시에 한 단계 끌어올릴 중요한 이정표입니다. 앞으로 NGC와 같은 기술이 발전함에 따라, 우리는 더욱 똑똑하고 효율적인 AI를 만나게 될 것입니다.

참고 자료

• 논문 링크 (arXiv)
• 코드 저장소 (GitHub)