Stabilizing Efficient Reasoning with Step-Level Advantage Selection

2026-04-26

AI Summary

This work addresses the instability and performance degradation of large language models during inference-time compression when trained with short-context post-training. To mitigate this, the authors propose a Step-level Advantage Selection mechanism within the GRPO reinforcement learning framework. By incorporating verifier feedback, the method assigns zero advantage to low-confidence steps in correct reasoning trajectories and high-confidence steps in failed trajectories, thereby stabilizing training without relying on length-aware objectives. Empirically, this approach improves average Pass@1 accuracy by 0.86 points across multiple mathematical and general reasoning benchmarks while reducing inference length by 16.3%, achieving a significantly better trade-off between accuracy and efficiency.

Technology Category

• Natural Language Processing: Safety and Robustness
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

Large language models (LLMs) achieve strong reasoning performance by allocating substantial computation at inference time, often generating long and verbose reasoning traces. While recent work on efficient reasoning reduces this overhead through length-based rewards or pruning, many approaches are post-trained under a much shorter context window than base-model training, a factor whose effect has not been systematically isolated. We first show that short-context post-training alone, using standard GRPO without any length-aware objective, already induces substantial reasoning compression-but at the cost of increasingly unstable training dynamics and accuracy degradation. To address this, we propose Step-level Advantage Selection (SAS), which operates at the reasoning-step level and assigns a zero advantage to low-confidence steps in correct rollouts and to high-confidence steps in verifier-failed rollouts, where failures often arise from truncation or verifier issues rather than incorrect reasoning. Across diverse mathematical and general reasoning benchmarks, SAS improves average Pass@1 accuracy by 0.86 points over the strongest length-aware baseline while reducing average reasoning length by 16.3%, yielding a better accuracy-efficiency trade-off.

Problem

• efficient reasoning
• short-context post-training
• reasoning compression
• training instability
• accuracy degradation

Innovation

• Step-level Advantage Selection
• efficient reasoning
• reasoning compression
• LLM inference optimization
• accuracy-efficiency trade-off