MDM-Prime-v2 : L'encodage binaire et la permutation d'indices permettent une mise à l'échelle optimale en calcul des modèles de langage par diffusion
MDM-Prime-v2: Binary Encoding and Index Shuffling Enable Compute-optimal Scaling of Diffusion Language Models
March 17, 2026
Auteurs: Chen-Hao Chao, Wei-Fang Sun, Junwei Qua, Chun-Yi Lee, Rahul G. Krishnan
cs.AI
Résumé
Les modèles de diffusion masquée (MDM) présentent une généralisation supérieure lorsqu'ils sont appris à l'aide d'un schéma de masquage partiel (Prime). Cette approche convertit les tokens en sous-tokens et modélise le processus de diffusion au niveau du sous-token. Nous identifions deux limitations du cadre MDM-Prime. Premièrement, nous manquons d'outils pour guider le choix des hyperparamètres de la granularité des tokens dans le sous-tokeniseur. Deuxièmement, nous constatons que la forme fonctionnelle du sous-tokeniseur dégrade significativement l'estimation de vraisemblance lorsqu'elle est associée aux tokeniseurs courants de type Byte-Pair-Encoding (BPE). Pour résoudre ces limitations, nous étudions la précision de la borne variationnelle dans MDM-Prime et développons MDM-Prime-v2, un modèle de langage par diffusion masquée qui intègre un encodage binaire et un mélange d'indices. Notre analyse d'échelle révèle que MDM-Prime-v2 est 21,8 fois plus efficace en calcul que les modèles autorégressifs (ARM). Dans des comparaisons en régime de calcul optimal, MDM-Prime-v2 atteint une perplexité de 7,77 sur OpenWebText, surpassant les ARM (12,99), les MDM (18,94) et les MDM-Prime (13,41). Lorsque la taille du modèle est étendue à 1,1 milliard de paramètres, notre modèle démontre en outre une précision zero-shot supérieure sur diverses tâches de raisonnement de sens commun.
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.