Entraînement conjoint de la prédiction multi-tokens dans l'apprentissage par renforcement via un calibrage optimal des coefficients

Résumé

L'apprentissage par renforcement à partir de récompenses vérifiables (RLVR) est devenu le paradigme standard pour améliorer la capacité de raisonnement des grands modèles de langage, tandis que la prédiction multi-tokens (MTP) est un module largement adopté en pré-entraînement. Les combiner est une approche naturelle, mais les pratiques actuelles de RL détachent les gradients de MTP car l'entraînement conjoint dégrade la performance. Nous réexaminons cet échec sous l'angle de l'optimisation. Nous montrons que l'effet par étape de MTP sur l'objectif de RL peut être décomposé en deux termes : une corrélation de premier ordre et une pénalité de perturbation de second ordre. Cette décomposition unifie trois régimes d'entraînement MTP : le détachement, la perte d'entropie croisée et la perte de politique, et explique pourquoi chacun réussit ou échoue. Une analyse plus poussée de la perte de politique révèle que, bien qu'elle soit conforme à l'intuition, la performance se dégrade encore : le terme de corrélation décroît tandis que la pénalité quadratique persiste. Guidés par cette analyse, nous proposons le calibrage optimal des coefficients (OCC), un schéma adaptatif qui suit le coefficient optimal en ligne via un proxy de log-probabilité à un coût négligeable. Sur six bancs de test de raisonnement mathématique de niveau compétition, OCC atteint ou dépasse systématiquement la ligne de base avec détachement, offrant une meilleure performance d'entraînement conjoint MTP-RL.

English

Reinforcement Learning from Verifiable Rewards (RLVR) has emerged as the standard paradigm for improving reasoning capability of large language models, while Multi-Token Prediction (MTP) has been a widely adopted module in pretraining. Combining them is a natural approach, yet current RL practices detach MTP gradients because joint training degrades the performance. We revisit this failure from an optimization perspective. We show that the per-step effect of MTP on the RL objective can be decomposed into two terms: a first-order correlation and a second-order perturbation penalty. This decomposition unifies three MTP training regimes: Detach, Cross-Entropy loss, and Policy loss, and explains why each succeeds or fails. Further analysis of policy loss reveals that, although it aligns with intuition, performance still degrades: the correlation term decays while the quadratic penalty persists. Guided by the analysis, we propose Optimal Coefficient Calibration (OCC), an adaptive scheme that tracks the optimal coefficient online via a log-probability proxy at negligible cost. Across six competition-level mathematical reasoning benchmarks, OCC consistently matches or exceeds the detach baseline, delivering improved joint MTP-RL training performance.