ZoomR: 다중 분해능 키-값 검색을 통한 메모리 효율적인 추론

ZoomR: Memory Efficient Reasoning through Multi-Granularity Key Value Retrieval

대규모 언어 모델(LLM)은 복잡한 추론 작업에서 뛰어난 성능을 보이지만, 종종 최종 답변에 도달하기 전에 긴 중간 단계를 생성해야 합니다. 생성 과정에서 LLM은 자기 회귀 디코딩을 위해 키-값(KV) 캐시에 의존합니다. 그러나 KV 캐시의 메모리 사용량은 출력 길이에 따라 증가합니다. 기존의 KV 캐시 최적화 연구는 주로 긴 입력 컨텍스트를 압축하는 데 중점을 두었지만, 디코딩에는 전체 KV 캐시를 유지합니다. 장문의 출력을 생성해야 하는 작업의 경우, 이는 계산 및 메모리 비용을 증가시킵니다. 본 논문에서는 LLM이 상세한 추론 단계를 요약본으로 적응적으로 압축하고, 이러한 요약본을 활용하는 동시에 세부 사항에 전략적으로 "집중"하는 동적 KV 캐시 선택 정책을 사용하는 새로운 접근 방식인 ZoomR을 소개합니다. ZoomR은 디코딩 과정에서 요약 키를 사용하여 전체 캐시에 대한 어텐션을 피하면서, 중요한 추론 단계에 대해서만 세부 정보를 검색합니다. 이러한 계층적 전략은 각 단계에서 전체 캐시에 대한 어텐션을 방지함으로써 메모리 사용량을 크게 줄입니다. 수학 및 추론 작업에서의 실험 결과는 제안하는 방식이 기준 모델과 경쟁력 있는 성능을 달성하면서, 추론에 필요한 메모리를 4배 이상 감소시킴을 보여줍니다. 이러한 결과는 다중 분해능 KV 선택이 특히 장문의 출력 생성을 위해 더 효율적인 디코딩을 가능하게 한다는 것을 입증합니다.

Large language models (LLMs) have shown great performance on complex reasoning tasks but often require generating long intermediate thoughts before reaching a final answer. During generation, LLMs rely on a key-value (KV) cache for autoregressive decoding. However, the memory footprint of the KV cache grows with output length. Prior work on KV cache optimization mostly focus on compressing the long input context, while retaining the full KV cache for decoding. For tasks requiring long output generation, this leads to increased computational and memory costs. In this paper, we introduce ZoomR, a novel approach that enables LLMs to adaptively compress verbose reasoning thoughts into summaries and uses a dynamic KV cache selection policy that leverages these summaries while also strategically "zooming in" on fine-grained details. By using summary keys as a coarse-grained index during decoding, ZoomR uses the query to retrieve details for only the most important thoughts. This hierarchical strategy significantly reduces memory usage by avoiding full-cache attention at each step. Experiments across math and reasoning tasks show that our approach achieves competitive performance compared to baselines, while reducing inference memory requirements by more than $4\times$. These results demonstrate that a multi-granularity KV selection enables more memory efficient decoding, especially for long output generation.