ReasonAlloc: 추론 모델을 위한 계층적 디코딩 시간 키-값(KV) 캐시 예산 할당

ReasonAlloc: Hierarchical Decoding-Time KV Cache Budget Allocation for Reasoning Models

대규모 언어 모델(LLM)의 추론 과정에서 발생하는 긴 사고 과정(chain-of-thought, CoT)은 빠른 키-값(KV) 캐시 증가로 인해 심각한 추론 병목 현상을 야기합니다. 현재 디코딩 시간 압축 방법들은 토큰 제거를 통해 이 문제를 완화하지만, 일반적으로 모든 레이어와 헤드에 대해 균일한 예산 분포를 가정합니다. 반면, 기존의 비균일 예산 할당 방법은 주로 정적 프롬프트 사전 채우기 단계에 맞춰 설계되었으며, 자기 회귀 추론의 단계별 컨텍스트 요구 사항을 반영하지 못합니다. 이러한 격차를 해소하기 위해, 우리는 학습이 필요 없는 프레임워크인 ReasonAlloc을 제안합니다. ReasonAlloc은 디코딩 시간 KV 압축을 계층적 예산 할당 문제로 재구성합니다. ReasonAlloc은 두 가지 상호 보완적인 수준에서 작동합니다. 오프라인 레이어별 사전 할당 전략은 아키텍처 기반의 요구 패턴인 `` extit{Reasoning Wave}'를 파악하며, 온라인 헤드별 전략은 실시간 유용성을 기반으로 정보가 풍부한 헤드에 리소스를 재할당합니다. DeepSeek-R1-Distill-Llama-8B, DeepSeek-R1-Distill-Qwen-14B 및 AceReason-14B를 사용하여 수학적 추론 벤치마크(MATH-500, AIME~2024)에서 ReasonAlloc은 균일 예산 R-KV, SnapKV 및 Pyramid-RKV (정적인 단조 감소 레이어 예산을 강제하는 기준 모델)보다 성능이 우수하며, 특히 작은 예산(128-512 토큰)에서 가장 큰 성능 향상을 보입니다. ReasonAlloc은 기존의 토큰 제거 정책과 호환되며, 추론 시간에 미미한 오버헤드를 발생시킵니다.

Long chain-of-thought (CoT) trajectories in large language model (LLM) reasoning cause severe inference bottlenecks due to rapid key-value (KV) cache growth. Current decoding-time compression methods mitigate this issue via token eviction, but typically assume a uniform budget distribution across all layers and heads. In contrast, existing non-uniform budget allocation methods are predominantly designed for the static prompt prefill phase, and they do not capture the stepwise context demands of autoregressive reasoning. To bridge this gap, we propose ReasonAlloc, a training-free framework that recasts decoding-time KV compression as a hierarchical budget allocation problem. ReasonAlloc operates at two complementary levels: an offline layer-wise preallocation strategy captures an architecture-driven demand pattern which we call ``\textit{Reasoning Wave}'', while an online head-wise strategy reallocates resources during decoding to information-rich heads based on real-time utility. Evaluations on mathematical reasoning benchmarks (MATH-500, AIME~2024) using DeepSeek-R1-Distill-Llama-8B, DeepSeek-R1-Distill-Qwen-14B, and AceReason-14B show that ReasonAlloc outperforms uniform-budget R-KV, SnapKV, and Pyramid-RKV (a baseline enforcing a static, monotonically decreasing layer budget), with the largest gains at small budgets (128-512 tokens). ReasonAlloc is plug-and-play with existing token-eviction policies and introduces negligible inference-time overhead.