Random Attention으로 KV 캐시 점수화를 제거한 추론 압축

왜 중요한가

긴 Chain-of-Thought는 생성 토큰 수만큼 KV 캐시를 키워 메모리 사용량을 압박합니다. Random Attention은 토큰별 중요도 점수 없이 Prompt 전체를 보존하고 나머지를 Head별 무작위로 샘플링해, 네 모델과 여섯 추론 과제에서 강한 기존 Evictor와 비슷한 정확도를 유지하면서 vLLM 처리량을 32–43% 높였습니다.

핵심 기여

Prompt 보존과 무작위 Eviction 결합

System Prompt, Chat Template, Question을 포함한 전체 Prefill 구간을 삭제하지 않고, 나머지 생성 흔적에는 KV Head별 독립 균등 무작위 점수를 적용해 상위 K개를 유지하는 Random Attention을 구성했습니다.

선택 점수보다 Prompt 보호가 중요

SnapKV, R-KV, VaSE, TriAttention을 같은 Prompt 보존 규칙 아래 비교한 결과, 기존 방법 간 정확도 차이의 상당 부분이 선택 점수 자체가 아니라 Prompt를 얼마나 남겼는지에서 발생했습니다.

추론 흔적의 이중 중복성 확인

모델은 필요한 중간값을 텍스트에서 반복해 쓰고 각 KV Head에 복사본을 저장했습니다. 실험에서 한 Head에만 남은 사실의 회수율은 낮았지만 두 Head에서는 60%, 세 Head에서는 83%, 여덟 Head에서는 99%로 증가했습니다.

서빙 처리량 개선

한 H200의 vLLM과 PagedAttention 환경에서 32k 토큰 생성 시 Random Attention은 TriAttention보다 32–43% 높은 출력 토큰 처리량을 기록했으며, 별도 점수 계산 없이 캐시 압축만 수행했습니다.

핵심 아이디어 이해하기

Transformer의 Attention은 현재 Query와 캐시된 모든 Key를 비교해 Value를 가중합합니다. 긴 추론에서는 생성되는 중간 토큰마다 Key-Value 쌍이 추가되어 KV 캐시가 길이에 비례해 커지고, 고정 예산을 넘으면 일부 쌍을 영구 삭제해야 합니다. 기존 Evictor는 이후에 필요할 토큰을 예측해 점수가 높은 항목을 남기지만, 한 번 삭제한 항목은 다시 읽을 수 없습니다.

이 논문은 캐시를 Prompt와 모델이 생성한 추론 흔적으로 나눕니다. Prompt는 질문이 한 번만 나오므로 삭제되면 복구할 수 없지만, 추론 흔적은 모델이 현재 사용하는 중간값을 텍스트에서 반복하고 여러 Attention Head에 복사합니다. 따라서 먼저 Prompt 전체를 보호하면 나머지 영역에서는 특정 토큰을 정밀하게 순위화하지 않아도 필요한 정보의 복사본이 일부 남습니다.

Random Attention은 Prompt 위치에 무한대 점수를 부여해 항상 보존하고, 나머지 위치에는 0과 1 사이의 독립 균등 난수를 배정합니다. 각 KV Head가 별도로 상위 K개를 선택하므로 서로 다른 Head가 서로 다른 추론 토큰을 남깁니다. 여러 Head에 남은 복사본을 합치면 단일 Head에서 사라진 중간값도 모델이 회수할 가능성이 커집니다.

이 구조가 모든 정보를 보존한다는 뜻은 아닙니다. 한 번만 언급되고 다시 쓰이지 않는 Passcode처럼 중복성이 없는 정보는 무작위 보존으로 잃기 쉽습니다. 다만 실제 추론 흔적에서는 반복 기록과 Head 간 복사로 중복성이 충분한 경우가 많아, 선택 점수의 추가 이득이 작아집니다.

방법론

Decode 단계에서 각 Attention Head는 캐시된 Key-Value 쌍을 고정 예산 K개와 최근 버퍼 r개로 관리합니다. 버퍼가 차면 최근 버퍼 밖의 후보 위치마다 Eviction 점수를 계산하고 상위 K개를 남깁니다. Random Attention은 이 점수 계산을 없애고 Prompt 길이 ℓp까지의 위치에는 무한대 점수, 그 이후 위치에는 독립 균등 난수를 할당한 뒤 Head별 top-K를 적용합니다.

기존 방법과의 비교에는 SnapKV, R-KV, VaSE, TriAttention을 사용했습니다. Qwen3-4B, Qwen3-14B, Qwen3-32B, Phi-4-reasoning에서 MATH500, GPQA-Diamond, AIME, HMMT, LiveCodeBench-v6를 평가했으며, 일반 과제는 약 4배, LiveCodeBench는 약 3배 압축 조건을 사용했습니다. 생성 최대 길이는 32,768토큰이고, 정확도는 최종 정답의 일치 여부로 계산했으며 paired clustered bootstrap 95% 신뢰구간과 exact sign test를 함께 적용했습니다.

Prompt의 영향을 분리하기 위해 각 방법에 동일한 Prompt 보호 규칙을 추가하고 보호 전후 정확도를 비교했습니다. 추론 흔적의 Head 간 중복성은 MATH500 생성 기록에 새로운 변수와 값을 심은 뒤 1,536토큰 후 질문을 붙이는 Planted-Fact Probe로 측정했습니다. 사실을 보존한 Head 수를 바꾸고 정답값 회수율과 Log-Probability 기반 Recall을 비교해 복사본 수와 정보 회수의 관계를 확인했습니다.

서빙 효율은 vLLM의 PagedAttention과 HuggingFace Transformers의 FlashAttention-2 환경에서 측정했습니다. vLLM 실험은 한 H200, 1k 토큰 Prompt, 32k 토큰 생성, K=2048 조건에서 처리량을 비교했고, 별도 점수 계산에 필요한 추가 캐시 순회와 동기화 비용을 측정했습니다.

주요 결과

주요 정확도 표의 60개 Baseline 비교 중 Random Attention은 31개에서 유의하게 앞섰고 한 비교에서만 유의하게 뒤처졌습니다. MATH500과 GPQA-Diamond에서는 Qwen3 및 Phi-4-reasoning 전반에 걸쳐 SnapKV와 VaSE보다 유의하게 높았으며, 어떤 Selector도 두 과제에서 Random Attention을 유의하게 앞서지 못했습니다.

LiveCodeBench에서는 Prompt 길이가 평균 557토큰으로 MATH500보다 6배 길고, 최대 길이는 K=3072 예산의 절반까지 차지했습니다. 이 조건에서 SnapKV는 모든 모델에서 Random Attention보다 20–35포인트 낮았고, TriAttention은 Qwen3-32B에서 약 3포인트 앞서 유일하게 유의한 Baseline 승리를 기록했습니다. 저자들은 이 차이를 선택 점수보다 긴 Prompt를 보존하는 비용과 연결했습니다.

Prompt를 공통으로 보호하자 SnapKV는 Phi-4-reasoning GPQA-Diamond에서 22.5포인트, Random Attention은 같은 조건에서 24.4포인트 상승했습니다. 보호 이후 세 Baseline은 모든 설정에서 서로 2.2포인트 이내에 모였고, 이는 원래의 큰 격차가 Prompt 생존 여부에서 비롯되었음을 뒷받침했습니다.

vLLM에서 Random Attention은 Full Attention 대비 1.58–2.67배 처리량을 기록했고 TriAttention보다 32–43% 높았습니다. 한 번의 Eviction 비용은 Random Attention이 0.30ms, TriAttention이 1.47–1.64ms였으며, 128개 동시 요청에서는 반복되는 점수 계산과 배치 동기화가 누적되어 차이가 커졌습니다.

기술 상세

Attention Head가 현재 Query q_t와 캐시된 Key k_i를 비교할 때, 위치 i의 Attention Weight는 원문 식에 따라 $\alpha_i^{(t)} = \frac{\exp(q_t^\top k_i / \sqrt{d})}{\sum_{j\leq N}\exp(q_t^\top k_j / \sqrt{d})}$ 로 계산됩니다. q_t와 k_i는 d차원 벡터이고 N은 캐시된 위치 수이며, 각 위치의 Weight를 Value v_i에 곱해 더한 $o_t = \sum_{i\leq N}\alpha_i^{(t)}v_i$ 가 현재 Head의 출력이 됩니다. Eviction은 이 계산에 참여할 Key-Value 쌍을 미리 줄이는 과정입니다.

기존 정책은 이 위치에 중요도 점수 s_i를 부여해 상위 K개를 남깁니다. H2O는 과거 Attention Weight의 누적합, SnapKV는 최근 w개 Query에서 얻은 Attention, R-KV는 SnapKV 점수와 Key 간 평균 Cosine Similarity, VaSE는 Value 좌표 범위와 SnapKV 기반 무작위 샘플링, TriAttention은 현재 위치와 Query의 거리를 사용합니다. Random Attention은 Prompt 위치에 +∞, 나머지 위치에 Uniform(0,1) 난수를 배정해 점수의 내용 의존성을 제거합니다.

Algorithm 1의 입력은 배치 크기 B, KV Head 수 H_kv, 캐시 위치 수 S, Prompt 길이 ℓp, 예산 K입니다. rand(B, H_kv, S)가 Head와 위치별 난수를 만들고, Prompt 구간을 +∞로 덮어쓴 뒤 topk(s, K)가 각 Head에서 K개 위치를 반환합니다. 따라서 한 Eviction 이벤트의 핵심 비용은 난수 생성과 top-k이며, 모델별 Calibration이나 Attention 통계 재계산이 필요하지 않습니다.

Head별 독립 추출은 추론 흔적의 복사본을 분산시키는 역할을 합니다. Qwen3-4B Planted-Fact Probe에서 한 Head만 사실을 보존했을 때 최고 회수율은 3%였지만, 두 Head는 60%, 세 Head는 83%, 여덟 Head는 99%의 회수율을 기록했습니다. 이는 단일 Head의 완전한 보존보다 여러 Head에 일부 복사본을 남기는 방식이 회수 가능성을 높인다는 뜻입니다.

반대로 중복성이 없는 정보에는 점수 기반 정책이 유리합니다. 57회의 압축 라운드 전에 한 번만 제시한 Passcode를 대상으로 Random Attention의 회수율은 0%, R-KV는 84%였습니다. R-KV는 전체 이력의 Attention 누적값을 사용해 이런 단일 사실을 찾았지만, 실제 추론 표에서는 TriAttention이 더 강한 결과를 냈고 Passcode Probe의 회수율은 낮아 과제별 강점이 달랐습니다.

한계점

논문은 긴 Code Prompt 전체를 고정하면 K=3072 예산의 큰 비율을 Prompt가 차지한다고 지적합니다. 또한 한 번만 언급되고 다시 반복되지 않는 정보는 Random Attention이 잃기 쉬워, Prompt 내부를 압축하거나 희귀 정보를 보존하는 선택 신호가 후속 과제로 남았습니다.

실무 활용

긴 추론을 수행하는 LLM을 GPU 메모리 예산 안에서 서빙할 때 적용할 수 있는 KV Cache Eviction 정책입니다. Prompt 전체를 보존한 뒤 생성 흔적을 Head별 무작위로 축약하므로 Calibration과 별도 중요도 모델 없이 구현할 수 있으며, 공개 GitHub 코드가 제공됩니다.

vLLM에서 32k 토큰 수준의 긴 추론 요청을 여러 개 동시 서빙하는 환경
GPU 메모리 한도 때문에 KV 캐시를 고정 예산으로 유지해야 하는 Reasoning Model 배포
새로운 KV Cache Selector의 효과를 비교하는 Null Hypothesis Baseline
Prompt 보존과 추론 흔적 압축을 분리해 실험하는 KV Cache 연구

코드 공개 여부: 공개

코드 저장소 보기

키워드

KV Cache 압축Random AttentionAttention Head추론 흔적Prompt 보존vLLM 처리량KV Cache Eviction

용어 해설

KV 캐시(KV Cache): — Transformer가 이전 토큰의 Key와 Value를 저장해 다음 토큰 생성 때 재사용하는 메모리 구조입니다. 생성 토큰이 늘면 저장 항목도 선형으로 증가하므로 긴 추론에서는 GPU 메모리 병목이 됩니다. KV 캐시 Eviction은 예산을 넘은 항목을 영구 삭제해 메모리를 제한합니다.
KV 캐시 Eviction(KV Cache Eviction): — 고정된 메모리 예산 안에 KV 캐시를 유지하려고 일부 토큰의 Key-Value 쌍을 삭제하는 기법입니다. 각 토큰의 향후 중요도를 점수화한 뒤 상위 항목을 보존하는 방식이 일반적이지만, 삭제된 항목은 이후 복구할 수 없습니다.
Attention Head: — Multi-Head Attention을 구성하는 독립적인 부분 모듈입니다. 각 Head는 토큰의 Key-Value 복사본을 별도로 저장하고 현재 Query와의 관련도를 계산합니다. 이 논문에서는 Head별 독립 보존이 추론 흔적의 중복성을 높이는 핵심 구조로 작동합니다.
추론 흔적(Reasoning Trace): — 모델이 문제를 풀면서 생성하는 중간 추론 토큰의 연속입니다. 논문에서는 모델이 현재 필요한 중간값을 텍스트 안에서 반복해 적고 여러 Attention Head에 복사해 저장하므로, 일부 캐시 항목이 삭제되어도 정보가 남는다고 정리합니다.
PagedAttention: — 긴 생성 과정의 KV 캐시를 페이지 단위로 관리해 여러 요청의 GPU 메모리 사용을 조정하는 방식입니다. 이 논문에서는 vLLM의 서빙 환경에서 캐시 압축과 함께 사용하며, 캐시 용량이 늘어나는 요청 수와 처리량에 직접 영향을 줍니다.

Random Attention으로 KV 캐시 점수화를 제거한 추론 압축

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

Grouped Query Attention의 작동 원리와 KV cache 최적화

KV cache를 검색 가능한 벡터 공간으로 라우팅

DeepSeek-V2의 핵심 기술인 Multihead Latent Attention(MLA)의 원리와 구현.

vLLM의 핵심 기술인 PagedAttention의 작동 원리와 메모리 효율화 방식

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

Grouped Query Attention의 작동 원리와 KV cache 최적화

KV cache를 검색 가능한 벡터 공간으로 라우팅

DeepSeek-V2의 핵심 기술인 Multihead Latent Attention(MLA)의 원리와 구현.

vLLM의 핵심 기술인 PagedAttention의 작동 원리와 메모리 효율화 방식