T²QuantTask- and time-aware quantizationPreprint

Different Tasks, Different TimestepsTemporal Task-Aware Quantization for Diffusion Language Models

1Technion – Israel Institute of Technology 2INSIGHT Lab, School of Electrical and Computer Engineering, Ben-Gurion University of the Negev *Equal contribution
Technion – Israel Institute of Technology Ben-Gurion University of the Negev INSIGHT Lab
Finding

Precision matters at different denoising steps for different tasks: GSM8K in the first eight steps, HumanEval in the middle of the trajectory.

Signal

Where extra bits pay off is predictable without labels: quantization divergence times each window's influence on the rest of the generation.

Result

Added to three diffusion-LM quantizers, T²Quant raises average accuracy on all nine base/model pairs, by up to 6.7 points at the same 4.00 average bits.

One generation, one precision budget
step 128 / 128window 16 / 16average 4.00 bits
Schedulebits per window
masked 3-bit 4-bit 6-bit 8-bit activations
A 512-token generation by LLaDA-8B-Instruct: 128 denoising steps, 4 tokens committed per step, 16 windows of 8 steps. Each token is shaded by the activation width of the step that committed it. Every schedule averages exactly 4.00 bits, the same as uniform W4A4. The schedules are calibrated per task (levels 3, 4, 6, 8); the order in which tokens unmask is illustrative.
01 · Overview

Precision as a resource spent over time

Diffusion Language Models (DLMs) introduce unique quantization challenges due to their iterative denoising process. In this work, we reveal that quantization sensitivity is both task- and time-dependent, with different downstream tasks exhibiting distinct precision requirements across denoising steps. We further show that these patterns can be predicted without ground-truth labels or task-specific evaluation metrics. Building on these findings, we introduce T²Quant, a task and time-aware mixed-precision quantization framework that uses a label-free sensitivity signal to optimally allocate activation precision across denoising steps under a fixed bit-width budget. Extensive experiments across three DLMs and six benchmarks demonstrate that T²Quant consistently improves the average accuracy of existing quantization methods, achieving gains of up to 6.7 percentage points under the same budget.

9 / 9base/model pairs gain average accuracy
42 / 54benchmark comparisons improve; GSM8K and ARC-C improve in all nine
4.00average activation bits, the same as uniform W4A4
16unlabeled prompts recover about 75% of the gain
02 · Finding

When precision matters depends on the task

Uniform activation quantization treats every denoising step alike. Measured at the positions each window commits, the damage is concentrated, and it is concentrated in different places for different tasks.

Quantization divergence over the denoising trajectory
GSM8K HumanEval MMLU-Pro

KL(pfp ‖ pW4A4) at the score-relevant positions committed in each 8-step window, per committed position, normalized to each task's peak (256 calibration prompts; 50 on HumanEval; full-precision trajectory, W4A4 base). The paper's Figure 1 measures the same quantity along the W4A4 trajectory. Hover for values.

Time. A uniform spread would give each of the 16 windows 6.25% of the task-relevant divergence. Along the W4A4 trajectory of LLaDA-8B-Instruct, GSM8K's first window carries 27% (4.4× uniform) and its first quarter 55%. On HumanEval, five mid-trajectory windows (steps 32–71) carry 52%.

Task. GSM8K is strongly front-loaded. MMLU-Pro is early but broad: its first four windows are all within 10% of its peak. HumanEval is insensitive early and peaks broadly around steps 40–63. The profiles carry over to LLaDA-1.5 (per-window correlation 0.99, 0.98 and 0.92), so they survive preference alignment.

Key takeawaySensitivity varies across denoising steps, and the pattern differs by task. A single fixed temporal schedule is unlikely to suit every task.
03 · Method

Profile, score, allocate

T²Quant does not change the quantizer, the weights or the sampler. It only decides which activation width each denoising window runs at, and it needs a few unlabeled prompts from the task to do so.

Profile local divergence

Generate a full-precision reference trajectory per calibration prompt and replay each saved state at every candidate width b. Holding the state fixed isolates the local effect of precision.

κ_τ,k(b) = Σ KL(p_fp ‖ p_b) over committed positions in window k / number committed

Weight by influence

A window matters only if its commitments steer what comes after. Influence is how much they move the full-precision model's beliefs about the positions still masked.

ι_i,k = Σ over still-masked j of KL(p_fp(·|after) ‖ p_fp(·|before))
g_τ,k(b) = ῑ_τ,k [κ_τ,k(b0) − κ_τ,k(b)]

Allocate exactly

Pick one width per window to maximize the predicted gain at the uniform budget. A dynamic program over windows and cumulative cost finds the global optimum, and returns the uniform schedule when no trade pays.

b* = argmax Σ g_k(b_k) s.t. Σ c_k(b_k) = C0
F_k(z) = max_b { F_k−1(z − c_k(b)) + g_k(b) }

Why divergence times influence. For one window, let Jw be the KL between its commitments at the upscaled and target widths, and Rw the range of the expected final score over those commitments. The effect of upscaling the window obeys |ΔM_w|² ≤ ½ E[J_w] E[R_w²]. Neither term is observable on one unlabeled trajectory, so T²Quant uses the divergence Kw and the influence Iw as surrogates and ranks windows by S_w = E[K_w] E[I_w].

Divergence alone is not enough. On HumanEval it ranks the final window highest, although tokens committed last cannot affect anything after them. The influence weight removes such windows.

At inference the schedule is a 16-entry lookup indexed by the denoising step. There are no extra kernels and no per-token decisions.

Does the proxy pick the right windows?
TaskPearson rSingle-window upscale (A4 → A8)
GSM8K0.98Proxy's top window is the best one: +7.0 points, 90% of the full-precision gap
MMLU-Pro0.96Proxy's top window is the best one: +4.4 points, the whole gap
HumanEval0.63Gains sit in the first two windows, where the proxy puts its mass; split-half seed noise caps r near 0.77

r: correlation across the 16 windows between the label-free proxy and the measured accuracy gain of upscaling that window alone (LLaDA-8B-Instruct). A proxy from 16 random unlabeled prompts matches the all-prompt proxy at r ≥ 0.96 on every task.

04 · Allocator

Re-solve the schedule yourself

These are the calibrated per-window gains that feed the dynamic program. Change the task or the average-bit budget and the page solves the allocation again, exactly, in your browser.

gain from 8 bits, gk(8) cost of 3 bits, gk(3)
Optimal scheduleat 4.00 bits
Uniformbaseline
← early (fully masked)denoisinglate (fully unmasked) →

Widths {3, 4, 6, 8}, 16 windows, cost = bits per window, so the budget fixes the sum of the 16 widths. Gains are from the calibration runs on the shared QuaRot + GPTQ W4A4 base (in these runs κ counts only score-relevant positions, an earlier variant of the paper's Eq. 3). At 4.00 bits the solver reproduces the stored schedule for all 18 model/task pairs.

05 · Results

Every base method improves, on every model

W4A4 accuracy on three diffusion LMs and six benchmarks. Each quantizer runs as published and with T²Quant on top; both use 4-bit weights and exactly 4.00 average activation bits, so a pair differs only in where along the trajectory the bits go.

Average accuracy, base method → with T²Quant
base method + T²Quant FP16
Table 1 · W4A4 accuracy (%)

Bold: best quantized result; underline: second best. Small blue numbers: change in average over the same base method. Held-out problems are disjoint from every calibration prompt; all rows share the sampler and prompts, so comparisons are paired.

Best quantized accuracy on every model. The two highest quantized averages are T²Quant rows on every model. They beat the strongest standalone baseline by 3.55, 1.13 and 1.40 points on Dream, LLaDA-1.5 and LLaDA-8B-Instruct, and sit 2.0, 2.5 and 1.6 points below FP16.

Largest gains where quantization hurts most. On Dream, GSM8K rises from 35.6 to 53.5 with DLLMQuant and from 46.1 to 58.8 with STaR-Quant. GSM8K gains reach 9.6 points on LLaDA-1.5 and 10.4 on LLaDA-8B-Instruct.

Placement, not budgetA random schedule with the same 4.00 average bits falls 14–19 points below the base methods on Dream. The same bits placed by T²Quant give the best quantized average on every model.
06 · Ablations

Cheap to calibrate, stable, and about placement

LLaDA-8B-Instruct on GSM8K (512 held-out problems), shared W4A4 base, levels {3, 4, 8}. Uniform W4A4 scores 44.1, T²Quant 49.4 and FP16 52.9.

T²Quant uniform W4A4 random schedule, same bits FP16
(a) Calibration prompts
(b) Tokens per step
(c) Average activation bits

(a) Mean over three random subsets per size, with the range; 256 is the full pool. (b) The schedule calibrated at 4 tokens per step, reused unchanged. (c) The random control spends exactly T²Quant's bits, in shuffled windows. Accuracy in %.

16 prompts already recover about three quarters of the gain (48.0 against 49.4 with 256). Twelve random subsets of 16–128 prompts yield five distinct schedules, all the same shape: high precision early, low precision in the last four to eight windows.

The schedule transfers to 2 tokens per step with the same +5.3 points (60.5 against 55.3). At 8 tokens per step every row collapses (FP16 31.6) and the gain shrinks to a tie.

3.75 bits beat uniform 4 bits by 4.3 points, and bits beyond 4 buy nothing. The same bits in random windows stay at the level of uniform A3 (28.7): one early 3-bit window costs as much as making every window 3-bit.

07 · Use it

Run T²Quant on one GPU

The repository is a minimal, tested implementation for LLaDA-8B-Instruct: QuaRot-style rotation, GPTQ 4-bit weights, switchable per-token activation quantization, label-free calibration and the exact allocator.

  1. Install. Python 3.10+ and one CUDA GPU with 40 GB or more. The model and datasets download from the Hugging Face Hub on first use.

  2. Run the notebook. T2Quant.ipynb loads the full-precision reference and the quantized base, calibrates on 16 unlabeled GSM8K training questions, solves the allocation and evaluates FP, uniform W4A4 and T²Quant. It takes about 4–5 hours at the paper's setting.

  3. Scale out. The command line shards calibration and evaluation across GPUs and reports accuracy with 95% bootstrap intervals and an exact sign test.

git clone https://github.com/OfekAvivi/T2Quant.git
cd T2Quant
pip install -r requirements.txt
jupyter notebook T2Quant.ipynb   # or headless: python run.py
python -m pytest tests -q        # CPU tests, seconds
P="--task gsm8k --out runs/gsm8k --gen_len 512 --steps 128 --windows 16 --cache_dir cache"
for i in 0 1 2 3; do CUDA_VISIBLE_DEVICES=$i python -m t2quant.cli calibrate $P --shard $i/4 & done; wait
for c in fp uniform t2quant; do
  for i in 0 1 2 3; do CUDA_VISIBLE_DEVICES=$i python -m t2quant.cli evaluate $P --config $c --shard $i/4 & done; wait
done
python -m t2quant.cli report --out runs/gsm8k
What the repository reproduces
GSM8K schedulelevels 3–8, 16 prompts
LLaDA-8B-Instruct, GSM8K (first 500 test problems)Full precisionUniform W4A4T²QuantFixed / broken
Accuracy (%), 4.00 average activation bits51.246.651.484 / 60

From the repository README. Fixed / broken: problems T²Quant solves and uniform W4A4 does not, and the reverse (exact two-sided sign test p = 0.055). The schedule moves precision to the first windows and pays for it at the end, recovering full-precision accuracy at the same average bits.

Cite

BibTeX

@misc{avivi2026t2quant,
  title         = {Different Tasks, Different Timesteps: Temporal Task-Aware Quantization for Diffusion Language Models},
  author        = {Avivi, Ofek and Rahimi, Eliron and Mendelson, Avi and Baskin, Chaim},
  year          = {2026},
  eprint        = {2610.31198},
  archivePrefix = {arXiv},
  primaryClass  = {cs.AI},
  url           = {https://arxiv.org/abs/2610.31198}
}