EN
Token-level credit improves diffusion model math training
DTCA: Denoising Token Credit Assignment for Reinforcement Learning in Diffusion Language Models · arXiv:2610.35201
In plain words
Diffusion language models fill in masked words over several rounds. A reward for the finished answer does not say which earlier word choices helped the other masked positions become easier to predict.
The paper introduces DTCA, which estimates those effects by grammatical category and uses the estimates to guide training.
On two math tests scored by final-answer correctness, DTCA beat the base model at a matched 256-token response limit and 256 denoising steps; the tests do not show whether its intermediate reasoning was correct.
The research
When a diffusion language model gets a reward for a finished answer, how can training tell which earlier token choices helped the remaining masked words become easier to predict?
Those choices interact during repeated denoising, so a response-wide score alone gives weak guidance about individual tokens.
DTCA estimates how token commitments affect later predictions, then uses those estimates to choose training steps and distribute learning signals in diffusion language models.
DTCA first uses offline denoising rollouts to estimate how committing a token changes confidence in positions still masked, averages estimates by grammatical category, then uses fixed category weights to select training steps and assign token-level learning signals. It also adjusts those signals using estimated changes in prediction uncertainty from updates.
On MATH500, a math test, and GSM8K, a math-question test, the authors score the fraction of questions with correct final answers. At a 256-token response limit and 256 denoising steps, DTCA beat the reproduced base checkpoint by 6.8 and 2.1 percentage points, respectively.
Replacing its contribution-based credits with uniform credits lowered accuracy while retaining timestep selection, supporting a role for those credits.
The evaluation covers mathematical reasoning with LLaDA-8B-Instruct. DTCA’s contribution estimates depend on part-of-speech tagging and alignment, and they use changes in prediction confidence as a proxy for a token’s effect on later denoising.
What the evidence shows
- The credit assignment problem
Denoising Token Credit Assignment (DTCA) is a framework for assigning token-level learning signals during reinforcement learning for diffusion language models. These models commit tokens over repeated denoising steps, and each commitment can affect later predictions at positions that remain masked.
A reward for the finished response does not identify which commitments helped, so the authors ask how to estimate their separate contributions.
- Commitments differ by category
To motivate a way to estimate contributions, the authors analyzed rollouts generated from the GSM8K training set. They measured how prediction confidence at remaining masked positions changed after a token commitment, then averaged the estimates by part of speech.
The categories differed: verbs had the largest average confidence gain, followed by nouns and proper nouns, while determiners and adpositions had smaller gains. This rollout analysis is distinct from the final model evaluation.
- From estimates to training signals
DTCA uses offline rollouts to estimate these effects before reinforcement learning, averages estimates within grammatical categories, and turns them into fixed weights. During training, the weights help select denoising steps and scale learning signals for optimization targets.
Those targets may include both tokens committed at a selected step and retained masked positions. DTCA also adjusts signals using estimated changes in prediction uncertainty after model updates.
- What the tests show
The authors tested mathematical reasoning on MATH500 and GSM8K, scoring final-answer correctness on the complete test splits. With 256-token response limits and 256 denoising steps, DTCA’s accuracy was higher than the reproduced base checkpoint on both tests.
An ablation kept DTCA’s timestep selection but replaced contribution-based credits with uniform credits; accuracy fell, supporting a role for contribution-based credits within that comparison.
Deep
- Original paper authors
- Hongze Chen, Miao Pan, Lizhou Liang, Qinfeng Li, Xinyin Ma
- Submitted to arXiv
- License
- CC BY 4.0