DTCA lifts a diffusion language model by up to 9 points on MATH500

DTCA is a new method for training diffusion language models with reinforcement learning, described in a paper on arXiv that has drawn 251 upvotes on Hugging Face daily papers. The authors, from the ZJU-OmniAI group, report that it lifts LLaDA-8B-Instruct by up to 9.0 points on the MATH500 benchmark and 5.3 points on GSM8K. The code is on GitHub.

The result matters because diffusion language models are one of the few serious alternatives to the way ChatGPT and Claude generate text, and they are much harder to train with reinforcement learning. DTCA is a fix for the hardest part of that problem.

What a diffusion language model is

Most chat models write one token at a time, left to right. A token is a word or a piece of a word. Each token is fixed once written.

A masked diffusion language model works differently. It starts with a response made entirely of blanks, called masks. Over several steps it fills in some of the blanks, the ones it is most sure about, and leaves the rest. Each step uses the tokens already filled in to make better guesses about the blanks that remain. The order is not left to right. The model might commit to a number first and a connecting word last.

LLaDA-8B-Instruct, the model used in the paper, is an open model of this kind with 8 billion parameters.

The problem: who gets the credit?

Reinforcement learning (RL) trains a model by rewarding good outputs. For math, the reward is simple: the final answer is right or wrong. The hard part is spreading that single reward back over the many tokens and steps that produced the answer. This is called credit assignment.

The authors argue that diffusion models make this worse. Every token that gets committed changes how well the model can fill in the blanks that are left. So a token matters twice: for what it says, and for how much it helps the later steps. A single reward at the end says nothing about either.

What the team did

DTCA, short for Denoising Token Credit Assignment, has three parts, according to the paper.

Measure each token's help. For a committed token, the method checks how much the model's confidence rises across the positions that are still masked. The authors group this signal by part of speech. In their analysis of GSM8K rollouts, verbs gave the largest confidence gain, followed by nouns and proper nouns. Determiners and prepositions gave smaller gains.

Pick the steps and weight the tokens. The contribution estimates come from offline rollouts, so training does not need extra passes through the model. Each part of speech gets a weight between 0.8 and 1.2. Low-confidence tokens in high-weight categories mark the denoising steps that count most. Training keeps the smallest set of steps that holds 70 percent of the total score, between 2 and 32 steps. Tokens in those steps receive credit in proportion to their weight.

Calibrate by entropy. Entropy here measures how uncertain the model is. The method estimates how a policy update will change it, and scales the credit to keep the model from drifting too fast.

Results

All numbers below are accuracy in percent, from the paper's runs. The authors state that they reproduced every baseline under their own protocol. The model was trained on the MATH training split and tested at generation lengths of 128, 256 and 512 tokens.

Against the base model, DTCA adds 9.0, 6.8 and 4.8 points on MATH500 at the three lengths, and 5.3, 2.1 and 2.5 points on GSM8K. It has the top score in every column of the table above. The paper also compares against d1, LLaDA-1.5, TraceRL and a supervised baseline, all of which score lower on MATH500.

The gains are largest at the shortest length. At 512 tokens they shrink to 4.8 points on MATH500.

Both parts matter

In the paper's ablations at 256 tokens, removing the token credit lowers MATH500 by 3.8 points and GSM8K by 2.4. Removing the entropy calibration lowers both by 1.4. At 512 tokens, removing calibration costs 4.2 points on MATH500 and 3.5 on GSM8K.

The authors also report that without calibration, the KL divergence, which measures how far the model drifts from its starting point, rises faster and late-stage reward ends lower. In a comparison with RLDF over updates 401 to 500, DTCA had a reward of 0.537 against 0.482 and a mean KL of 0.36 against 0.89. The authors note this comparison does not isolate single modules.

Limits

The authors list three. The method depends on the quality of part-of-speech tagging, which they did with the spaCy library. The idea of "contribution," meaning confidence gained at the remaining blanks, is a heuristic, and they say it needs firmer theory. And the experiments cover only math reasoning on one model, LLaDA-8B-Instruct. They suggest other tasks such as code may need different token groupings, and other diffusion architectures remain untested.

There is a further caveat from the setup. Scores come from one response per question at temperature 0, so the paper does not show how much the numbers vary between runs. The paper does not say whether trained weights are released.

What to watch

The next test is whether the same idea holds on code and on larger diffusion models. The code is out, so others can try it, and the project page is at zju-omniai.github.io/DTCA.