SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yefei He $^{1}$  Luping Liu $^{1}$  Jing Liu $^{2}$  Weijia Wu $^{1}$  Hong Zhou $^{1\dagger}$  Bohan Zhuang $^{2\dagger}$
<sup>1</sup>Zhejiang University, China  
<sup>2</sup>ZIP Lab, Monash University, Australia
Diffusion models have recently dominated image synthesis and other related generative tasks. However, the iterative denoising process is expensive in computations at inference time, making diffusion models less practical for low-latency and scalable real-world applications. Post-training quantization of diffusion models can significantly reduce the model size and accelerate the sampling process without requiring any re-training. Nonetheless, applying existing post-training quantization methods directly to low-bit diffusion models can significantly impair the quality of generated samples. Specifically, for each denoising step, quantization noise leads to deviations in the estimated mean and mismatches with the predetermined variance schedule. Moreover, as the sampling process proceeds, the quantization noise may accumulate, resulting in a low signal-to-noise ratio (SNR) during the later denoising steps. To address these challenges, we propose a unified formulation for the quantization noise and diffusion perturbed noise in the quantized denoising process. Specifically, we first disentangle the quantization noise into its correlated and residual uncorrelated parts regarding its full-precision counterpart. The correlated part can be easily corrected by estimating the correlation coefficient. For the uncorrelated part, we subtract the bias from the quantized results to correct the mean deviation and calibrate the denoising variance schedule to absorb the excess variance resulting from quantization. Moreover, we introduce a mixed-precision scheme for selecting the optimal bitwidth for each denoising step, which prioritizes lower bitwidths to expedite early denoising steps, while ensuring that higher bitwidths maintain a high signal-to-noise ratio (SNR) in the later steps. Extensive experiments demonstrate that our method outperforms previous post-training quantized diffusion models in generating high-quality samples, with only a 0.06 increase in FID score compared to full-precision LDM-4 on ImageNet  $256 \times 256$ , while saving  $19.9 \times$  bit operations. Code is available at https://github.com/ziplab/PTQD.
Diffusion models have demonstrated remarkable ability in generating high-quality samples in multiple fields [10, 5, 58, 18, 38, 30, 13, 52, 8, 49]. Compared to generative adversarial networks (GANs) [15] and variational autoencoders (VAEs) [27], diffusion models do not face the issue of mode collapse and posterior collapse, thus training is more stable. Nonetheless, the application of diffusion models is limited by two major bottlenecks. Firstly, diffusion models typically require hundreds of denoising steps to generate high-quality samples, making the process significantly slower than that of GANs. To address this, many studies [50, 36, 2, 28, 34] have proposed advanced training-free sampler to reduce the number of denoising iterations. Among them, a recent study DPM-solver [36] curtails the denoising process to ten steps by analytically computing the linear part of the diffusion ordinary
†Corresponding author. Email: zhouh@mail.bme.zju.edu.cn, bohan.zhuang@gmail.com
Q-Diffusion (W4A8)
PTQD (W4A8)
Figure 1: The comparisons of samples generated by Q-Diffusion [31], PTQD and full-precision LDM-4 [45] on CelebA-HQ  $256 \times 256$  dataset. Here, WxAy indicates the weights are quantized to  $x$ -bit while the activations are quantized to  $y$ -bit.
differential equations (ODEs). Nevertheless, diffusion models with these fast samplers are not yet ready for real-time applications. For instance, even when executed on a high-performance platform such as the RTX 3090, Stable Diffusion [45] with the DPM-Solver [36] sampler still takes over a second to generate a  $512 \times 512$  image. Second, the application of diffusion models on various devices is constrained by the massive parameters and computational complexity. To illustrate, executing Stable Diffusion [45] requires 16GB of running memory and GPUs with over 10GB of VRAM, which is infeasible for most consumer-grade PCs, not to mention resource-constrained edge devices.
Model quantization, which employs lower numerical bitwidth to represent weights and activations, has been widely studied to reduce memory footprint and computational complexity. For instance, employing 8-bit models can result in a significant speed-up of  $2.2 \times$  compared to floating-point models on ARM CPUs [21]. Adopting 4-bit quantization can further deliver a throughput increase of up to  $59\%$  compared to 8-bit quantization [3]. To facilitate the quantization process without the need for re-training, post-training quantization (PTQ) has emerged as a widely used technique, which is highly practical and easy to implement. While PTQ on traditional models have been widely studied [40, 32, 20, 33, 57], its application on diffusion models incurs two new challenges at the fundamental level. First, with the noise prediction network quantized, its quantization noise inevitably introduces bias in the estimated mean and brings additional variance that collides with the predetermined variance schedule in each denoising step. Additionally, the quantization noise accumulates as the iterative sampling process progresses, leading to a significant drop in the signal-to-noise ratio (SNR) of the noise prediction network in the later denoising steps. This diminished SNR severely impedes the denoising capability, resulting in a noticeable degradation in the quality of the generated images.
To tackle the aforementioned challenges, we present PTQD, a novel post-training quantization framework for diffusion models. To address the mean deviation and additional variance in each denoising step, we model the quantization noise by disentangling it into its correlated and residual uncorrelated parts regarding its full-precision counterpart, and designs separate correction methods for them. By estimating the correlation coefficient, the correlated part can be easily rectified. For the residual uncorrelated part, we subtract the bias from the estimated mean and propose variance schedule calibration, which absorbs the additional variance into the diffusion perturbed noise. To overcome the issue of low SNR that diminishes denoising capability in later denoising steps, we introduce a step-aware mixed precision scheme, which adaptively allocates different bitwidths for synonymous steps to maintain a high SNR for the denoising process.
- We present PTQD, a novel post-training quantization framework for diffusion models, which provides a unified formulation for quantization noise and diffusion perturbed noise.
- We disentangle the quantization noise into correlated and uncorrelated parts regarding its full-precision counterpart. Then we correct the correlated part by estimating the correlation coefficient, and propose variance schedule calibration to rectify the residual uncorrelated part.
- We introduce a step-aware mixed precision scheme, which dynamically selects the appropriate bitwidths for synonymous steps, preserving SNR throughout the denoising process.
- Our extensive experiments demonstrate that our method reaches a new state-of-the-art performance for post-training quantization of diffusion models.
Efficient diffusion models. While diffusion models can produce high-quality samples, their slow generation speed hinders their large-scale applications in downstream tasks. To explore efficient diffusion models, many methods have been proposed to expedite the sampling process. These methods can be classified into two categories: methods that necessitate re-training and advanced samplers for pre-trained models that do not require training. The first category of methods comprises knowledge distillation [37, 47], diffusion scheme learning [7, 12, 64, 39], noise scale learning [26, 43], and sample trajectory learning [55, 29]. Although these methods can accelerate sampling, re-training a diffusion model can be resource-intensive and time-consuming. On the other hand, the second category of methods designs advanced samplers directly on pre-trained diffusion models, eliminating the need for re-training. The primary methods in this category are implicit sampler [50, 28, 63, 56], analytical trajectory estimation [2, 1], and differential equation (DE) solvers such as customized SDE [51, 23, 25] and ODE [36, 34, 62]. Although these methods can reduce the sampling iterations, the diffusion model's massive parameters and computational complexity restrict their use to high-performance platforms. Conversely, our proposed low-bit diffusion model can significantly reduce the model's computational complexity while speeding up the sampling and reducing the demand for hardware computing resources in a training-free manner.
Model quantization. Quantization is a dominant technique to save memory costs and speed up computation. It can be divided into two categories: quantization-aware training (QAT) [14, 35, 22, 65, 61] and post-training quantization (PTQ) [32, 40, 20, 57, 33]. QAT involves simulating quantization during training to achieve good performance with lower precision, but it requires substantial time, computational resources, and access to the original dataset. In contrast, PTQ does not require fine-tuning and only needs a small amount of unlabeled data to calibrate. Recent studies have pushed the limits of PTQ to 4-bit on traditional models by using new rounding strategies [40], layer-wise calibration [20, 53], and second-order statistics [32, 57]. Additionally, mixed precision (MP) [19, 11, 4, 59, 6] allows a part of the model to be represented by lower bitwidths to accelerate inference. Common criteria for determining quantization bitwidths include Hessian spectrum [11, 6] or Pareto frontier [4]. In contrast, we propose a novel mixed-precision scheme for diffusion models that adapts different bitwidths for synonymous denoising steps.
Until now, there have been few studies specifically focusing on quantizing a pre-trained diffusion model without re-training. PTQ4DM [48] is the first attempt to quantize diffusion models to 8-bit, but its experiments are limited to small datasets and low resolution. Q-Diffusion [31] applies advanced PTQ techniques proposed by BRECQ [32] to improve performance and evaluate it on a wider range of datasets. Our paper aims to analyze systematically the quantization effect on diffusion models and establish a unified framework for accurate post-training diffusion quantization.
Diffusion models [50, 17] gradually apply Gaussian noise to real data  $\mathbf{x}_0$  in the forward process and learn a reverse process to denoise and generate high-quality images. For DDPMs [17], the forward process is a Markov chain, which can be formulated as:
$$
q \left(\mathbf {x} _ {1: T} \mid \mathbf {x} _ {0}\right) = \prod_ {t = 1} ^ {T} q \left(\mathbf {x} _ {t} \mid \mathbf {x} _ {t - 1}\right), \quad q \left(\mathbf {x} _ {t} \mid \mathbf {x} _ {t - 1}\right) = \mathcal {N} \left(\mathbf {x} _ {t}; \sqrt {\alpha_ {t}} \mathbf {x} _ {t - 1}, \beta_ {t} \mathbf {I}\right) \tag {1}
$$
where  $\alpha_{t},\beta_{t}$  are hyperparameters and  $\beta_{t} = 1 - \alpha_{t}$
In the reverse process, since directly estimating the real distribution of  $q(\mathbf{x}_{t - 1}|\mathbf{x}_t)$  is intractable, diffusion models approximate it via variational inference by learning a Gaussian distribution  $p_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t - 1};\pmb {\mu}_\theta (\mathbf{x}_t,t),\pmb {\Sigma}_\theta (\mathbf{x}_t,t))$  and reparameterize its mean by a noise prediction network  $\epsilon_{\theta}(\mathbf{x}_t,t)$ :
$$
\boldsymbol {\mu} _ {\theta} \left(\mathbf {x} _ {t}, t\right) = \frac {1}{\sqrt {\alpha_ {t}}} \left(\mathbf {x} _ {t} - \frac {\beta_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \boldsymbol {\epsilon} _ {\theta} \left(\mathbf {x} _ {t}, t\right)\right) \tag {2}
$$
where  $\bar{\alpha}_t = \prod_{s=1}^t \alpha_s$ . The variance  $\Sigma_\theta(\mathbf{x}_t, t)$  can either be reparameterized or fixed to a constant schedule  $\sigma_t$ . When it uses a constant schedule, the sampling of  $\mathbf{x}_{t-1}$  can be formulated as:
$$
\mathbf {x} _ {t - 1} = \frac {1}{\sqrt {\alpha_ {t}}} \left(\mathbf {x} _ {t} - \frac {\beta_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \boldsymbol {\epsilon} _ {\theta} (\mathbf {x} _ {t}, t)\right) + \sigma_ {t} \mathbf {z}, \text {w h e r e} \mathbf {z} \sim \mathcal {N} (\mathbf {0}, \mathbf {I}). \tag {3}
$$
Our method focuses on post-training quantization of diffusion models without the need of training. Instead, we use pre-trained diffusion models and inherit their hyperparameters and variance schedules for inference. Although the derivations presented in this paper are based on DDPM, they can be readily extended to other fast sampling methods, such as DDIM [50]. Additional information can be found in the supplementary material.
3.2 Model Quantization
We use uniform quantization in our study and all the experiments. For uniform quantization, given a floating-point vector  $\mathbf{x}$ , the target bitwidth  $b$ , the quantization process can be defined as:
$$
\hat {\mathbf {x}} = \Delta \cdot \left(\operatorname {c l i p} \left(\left\lfloor \frac {\mathbf {x}}{\Delta} \right\rceil + Z, 0, 2 ^ {b} - 1\right) - Z\right), \tag {4}
$$
where  $\lfloor \cdot \rfloor$  is the round operation,  $\Delta = \frac{\max(\mathbf{x}) - \min(\mathbf{x})}{2^b - 1}$  and  $Z = -\lfloor \frac{\min(\mathbf{x})}{\Delta} \rfloor$ .
To ensure clarity and consistency, we introduce notation to define the variables used in the paper. Let  $X$  be a tensor (weights or activations) in the full-precision model, the result after normalization layers is denoted as  $\overline{X}$ . The corresponding tensor of the quantized model is represented as  $\hat{X}$ . The quantization noise is depicted by  $\Delta_{X}$ , which is the difference between  $\hat{X}$  and  $X$ .
Model quantization discretizes the weights and activations, which will inevitably introduce quantization noise into the result. As per Eq. (3), during the reverse process of the quantized diffusion model, the sampling of  $\mathbf{x}_{t-1}$  can be expressed as:
$$
\begin{array}{l} \mathbf {x} _ {t - 1} = \frac {1}{\sqrt {\alpha_ {t}}} \left(\mathbf {x} _ {t} - \frac {\beta_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \hat {\epsilon} _ {\theta} \left(\mathbf {x} _ {t}, t\right)\right) + \sigma_ {t} \mathbf {z} \tag {5} \\ = \frac {1}{\sqrt {\alpha_ {t}}} \left(\mathbf {x} _ {t} - \frac {\beta_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \left(\boldsymbol {\epsilon} _ {\theta} (\mathbf {x} _ {t}, t) + \Delta_ {\boldsymbol {\epsilon} _ {\theta} (\mathbf {x} _ {t}, t)}\right)\right) + \sigma_ {t} \mathbf {z}. \\ \end{array}
$$
Here,  $\hat{\epsilon}_{\theta}(\mathbf{x}_t,t)$  is the output of the quantized noise prediction network and  $\Delta_{\epsilon_{\theta}(\mathbf{x}_t,t)}$  refers to the quantization noise. The additional quantization noise will inevitably alter the mean and variance of  $\mathbf{x}_{t - 1}$ , decreasing the signal-to-noise ratio (SNR) and adversely affecting the quality of the generated samples. Therefore, to mitigate the impact of quantization, it is necessary to correct the mean and variance to restore the SNR at each step of the reverse process.
4.1 Correlation Disentanglement
We begin by making an assumption that a correlation exists between the quantization noise and the result of the full-precision noise prediction network. While other factors, such as nonlinear operations, may contribute to this correlation, Proposition 1 demonstrates that normalization layers are responsible for a part of it.
Proposition 1. Given  $Y$  and  $\hat{Y}$  as inputs to a normalization layer in a full-precision model and its quantized version, where the quantization noise  $\Delta_{Y} = \hat{Y} - Y$  is initially uncorrelated with  $Y$ , a correlation between the quantization noise and the output of the full-precision model after the normalization layer will exist.
The proof is based on the fact that the mean and variance of  $\hat{Y}$  will differ from that of  $Y$  (depending on the specific quantization scheme). Therefore, the quantization noise after normalization layer can be expressed as:
$$
\Delta_ {\bar {Y}} = \frac {\hat {Y} - \mu_ {\hat {Y}}}{\sigma_ {\hat {Y}}} - \frac {Y - \mu_ {Y}}{\sigma_ {Y}} = \frac {\sigma_ {Y} \Delta_ {Y} - (\sigma_ {\hat {Y}} - \sigma_ {Y}) Y + \sigma_ {\hat {Y}} \mu_ {Y} - \sigma_ {Y} \mu_ {\hat {Y}}}{\sigma_ {\hat {Y}} \sigma_ {Y}}. \tag {6}
$$
Figure 2: The correlation between the quantization noise (Y-axis) and the output of the full-precision noise prediction network (X-axis). Each data point on the plot corresponds to specific entries within these vectors. Data were collected by generating samples with 4-bit LDM-8 [45] for 200 steps on LSUN-Churches [60].
Here, we omit the affine transform parameters in normalization layers for simplicity. It can be observed from Eq. (6) that the second term in the numerator is related to  $Y$ , while the other three terms are uncorrelated. Therefore, after normalization layers, the quantization noise  $\Delta_{\overline{Y}}$  will be correlated with  $Y$ .
The empirical observation illustrated in Figure 2 confirms a strong correlation between the quantization noise and the output of the full-precision noise prediction network, which further verifies our assumption. Based on the assumption and observation, the quantization noise of the quantized noise prediction network can be disentangled into two parts:
$$
\Delta_ {\epsilon_ {\theta} \left(\mathbf {x} _ {t}, t\right)} = k \epsilon_ {\theta} \left(\mathbf {x} _ {t}, t\right) + \Delta_ {\epsilon_ {\theta} \left(\mathbf {x} _ {t}, t\right)} ^ {\prime}. \tag {7}
$$
The first part, denoted by  $k\epsilon_{\theta}(\mathbf{x}_t,t)$ , is linearly related to  $\epsilon_{\theta}(\mathbf{x}_t,t)$ . The second part, expressed by  $\Delta_{\epsilon_{\theta}(\mathbf{x}_t,t)}'$ , represents the residual component of the quantization noise, and is assumed to be uncorrelated with  $\epsilon_{\theta}(\mathbf{x}_t,t)$ . Here,  $k$  is the correlation coefficient, which can be estimated by applying linear regression on the quantization noise  $\Delta_{\epsilon_{\theta}(\mathbf{x}_t,t)}$  and the original value  $\epsilon_{\theta}(\mathbf{x}_t,t)$ . Details can be found in Section 5.1.
With the disentanglement presented in Eq. (7), the sampling of  $\mathbf{x}_{t - 1}$  can be further expressed as:
$$
\begin{array}{l} \mathbf {x} _ {t - 1} = \frac {1}{\sqrt {\alpha_ {t}}} \left(\mathbf {x} _ {t} - \frac {\beta_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \left(\epsilon_ {\theta} \left(\mathbf {x} _ {t}, t\right) + \Delta_ {\epsilon_ {\theta} \left(\mathbf {x} _ {t}, t\right)}\right)\right) + \sigma_ {t} \mathbf {z} \tag {8} \\ = \frac {1}{\sqrt {\alpha_ {t}}} \left(\mathbf {x} _ {t} - \frac {\beta_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \left((1 + k) \boldsymbol {\epsilon} _ {\theta} (\mathbf {x} _ {t}, t) + \Delta_ {\boldsymbol {\epsilon} _ {\theta} (\mathbf {x} _ {t}, t)} ^ {'}\right)\right) + \sigma_ {t} \mathbf {z}. \\ \end{array}
$$
Consequently, the bias and additional variance arise from both the correlated and uncorrelated parts of quantization noise. In the following section, we will provide a detailed explanation of how these two parts of quantization noise can be separately corrected.
4.2 Quantization Noise Correction
4.2.1 Correlated Noise Correction
Based on Eq. (8), the correlated part of the quantization noise can be rectified by dividing the output of the quantized noise prediction network  $\hat{\epsilon}_{\theta}(\mathbf{x}_t,t)$  by  $1 + k$ , resulting in:
$$
\mathbf {x} _ {t - 1} = \frac {1}{\sqrt {\alpha_ {t}}} \left(\mathbf {x} _ {t} - \frac {\beta_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \boldsymbol {\epsilon} _ {\theta (\mathbf {x} _ {t}, t)}\right) + \sigma_ {t} \mathbf {z} - \frac {\beta_ {t}}{\sqrt {\alpha_ {t}} \sqrt {1 - \bar {\alpha} _ {t}} (1 + k)} \Delta_ {\epsilon_ {\theta} (\mathbf {x} _ {t}, t)} ^ {\prime}. \tag {9}
$$
Consequently, only the uncorrelated quantization noise remains. Moreover, for values of  $k \geq 0$ , it can be deduced that the mean and variance of the uncorrelated quantization noise are diminished by  $\frac{1}{1 + k}$ . In practice, we enforce the non-negativity of  $k$ , and reset it to zero if it is negative. In the following, we will explain how to handle the uncorrelated quantization noise that persists in Eq. (9).
4.2.2 Uncorrelated Noise Correction
The presence of uncorrelated quantization noise introduces additional variance at each step, resulting in a total variance that exceeds the scheduled value  $\sigma_t^2$ . To address this, we propose to calibrate the variance schedule for quantized diffusion models, which is denoted as  $\sigma_t'^2$  and smaller than the original schedule  $\sigma_t^2$ . To estimate  $\sigma_t'^2$ , we further make an assumption and model the uncorrelated quantization noise as a Gaussian distribution with a mean of  $\mu_q$  and a variance of  $\sigma_q^2$ :
$$
\Delta_ {\epsilon_ {\theta} \left(\mathbf {x} _ {t}, t\right)} ^ {\prime} \sim \mathcal {N} \left(\mu_ {q}, \sigma_ {q}\right). \tag {10}
$$
To verify this assumption, we conduct statistical tests (refer to the supplementary material) and present the distribution of the uncorrelated quantization noise in Figure 3. The values of mean and variance can be estimated by generating samples with both quantized and full-precision diffusion models and collecting the statistics of the uncorrelated quantization noise. Following prior work [41], the mean deviation can be rectified through Bias Correction (BC), where we collect the channel-wise means of uncorrelated quantization noise and subtract them from the output of quantized noise prediction network. For the variance of the uncorrelated quantization noise, we propose Variance Schedule Calibration (VSC), where the uncorrelated quantization noise can be absorbed into Gaussian diffusion noise with the above assumption. By substituting the calibrated variance schedule  $\sigma_t^{2}$  into Eq. (9) while keeping the variance of each step unaltered, we can solve for the optimal variance schedule using the following approach:
$$
\sigma_ {t} ^ {\prime 2} + \frac {\beta_ {t} ^ {2}}{\alpha_ {t} (1 - \bar {\alpha} _ {t}) (1 + k) ^ {2}} \sigma_ {q} ^ {2} = \sigma_ {t} ^ {2}, \tag {11}
$$
$$
\sigma_ {t} ^ {\prime 2} = \left\{ \begin{array}{l l} \sigma_ {t} ^ {2} - \frac {\beta_ {t} ^ {2}}{\alpha_ {t} (1 - \bar {\alpha} _ {t}) (1 + k) ^ {2}} \sigma_ {q} ^ {2}, & \text {i f} \sigma_ {t} ^ {2} \geq \frac {\beta_ {t} ^ {2}}{\alpha_ {t} (1 - \bar {\alpha} _ {t}) (1 + k) ^ {2}} \sigma_ {q} ^ {2} \\ 0, & \text {o t h e r w i s e .} \end{array} \right. \tag {12}
$$
It can be observed that if the additional variance of quantization noise is smaller than the noise hyperparameter  $\sigma_t^2$ , the increase in variance caused by quantization can be eliminated. According to Eq. (12), the coefficient for the variance of the quantization noise can be calculated as  $\frac{\beta_t^2}{\alpha_t(1 - \bar{\alpha}_t)(1 + k)^2}$  which is generally small enough to ensure that the quantization noise can be fully absorbed, except for cases of deterministic sampling where  $\sigma_t$  is zero. In this case, there is no analytical solution for  $\sigma_t'^2$ , and we use the optimal solution that is  $\sigma_t'^2 = 0$ . Overall, the process quantization noise correction is summarized in Algorithm 1.
Algorithm 1: Quantization noise correction.
Statistics collection before sampling:
1) Quantize diffusion models with BRECQ [32] (or other PTQ methods);
2) Generate samples with both quantized and FP models and collect quantization noise;
3) Calculate the correlated coefficient  $k$  based on Eq. (7), and the mean and variance of the uncorrelated quantization noise as per Eq. (10);
Noise correction for each sampling step: