Condition contrast improves teaching in on-policy language model training
Condition-Contrastive Extrapolation for On-Policy Distillation of Language Models
Artificial Intelligence
Summary
Language models can learn better by training on their own generated responses guided by a stronger AI teacher, but sometimes the teacher mistakenly favors wrong answers when the AI is given its own previous words. The authors introduce a new method that compares the teacher’s judgment on the full question and on altered versions missing key parts, which helps the teacher avoid encouraging wrong reasoning. This approach reduces errors and helps smaller language models learn more accurately, especially in math reasoning tasks. Their tests show consistent improvement on multiple benchmarks compared to previous methods.
What this means in practice
- •For language model engineers: Improve training of smaller language models by using enhanced teacher feedback that avoids reinforcing incorrect reasoning.
- •For automated math tutoring developers: Enhance reasoning accuracy in AI tutors by applying condition-contrastive on-policy distillation to reduce invalid mathematical answers.
Authors
Wenxuan Luo, Chenxi Li, Haoyue Liu, Yongxin Guo, Xiaoying Tang
Abstract
On-policy distillation (OPD) improves a student language model by training on its own generated trajectories with supervision from a stronger teacher. However, when conditioned on student-generated prefixes, the teacher can favor invalid continuations that violate the question's premises, causing OPD to encourage erroneous reasoning steps. As a remedy, we introduce Condition-Contrastive Extrapolated On-Policy Distillation (CCE-OPD), which makes teacher supervision sensitive to question conditions. CCE-OPD incorporates (1) a fixed-prefix comparison in which the same frozen teacher evaluates the complete question and a counterpart with a selected condition masked, and (2) a bounded log-probability contrast that reweights the teacher distribution into an auxiliary distillation target. On 200 diagnostic reasoning states, condition contrast reduces the rate of preferring invalid continuations from 27.0% to 16.0%. Extensive experiments across six mathematical benchmarks show that CCE-OPD improves avg@8 over OPD by an average of 1.60 and 1.95 percentage points when distilling Qwen3-8B into 0.6B and 1.7B students, respectively. It also achieves the strongest results on four transfer benchmarks.