Chengyue Wu Hao Zhang Shuchen Xue Shizhe Diao Yonggan Fu Zhijian Liu Pavlo Molchanov Ping Luo Song Han Enze Xie

Résumé

Les modèles linguistiques à grande échelle (LLM) autoregressifs (AR) ont atteint des performances remarquables sur un large éventail de tâches du langage naturel, mais leur décodage séquentiel inhérent limite l'efficacité du processus d'inférence. Dans ce travail, nous proposons Fast-dLLM v2, un modèle linguistique par diffusion par blocs (dLLM) soigneusement conçu, qui permet d'adapter efficacement des modèles AR préentraînés en dLLM pour une génération de texte parallèle, en nécessitant uniquement environ 1 milliard de tokens pour le fine-tuning. Cette approche représente une réduction de 500 fois du volume de données d'entraînement par rapport aux dLLM à attention complète tels que Dream (580 milliards de tokens), tout en préservant les performances du modèle original. Notre méthode introduit une nouvelle recette d'entraînement combinant un mécanisme de diffusion par blocs avec un masque d'attention complémentaire, permettant une modélisation contextuelle bidirectionnelle par blocs sans compromettre les objectifs d'entraînement AR. Pour accélérer davantage le décodage, nous avons conçu un mécanisme hiérarchique de mise en cache : un cache au niveau des blocs, qui stocke les représentations contextuelles historiques à travers les blocs, et un cache sous-bloc, qui permet une génération parallèle efficace au sein des blocs partiellement décodés. Associé à notre pipeline de décodage parallèle, Fast-dLLM v2 atteint une accélération allant jusqu'à 2,5 fois par rapport au décodage AR standard, sans compromettre la qualité de génération. Des expérimentations étendues sur diverses benchmarks démontrent que Fast-dLLM v2 égale ou dépasse les modèles de base AR en précision, tout en offrant une efficacité au niveau de l'état de l'art parmi les dLLM — marquant ainsi une avancée significative vers le déploiement pratique de LLM rapides et précis. Le code source et les modèles seront publiés publiquement.