MDM-Prime-v2: la codifica binaria e la permutazione degli indici abilitano un ridimensionamento computazionalmente ottimale per i modelli linguistico-diffusivi
MDM-Prime-v2: Binary Encoding and Index Shuffling Enable Compute-optimal Scaling of Diffusion Language Models
March 17, 2026
Autori: Chen-Hao Chao, Wei-Fang Sun, Junwei Qua, Chun-Yi Lee, Rahul G. Krishnan
cs.AI
Abstract
I modelli di diffusione mascherata (MDM) mostrano una generalizzazione superiore quando appresi utilizzando uno schema di mascheramento parziale (Prime). Questo approccio converte i token in sub-token e modella il processo di diffusione a livello di sub-token. Identifichiamo due limitazioni del framework MDM-Prime. In primo luogo, mancano strumenti per guidare la scelta degli iperparametri della granularità del token nel sottotokenizzatore. In secondo luogo, rileviamo che la forma funzionale del sottotokenizzatore degrada significativamente la stima della likelihood quando abbinata ai tokenizzatori Byte-Pair-Encoding (BPE) comunemente utilizzati. Per affrontare queste limitazioni, studiamo la tensione del limite variazionale in MDM-Prime e sviluppiamo MDM-Prime-v2, un modello linguistico di diffusione mascherata che incorpora la codifica binaria e l'ordinamento casuale degli indici. La nostra analisi di scalabilità rivela che MDM-Prime-v2 è 21,8 volte più efficiente in termini di calcolo rispetto ai modelli autoregressivi (ARM). In confronti ottimali per il calcolo, MDM-Prime-v2 raggiunge una perplexity di 7,77 su OpenWebText, superando ARM (12,99), MDM (18,94) e MDM-Prime (13,41). Estendendo la dimensione del modello a 1,1 miliardi di parametri, il nostro modello dimostra un'ulteriore superiorità in accuratezza zero-shot su varie attività di ragionamento di senso comune.
English
Masked diffusion models (MDM) exhibit superior generalization when learned using a Partial masking scheme (Prime). This approach converts tokens into sub-tokens and models the diffusion process at the sub-token level. We identify two limitations of the MDM-Prime framework. First, we lack tools to guide the hyperparameter choice of the token granularity in the subtokenizer. Second, we find that the function form of the subtokenizer significantly degrades likelihood estimation when paired with commonly used Byte-Pair-Encoding (BPE) tokenizers. To address these limitations, we study the tightness of the variational bound in MDM-Prime and develop MDM-Prime-v2, a masked diffusion language model which incorporates Binary Encoding and Index Shuffling. Our scaling analysis reveals that MDM-Prime-v2 is 21.8times more compute-efficient than autoregressive models (ARM). In compute-optimal comparisons, MDM-Prime-v2 achieves 7.77 perplexity on OpenWebText, outperforming ARM (12.99), MDM (18.94), and MDM-Prime (13.41). When extending the model size to 1.1B parameters, our model further demonstrates superior zero-shot accuracy on various commonsense reasoning tasks.