What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models
[Sztuczna inteligencja] Autorzy: Jinya Sakurai, Xun Xu, Jaehong Yoon
Masked diffusion language models generate tokens in parallel, but their efficiency depends on which predictions are committed. Learned unmasking policies improve the speed--accuracy trade-off, yet differ in supervision, inputs, and capacity, obscuring which choices matter. We view supervised policies as distilling a decoding procedure through a teacher and examine independently what the teacher labels as committable, on which states the policy imitates it, what the policy observes, and how much the number of function evaluations (NFE) must decrease to offset policy overhead. We find that labeling predictions matching the final output is overly permissive at low NFE, and hidden states are weak alone but can complement confidence. About a 10% NFE reduction suffices to offset policy overhead for policies up to 7.5M parameters across tested settings. We therefore train a policy on trajector…
- Kategoria arxiv
- cs.AI
- Kategoria
- Sztuczna inteligencja