Reproduce “Joint Training of Multi-Token Prediction in Reinforcement Learning via Optimal Coefficient Calibration” on OpenResearch

About this paper

Zili Wang, Jiajun Chai, Lin Chen, Xiaohan Wang, Shiming Xiang, Guojun Yin

The authors introduce OCC to adjust multi-token prediction updates during reinforcement learning using estimated alignment with reinforcement learning gradients and auxiliary-update size. The method uses changes in log probabilities as proxies for full gradients. On MiMo-7B-RL trained with DAPO, OCC averages 61.7% across six math benchmarks, versus 58.9% for Detach; evidence is limited to tested mathematical reasoning with verifiable rewards.

Key insights

• The observed performance degradation in joint MTP-RL training stems from a 'phase transition' where the beneficial correlation between RL and MTP gradients decays over time, while the detrimental perturbing magnitude of the MTP gradient persists.
• An optimal MTP coefficient can be theoretically derived by maximizing the per-step improvement of the RL objective, found to be proportional to the ratio of gradient correlation to the squared MTP gradient magnitude.
• Log-probability changes effectively serve as low-cost, high-fidelity proxies for full-model gradients, enabling efficient and scalable online calibration of the MTP coefficient without significant computational overhead.

Summary by alphaXiv.