DTCA: Denoising Token Credit Assignment for Reinforcement Learning in Diffusion Language Models
Abstract
DTCA improves reinforcement learning for diffusion language models by assigning token-level credit via confidence improvements and entropy calibration, boosting mathematical reasoning.
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.
Community
Which tokens contribute most to a diffusion language model’s response? As tokens are revealed during denoising, they help shape predictions at the remaining masked positions. However, final-answer rewards provide little information about these individual contributions. People often prioritize major word classes such as nouns and verbs when thinking about a topic, and diffusion language models may show a similar pattern with different POS contributing unequally. DTCA thus derives POS-based token-level credit from confidence gains during denoising, so that learning signals reflect each token’s estimated contribution to the rest of the response.
Hello, I am the first author of this paper. We noticed an unexpected increase in the number of upvotes within a short period of time. Our team has not purchased upvotes or used any automated voting service. We do not know the source of these votes and have reported the situation to the Hugging Face team for review. We are posting this comment to clarify the situation and maintain transparency.
Nice idea to use the denoising process itself as the credit signal, and grouping the confidence gains by POS category is a pragmatic way to keep the estimates stable without paying for extra forward passes during RL. The ablation is informative: at the 256 budget, dropping contribution credit costs 3.8 points on MATH500, while at 512 dropping entropy calibration hurts more (MATH500 44.0 to 39.8). I was also curious to see that uniform credit gives similar reward curves but consistently larger gradient norms. Since the POS contribution table is built offline once from base-model rollouts, did you check whether recomputing it partway through training changes the weights much, and how would you group tokens for code, where POS tags fit less naturally?
Thank you for the careful read and the constructive suggestions.
On whether the offline POS contribution table remains useful, you raise a good point about recomputing it partway through training. We have made some preliminary attempts to recompute the table during training. However, we have not yet systematically quantified how much the table changes, and these attempts have not led to a firm conclusion. We plan to recompute the table at several training checkpoints, compare changes in the category weights and their rankings, and examine how these changes relate to final performance.
On code generation, we agree that natural-language POS tags fit code less naturally and do not directly capture the different roles of code tokens. This is a limitation of the current work. As you suggested, one could replace natural-language POS categories with code-specific groups, such as keywords, identifiers, operators, and literals, defined through rules and syntax parsing. The same denoising confidence gains could then be used to estimate contributions for these groups. Your suggestion provides a useful starting point for exploring grouping strategies better suited to code.
Thank you again for the thoughtful comments and for these helpful suggestions.
Get this paper in your agent:
hf papers read 2610.35201 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper