Chen-Hao Chao
Author directory2026
MDM-Prime-v2: Binary Encoding and Index Shuffling Enable Scaling of Diffusion Language Models
Chen-Hao Chao | Wei-Fang Sun | Junwei Quan | Chun-Yi Lee | Rahul G Krishnan
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Chen-Hao Chao | Wei-Fang Sun | Junwei Quan | Chun-Yi Lee | Rahul G Krishnan
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
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 log-likelihood value 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.
Pivot-SD: Efficient Self-Distillation for Masked Diffusion Language Models
Seo Hyun Kim | Sunwoo Hong | Younwoo Choi | Chen-Hao Chao | Se-Young Yun | Rahul G Krishnan
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Seo Hyun Kim | Sunwoo Hong | Younwoo Choi | Chen-Hao Chao | Se-Young Yun | Rahul G Krishnan
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Masked diffusion language models (dLMs) offer a promising parallel alternative to autoregressive models for complex reasoning. However, they face a unique credit-assignment challenge: reasoning failures typically stem from a few critical, early commitments that cascade downstream. Existing post-training overlooks this dynamic: supervised fine-tuning (SFT) weights all tokens uniformly, while online reinforcement learning (RL) attributes a single trajectory-level reward to every position. We introduce Pivot-SD, an efficient offline framework supervising only these high-impact bottlenecks (pivots). Pivot-SD isolates pivots using an information-gain metric measuring uncertainty reduction over the remaining masked block. It applies bidirectional supervision: successful pivots are reinforced with cross-entropy, while critical errors in failed trajectories are suppressed via targeted unlikelihood, preserving valid surrounding sub-steps. Using only 200 questions and four rollouts each, Pivot-SD improves LLaDA-8B-Instruct over full-sequence SFT and budget-matched diffusion RL baselines across math and code benchmarks.