Quantizing Diffusion Models from a Sampling-Aware Perspective

TL;DR

Diffusion models excel in visual generation but suffer from slow sampling and heavy noise estimators. The authors show that quantization noise disrupts directional estimates, especially for high-order samplers, potentially degrading the sampling trajectory. They introduce a sampling-aware quantization framework centered on Mixed-Order Trajectory Alignment to enforce a more linear probability flow, enabling parallel improvements in sampling speed and model efficiency. Across multiple benchmarks with sparse-step sampling, the proposed SA-PTQ and SA-QLoRA methods achieve strong generation quality while maintaining fast convergence, narrowing the gap to full-precision baselines.

Abstract

Diffusion models have recently emerged as the dominant approach in visual generation tasks. However, the lengthy denoising chains and the computationally intensive noise estimation networks hinder their applicability in low-latency and resource-limited environments. Previous research has endeavored to address these limitations in a decoupled manner, utilizing either advanced samplers or efficient model quantization techniques. In this study, we uncover that quantization-induced noise disrupts directional estimation at each sampling step, further distorting the precise directional estimations of higher-order samplers when solving the sampling equations through discretized numerical methods, thereby altering the optimal sampling trajectory. To attain dual acceleration with high fidelity, we propose a sampling-aware quantization strategy, wherein a Mixed-Order Trajectory Alignment technique is devised to impose a more stringent constraint on the error bounds at each sampling step, facilitating a more linear probability flow. Extensive experiments on sparse-step fast sampling across multiple datasets demonstrate that our approach preserves the rapid convergence characteristics of high-speed samplers while maintaining superior generation quality. Code will be made publicly available soon.

Figures (12)

• Figure 1: Comparison of generated samples on the ImageNet 256×256 dataset between full-precision LDM-4 and its quantized versions using PTQ4DM, Q-diffusion, PTQD, EfficientDM and our proposed SA-QLoRA).
• Figure 2: Direction estimation in reverse diffusion sampling. (a) The first-order sampler performs a single direction estimation at the beginning of the sampling interval. (b) The second-order sampler refines the direction estimation by evaluating additional intermediate steps within the interval. (c) Quantization errors lead to deviations in direction estimation, causing the intermediate steps in high-order samplers to drift over time, ultimately impacting the final direction estimation. (d) Our proposed Mixed-Order Trajectory Alignment achieves a more linearized probability flow.
• Figure 3: Sampling-aware quantization workflow. (a) Module-level reconstruction process employed in SA-PTQ, where $\hat{f}_i(\cdot)$ denotes the module undergoing quantization and reconstruction. (b) Basic fine-tuning workflow in SA-QLoRA, where LoRA weights $W_{LoRA}$ and quantization parameters $s, z$ are iteratively updated after each sampling step.
• Figure 4: Visualization of the generative performance of our SA-QLoRA under $W4A8$ and $W4A4$ quantization settings.
• Figure 5: Latent space feature trajectories of LDM4 under 20-step sampling on the ImageNet 256$\times$256 dataset. (a) Feature trajectories sampled using DPM-Solver-1. (b) Intermediate-step feature trajectories sampled using DPM-Solver-2.