FAIR-Calib: Frontier-Aware Instability-Reweighted Calibration for Post-Training Quantization of Diffusion Large Language Models
Abstract
Diffusion Large Language Models (dLLMs) refine tokens iteratively but commit them irreversibly, leading to a "stability lag" where early decisions remain fragile even after being written. We reveal that Post-Training Quantization (PTQ) error easily flips these borderline decisions at the write frontier, which are then permanently locked in and amplified. To address this, we propose Frontier-Aware Instability-Reweighted Calibration (FAIR-Calib), a two-stage PTQ framework for dLLMs. Stage I probes a full-precision teacher to estimate a position prior that combines frontier hits and masked-stage reliability. Stage II performs off-policy, layer-wise calibration by minimizing a reweighted hidden-state MSE, effectively prioritizing the protection of fragile frontier states without requiring expensive end-to-end diffusion rollouts. We further theoretically justify our weighted objective as a surrogate for output KL divergence. Empirically, FAIR-Calib consistently outperforms state-of-the-art baselines on LLaDA and Dream (W4A4), significantly reducing frontier decision flips and suppressing post-commit mismatches across diverse benchmarks.
Lay Summary
Diffusion large language models generate text by repeatedly refining masked token positions and then irreversibly committing some tokens during decoding. We find that commitment does not always mean stabilization: some committed tokens remain fragile for later refinement steps, making them especially sensitive to low-bit quantization. Small numerical perturbations can therefore flip fragile commit decisions; once an incorrect token is committed, it becomes fixed context and can affect later refinement steps. We propose FAIR-Calib, a two-stage post-training calibration method that first estimates a theoretically motivated, frontier-aware position prior from the full-precision model and then uses this prior to reweight hidden-state calibration for the quantized model. This makes diffusion language models more memory-efficient while better preserving their generation quality.