EDA-DM: Enhanced Distribution Alignment for Post-Training Quantization of Diffusion Models

TL;DR

This paper introduces EDA-DM, a post-training quantization method for diffusion models that effectively reduces model size and speeds up inference with minimal performance loss by addressing distribution mismatch issues.

Contribution

The paper proposes EDA-DM, a novel PTQ approach that aligns activation distributions at multiple levels, improving quantization of diffusion models without fine-tuning.

Findings

Achieves 1.83x speedup and 4x compression on Stable-Diffusion

Maintains high image generation quality with only 0.05 loss in CLIP score

Outperforms existing PTQ methods across various models and datasets

Abstract

Diffusion models have achieved great success in image generation tasks. However, the lengthy denoising process and complex neural networks hinder their low-latency applications in real-world scenarios. Quantization can effectively reduce model complexity, and post-training quantization (PTQ), which does not require fine-tuning, is highly promising for compressing and accelerating diffusion models. Unfortunately, we find that due to the highly dynamic activations, existing PTQ methods suffer from distribution mismatch issues at both calibration sample level and reconstruction output level, which makes the performance far from satisfactory. In this paper, we propose EDA-DM, a standardized PTQ method that efficiently addresses the above issues. Specifically, at the calibration sample level, we extract information from the density and diversity of latent space feature maps, which guides the selection of calibration samples to align with the overall sample distribution; and at the reconstruction output level, we theoretically analyze the reasons for previous reconstruction failures and, based on this insight, optimize block reconstruction using the Hessian loss of layers, aligning the outputs of quantized model and full-precision model at different network granularity. Extensive experiments demonstrate that EDA-DM significantly outperforms the existing PTQ methods across various models and datasets. Our method achieves a 1.83 times speedup and 4 times compression for the popular Stable-Diffusion on MS-COCO, with only a 0.05 loss in CLIP score. Code is available at http://github.com/BienLuky/EDA-DM .