Improving student model training with better reward shaping methods
Revisiting On-policy Adversarial Black-Box Distillation: Calibrating Groupwise Reward Geometry for Effective Advantage Construction
Machine LearningComputer Vision and Pattern Recognition
Summary
Trying to teach smaller AI models by copying bigger, powerful ones can be tricky because the feedback used to train the smaller models can be inconsistent or unclear. The authors show a way to improve this feedback by carefully shaping how rewards are given during training, making it easier for the smaller model to learn effectively from the bigger model's examples. They introduce a two-step process that adjusts rewards both when training the critic (which rates answers) and when training the student model itself. Their method improves performance without adding much extra work.
What this means in practice
- •For ai development teams: Train smaller language models more efficiently using limited access to large teacher models by improving reward feedback during policy optimization.
- •For software engineers working on conversational agents: Improve response quality in chatbot systems by distilling knowledge from large APIs into deployable smaller models with more stable training signals.
Authors
Xiao Cui, Mo Zhu, Yulei Qin, Yuze Wu, Wengang Zhou, Houqiang Li
Abstract
Black-box distillation is a practical route for transferring capabilities from API-accessible large language models that expose only text outputs into smaller student models. Recent on-policy adversarial methods such as GAD improve over SeqKD by forming an adversarial loop between a critic and a student, where the critic provides rewards for GRPO-based student policy optimization over the student's sampled responses. However, GRPO computes advantages from the within-group relative rewards of student samples for the same prompt, whereas the critic is trained primarily to distinguish teacher responses from student responses. This objective mismatch can produce reward groups with collapsed scale or fragile margins, leading to brittle grouped optimization signals. We propose Groupwise Reward Geometry Conditioning (GRGC), a two-stage framework that improves advantage construction by shaping student-side reward groups during both critic training and policy optimization. To improve critic-side conditioning, Gaussian groupwise Optimal Transport calibration regularizes the critic during training to produce reward groups with non-collapsed spread and smooth rank-wise gaps by matching sorted prompt-wise rewards to group-centered Gaussian quantiles. Building on this conditioned reward geometry, policy-side group power modulation reshapes the prompt-wise reward groups before they are converted into advantages, preserving the critic-induced ordering while increasing optimization-relevant margin separability. Extensive experiments across diverse teachers, student model families and scales, and training datasets demonstrate the effectiveness of GRGC on both in-distribution and out-of-distribution evaluations, while introducing negligible overhead over GAD. The code is available at https://github.com/2018cx/GRGC.