AI 에이전트의 메모리 문제와 KV 캐시 압축: Baseten 연구팀의 STILL 접근

핵심 메시지

• 메모리에는 선형으로 늘어나는 무손실 KV 캐시와 파인튜닝·RAG 같은 고압축 방식이 있는데, 그 사이를 메울 중간 메모리 층이 빠져 있다.
• 긴 호흡의 에이전트에는 압축 메모리, 루프 안에서의 반복 사용, 경험을 가중치로 되돌리는 투영이라는 세 요소가 필요하다.
• 스파스 오토인코더처럼 추론 비용을 학습 단계로 분산(amortize)해, 전체 KV 캐시를 압축 캐시로 바꾸는 함수를 학습한다.
• 압축된 KV 캐시는 사실상 하나의 MLP이며, 이는 그래디언트 디센트가 아니라 압축으로 가중치를 만들어 내는 새로운 방식이다.

쉽게 이해하기

Baseten 연구팀은 추론 시점에 필요한 맥락을 손실 없이 불러오는 문제, 곧 '메모리 문제'를 다룬다. 그들은 메모리를 크게 두 종류로 본다. 하나는 전체 KV 캐시처럼 비교적 무손실이지만 메모리 요구량이 선형으로 늘어 긴 호흡의 에이전트에서 병목이 되는 방식이고, 다른 하나는 파인튜닝·RAG·자연어 요약 파일처럼 강하게 압축된 방식이다. 문제는 이 둘 사이를 메울 '중간 메모리 층'이 비어 있다는 점이다.

신뢰할 수 있는 장기 에이전트를 만들려면 세 가지가 필요하다고 본다. 첫째는 효과적인 압축 메모리인데, 이들에게는 KV 캐시 공간에서의 압축을 뜻한다. 둘째는 메모리를 루프 안에서 반복적으로 사용하는 능력으로, 이것이 긴 작업을 가능하게 한다. 셋째는 그 경험을 다시 가중치로 투영해 지속적인 지식으로 바꾸는 장치다.

KV 캐시 압축 방법을 그들은 두 축의 2×2 행렬로 정리한다. 한 축은 원본 KV의 일부를 고르는 '선택'과 역전파로 학습된 표현을 만드는 '합성'이고, 다른 축은 맥락마다 비용을 치르는 방식과 학습 단계에서 비용을 미리 치러 분산하는 방식이다. 맥락별 합성의 예로는 Cartridges 논문이, 사전학습 단계의 선택·합성 예로는 DeepSeek의 희소 어텐션이 언급된다.

이들의 방법 STILL은 스파스 오토인코더의 발상에서 출발한다. 입력마다 반복 최적화를 돌리는 대신, 모든 학습 데이터에 적용할 일반 함수를 학습해 추론 비용을 학습 단계로 분산하는 것이다. 같은 아이디어를 KV 캐시 압축에 적용해, 전체 KV 캐시를 압축 캐시로 내보내는 함수를 학습한다. 토큰 수가 가변적이므로 어텐션을 쓰며, 학습된 질의(query)가 캐시의 모든 토큰에 교차 어텐션해 일반적인 질문을 던지고, 전체 캐시를 쓴 모델과 압축 캐시를 쓴 모델의 출력이 같아지도록 KL 발산을 손실로 삼는다.

단일 압축을 반복하는 '무한 맥락'으로 확장하면 몇 번의 반복 뒤 붕괴하는 문제가 생긴다. 연구팀은 바로 다음 블록이 아니라 그다음 블록에서 KL 발산을 계산하도록 바꿔 16~32회 반복 압축에도 높은 정확도를 유지했다. 더 흥미로운 통찰은 압축된 KV 캐시가 사실상 MLP라는 점이다. 질의가 키와 곱해지고 소프트맥스 비선형을 거쳐 값과 다시 곱해지는 과정이 MLP의 업·다운 프로젝션과 같기 때문이다. 즉 압축은 그래디언트 디센트와는 다른 방식으로 맥락에서 가중치를 만들어 내며, 한 번 말한 사실을 그대로 가중치로 담는 표본 효율적인 학습의 길을 연다.

주요 인사이트

• 메모리를 무손실이냐 고압축이냐의 양극단으로만 보지 않고, 그 사이의 '중간 층'을 KV 캐시 압축으로 채우려는 발상이 장기 에이전트의 핵심 과제로 떠오른다.
• 스파스 오토인코더가 입력마다의 반복 최적화를 일반 함수 학습으로 대체했듯, KV 압축도 추론 비용을 학습 단계로 분산하면 '옵티마이저를 학습'하는 셈이 된다.
• 압축된 KV 캐시가 MLP와 같다는 등치는, 맥락을 곧 가중치로 바꾸는 새로운 가중치 생성 경로를 시사한다.
• 한 번 말한 내용을 캐시에 담아 곧바로 가중치로 만들 수 있다면, 여러 예시가 필요한 그래디언트 디센트보다 훨씬 표본 효율적인 학습이 가능해진다.

자주 묻는 질문

관련 AI 소식