The Role of Feedback Alignment in Self-Distillation

Abstract: Conditioning a LLM on additional context, such as feedback on a previous attempt, typically improves its response. Self-distillation trains the model to retain this improvement when the context is not present. The method works by matching the model's output distribution under two settings: a student that sees only the question, and a self-teacher that also sees the context. What the model learns therefore depends on what context the self-teacher receives, yet the design of this context remains largely unexplored. We study context design for self-distillation by training a solver on feedback from a frozen critic. We compare three conditions: (i) a binary reward (GRPO), (ii) the reference solution, and (iii) a step-by-step critique aligned to the solver's reasoning trace. Step-aligned critique yields the largest gains, outperforming GRPO by 16.11 points and reference-solution-conditioned self-distillation by 5.27 points (Avg@12). Per-token advantage analysis reveals why: step-aligned feedback targets only the tokens where reasoning fails, leaving correct behavior intact. Conditioning on the reference solution, by contrast, pressures the model to change its behavior at every token (even correct steps) because an alternative derivation inevitably differs in phrasing and approach. This suggests that structural alignment between feedback and the solver's reasoning is a key driver of self-distillation effectiveness.

• The paper shows that step-aligned feedback enhances token-level credit assignment by aligning gradient signals with correct and incorrect reasoning steps.
• The study contrasts GRPO, RefSol, and StepAlignFB, with StepAlignFB achieving up to 35.83 Avg@12 and outperforming reference-based methods by over 5 points.
• The research highlights that careful prompt engineering and structured feedback yield sharper error localization and enable more resource-efficient training.

The Role of Feedback Alignment in Self-Distillation

Introduction and Motivation

Self-distillation has emerged as a potent strategy for improving the performance of LLMs, enabling dense, token-level credit assignment without requiring access to external teacher logits or expensive step-level human annotations. Unlike conventional reinforcement learning from verifiable rewards (RLVR), which relies on sparse, trajectory-level feedback, self-distillation leverages additional context, such as feedback or reference solutions, to optimize the policy via context-dependent supervision. However, the structure of the context used for self-distillation—specifically, the alignment between the feedback and the model's own stepwise reasoning—has remained underexplored.

This paper systematically analyzes how the feedback structure in self-distillation governs learning dynamics and performance. By comparing binary reward (GRPO), reference-solution-based self-distillation (RefSol), and step-aligned feedback (StepAlignFB), the study demonstrates that structural feedback alignment is a principal axis determining distillation effectiveness. Extensive empirical analysis is supported by mechanistic interpretations of the gradient signals at the token level, elucidating the learning advantages conferred by localized, step-aligned critique.

Methodology

The experiments are conducted in the domain of mathematical reasoning, using a solver–critic framework. The solver, a Qwen3-1.7B model, generates step-tagged reasoning traces for each math problem. A frozen, capable critic model provides feedback, which is then used as context in the self-distillation process. The study contrasts three feedback settings:

• GRPO (Binary Reward): Uses group-normalized binary reward across rollouts as the advantage signal, as in classical RLVR.
• RefSol (Reference Solution): The teacher conditions on a canonical reference solution from the dataset, inducing self-distillation via imitation of an alternative, fully correct solution.
• StepAlignFB (Step-Aligned Feedback): The critic inspects both the model's generated solution and the canonical answer, producing fine-grained, step-by-step feedback that precisely aligns with the student's reasoning trace, copying correct steps verbatim and correcting only erroneous steps.

The only variable across experimental conditions is the content of the feedback context; all other aspects of model architecture, objectives, and hyperparameters remain invariant, enabling an ablation-style isolation of feedback structure effects.

Experimental Results

StepAlignFB demonstrates pronounced and consistent improvements over both GRPO and RefSol across all evaluated metrics on a challenging subset of the OpenMathReasoning dataset. Specifically, step-aligned feedback achieves maximum Avg@12 (35.83), exceeding RefSol by 5.27 points and GRPO by 16.11. Gains are robust across majority-vote and pass-rate metrics as well, indicating not only enhanced accuracy but also sharper concentration of predictive probability mass on correct solutions.

Figure 1: Metric Avg@12 as a function of training step, demonstrating superior final performance of StepAlignFB over both RefSol and GRPO baselines.

This dominance extends across the training trajectory, with both distillation approaches outperforming binary-reward RLVR (GRPO) throughout. Notably, the self-distillation variants benefit from earlier stopping, highlighting a potential avenue for resource-efficient training.

Token-Level Credit Assignment and Mechanistic Insights

A key contribution is the detailed analysis of per-token advantages, directly reflecting the gradients supplied by different feedback strategies. Under StepAlignFB, negative advantage spikes precisely co-localize with erroneous steps, while correct reasoning segments receive positive reinforcement, analogous to a process reward model. This property emerges automatically from the alignment between the critic's feedback and the student's trace, obviating the need for explicit per-step human or automated scalar labels.

Figure 2: Per-token advantages $A_t^{\mathrm{SD}}$ for a fully-correct rollout; StepAlignFB (top) supplies uniformly positive signal along the correct trajectory, while RefSol (bottom) exhibits diffuse, negative suppression due to phrase-level and stylistic divergences despite overall correctness.

Figure 3: For a rollout with an incorrect step, StepAlignFB (top) applies sharp negative advantage only at tokens post-error (vertical dashed line); RefSol (bottom) produces broadly negative signal, diluting error-localization by conflating divergence in phrasing and notation.

The RefSol context, in contrast, yields diffuse suppressive gradients even for entirely correct rollouts, due to inevitable mismatches in derivation surface form. As a result, the model is pressured to alter correct behaviors to better match teacher phrasing, resulting in suboptimal generalization and less efficient credit targeting.

StepAlignFB effectively operationalizes token-level process supervision with only natural language feedback. This mirrors the performance advantage of process reward models (PRMs) over scalar outcome-based RL, but via dense, textual credit assignment.

The Role of Induction-Head Copying and Prompt Engineering

The token-level advantage structure is sharply influenced by the critic's prompt design. Full verbatim repetition of the student's trace, including errors, leads to induction-head copying that reinforces all student tokens, masking mistakes and preventing corrective gradients. Conversely, omitting the student's reasoning in the feedback causes the teacher's distribution to drift, penalizing even correct student steps.

Figure 4: Full verbatim repetition of an incorrect step, followed by correction. The teacher's distribution is dominated by copying, yielding uniformly positive advantages, which fails to penalize the error.

Figure 5: Omission of the student solution from the feedback context. This results in negative advantages across the whole correct trajectory, as the teacher cannot anchor on the correct steps.

The only effective regime is the "faithful scribe" convention: repeating the correct prefix verbatim up to the first error, then substituting a correction. This leverages induction heads as a mechanism for targeted credit assignment—anchoring the teacher on correct tokens while allowing errors to be corrected locally.

Implications and Future Directions

The findings establish that structural alignment between feedback and student reasoning is as critical as feedback "quality." Highly capable reference solutions may provide strong theoretical supervision, but unless their stepwise form matches the student's trace, their supervisory effect is diluted by distributional drift at every token. Step-aligned feedback concentrates the learning signal precisely where errors occur, providing implicit process supervision and addressing the long-standing credit assignment problem in LLM reasoning.

However, step-aligned critique introduces practical trade-offs. Generating high-fidelity, locally anchored feedback necessitates a strong, expensive critic model and careful prompt engineering. The additional system complexity and computational cost may limit scalability compared to cheaper, reference-solution-based approaches with broader applicability.

Generalization beyond mathematical reasoning tasks, transfer to larger models, and adaptation to alternative domains are open research questions. There is a need to investigate whether similarly structured, but more trivially obtainable, forms of feedback can recapitulate these benefits.

Conclusion

This paper demonstrates that the structure of feedback context is a primary driver of self-distillation efficacy in LLMs. Step-aligned, trace-faithful feedback delivers superior accuracy and sharper process-level supervision compared to both outcome-level rewards and reference-solution-based contexts. These results suggest that effective self-distillation is less about access to "more" feedback and more about the right feedback structure—specifically, alignment with the model's own reasoning trajectory. The mechanistic correspondence with process reward models points toward promising hybrid approaches and compels further exploration into cost-effective, scalable methods for generating step-aligned critique in broader domains.