Обрезка непредсказуемого: эффективное рассуждение о коде через удивление первого токена

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

Аннотация

Недавно крупные модели рассуждений (LRMs) продемонстрировали впечатляющие способности в анализе кода за счет масштабирования длины цепочки рассуждений (Chain-of-Thought, CoT). Однако чрезмерно длинные траектории рассуждений создают значительные трудности с точки зрения затрат на обучение, задержки при выводе и практической реализации. Хотя для решения этой проблемы появились различные подходы к сжатию CoT, они сталкиваются с неизбежными компромиссами: методы на уровне токенов часто нарушают синтаксическую и логическую согласованность, а методы на уровне шагов, основанные на перплексии, не могут надежно выделить логически важные этапы рассуждений. В данной статье мы предлагаем ASAP (Anchor-guided, Surprisal-based Pruning), новый двухэтапный фреймворк для сжатия CoT. ASAP сначала выполняет якорно-ориентированное сокращение, чтобы сохранить основную структуру рассуждений, что эффективно уменьшает пространство поиска для последующей обработки. Затем он обеспечивает логически осознанное сокращение, выбирая логически важные этапы рассуждений на основе новой метрики удивления первого токена. Наконец, ASAP обучает модели автономно генерировать и использовать эти сжатые CoT во время вывода, что позволяет эффективно выполнять рассуждения в задачах программирования. Эксперименты показывают, что ASAP достигает наивысшей точности на нескольких бенчмарках генерации кода, значительно снижая затраты на обучение и вывод. На сложном бенчмарке LiveCodeBench v4_v5 наш подход сокращает генерацию токенов на 23,5% и задержку вывода на 43,5% по сравнению с самым сильным базовым методом, при этом достигая конкурентоспособной точности 36,19% в метрике Pass@1. Наши результаты указывают на перспективное направление для создания мощных и эффективных LRMs.

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.