TFMQ-DM: Temporal Feature Maintenance Quantization for Diffusion Models

The Diffusion model, a prevalent framework for image generation, encounters significant challenges in terms of broad applicability due to its extended inference times and substantial memory requirements. Efficient Post-training Quantization (PTQ) is pivotal for addressing these issues in traditional models. Different from traditional models, diffusion models heavily depend on the time-step t to achieve satisfactory multi-round denoising. Usually, t from the finite set {1, …, T} is encoded to a temporal feature by a few modules totally irrespective of the sampling data. However, existing PTQ methods do not optimize these modules separately. They adopt inappropriate reconstruction targets and complex calibration methods, resulting in a severe disturbance of the temporal feature and denoising trajectory, as well as a low compression efficiency. To solve these, we propose a Temporal Feature Maintenance Quantization (TFMQ) framework building upon a Temporal Information Block which is just related to the time-step t and unrelated to the sampling data. Powered by the pioneering block design, we devise temporal information aware reconstruction (TIAR) and finite set calibration (FSC) to align the full-precision temporal features in a limited time. Equipped with the framework, we can maintain the most temporal information and ensure the end-to-end generation quality. Extensive experiments on various datasets and diffusion models prove our state-of-the-art results. Remarkably, our quantization approach, for the first time, achieves model performance nearly on par with the full-precision model under 4-bit weight quantization. Additionally, our method incurs almost no extra computational cost and accelerates quan-tization time by 2.0× on LSUN-Bedrooms 256 × 256 compared to previous works. Our code is publicly available at https://github.com/ModelTC/TFMQ-DM.

1. Introduction

Generative modeling plays a crucial role in machine learning, particularly in applications like image [13], [14], [17], [38], [44], voice [37], [42], and text synthesis [2], [51]. Diffusion models have showcased impressive capabilities in producing high-quality samples across diverse domains. In comparison to generative adversarial networks (GANs) [8] and variational autoencoders (VAEs) [20], diffusion models successfully sidestep issues such as model collapse and pos-terior collapse, resulting in a more stable training process. However, the substantial computational cost poses a critical bottleneck hampering the widespread adoption of diffusion models. Furthermore, the computational cost for diffusion models can be attributed to two primary factors. First, these models typically require hundreds of denoising steps to generate images, rendering the procedure considerably slower than that of GANs. Prior efforts [21], [27], [29], [44] have addressed this challenge by seeking shorter and more efficient sampling trajectories, thereby reducing the number of denoising steps. Second, the substantial network architecture of diffusion models demands considerable time and memory resources, particularly for foundational models pretrained on large-scale datasets, e.g., LDM [38] and Stable Diffusion. Our work aims to tackle the latter challenge, focusing on the compression of diffusion models.