Value augmented self-distillation improves reasoning in AI language models
VASPO: Self-Distillation with Privileged Value Estimation
Computation and LanguageArtificial IntelligenceMachine Learning
Summary
Predicting answers by AI can be improved by teaching the model to better guess the value of its future words, not just the next best word choice. The authors studied this and created a method called VASPO, which combines a teacher model’s guidance and a critic that estimates the value of continuing a response. This approach helps the AI make smarter decisions when generating long and complex reasoning tasks, leading to better accuracy on math problems.
What this means in practice
- •For ai system developers: Improve reasoning accuracy of language models in complex problem solving by integrating value-augmented self-distillation during training.
- •For automated reasoning software teams: Enhance mathematical reasoning capabilities of AI tools used in automated theorem proving or logic puzzles with better reward predictions.
Authors
Fei Wu, Zhenrong Zhang, Jianshu Zhang, Jun Du
Abstract
On-policy self-distillation provides dense supervision for reasoning models via a privileged self-teacher. However, this supervision reflects local token preferences rather than the expected return of an unguided student continuation. Recent methods combine this signal with GRPO. They use GRPO for update directions and the teacher for signal scaling. However, we demonstrate a more effective integration within PPO. Specifically, we identify complementary positional dynamics: privileged information enhances a critic's value-prediction accuracy toward the response's end, whereas self-teacher guidance is strong initially, drops mid-generation, and rebounds finally. These insights motivate Value-Augmented Self-Distillation Policy Optimization (VASPO), wherein a self-teacher and critic evaluate identical privileged information at student-generated prefixes. The critic computes token-level PPO advantages from predicted continuation returns, while the self-teacher modulates selected positive advantages via bounded probability ratios. Across six mathematical reasoning benchmarks, VASPO improves mean accuracy over PPO by 2.59 and 1.81 percentage points on Qwen3-4B and OLMo3-7B variants, respectively, and achieves the highest mean accuracy among evaluated methods on Knights and Knaves. Our code is available at https://github.com/cii030/VASPO.