Backtracking When It Strays: Mitigating Dual Exposure Biases in LLM Reasoning Distillation

TL;DR

The paper introduces MOTAB, a novel distillation pipeline for LLM reasoning that dynamically backtracks to correct errors, effectively mitigating dual exposure biases and improving reasoning performance.

Contribution

It proposes a new method that monitors and backtracks student reasoning trajectories to address both off-policy and on-policy exposure biases in LLM distillation.

Findings

MOTAB achieves approximately 3% performance improvement on reasoning tasks.

The approach effectively alleviates dual exposure biases in LLM reasoning distillation.

Experimental results on LIMO-v2 and AceReason datasets validate the method's effectiveness.

Abstract

Large language models (LLMs) have achieved remarkable success in complex reasoning tasks via long chain-of-thought (CoT), yet their immense computational overhead hinders real-world deployment. LLM reasoning distillation addresses this by transferring reasoning capabilities from formidable teacher models to compact student models. However, existing distillation paradigms face a fundamental dilemma. Typical off-policy distillation strictly utilizes teacher-generated golden trajectories, suffering from an exposure bias due to the mismatch between training distributions and student-generated inference contexts, which leads to error cascades in long CoT reasoning. To address this, on-policy distillation allows students to explore their own trajectories, but we demonstrate that it inherently introduces a reciprocal reversed exposure bias: the teacher model also struggles to provide positive guidance when conditioned on student-generated sub-optimal contexts. To resolve this dual exposure biases problem, we propose Monitoring Trajectories and Backtracking when it strays (MOTAB), a new LLM reasoning distillation pipeline. Specifically, MOTAB dynamically monitors the student's on-policy generation against an adaptive safety boundary. When the generation strays and exceeds this threshold, MOTAB backtracks to the last safe state and leverages teacher intervention to correct the course. This approach inherently tolerates minor student errors to mitigate exposure bias, while preventing sub-optimal contexts to circumvent reversed exposure bias. Extensive experiments on the LIMO-v2 and AceReason datasets demonstrate that MOTAB effectively alleviates the dual exposure biases, yielding a roughly 3% average performance improvement in reasoning tasks.