Yinjie Wang Ling Yang Bowen Li Ye Tian Ke Shen Mengdi Wang

Résumé

Nous proposons TraceRL, un cadre d'apprentissage par renforcement conscient des trajectoires pour les modèles linguistiques à diffusion (DLM), qui intègre des trajectoires d'inférence préférées durant le post-entraînement et s'applique à différentes architectures. Grâce à un modèle de valeur basé sur la diffusion, qui améliore la stabilité de l'entraînement, nous démontrons une performance accrue en raisonnement sur des tâches complexes de mathématiques et de codage. En outre, cette approche peut également être utilisée pour adapter des modèles spécifiques à des blocs plus grands, ce qui améliore la flexibilité de l'échantillonnage. En utilisant TraceRL, nous avons développé une série de modèles linguistiques à diffusion d'avant-garde, appelés TraDo. Bien que plus petits que les modèles AR de taille 7B, TraDo-4B-Instruct surpasse de manière cohérente ces derniers sur des tâches complexes de raisonnement mathématique. TraDo-8B-Instruct atteint une amélioration relative de précision de 6,1 % par rapport à Qwen2.5-7B-Instruct et de 51,3 % par rapport à Llama3.1-8B-Instruct sur des benchmarks de raisonnement mathématique. Grâce à un apprentissage par curriculum, nous avons également obtenu le premier DLM à longue chaîne de raisonnement (long-CoT), qui dépasse Qwen2.5-7B-Instruct sur MATH500 avec une amélioration relative de précision de 18,1 %. Afin de faciliter la recherche reproductible et les applications pratiques, nous mettons à disposition un cadre open-source complet pour la conception, l'entraînement et le déploiement de modèles linguistiques à diffusion sur diverses architectures. Ce cadre intègre des techniques accélérées de mémoire tampon KV (KV-cache) et des moteurs d'inférence adaptés à la fois à l'inférence et à l'apprentissage par renforcement, et inclut des implémentations de différentes méthodes de fine-tuning supervisé et d'apprentissage par renforcement pour des tâches mathématiques, de codage et générales.Code et modèles : https://github.com/Gen-Verse/dLLM-RL