Model improves following multiple rules by using counterfactual distillation
Counterfactual Constraint-Conditioned On-Policy Distillation for Multi-Constraint Instruction Following
Machine Learning
Summary
Following many rules at once is hard for AI models, which often break some rules despite training. The authors created a new training method called CC-OPD that checks how each rule impacts the model's decisions by leaving out one rule at a time and comparing outcomes. This approach helps the model learn exactly which parts of its output are influenced by each rule without needing extra evaluators. Tests showed this method helps smaller AI models perform better, even outperforming larger ones trained differently.
What this means in practice
- •For ai model trainers: Train language models to better respect multiple simultaneous constraints without external verifiers during training.
- •For customer support platform developers: Develop chatbots that reliably adhere to complex company policies by improving multi-constraint instruction following.$Commercial implications: Enables commercial chatbot products that ensure policy compliance by training models with precise constraint conditioning.
Authors
Yanzhao Zheng, Yuanqiang Yu, Tianze Xu, Chao Ma, Zhentao Zhang, Jihuai Zhu, Baohua Dong, Hangcheng Zhu, Ruohui Huang
Abstract
Multi-constraint instruction following requires a model to respond to a query under many simultaneously active constraints. Even strong instruction-tuned models still routinely violate some of them. Existing approaches either augment supervision with sequence- or token-level RL rewards from external verifiers or learned graders, or use on-policy distillation (OPD) against a single full-context teacher whose probability mass becomes diluted as more constraints become simultaneously active. We propose CC-OPD (Counterfactual Constraint-Conditioned On-Policy Distillation), which inverts the standard supervision-generation direction in distillation. Rather than enriching the teacher with information beyond what the student sees, CC-OPD ablates each constraint from the teacher's conditioning in turn, and constructs the per-constraint signal from the resulting per-token probability differentials. The resulting per-token leave-one-out log-likelihood shifts are summed, clipped, and added to the vanilla OPD reward as a token-level shaping term. All shaping terms are obtained from the frozen teacher, without an external verifier during distillation, and the reward equals vanilla OPD wherever the aggregate shift is zero. Across two Qwen model pairs and seven benchmarks, CC-OPD achieves the highest average among all evaluated student-training methods. A 1.5B student trained with CC-OPD surpasses its own 7B RL-trained teacher on the MulDimIF benchmark.