Different Tasks, Different Timesteps: Temporal Task-Aware Quantization for Diffusion Language Models

Different Tasks, Different Timesteps. Temporal Task-Aware Quantization for Diffusion Language Models. By Ofek Avivi and three others, from the Technion, Israel Institute of Technology and Ben-Gurion University of the Negev. Published on the twenty-second of October. Its code repository is public and real: sixteen source files.

Diffusion Language Models produce text through an iterative denoising process, and that process brings its own challenges for quantization. A model that works over many denoising steps does not necessarily need the same precision at every one of them. The authors show that sensitivity to quantization depends both on the task and on the timestep. Different downstream tasks have distinct precision requirements, and those requirements change across the denoising steps.

The second observation is what makes this usable. These patterns can be predicted without ground-truth labels, and without any evaluation metric specific to the task. If sensitivity can be read off without labels, then precision can be assigned where it matters before anything is scored.

That is the basis of T 2Quant, a mixed-precision quantization framework aware of both task and time. It uses a label-free sensitivity signal to allocate activation precision across the denoising steps, and it does so under a fixed bit-width budget. The total precision spent stays the same. What changes is which steps receive it, an allocation the authors describe as optimal.

Experiments span three Diffusion Language Models and six benchmarks. Applied on top of existing quantization methods, T 2Quant consistently improves their average accuracy, with gains of up to six point seven percentage points under the same budget.

The point is that a diffusion model's precision needs are not one number. They move with the task and with the step, and once that movement can be predicted without labels, a fixed budget can simply be spent where it buys the most.

Different Tasks, Different Timesteps. Temporal Task-Aware Quantization for Diffusion Language Models. By Ofek Avivi and three others, from the Technion, Israel Institute of Technology and Ben-Gurion University of the Negev. Published on the twenty-second of October. This is the long version: the whole paper rather than the abstract.

Diffusion language models write a whole block of text at once. They start from a sequence of masked positions and, over many denoising steps, commit tokens a few at a time, running the full network again at every step. That repetition makes quantization, running the network in fewer bits, especially attractive. Existing methods for these models keep activation precision fixed across the whole trajectory. The authors, from the Technion and Ben-Gurion University of the Negev, ask a plainer question. When does precision actually matter? Their answer is that it depends on the step, and on the task, and that both can be measured without a single label.

The diagnostic uses LLaDA-8B-Instruct, generating over one hundred and twenty-eight denoising steps, with four positions committed per step. The trajectory is cut into sixteen windows of eight steps each, and the quantized model uses four-bit weights and four-bit activations. For every window, the authors measure how far the quantized model's predictions drift from full precision at the positions the task's scorer actually reads. If that damage were spread evenly, each window would carry six point two five percent of it. It is not. On GSM8K, the math word problems, the first window alone carries twenty-seven percent, and the first quarter of the trajectory carries fifty-five percent. HumanEval, the coding task, is insensitive early and peaks in the middle, where five windows hold fifty-two percent. MMLU-Pro is early but broad. A single precision schedule shared by every task is unlikely to fit any of them well.

The first piece is a label-free signal. The true gain from a window is the accuracy change when only that window is upgraded from four bits to eight, which needs labels. The proxy instead multiplies two quantities computed from unlabeled prompts. The first is divergence: the KL divergence between full-precision and quantized predictions, summed over the positions committed in that window. It says how much extra precision could change the window's decisions. The second is influence: take the states just before and just after the window, and measure how much the full-precision model's beliefs about the still-masked positions shift between them. It says how consequential those decisions are. A proposition bounding accuracy change by commitment perturbation and score sensitivity motivates the product, though the authors stress it is a proxy, not a bound.

T2Quant turns that signal into a schedule. For each task, a set of unlabeled calibration prompts is generated once at full precision, and every saved state is replayed under each candidate activation width: three, four, six and eight bits. For each window and width, the benefit is the window's influence times the drop in divergence relative to four bits, so negative values price a downgrade. A dynamic program over windows and cumulative budget then picks one width per window to maximise total benefit, with the average held at exactly four bits. With these levels, one eight-bit window is paid for by four three-bit windows. The program is exact, and may return the uniform schedule. At inference each step looks up its width by window index, and since the quantizer is never modified, the schedule sits on top of existing methods.

First, does the proxy predict? On GSM8K and MMLU-Pro it picks the most useful window to upgrade, recovering ninety and one hundred percent of the gap to full precision, and it tracks the measured window gains with correlations of zero point nine eight and zero point nine six. HumanEval is noisier, at zero point six three. The main table covers three models, six benchmarks, and three existing quantization methods, DLLMQuant, FAIR-Calib and STaR-Quant, each with and without T2Quant at the same budget. Average accuracy rises in all nine pairings, by up to six point seven points on Dream. Not every cell improves: forty-two of fifty-four benchmark comparisons do.

The comparison that most directly tests the thesis is the random schedule: the same bits, moved around at random. On Dream it falls fourteen to nineteen points below the plain baselines. Mixed precision is not the gain in itself. Where the bits go is. The budget ablation on GSM8K makes the same point from the other side. With only three point seven five average bits, T2Quant reaches forty-eight point four percent accuracy, against forty-four point one for uniform four-bit activations, and raising the budget beyond four bits adds almost nothing.

The paper has a short limitations section, naming hardware efficiency and model safety as future work. The first matters more than its length suggests: quantization here is simulated, so no speedup or memory saving is measured. Most experiments use a single quantization seed, HumanEval has one hundred held-out problems, and the authors note that the motivating bound is vacuous under the greedy decoding they use. Gains vanish at eight tokens per step. And on GSM8K, the appendix reports that each schedule in the family the allocator chooses from coincides with a calibration-free rule: eight bits on the first few windows, three bits on the rest. There, the calibration's real work is deciding how many windows to upgrade.

The idea to carry away is that in a diffusion language model, a bit spent at one step does not buy the same as a bit spent at another, and the exchange rate depends on what the model is writing. A window matters when its quantized predictions drift and when what it commits shapes what comes after. Both can be read from the model's own beliefs. Even the model changes the map: on Dream, HumanEval's damage arrives late, around steps ninety-six to one hundred and four, rather than in the middle as on LLaDA. A fixed precision schedule is a guess. T2Quant replaces the guess with a measurement, then spends the same budget where the measurement says it pays.