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

AI summary

Overview

Research area: Model compression / efficient inference for diffusion language models (DLMs) — specifically post-training quantization of activations, combined with temporal (per-denoising-step) mixed-precision scheduling.

Technical level: Intermediate. The paper is readable with a basic grasp of quantization (bit-widths, W4A4, calibration) and of how diffusion language models denoise; the core idea is intuitive even though the machinery (KL divergence, dynamic programming) is technical.

Scope: The paper shows that sensitivity to activation quantization in DLMs varies jointly across downstream tasks and denoising steps, and introduces T²Quant, a label-free framework that allocates activation precision across denoising windows under a fixed average bit budget.

What This Paper Is About

Diffusion language models generate text by iteratively denoising a masked sequence, which means the model is run many times per output — making quantization attractive for reducing memory and compute. Existing DLM quantization methods, however, keep a single fixed activation precision for the entire denoising trajectory. The authors show that this is wasteful: different tasks need high precision at different points in the denoising process, and the most sensitive points can be predicted without any labels. Their goal is a method that spends a fixed bit budget unevenly over time, per task, on top of existing quantizers.

Key Contributions

1. Task- and time-dependent precision sensitivity. The authors demonstrate that activation-quantization sensitivity in DLMs is strongly non-uniform across denoising steps and that the temporal profile differs by downstream task (GSM8K, HumanEval, MMLU-Pro behave differently).

2. A label-free sensitivity predictor. They build a signal that combines quantization-induced divergence at committed positions with the downstream influence of each denoising window, so precision importance can be estimated from unlabeled calibration prompts — no ground-truth answers, no task-specific metrics.

3. T²Quant, a task- and time-aware mixed-precision scheduler. Given a finite set of activation widths and a fixed average bit budget, the method selects a per-window precision schedule by budget-constrained dynamic programming, with a proof that the solver is globally optimal for the calibrated additive surrogate (Proposition 2).

4. Quantizer-agnostic integration with consistent empirical gains. T²Quant is applied on top of DLLMQuant, FAIR-Calib and STaR-Quant across three DLMs and six benchmarks at the same 4.00 average activation bits, improving average accuracy in all nine baseline/model combinations, by up to 6.7 percentage points.

Main Findings

• Sensitivity is concentrated, not uniform. With 128 denoising steps split into 16 non-overlapping windows of eight steps (6.25% each under uniformity), GSM8K's first window alone carries 27% of the task-relevant divergence (4.4× uniform) and its first quarter 55%; on HumanEval, five mid-trajectory windows (steps 32–71) carry 52%; MMLU-Pro, the most evenly spread task, still varies by nearly 2× between its most and least sensitive windows.

• The temporal profile is task-specific. GSM8K is front-loaded with a peak in the first eight steps and rapid decay; MMLU-Pro is also early but broad, with its first four windows (steps 0–31) all within 10% of its peak; HumanEval is insensitive early and peaks broadly around steps 40–63.

• The profile survives fine-tuning and transfers across architectures, but shifts. LLaDA-1.5's per-window profiles correlate with LLaDA-8B-Instruct's at 0.99 (GSM8K), 0.98 (HumanEval) and 0.92 (MMLU-Pro). On Dream-7B, tasks peak at different points than on LLaDA — HumanEval peaks late at steps 96–104 rather than mid-trajectory.

• A label-free proxy tracks where precision pays off. The product of expected divergence and expected influence correlates with measured window-level accuracy gains at r = 0.98 (GSM8K), 0.96 (MMLU-Pro) and 0.63 (HumanEval), and identifies the single most beneficial window, recovering 90% and 100% of the full-precision gap on GSM8K and MMLU-Pro respectively.

• The influence term is necessary. Using divergence alone fails on HumanEval (r = -0.08), because it wrongly favors the final window, whose commitments cannot affect future tokens.

• Very few unlabeled prompts suffice. Proxy estimates from 16 unlabeled calibration prompts match all-prompt estimates at r ≥ 0.96 across tasks, and 16 prompts recover roughly 75% of the full gain.

• Consistent gains over three quantization baselines. At 4-bit weights and 4.00 average activation bits, T²Quant improves average accuracy for all nine baseline/model combinations, with gains of up to 6.7 points on Dream, 3.4 on LLaDA-1.5 and 3.5 on LLaDA-8B-Instruct; 42 of 54 benchmark-level comparisons improve, and GSM8K and ARC-Challenge improve in every combination.

• Best quantized results approach full precision. T²Quant takes the two highest quantized averages on every model, beating the strongest standalone baseline by 3.55, 1.13 and 1.40 points on Dream, LLaDA-1.5 and LLaDA-8B-Instruct, leaving gaps of only 2.0, 2.5 and 1.6 points to FP16.

• The allocation itself is what matters. A random schedule at the same bit budget performs far worse — 14–19 points below the original baselines on Dream. And T²Quant reaches 48.4% on LLaDA-8B-Instruct/GSM8K with only 3.75 average bits, versus 44.1% for uniform W4A4 at 4 bits (FP16: 52.9), while raising precision beyond 4 bits gives negligible further benefit.

• Robust to granularity and precision set, and useful at tighter budgets. Accuracy stays within 1.7 points (47.7 to 49.4) when moving from 16 steps per window down to a single step; extending the width set from {3,4,6,8} to {3,4,5,6,7,8} raises the mean over 36 composed cells from 45.7 to 46.3. At W4A3 (3.00 average bits) the schedule recovers most of the loss — up to 32 points on Dream ARC-C and 8.6 points on LLaDA-1.5 GSM8K — while at W8A8 it is essentially a no-op (one gain of +3.5, p = 0.04, otherwise ties).

Methodology in Plain English

The pipeline has four steps, and none of them modifies the quantizer itself.

1. Profile with unlabeled prompts. For a target task, take a small set of unlabeled prompts and run the model at full precision, saving the intermediate state before each denoising step and the set of positions the sampler commits at that step.

2. Replay states at several bit-widths. Re-run each saved state under every candidate activation precision, keeping the input state fixed so only local quantization effects are measured. Compare the quantized predictive distributions against the full-precision ones with KL divergence, averaged per committed position — this says how much precision could change a window's decisions.

3. Measure downstream leverage. For each window, measure how much the full-precision model's beliefs about the still-masked positions change from just before to just after the window — this says how consequential those decisions are. Multiply the precision-induced divergence reduction by this influence to score each (window, precision) pair. Proposition 1 motivates this pairing with a bound separating commitment perturbation from downstream score sensitivity.

4. Solve for a schedule. Choose one precision per window to maximize total score subject to matching the uniform-precision budget, via dynamic programming over windows and cumulative budget; backtracking yields the schedule, and the all-uniform solution remains available if reallocation isn't estimated to help. At inference, each step simply looks up its assigned precision — no extra diagnostic passes.

The key design choice is that schedules are computed per task, because the sensitive windows differ by task, and that everything is estimated from internal model signals rather than labels or task metrics.

Why This Matters

Impact on research. The paper reframes activation precision in DLMs as a temporal resource to be allocated, rather than a global setting, and shows the allocation should be task-conditioned. It connects prior task-aware quantization work (which allocates spatially across layers or weights in autoregressive models) with temporal scheduling work from image diffusion, and argues DLMs need a distinct treatment because token commitments are discrete and often irreversible, so latent reconstruction error does not directly reflect downstream correctness. The label-free calibration result also lowers the practical cost of task-specific deployment.

Real-world applications (plausible uses of the efficiency gains; the paper evaluates accuracy, not deployment):

• Serving diffusion language models where repeated denoising passes dominate cost and memory is the binding constraint.
• On-device or edge text generation, where a smaller bit budget matters more than squeezing out the last point of accuracy.
• Task-specialized deployments — a math-reasoning endpoint, a code-generation endpoint, a multiple-choice QA endpoint — each getting its own cheap calibration and schedule.
• Tight-budget regimes: at 3.00 average activation bits the schedule recovers much of the collapse that uniform quantization suffers.

Industry relevance. Because T²Quant sits on top of existing post-training quantizers (DLLMQuant, FAIR-Calib, STaR-Quant) without changing them, it is an incremental add-on rather than a new quantization pipeline, and calibration requires only a few unlabeled prompts. The paper reports gains at matched bit budgets, which is the form that matters for hardware provisioning.

Future Directions

• Hardware efficiency is unmeasured. The authors explicitly leave inference speed and energy consumption to future work; whether a non-uniform schedule maps cleanly onto real kernels is an open engineering question.
• Safety and robustness effects. The paper notes that quantization can alter behavior beyond benchmark accuracy, including safety-related responses, and calls for examining task- and time-aware allocation in that light.
• Broader model and task coverage. Results cover three DLMs (Dream-7B-Instruct, LLaDA-1.5-8B, LLaDA-8B-Instruct) and six benchmarks; other architectures and samplers, and the sampler-sensitivity visible in Dream's EOS handling, remain open.
• Whether sensitivity can be predicted without any calibration at all. The paper observes that for every tested number of 8-bit windows the chosen schedule coincided with a calibration-free prior (8 bits on the first m windows, 3 bits on the rest), suggesting that on some tasks the calibration earns its keep mainly by choosing m.

Target Audience

Graduate students and researchers in efficient machine learning and model compression; practitioners deploying diffusion language models under memory or compute constraints; and engineers working on post-training quantization who want a task-aware, temporally resolved complement to existing quantizers. Readers new to quantization will follow the high-level story but may need background for the divergence/influence formulation and the dynamic program.

Authors’ 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.