Full-Distribution Self-Distillation for Optimizing Soft Latent Reasoning

Abstract
Soft latent reasoning extends chain-of-thought reasoning beyond discrete token sequences by using continuous concept tokens. However, effectively training models in this space remains challenging: conventional reinforcement learning methods mainly optimize sampled rollout-action likelihoods and therefore under use the full distributions produced during soft rollouts. We propose Full-Distribution Self Distillation (FDSD), an off-policy training framework that treats stochastic soft rollout distributions as dense pseudo-labels. FDSD aligns the current model distribution with advantage weighted rollout targets via KL minimization and importance correction, enabling the model to learn from the full rollout distribution rather than only from sampled-action likelihoods. Experiments on GSM8K, MATH-500, OlympiadBench, and out-of-domain benchmarks show that FDSD achieves consistent average improvements in mathematical reasoning across model scales while preserving OOD generalization, with only modest computational overhead. These results suggest that full-distribution alignment is a promising training direction for soft latent reasoning.
Authors
Ziyong Lin, Yanting Wang, Zixia Jia, Zilong Zheng✉
Publication Year
2026
http://eng.bigai.ai/wp-content/uploads/sites/7/2026/10/7617_Full_Distribution_Self_Di.pdf
Publication Venue
EMNLP
Scroll to Top