Mask-Aware Policy Gradients for Diffusion Language Models
Reinforcement learning has proven effective for improving reasoning in large language models, but extending it to Masked Diffusion Language Models (MDLMs) remains challenging due to the intractability of the log-likelihood estimation. Existing approaches approximate this log-likelihood by modeling only the token predictions, ignoring the order in which positions are unmasked during generation. We observe that MDLM generation involves two decisions at each step: what tokens to place at each masked position and which positions to remask. We formalize this as a two-stage action MDP, showing that
Lineage graph
Paper → model → repo connections mined from source citations (Tier-1 exact match).
Why these links exist
Every edge carries a method, confidence, and the source snippet that justified it — so bad links are debuggable.
- FuzzySimilar title/name (fuzzy) · 59%CompVis/stable-diffusion-v1-4 →
“Fuzzy title match (0.73): “Mask-Aware Policy Gradients for Diffusion Language Models” ≈ “CompVis/stable-diffusion-v1-4””
- FuzzySimilar title/name (fuzzy) · 59%stabilityai/stable-diffusion-3.5-large →
“Fuzzy title match (0.73): “Mask-Aware Policy Gradients for Diffusion Language Models” ≈ “stabilityai/stable-diffusion-3.5-large””
- PossiblePossibly related (embedding) · 61%Learning Unmasking Policies for Diffusion Language Models - Apple Machine Learning Research →
- LinkedLinked via arxiv author · 85%Haran Raajesh →
“Mask-Aware Policy Gradients for Diffusion Language Models”
- LinkedLinked via arxiv author · 85%Kulin Shah →
“Mask-Aware Policy Gradients for Diffusion Language Models”
- LinkedLinked via arxiv author · 85%Adam Klivans →
“Mask-Aware Policy Gradients for Diffusion Language Models”
- LinkedLinked via arxiv author · 85%Philipp Krähenbühl →
“Mask-Aware Policy Gradients for Diffusion Language Models”