Sampling-Aware Quantization for Diffusion Models

TL;DR

This paper introduces a sampling-aware quantization method for diffusion models that improves sampling speed and quality by aligning quantization with the sampling trajectory, enabling faster and more accurate visual generation.

Contribution

It proposes a novel sampling-aware quantization strategy with trajectory alignment to enhance diffusion model sampling efficiency and fidelity.

Findings

Preserves rapid convergence of high-speed samplers.

Maintains superior generation quality with sparse-step sampling.

Achieves dual acceleration in diffusion sampling.

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 is publicly available at: https://github.com/TaylorJocelyn/Sampling-aware-Quantization.