How Much Training-Inference Mismatch Can On-Policy Distillation Tolerate?

How Much Training-Inference Mismatch Can On-Policy Distillation Tolerate?. By Cong Yu and seven others, from Aalto University, Nanyang Technological University, and four other institutions. Published on the eighth of November.

On-policy distillation, or OPD, trains a student language model on responses the student generates itself, with a fixed teacher providing guidance. Producing those responses is the expensive part, so training systems cut the cost with quantized inference and with asynchronous rollouts. Both shortcuts open a gap between the rollout engine and the model actually being trained, a training-inference mismatch. In reinforcement learning, that kind of mismatch can destabilize training. What it does to OPD had been unclear.

The answer here is that OPD tolerates a substantial amount of mismatch without meaningful loss in training quality. Push the mismatch higher, though, and quality does degrade. Between those two regimes sits what the authors call an empirical mismatch budget: a quantity of mismatch that can be spent safely.

Every source of mismatch is measured on a single scale, which is what makes a budget expressible at all. Quantization is examined across three student-teacher pairs. On one of those pairs, weight staleness and injected noise are tested as well. Two dependencies show up. The quality cost of a given amount of mismatch depends on the model, and so does the amount of mismatch that quantization introduces in the first place, with some layers particularly sensitive to reduced precision.

Set against reinforcement learning in controlled comparisons, OPD stays stable under rollout configurations that destabilize RL methods. Weight noise reveals a subtler effect. There, mismatch also shifts which prefixes generation reaches, and the errors this introduces are of a kind that per-token corrections cannot address.

The findings are then put to work. A lightweight admission check screens candidate rollout configurations against a calibrated budget, without a full training run for each candidate. Model-specific FP8 recipes keep training quality intact while making room for at least one point nine three times as many KV-cache tokens as BF16. And pairing FP8 with asynchronous rollouts reaches the same training quality one point six two times sooner than synchronous BF16 training.

The boundaries need care. Quantization was studied across three student-teacher pairs; staleness and injected noise on only one. The tolerance is not unlimited: there is a level past which OPD degrades. And it is not portable as a single number, because both the cost of mismatch and the mismatch quantization produces vary from model to model. That is why the recipes are model-specific and why the budget has to be calibrated rather than assumed. Finally, the prefix-shift errors seen under weight noise mark a limit on the remedy side: per-token corrections do not reach them.

Read as a systems result, the mismatch that reinforcement learning has to fight is, for on-policy distillation, partly a resource. Within the calibrated budget it can be traded for cheaper precision and for asynchrony, and the trade shows up as more cache capacity and a faster path to the same quality.