CORTEXA
← Browse
arxivcs.CLcs.AIcs.LG2026-07-16

Mask-Aware Policy Gradients for Diffusion Language Models

Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl

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 the policy gradient naturally decomposes into a token term and a masking term. Combining optimization of both terms leads to state-of-the-art outcomes on mathematical reasoning and coding benchmarks, with scores of 87.1% on GSM8K and 53.4% on MBPP.

View free PDFSource page

Related papers

arxivcs.CLcs.AIcs.LG2026-07-18

Trace-Based On-Policy Distillation for Masked Diffusion Language Models

Haolin Ren, Ziyang Huang, Chenhao Yuan, Jun Zhao, Kang Liu

Diffusion large language models (dLLMs) are a promising alternative to autoregressive generation. However, reasoning-oriented post-training for dLLMs remains challenging. Supervised fine-tuning (SFT) for dLLMs requires dense but often off-policy masked states, while reinforcement…

View free PDFSource page
arxivcs.CLcs.AIcs.LG2026-07-17

Induction in Both Directions: A Mechanistic Analysis of In-Context Learning in Masked Diffusion Language Models

Andy Catruna, Emilian Radoi

While the internal mechanisms of autoregressive (AR) transformers have been studied extensively, much less is known about diffusion language models (DLMs), an emerging alternative that generates text by iterative denoising. In this work, we study how DLMs implement induction, a m…

View free PDFSource page
arxivcs.LGcs.AIcs.CL2026-07-14

Accelerating Masked Diffusion Large Language Models: A Survey of Efficient Inference Techniques

Daehoon Gwak, Minhyung Lee, Junwoo Park, Jaegul Choo

Diffusion large language models (dLLMs) offer a theoretical advantage in parallel generation over standard autoregressive models. However, parallel generation alone does not guarantee practical speedups. Realizing this efficiency requires specialized inference mechanisms, such as…

View free PDFSource page
arxivcs.CLcs.AIcs.LG2026-07-07Cited by 14

Ad Headline Generation using Self-Critical Masked Language Model

Yashal Shakti Kanungo, Sumit Negi, Aruna Rajan

For any E-commerce website it is a nontrivial problem to build enduring advertisements that attract shoppers. It is hard to pass the creative quality bar of the website, especially at a large scale. We thus propose a programmatic solution to generate product advertising headlines…

View free PDFSource page
arxiveess.AScs.AIcs.CLcs.LGcs.SD2026-07-15

Improving Text-to-Audio Instruction Following via Fine-Grained Feedback from Audio-Aware Large Language Models

Chun-Yi Kuan, Siwon Kim, Byeonggeun Kim, Suyoun Kim, Bo-Ru Lu, Qinming Tang, et al.

Recent text-to-audio models generate high-quality audio, but often fail to follow instructions involving multiple sound events and temporal order. This gap arises because existing evaluation and training signals mainly emphasize global similarity or perceptual quality, with limit…

View free PDFSource page
arxivcs.LGcs.AIcs.CL2026-07-02

Training Hybrid Block Diffusion Language Models with Partial Bidirectionality

Pranshu Chaturvedi, Parth Shroff, Tarun Suresh, Hangoo Kang, Kaiyue Wen

High-throughput long-context generation is one of the central challenges for large language models. Generation is typically memory-bandwidth-bound rather than compute-bound: each decoding step must stream the accumulated key/value (KV) cache from memory, so bandwidth demand grows…

View free PDFSource page