💻 computer science

Thought-Aware KV Cache Compaction for Reasoning via Adaptive Attention Matching

이 논문은 적응형 예산 할당과 핵심 토큰 보호를 통해 사고 사슬(chain-of-thought) 추론의 계층적 구조를 활용함으로써, 균등 압축 방식에 비해 정확도를 유지하거나 오히려 향상시키면서도 메모리 사용량을 크게 줄이는 새로운 KV 캐시 압축 방법인 Thought-Aware Attention Matching(TAM)을 제안한다.

원저자: Yang Liu, Bin Chong, Chongyang Zhang, Hao Zheng, Jiayu Liang, Xu Kefu

게시일 2026-08-14

📖 5 분 읽기🧠 심층 분석

CC BY 4.0

원저자: Yang Liu, Bin Chong, Chongyang Zhang, Hao Zheng, Jiayu Liang, Xu Kefu

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

당신은 아주 어려운 퍼즐을 풀려고 노력 중이라고 상상해 보세요. 하지만 당신에게는 매우 엄격한 규칙이 하나 있습니다. 바로 눈앞에 아주 적은 수의 포스트잇(sticky notes)만을 두어 단서들을 기억해야 한다는 것입니다. 작업을 진행하면서 당신의 뇌는 긴 생각의 사슬, 즉 "사고의 사슬(chain of thought)"을 생성하며, 여기에는 모든 단계, 모든 추측, 그리고 모든 막다른 길(dead end)을 기록합니다. 인공지능의 세계에서 이 "포스트잇"은 **KV 캐시(KV cache)**라고 불립니다. 이것은 컴퓨터가 지금까지 말한 모든 것을 기억하여 대화를 계속 이어갈 수 있게 해주는 방식입니다.

문제는 정말 똑똑한 AI 모델들이 어려운 수학 문제를 풀려고 할 때, 이 사고의 사슬이 믿기지 않을 정도로 길어진다는 점입니다. 포스트잇 더미가 너무 커져서 컴퓨터의 메모리가 바닥나고, 이로 인해 AI가 충돌하거나 속도가 현저히 느려지게 됩니다. 이를 해결하기 위해 과학자들은 캐시를 "압축(compact)"하는 방법을 시도해 왔습니다. 기본적으로 중요도가 낮은 메모를 버려서 공간을 확보하는 방식이죠. 하지만 여기에는 함정이 있습니다. 대부분의 기존 방식은 모든 메모를 똑같이 중요한 것처럼 취급합니다. 그들은 마치 방을 청소할 때 TV처럼 보이지 않는 것은 무엇이든 다 던져버리는 것처럼, 그냥 한 움큼의 메모를 집어 들고 나머지는 버려버립니다. 이 방식은 퍼즐을 풀기 위해 필요한 결정적인 단서들을 버리게 되어, 결국 AI를 혼란에 빠뜨리고 작업을 끝내지 못하게 만듭니다.

이 논문은 **사고 인식 어텐션 매칭(Thought-Aware Attention Matching, TAM)**이라 불리는 더 똑똑한 메모리 정리 방식을 소개합니다. TAM은 AI의 생각을 단순히 평평하고 지루한 단어의 목록으로 보는 대신, 추론에는 구조가 있다는 점을 깨달았습니다. 그것은 마치 이야기와 같습니다. 어떤 장(chapter)은 흥미진진한 반전과 핵심 사실을 담고 있는 반면, 어떤 장은 캐릭터가 숲속에서 길을 잃고 헤매는 과정일 뿐입니다. TAM은 어떤 부분이 "반전"이고 어떤 부분이 "헤매는 과정"인지 파악하여, 오직 헤매는 부분만을 버립니다. 이렇게 함으로써 가장 중요한 기억은 안전하게 지키면서 나머지 부분은 축소하여, AI가 메모리 부족 없이 복잡한 문제를 해결할 수 있도록 합니다.

문제: AI의 뇌에서 발생하는 메모리 누수

AI 모델이 수학 문제를 풀려고 할 때, 모델은 단순히 답을 내뱉는 것이 아닙니다. 모델은 "사고의 사슬(chain of thought)"이라고 알려진 긴 단계들을 생성하며 생각을 겉으로 드러냅니다. 이 생각을 추적하기 위해 모델은 KV 캐시라는 방대한 양의 데이터를 저장합니다. 이 캐시를 AI가 단어를 쓸 때마다 점점 무거워지는 배낭이라고 생각해 보세요. 만약 AI가 어려운 문제를 풀고 있다면, 배낭은 너무 무거워져서 컴퓨터의 메모리를 망가뜨리고 AI를 멈추게 할 수 있습니다.

과학자들은 이 배낭을 "압축"하여(즉, 일부 물건을 버려 가볍게 만들어) 이를 해결하려고 노력해 왔습니다. 하지만 이전의 방법들은 서투른 청소부와 같았습니다. 그들은 배낭을 보고 "좋아, 아이템의 10%만 남기고 나머지는 버리자"라고 말할 뿐, 그 아이템이 실제로 무엇인지에는 신경 쓰지 않았습니다. 그들은 결정적인 수학 공식과 아무 의미 없는 "음, 생각 좀 해볼게요"라는 멈춤을 똑같이 취급했습니다. 이러한 "균등한(uniform)" 접근 방식은 종종 가장 중요한 단서들을 버리게 하여, AI가 실수를 하거나 문제를 아예 풀지 못하게 만들었습니다.

해결책: 스마트한 사서

이 논문의 저자들은 **사고 인식 어텐션 매칭(Thought-Aware Attention Matching, TAM)**이라는 새로운 방법을 제안합니다. 서투른 청소부 대신, TAM은 어떤 책이 고전이고 어떤 것이 그저 오래된 잡지인지 정확히 아는 스마트한 사서처럼 행동합니다.

TAM은 AI의 사고 과정이 단순한 단어의 나열이 아니라 구조화된 여정임을 이해함으로써 작동합니다. TAM은 AI의 출력을 "사고 세그먼트(thought segments)"—마치 책의 장(chapter)처럼—로 나눕니다. 어떤 장은 필수적이지만(예: 문제 정의 또는 핵심 숫자 찾기), 어떤 장은 막다른 길(예: 잘못된 경로를 시도하다가 실패를 깨닫는 과정)입니다.

TAM이 마법을 부리는 방식은 다음 세 단계로 이루어집니다:

이야기 분절하기(Segmenting the Story): TAM은 AI의 출력을 살펴보고 아이디어 사이의 자연스러운 끊김을 찾습니다. TAM은 긴 사고의 사슬을 관리 가능한 덩어리로 나누기 위해, 줄바꿈 두 번(AI가 새 단락을 시작하는 지점)을 찾는 것과 같은 간단한 규칙을 사용합니다.
적응형 예산 책정(Adaptive Budgeting): 이것이 영리한 부분입니다. TAM은 "이 덩어리가 얼마나 중요한가?"라고 묻습니다. TAM은 AI의 현재 생각이 각 세그먼트에 얼마나 의존하는지를 측정합니다. 만약 어떤 세그먼트가 AI가 이미 지나쳐온 "막다른 길"이라면, TAM은 그것에 아주 작은 예산을 할당합니다. 즉, 세부 사항을 대폭 압축하여 대부분을 버립니다. 만약 어떤 세그먼트가 "핵심 앵커(key anchor)"(예: 원래의 문제 문구)라면, TAM은 그것에 거대한 예산을 할당하여 거의 모든 세부 사항을 안전하게 보관합니다. 이것은 여행을 위해 짐을 싸는 것과 같습니다. 여권과 지갑은 안전하게 보관하지만, 양말이나 티셔츠는 공간을 아끼기 위해 꾹꾹 눌러 담는 것과 같습니다.
앵커 보호하기(Protecting the Anchors): 때때로 특정 단어들은 너무 중요해서 절대 건드려서는 안 됩니다. TAM은 AI가 계속해서 되돌아보는 단어들, 즉 상수(constant)나 결정적인 정의와 같이 지속적으로 참조하는 "중추적인 토큰(pivotal tokens)"을 식별하고, 이를 삭제할 수 없는 특별한 안전 구역에 잠금 처리합니다.

연구 결과: 단순히 작아지는 것이 아니라 더 똑똑하게

연구진은 이 새로운 방법을 두 가지 어려운 수학 벤치마크인 AIME 2024(30개의 어려운 문제로 구성된 경연 대회)와 MATH-500(500개의 문제 세트)에서 테스트했습니다. 그들은 TAM이 기존의 방식보다 적은 메모리를 사용하면서도 이 문제들을 해결할 수 있는지 확인하기 위해 Qwen3-4B라는 모델을 사용했습니다.

결과는 유망했습니다. TAM을 기존의 "균등한(uniform)" 방식(무작위로 덩어리를 버리는 방식)과 비교했을 때, TAM은 일관되게 더 높은 점수를 얻었습니다.

AIME 2024 테스트에서, 기존의 균등한 방식은 약 **56.7%**의 정답률을 보였으나, TAM은 이를 **60.0%**로 향상시켰습니다.
MATH-500 테스트에서, 균등한 방식은 **64.6%**를 기록한 반면, TAM은 **67.8%**에 도달했습니다.

더욱 인상적인 것은 메모리 절감 효과였습니다. (끝까지 기다리지 않고 매 1,024단어마다 메모리를 정리하는) "주기적(periodic)" 버전의 TAM을 사용했을 때, 피크 메모리 사용량을 3.1~3.2 GB로 줄일 수 있었습니다. 이는 압축을 전혀 하지 않았을 때 사용될 약 9.2 GB와 비교하여 65%의 감소입니다. 결정적으로, 그들은 AI의 정확도를 높게 유지하면서도 메모리 점유율을 훨씬 작게 만들었습니다.

트레이드오프와 한계

논문은 또한 이 "스마트한 청소"에 시간이 얼마나 걸리는지도 살펴보았습니다. 연구진은 어떤 부분이 중요한지 판단하기 위해 TAM이 수행하는 추가 작업이 매우 빠르며, 프로세스에 추가되는 시간은 약 0.15초에 불과하다는 것을 발견했습니다. 이는 텍xt 자체를 생성하는 데 걸리는 시간에 비하면 아주 미미한 수준입니다.

하지만 저자들은 이것이 모든 상황에 적용되는 마법의 해결책은 아니라는 점을 주의 깊게 명시하고 있습니다. 그들의 방법은 AI의 출력이 명확한 구조(예: 단락)를 가지고 있다는 것에 의존합니다. 만약 AI의 사고 과정이 무질서하고 명확한 끊김이 없거나, 앞뒤로 혼란스럽게 왔다 갔다 한다면 TAM은 올바른 세그먼트를 찾는 데 어려움을 겪을 수 있습니다. 또한, 이들은 특정 모델 하나를 대상으로 수학 문제만을 테스트했습니다. 결과는 강력하지만, 이것이 이야기를 쓰거나 소프트웨어를 코딩하는 데도 동일한 방식으로 작동할지, 혹은 훨씬 더 큰 규모의 AI 모델에서도 작동할지는 아직 알 수 없습니다.

요약하자면, 이 논문은 AI의 생각을 무질서한 단어의 더미가 아닌 구조화된 이야기로 취급함으로써, 명확한 사고 능력을 잃지 않으면서도 엄청난 양의 메모리를 절약할 수 있음을 시사합니다. 이는 스마트한 AI 모델이 가장 중요한 여정을 잊지 않으면서도 더 작고 저렴한 컴퓨터에서 실행될 수 있도록 만드는 한 단계입니다.

기술 요약: 적응형 어텐션 매칭(TAM)을 통한 추론용 사고 인지형 KV 캐시 압축

1. 문제 정의

추론 능력을 갖춘 언어 모델은 최종 답변을 생성하기 전 긴 사고 사슬(Chain-of-Thought, CoT) 시퀀스를 생성합니다. 이러한 토큰들의 어텐션 상태를 저장하기 위해 필요한 키-값(Key-Value, KV) 캐시는 시퀀스 길이에 따라 선형적으로 증가하며, 이는 디코딩 처리량을 제한하고 자원이 제한된 장치에 추론 모델을 배포하는 데 있어 심각한 메모리 병목 현상을 초래합니다.

기존의 KV 캐시 압축 방식들은 추론 워크로드에 적용될 때 두 가지 주요 한계에 직면합니다:

균일한 압축 (Uniform Compression): 대부분의 방법(예: 제거 기반 또는 잠재 공간 최적화)은 추론 궤적을 평면적인 토큰 시퀀스로 취급하여 균일한 압축 정책을 적용합니다. 이는 CoT 추론의 계층적 구조(예: 문제 재진술, 중간 계산, 탐색적 막다른 길, 결론 등)를 간과하며, 각 단계가 미래의 생성에 기여하는 중요도가 극명하게 다르다는 점을 무시합니다.
중간 궤적 처리의 미흡 (Inadequate Mid-Trajectory Handling): 일부 방법은 긴 입력 프롬프트(prefill)를 다루지만, 디코딩 중 발생하는 생성된 토큰의 효율적인 압축(mid-trajectory)은 여전히 충분히 연구되지 않았습니다. 기존의 중간 궤적 접근 방식은 단순한 토큰 제거(eviction)에 의존하는 경우가 많으며, 이는 높은 압축률에서 "어텐션 질량(attention mass, 총 어텐션 가중치)"을 체계적으로 과소평가하여 상당한 정확도 저하를 야기합니다.

2. 방법론: 사고 인지형 어텐션 매칭 (Thought-Aware Attention Matching, TAM)

저자들은 균일한 압축을 계층적 프로세스로 변환하는 구조 인식 압축 프레임워크인 **사고 인지형 어텐션 매칭(TAM)**을 제안합니다. TAM은 어텐션 매칭(AM) 최적화 파이프라인 상단에서 작동하며 세 가지 핵심 메커니즘을 활용합니다.

A. 사고 세분화 (Thought Segmentation)

TAM은 추론 궤적을 평면적인 스트림으로 취급하는 대신, 일관된 추론 블록(세그먼트)으로 분해합니다.

휴리스틱 접근법: 주요 방법은 추론 모델이 의미적 단계를 구분하기 위해 사용하는 더블 뉴라인(double-newline) 문자를 사용하여 구조적 경계를 감지합니다. 과도한 파편화를 방지하기 위해 짧은 세그먼트들은 병합됩니다.
대안적 접근법: 어텐션 강도 기울기가 크게 변화하는 경계 지점을 감지하는 어텐션 기반 방법이 있으나, 속도와 견고함 측면에서 휴리스틱 접근법이 선호됩니다.

B. 적응형 예산 할당 (Adaptive Budget Allocation)

세분화가 완료되면, TAM은 각 세그먼트에 할당되는 키(key)의 수를 중요도와 크기에 따라 균등하게 배분하지 않고 결정합니다.

중요도 점수 산출: 세그먼트의 중요도( $w_i$ )는 "선택 윈도우(selector window, 가장 최근의 $R$ 개 토근)"로부터 받는 평균 어텐션 질량으로 계산되며, 이는 미래의 쿼리에 대한 경량화된 대리 지표 역할을 합니다.
할당 규칙: 세그먼트 $i$ 에 대한 압축된 키의 수( $t_i$ )는 다음과 같이 할당됩니다:
$t_i \propto \sqrt{w_i} \cdot n_i$
여기서 $n_i$ 는 세그먼트 크기입니다. 이론적 분석(Proposition C.1)에 따르면, 볼록 오차 모델(convex error model) 하에서 이 특정 가중치 방식은 전체 가중치 근사 오차를 최소화합니다. 이를 통해 중요도가 높은 큰 세그 dt는 더 많은 토큰을 유지하고, 관련 없는 "막다른 길" 탐색은 공격적으로 압축합니다.

C. 핵심 토큰 보호 (Pivotal Token Protection)

TAM은 추론의 앵커 역할을 하는 특정 "핵심 토큰"(예: 문제의 상수, 중요한 정의, 주요 중간 결과)을 식별하고 보호합니다.

식별: 선택 윈도우로부터 지속적으로 높은 어텐션을 받는 토큰(임계값 $\delta = c \cdot \bar{\alpha}_{mean}$ 이상)은 핵심 토큰으로 표시됩니다.
보호: 이 토낀들은 편향 수정 없이 그대로(zero bias modification) 압축된 캐시에 유지됩니다. 나머지 예산은 비핵심 토큰들에게 분배됩니다. 이론적 분석(Proposition C.2)은 이 방식이 보호된 토큰의 어텐션 질량에 비례하여 최악의 경우 근사 오차를 줄임을 보여줍니다.

D. 주기적 압축 (Periodic Compaction)

피크 메모리를 제한하기 위해, TAM은 최대 시퀀스 길이에 도달했을 때만 실행되는 것이 아니라 주기적으로(매 $P$ 개의 새로운 토큰마다) 트리거될 수 있습니다. 이를 통해 캐시 크기가 목표 크기 $t$ 와 $t+P$ 사이에서 진동하도록 하여, 긴 생성 과정 동안 메모리가 무한히 증가하는 것을 방지합니다.

3. 주요 기여

TAM 프레임워크: 추론 궤적을 세분화하고, 중요도와 크기에 따라 예산을 적응적으로 할당하며, 핵심 추론 앵커를 보호하는 구조 인식형 중간 궤적 압축 방법론을 도입했습니다.
이론적 토대:
- 할당 규칙( $t_i \propto \sqrt{w_i} \cdot n_i$ )이 볼록 오차 모델 하에서 최적임을 증명했습니다.
- 핵심 토큰 보호가 근사 오차를 감소시킨다는 것을 입증했습니다.
- 순차적(주기적) 압축 하에서의 누적 오차가 유계(bounded)임을 증명했습니다.
실증적 검증: Qwen3-4B를 사용하여 수학적 추론 벤치마크(AIME 2024, MATH-500)에서 평가를 진행하였으며, 동일한 메모리 풋프린트에서 균일 압축 베이스라인 대비 일관된 정확도 향상을 입증했습니다.

4. 실험 결과

Qwen3-4B 모델을 사용하여 AIME 2024(30문제) 및 MATH-500(500문제)에 대해 실험을 수행했습니다.

정확도 vs 메모리: 목표 압축 비율이 0.1(KV 엔트리의 10% 유지)일 때, TAM은 AIME 2024에서 60.0%, MATH-500에서 **67.8%**의 정확도를 달성했습니다. 이는 균일 어텐션 매칭 베이스라인인 PAM(각각 56.7%, 64.6%)보다 우수한 성능입니다.
메모리 감소: TAM (Periodic) 설정(1024 토큰마다 트리거)은 피크 메모리를 3.1–3.2 GB로 제한하였으며, 이는 "압축 없음(No Compaction)" 베이스라인(AIME 기준 9.2 GB) 대비 65%의 감소를 나타내면서도 경쟁력 있는 정확도를 유지했습니다.
제거(Eviction) 방식과의 비교: 단순 제거 기반 방식은 상당한 정확도 저하(예: AIME에서 46.7%)를 겪었으며, 이는 제거 전용 접근 방식의 질량 결핍(mass deficit)이라는 이론적 한계를 확인시켜 줍니다.
절제 연구 (Ablation Studies):
- 적응형 할당과 핵심 토큰 보호 모두가 독립적으로 성능 향상에 기여했습니다.
- 휴리스틱 세분화(더블 뉴라인)는 어텐션 기반 세분화만큼 정확하면서도 훨씬 빨랐습니다(45ms 대비 2ms).
- 압축 비용(4k-토큰 궤적 기준 스텝당 약 5.6초)은 전체 추론 시간에 비해 무시할 수 있는 수준이었습니다.

5. 의의 및 주장

본 논문은 추론은 평면적인 시퀀스가 아니라 계층적인 과정임을 인정함으로써, TAM이 추론 모델 배포의 핵심적인 격차를 해결한다고 주장합니다. CoT 궤적의 구조적 이질성을 활용함으로써, TAM은 균일 압축 방식보다 메모리 사용량과 추론 정확도 사이에서 더 효율적인 트레이드오프를 달상합니다.

저자들은 다음을 강조합니다:

구조가 중요하다: 추론의 계층적 특성을 무시하면 최적이 아닌 압축과 정확도 손실이 발생합니다.
최적성: 제안된 할당 규칙은 이론적으로 근거가 있으며, 특정 가정 하에서 오차를 최소화하는 데 최적입니다.
실용성: 주기적 압축은 복잡한 수학 문제를 해결하는 모델의 능력을 희생하지 않으면서도, 엄격한 메모리 제약이 있는 장치(예: 피크 메모리 65% 감소)에 추론 모델을 배포할 수 있게 합니다.

인정된 한계점: 논문은 TAM이 명확한 구조적 경계가 없거나 복잡한 백트래킹을 포함하는 추론 궤적에서는 실패할 수 있는 휴리스틱 세분화에 의존한다는 점을 언급합니다. 또한, 실증적 검증이 현재 수학적 추론 벤치마크와 단일 모델 아키텍처에 국한되어 있음을 밝히고 있습니다.