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

Ofek Avivi; Eliron Rahimi; Avi Mendelson; Chaim Baskin

cs.AIOct 22, 2026

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

Authors: Ofek Avivi, Eliron Rahimi, Avi Mendelson, Chaim Baskin

Organizations: Technion – Israel Institute of Technology · INSIGHT Lab, Ben-Gurion University of the Negev

Abstract

Diffusion Language Models (DLMs) introduce unique quantization challenges due to their iterative denoising process. In this work, we reveal that quantization sensitivity is both taskand time-dependent, with different downstream tasks exhibiting distinct precision requirements across denoising steps. We further show that these patterns can be predicted without groundtruth labels or task-specific evaluation metrics. Building on these findings, we introduce T 2Quant, a task and time-aware mixedprecision 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 2Quant consistently improves the average accuracy of existing quantization methods, achieving gains of up to 6.7 percentage points under the same budget.

Figures & tables

Figure 1: Task-dependent temporal sensitivity to activation quantization. Each curve shows, per eight-step denoising window, the KL divergence

D_{\mathrm{KL}}(p_{\mathrm{fp}}\,\|\,p_{\mathrm{q}})

summed over the positions committed in that window that the task’s scoring function reads.

Figure 2: T 2 Quant overview. (1) A few unlabeled prompts per task are profiled with the full-precision and quantized models. (2) Divergence

\times

influence gives the predicted gain of upscaling each denoising window. (3) An exact dynamic program selects the windows to upscale under the bit budget. (4) The result is a task-specific schedule at the same average precision as uniform

A_{\mathrm{target}}

Figure 3: KL

\times

influence predicts where additional precision pays off. Bars show the accuracy gain from upscaling a single denoising window from

\mathrm{A}_{\mathrm{target}}

\mathrm{A}_{\mathrm{up}}

, normalized by the full-precision gap (dashed line, 1.0 = full-precision accuracy), with 95% bootstrap intervals. The black curve shows the label-free proxy

K_{w}I_{w}

, normalized to its peak window.

r

is the Pearson correlation between the two across windows.

Dream-7B-Instruct
Method	Avg.	Evaluation Benchmarks $\uparrow$
		GSM8K	Hum.	ARC-C	MMLU-Pro	DROP	BBH
FP16	44.05	62.50	19.00	87.30	29.69	37.89	27.93
RTN	0.00	0.00	0.00	0.00	0.00	0.00	0.00
DLLMQuant	33.99	35.55	9.00	79.30	24.80	38.09	17.19
FAIR-Calib	34.81	44.92	4.00	81.64	20.70	41.60	16.02

Table 1: W4A4 accuracy (%) of three diffusion-LM quantization methods, each with and without T 2 Quant, on three masked-diffusion LMs and six benchmarks at the same 4.00 average activation bits (bold: best, underline: second-best quantized result).

Figure 4: Ablations on LLaDA-8B-Instruct, GSM8K (

n=512

; bars are 95% intervals): (a) calibration prompts (mean over three random subsets per size), (b) tokens per step with the schedule calibrated at 4, (c) average activation bits, with uniform W4A4 at 4 bits and, at every budget, a random schedule with the same bits as T 2 Quant.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: A few unlabeled prompts suffice. Proxy computed from

N

random calibration prompts. Top: correlation with the all-prompt proxy. Bottom: correlation with held-out accuracy gains.

Figure 6: Label-free diagnostic for LLaDA-1.5 under W4A4, computed as in Figure 1 : score-relevant KL per 8-step window, normalized to each task’s peak, with 95% problem-bootstrap bands.

Figure 7: Label-free diagnostic for Dream-v0-Instruct-7B under W4A4, computed as in Figure 1 : score-relevant KL per 8-step window, normalized to each task’s peak, with 95% problem-bootstrap bands.

Figure 8: Further ablations on LLaDA-8B-Instruct, GSM8K (

n=512

; bars are 95% intervals): (a) steps per window, (b) number of 8-bit windows

m

, the circle marking the allocation’s own choice. Orange: uniform W4A4; dashed: FP16.

Dream-7B-Instruct
Method	Avg.	GSM8K	HumanEval	ARC-C	MMLU-Pro
FP16	49.62	62.50	19.00	87.30	29.69
RTN	0.00	0.00	0.00	0.00	0.00
DLLMQuant	37.16	35.55	9.00	79.30	24.80
FAIR-Calib	37.81	44.92	4.00	81.64	20.70
STaR-Quant	41.57	46.09	11.00	83.01	26.17

Appendix

Table 2: Precision-set ablation: W4A4 accuracy (%) with T 2 Quant allocating over the six activation widths

\{3,4,5,6,7,8\}

instead of the four of Table 1 ,

\{3,4,6,8\}

, at the same 4.00 average activation bits. Avg. is over the four benchmarks shown (bold: best, underline: second-best quantized result).

Proxy	GSM8K	MMLU-Pro	HumanEval
$\mathbb{E}[K_{w}]\,\mathbb{E}[I_{w}]$	0.98	0.96	0.63
$\mathbb{E}[\sqrt{K_{w}I_{w}}]$	0.98	0.85	0.61
$\mathbb{E}[K_{w}]$ only	0.97	0.68	$-$ 0.08
$\mathbb{E}[I_{w}]$ only	0.98	0.93	0.66

Appendix

Table 3: Label-free proxies vs. measured window gains (Figure 3 ): Pearson

r

over the 16 windows.

W8A8, levels $\{4,8,16\}$ , 8.00 bits
	LLaDA-1.5-8B			Dream-7B-Instruct
Row	GSM8K	HumanEval	ARC-C	GSM8K	HumanEval	ARC-C
FP16	57.0	26.0	82.6	62.5	19.0	87.3
uniform W8A8	58.0	24.0	82.4	62.3	15.0	86.1
random schedule	58.6 (+0.6)	25.0 (+1.0)	82.4 (+0.0)	62.3 (+0.0)	16.0 (+1.0)	85.7 (-0.4)
T 2 Quant	61.5 (+3.5 ∗ )	24.0 =	83.0 (+0.6)	60.4 (-2.0)	14.0 (-1.0)	85.9 (-0.2)

Appendix

Table 4: Other precisions (LLaDA-1.5-8B and Dream-7B-Instruct; GSM8K / HumanEval / ARC-C,

n=512

/ 100 / 512): accuracy (%), in parentheses the paired difference to the uniform row ( ∗ :

p<0.05

); = : the allocation returned the uniform schedule.

Benchmark	Eval.	Calib.	Tokens	Steps
GSM8K	512	256	512	128
HumanEval	100	50	512	128
ARC-Challenge	512	256	128	32
MMLU-Pro	512	256	512	128
DROP	512	256	64	16
BBH	512	256	256	64

Appendix

Table 5: Benchmarks: held-out problems, calibration prompts, output tokens and denoising steps (4 tokens per step).

Neural Network Quantization

Neural network quantization represents weights, activations, and states with fewer bits to reduce memory use, computation, and deployment costs.

See all 404 papers

TMPDiff: Temporal Mixed-Precision for Diffusion Models
Mar 14, 2026Basile Lewandowski, Simon Kurz, Aditya Shankar +3Diffusion Model Quantization Diffusion Model Inference Acceleration
University of Neuchâtel, Switzerland · TU Dortmund University, Germany · LAMARR Institute, Germany +2
SoloQ: Calibration-Free Quantization for Diffusion Language Models
Oct 5, 2026Donghyun Lee, Arkapravo Ghosh, Varun Manjunath +5Diffusion Model Quantization KV-Cache Quantization
University of Southern California · Yale University · Korea University
STaR-Quant: State-Time Consistent Post-Training Quantization for Diffusion Large Language Models
Jun 3, 2026Xin Yan, Aqiang Wang, Zhenglin Wan +2LLM Quantization Diffusion Model Quantization
School of Artificial Intelligence, Beijing Normal University, Beijing, China · Department of Computer Science, National University of Singapore, Singapore · Centre for Frontier AI Research, Agency for Science, Technology and Research (A*STAR), Singapore

Mar 14, 2026J/K move · Enter open · S save

TMPDiff: Temporal Mixed-Precision for Diffusion Models

Basile Lewandowski, Simon Kurz, Aditya Shankar, Robert Birke+2

University of Neuchâtel, Switzerland · TU Dortmund University, Germany · LAMARR Institute, Germany+2

Diffusion Model Quantization Diffusion Model Inference Acceleration Mixed-Precision Quantization

Open paper AskPDF arXiv

Diffusion models are the go-to method for Text-to-Image generation, but their iterative denoising processes has high inference latency. Quantization reduces compute time by using lower bitwidths, but applies a fixed precision across all denoising timesteps, leaving an entire optimization axis unexplored. We propose TMPDiff, a temporal mixed-precision framework for diffusion models that assigns different numeric precision to different denoising timesteps. We hypothesize that quantization errors accumulate additively across timesteps, which we then validate experimentally. Based on our observations, we develop an adaptive bisectioning-based algorithm, which assigns per-step precisions with linear evaluation complexity, reducing an otherwise exponential search problem. Across four state-of-the-art diffusion models and three datasets, TMPDiff consistently outperforms uniform-precision baselines at matched speedup, achieving 10 to 20% improvement in perceptual quality. On FLUX.1-dev, TMPDiff achieves 90% SSIM relative to the full-precision model at a speedup of 2.5x over 16-bit inference.

Abstract

Neural Network Quantization

Explore similar work

TMPDiff: Temporal Mixed-Precision for Diffusion Models

SoloQ: Calibration-Free Quantization for Diffusion Language Models

STaR-Quant: State-Time Consistent Post-Training Quantization for Diffusion Large Language Models