MDM-Prime-v2: Codificação Binária e Embaralhamento de Índices Permitem Escalonamento Computacionalmente Ótimo de Modelos de Linguagem de Difusão
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
Resumo
Os modelos de difusão mascarada (MDM) exibem generalização superior quando aprendidos usando um esquema de mascaramento parcial (Prime). Esta abordagem converte tokens em sub-tokens e modela o processo de difusão ao nível do sub-token. Identificamos duas limitações da estrutura MDM-Prime. Primeiro, carecemos de ferramentas para orientar a escolha do hiperparâmetro da granularidade do token no sub-tokenizador. Segundo, descobrimos que a forma funcional do sub-tokenizador degrada significativamente a estimativa de verosimilhança quando emparelhada com tokenizadores Byte-Pair-Encoding (BPE) comumente usados. Para resolver estas limitações, estudamos o rigor do limite variacional no MDM-Prime e desenvolvemos o MDM-Prime-v2, um modelo de linguagem de difusão mascarada que incorpora Codificação Binária e Embaralhamento de Índices. A nossa análise de escalabilidade revela que o MDM-Prime-v2 é 21,8 vezes mais eficiente em termos computacionais do que os modelos autoregressivos (ARM). Em comparações computacionalmente ótimas, o MDM-Prime-v2 atinge 7,77 de perplexidade no OpenWebText, superando o ARM (12,99), o MDM (18,94) e o MDM-Prime (13,41). Ao estender o tamanho do modelo para 1,1 mil milhões de parâmetros, o nosso modelo demonstra ainda uma precisão zero-shot superior em várias tarefas de raciocínio de senso comum.
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.