Une politique de 856K paramètres accélère le LLM à diffusion LLaDA-8B

Les modèles de langage à diffusion promettent ce que les modèles autorégressifs classiques ne savent pas faire : générer plusieurs tokens en une seule passe. En pratique, tout se joue sur une question : parmi les tokens prédits, lesquels peut-on figer sans risque dès maintenant, et lesquels vaut-il mieux laisser sous le masque jusqu'à l'étape suivante ? Une nouvelle étude contrôlée décompose cette question en ses différents éléments et en tire une politique de seulement 856 000 paramètres. Sur LLaDA-8B-Instruct, elle atteint la meilleure frontière « précision — nombre d'évaluations du modèle (NFE) » et le débit le plus élevé sur GSM8K.

La politique comme distillation du décodage

Les politiques d'unmasking apprenables existent déjà, et elles améliorent bel et bien le compromis entre vitesse et qualité. Mais les travaux antérieurs différaient sur plusieurs axes à la fois : la façon d'annoter le jeu d'entraînement, les entrées que voit la politique, sa taille. Il était donc presque impossible de comparer les approches entre elles, et on ne savait pas quel choix comptait vraiment.

Les auteurs proposent un cadre unifié. Ici, une politique entraînée de façon supervisée distille une procédure de décodage via un « enseignant », c'est-à-dire via l'annotation. Ils étudient ensuite quatre éléments séparément :

• ce que l'enseignant considère comme un token pouvant être figé ;
• sur quels états de génération la politique imite l'enseignant ;
• ce que la politique observe en entrée ;
• de combien il faut réduire le NFE pour amortir le surcoût de la politique elle-même.

Où l'annotation naïve atteint ses limites

La méthode d'annotation la plus évidente se formule ainsi : un token est jugé prêt à être figé s'il coïncide avec ce qui figurera dans la réponse finale. L'étude montre qu'à faible NFE, ce critère est trop permissif. La politique apprend à accepter des tokens qui coïncident par hasard avec le résultat, mais qui écartent en réalité la génération du chemin de référence. Quand le budget de calcul est réduit, le coût de ces erreurs grimpe.

Côté entrées, le tableau est lui aussi nuancé. Les états cachés du modèle fournissent, seuls, un signal faible pour décider de figer un token. En revanche, ils fonctionnent comme complément de la confiance (confidence) de la prédiction : ensemble, les deux sources d'information apportent plus que chacune prise isolément.

Les auteurs ont aussi chiffré le coût de la politique elle-même. Toute couche ajoutée au modèle ajoute du calcul, et le gain issu d'un nombre d'étapes réduit peut être absorbé par son fonctionnement. Selon leur estimation, une réduction du NFE d'environ 10 % suffit à compenser le surcoût de politiques pouvant atteindre 7,5 millions de paramètres, dans toutes les configurations testées.

La trajectoire plutôt que la coïncidence avec le résultat final

Ces observations ont conduit au schéma final. La politique est entraînée sur des objectifs de trajectory-contraction : un token n'est figé que s'il reste sur la trajectoire de décodage de référence. En entrée, la politique reçoit la confiance des prédictions et les états cachés.

Résultats clés :

• taille de la politique : 856K paramètres, soit moins de 0,02 % des 8 milliards de paramètres de LLaDA-8B-Instruct ;
• GSM8K : meilleure frontière précision — NFE et débit le plus élevé parmi les approches comparées ;
• MATH-500 et HumanEval : résultats compétitifs pour des cibles d'accélération modérées ;
• Dream-7B-Instruct : la politique se transfère en grande partie à un autre modèle sans refonte complète.

Une réserve ressort des résultats eux-mêmes : sur MATH-500 et HumanEval, l'avantage n'est revendiqué que pour des cibles modérées. Autrement dit, avec une réduction agressive du nombre d'étapes, la première place n'est pas garantie.

Ce qu'il faut en retenir en pratique

Pour les équipes qui font tourner des LLM à diffusion en production, l'enseignement principal est d'ordre technique. Accélérer le décodage n'exige pas forcément une lourde couche supplémentaire : une politique de moins d'un million de paramètres est rentabilisée si elle réduit le NFE d'au moins une dizaine de pour cent. C'est un repère pratique pour le budget : si le gain en nombre d'étapes est inférieur, le surcoût de la politique l'absorbera entièrement.

La seconde recommandation concerne l'annotation. Si vous distillez une stratégie de décodage, mieux vaut ne pas considérer « coïncide avec la réponse finale » comme un critère suffisant, en particulier avec un faible nombre d'évaluations. L'ancrage sur la trajectoire de référence fournit des signaux d'entraînement plus stricts. Quant aux états cachés, il vaut mieux les ajouter à la confiance que les utiliser comme unique entrée.

Enfin, le transfert vers Dream-7B-Instruct laisse penser que ces politiques peuvent être rendues relativement universelles au sein de la famille des modèles de diffusion à masquage. Pour une industrie où les modèles de langage à diffusion se disputent une place aux côtés des modèles autorégressifs sur des tâches à contraintes de latence strictes, c'est un argument en faveur de couches additionnelles peu coûteuses et transférables, plutôt que du réentraînement des modèles eux-mêmes.