TLDR: DMQ is a new post-training quantization (PTQ) method for diffusion models that addresses performance degradation at low bit-widths by effectively handling outliers. It combines Learned Equivalent Scaling (LES) to redistribute quantization difficulty and channel-wise Power-of-Two Scaling (PTS) with a robust voting algorithm to manage extreme outliers. This approach significantly improves image generation quality and model stability, outperforming existing methods, especially at ultra-low bit-widths like W4A6 and W4A8.
Diffusion models have rapidly become a cornerstone in the field of artificial intelligence, particularly for their remarkable ability to generate high-fidelity images. From creating realistic photos to enabling advanced image editing and even 3D and video generation, their capabilities are truly impressive. However, this power comes at a significant computational cost. The iterative nature of their denoising process often requires hundreds or even thousands of steps, making them challenging to deploy in environments with limited resources, such as mobile devices or edge computing.
The Challenge of Quantization
To address these computational demands, researchers often turn to quantization. This technique reduces the memory and processing power required by converting high-precision floating-point values within neural networks into lower-bit integer approximations. While effective for many neural networks, applying quantization to diffusion models presents unique hurdles. Their iterative process leads to highly varied activation distributions across different time steps, and crucially, quantization errors can accumulate as the denoising progresses, leading to a noticeable drop in the quality of the final output, especially at very low bit-widths.
Existing post-training quantization (PTQ) methods for diffusion models have tried various approaches, such as carefully composing calibration data or adapting quantization parameters. However, many of these methods often overlook a critical issue: outliers. These are extreme values in certain channels of the network that can stretch the quantization range, making it difficult to accurately quantize the more common, non-outlier values. This oversight often results in significant performance degradation when trying to achieve very low bit-widths, such as 4-bit weights and 6-bit activations (W4A6).
Introducing DMQ: A New Approach to Outliers
A new research paper, titled DMQ: Dissecting Outliers of Diffusion Models for Post-Training Quantization, proposes an innovative solution to these challenges. Developed by Dongyeun Lee, Jiwan Hur, Hyounguk Shon, Jae Young Lee, and Junmo Kim from KAIST, DMQ is a novel post-training quantization framework specifically designed for diffusion models. It combines two key techniques: Learned Equivalent Scaling (LES) and channel-wise Power-of-Two Scaling (PTS), to ensure accurate quantization even under stringent low-bit constraints.
Learned Equivalent Scaling (LES)
The first component, Learned Equivalent Scaling (LES), is designed to optimize channel-wise scaling factors. Imagine you have a very wide range of numbers you need to fit into a smaller box. LES intelligently adjusts these numbers so they fit better, effectively redistributing the ‘difficulty’ of quantization between the network’s weights and activations. This process minimizes the overall quantization error. Unlike simpler methods that might use fixed scaling factors, LES learns these factors by minimizing the difference between the original and quantized outputs.
A crucial insight behind LES is the recognition that even small quantization errors in the early stages of the denoising process can have a significant, cumulative impact on the final image quality. To address this, DMQ incorporates an adaptive timestep weighting scheme. This scheme prioritizes these critical early steps during the learning process, ensuring that the model focuses on accuracy where it matters most for the final output.
For practical deployment, DMQ ensures that these learned scaling factors don’t add extra computational burden during inference. The scaling factors for activations are cleverly integrated into the existing quantization scale, while weight scaling factors are pre-computed and fused directly into the weights. This means DMQ can apply its scaling without any additional overhead during the actual image generation process.
Power-of-Two Scaling (PTS)
While LES is highly effective, some layers in diffusion models, particularly ‘skip connections’ which lack normalization, can exhibit extremely large outliers. These outliers pose a significant challenge that LES alone cannot fully resolve. This is where channel-wise Power-of-Two Scaling (PTS) comes in.
PTS directly tackles these extreme activation outliers by scaling them using factors that are powers of two (e.g., 2, 4, 8, 1/2, 1/4). The beauty of using power-of-two factors is that they can be implemented very efficiently on hardware using simple bit-shifting operations, rather than complex multiplications. This means the quantization difficulty is not just transferred but effectively removed with minimal computational cost.
To ensure robust selection of these PTS factors, especially with small calibration datasets, DMQ introduces a clever voting algorithm. Instead of simply picking the factor that minimizes error for a single sample, the algorithm considers multiple samples and selects the factor that has the most consensus across them. This conservative approach prevents overfitting to anomalies and ensures reliable scaling, leading to better generalization.
Also Read:
- POT-PTQ: Enhancing Large Language Model Efficiency with Two-Step Power-of-Two Quantization
- Discrete Latent Codes: A New Approach to High-Fidelity and Creative Image Generation
Impressive Results Across the Board
Extensive experiments demonstrate that DMQ consistently outperforms existing methods across various datasets and image generation tasks. Whether it’s unconditional image generation (like LSUN-Bedroom or FFHQ), class-conditional generation (ImageNet), or text-guided image generation (Stable Diffusion on MS-COCO), DMQ maintains high image generation quality and model stability, even at challenging low bit-widths like W4A6 (4-bit weights, 6-bit activations), where many previous methods struggle or fail significantly.
For instance, in unconditional generation, DMQ showed superior FID and sFID scores. In class-conditional generation, it not only achieved better FID and sFID but also excelled in metrics like LPIPS, SSIM, and PSNR, indicating that its generated images are perceptually and structurally closer to those produced by full-precision models. For text-guided generation, DMQ also achieved top scores in CLIP, LPIPS, SSIM, and PSNR, confirming strong semantic alignment with text prompts and high visual fidelity.
The ablation studies further confirm the effectiveness of each component: LES significantly improves performance, adaptive timestep weighting provides an additional boost by focusing on critical steps, and PTS, particularly when applied to specific layers like skip connections, yields the best overall results. The robust voting algorithm for PTS factors also proved crucial for stable performance on unseen data.
In conclusion, DMQ offers a powerful and efficient solution for quantizing diffusion models, making them more accessible for deployment in resource-constrained environments without sacrificing the remarkable quality of their generated outputs. By intelligently handling outliers and prioritizing critical denoising steps, DMQ pushes the boundaries of what’s possible with low-bit diffusion models.
