TRACES: Tagging Reasoning Steps for Adaptive Cost-Efficient Early-Stopping

TL;DR

TRACES is a lightweight framework that tags reasoning steps in real-time, enabling adaptive early stopping of language models to reduce inference costs while maintaining accuracy.

Contribution

It introduces a novel real-time tagging framework for reasoning steps, allowing effective early stopping and cost reduction in language reasoning models.

Findings

Achieves 20-50% token reduction with comparable accuracy.

LRMs tend to shift reasoning behavior after reaching a correct answer.

Monitoring reasoning steps enables effective early stopping criteria.

Abstract

The field of Language Reasoning Models (LRMs) has been very active over the past few years with advances in training and inference techniques enabling LRMs to reason longer, and more accurately. However, a growing body of studies show that LRMs are still inefficient, over-generating verification and reflection steps. Additionally, the high-level role of each reasoning step and how different step types contribute to the generation of correct answers, is largely underexplored. To address this challenge, we introduce TRACES (Tagging of the Reasoning steps enabling Adaptive Cost-Efficient early-Stopping), a lightweight framework that tags reasoning steps in real-time, and enable adaptive, cost-efficient early stopping of large-language-model inferences. Building on this framework we monitor reasoning behaviors during inferences, and we find that LRMs tend to shift their reasoning behavior after reaching a correct answer. We demonstrate that the monitoring of the specific type of steps can produce effective interpretable early stopping criteria. We evaluate the TRACES framework on three mathematical reasoning benchmarks, namely, MATH500, GSM8K, AIME and two knowledge and reasoning benchmarks, MMLU and GPQA respectively. We achieve 20 to 50% token reduction while maintaining comparable accuracy to standard generation.