Technical Summary: The Role of Feedback Alignment in Self-Distillation

Problem Statement

Post-training large language models (LLMs) for reasoning typically relies on Reinforcement Learning from Verifiable Rewards (RLVR), such as Group Relative Policy Optimization (GRPO). While effective, RLVR suffers from sparse credit assignment: a single scalar reward for the final answer does not indicate where in the reasoning trace the model erred. Conversely, knowledge distillation offers dense, token-level supervision but requires access to a stronger teacher model's logits, which is often unavailable or costly.

Self-distillation attempts to bridge this gap by using the same model as both student and teacher under different contexts. The student conditions only on the question, while the "self-teacher" conditions on the question plus additional context (e.g., reference solutions or feedback). However, prior work has largely treated this context as a fixed choice without investigating how the structure of the feedback influences the learning signal. This paper asks: What form of feedback produces the most effective self-teacher for self-distillation?

Methodology

The authors implement a solver–critic setup to study self-distillation in the mathematical reasoning domain.

1. Setup: A trainable solver ($\pi_\theta$) generates reasoning traces for math problems. A frozen critic ($\pi_{critic}$) analyzes these traces and produces feedback. The solver is then updated via self-distillation, minimizing the divergence between its output distribution (conditioned on the question alone) and the self-teacher's distribution (conditioned on the question plus the critic's feedback).
2. Experimental Conditions: The study compares three distinct feedback contexts:
   • GRPO: A standard RLVR baseline using binary rewards (correct/incorrect) without self-distillation.
   • RefSol (Reference Solution): The self-teacher is conditioned on a ground-truth reference solution (chain-of-thought) from a stronger model.
   • StepAlignFB (Step-Aligned Critique): The self-teacher is conditioned on a critique generated by the critic. Crucially, this critique is structurally aligned with the solver's specific reasoning trace. The critic is prompted to copy the solver's correct steps verbatim and only modify the specific steps where errors occur, preserving the solver's style and notation for correct portions.
3. Mechanism Analysis: The authors analyze the "per-token advantage" (the log-ratio of probabilities between the teacher and student) to understand how different feedback types shift the model's behavior. They also investigate the role of "induction heads" (mechanisms that copy in-context patterns) in how the model processes verbatim repetitions in the feedback.

Key Contributions & Findings

1. Step-Aligned Feedback Outperforms Reference Solutions

The StepAlignFB condition yields the largest performance gains, outperforming both GRPO and RefSol-conditioned self-distillation.

• Results: On the OpenMathReasoning dataset (Avg@12 metric), StepAlignFB achieved 35.83, compared to 30.56 for RefSol and 19.72 for GRPO.
• Significance: Despite never seeing the ground-truth derivation, StepAlignFB outperforms RefSol. This suggests that the alignment of the feedback to the solver's specific trajectory is more critical than the mere presence of a correct solution.

2. Implicit Process Supervision via Structural Alignment

Per-token advantage analysis reveals the mechanism behind StepAlignFB's success:

• Localization: Step-aligned feedback concentrates distributional shifts (negative advantages) specifically at tokens where reasoning fails. It reinforces correct steps by preserving them verbatim (producing positive advantages). This mimics the behavior of a Process Reward Model (PRM) but is derived entirely from natural language feedback without training a separate reward model.
• Diffusion in RefSol: In contrast, conditioning on a reference solution creates a "diffuse suppressive signal." Even when the solver's step is correct, the reference solution often uses different phrasing or notation. This causes the self-teacher to penalize correct tokens simply because they differ from the reference, diluting the learning signal with stylistic disagreement.

3. The Role of Induction-Head Copying

The paper identifies a critical design pattern for the critic's feedback, termed the "faithful-scribe convention":

• Full Verbatim Repetition: If the critic repeats the solver's entire attempt (including errors) verbatim, the model's induction heads anchor on the incorrect tokens, causing the teacher to reproduce the error.
• Omission: If the critic omits the solver's correct steps entirely, the teacher's distribution drifts away from the solver's correct reasoning.
• Partial Repetition (The Solution): The optimal strategy is to repeat the solver's correct steps verbatim up to the error, replace the erroneous step with a correction, and continue. This leverages induction-head copying to reinforce correct prefixes while allowing the correction to govern the error location.

Significance and Claims

The paper argues that the primary advantage of self-distillation is not merely access to additional context, but access to context whose structure matches the student's reasoning trajectory.

• Structural Alignment > Feedback Quality: A complete, correct reference derivation is a strong signal, but in self-distillation, it fails because it diverges from the solver's trace in surface form. Step-aligned critique succeeds because it addresses the solver's actual trace, concentrating the signal exactly where reasoning goes wrong.
• Implicit PRM: Step-aligned feedback provides a form of implicit process supervision, localizing errors without the computational cost of training a step-level reward model.
• Practical Trade-off: While StepAlignFB is superior, the authors note it requires a capable critique model to generate high-quality, step-aligned feedback, which increases training costs and system complexity compared to simple reference solution conditioning or binary rewards.

The study is limited to the OpenMathReasoning dataset and mathematical reasoning; the authors explicitly state that further work is needed to determine if these benefits transfer to other domains, model sizes, or solver-critic configurations.