In a diffusion language model, a final answer is assembled through successive token commitments, but a single outcome reward does not reveal which commitments helped later denoising. DTCA turns that hidden interaction into a training signal: it asks whether committing a token makes the model more confident about positions that are still masked, then uses the estimate to guide credit assignment.
Key Findings
- Contribution estimates are grouped by part of speech, making token-level signals more stable across rollout samples and helping select informative denoising steps.
- The method calibrates learning signals using token-level entropy changes caused by policy updates, addressing the risk that an update changes confidence without improving useful behavior.
- On LLaDA-8B-Instruct, reported gains over the base model reach 9.0 percentage points on MATH500 and 5.3 on GSM8K at a generation budget of 128; improvements remain positive at budgets of 256 and 512. This suggests the approach is particularly helpful when denoising computation is constrained.
Method and Scope
The central idea is to measure refinement, not merely reward: after a token is committed, DTCA compares prediction confidence on positions that remain masked. Offline rollouts provide estimates that are averaged by POS category, avoiding the extra forward-pass cost of repeatedly calculating individual token effects during policy optimization. The approach therefore depends on the quality of POS grouping and on confidence change being a useful proxy for contribution.
Who It’s For
A good fit for researchers studying reinforcement learning, credit assignment, or training methods for diffusion language models, especially those evaluating mathematical reasoning under limited generation budgets. Look elsewhere if you need evidence across diverse tasks or model families: the reported experiments center on mathematical benchmarks and LLaDA-8B-Instruct, and broader applicability remains unestablished.