Modelos de Linguagem de Difusão Multi-Bloco

Resumo

Modelos de Linguagem de Difusão por Blocos (BD-LMs) melhoram a geração de texto baseada em difusão com cache de pares chave-valor e geração de comprimento flexível. Um passo natural é estendê-los da Difusão de Bloco Único (SingleBD) para a Difusão de Blocos Múltiplos (MultiBD), onde um conjunto dinâmico de blocos consecutivos é decodificado simultaneamente para paralelismo entre blocos. No entanto, os BD-LMs existentes são majoritariamente treinados sob forçamento do professor, onde o modelo observa apenas um bloco ruidoso condicionado a um prefixo limpo. Embora a estratégia recente de forçamento de difusão introduza visibilidade entre múltiplos blocos ruidosos, seus estados de treinamento ainda diferem da inferência MultiBD, onde a decodificação opera em um conjunto dinâmico limitado com padrões de ruído heterogêneos por posição. Para preencher essa lacuna, propomos os Modelos de Linguagem de Difusão por Blocos Múltiplos (MBD-LMs), obtidos por pós-treinamento de BD-LMs com Forçamento do Professor para Múltiplos Blocos (MultiTF). O MultiTF integra o forçamento do professor e o forçamento de difusão ao treinar em grupos de ruído limitados condicionados a prefixos limpos, com agendadores de ruído aleatorizados que melhor correspondem aos estados de inferência MultiBD. Para tornar a MultiBD praticamente executável, introduzimos ainda um algoritmo de decodificação otimizado baseado no mecanismo de Buffer de Blocos, que preserva o reuso do cache de prefixo, mantém as formas de entrada estáticas e traduz o aumento do paralelismo de decodificação em aceleração de tempo real. Empiricamente, o MBD-LLaDA2-Mini aumenta a média de Tokens por Forward (TPF) de 3,47 para 6,19 e melhora a acurácia média de 79,95% para 81,03%; quando combinado com DMax, o MBD-LLaDA2-Mini-DMax atinge um TPF médio de 9,34 com uma queda de apenas 1,02% na acurácia em benchmarks de matemática e código.

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.