Towards Accurate Post-training Quantization for Diffusion Models

Diffusion models leverage a forward pass to impose noise on images and a reverse pass to transform Gaussian noise into an image for generation. Denoting the real data as $\bm{x}_{0}$ and the latent image at the $t_{th}$ step as $\bm{x}_{t}$, the probability of the forward process can be represented as follows:

where $\beta_{t}$ means the variance schedule at the $t_{th}$ step that indicates imposed Gaussian noise to the latent image. When the total number of forward steps denoted as $T$ becomes large enough, the latent image $\bm{x}_{T}$ can be regarded as the standard Gaussian noise. We leverage an approximated condition distribution $p_{\theta}(\bm{x}_{t-1}|\bm{x}_{t})$ to generate the latent image in reverse process due to the intractability of true distribution $q(\bm{x}_{t-1}|\bm{x}_{t})$, where the approximated distribution is parameterized by neural networks with weight $\theta$:

The training process of diffusion model aims to minimize the negative log-likelihood with the evidence lower bound optimization in variational inference:

In the practical applications, iterative noise estimation process is implemented with the diffusion model for content generation, and the heavy computational cost of the reverse phase disables the deployment in resource-constrained devices such as mobile phones and robots. To accelerate the denoising process for each reverse step, post-training quantization leverages a small calibration set to learn the rounding function parameters for weights and activations of the decoder, where the quantization function can be represented as follows:

where $x$ and $\hat{x}$ represent real-valued and quantized matrix respectively. $[\cdot]$ means the rounding function to the nearest integer and $\Phi$ is the clipping operation that regularizes the element into the range from $z_{min}$ to $z_{max}$. The quantization scaling parameter $s$ indicates the interval between adjacent rounding points. As empirically demonstrated in [31], the activation distribution varies significantly across different timesteps during the reverse process, and the shared rounding functions usually cause severe quantization errors for image generation. Moreover, randomly selecting the timestep to generate latent images for calibration set construction fails to provide sufficient information for generalizable quantization function learning.

Since the activation distribution changes significantly across timesteps, discretizing the full-precision intermediate features in similar data distribution with the same quantization functions can reduce the quantization errors. We first describe the distribution-aware quantization scheme and then illustrate the differentiable group assignment of timesteps.

Shared quantization functions may cause large clipping errors for widely distributed activations and rounding errors for narrowly distributed ones. Directly assigning specific rounding functions for network activations in each timestep leads to overfitting in optimization because of the limited calibration samples, and quantizing activations in the timestep where optimal quantization range is similar with the same rounding functions can achieve better trade-off between the quantization accuracy and rounding function generalizability. Assuming partitioning all $T$ timesteps into $C$ groups, the quantization strategy can be written in the following:

where $g(i)$ represents the assigned group for activations in the $i_{th}$ timestep. Meanwhile, $s_{g(i)}$, $z_{min}^{g(i)}$ and $z_{max}^{g(i)}$ respectively stand for scale parameters, the lower bound and the upper bound of quantization for activations in the $i_{th}$ timestep. Assigning the optimal group indexes for different timesteps is critical in distribution-aware quantization to reduce the quantization errors without obvious computation overhead. Since enumerating assignment permutation is NP-hard to find the optimal solution, we extend the differentiable search framework to efficiently partition timesteps with minimal quantization errors. In the differentiable search, the latent images are quantized by all quantization functions, whose output values are summed with learnable importance weights to form the input for the next layer in the diffusion model:

where $\sigma_{g}$ means the importance weight of the quantization function for the $c_{th}$ group with the normalization $\sum_{g=1}^{G}\sigma_{g}=1$. When the training process completes, the rounding function with the largest importance weight is selected to be the search results for group-wise quantization. Despite the noise estimation loss of diffusion models, we also enforce the importance weights to approach zero or one by minimizing the entropy to avoid discretization errors in rounding function selection. The overall optimization objective $J$ is written as follows, where we denote $\bm{\epsilon}_{\theta}(\sqrt{\overline{\alpha}_{t}}\bm{x}_{0}+\sqrt{1-\overline{\alpha}_{t}}\bm{\epsilon},t)$ as $\bm{\epsilon}_{\theta}$ for simplicity:

where $J_{d}$ and $J_{e}$ respectively represent the simplified variational lower bound of diffusion model objective and the discretization minimization loss in differentiable search, and the hyperparameter $\lambda$ balances the importance of different terms. $\sigma_{g}^{t}$ demonstrates the importance weights of the quantization function for the $g_{th}$ group in the $t_{th}$ timestep. The noise $\bm{\epsilon}$ from standard Gaussian distribution is approximated by the predicted noise $\bm{\epsilon}_{\theta}$ in the optimization objective. The diffusion parameter $\overline{\alpha}_{t}=\prod_{i-1}^{t}1-\beta_{i}$ controls the strength of noise in diffusion. We jointly update the parameters in quantization functions and the importance weights until convergence or achieving the maximal iteration steps, and the discretized hypernetwork is directly employed to generate images efficiently.

The pipeline of post-training quantization for iterative reverse process in diffusion models differs significantly from that in conventional vision models. Leveraging latent images in all timesteps leads to unbearable training cost for quantization function learning, and latent images in adjacent timesteps can only offer redundant information for parameter optimization. On the contrary, randomly select part of the timesteps usually fails to provide sufficient supervision that is representative to demonstrate the real distribution of the latent images. Therefore, it is desirable to actively sample the timesteps to generate latent images for quantization parameter learning with effective guidance. We generate representative samples by structural risk minimization, which minimize the distance between selected and real distribution. We first introduce the extension of SRM principle to active timestep selection, and then formulate the selection criteria that can be feasibly computed.

Structural risk minimization principle minimizes the upper bound of the true risk on unseen data distribution, where the bound can be written as follows for a dataset containing $n$ samples with the probability at least $1-\delta$[4]:

where $E(J(x))$ and $\overline{E(J(x))}$ respectively illustrate the true expectation of the risk $J$ for real data distribution $x$ and the empirical expectation of that for sampled data from $x$, and $R_{n}(\mathcal{F})$ is the Rademacher complexity over the function class $\mathcal{F}$. The SRM principle requires the data to be sampled from i.i.d. distribution, while the latent images in selected timesteps should be more informative and representative. Therefore, we rewrite the SRM principle in the following way, where the detailed formulation is in the supplementary:

where we omit the data distribution $x$ for simplicity. $\overline{E_{S}(J)}$ denotes the empirical risk of the latent images of selected timesteps for noise estimation, and $\mathcal{R}_{0}$ demonstrates the complexity of the diffusion model in the reverse process. $X$ and $X_{s}$ stand for the distribution of latent images generated in all timesteps and the selected ones. The maximal mean discrepancy between two distributions $p(X)$ and $p(X_{s})$ is represented as $MMD(p(X),p(X_{s}))$, which demonstrates the generalization ability of the calibration sets for quantization learning. The first criteria acquired by worst-case empirical risk of sampled latent images is formulated as follows for timesteps selection :

where $\mathcal{S}$ represents the images selected in the calibration set, and $J$ is the optimization objective defined in (7). This formula aims to train the highly quantified diffusion model with the constructed calibration sets. Since the original latent $\bm{x}_{T}$ can be regarded as the standard Gaussian noise without bias, the objective $J_{d}$ does not affect the worst-case empirical risk with different timesteps. The variance of $J_{e}$ influences the worst-case empirical risk across timesteps, because the entropy of the importance weights of distribution-aware quantization functions changes with the timesteps. Therefore, the criteria $s_{1}$ from the empirical risk minimization can be transformed to selecting the timestep with the highest entropy of importance weights as $s_{1}=\sum_{g=1}^{G}-\sigma_{g}^{t}\log\sigma_{g}^{t}$.

Meanwhile, The definition of maximal mean discrepancy can be written as follows, where we denote $MMD(p(X),p(X_{s}))$ as $M$ for simplicity:

where $U$ means the full set containing all original latent and timesteps for calibrating image selection, and $|\cdot|$ represents the number of elements in the set. $\varphi$ is a constant in timestep sampling and $N_{t}$ denotes the number of sampling times for the $t_{th}$ timestep in calibration set construction. Because the estimated noise of the samples in the full set is intractable, we utilize the number of sampling times to optimize the maximal mean discrepancy based on upper confidence bound (UCB) principle [1] which achieves exploitation-exploration trade-off when sampling. A detailed derivation of the formula (11) can be found in the supplementary material. For the timestep when we sample a large number of latent images for calibration set construction, the maximal mean discrepancy becomes low as we acquire sufficient information of the latent image distribution in this timestep. Therefore, we explore latent images in the timestep with few sampling times to further minimize the maximal mean discrepancy with high marginal benefits. The criteria $s_{2}$ from the maximal mean discrepancy is designed as $s_{2}=1/(N_{t}+1)$. The overall timestep selection criteria can be written as follows:

where $\eta$ is a hyperparameter to balance the importance of empirical risk and maximal mean discrepancy. The overall pipeline of our method is depicted in Figure 2. With the optimally selected timesteps, the generated calibration images can provide effective supervision for quantization function learning, which can be well generalized in deployment.

We utilize the diffusion frameworks for post-training quantization including DDIM [32] and LDMs [29] with their pre-trained weights, which require 100 iterative denoising timesteps for image generation in most experiments. We set the bitwidth of quantized weight and activation to 6 and 8 to evaluate our method in different quality-efficiency trade-offs and utilized the uniform quantization scheme where the interval between adjacent rounding points was equal. For distribution-aware quantization across different timesteps, we partitioned all timesteps into eight groups in most experiments. We followed the initialization of the quantization function parameters in [20] for the baseline methods and our APQ-DM, where we minimized the $l_{p}$ distance [27, 37] between the full-precision and quantized activations to optimize the value range for clipping. The hyperparameters $\lambda$ in the objective of differentiable search and $\eta$ in the timestep selection criteria were set to 0.8 and 1.5 respectively.

For the parameter learning in differentiable search, we generated 1024 images for hyper-network learning where the batchsize was assigned with 64 for calibration set construction. The learning rate was initialized to $3e^{-3}$ and $5e^{-3}$ for 6 and 8-bit diffusion models and ended up with $1e^{-5}$ for all bitwidth settings with 0.05 decaying strategy. The quantization function parameters and the importance weights were jointly updated for 10 epochs in the differentiable search, and the acquired distribution-aware quantization function was directly employed for image generation.

In order to investigate the influence of the distribution-aware quantization for network activations across different timesteps, we vary the number of groups with different trade-offs between quantization accuracy and rounding function generalizability. To show the effectiveness of our active timestep selection for calibration set generation, we compare our strategy with various sampling techniques. Meanwhile, we modified the hyperparameter $\lambda$ and $\eta$ to demonstrate the effect of the discretization loss in rounding function selection and the maximal mean discrepancy in timestep selection criteria. All experiments in the ablation study were conducted with the $32\times 32$ cifar-10 dataset and the DDIM diffusion framework.

Performance w.r.t. the number of timestep groups: Dividing the timesteps into more groups can significantly reduce the clipping and rounding errors for differently distributed activations in quantization function learning, while may resulting in the rounding function overfitting due to the limited calibration samples and large-scale learnable parameters. Table 1 illustrates the quantization errors, Inception Score (IS) and FID score for our method that partitions the timesteps into different numbers of groups. Observing the FID and IS for different group partition settings, we conclude that dividing the timesteps into several groups outperforms both the shared quantization policy and the step-wise rounding functions. Therefore, we assign the number of groups for the timestep partition to 8 in the rest of experiments to achieve the optimal trade-off between the quantization accuracy and rounding function generalizability.

Performance w.r.t. different timestep sampling strategies for calibration set construction: We compare our active timestep sampling strategy for calibration set generation with random and heuristic sampling policies [31]. Random sampling assigns an integer number drawn from uniform distribution from zero to the maximal time steps $T$, and heuristic sampling determines the timestep from a Gaussian distribution $\mathcal{N}(\mu,\frac{T}{2})$ where $\mu$ is less than $\frac{T}{2}$. Table 2 shows the generation quality for different timestep sampling methods across various sizes of the calibration set. Our active sampling strategy outperforms the random and heuristic sampling policies by a large margin, and the advantage is more obvious for calibration sets with small sizes because informative samples are extremely important for post-training quantization in the scenario without sufficient images.

Visualization of importance weight in differentiable search: Figure 3(a) depicts the evolution of importance weights during the differentiable search for group assignment, where different colors represent disparate groups. At the early stage of the differentiable search, the differences between importance weights are not obvious because of insufficient calibration images. When the network gradually converges, the principal impacts on the performance result from the quantization functions with different rounding and clipping errors and differentiate group assignments between different data distributions.

Performance w.r.t. hyperparameters $\lambda$ and $\eta$: The hyperparameter $\lambda$ controls the importance of the discretization loss in distribution-aware quantization function in the objective of differentiable search, and $\eta$ balances the empirical risk and the maximal mean discrepancy in the timestep selection. Figure 3(b) depicts the FID for different hyperparameter settings, where the medium value for both parameters achieves the highest generation quality. The model performance is more sensitive to the hyperparameter $\lambda$ because the importance weights of quantization functions in different groups usually have similar distribution as one-hot vector, and slight change to $\lambda$ leads to large perturbation to the overall objective in differentiable search due to the logarithm.

In this section, we compare our proposed method with the state-of-the-art data-free post-training quantization frameworks including LSQ [9] and those specifically designed for diffusion models including PTQ4DM [31] and Q-diffusion [20]. The IS, FID, and sFID scores of the baseline methods are acquired by implementing the officially released code or our re-implementation. For fair comparison of all listed methods, we leverage the rounding function in LSQ for quantization and de-quantization, and generate latent images with 100 iterative timesteps.

Results on unconditional generation: Unconditional generation samples a random variable for diffusion models to yield images with similar distribution of the training datasets. We evaluate our data-free post-training quantization methods on $32\times 32$ Cifar-10 [15], $64\times 64$ CelebA [23], $256\times 256$ LSUN-Church Outdoor and LSUN-Bedroom datasets [39] for DDIM frameworks, and evaluate for LDMs diffusion frameworks on $256\times 256$ CelebA-HQ [16], LSUN-Church Outdoor and LSUN-Bedroom datasets, where the generalization quality and efficiency are reported in Table 3 and Table 4 respectively. LSQ learns the optimal quantization step sizes with minimized discretization errors, while the shared quantization policy across timesteps and randomly constructed calibration set in diffusion models leads to significant quantization loss. PTQ4DM and Q-diffusion employ the step-wise quantization functions to minimize the quantization errors of the diversely distributed activations across timesteps, and presents heuristic timesteps selection criteria for calibration image generation. However, the optimization of large-scale learnable parameters faces the challenges of overfitting due to the very limited quantity of calibration samples, and the data-independent calibration set construction cannot guarantee the optimality of the calibration images. As a result, our method outperforms PTQ4DM by 0.32 (2.55 vs. 2.23) and 1.02 (6.46 vs. 7.48) for IS and FID in LSUN-Bedroom respectively. The computational cost remains the same for baseline methods and our APQ-DM due to the stored rounding parameters. The advantage of our method becomes more obvious for 6-bit diffusion models because quantization errors and calibration sample informativeness are more important for networks with low capacity.

Results on conditioned image generation: Conditioned image generation synthesizes samples according to text including class names or descriptions. For class-conditional image generation, we discretize the LDMs model that is pre-trained on the $256\times 256$ class-conditional ImageNet [8] dataset, where the guidance strength is set to 3.0 to balance the generation quality and diversity. Table 4 shows the quantitative experimental results for different post-training quantization methods, while our method increases the FID and IS by 1.01 (11.58 vs.12.59) and 17.38 (179.13 vs. 161.75) respectively compared with the state-of-the-art method PTQ4DM. For description-conditional image generation, Figure 4 demonstrates some examples of the text prompts images that are generated by different quantized Stable Diffusion models, where our method can still acquire plausible images with high-quality details with weights and activations in low bitwidths. Since conditioned image generation is widely adopted in many realistic multimedia applications, our method brings potential to deploy large pre-trained diffusion models on mobile devices or robots under limited resource constraints with satisfying generation quality.