🤖 machine learning

Back from the Future: Key-Value Cache Management by Counter-Causal Surprise

이 논문은 과거의 토큰들이 미래의 문맥으로부터 얼마나 잘 예측될 수 있는지를 측정하기 위해 역인과적 주의(counter-causal attention)를 활용하여 중복된 토큰을 식별하고 제거함으로써, 다양한 거대 언어 모델 전반에서 경쟁력 있는 성능을 유지하면서도 메모리 사용량과 추론 지연 시간을 줄이는 학습이 필요 없는 Key-Value 캐시 제거 전략인 "Back from the Future"를 소개한다.

원저자: Stephen Gould, Anton van den Hengel

게시일 2026-07-31

📖 6 분 읽기🧠 심층 분석

CC BY 4.0

원저자: Stephen Gould, Anton van den Hengel

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 다음 이야기를 들려주기 위해 길고 복잡한 이야기를 기억하려고 노력하고 있다고 상상해 보세요. 당신의 뇌는 놀랍지만, 한 번에 담을 수 있는 양에는 한계가 있습니다. 이야기가 너무 길어지면, 새로운 내용을 담기 위해 일부 부분을 잊어야만 합니다. 이것이 바로 현대의 "대규모 언어 모델(Large Language Models)"(우리가 사용하는 초스마트 AI 챗봇들)이 직면한 문제입니다. 이 모델들은 지금까지 읽은 모든 것을 살펴보고 다음 단어를 추측하는 방식으로 작동합니다. 이를 빠르게 수행하기 위해, 이들은 컴퓨터 메모리에 키-값(Key-Value, KV) 캐시라고 불리는 "연습장"을 유지합니다. 이 캐시는 모델이 지금까지 읽은 이야기에서 가장 중요한 단서들을 적어두는 정신적 메모장과 같습니다.

문제는 이야기가 길어질수록 이 메모장이 점점 더 커진다는 것입니다. 결국, 이는 컴퓨터의 메모리를 가득 채워 AI를 충돌시키거나 속도를 매우 느리게 만듭니다. 과학자들은 이 문제를 해결하기 위해 메모장에 적힌 노트 중 무엇을 유지하고 무엇을 버릴지 결정하는 방법을 찾아내기 위해 노력해 왔습니다. 어떤 방식은 가장 오래된 노트를 그냥 버리기도 하고(슬라이딩 윈도우 방식처럼), 다른 방식은 모델이 가장 많이 "쳐다보는" 노트를 유지하려고 합니다. 하지만 이러한 기존 방식에는 결함이 있습니다. 만약 모델이 특정 노트를 너무 자주 본다면, 그 노트를 계속해서 더 많이 보게 되는 루프가 발생하여, 중요하지만 조용한 사실들이 충분히 "시끄럽지" 않다는 이유로 삭제되는 일이 생깁니다.

이 논문은 무엇을 남길지 결정하는 영리하고 새로운 방법인 **역인과적 놀라움(Counter-Causal Surprise)**을 소개합니다. 이 방법은 "모델이 무엇을 가장 많이 보았는가?"라고 묻는 대신, "만약 내가 이 노트를 가져간다면, 모델이 미래의 노트들을 보고 이 노트를 여전히 추측할 수 있을까?"라고 묻습니다. 만약 모델이 뒤따르는 단어들만 보고도 과거의 단어를 쉽게 추측할 수 있다면, 그 과거 단어는 그리 특별하지 않습니다. 즉, 중복된 정보라는 뜻입니다. 하지만 모델이 완전히 놀라며 미래의 단어들로부터 과거의 단어를 추측할 수 없다면, 그 단어는 독특하고 필수적인 정보를 담고 있는 것입니다. 저자들은 다양한 AI 모델을 대상으로 테스트를 진행했으며, 이 "놀라움" 방식이 기존 방식보다 메모리가 부족한 상황에서도 모델을 더 똑똑하고 정확하게 유지한다는 것을 발견했습니다. 또한, 계산을 훨씬 빠르게 수행하는 "패스트 모드(fast mode)"를 찾아내어, 속도를 크게 늦추지 않고도 실생활에서 사용할 수 있도록 만들었습니다.

"뒤를 돌아보는" 시선의 마법

그렇다면 이 "역인과적 놀라움"은 실제로 어떻게 작동할까요? AI가 추리 소설을 읽고 있다고 상상해 봅시다. 보통 모델은 일반적인 사람처럼 왼쪽에서 오른쪽으로 읽습니다. "집사"를 보고, 그다음 "집어 들었다"를 보고, 그다음 "양초"를 봅니다. 다음 단어를 예측하기 위해, 모델은 지금까지 본 모든 것을 사용합니다. 이것이 AI가 작동하는 표준 방식입니다.

하지만 메모리에서 무엇을 버릴지 결정하기 위해, 이 새로운 방법은 이상한 행동을 합니다. 바로 뒤를 돌아보는 것입니다. 모델은 이미 읽은 이야기의 한 덩어리를 가져와서 이렇게 묻습니다. "만약 내가 메모리에서 '양초'라는 단어를 숨긴다면, '집사가 양초를 집어 들었다...'라는 문맥만 보고도 여전히 '양초'를 추측할 수 있을까?"

낮은 놀라움 (버리기): 만약 문장이 "집사가 [양초]를 집어 들었다"였고, 그다음 단어들이 "그리고 주방으로 걸어갔다"라면, 모델은 주방에 있는 집사라는 문맥을 통해 "양초"를 쉽게 추측할 수 있을 것입니다. 모델이 이를 쉽게 추측할 수 있다면, 그 단어는 새로운 정보를 별로 더해주지 않은 것입니다. 이는 문장에서 "the"라는 단어를 기억하는 것과 같습니다. "the"는 어디에나 있기 때문에 이를 위해 특별한 메모를 남길 필요가 없습니다. 논문은 이러한 예측하기 쉬운 단어들을 캐시에서 삭제하여 공간을 절약할 것을 제안합니다.
높은 놀라움 (유지하기!): 이제 문장이 "집사가 [찻주전자]를 집어 들었다"라고 가정해 봅시다. 만약 그다음 단어들이 "그리고 주방으로 걸어갔다"라면, 모델은 완전히 혼란에 빠질 수 있습니다. "찻주전자라고?" 왜? 이건 놀라운 일입니다! 집사가 찻주전자를 집어 들었다는 사실은 미래의 단어들이 예측하지 못한 독특한 디테일입니다. 이 "놀라움"은 그 단어가 나머지 이야기에는 아직 알려지지 않은 비밀을 담고 있다는 것을 의미합니다. 논문은 이러한 "놀라운" 단어들이 가장 가치 있는 것이며, 메모리 캐시에 유지되어야 한다고 주장합니다.

"패스트 모드"라는 지름길

모든 긴 이야기의 모든 단어에 대해 이 "뒤를 돌아보는 작업"을 수행하는 것은 힘든 일입니다. 마치 책 한 권을 다 읽은 다음, 메모를 확인하기 위해 책을 다시 거꾸로 읽는 것과 같습니다. 저자들은 이것이 많은 컴퓨터 자원을 소모한다는 것을 깨달았습니다. 그래서 그들은 **패스트 싱글 레이어 근사치(Fast Single-Layer Approximation)**를 고안했습니다.

깊은 신경망(AI의 뇌)을 다층 건물이라고 생각해 보세요. 정보는 최종 답변이 나오기 전까지 여러 층(레이어)을 거쳐 올라갑니다. 전체 방법은 모든 층을 확인하여 무엇이 놀라운지 체크합니다. "패스트 모드"는 이렇게 말합니다. "이봐, 그냥 맨 꼭대기 층만 확인하자." 저자들은 AI 뇌의 마지막 레이어만 확인해도 전체 건물을 확인하는 것과 거의 동일한 결과를 얻을 수 있다는 것을 발견했으며, 이는 7~9배 더 빠릅니다.

테스트 결과, 이 빠른 버전은 512 토큰(작은 텍스트 덩어리)의 캐시를 새로고침하는 데 단 **7.9 밀리초(ms)**가 걸린 반면, 전체 체크 방식은 54 밀리초가 걸렸습니다. 4,096 토큰의 거대한 캐시에 대해서도, 빠른 버전은 52.6 밀리초에 불과했던 반면, 전체 버전은 496 밀리초가 걸렸습니다. 이는 이 방법이 실시간으로 사용되면서도 AI를 느릿하게 만들지 않도록 만드는 엄청난 속도 향상입니다.

실제로 효과가 있을까?

저자들은 단순히 상상만 한 것이 아니라, Qwen2.5 및 LLaMA 3.1과 같은 가장 똑똑한 오픈 소스 AI 모델들을 사용하여 복잡한 수학 문제 풀이, 긴 의료 기록 읽기, 긴 대화 따라하기와 같은 까м 까다로운 과제들을 테스트했습니다.

수학 문제: AI가 복잡한 수학 문제를 풀어야 하는 MATH500 벤치마크에서, 이 새로운 방법은 AI가 흐름을 놓치지 않게 하는 데 가장 뛰어났습니다. Qwen2.5-7B 모델의 경우, 새 방법은 **74.4%**의 정확도를 기록했습니다 (실제로 H2O가 76.2%로 약간 더 높았지만, 새 방법은 매우 근접했으며 3B 및 14B 버전에서는 더 나은 성능을 보였습니다). Llama-3.1-8B 모델에서 새 방법은 **48.2%**를 기록하여, 모든 "버리기" 방식 중 최고였으며, 완벽한 "제한 없음(no-limit)" 기준점인 **48.8%**에 매우 근접했습니다.
긴 대화: 이것은 기존 방식들이 정말 어려워했던 부분입니다. 긴 대화를 다루는 LoCoMo 데이터셋에서, 기존의 "어텐션 기반(attention-based)" 방식들(H2O 등)은 실패하기 시작했습니다. 그들은 대화 초기에 있었던 독특한 사실들을 버려버림으로써 혼란에 빠져 질문을 반복하거나 관련 없는 그림에 대해 이야기하곤 했습니다. 하지만 새로운 "역인과적" 방법은 이런 실수를 하지 않았습니다. 독특하고 놀라운 사실들을 유지함으로써, 긴 시간이 흐른 뒤에도 AI가 질문에 올바르게 답할 수 있게 해주었습니다.
사고 모드: 그들은 또한 AI가 답변하기 전에 오랫동안 "생각"해야 하는 AIME 수학 문제로도 테스트했습니다. 기존 방식들은 너무 많은 정보를 버려서 AI가 사고 과정을 제대로 마치지 못하고 혼란에 빠지는 경우가 많았습니다. 새 방법은 추론 체인을 훨씬 더 잘 유지하며 훨씬 더 높은 정확도를 달보았습니다.

이것이 왜 중요한가

핵심적인 결론은, 무엇을 잊을지 결정하는 기존 방식에 결함이 있었다는 것입니다. 기존 방식은 AI가 정보를 얼마나 많이 "쳐다보는지"에 의존했는데, 이는 인기 있는 단어는 더 인기 있게 만들고, 조용하지만 중요한 사실은 삭제해 버리는 편향을 만들었습니다. 이 새로운 방법은 이 판도를 뒤집습니다. "이 정보는 예측 가능한가?"라고 묻습니다. 만약 그렇다면, 그것은 필요하지 않습니다. 만약 놀랍다면, 그것은 금과 같습니다.

저자들은 이 접근 방식이 AI 모델을 다시 학습시킬 필요 없이, 현재 우리가 가진 모델들을 그대로 활용하면서도 메모리를 관리할 수 있는 원칙적이고 견고한 방법이라고 제안합니다. "놀라움"을 계산하는 데 약간의 추가 시간이 걸리긴 하지만, 패스트 버전은 매우 빠르기 때문에 AI를 느려지게 하지 않습니다. 이것은 마치 사람들이 찾는 책을 단순히 보관하는 것이 아니라, 아무도 예측할 수 없는 비밀을 담고 있는 책을 보관하여 이야기가 가장 중요한 반전을 잃지 않도록 보장하는 사서와 같습니다.

기술 요약: 미래로부터의 귀환: 역인과적 놀라움(Counter-Causal Surprise)을 통한 키-값(Key-Value) 캐시 관리

문제 정의

대규모 언어 모델(LLM)은 이전에 계산된 키(Key)와 값(Value)을 저장하여 자기 회귀적 생성 과정에서 효율적인 스케일드 닷 프로덕트 어텐션(scaled dot-product attention)을 가능하게 하는 키-값(KV) 캐시에 의존합니다. 그러나 KV 캐시의 크기는 컨텍스트 길이에 따라 선형적으로 증가하며, 긴 프롬프트나 확장된 생성 과정에서 가용 GPU 메모리를 빠르게 고갈시킵니다. 슬라이딩 윈도우(가장 오래된 토큰을 제거)나 압축과 같은 전략들이 존재하지만, 이들은 종종 중요한 정보를 유지하는 데 실패합니다.

기존의 동적 제거(dynamic eviction) 방식(예: H2O, TOVA)의 중대한 한계는 어텐션 점수에 의존한다는 점입니다. 이러한 방식은 자기 강화 편향(self-reinforcing bias) 문제를 겪습니다. 즉, 캐시에 유지되는 토큰은 후속 디코딩 단계에서도 계속해서 어텐션을 받게 되어 누적 점수가 부풀려지고, 이로 인해 제거될 가능성이 더욱 낮아집니다. 반대로, 사실 관계 측면에서 매우 중요하지만 단 한 번만 언급된 저빈도 토큰은 충분한 어텐션을 받지 못해 제거될 수 있으며, 이는 멀티홉 추론이나 정밀한 정보 검색이 필요한 작업에서 정확도 저하를 초래합니다.

방법론

저자들은 **역인과적 놀라움(Counter-Causal Surprise)**에 기반한 새로운 학습 불필요(training-free) 제거 전략을 제안합니다. 핵심 통찰은 만약 과거의 토큰이 뒤따르는 토큰들에 의해 잘 예측된다면, 그 토큰이 가진 정보는 중복된 것이므로 제거할 수 있다는 것입니다. 반대로, 미래의 문맥으로부터 예측하기 어려운 토큰은 독특하고 필수적인 정보를 포함하고 있으므로 유지해야 합니다.

역인사적 메커니즘(The Counter-Causal Mechanism)

이를 실행하기 위해, 저자들은 원래의 토큰 순서로 진행되지만 역인과적(상삼각, upper-triangular) 어텐션 마스크를 적용하는 스코어링 패스를 도입합니다.

과정: 모델은 표준 순방향 패스 순서로 토큰 $x_1, \dots, x_t$ 를 처리합니다.
마스킹: 표준 인과적 마스크(과거를 참조) 대신, 역인과적 마스크를 적용하여 위치 $i$ 가 오직 미래 위치 $j > i$ 만을 참조할 수 있도록 합니다.
재사용: 이 방법은 초기 순방향 패스에서 이미 저장된 KV 표현(key, value)을 재사용합니다. 상대적 오프셋을 정확하게 계산하기 위해 원래의 위치 ID를 사용하여 새로운 쿼리(query)를 생성합니다.
스코어링: 토큰 $x_i$ 에 대한 놀라움 점수 $s_i$ 는 $1 - P(x_i | x_{i+1:t})$ 로 정의됩니다. 실제 구현에서는 $x_i$ 의 로짓(logit)이 이 확률의 단조 근사치(monotone proxy)로 사용됩니다.
제거: 메모리 리프레시 주기 동안, 놀라움 점수가 가장 낮은(가장 예측 가능한) 토큰은 제거되고, 놀라움이 높은 토큰은 유지됩니다. 마지막 토큰은 미래 문맥이 없으므로 항상 유지됩니다.

빠른 단일 레이어 근사(Fast Single-Layer Approximation)

전체 역인과적 패스의 계산 비용( $L$ 이 레이어 수, $n$ 이 캐시 크기일 때 $O(Ln^2)$ ) 문제를 해결하기 위해, 저자들은 마지막 트랜스포머 레이어에만 역인과적 패스를 제한하는 빠른 근사법을 제안합니다.

이 방법은 초기 순방향 패스 중에 바로 이전 레이어의 은닉 활성화 값 $H^{(L-1)}$ 을 기록합니다.
리프레시 중에, 저장된 이 활성화 값들로부터 쿼리를 계산하고 마지막 레이어에서만 역인과적 마스크를 적용합니다.
이를 통해 복잡도를 $O(n^2)$ 로 줄여, 미미한 정확도 손실을 유지하면서 리프레시 주기당 7~9배의 속도 향상을 달면합니다.

시스템은 $h$ 개의 토큰을 처리하는 추론 단계와 캐시를 크기 $J$ 로 평가 및 프루닝(pruning)하는 메모리 리프레시 단계를 교대로 수행하는 **청크 단위 사이클(chunked cycles)**로 작동합니다.

주요 기여

역인과적 놀라움 지표: 미래 문맥으로부터의 예측 가능성을 바탕으로 토큰 유용성을 추정하는 원칙적이고 학습이 필요 없는 방법을 제공하여, 어텐션 기반 점수의 자기 강화 편향을 피합니다.
인-디스트리뷰션 스코어링(In-Distribution Scoring): 기존의 캐시된 표현을 재사용하고 원래의 토큰 순서로 실행되므로, 추가적인 학습이나 분포 외(out-of-distribution) 가정을 요구하지 않고 실제 캐시 내용과 직접 연결된 평가를 보장합니다.
효율적인 근사: 단일 레이어 근사법을 통해 리프레시 지연 시간을 획기적으로 줄이면서(512-토큰 캐시 기준 수백 밀리초에서 10ms 미만으로) 경쟁력 있는 성능을 유지합니다.
포괄적인 평가: 다양한 모델 제품군(Qwen2.5, LLaMA 3.1)과 다양한 작업(수학적 추론, 롱 컨텍스트 QA, 멀티홉 대화)에 걸친 광범위한 벤치마크를 수행했습니다.

실험 결과

방법론은 MATH500, AIME (thinking mode), LongHealth, Qasper, LoCoMo 벤치마크에서 평가되었습니다.

정확도: MATH500에서 역인과적 방식은 네 가지 모델 중 세 가지 모델(Qwen2.5-3B, 7B, 14B)에서 슬라이딩 윈도우, 중요도 샘플링(importance sampling), H2O보다 우수한 성능을 보였습니다. Llama-3.1-8B의 경우, 모든 제거 전략 중 가장 높은 48.2%의 정확도를 기록했으며, 이는 풀 캐시(full-cache) 베이스라인(48.8%)에 거의 근접한 수치입니다.
추론 일관성: 긴 사고 체인(chain of thought)을 요구하는 AIME "thinking mode" 벤치마크에서, 역인과적 방식은 제거 전략 중 가장 높은 정확도(36.7%)를 달 기록했으며, H2O(33.3%)와 Importance(26.7%)를 크게 앞질렀습니다. 또한 "pred=None"(추론 완료 실패) 발생률이 가장 낮아, 추론의 일관성을 더 잘 보존함을 보여주었습니다.
롱 컨텍스트 작업: LoCoMo(멀티홉 대화)에서 H2O는 작은 캐시 크기에서 질문을 반복하거나 무관한 이미지 캡션을 삽입하는 등 실패하는 모습을 보였는데, 이는 높은 어텐션 토큰에 치우친 편향 때문입니다. 역인과적 제거는 단 한 번 언급된 독특한 사실들을 성공적으로 유지하여 이러한 실패 모드를 방지했습니다.
효效率: 빠른 단일 레이어 근사법은 리프레시 지연 시간을 $n=4096$ 일 때 496ms에서 52.6ms로, $n=512$ 일 때 54ms에서 7.9ms로 단축했습니다. 리프레시 오버헤드에도 불구하고, 감소된 캐시 크기가 디코드 단계의 어텐션 비용을 낮추기 때문에 엔드 투 엔드 생성 시간은 풀 캐시 베이스라인과 유사한 수준을 유지했습니다.

의의 및 주장

본 논문은 역인과적 놀라움이 어텐션 기반 제거 전략에 내재된 자기 강화 편향에 강건한, 토큰 정보량에 대한 원칙적인 척도를 제공한다고 주장합니다. 미래 문맥을 통해 중복되는 토큰을 식별함으로써, 이 방법은 모델 재학습이나 복잡한 보조 네트워크 없이 효과적으로 메모리를 관리합니다.

저자들은 전체 역인과적 패스가 계산 오버헤드를 추가할 수 있지만, 빠른 단일 레이어 근사법이 미미한 정확도 저하를 유지하면서도 상당한 속도 향상을 달성하는 실용적인 절충안을 제공한다고 언급했습니다. 이 방법은 오토레그레시브 모델을 사용하는 모든 모델에 적용 가능한 일반적인 솔루션으로 제시되며, 특히 사실적 정확성과 추론 일관성을 유지하는 것이 중요한 롱 컨텍스트 생성 작업에 유용합니다. 결론적으로, 토큰 수준의 스코어링이 자연스럽지만, 개별 토큰은 쉽게 예측되더라도 구절(phrase) 전체가 중요한 경우를 다루기 위해 향-후 연구에서 더 거친 입도(구절 또는 문장 단위)의 전략을 탐색할 수 있다고 제언합니다.