Diffusion Product Quantization

TL;DR

This paper introduces a product quantization method for diffusion models that significantly reduces model size while maintaining performance, enabling extreme compression down to 1 bit.

Contribution

The authors develop a novel product quantization approach with importance-based codebook compression and end-to-end calibration for diffusion models.

Findings

Achieves over 24x model size reduction to 1 bit.

Outperforms existing quantization methods on ImageNet.

Maintains competitive generative quality after compression.

Abstract

In this work, we explore the quantization of diffusion models in extreme compression regimes to reduce model size while maintaining performance. We begin by investigating classical vector quantization but find that diffusion models are particularly susceptible to quantization error, with the codebook size limiting generation quality. To address this, we introduce product quantization, which offers improved reconstruction precision and larger capacity -- crucial for preserving the generative capabilities of diffusion models. Furthermore, we propose a method to compress the codebook by evaluating the importance of each vector and removing redundancy, ensuring the model size remaining within the desired range. We also introduce an end-to-end calibration approach that adjusts assignments during the forward pass and optimizes the codebook using the DDPM loss. By compressing the model to as low as 1 bit (resulting in over 24 times reduction in model size), we achieve a balance between compression and quality. We apply our compression method to the DiT model on ImageNet and consistently outperform other quantization approaches, demonstrating competitive generative performance.