Des chercheurs suppriment le bruit de l'entraînement MTP des LLM en une seule ligne de code

Presque tous les grands modèles de langage récents résolvent au préentraînement une tâche un peu plus difficile que la normale. Ils ne prédisent pas seulement le token suivant, mais aussi plusieurs tokens à l'avance. Cette technique s'appelle multi-token prediction, ou MTP.

Elle sert à deux choses : le modèle gère mieux les tâches à long horizon, comme les mathématiques et le code, et les « têtes » supplémentaires peuvent servir de rédacteurs de brouillons rapides pour le décodage spéculatif. Un brouillon de plusieurs tokens s'esquisse à moindre coût, puis le modèle principal le vérifie en une seule passe. Les chercheurs du projet Soft-MTP ont constaté que l'entraînement standard de ces têtes charrie une part de bruit, et ils proposent un moyen de l'éliminer.

Il suffit pour cela d'une ligne de code supplémentaire.

Ce qui cloche dans l'entraînement MTP classique

Les têtes MTP sont entraînées avec l'entropie croisée standard contre les données réelles. Les auteurs montrent que le gradient de cette fonction de perte se décompose en deux parties, et ce de façon exacte, sans approximation.

• La première partie est une autodistillation KL. La divergence KL mesure à quel point une distribution de probabilités diffère d'une autre. Ici, elle tire chaque tête MTP vers la prédiction du modèle lui-même pour le même token. C'est cette partie qui augmente la part de brouillons acceptés lors du décodage spéculatif : un token proposé par une tête n'est accepté que s'il coïncide avec la prédiction du modèle principal. • La seconde partie est le résidu. C'est le gradient de l'entropie croisée classique du token suivant contre les données, arrivé par un autre chemin. Son rapport signal/bruit est faible, et il étouffe le signal KL utile.

Pourquoi le résidu est-il bruité ? Le token cible dans les données n'est qu'un échantillon précis de la distribution réelle : un vecteur « one-hot », avec un 1 à une seule position. La variance d'un tel échantillon ne disparaît ni avec la taille du modèle, ni avec un entraînement prolongé. Les auteurs l'ont mesuré directement. Le rapport signal/bruit du résidu est nettement inférieur à celui de la composante KL, et l'écart se creuse avec la taille du modèle. De plus, le résidu va à l'encontre de la mise à jour que reçoit la tâche principale de prédiction du token suivant. Il nuit donc à la fois à l'entraînement des têtes MTP et à celui du modèle lui-même.

Une ligne de code au lieu de matrices à la taille du vocabulaire

On peut obtenir le gradient KL pur de manière frontale : construire les distributions du professeur et de l'élève, puis minimiser l'écart entre elles. Mais avec un vocabulaire de plusieurs dizaines de milliers de tokens, il faudrait garder en mémoire d'énormes matrices de logits. De plus, cette variante ne fonctionne pas avec les noyaux d'entropie croisée « fusionnés » standard (fused kernels), conçus précisément pour éviter de matérialiser ces matrices. C'est pourquoi, en pratique, les modèles de pointe entraînent leurs têtes MTP avec l'entropie croisée classique.

Les auteurs proposent une astuce baptisée KL gradient-recovery trick. Lorsque deux distributions partagent une même tête de sortie linéaire, le gradient de leur KL peut être retrouvé exactement en combinant deux gradients d'entropie croisée déjà calculés. Techniquement, il s'agit d'un routage de gradient de type straight-through : on ajoute à l'état caché de la tête principale l'expression hk.detach() - hk. La valeur ne change pas, de sorte que l'entraînement de la tâche principale se déroule comme avant. En revanche, un gradient supplémentaire s'infiltre dans la représentation MTP et, additionné au gradient habituel de la perte MTP, il soustrait le résidu bruité. Au final, seul le signal KL parvient aux têtes MTP et au tronc commun du modèle.

Voilà comment fonctionne Soft-MTP. La méthode n'ajoute aucun paramètre et ne crée pas de surcoût notable. Elle fonctionne aussi bien avec le MTP de style DeepSeek qu'avec la variante de Meta, et elle est compatible avec le parallélisme utilisé pour le préentraînement des grands modèles. Le code est publié sur GitHub : JaydenTeoh/Soft-MTP.

Ce que montrent les expériences de 167M à 1,2B

Les auteurs ont entraîné des modèles de 167 millions à 1,2 milliard de paramètres et comparé Soft-MTP au MTP standard.

• La perplexité sur la prédiction du token suivant est plus basse à toutes les tailles, et l'écart se creuse avec l'échelle. • La perplexité des têtes MTP est plus basse à partir du modèle de 689 millions de paramètres. • L'avantage de Soft-MTP en précision sur les tâches sans ajustement (zero-shot) se maintient à chaque échelle. • Le décodage auto-spéculatif est accéléré dans tous les domaines testés. • Après ajustement, le modèle de 1,2 milliard de paramètres résout mieux les tâches de mathématiques et de code.

Les auteurs ont aussi comparé la méthode à la minimisation KL directe, qui exige un calcul explicite des logits. Sur la plupart des métriques, Soft-MTP ne lui cède pas et la surpasse par endroits. L'explication probable : Soft-MTP continue d'entraîner la tête de sortie commune avec l'entropie croisée classique sur les données. Avec un KL pur, elle dériverait vers les prédictions propres du modèle.

Coût et limites

La méthode implique un vrai compromis. En supprimant le résidu bruité, Soft-MTP introduit un biais : les têtes MTP sont tirées vers la distribution du modèle lui-même, et non vers la distribution réelle des données. Les auteurs estiment que cet échange devient de plus en plus avantageux à mesure que l'échelle augmente : l'erreur du modèle-professeur diminue, alors que le bruit d'échantillonnage demeure. En outre, les têtes MTP servent avant tout à accélérer l'inférence, où la concordance avec la prédiction du modèle principal est justement ce qu'on cherche.

Pour les équipes qui préentraînent leurs propres LLM avec MTP, la conclusion pratique est simple. La modification est minuscule, ne coûte rien en mémoire ni en vitesse et ne casse pas les fused kernels habituels. Les expériences se limitent à des modèles allant jusqu'à 1,2 milliard de paramètres : à l'échelle de DeepSeek-V3, le résultat reste à vérifier. Mais la tendance, où le gain croît avec la taille du modèle, est encourageante.