추론 시점 하이퍼 스케일링 및 KV 캐시 압축

초록

추론 시간 스케일링은 더 길거나 더 병렬적인 시퀀스를 생성함으로써 효율성을 희생시키는 대신 추론 정확도를 높입니다. 그러나 트랜스포머 대형 언어 모델(LLM)에서는 생성 비용이 생성된 토큰의 수보다는 키-값(KV) 캐시의 크기에 의해 병목 현상이 발생합니다. 따라서 우리는 추론 시간 초스케일링을 탐구합니다: KV 캐시를 압축함으로써 동일한 컴퓨팅 예산 내에서 더 많은 토큰을 생성하고 스케일링된 추론의 정확도를 더욱 향상시킬 수 있습니다. 그러나 이 접근법의 성공은 높은 압축률에서도 정확도를 유지할 수 있는 압축 방법의 능력에 달려 있습니다. 초스케일링을 실용적으로 만들기 위해, 우리는 동적 메모리 희소화(DMS)라는 새로운 방법을 도입했습니다. 이 방법은 단 1,000번의 훈련 단계로 8배 압축을 달성하면서도 훈련이 필요 없는 희소 주의력보다 더 나은 정확도를 유지합니다. DMS는 캐시된 토큰을 조기에 버리는 대신 토큰 제거를 지연시켜 표현을 암묵적으로 병합하고 중요한 정보를 보존합니다. 우리는 DMS를 사용한 추론 시간 초스케일링의 효과를 여러 LLM 계열에서 입증하며, 이 방법이 비슷한 추론 실행 시간과 메모리 부하에서 정확도를 향상시킨다는 것을 보여줍니다. 예를 들어, 우리는 Qwen-R1 32B 모델을 AIME 24에서 평균 9.1점, GPQA에서 7.6점, LiveCodeBench에서 9.6점 향상시켰습니다.

English

Inference-time scaling trades efficiency for increased reasoning accuracy by generating longer or more parallel sequences. However, in Transformer LLMs, generation cost is bottlenecked by the size of the key-value (KV) cache, rather than the number of generated tokens. Hence, we explore inference-time hyper-scaling: by compressing the KV cache, we can generate more tokens within the same compute budget and further improve the accuracy of scaled inference. The success of this approach, however, hinges on the ability of compression methods to preserve accuracy even at high compression ratios. To make hyper-scaling practical, we introduce Dynamic Memory Sparsification (DMS), a novel method for sparsifying KV caches that only requires 1K training steps to achieve 8times compression, while maintaining better accuracy than training-free sparse attention. Instead of prematurely discarding cached tokens, DMS delays token eviction, implicitly merging representations and preserving critical information. We demonstrate the effectiveness of inference-time hyper-scaling with DMS on multiple families of LLMs, showing that it boosts accuracy for comparable inference runtime and memory load. For instance, we enhance Qwen-R1 32B by an average of 9.1 points on AIME 24, 7.6 on GPQA, and 9.6 on LiveCodeBench across compute budgets.

추론 시점 하이퍼 스케일링 및 KV 캐시 압축

Inference-Time Hyper-Scaling with KV Cache Compression

초록

Support