QNCD: Quantization Noise Correction for Diffusion Models

We consider the quantization noise within the denoising network UNet in a single sampling step as intra quantization noise, which is strongly affected by embedding. As shown in Fig. 2(a), intra quantization noise of the diffusion model exhibits periodic changes during a single step. It escalates when embedding is incorporated into features but then decreases during the fusion of UNet’s low-level and high-level features. This cyclical behavior underscores that the primary culprit of quantization noise is the embedding integration phase. Take the embedding fusion stage of DDIM as an example:

These outliers are several magnitudes larger than the majority of the data, leading to a skew in the maximum magnitude measurement during quantization. This dominance by outliers could possibly result in reduced precision for the majority of non-outlier values. Further complications arise with the incorporation of embedding, as shown in the middle part of Fig. 2(b). Embedding separates a $1\!\times c$ dimensional scale vector $scale_{t}$ which scales the feature $h_{t}$ on a channel-by-channel basis. This channel-specific scaling alters the distribution of the activation $h_{t}$ in a manner that certain channels, especially those with problematic outliers, experience amplification. Since activations are typically per-tensor quantized, the combined effect of embedding magnification and existing outliers makes the quantization of activations less efficient.

In summary, after normalization, feature $h_{t}$ is easily quantifiable, but with the introduction of $emb_{t}$, it becomes challenging to quantify , indicating an increase in intra quantization noise.

Diffusion models attain their final outputs through iterative denoising network calls. During this procedure, the inter quantization noise, expressed as $q_{\theta}(\widetilde{x}_{t},t)$ in Eq. 4, assimilates into $x_{t-1}$ and advances to the subsequent denoising step, exerting an influence over the entire sampling process.

As shown in Fig. 3, we plot the variation curves of the output Mean and Std during sampling steps. The accumulated inter quantization noise changes the distribution of synthesized data. With continuous sampling, the data distribution of the quantization model further deviates from that of the full-precision model.

As shown in Fig.2, during the single-step sampling process, embedding amplifies outliers of activation, leading to an imbalance among channels. Ultimately, this induces a periodic increase in intra quantization noise. For reducing intra quantization noise, we propose the utilization of a channel-specific smoothing factor $S$. By dividing activation with their respective $S$ values, channels are balanced out and more adaptable to quantization. We then incorporate the filtered factor into weights, thus maintaining mathematical equivalence of the convolution, as follows:

Ultimately, we can transfer the quantization challenges presented by embedding from activations to weights, which are more robust to quantization.

Since embedding operates on a channel-by-channel basis, our goal is to derive a factor $S$ for each channel from the embedding, making $h_{t}=h_{t}/S$ easier to quantize. As evident from Eq. 5, the term $1+scale_{t}$, derived from the separation from embedding, accentuates the discrepancies among the activation channels. However, $1+scale_{t}$ is dynamic and fluctuates based on the time step $t$. Consequently, we examine the embedding across all $t$ scenarios to ascertain the mean value of $1+scale_{t}$ and employ it as a static factor $S$:

As shown in Fig. 2(b), The static factor $S$ we obtained serves to calibrate these unbalanced channels, harmonizing the distribution across each one, rendering the eventual activations more conducive to per-tensor quantization. Different $scale_{t}$ and $S$ are highly similar, with only minor amplitude differences present on some channels.

In addition, in LDM-type diffusion models, embedding is incorporated differently than in Eq. 5:

where $\alpha$ and $\beta$ are pretrained affine transform parameters in the Group-Norm operation. At this point, the distribution of the final activation $h_{t}$ is jointly determined by $emb_{t}$ and the coefficients $\alpha$ of group norms. Thus, our smoothing factor $S$ is calculated as follows:

In Fig. 4, we present a visualization of $scale_{t}$ within $emb_{t}$ across various modules, time steps, and categories. The key observation is that the distribution of $scale_{t}$ exhibits imbalance, unevenly scaling the input activations by channel, which renders the activations less suitable for quantization. Despite this, the aggregate distribution of different $scalet_{t}$ remains relatively stable, with only minor amplitude discrepancies across certain channels.

Our smoothing factor $S$ is thus capable of effectively representing the $scalet_{t}$ encountered during the sampling phase. By filtering through $S$, our QNCD method mitigates the enhancing impact of $salet_{t}$ on activation outliers, leading to a reduction in intra quantization noise. Crucially, $S$ is pre-computed prior to the formal inference stage for each Resblock, ensuring no additional computational load during the actual inference.

Eq. 4 shows the process of a single-step sampling in diffusion models, where UNet outputs the filtered noise ${\widetilde{\epsilon}}_{\theta}(\widetilde{x}_{t},t)$. It consists of two parts, the de-noising noise ${\epsilon}_{\theta}(\widetilde{x}_{t},t)$, and inter quantization noise $q_{\theta}(\widetilde{x}_{t},t)$, which keeps accumulating in $\widetilde{x}_{t}$. Ideally, we filter out $q_{\theta}(\widetilde{x}_{t},t)$, thus avoiding the accumulation of quantization noise, but it is impractical to separate $q_{\theta}(\widetilde{x}_{t},t)$ from ${\widetilde{\epsilon}}_{\theta}(\widetilde{x}_{t},t)$.

The training stage of the diffusion model gives us a possibility to separate $q_{\theta}(\widetilde{x}_{t},t)$:

Eq. 10 does a single-step diffusion operation, where a Gaussian noise $z_{1}$ is added to $x_{t-1}$ to get $x_{t}$.

The training process of the diffusion model (Eq. 11) drives the denoising network $\epsilon_{\theta}$ to learn the distribution of Gaussian noise $z_{1}$, which means $\epsilon_{\theta}(x_{t},t)\approx z_{1}$ is satisfied in the pre-trained diffusion model.

This property is still guaranteed in the quantized pre-trained diffusion model:

As shown in Fig. 5 and Eq. 12, we add the standard Gaussian noise $z_{1}$ to $\widetilde{x}_{t-1}$ to get $\hat{x}_{t}$, and feed it into the quantized model $\widetilde{\epsilon}_{\theta}$. The output of the quantized model contains both the Gaussian noise $z_{1}$ to be filtered out and the newly introduced quantization noise $q_{\theta}(\hat{x}_{t},t)$. $\hat{x}_{t}$ is obtained by a single-step denoising and diffusion process on $\widetilde{x}_{t}$, thus their distributions remain highly similar as well as the corresponding quantization noise:

Finally, the quantization noise $q_{\theta}(\widetilde{x}_{t},t)$ can be determined, as the Gaussian noise $z_{1}$ is manually designed and $\widetilde{\epsilon}_{\theta}(\hat{x}_{t},t)$ is the output of the noise predicting network, both of which are ascertainable.

Based on the above analysis, we can estimate the quantization noise by simulating the diffusion model training process. We first add the deterministic noise $z_{1}$ to $\widetilde{x}_{t-1}$, which can be filtered out in the diffusion sample process, and then the rest of the indeterministic noise is the quantization noise $q_{\theta}(\widetilde{x}_{t},t)$. Besides, the quantization noise is obtained through estimation and doesn’t align perfectly with the actual noise in terms of pixel dimension, whereas it is identical at the level of the overall distribution. Thus we get the distribution of the quantization noise at stage $t\!-\!1$ and correct the output sample in Eq. 4.

The above procedure is for single-step quantization noise estimation. In the diffusion model, the distributions of samples from neighboring steps are very similar, as well as the corresponding quantization noise distribution ($q_{\theta}(\widetilde{x}_{t+1},t+1)\approx q_{\theta}(\widetilde{x}_{t},t)$). Therefore, in actual sampling process, we divide entire sampling steps into multiple stages, estimating the distribution of quantization noise one time per stage. For example our method run only 4 times to estimate the quantization noise during a sampling process of 100 time steps, which brings a negligible increase in sampling duration. As shown in Fig. 3, our QNCD periodically estimates the distribution of inter quantization noise and corrects distribution deviations caused by this noise. This noise correction enables the distribution of sample outputs to closely align with the full-precision models.

Given the output of the floating-point noise prediction network and the corresponding quantization noise, we plot the change curve of the signal-to-noise ratio (SNR) of the quantized diffusion model, as shown in Fig. 8, as well as the cosine similarity between the corresponding output and the FP output. We can find that: 1, with the increase of denoising steps, the cosine similarity between the quantized model output and the FP output is continuously decreasing, which also means that the overall quantization noise is continuously accumulating, and the corresponding SNR is continuously decreasing. 2, our method can estimate and filter out the quantization noise on a global scale, resulting in a better SNR, and a sampling process closer to that of the floating-point model.

As shown in Tab. 3, we conducted ablation experiments on Stable Diffusion (step=50) with MS-COCO 512×512 dataset to demonstrate the effectiveness of our QNCD, which includes intra quantization noise correction (Intra-QNCD) and inter quantization noise correction (Inter-QNCD). Intra-QNCD reduces FID by 0.43 compared to Q-Diffusion, while Inter-QNCD reduces FID by 0.24 and improves the CLIP Score by 0.06. Combining both blocks, QNCD achieves a 0.51 reduction in FID, showing their collaborative effectiveness in improving performance. These results validate our noise correction techniques in post-training quantization of diffusion models.

For fair comparison, we provide end-to-end inference times in Tab. 4. Inference times are based on the UNet of Stable Diffusion V1.4, which denote whole denoising process of diffusion models. The experimental background is A100, TensorRT-8.6 and CUDA-11.7. Similarly to our QNCD, Q-Diffusion introduces the Short-Cut split operation in pursuit of better model performance, which also imposes an additional inference burden (26.5ms compared to original PTQ). Our method runs at a similar speed to Q-Diffusion, but with higher image quality.

As shown in Tab. 2, the FID metrics of PTQD and Q-Diffusion on ImageNet dataset are 11.94 and 10.92, superior to the FP’s score of 12.45 under the W8A8 setting. This same pattern extends to the LSUN-Bedrooms and CIFAR datasets, which is unexpected and implies that the FID metric may not be an optimal indicator of image quality. This is because FID focuses more on the overall distribution similarity rather than the specific quality of each image. For a more comprehensive comparison, we further refer to objective Image Quality Assessment(IQA) metrics proposed in CLIP-IQA (clipiqa, ) to evaluate 5000 synthesized images. Our method achieves an IQA metric of 0.793, which is better than Q-Diffusion (0.775), but still falls short compared to FP (0.847).