Different Tasks, Different Timesteps: Temporal Task-Aware Quantization for Diffusion Language Models
Ofek Avivi, Eliron Rahimi, Avi Mendelson, Chaim Baskin
cs.AI
2026-10-22
T2Quant spends a fixed 4-bit budget on task-specific denoising steps, lifting average accuracy by up to 6.7 points across three DLMs and six benchmarks.
A diffusion language model (DLM) rewrites masked positions in parallel under bidirectional context, and the sentence exists only after the full denoising trajectory. One model call per step means activation quantization, fewer bits for those activations, sets the memory and compute cost of every step. DLLMQuant, STaR-Quant, and FAIR-Calib keep a single activation bit-width for the whole trajectory. Image diffusion already varies precision by timestep, and autoregressive LLMs already spend bits by layer or by task-relevant weight circuits. The open axis is the task: on one trajectory, the steps that cannot tolerate low precision differ by downstream task.
The diagnostic is LLaDA-8B-Instruct, 512 tokens, 128 steps, 16 non-overlapping windows of eight steps, 4-bit GPTQ weights with QuaRot, and dynamic per-token 4-bit activations. What is measured is the KL divergence between full-precision and quantized predictions, only on committed tokens the task score actually reads. A commitment is the step that writes a concrete token into a mask, and that write is often irreversible. Spread flat, each window would hold 6.25% of this task-relevant divergence. On GSM8K the first window holds 27%, about 4.4 times a uniform share, and the first quarter holds 55%. HumanEval is quiet early: five windows from step 32 to 71 hold 52%. MMLU-Pro is the flattest task, and its peak and trough still differ by nearly 2 times. One timetable that ignores the task does not fit these three shapes.
T2Quant does not introduce a quantizer. It sits on post-training quantization (PTQ): no retraining, parameters fit from a small calibration set. The only decision is which activation precision each denoising window uses, at the same average bit-width as uniform Atarget. Calibration prompts are unlabeled. No answers, and no task metric.
Each window is scored from internal signals. Kw is the KL between full-precision and quantized predictions at positions committed in that window, a read on whether extra bits would change the decision. Iw is how far the full-precision model’s predictions on still-masked positions move from before the window to after it, a read on whether those decisions still affect later steps. The proxy is the product of the two expectations. Divergence alone ranks HumanEval’s last window first, and tokens committed there cannot change anything that follows.
Scheduling replays every candidate bit-width on states saved from a full-precision trajectory, with the input state held fixed, so the number is local quantization error. A window’s gain is its average influence times how much divergence falls relative to the reference precision. The reference scores zero. A downgrade scores negative. Dynamic programming then selects windows whose costs sum to the uniform reference budget. For integer costs it is optimal for this additive surrogate. Real accuracy is not additive, and the guarantee stops there. Inference is a lookup by window index.
Sixteen unlabeled prompts match the all-prompt proxy at correlation r of at least 0.96. On GSM8K and MMLU-Pro the proxy names the single window with the largest accuracy gain and recovers 90% and 100% of the gap to full precision, with window-level correlations 0.98 and 0.96. On HumanEval both the proxy and the measured gains pile into the first two windows, and noise pulls the correlation down to 0.63. The diagnostic KL for code peaks mid-trajectory, while the single-window accuracy gain is early. The influence term is what accounts for that mismatch.
The main table is W4A4: 4-bit weights, 4.00 average activation bits. The models are Dream-7B-Instruct, LLaDA-1.5-8B, and LLaDA-8B-Instruct, across six benchmarks, with eval prompts disjoint from calibration. DLLMQuant, FAIR-Calib, and STaR-Quant each run alone and again with T2Quant on top. A random schedule at the same budget checks mixed precision by itself.
| Model | Best original | Best + T2Quant | FP16 | Largest average gain |
| Dream-7B | STaR 38.55 | STaR+T2 42.10 | 44.05 | DLLMQuant 33.99→40.69 |
| LLaDA-1.5 | STaR 44.09 | STaR+T2 45.22 | 47.73 | DLLMQuant 39.09→42.47 |
| LLaDA-8B | STaR 43.36 | FAIR+T2 44.76 | 46.32 | DLLMQuant 37.35→40.88 |
Average accuracy rises in all nine method-model pairs, by up to 6.7, 3.4, and 3.5 points on Dream, LLaDA-1.5, and LLaDA-8B. Of 54 benchmark-level comparisons, 42 improve. The best quantized averages trail FP16 by 2.0, 2.5, and 1.6 points. GSM8K moves the most: on Dream, DLLMQuant goes from 35.55 to 53.52 and STaR-Quant from 46.09 to 58.79. The largest GSM8K gains are about 9.6 points on LLaDA-1.5 (35.94 to 45.51) and about 10.4 on LLaDA-8B-Instruct (31.84 to 42.19). Round-to-nearest falls to near zero on all three models.
On Dream the random schedule lands 14 to 19 points below the original baselines. Which steps receive the bits matters more than mixing precisions at random.
A separate ablation on LLaDA-8B-Instruct GSM8K uses uniform W4A4 (n=512): uniform 44.1, T2Quant 49.4, FP16 52.9. Sixteen calibration prompts recover about 75% of the gain, and the schedule keeps favoring early steps. A schedule fit at 4 tokens per step still leads by about 5 points at 2 tokens per step. At 8 tokens per step accuracy collapses and the gain goes with it. At 3.75 average bits T2Quant still scores 48.4, above uniform 4-bit at 44.1. Bits past 4 add almost nothing.
The quantizers already exist. The open choice is how to split one bit budget along the denoising clock. T2Quant leaves the quantizer and the sampler unchanged. Inference is a window-index lookup. Tasks do not share a curve. Math puts the extra precision at the start. Single-window accuracy gains on code also sit in the first two windows, even though the KL on tokens the scorer reads peaks mid-trajectory.
The improvement is incremental. Average activation width stays at 4 bits. What changes is which steps hold those bits. A 3.75-bit schedule beating uniform 4-bit means the budget can go lower when the allocation is right.
Speed, energy, and effects on safety-related answers are left unmeasured. Precision that varies over time could concentrate behavioral drift in a few steps. The paper does not check that.
Dynamic programming is optimal only for the additive surrogate. Single-window upgrades track the full-precision gap. Whether the ranking still holds when many windows move together is not tested on its own. HumanEval’s correlation is 0.63. The table shows local regressions: STaR-Quant HumanEval falls from 20 to 17 on LLaDA-1.5 and from 25 to 21 on LLaDA-8B-Instruct. Averages still rise. GSM8K and ARC-Challenge improve in all nine pairs, and those HumanEval drops do not pull the mean under the baseline.
Calibration assumes unlabeled prompts that resemble the target task. If the task is unknown when a request arrives, there is no schedule to look up. Changing how many tokens are committed per step can erase the gain. The models are 7B to 8B masked diffusion models, and weight bits are not scheduled together with activations.