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

1 Introduction

• The paper identifies as unexplored how quantization sensitivity varies jointly across DLM tasks and denoising steps, and how those variations can guide mixed-precision allocation.Existing methods generally maintain fixed activation precision throughout denoising.Source: §1 · p. 1
• Different downstream tasks exhibit distinct temporal sensitivity patterns, with precision more important early for some tasks and at intermediate stages for others.Figure 1 reports task-relevant divergence across denoising windows.Source: §1 · p. 1, Figure 1 · p. 1
• A label-free signal combines quantization-induced divergence with each denoising stage’s downstream influence to predict where additional precision matters, using unlabeled calibration prompts.It does not require ground-truth labels or task-specific metric evaluations.Source: §1 · p. 2
• Across three DLMs and six benchmarks, T2Quant improves average accuracy over quantized baselines by up to 6.7 percentage points under a fixed budget.It is complementary to existing post-training quantization methods and does not require modifying their underlying quantizers.Source: §1 · p. 2
• T2Quant uses task-specific sensitivity patterns to allocate activation precision across denoising steps under a fixed average bit-width budget.Its overview shows profiling, estimating gains from divergence × influence, and selecting windows with dynamic programming.Source: §1 · p. 2, Figure 2 · p. 2

DLM Quantization

• Existing DLM quantization methods address diffusion-state and temporal-error challenges, but keep activation bit-width fixed across the denoising trajectory.The paper instead treats activation precision as a resource allocated across steps under a fixed average bit budget.Source: §2 · p. 2

Task-Aware Quantization

• Task-aware quantization methods allocate precision spatially across layers or weights, whereas T2Quant allocates activations temporally across denoising steps according to task-specific sensitivity.The cited task-aware approaches concern autoregressive LLMs.Source: §2 · pp. 2–3

Temporal Precision Scheduling

• Image-diffusion methods schedule precision by timestep, while T2Quant accounts for discrete token commitments and task-relevant downstream influence to make schedules task-specific.In DLMs, latent reconstruction error does not directly reflect downstream correctness, and a timestep’s importance depends on which task-relevant tokens it commits.Source: §2 · p. 3

3.1 Temporal Sensitivity Diagnostic

• On GSM8K, the first eight-step window carries 27% of task-relevant divergence and the first quarter carries 55%; on HumanEval, five mid-trajectory windows carry 52%.MMLU-Pro’s most and least sensitive windows differ by nearly 2×, showing sensitivity is nonuniform across the trajectory.Source: §3.1 · p. 3
• Sensitivity profiles differ by task: GSM8K peaks in the first eight steps, MMLU-Pro is broadly sensitive early, and HumanEval peaks mid-trajectory.This variation challenges a single task-independent temporal precision schedule.Source: §3.1 · pp. 3–4

3.2 A Label-Free Signal for Precision Allocation

• The proposed label-free proxy ranks a window’s value using divergence and downstream influence, measured from an Atarget trajectory without labels.Divergence captures how additional precision can change window decisions; influence captures changes in full-precision beliefs about positions still masked afterward.Source: §3.2 · p. 4
• On GSM8K and MMLU-Pro, the proxy identifies the most beneficial window and recovers 90% and 100% of the full-precision gap, respectively; HumanEval’s correlation is r = 0.63.Influence matters: divergence alone incorrectly favors HumanEval’s final window, whose commitments cannot affect future tokens.Source: §3.2 · p. 4
• Estimates from 16 unlabeled calibration prompts closely match estimates computed from all prompts, with r ≥0.96 across tasks.The proxy is a motivated predictor, not a formal bound, and its predictive ability is evaluated empirically.Source: §3.2 · p. 4

4 Method

• T2Quant extends single-window sensitivity analysis into a task-specific activation-precision schedule under a fixed budget, using unlabeled prompts to estimate precision value.The schedule is designed to build on the label-free signal that combines quantization divergence with downstream influence.Source: §4 · pp. 4–5

4.1 Problem Setup

• The objective assigns supported activation precisions to ordered denoising windows to maximize the task metric while matching the uniform-reference budget.The ideal objective requires labeled, closed-loop generations, so the method instead constructs a label-free surrogate.Source: §4.1 · p. 5

4.2 Label-Free Calibration

• Calibration replays saved full-precision states at candidate precisions, averaging divergence over committed positions to isolate local quantization effects.Using fixed replay states separates local quantization effects from changes in the generated trajectory, without task-specific token selection.Source: §4.2 · pp. 5–6
• The calibration score combines local divergence changes with downstream influence; positive scores favor upgrades, while negative scores estimate the cost of reducing precision.The score is a surrogate for marginal downstream benefit, not the true change in task accuracy.Source: §4.2 · p. 6

4.3 Optimal Budget-Constrained Scheduling

• Dynamic programming jointly allocates precision across windows at the reference budget, and backtracking returns the selected schedule.The solver can retain the uniform-precision schedule when estimated gains from reallocating precision are insufficient.Source: §4.3 · p. 6
• The recurrence is globally optimal when the budget is feasible and costs are finite, nonnegative, integer-valued, but the guarantee applies to the additive surrogate.The true downstream metric is generally non-additive; calibration produces one schedule per task, with no additional diagnostic forward passes at inference.Source: §4.3 · p. 6

5.1 Experimental Setup

• Evaluation covers three open DLMs and six benchmarks, using held-out prompts disjoint from calibration prompts.Comparisons include FP16, RTN, and three post-training quantization methods, with and without T2Quant.Source: §5.1 · pp. 6–7, §5.2 · p. 7

5.2 Main Results

• T2Quant improves average accuracy in all nine baseline/model combinations, with gains up to 6.7 points on Dream.Overall, 42 of 54 benchmark comparisons improve, and GSM8K and ARC-Challenge improve in every combination.Source: §5.2 · p. 7
• T2Quant achieves the two highest quantized averages on every model, beating the strongest standalone baselines by 3.55, 1.13, and 1.40 points on Dream, LLaDA-1.5, and LLaDA-8B-Instruct.The best configurations are within 2.0, 2.5, and 1.6 points of FP16, respectively.Source: §5.2 · p. 7, Table 1 · p. 7

5.3 Ablation Study

• On LLaDA-8B-Instruct and GSM8K, 16 calibration prompts recover roughly 75% of the full gain, with stable schedules prioritizing early denoising steps.Figure 4 reports mean results over three random subsets per calibration-prompt size.Source: §5.3 · p. 8, Figure 4 · p. 8
• The schedule transfers from 4 to 2 tokens per step with an approximately 5-point improvement, while gains vanish at 8 tokens per step as accuracy deteriorates.Source: §5.3 · p. 8, Figure 4 · p. 8
• At 3.75 average activation bits, T2Quant reaches 48.4% accuracy versus 44.1% for uniform W4A4, and it consistently beats random schedules at matched budgets.Increasing precision beyond 4 bits provides negligible additional benefit.Source: §5.3 · p. 8, Figure 4 · p. 8

• Across three DLMs and six benchmarks, T2Quant improves average accuracy with existing quantization methods, gaining up to 6.7 percentage points under W4A4.Source: §6 · p. 8
• The authors identify task requirements and temporal generation dynamics as relevant considerations when quantizing DLMs.Source: §6 · p. 8
• Future work should examine safety, robustness, and real-world inference efficiency; quantization can also alter behavior beyond benchmark accuracy.The discussion notes that deployment may potentially reduce memory and energy requirements, but does not report these as measured outcomes.Source: §6 · p. 8

• Hardware efficiency, including inference speed and energy consumption, and potential effects on model safety remain unevaluated.The reported results cover three DLMs, six benchmarks, and three quantization methods.Source: Limitations · p. 9

• LLMs were used mainly for grammar checking and improving manuscript clarity and naturalness; authors manually reviewed and edited all suggestions.Source: LLM Usage Statement · p. 9
• LLMs did not play an important role in developing the methods, experiments, analysis, or conclusions.Source: LLM Usage Statement · p. 9

Ethical Considerations

• Quantization’s effects on model safety and robustness are left for future work.The work aims to improve diffusion language model efficiency and accessibility.Source: Ethical Considerations · p. 9

Appendix Overview

• The appendix provides supplementary analyses, experiments, ablations, theoretical proofs, algorithms, and implementation details.Its sections cover further results, ablations, proofs, algorithms, and experimental protocols.Source: Appendix Overview · p. 12

A.1 Temporal Sensitivity Diagnostic

• LLaDA-1.5’s temporal profiles closely match LLaDA-8B-Instruct, with per-window correlations of 0.99 on GSM8K, 0.98 on HumanEval, and 0.92 on MMLU-Pro.LLaDA-1.5 is preference-optimized but retains the same architecture, tokenizer, sampler, and EOS rules; this comparison does not test other architectures.Source: §A.1 · p. 12
• Dream’s profiles remain task-dependent, with GSM8K damage concentrated in the first window, HumanEval peaking at steps 96–104, and MMLU-Pro rising toward the end.HumanEval’s late peak differs from its mid-trajectory peak on LLaDA.Source: §A.1 · p. 13

B Additional Ablation Studies

• Using windows from 16 steps down to one keeps GSM8K accuracy within 1.7 points (47.7 to 49.4), so coarse windows perform as well as fine ones.Figure 8a reports the LLaDA-8B-Instruct GSM8K ablation with 512 problems and 95% intervals.Source: B Additional Ablation Studies · p. 13, Figure 8 · p. 13
• At a fixed 4.00-bit average, choosing m = 1 or m = 3 instead of m = 2 costs 1.0 and 3.5 points, respectively; calibration selects the optimum in this family.For each m, the schedule matches a calibration-free prior, so on GSM8K calibration earns its keep in choosing m.Source: B Additional Ablation Studies · p. 13
• Averaging across 36 composed cells, expanding the precision set raises mean accuracy from 45.7 to 46.3, while a T2Quant row is best in 11 of 12 model–benchmark cells.The additional widths change accuracy little; the main experiments retain {3, 4, 6, 8} because each added width requires another calibration-trajectory replay.Source: B Additional Ablation Studies · p. 13, Table 2 · p. 14
• At W4A3, T2Quant recovers up to 32 points on Dream ARC-C and 8.6 points on LLaDA-1.5 GSM8K, while W8A8 is already at full precision under uniform quantization.At W8A8, the schedule is a no-op except for a +3.5 gain on LLaDA-1.5 GSM8K (p = 0.04).Source: B Additional Ablation Studies · p. 14
• The proxy’s divergence-only alternative fails on HumanEval (r = −0.08), while influence alone ranks windows nearly as well as the combined signal.The divergence term is still needed to price each precision level; the comparison evaluates proxy rankings against measured window gains.Source: B Additional Ablation Studies · p. 14, Table 3 · p. 14

C.1 Proof of Proposition 1

• The proposition bounds expected task-score change using commitment-law divergence and downstream score sensitivity, but the practical surrogates do not bound the change or guarantee its ranking.Under greedy decoding, the bound is vacuous because commitment laws are point masses whenever commitments differ.Source: §C.1 · p. 15

C.2 Proof of Proposition 2

• Induction establishes that the dynamic program’s recurrence returns the best feasible assignment at each budget, and backtracking recovers an optimal assignment.Its optimality is for the additive surrogate objective, not downstream task accuracy.Source: §C.2 · p. 15, D Algorithms · p. 16

D Algorithms

• T2Quant calibration combines quantization divergence at committed positions with downstream influence, then allocates precision under a fixed budget.The dynamic program has O(KC0|B|) time and O(KC0) memory.Source: D Algorithms · p. 16

E.1 Model, Sampler, and Quantization

• The default experiments use deterministic LLaDA-8B-Instruct generation over 512 tokens and 128 denoising steps, with 4-bit weights and per-token 4-bit activations.The setup uses temperature 0, no classifier-free guidance, and a standard low-confidence remasking sampler.Source: §E.1 · pp. 16–17

E.2 Main Results (Section 5.2)

• The experiments cover Dream-7B-Instruct, LLaDA-1.5, and LLaDA-8B-Instruct, which differ in training or alignment and, for Dream, architecture and initialization.Dream is described as the most sensitive of the three to activation quantization.Source: §E.2 · p. 17
• Each benchmark uses a calibration pool seen by calibrated methods and a shared held-out evaluation pool, making paper comparisons paired by problem.Accuracy is based on normalized exact-match answers, with DROP using its official answer normalization.Source: §E.2 · p. 17
• Rows use each model’s upstream sampler unchanged, so they differ only in the quantization applied to the network.The sampler unmasks four tokens per step with low-confidence remasking at temperature 0.Source: §E.2 · p. 18
• T2Quant keeps the shared rotated GPTQ W4 base and varies activation width across 16 trajectory windows under an exact 4.00-bit average.The random-schedule comparator uses the same average-bit budget; baseline-plus-T2Quant retains each baseline’s calibration.Source: §E.2 · p. 18