Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Gradient-Aligned Calibration for Post-Training Quantization of Diffusion Models

Dung Anh Hoang, Cuong Pham anh Trung Le, Jianfei Cai, Toan Do

TL;DR

The paper addresses the gradient-conflict problem in post-training quantization of diffusion models by learning timestep-aware calibration weights that align gradient directions across denoising timesteps. It introduces a meta-learning framework with a bi-level objective and a surrogate gradient-matching loss to optimize per-sample weights, backed by theoretical guarantees. Empirically, the method achieves state-of-the-art FID and sFID across CIFAR-10, LSUN-Bedrooms, and ImageNet under aggressive quantization, with robust performance across various timesteps and validation-set sizes. The approach offers practical gains for deploying diffusion models on resource-constrained devices without retraining, preserving quality while reducing memory and compute during inference.

Abstract

Diffusion models have shown remarkable performance in image synthesis by progressively estimating a smooth transition from a Gaussian distribution of noise to a real image. Unfortunately, their practical deployment is limited by slow inference speed, high memory usage, and the computational demands of the noise estimation process. Post-training quantization (PTQ) emerges as a promising solution to accelerate sampling and reduce memory overhead for diffusion models. Existing PTQ methods for diffusion models typically apply uniform weights to calibration samples across timesteps, which is sub-optimal since data at different timesteps may contribute differently to the diffusion process. Additionally, due to varying activation distributions and gradients across timesteps, a uniform quantization approach is sub-optimal. Each timestep requires a different gradient direction for optimal quantization, and treating them equally can lead to conflicting gradients that degrade performance. In this paper, we propose a novel PTQ method that addresses these challenges by assigning appropriate weights to calibration samples. Specifically, our approach learns to assign optimal weights to calibration samples to align the quantized model's gradients across timesteps, facilitating the quantization process. Extensive experiments on CIFAR-10, LSUN-Bedrooms, and ImageNet demonstrate the superiority of our method compared to other PTQ methods for diffusion models.

Gradient-Aligned Calibration for Post-Training Quantization of Diffusion Models

TL;DR

The paper addresses the gradient-conflict problem in post-training quantization of diffusion models by learning timestep-aware calibration weights that align gradient directions across denoising timesteps. It introduces a meta-learning framework with a bi-level objective and a surrogate gradient-matching loss to optimize per-sample weights, backed by theoretical guarantees. Empirically, the method achieves state-of-the-art FID and sFID across CIFAR-10, LSUN-Bedrooms, and ImageNet under aggressive quantization, with robust performance across various timesteps and validation-set sizes. The approach offers practical gains for deploying diffusion models on resource-constrained devices without retraining, preserving quality while reducing memory and compute during inference.

Abstract

Diffusion models have shown remarkable performance in image synthesis by progressively estimating a smooth transition from a Gaussian distribution of noise to a real image. Unfortunately, their practical deployment is limited by slow inference speed, high memory usage, and the computational demands of the noise estimation process. Post-training quantization (PTQ) emerges as a promising solution to accelerate sampling and reduce memory overhead for diffusion models. Existing PTQ methods for diffusion models typically apply uniform weights to calibration samples across timesteps, which is sub-optimal since data at different timesteps may contribute differently to the diffusion process. Additionally, due to varying activation distributions and gradients across timesteps, a uniform quantization approach is sub-optimal. Each timestep requires a different gradient direction for optimal quantization, and treating them equally can lead to conflicting gradients that degrade performance. In this paper, we propose a novel PTQ method that addresses these challenges by assigning appropriate weights to calibration samples. Specifically, our approach learns to assign optimal weights to calibration samples to align the quantized model's gradients across timesteps, facilitating the quantization process. Extensive experiments on CIFAR-10, LSUN-Bedrooms, and ImageNet demonstrate the superiority of our method compared to other PTQ methods for diffusion models.
Paper Structure (28 sections, 5 theorems, 25 equations, 4 figures, 8 tables, 2 algorithms)

This paper contains 28 sections, 5 theorems, 25 equations, 4 figures, 8 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related works
  3. Post-training quantization on diffusion Models.
  4. Data optimization for diffusion model quantization.
  5. Preliminary analysis
  6. Proposed method
  7. Problem definition
  8. Calibration data optimization
  9. Regarding the loss $\mathcal{L}_{MSE}$ in Eq. (\ref{['eq:overall_optimization']}).
  10. Overall optimization framework
  11. Experiments
  12. Experimental setup
  13. Comparison with the state-of-the-art quantization methods for diffusion models
  14. Ablation studies
  15. Ablation studies for the temperature parameter $\tau$.
  16. ...and 13 more sections

Key Result

Theorem 4.1

The optimization in Algorithm alg:data_optimization implicitly lead to the minimization of the target objective $\mathcal{L}_{\text{VAL}}(\cdot)$ in Eq. (eq:overall_optimization).

Figures (4)

  • Figure 1: Timestep-wise behavior in quantized diffusion models. (a) Gradient dissimilarity matrix constructed by computing pairwise cosine distance between gradient vectors of the quantization loss, with respect to model parameters, for calibration samples drawn from different timesteps. Higher value indicates higher divergent between timesteps. (b) Quantization loss evaluated separately for calibration samples grouped by timestep, highlighting the uneven performance of the quantized model across different timesteps.
  • Figure 2: Visualization of the correlation between optimized sample weights and gradient alignments. All samples are sorted in descending order ot their sample weights, and divided uniformly into 50 groups. The blue line represents the average sample weight per group, while the red line indicates the average gradient alignment between samples in each group and the validation set. This demonstrates the positive correlation between gradient alignment and sample weight.
  • Figure 3: Generated samples from (a) full-precision LDM-4, (b) Q-Diffusion (W4A32), (c) TFMQ-DM (W4A32), and (d) our proposed method (W4A32) on LSUN-Bedrooms $256 \times 256$ dataset with a fixed random seed.
  • Figure 4: Visualization of loss differences across timestep groups on the CIFAR-10 calibration set (4/32 setting) for quantized models trained with and without our sample-weighting method. Data are grouped by timesteps into five categories during sample weight optimization and are shown in ascending order of loss. Orange bars indicate loss reduction with our method, while green bars indicate an increase. The results demonstrate that our approach effectively reduces loss for under-optimized timesteps, addressing the gradient conflict issue that leads to the neglect of certain timestep

Theorems & Definitions (7)

  • Theorem 4.1
  • Lemma 4.2
  • Lemma 4.3
  • Lemma A.1: Restated from Lemma 4.3
  • proof : Proof of Lemma 1.1
  • Lemma A.2: Restated from Lemma 4.2
  • proof : Proof of Lemma 1.2