[Paper Review] EfficientDM: Efficient Quantization-Aware Fine-Tuning of Low-Bit Diffusion Models

Yefei He, Jing Liu+3 more|arXiv (Cornell University)|Oct 5, 2023

Advanced Data Compression Techniques4 citations

TL;DR

EfficientDM proposes a data-free, quantization-aware, and parameter-efficient fine-tuning framework for low-bit diffusion models that achieves QAT-level performance with PTQ-level efficiency. By introducing a quantization-aware LoRA variant (QALoRA), scale-aware optimization, and temporal learned step-size quantization, it enables joint quantization of adapter and model weights, reducing sFID by 1.56 points on ImageNet 256×256 at 4-bit precision while cutting quantization time by 16.2× compared to QAT.

ABSTRACT

Diffusion models have demonstrated remarkable capabilities in image synthesis and related generative tasks. Nevertheless, their practicality for real-world applications is constrained by substantial computational costs and latency issues. Quantization is a dominant way to compress and accelerate diffusion models, where post-training quantization (PTQ) and quantization-aware training (QAT) are two main approaches, each bearing its own properties. While PTQ exhibits efficiency in terms of both time and data usage, it may lead to diminished performance in low bit-width. On the other hand, QAT can alleviate performance degradation but comes with substantial demands on computational and data resources. In this paper, we introduce a data-free and parameter-efficient fine-tuning framework for low-bit diffusion models, dubbed EfficientDM, to achieve QAT-level performance with PTQ-like efficiency. Specifically, we propose a quantization-aware variant of the low-rank adapter (QALoRA) that can be merged with model weights and jointly quantized to low bit-width. The fine-tuning process distills the denoising capabilities of the full-precision model into its quantized counterpart, eliminating the requirement for training data. We also introduce scale-aware optimization and temporal learned step-size quantization to further enhance performance. Extensive experimental results demonstrate that our method significantly outperforms previous PTQ-based diffusion models while maintaining similar time and data efficiency. Specifically, there is only a 0.05 sFID increase when quantizing both weights and activations of LDM-4 to 4-bit on ImageNet 256x256. Compared to QAT-based methods, our EfficientDM also boasts a 16.2x faster quantization speed with comparable generation quality. Code is available at \href{https://github.com/ThisisBillhe/EfficientDM}{this hrl}.

Motivation & Objective

• To address the efficiency-quality trade-off in low-bit quantization of diffusion models, where PTQ lacks performance and QAT demands excessive data and compute.
• To eliminate dependency on original training data during quantization, enabling practical deployment on resource-constrained devices.
• To maintain high generation quality at 4-bit precision for both weights and activations, where prior PTQ methods fail to preserve denoising capability.
• To achieve QAT-level performance with the time and data efficiency of PTQ, making low-bit diffusion models viable for real-world applications.

Proposed method

• Introduces QALoRA, a quantization-aware variant of low-rank adapters that can be jointly quantized with model weights, enabling efficient and low-bit inference.
• Employs a data-free distillation framework that minimizes mean squared error between noise predictions of full-precision and quantized models, eliminating the need for training data.
• Applies scale-aware optimization to adaptively adjust gradient scales across layers, improving optimization effectiveness in low-bit settings with varying quantization scales.
• Introduces temporal learned step-size quantization (TLSQ) to handle distribution shifts in activations across denoising steps, improving robustness during inference.
• Merges QALoRA weights with the quantized model, enabling parameter-efficient fine-tuning with no additional storage or computation overhead.
• Uses a two-stage training process: first distilling full-precision model knowledge into the quantized model via QALoRA, then refining with TLSQ and scale-aware optimization.

Experimental results

Research questions

• RQ1 Can a data-free fine-tuning method achieve QAT-level performance in low-bit diffusion model quantization without requiring access to the original training dataset?
• RQ2 How can quantization-aware low-rank adapters be designed to jointly quantize with model weights while preserving performance at 4-bit precision?
• RQ3 What optimization strategies are effective in mitigating performance degradation caused by varying quantization scales across network layers in low-bit settings?
• RQ4 Can temporal adaptation of quantization parameters across denoising steps improve model robustness and generation quality in low-bit diffusion models?
• RQ5 Is it possible to achieve near-full-precision generation quality at 4-bit precision with significantly reduced time and memory costs compared to QAT?

Key findings

• EfficientDM achieves an sFID of 7.75 on ImageNet 256×256 at 4-bit precision (W/A), with only a 0.05 increase over the full-precision baseline, outperforming prior PTQ methods that fail to denoise at this bit-width.
• The method reduces quantization time by 16.2× compared to QAT-based approaches, achieving 3.05 GPU hours on ImageNet 256×256 versus 54.5 GPU hours for QAT.
• On CIFAR-10, EfficientDM achieves a FID of 10.48 with 0.97 GPU hours, matching QAT performance (FID 7.30) while being 16.8× faster than QAT and 1.1× faster than PTQ.
• The scale-aware optimization reduces FID from 24.18 (QALoRA only) to 13.23, demonstrating its effectiveness in improving optimization across layers with varying quantization scales.
• TLSQ further reduces FID from 13.23 to 7.75, showing that step-specific quantization improves performance by adapting to activation distribution shifts across denoising steps.
• EfficientDM uses only 19.8GB GPU memory on ImageNet 256×256, significantly less than QAT’s OOM on RTX 3090, and avoids data dependency, making it practical for real-world deployment.