Source-linked AI summary

The Role of Feedback Alignment in Self-Distillation

Semih Kara, Oğuzhan Ersoy

arXiv:2606.11173v1cs.AIcs.LG

TL;DR

The paper asks which feedback structure makes self-distillation most effective. It compares binary rewards, reference solutions, and step-aligned critiques in a solver–critic setup, finding that step-aligned critique performs best across accuracy metrics, including a 5.27-point Avg@12 gain over reference-solution conditioning.

  • Problem

    The study examines which feedback structure most effectively serves as context for self-distillation, a question whose evidence remains limited.

  • Method

    The authors train a solver with feedback from a frozen critic and compare binary rewards, reference solutions, and critiques aligned to the solver’s reasoning trace.

  • Results

    Step-aligned feedback consistently performs best across accuracy metrics, exceeding reference-solution conditioning by 5.27 points on Avg@12.

  • Takeaways & Limitations

    Feedback alignment concentrates distributional shifts near reasoning errors, preserving correct behavior more effectively than solution-level feedback.

  • Takeaways & Limitations

    All experiments use the OpenMathReasoning dataset, so transfer to other benchmarks, reasoning domains, and solver–critic configurations remains unclear.

Abstract

from arXiv · show

Conditioning a language model 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.

1. Introduction

The paper studies how feedback context design shapes self-distillation in mathematical reasoning. It finds that critique aligned to the solver’s reasoning trace outperforms binary rewards and reference solutions by concentrating learning on error-adjacent tokens.

  • Method: Self-distillation trains a question-only student to match a self-teacher conditioned on the question and additional context.The context can include execution traces, follow-up prompts, reference solutions, or feedback from another model.
  • Motivation: Prior work treats self-distillation context as fixed, leaving the effect of context design on learned behavior unstudied.Existing contexts include code execution traces, ground-truth solutions, and feedback from models or users.
  • Method: The study compares binary reward (GRPO), the reference solution (RefSol), and step-by-step critique aligned to the solver’s reasoning trace (StepAlignFB).The comparison uses a solver–critic setup in mathematical reasoning.
  • Results: StepAlignFB outperforms GRPO and RefSol on all aggregation-style accuracy metrics evaluated.The comparison is reported for the mathematical reasoning solver–critic setup.
  • Analysis: Per-token advantages from StepAlignFB concentrate distributional shifts at error-adjacent tokens, whereas solution-level feedback produces a more diffuse signal.This localization resembles process supervision from a PRM without training a reward model or collecting per-step scalar labels.
  • Interpretation: Feedback alignment matters because reference derivations diverge from solver traces even at correct steps, while StepAlignFB targets tokens where reasoning goes wrong.Thus, a complete and correct reference solution can diffuse its signal across the solver’s rollout.

2. Background

The background contrasts standard and on-policy distillation with teacher-free RLVR, then introduces self-distillation as an on-policy framework whose effectiveness depends on feedback-context structure. Self-distillation provides dense, token-level credit assignment by comparing the same model under different prompting contexts.

  • Distillation: Standard distillation trains a student to match a stronger teacher with per-token supervision, but requires access to the teacher’s logits.It is also off-policy because the student trains on teacher rollouts rather than its own.
  • Distillation: On-policy distillation uses student-generated rollouts and token-level teacher supervision to reduce distribution drift, but still requires a stronger teacher’s logits.This addresses the off-policy mismatch described for standard distillation.
  • RLVR: RLVR removes the need for a teacher by scoring sampled rollouts with binary rewards, while GRPO estimates advantages through group normalization.The resulting GRPO advantage is constant across all tokens in a rollout.
  • Self-distillation: Self-distillation uses the same model as student and self-teacher, conditioning them respectively on the question alone and the question augmented with context.Training minimizes a per-token divergence between the two resulting distributions.
  • Self-distillation: Self-distillation advantages vary at every token according to how context shifts next-token predictions, providing dense credit assignment whose quality depends on context structure.When feedback comes from another model, practitioners control that feedback’s structure, which is the paper’s focus.

3. Methodology

The methodology isolates feedback structure in a solver–critic self-distillation setup by varying only the critic-produced context. It compares binary-reward training, reference-solution conditioning, and step-aligned critique under a fixed training framework.

  • Experimental control: Only the feedback function varies across experimental conditions; the solver, loss, divergence, and hyperparameters remain fixed.This design isolates feedback structure as the independent variable.
  • Solver–critic setup: The solver–critic loop trains only the solver, while a critic generates feedback from the question and each solver response.Self-distillation uses the critic feedback as context for the solver’s training.
  • Feedback conditions: GRPO uses group-normalized binary rewards without a critic or self-distillation, whereas RefSol conditions on a stronger model’s chain-of-thought reference solution.RefSol references may be generated on-the-fly or precomputed during dataset construction.
  • Feedback conditions: StepAlignFB produces per-step critique that copies correct steps verbatim and repairs incorrect or incomplete steps while staying close to the solver’s reasoning trace.The critic receives both the solver’s step-tagged response and the dataset reference solution.
  • Self-distillation recipe: Training uses Qwen3-1.7B with on-policy sampling, forward KL, a 2048-token cap, and a fixed base-policy teacher paired with Thinking-Mode-Off students.Self-distillation uses group size G=1, while GRPO uses G=8.
  • Dataset: The dataset contains 312 filtered OpenMathReasoning problems: 30 evaluation samples and 282 training samples, selected for difficulty and formatting constraints.The filters require Avg@16 < 5/16 and formatting accuracy greater than 0.9.

4. Results

Step-aligned feedback outperforms reference-solution conditioning and GRPO, with targeted advantages that reinforce correct reasoning while suppressing errors. Its effectiveness depends on feedback that preserves correct solver tokens and rewrites only incorrect steps.

  • Comparative performance: StepAlignFB exceeds RefSol by +5.27 Avg@12, +13.33 Maj@12, and +2.33 Pass@12 at per-metric best values.The Majority-Vote gain suggests sharper concentration on correct answers, benefiting test-time aggregation.
  • Evaluation protocol: Self-distillation benefits from early stopping: 5–6 epochs reached peak performance despite seven training epochs, making per-checkpoint selection necessary for fair evaluation.Fixed end-of-run evaluation can understate self-distillation’s performance ceiling.
  • Comparative performance: Both self-distillation methods remain above GRPO in accuracy throughout training, with a final Avg@12 gap on the order of 8 points.GRPO retains token efficiency, while self-distillation uses a single rollout versus GRPO’s eight-rollout advantage estimate at matched compute.
  • Mechanism: Step-aligned feedback creates process-reward-like credit assignment by producing large negative advantages at incorrect steps while preserving and reinforcing correct reasoning steps.The critic targets errors without suppressing valid portions of the solver’s trace.
  • Mechanism: Reference-solution feedback produces diffuse negative advantages because alternative derivations differ in reasoning path, notation, and phrasing even where the solver is correct.This mixes error correction with stylistic disagreement and dilutes the learning signal.
  • Feedback design: The faithful-scribe convention repeats correct steps verbatim and rewrites only incorrect ones, preserving positive advantages on the correct prefix.Full repetition anchors the teacher to erroneous tokens, whereas omitting correct steps causes drift away from the student’s correct trace.

5. Related Work

Prior self-distillation work varies in its teacher context and training objective, including execution traces, reference solutions, stronger-model solutions, ground-truth solutions, and text feedback. Related research also emphasizes step-level supervision, inter-model judgments, and structured reasoning as ways to improve mathematical and multi-step reasoning.

  • Self-distillation for LLMs: Self-distillation studies use diverse teacher contexts, including code execution traces, reference solutions, stronger models’ solutions, and ground-truth solutions.SDPO used execution traces, dataset references, and stronger-model solutions; OPSD used ground-truth solutions as teacher context.
  • Self-distillation for LLMs: RLTF learns from text feedback through self-distillation and feedback-modeling objectives while re-sampling in-context solver rollouts.SDFT instead applies self-distillation to continual learning from demonstrations.
  • Process vs. outcome supervision: Per-step process supervision localizes errors and provides denser credit assignment than outcome supervision, while discouraging flawed intermediate reasoning.Process-based feedback and PRMs consistently outperform outcome supervision for mathematics.
  • Multi-agent debate and LLM-as-judge: Multi-agent debate and LLM-as-judge established inter-model interaction for improving inference-time reasoning, and generated judgments are now widely used as training signals.The solver–critic setup follows this inter-model interaction paradigm.
  • Structured reasoning: Chain-of-thought and scratchpads improve multi-step reasoning by exposing intermediate computations, while STaR turns model-generated rationales into a training signal.Prompting alone can often elicit intermediate reasoning; STaR retains rationales that produce correct answers.

6. Conclusion

Feedback structure is central to self-distillation effectiveness: step-aligned feedback consistently outperforms sparse binary rewards and reference-solution conditioning. Its benefits come with higher critique-generation costs, while the findings remain limited to OpenMathReasoning and require broader validation.

  • Conclusion: Step-aligned feedback consistently performs best on all accuracy metrics, outperforming GRPO’s sparse binary reward and reference-solution self-distillation.The comparison identifies feedback structure as a central determinant of self-distillation effectiveness.
  • Conclusion: High-quality step-aligned critiques require a capable critique model, increasing training cost and system complexity.Reference solutions can be reused across many training runs because they are not tailored to individual rollouts.
  • Conclusion: The study’s conclusions may not transfer beyond OpenMathReasoning, other mathematical benchmarks, broader reasoning domains, or different solver–critic configurations.The paper identifies testing these settings as a direction for future work.

A. Experiment details … Critic Prompt (Qwen/QwQ-32B)

The experiments use a step-formatted solver prompt and compare reference-solution conditioning with critic-generated, step-aligned feedback. A frozen Qwen/QwQ-32B critic classifies solver rollouts and emits case-specific feedback that preserves correct work while targeting errors or incompleteness.

  • A.1. Prompt templates: The solver sees only the problem statement under the same step-formatted instruction used for evaluation, aligning on-policy training and inference prompts.This prompt is shown in Figure 8.
  • Teacher Prompt — RefSol (reference-solution conditioning): RefSol conditions the teacher on the dataset’s privileged reference solution, presented solely to ensure correctness.The teacher prompt explicitly inserts the reference solution as context.
  • Teacher Prompt — RefSol (reference-solution conditioning): StepAlignFB replaces the reference solution with a critic-generated critique produced by Qwen/QwQ-32B.Figure 10 describes the replacement, and the accompanying passage identifies Qwen/QwQ-32B as the producer.
  • Critic Prompt (Qwen/QwQ-32B): The critic acts as a math grader and defaults to reproducing the student’s solution wherever the work is correct.This faithful-scribe behavior is the basis for preserving correct reasoning in the feedback.
  • Critic Prompt (Qwen/QwQ-32B): The critic first classifies each rollout into one of four cases: fully correct, unjustified step, incomplete derivation, or incorrect answer.Cases A–D respectively cover full correctness, missing justification, running out of tokens, and an incorrect final answer.
  • Critic Prompt (Qwen/QwQ-32B): It then identifies the pivotal step as the unjustified claim, stopping point, or earliest erroneous step, depending on the case.The prompt defines pivotal step N separately for Cases B, C, and D.
  • Critic Prompt (Qwen/QwQ-32B): For missing-justification cases, feedback reproduces the preceding and subsequent correct steps while appending the required justification at step N.The schema names the missing theorem, formula, or equation and retains the remainder of the student’s solution.
  • Critic Prompt (Qwen/QwQ-32B): For incorrect-answer cases, the critic preserves prior correct work, replaces only the erroneous step in the student’s style, and continues the derivation correctly.The corrected step matches the student’s notation, step granularity, and level of detail.

A.2. Training configuration (full)

This appendix provides the complete hyperparameter configurations for the reference-solution self-distillation and GRPO runs reported in the main paper.

  • A.2. Training configuration (full): The appendix lists complete hyperparameters for the reference-solution self-distillation and GRPO runs.The corresponding values are provided in Tables 2, 3, 4, and 5.

B. Sample solution, critique, and self-distillation advantage plots · B.1. Example correct student solution

The example correct student solution tests constant and power-function forms for the functional equation, verifies f(n)=1, and argues that this is the only solution. The accompanying analysis also frames the problem through a multiplicative exponent structure and per-token advantage visualization.

  • B.1. Example correct student solution: Setting m=n yields f(4m)=f(m)^4042, exposing a useful constraint on even arguments.
  • B.1. Example correct student solution: The candidate f(n)=1 satisfies both sides of the equation identically, confirming it as a solution.
  • B.1. Example correct student solution: The solution argues uniqueness by stating that nonconstant recursive relationships would create inconsistencies, so f must be constant and therefore equal to 1.
  • B.1. Example correct student solution: The problem asks for all functions f: Z_{>0} → Z_{>0} satisfying the stated functional equation for positive integers m and n.
  • B.1. Example correct student solution: For a constant function f(n)=c, the equation reduces to c=c^4042, forcing c=1 among positive integers.
  • B.1. Example correct student solution: The power-function trial f(n)=n^k fails because the two sides produce incompatible expressions, 2^k(m+n)^k and (mn)^(2021k).
  • B.1. Example correct student solution: Recasting f(n)=c^α(n) produces a recursive exponent equation, with the supplied derivation concluding that only α(n)=0 is compatible with positive-integer outputs.
  • B. Sample solution, critique, and self-distillation advantage plots: Figure 12 presents per-token advantages under each teacher context for the correct student solution.

B.2. Example incorrect student solution

The student correctly derives x = 380 - 3k and y = 2k + 1 but incorrectly allows unrestricted integer k, overlooking the non-negative constraints. The correction restricts k to 0 through 126 and yields 127 solutions.

  • Correct derivation: The student correctly establishes the parity condition and derives x = 380 - 3k with y = 2k + 1.The derivation requires y to be odd so that 763 - 3y is even.
  • Incorrect constraint handling: The key error is treating k as any integer and concluding that infinitely many integer solutions exist.The critique identifies Step 4 as incorrect because it fails to enforce non-negative constraints on x and y.
  • Correction: The corrected reasoning requires both x and y to be non-negative integers.This changes Step 4 from unrestricted integrality to bounded feasibility.
  • Correction: The constraints imply k ≤ 126 from x = 380 - 3k and k ≥ 0 from y = 2k + 1.Together, these bounds give 0 ≤ k ≤ 126.
  • Correct result: 127 solutions result from counting the integers from 0 through 126 inclusive.The count is 126 - 0 + 1 = 127.
Loading 2606.11173v1…