Reinforcement learning improves token credit in diffusion language models
DTCA: Denoising Token Credit Assignment for Reinforcement Learning in Diffusion Language Models
Computation and Language
Summary
It can be hard for AI language models to understand which words are most important when they generate sentences step-by-step. The authors propose a new way to give credit to individual words during training, based on how confident the model is about other missing words after fixing one word. They also group words by their type, like nouns or verbs, to better assign learning signals. This helps the model get better at solving math problems. Their tests show improved math reasoning using this method.
What this means in practice
- •For language model engineers: Improve reasoning capabilities of diffusion-based language models by integrating token-level credit assignment in reinforcement learning training.
- •For ai math tutoring developers: Enhance accuracy of AI systems that solve math problems by using refined token-level learning signals during training.
Authors
Hongze Chen, Miao Pan, Lizhou Liang, Qinfeng Li, Xinyin Ma, Yongliang Shen, Weishi Wang, Daniel Dahlmeier, Xuhong Zhang, Wenqi Zhang
Abstract
Diffusion language models (DLMs) generate text through iterative denoising, enabling bidirectional context modeling and parallel token refinement. Recent studies have applied reinforcement learning (RL) to DLMs and improved their reasoning performance. However, the iterative denoising process makes credit assignment more challenging, as final outcome rewards provide limited information about how different tokens contribute to the final output. In this paper, we propose Denoising Token Credit Assignment (DTCA), a framework for token-level credit assignment in reinforcement learning for diffusion language models. DTCA estimates token contributions by measuring the improvement in prediction confidence of remaining masked positions after token commitment and groups these signals according to part-of-speech (POS) categories. These contribution weights are used to select informative denoising timesteps and assign token-level learning signals during RL training. DTCA further calibrates these learning signals according to token-level entropy changes induced by policy updates. Experiments on LLaDA-8B-Instruct show that DTCA consistently improves mathematical reasoning performance across different generation budgets, with gains of up to 9.0 points on MATH500 and 5.3 points on GSM8K over the base model.