Abstract: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 find that the functional form of the subtokenizer significantly increases the cross-entropy loss in the objective when paired with commonly used Byte-Pair-Encoding (BPE) tokenizers. Second, we lack tools to guide the hyperparameter choice of the token granularity in the subtokenizer. To address these limitations, we analyze the optimal design of the subtokenizer that minimizes MDM-Prime training objective and develop MDM-Prime-v2, a masked diffusion language model which incorporates Binary Encoding and Index Shuffling. Our analysis characterizes how token granularity and sub-token entropy influence the training objective and downstream performance, providing principled criteria for subtokenizer design. When extending the model size to 1.1B parameters, MDM-Prime-v2 demonstrates superior average zero-shot accuracy across eight commonsense reasoning benchmarks, outperforming similar-sized baselines including GPT-Neo, OPT, Pythia, Bloom, SMDM, and TinyLLaMA.
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2603.16077 [cs.LG] |
| (or arXiv:2603.16077v3 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2603.16077 arXiv-issued DOI via DataCite |
From: Chen-Hao Chao [view email]
[v1]
Tue, 17 Mar 2026 02:54:16 UTC (8,663 KB)
[v2]
Mon, 30 Mar 2026 18:21:44 UTC (1 KB) (withdrawn)
[v3]
Thu, 21 May 2026 03:27:43 UTC (9,296 KB)
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。