더 적은 토큰, 더 작은 캐시: 보상 기반 효율적인 추론

Fewer Tokens, Smaller Cache: Reward-Coordinated Efficient Reasoning

대규모 추론 모델(LRM)은 긴 연쇄적 사고(CoT) 방식을 통해 복잡한 작업에서 뛰어난 성능을 보이지만, 이러한 방식의 긴 중간 단계는 과도한 계산으로 인해 추론 비용을 증가시킵니다. KV-캐시 압축은 일반적인 해결책이지만, 기존의 추론 관련 방법들은 전체 경로에 대해 균일한 정책을 적용하고 캐시에서 제거되는 정보만을 기준으로 압축 여부를 판단합니다. 그러나 두 가지 관찰 결과는 다른 가능성을 제시합니다. 첫째, 추론 과정 중 특정 단계에서의 문맥 손실 허용 범위가 다르며, 보상(reward)은 이러한 변화를 반영합니다. 즉, 높은 보상을 받는 단계에서 토큰을 삭제하는 것은 무작위로 동일한 양의 토큰을 삭제하는 것보다 정확도를 훨씬 더 잘 유지합니다. 둘째, 압축은 생성 과정에도 영향을 미치므로, 캐시 크기가 작아지면 모델이 더 많은 토큰을 생성하게 되어, 저장 공간 절감 효과를 부분적으로 상쇄할 수 있습니다. 이러한 점들을 고려하여, 본 연구에서는 보상 기반으로 전체 과정을 조율하는 방법을 제안합니다. 우리는 ReCo(Reward-Coordinated Compression)라는 단계별 프레임워크를 제시하며, 이는 경량의 과정-보상 추정기를 사용하여 각 완료된 단계를 평가하고 다음과 같은 세 가지 구성 요소를 제어합니다: (1) 보상에 따라 적응적으로 KV-캐시를 압축하여, 높은 보상을 받는 단계에서는 더 강하게, 낮은 보상을 받는 단계에서는 덜 압축하도록 합니다. (2) 반사 토큰(reflection tokens)에 대해 보상 범위 기반의 패널티를 적용하여 불필요한 생성을 억제합니다. (3) 추론 결과에 대한 신뢰도를 기반으로 조기 종료 기능을 활성화하여, 정확도가 높다고 판단될 경우 추론을 중단합니다. 세 가지 추론 모델과 여섯 가지 벤치마크에서 ReCo는 생성된 토큰 수를 37%에서 65% 줄이고 전체 지연 시간을 2.08배에서 2.35배 단축했으며, 동시에 정확도를 크게 유지했습니다.

Large Reasoning Models (LRMs) excel on complex tasks through long chain-of-thought (CoT) reasoning, but their lengthy intermediate steps cause severe overthinking that inflates inference cost. KV-cache compression is a common solution, yet existing reasoning-oriented methods apply a uniform policy across the trajectory and judge compression only by what it removes from the cache. Two observations point the other way. First, a reasoning state's tolerance to context loss varies along the trajectory, and process reward tracks it: deleting tokens at high-reward steps preserves accuracy far better than deleting the same budget at random. Second, compression is not free on the generation side, since a smaller cache leads the model to generate more tokens, partly canceling the saving. Together these motivate coordinating both sides under a single process reward. We propose ReCo (Reward-Coordinated Compression), a step-wise framework in which a lightweight process-reward estimator scores each completed step and drives three components: (1) reward-adaptive KV-cache compression that shrinks the retained cache harder at high-reward steps and less at low-reward ones, (2) a reward-banded penalty on reflection tokens that curbs redundant generation, and (3) confidence-based early stopping that triggers when the reasoning is reliable. Across three reasoning models and six benchmarks, ReCo reduces generated tokens by 37%-65% and end-to-end latency by 2.08x-2.35x over Full CoT, all while largely preserving accuracy.