How to Distill a Good Student from a Biased Teacher, Without Grading the Target
When you distill a large language model into a smaller one, you're trusting the teacher's judgment on every token. If the teacher is systematically wrong about certain types of questions, the student learns those errors too. This is especially problematic under covariate shift, where the target questions differ from the ones the teacher was trained on and where reliable reward feedback is available only on source questions, not on the target distribution. Building a reward model for new question types requires expert-designed rubrics, labeled answers, and careful validation. In many practical settings, that infrastructure simply doesn't exist for the target domain.
A new paper from MIT and UIUC proposes Coupled Calibration and Learning, an algorithm that iteratively corrects teacher bias using source reward feedback and trains the student on target questions using the corrected teacher. The key insight is that calibration and learning should be coupled: each updated student provides better alternatives for the next round of calibration, and each improved calibration provides better supervision for the next round of student training. The authors prove that this process converges to the oracle student, the best policy within the student class, at a polynomial rate, without any reward feedback on target questions.
They also prove a separation result: direct matching to the teacher can remain strictly suboptimal even when the teacher achieves higher regularized target reward than every policy in the student class. In other words, the teacher might be better than any student, but naively imitating it doesn't produce the best student.
The Problem: Direct Imitation Propagates Bias
Knowledge distillation works by training a student to match a teacher's output distribution. The teacher provides rich synthetic supervision, which is especially attractive when collecting task-specific demonstrations is expensive. But matching the teacher's predictions also reproduces its systematic errors. Image classification experiments have shown that distillation amplifies teacher errors on difficult classes even when average student accuracy improves. With LLMs, experiments show substantial accuracy losses when in-context demonstrations come from different topic domains than the evaluation questions.
The challenge intensifies under covariate shift. Strong source performance doesn't certify the teacher's accuracy on target questions. A reward model that works well on source questions may be unreliable on target questions: multilingual evaluations find lower reward-model accuracy and inconsistent preferences across languages. Constructing new reward feedback is expensive, involving extensive human annotations for mathematical process verifiers or task-specific tests for code execution.
The practical setting the authors target is this: trusted reward feedback exists for source questions, but only the questions themselves are available in the target dataset. The learner must use source feedback to address target teacher bias, without ever evaluating target answers against their true rewards.
Token-Level Branching: The Calibration Mechanism
CCL's calibration mechanism works through token-level branching. Each iteration selects a random source question and a random position in a teacher-generated answer. At that prefix, the algorithm pairs the teacher's next token with an alternative sampled from the current student. This creates a comparison between two token choices in the same context.
The comparison is evaluated using a reference-weighted branch. One of the two tokens is selected with probability proportional to the reference policy's probability for that token, the selected branch is completed with the reference policy, and the resulting answer is evaluated using the source reward oracle. A reward-dependent acceptance rule turns this into a statistical signal. The acceptance indicator satisfies a specific probability: the probability that a branch is accepted equals the exponential of the reward minus one, divided by the regularization parameter lambda.
The key mathematical identity is that, under a realizability assumption, the probability that the accepted branch index equals one equals the sigmoid of a feature difference dotted with the optimal calibration parameter. This means source reward feedback provides logistic supervision for estimating the calibration coefficients. The algorithm updates the calibration parameter by projecting a stochastic gradient step onto a compact convex set.
The feature difference between the teacher's token and the student's alternative is the crucial signal. The log-ratio of the calibrated teacher's probabilities for the two tokens is linear in this feature difference and the calibration parameter. Learning the optimal relative probabilities of two tokens at a given prefix constrains the calibration coefficients along that feature direction. Sampling branching positions across all H token steps collects the joint source information needed to identify the full calibration parameter.
The realizability assumption requires that a single calibration parameter works for both source and target prompts. This is the shared structure that enables transfer: source and target prompts use the same optimal calibration coefficients, evaluated through their respective prefix features. The assumption is imposed on the calibration class and allows the student class to remain misspecified.
From Calibrated Teacher to Student Training
Once the teacher is calibrated, the student trains by minimizing the KL divergence between its output distribution and the calibrated teacher's distribution on target questions. The student's training objective is the average KL divergence over target prompts, scaled by the regularization parameter lambda. This objective is estimated using fresh rollouts from the current student on randomly selected target questions.
The gradient of this objective is unbiased: the score function multiplied by the log-ratio cost gives an unbiased gradient estimate. The student takes a projected gradient step, and then a validation step compares this proposal against a randomly sampled alternative from the parameter ball. The candidate with the lower estimated cost is selected as the next student.
This alternating structure means the student's updated parameters become the source of alternatives for the next calibration round. Better student alternatives produce more informative comparisons for calibration, and better calibration produces more accurate supervision for student training. The two processes reinforce each other.
The Convergence Guarantee
The main theoretical result establishes that the expected average KL divergence from the CCL student to the oracle student converges to zero at a polynomial rate. The oracle student is defined as the policy within the student class that maximizes the true reference-regularized target reward. It need not represent the unrestricted optimal policy, which may lie outside the student class.
The analysis uses three mechanisms. Near the oracle, projected student gradient updates provide progress: the gradient of the true objective is bounded away from zero at any suboptimal point, and the smoothness of the objective ensures that gradient steps make progress. Away from the oracle, the random exploration candidate provides progress by sampling from the full parameter ball, which has nonzero probability of landing in a region closer to the oracle. Calibration error is controlled through the source information matrix and the number of iterations, ensuring that the calibrated teacher's error relative to the true calibration parameter shrinks over time.
The convergence rate depends on the source information matrix, the smoothness of the objective, the diameter of the parameter sets, and the regularization parameter. The polynomial rate means that doubling the number of iterations reduces the error by a constant factor, which is sufficient for practical convergence in many settings.
Separation from Direct Matching
The paper's most striking result is the separation from regularized direct matching. Direct matching trains the student to minimize KL divergence to the teacher's distribution, without any calibration step. The authors prove that the direct-matching student's error relative to the oracle can remain bounded away from zero even when the teacher achieves higher regularized target reward than every policy in the student class.
This means a teacher that is "better" than any student, in the sense of achieving higher true regularized reward, can still produce a student that is strictly suboptimal through direct imitation. The teacher's bias on target questions persists in the student even when the teacher has useful knowledge to transfer. CCL's calibration step eliminates this persistent error by correcting the teacher's predictions using source reward feedback before using them for student training.
The separation holds because direct matching optimizes the wrong objective. It minimizes KL divergence to the teacher, which includes the teacher's biases and errors. CCL instead optimizes KL divergence to the calibrated teacher, which approximates the true regularized optimal policy. The calibration step is what makes the difference: it aligns the distillation objective with the actual target optimization goal.
Practical Implications and Limitations
For practitioners, the results suggest that naive distillation from a powerful teacher to a smaller student is insufficient when the teacher's reliability on target questions is uncertain. The calibration step adds computational overhead, but it provides a principled way to correct for teacher bias without requiring target-domain reward feedback.
The algorithm requires a source reward oracle, which in practice means a verifier or scoring function that works on source questions. This is a realistic assumption for many applications: mathematical reasoning has automated graders, code execution has test suites, and content generation can be evaluated against source-domain criteria. The restriction is that this oracle doesn't work on target questions, which is the practical constraint the algorithm addresses.
The theoretical analysis uses linear-softmax policy classes and fixed feature maps. Extending the results to nonlinear function approximation, such as neural network policies, is an open direction. The polynomial convergence rate is a worst-case guarantee; practical convergence may be faster in specific settings.
The algorithm's iterative structure means it requires multiple rounds of teacher calibration and student training. Each round involves sampling teacher answers, branching at random positions, evaluating source rewards, and updating both calibration and student parameters. The batch sizes for gradient estimation and validation affect both the variance of the updates and the computational cost per iteration.
The separation result has a direct practical consequence: if you're distilling an LLM and you suspect the teacher may be biased on your target distribution, don't just match the teacher's outputs. Calibrate the teacher using whatever reward feedback you have on source questions, then train the student on the calibrated predictions. The extra step is cheap compared to the cost of building a target-domain reward model, and the theory guarantees it recovers the best achievable student.
The Broader Context
This work connects to several active research areas. In reinforcement learning, the problem of learning from a biased expert is well-studied, but the LLM setting introduces specific constraints: autoregressive generation, token-level decisions, and the practical unavailability of target reward feedback. In transfer learning, covariate shift theory characterizes when source supervision supports target prediction, but typically in the context of supervised classification, not sequential decision-making.
The algorithm's structure also relates to recent work on LLM self-improvement and preference optimization. Methods that use reward models to score and select among model outputs face similar challenges when the reward model is imperfect. CCL's approach of calibrating the reward signal using source feedback before applying it to target data is a principled alternative to reward hacking or overoptimization.
For teams building distilled LLM systems, the practical takeaway is clear: if you can't trust your teacher on the target distribution, calibrate it before distilling. The source reward feedback you already have is enough, provided you have a realistic calibration class that can capture the teacher's systematic biases. The coupled structure of CCL ensures that calibration improves as the student improves, creating a virtuous cycle that converges to the best achievable student policy.
Read the paper on arXiv