北京通用人工智能研究院BIGAI

Full-Distribution Self-Distillation for Optimizing Soft Latent Reasoning

Basic Information

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.