Efficient Quantization Strategies for Latent Diffusion Models

DMs [39, 13] involve two processes: forward diffusion and reverse diffusion. Forward diffusion process adds noise to input image, $\mathbf{x}_{0}$ to $\mathbf{x}_{T}$ iteratively and reverse diffusion process denoises the corrupted data, $\mathbf{x}_{T}$ to $\mathbf{x}_{0}$ iteratively. LDMs involve the same two processes but in the latent space $\mathbf{z}_{i}$ instead of original data space. We reiterate the diffusion processes in the latent space.

The forward process adds a Gaussian noise, $\mathcal{N}(\mathbf{z}_{t-1};\sqrt{1-\beta_{t}}\mathbf{z}_{t-1},\beta_{t}\mathbb{I})$, to the example at the previous time step, $\mathbf{z}_{t-1}$. The noise-adding process is controlled by a noise scheduler, $\beta_{t}$.

The reverse process aims to learn a model to align the data distributions between denoised examples and uncorrupted examples at time step $t-1$ with the knowledge of corrupted examples at time step $t$.

To simplify the optimization, [13] proposes only approximate the mean noise, $\theta(\mathbf{z}_{t},t)\sim\mathcal{N}(\mathbf{z}_{t-1};\sqrt{1-\beta_{t}}\mathbf{z}_{t},\beta_{t}\mathbb{I})$, to be denoised at time step $t$ by assuming that the variance is fixed. So the reverse process or the inference process is modeled as:

Post-training quantization (PTQ) reduces numerical representations by rounding elements $v$ to a discrete set of values, where the quantization and de-quantization can be formulated as:

where $s$ denotes the quantization scale parameters. round(·) represents a rounding function [4, 47]. $c_{min}$ and $c_{max}$ are the lower and upper bounds for the clipping function $clip(\cdot)$. Calibrating parameters in the PTQ process using weight and activation distribution estimation is crucial. A calibration relying on min-max values is the most simple and efficient method but it leads to significant quantization loss (see examples in Figure 2), especially when outlier values constitute a minimal portion of the range [46]. More complicated optimized quantization is applied to eliminate the outliers for better calibration. For example, PTQ4DM [37] and Q-Diffusion [21] apply reconstruction-based PTQ methods [23] to diffusion models. PTQD [11] further decomposes quantization noise and fuses it with diffusion noise. Nevertheless, all these methods demand a significant increase in computational complexity. In this study, our exclusive experimentation revolves around min-max quantization due to its simplicity and efficiency. Additionally, we aim to demonstrate how our proposed strategy proficiently mitigates the quantization noise inherent in min-max quantization.

In general, a homogeneous quantization strategy is applied to all parts in a network for its simplicity. In other words, the same precision quantization is applied to all modules. Though different quantization schemes coincide in mixed-precision quantization [6, 8, 44] and non-uniform quantization [22, 50, 14], additional training or optimization is required to configure more challenging quantization parameters. LDMs consist of many downsampling and upsampling blocks (Figure 3(b)). The concatenated blocks in LDMs pose challenges for end-to-end quantization, given dynamic activation ranges and recursive processing across steps and blocks. In contrast to prior attempts to uniformly quantize all blocks, our approach suggests an efficient quantization strategy by initiating a hybrid quantization at a specific block and subsequent blocks.

We attentively discover that quantizing up to a specific block in LDM imposes minimum loss in image generating quality. In Figure 3(a), as we measure FID for a progressive homogeneous quantization on LDM 1.5, blocks nearer to the output exhibit higher sensitivity to quantization. Identifying these sensitive blocks beforehand at a low cost facilitates the implementation of a decisive strategy. However, computing FID is computationally expensive and time-consuming. An efficient metric is required to pinpoint sensitive blocks (and modules) for the development of an improved quantization strategy.

We denote $\theta_{fp}$ and $\theta_{q}$ are the full-precision model and quantized model respectively. If we simplify the iterative denoising process as (3), where the timestep embedding and text embedding are omitted for simplicity, intuitively each parameterized module contributes quantization noise and the recursive process will only accumulate these noise. This behaviour is modeled as (4) where $\Phi$ is an accumulative function. This is akin to conducting sensitivity analysis to scrutinize how the accumulation of quantization impact on each block or module and influence the quality of the generated images.

We need a computationally efficient $metric(.,.)$ function to identify sensitive blocks in $\theta$ if the accumulated $\Delta_{block,t:T\rightarrow 1}$ is too high, and identify sensitive modules if the relative $\Delta_{module}$ is too high. In other words, the $metric(.,.)$ necessitates three key properties:

• •

  Accumulation of iterative quantization noise needs to be ensured as the global evaluation of the quantization noise at the output directly reflect image quality.

• •

  Relative values across different modules provide qualitative comparisons to identify local sensitivity.

• •

  Convenient computation allows an efficient evaluation and analysis which expedites quantization process.

Hence metrics like MSE is not sufficient to make realtive and fair local comparisons, and FID is not efficient to facilitate the quantization process.

We adapt a relative distance metric, $\mathbf{SQNR}$ [17], as the $metric(.,.)$:

where $\xi$ represents the whole model for global sensitivity evaluation and a module for local sensitivity evaluation. Other works have associated $\mathbf{SQNR}$ with the quantization loss. [38] evaluates $\mathbf{SQNR}$ for quantization loss on a convolutional neural network (CNN). [18] optimizes $\mathbf{SQNR}$ as the metric to efficiently learn the quantization parameters of the CNN. [31] computes $\mathbf{SQNR}$ at the end of inference steps to quantify the accumulated quantization loss on softmax operations in diffusion models. [32] assesses $\mathbf{SQNR}$ to fairly compare the bit options for the quantized weight in a deep neural network. Therefore $\mathbf{SQNR}$ is an optimal candidate for $metric(.,.)$ and also it offers three major advantages:

• •

  At global levels, instead of modeling an accumulative function $\Phi$, we evaluate the time-averaged $\mathbf{SQNR}$ at the output of the $\theta$, i.e. $\Delta_{block_{N},t:T\rightarrow 1}=\mathbb{E}_{t=1}^{T}[\mathbf{SQNR}_{\theta,t}]$. Averaging the relative noise at the output comprises an end-to-end accumulation nature. A low $\Delta_{block_{N},t:T\rightarrow 1}$ indicates a high quantization noise, which leads to poorly generated images. We justify this correlation as qualitative measure $\mathbb{E}_{t=1}^{T}[\mathbf{SQNR}_{\hat{\theta},t}]$ and FID and show the correlation in Figure 4.

• •

  At local levels, $\mathbf{SQNR}_{\xi,t}$ is a relative metric that measures the noise ratio between the full-precision and quantized values. This aids a fair comparison between local modules to identify the sensitive modules.

• •

  A small number of examples is sufficient to compute the $\mathbf{SQNR}$ instantaneously with high confidence (Supplementary C). This swift computation enables an efficient assessment of quantization noise and the quality of the generated image, in comparison to FID and IS.

Conventionally, the quantization is applied with a consistent precision on the whole model. But the quantization fails with low-bit precision constantly [21, 11] and there is no substantiate evidence to identify the root cause for these failures. We measure the time-averaged $\mathbf{SQNR}$ at the output of $\theta$ which is denoted as $\mathbf{\overline{SQNR}}_{\theta}$ shown in Eq (6). Progressively quantizing each block in $\theta$ examines the sensitivity of each block based on $\mathbf{\overline{SQNR}}_{\theta}$. Once blocks with high sensitivity are identified, a higher-precision quantization can be applied to alleviate quantization noise from these blocks. The proposed global hybrid quantization is illustrated in Figure 5(b). We first explore the sensitivity of each block in the UNet by quantizing first $N$ blocks, and leave the rest blocks unquantized. In Figure 5(a), the block sensitivity of different LDMs at different quantization bits consistently indicates that latter blocks in up blocks are severely sensitive to quantization. Using the time-averaged $\mathbf{SQNR}$, sensitive blocks can be identified efficiently and a higher precision quantization can be applied on these blocks collectively.

Similarly to block sensitivity analysis, we evaluate $\mathbf{SQNR}_{\xi}$ at each module and visualize the sensitivity at different inference steps. We intend to formulate a local sensitivity first and identify local sensitive modules as illustrated in Figure 6. We visualize the $\mathbf{SQNR}_{\xi}$ for each module at two inference steps, $[1,50]$, and the average over all 50 sampling steps when quantizing LDM 1.5 with 8A8W in Figure 6 (Illustration for other LDMs in Supplementary A). A low local $\mathbf{SQNR}_{xi}$ only suggests that one module is sensitive to quantization noise, and vice versa. From the module sensitivity across inference time and across the whole model, two types of root causes for the local sensitivity can be identified and will be addressed respectively.

We present qualitative examples for both global hybrid strategy Figure 9 and local noise correction strategy Figure 10. Qualitatively, LDM XL is more robust to min-max quantization but LDM 1.5 generates extensively corrupted images with naive min-max quantization. Nonetheless, the local noise corrections can significantly improve the quality of the images by just identifying 10% of most sensitive modules.

As mentioned in Section 3.3, we examine the different hybrid quantization configurations to verify the block sensitivity shown in Figure 5(a). We explore all hybrid quantization configurations on LDM 1.5 and evaluate FID and computational efficiency in Table 3. Though FID suggests that higher-precision quantization generates less presentable images, which contrasts to the expectation. We argue that FID is not an optimal metric to quantize the quantization loss when the quantization noise is low. Only when significant quantization noise corrupts the image, a high FID indicates the generated images are too noisy.

As we obtain the local sensitivity for the LDM, we can identify the sensitive local modules from most sensitive to least sensitive based on $\mathbf{SQNR}_{\xi}$. We explore the portion of the modules to adapt SmoothQuant and evaluate $\mathbf{\overline{SQNR}}_{\theta}$. Note that both statistic preparation and computation on $diag(\mathbf{s})^{-1}$ proportionally increase with the number of identified modules. As the 10% most sensitive modules identified, SmoothQuant can effectively address the activation quantization challenge as $\mathbf{\overline{SQNR}}_{\theta}$ increases and saturate till 90%. And smoothing least sensitive modules results in deterioration due to unnecessary quantization mitigation. With different mitigation factor, $\alpha$, the behaviour remains consistent, only higher $\alpha$ reduces quantization errors by alleviating activation quantization challenge.

We also adapt SmoothQuant to the hybrid quantization to evaluate the benefit of such activation quantization noise correction. 10% most sensitive modules are identified and applied with SmoothQuant ($\alpha=0.7$). FID and $\mathbf{\overline{SQNR}}_{\theta}$ are reported. 1-step calibration significantly reduces the quantization noise for the hybrid quantization. However, SmoothQuant does not make visible difference. This is because identified sensitive modules reside mostly in the upsampling blocks, which are replaced with higher-precision quantization and therefore the quantization noise is already reduced by the hybrid strategy.

As we discuss in 3.4.2, we notice that calibrate on multiple sampling steps is problematic for quantization due to gradually changing activation ranges. And quantization parameters are most robust to the added diffusion noise at first sampling step (last forward diffusion step). Hence calibrating on fewer sampling steps leads to less vibrant activation ranges, which results in better quantization calibration. We evaluate the effect on the number of calibration steps on the relative quantization error, $\mathbf{\overline{SQNR}}_{\theta}$, and FID. It is obvious that the proposed 1-step calibration significantly improves the calibration efficiency ($\times{50}$) compared to calibration methods proposed in [21, 11, 37], where calibration data are sampled from all 50 sampling steps.

We develop an efficient computation of $\mathbf{SQNR}$ by forward passing two same inputs to all input modules. And we rearrange the output of all modules so that first half of the output is full-precision output, and the second half of the output is quantized output. In this way, $\mathbf{SQNR}$ can be computed at all modules with a single forward pass and with a small number of examples. Other complicated optimization like Q-Diffusion or PTQ4DM collects thousands of calibration data over all 50 sampling steps. Detailed description is included in Supplementary C.