提出新策略梯度方法，提升掩码扩散语言模型的推理能力。
Mask-Aware Policy Gradients for Diffusion Language Models
- 将生成过程建模为分两阶段的动作马尔可夫决策过程
- 在GSM8K上达到87.1%准确率，MBPP上达53.4%
- 适合研究推理增强与扩散语言模型的学者
强化学习在提升大语言模型推理能力方面表现良好，但将其扩展至掩码扩散语言模型（MDLMs）仍具挑战性，主要源于对数似然估计不可行。现有方法仅建模词元预测来近似该似然，忽略了生成过程中遮蔽位置的解封顺序。我们观察到，MDLM生成每一步涉及两个决策：在每个遮蔽位置填入什么词元，以及选择哪些位置继续遮蔽。我们将此形式化为两阶段动作马尔可夫决策过程，发现策略梯度自然分解为词元项和遮蔽项。联合优化两项，在数学推理和编程基准上取得当前最佳性能，其中在GSM8K上达87.1%，在MBPP上达53.4%。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。