놀라움 없는 요소 제거: 첫 토큰 놀라움을 통한 효율적인 코드 추론
Pruning the Unsurprising: Efficient Code Reasoning via First-Token Surprisal
August 8, 2025
저자: Wenhao Zeng, Yaoning Wang, Chao Hu, Yuling Shi, Chengcheng Wan, Hongyu Zhang, Xiaodong Gu
cs.AI
초록
최근 대규모 추론 모델(Large Reasoning Models, LRMs)은 사고의 연쇄(Chain-of-Thought, CoT) 길이를 확장함으로써 코드 추론에서 놀라운 능력을 보여주고 있습니다. 그러나 지나치게 긴 추론 흔적은 훈련 비용, 추론 지연 시간, 배포 가능성 측면에서 상당한 어려움을 야기합니다. 이러한 문제를 해결하기 위해 다양한 CoT 압축 방법이 등장했지만, 이들은 본질적인 절충점을 안고 있습니다: 토큰 수준의 방법은 구문적 및 논리적 일관성을 해치는 반면, 혼란도(perplexity)에 기반한 단계 수준의 방법은 논리적으로 중요한 추론 단계를 안정적으로 포착하지 못합니다. 본 논문에서는 CoT 압축을 위한 새로운 coarse-to-fine 프레임워크인 ASAP(Anchor-guided, Surprisal-based Pruning)을 제안합니다. ASAP은 먼저 핵심 추론 구조를 보존하기 위해 앵커 기반 가지치기를 수행하여 후속 처리의 탐색 공간을 효율적으로 줄입니다. 그런 다음, 새로운 첫 토큰 혼란도 메트릭을 기반으로 논리적으로 필수적인 추론 단계를 선택함으로써 논리 인식 가지치기를 가능하게 합니다. 마지막으로, ASAP은 모델이 추론 시 이러한 간결한 CoT를 자율적으로 생성하고 활용하도록 가르쳐 코딩 작업에서 효율적인 추론을 가능하게 합니다. 실험 결과, ASAP은 여러 코드 생성 벤치마크에서 최첨단 정확도를 달성하면서도 훈련 및 추론 비용을 크게 절감했습니다. 특히 도전적인 LiveCodeBench v4_v5 벤치마크에서, 우리의 접근 방식은 가장 강력한 베이스라인 대비 토큰 생성을 23.5% 줄이고 추론 지연 시간을 43.5% 단축하면서도 Pass@1에서 36.19%의 경쟁력 있는 정확도를 달성했습니다. 이러한 결과는 강력하고 효율적인 LRM 구축을 위한 유망한 방향성을 제시합니다.
English
Recently, Large Reasoning Models (LRMs) have demonstrated remarkable
capabilities in code reasoning by scaling up the length of Chain-of-Thought
(CoT). However, excessively long reasoning traces introduce substantial
challenges in terms of training cost, inference latency, and deployment
feasibility. While various CoT compression approaches have emerged to address
this challenge, they face inherent trade-offs: token-level methods often
disrupt syntactic and logical coherence, while step-level methods based on
perplexity fail to reliably capture the logically critical reasoning steps. In
this paper, we propose ASAP (Anchor-guided, Surprisal-based Pruning), a novel
coarse-to-fine framework for CoT compression. ASAP first performs anchor-guided
pruning to preserve the core reasoning structure, which efficiently reduces the
search space for subsequent processing. It then enables a logic-aware pruning
by selecting logically essential reasoning steps based on a novel first-token
surprisal metric. Finally, ASAP teaches models to autonomously generate and
leverage these concise CoTs at inference time, enabling efficient reasoning in
coding tasks. Experiments show that ASAP achieves state-of-the-art accuracy
across multiple code generation benchmarks while substantially reducing
training and inference costs. On the challenging LiveCodeBench v4_v5 benchmark,
our approach reduces token generation by 23.5% and inference latency by 43.5%
compared to the strongest baseline, while achieving a competitive accuracy of
36.19% in Pass@1. Our results highlight a promising direction for building
powerful and efficient LRMs.