Modèles de langage à diffusion multi-bloc

Résumé

Les modèles de langage à diffusion par blocs (BD-LM) améliorent la génération de texte basée sur la diffusion grâce à la mise en cache KV et à la génération de longueur flexible. Une étape naturelle consiste à les étendre de la diffusion à bloc unique (SingleBD) à la diffusion multi-blocs (MultiBD), où un ensemble courant de blocs consécutifs est décodé simultanément pour un parallélisme inter-blocs. Cependant, les BD-LM existants sont principalement entraînés sous teacher forcing, où le modèle n'observe qu'un seul bloc bruité conditionné par un préfixe propre. Bien que la stratégie récente de diffusion forcing introduise une visibilité entre plusieurs blocs bruités, ses états d'entraînement diffèrent encore de l'inférence MultiBD, où le décodage opère sur un ensemble courant borné avec des motifs de bruit hétérogènes par emplacement. Pour combler cet écart, nous proposons les modèles de langage à diffusion multi-blocs (MBD-LM), obtenus par post-entraînement des BD-LM avec le teacher forcing multi-blocs (MultiTF). MultiTF intègre le teacher forcing et le diffusion forcing en s'entraînant sur des groupes de bruit bornés conditionnés par des préfixes propres, avec des planificateurs de bruit aléatoires qui correspondent mieux aux états d'inférence MultiBD. Pour rendre MultiBD exécutable en pratique, nous introduisons en outre un algorithme de décodage optimisé basé sur le mécanisme de tampon de blocs, qui préserve la réutilisation du cache de préfixe, maintient les formes d'entrée statiques et traduit un parallélisme de décodage accru en une accélération du temps réel. Empiriquement, MBD-LLaDA2-Mini augmente le nombre moyen de jetons par passe avant (TPF) de 3,47 à 6,19 et améliore la précision moyenne de 79,95 % à 81,03 % ; combiné à DMax, MBD-LLaDA2-Mini-DMax atteint un TPF moyen de 9,34 avec seulement une baisse de précision de 1,02 % sur les benchmarks de mathématiques et de code.

English

Block Diffusion Language Models (BD-LMs) improve diffusion-based text generation with KV caching and flexible-length generation. A natural next step is to extend them from Single-Block Diffusion (SingleBD) to Multi-Block Diffusion (MultiBD), where a running-set of consecutive blocks is decoded concurrently for inter-block parallelism. However, existing BD-LMs are mostly trained under teacher forcing, where the model observes only one noisy block conditioned on a clean prefix. While the recent diffusion forcing strategy introduces visibility among multiple noisy blocks, its training states still differ from MultiBD inference, where decoding operates on a bounded running-set with heterogeneous slot-wise noise patterns. To bridge this gap, we propose Multi-Block Diffusion Language Models (MBD-LMs), obtained by post-training BD-LMs with Multi-block Teacher Forcing (MultiTF). MultiTF integrates teacher forcing and diffusion forcing by training on bounded noise-groups conditioned on clean prefixes, with randomized noise-schedulers that better match MultiBD inference states. To make MultiBD practically executable, we further introduce an optimized decoding algorithm based on the Block Buffer mechanism that preserves prefix-cache reuse, keeps input shapes static, and translates increased decoding parallelism into wall-clock acceleration. Empirically, MBD-LLaDA2-Mini increases average Tokens Per Forward pass (TPF) from 3.47 to 6.19 and improves average accuracy from 79.95% to 81.03%; when combined with DMax, MBD-LLaDA2-Mini-DMax reaches an average TPF of 9.34 with only a 1.02% accuracy drop on math and code benchmarks.