这篇论文给扩散语言模型加了个策略梯度法子,数学推理和编程测试成绩都刷了新高,搞RL训练模型的可以看看。
论文提出掩码感知策略梯度方法,解决强化学习应用于掩码扩散语言模型(MDLM)时对数似然估计难的问题。方法将MDLM生成建模为两阶段动作MDP,策略梯度分解为token项和掩码项。联合优化两项后,模型在GSM8K上达87.1%准确率,在MBPP上达53.4%得分,均达到当前最优。
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 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.