MDM-Prime-v2: La codificación binaria y la permutación de índices permiten un escalado óptimo en cómputo para modelos de lenguaje de difusión
MDM-Prime-v2: Binary Encoding and Index Shuffling Enable Compute-optimal Scaling of Diffusion Language Models
March 17, 2026
Autores: Chen-Hao Chao, Wei-Fang Sun, Junwei Qua, Chun-Yi Lee, Rahul G. Krishnan
cs.AI
Resumen
Los modelos de difusión enmascarada (MDM) exhiben una generalización superior cuando se aprenden mediante un esquema de enmascaramiento parcial (Prime). Este enfoque convierte los tokens en sub-tokens y modela el proceso de difusión a nivel de sub-token. Identificamos dos limitaciones del marco MDM-Prime. Primero, carecemos de herramientas para guiar la elección del hiperparámetro de la granularidad del token en el sub-tokenizador. Segundo, encontramos que la forma funcional del sub-tokenizador degrada significativamente la estimación de verosimilitud cuando se combina con tokenizadores comúnmente utilizados como la Codificación por Pares de Bytes (BPE). Para abordar estas limitaciones, estudiamos la estrechez del límite variacional en MDM-Prime y desarrollamos MDM-Prime-v2, un modelo de lenguaje de difusión enmascarada que incorpora Codificación Binaria y Mezcla de Índices. Nuestro análisis de escalado revela que MDM-Prime-v2 es 21.8 veces más eficiente en cómputo que los modelos autoregresivos (ARM). En comparaciones óptimas de cómputo, MDM-Prime-v2 alcanza una perplejidad de 7.77 en OpenWebText, superando a ARM (12.99), MDM (18.94) y MDM-Prime (13.41). Al extender el tamaño del modelo a 1.1B de parámetros, nuestro modelo demuestra además una precisión cero-shot superior en varias tareas de razonamiento de sentido común.
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.