Learning from Partial Chain-of-Thought via Truncated-Reasoning Self-Distillation

TL;DR

This paper introduces Truncated-Reasoning Self-Distillation (TRSD), a method that trains language models to produce accurate answers from partial reasoning traces, reducing inference costs and improving robustness to truncated reasoning.

Contribution

TRSD is a novel post-training technique that enables models to generate correct answers from incomplete reasoning, reducing computational costs and enhancing robustness.

Findings

TRSD improves model robustness to truncated reasoning across benchmarks.

TRSD-trained models inherently produce shorter reasoning traces.

Significant reduction in inference-time costs without explicit regularization.

Abstract

Reasoning-oriented language models achieve strong performance by generating long chain-of-thought traces at inference time. However, this capability comes with substantial and often excessive computational cost, which can materialize in redundant or inefficient reasoning. We study this setting and introduce Truncated-Reasoning Self-Distillation (TRSD), a lightweight post-training procedure that encourages models to produce correct predictions from partial reasoning traces. In TRSD, a frozen teacher model first generates a full reasoning trace and evaluates the corresponding answer distribution conditioned on the prompt and the complete reasoning to construct a synthetic training target. A student model with the same architecture is then trained to match the teacher's answer distribution while being conditioned only on a truncated prefix of its reasoning trace. Across multiple reasoning benchmarks and token budgets, we demonstrate that TRSD improves robustness to truncated inference, with far reduced accuracy tradeoffs when applied to a diverse set of reasoning models. Moreover, although never explicitly regularized for shorter generation during training, we also find that TRSD-trained models inherently output shorter reasoning traces without truncation, significantly reducing inference-time costs even without artificial interventions.