The Role of Feedback Alignment in Self-Distillation
Abstract
Process supervision (assigning credit to individual steps of a reasoning trace) has been shown to outperform outcome-only rewards for mathematical reasoning. The standard implementation trains a process reward model (PRM) on step-labeled data, which is expensive to collect. We show that using step-aligned critique as context in self-distillation provides PRM-like credit assignment. It correctly signs and localizes per-token advantages at error-prone steps in the solver's trace, without training a reward model. We compare three training signals in a solver--critic setup on a filtered subset of OpenMathReasoning: (i) a binary reward (\grpo), (ii) self-distillation with the reference solution, and (iii) self-distillation with a step-by-step critique aligned to the solver's reasoning trace. Step-aligned critique yields the largest gains, beating \grpo by 16.11 points in Avg@12 on the evaluation set and reference-solution-conditioned self-distillation by 5.27 points. Per-token advantage analysis explains why: step-aligned feedback concentrates distributional shifts at the tokens where reasoning fails, acting as implicit process supervision. In contrast, the reference solution spreads signal across all tokens, because an alternative derivation diverges from the solver's trace even at correct steps. Self-distillation with step-aligned critique therefore recovers the localization benefit of process supervision within a simpler training pipeline.