SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Junhyuk So $^{1*}$  Jungwon Lee $^{2*}$  Daehyun Ahn $^{3}$  Hyungjun Kim $^{3}$  Eunhyeok Park $^{1,2}$
$^{1}$  Department of Computer Science and Engineering, POSTECH
2 Graduate School of Artificial and Intelligence, POSTECH
3 SqueezeBits. Inc
The diffusion model has gained popularity in vision applications due to its remarkable generative performance and versatility. However, high storage and computation demands, resulting from the model size and iterative generation, hinder its use on mobile devices. Existing quantization techniques struggle to maintain performance even in 8-bit precision due to the diffusion model's unique property of temporal variation in activation. We introduce a novel quantization method that dynamically adjusts the quantization interval based on time step information, significantly improving output quality. Unlike conventional dynamic quantization techniques, our approach has no computational overhead during inference and is compatible with both post-training quantization (PTQ) and quantization-aware training (QAT). Our extensive experiments demonstrate substantial improvements in output quality with the quantized diffusion model across various datasets.
Generative modeling is crucial in machine learning for applications such as image [1, 2, 3, 4, 5, 6, 7, 8], voice [9, 10], and text synthesis [11, 12, 13]. Diffusion models [3, 2, 4], which progressively refine input images through a denoising process involving hundreds of iterative inferences, have recently gained prominence due to their superior performance compared to alternatives like GANs [14]. However, the high cost of diffusion models presents a significant barrier to their widespread adoption. These models have large sizes, often reaching several gigabytes, and demand enormous computation of iterative inferences for a single image generation. Consequently, executing diffusion models on resource-limited mobile devices is practically infeasible, thus most applications are currently implemented on expensive, high-performance servers.
To fully exploit the potential of diffusion models, several methods to reduce the computational cost and memory requirement of diffusion models while preserving the generative performance have been proposed. For example, J. Song et al. [2] and L. Liu et al. [15] proposed a more efficient sampling scheduler, and T. Salimans & J. Ho proposed to reduce the number of sampling steps using knowledge distillation technique [16, 17]. As a result, high-fidelity images can be generated with fewer sampling steps, making the use of diffusion models more affordable and feasible. Despite these advancements, the denoising process of diffusion models still demands a substantial computational cost, necessitating further performance enhancements and model compression.
While the majority of previous approaches have focused on reducing the number of sampling steps to accelerate the denoising process, it is also important to lighten the individual denoising steps. Since the single denoising step can be regarded as a conventional deep learning model inference, various model compression techniques can be used. Quantization is a widely used compression technique where both weights and activations are mapped to the low-precision domain. While advanced quantization schemes have been extensively studied for conventional Convolution Neural Networks (CNNs) and
Figure 1: The forward and reverse processes of diffusion models.
language models, their application to diffusion models has shown significant performance degradation. Due to the unique property of diffusion models, such as the significant changes in the activation distribution throughout the iterative inferences, the output is heavily distorted as the activation bit-width decreases [18]. Existing quantization techniques, including both quantization-aware training (QAT) [19, 20, 21] and Post-training quantization (PTQ) [22, 23, 24, 25], are designed to address a specific distribution in existing DNNs, therefore cannot deal with time-variant activation distributions in diffusion models.
To tackle the unique challenges of diffusion model quantization, we introduce a novel design called Temporal Dynamic Quantization (TDQ) module. The proposed TDQ module generates a time-dependent optimal quantization configuration that minimizes activation quantization errors. The strong benefit of the TDQ module is the seamless integration of the existing QAT and PTQ algorithms, where the TDQ module extends these algorithms to create a time-dependent optimal quantization configuration that minimizes activation quantization errors. Specifically, the module is designed to generate no additional computational overhead during inference, making it compatible with existing acceleration frameworks without requiring modifications. The TDQ module significantly improves quality over traditional quantization schemes by generating optimal quantization parameters at each time step while preserving the advantages of quantization.
2 Backgrounds and Related Works
Diffusion models have been first introduced in 2015 [26] and revolutionized image generation by characterizing it as a sequential denoising process. As shown in Fig. 1, the forward diffusion process gradually transforms the image  $(x_0)$  into a random data  $(x_{T})$  which follows standard normal distribution by adding small Gaussian noise at each time step. The reverse diffusion process generates a clean image  $(x_0)$  from a random data  $(x_{T})$  by gradually removing noise from the data through iterative denoising steps. Therefore, a diffusion model learns the reverse process which is estimating the noise amount on a given noisy data at each time step  $(x_{t})$ . The forward  $(q)$  and reverse  $(p_{\theta})$  processes can be described as
$$
q \left(x _ {t} \mid x _ {t - 1}\right) = N \left(x _ {t}; \sqrt {1 - \beta_ {t}} x _ {t - 1}, \beta_ {t} \mathbf {I}\right), \tag {1}
$$
$$
p _ {\theta} \left(x _ {t - 1} \mid x _ {t}\right) = N \left(x _ {t - 1}; \mu_ {\theta} \left(x _ {t}, t\right), \sigma_ {t} ^ {2} \mathbf {I}\right), \tag {2}
$$
where  $\beta_{t}$  denotes the magnitude of Gaussian noise.
[3] introduced a reparameterization trick for  $\mu_{\theta}$  and the corresponding loss function, which facilitates the training of diffusion models.
$$
\mu_ {\theta} \left(x _ {t}, t\right) = \frac {1}{\sqrt {\alpha_ {t}}} \left(x _ {t} - \frac {\beta_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \epsilon_ {\theta} \left(x _ {t}, t\right)\right) \tag {3}
$$
$$
L _ {\text {s i m p l e}} = E _ {t, x _ {0}, \epsilon} [ \| \epsilon - \epsilon_ {\theta} (x _ {t}, t) \| ] \tag {4}
$$
While diffusion models can produce high-quality images, the iterative denoising process makes diffusion models difficult to be used for real-world scenarios. The early works on diffusion models such as DDPM [3] required hundreds to thousands of iterative inferences to generate single image, resulting extremely slow sampling speed. Therefore, numerous studies have been investigating algorithmic enhancements and various optimizations to boost performance, aiming to make diffusion
models more efficient and suitable for real-world applications. DDIM [2] introduced an implicit probabilistic model that reinterprets the Markov process of the DDPM method achieving competitive image quality with one-tenth of denoising steps. Distillation-based methods [16, 17] proposed to reduce the number of denoising steps using knowledge distillation techniques.
On the other hand, to address the significant computational overhead associated with generating high-resolution images, [4] proposed a Latent Diffusion Model (LDM) where the diffusion model takes latent variables instead of images. Especially, the large-scale diffusion model (e.g., Stable Diffusion [4]) has leveraged LDM and learned from a large-scale dataset (LAION-Dataset [27]), enabling the creation of high-quality, high-resolution images conditioned on textual input.
2.2 Quantization
Quantization is a prominent neural network optimization technique that reduces storage requirements with low-precision representation and performance improvement when the corresponding low-precision acceleration is available. With  $b$ -bit precision, only  $2^{b}$  quantization levels are accessible, which considerably limits the degree of freedom of the quantized model compared to floating-point representation. Therefore, the final quality significantly varies depending on the tuning of quantization hyperparameters, such as the quantization interval or zero offset. To preserve the quality of output in low-precision, it is crucial to update the quantization parameters as well as the model parameters, taking into account the specific features and requirements of the target task.
2.3 Quantization-aware Training and Post-training Quantization
Quantization algorithms can be broadly categorized into two types: Quantization-Aware Training (QAT) [19, 28, 29, 20, 30] and Post-Training Quantization (PTQ) [23, 24, 25]. QAT applies additional training after introducing quantization operators, allowing the update of network parameters toward minimizing the final loss value considering the effect of the quantization operators. On the other hand, PTQ does not apply end-to-end forward/backward propagation after quantization. Instead, it focuses on reducing block-wise reconstruction errors induced by quantization. QAT typically outperforms PTQ in low-precision scenarios, but it may not always be applicable due to limitations such as the deficiencies of datasets, training pipelines, and resource constraints. Due to its practical usefulness, PTQ has been actively researched recently. In the literature of diffusion models, [18] introduced a dedicated 8-bit post-training quantization method, demonstrating high-fidelity image generation performance.
On the other hand, while most QAT and PTQ algorithms focus on static quantization, recent research has highlighted the benefits of input-dependent dynamic quantization [31, 32, 33, 34]. Dynamic quantization enables the adjustment of quantization intervals based on the varying input-dependent distribution of activations, which has the potential to reduce quantization error. However, implementing these dynamic approaches often incur additional costs for extracting statistical information from activations, making it challenging to achieve performance improvements in practice. For instance, [35] pointed out that input-dependent activation functions can cause significant latency increases despite having relatively little computation.
While a number of previous works proposed different methods to accelerate sampling process of diffusion models, there have been limited works tried to exploit dynamic nature of diffusion models. In this paper, we propose a novel quantization scheme for diffusion models which minimizes activation quantization errors by enabling the generation of suitable quantization interval based on the time step information, all without incurring additional inference costs.
3 Temporal Dynamic Quantization
3.1 Quantization Methods
Before elaborating the proposed method, we define the quantization function used in our work. In this study, we focus on  $b$ -bit linear quantization, where the  $2^{b}$  possible quantization levels are evenly spaced. Linear quantization involves two key hyperparameters: the quantization interval  $s$  and the zero offset  $z$ . Given the full-precision data  $x$ , the quantized data  $\hat{x}$  can be calculated as follows:
$$
\bar {x} = \operatorname {c l i p} (\lfloor x / s \rceil + z, n, p), \quad \hat {x} = s \cdot (\bar {x} - z), \tag {5}
$$
Figure 2: Evolution of Activation Distribution of Diffusion Model over Time Steps. Boxplot depicting the maximum and minimum values of activation for the DDIM model trained on the CIFAR-10 dataset across different time steps.
(a) The interval is calibrated at  $\mathrm{T} = 0$
(b) The interval is calibrated at  $T = 100$
Figure 3: Limitation of Static Quantization for Diffusion Models. Assume that the activation distribution enlarges as time step progresses. (a) Large truncation error due to small quantization interval and (b) Large rounding error due to large quantization interval.
where  $n(p)$  is the smallest(largest) quantization index,  $\text{clip}(\cdot)$  is a clipping function and  $\lfloor \cdot \rfloor$  is a rounding function. For practical acceleration purposes, we utilize symmetric quantization for weights, where  $z = 0$  and  $p = -n = 2^{(b-1)} - 1$ . On the other hand, for activation quantization, we assume asymmetric quantization with  $z \neq 0$ ,  $n = 0$ , and  $p = 2^b - 1$ , which is a commonly adopted approach [36].
3.2 Challenges of Diffusion Model Quantization
The biggest challenge of quantizing diffusion models is to find optimal quantization parameters ( $s$  and  $z$ ) for activation that minimize the quantization error. As shown in Fig. 2, activation distribution of diffusion models have a unique property due to the iterative denoising process, whose distribution highly varies depending on the time step ( $t$ ) regardless of the layer index.
Therefore, using static values for quantization parameters causes significant quantization error for different time steps, as shown in Fig. 3. Previous studies [37, 18] has also reported the dynamic property of activations in diffusion models and attempted to address it by sampling the calibration dataset across overall time frames. However, despite these efforts, these studies still relied on static parameters, resulting in sub-optimal convergence of minimizing quantization error. To tackle this problem fundamentally, it is crucial to enable the update of quantization parameters considering the dynamics in the input activation distribution.
3.3 Implementation of TDQ Module
To address the rapid changes in input activation, one easy way of think is input-dependent dynamic quantization. Previous studies have demonstrated that incorporating a quantization module that generates the quantization parameters based on input features such as minimum, maximum, and average values can lead to significant improvements in accuracy for specific applications [38, 32]. According to our observation, as will be presented in the Appendix, our own implementation of input-dependent dynamic quantization offers notable quality improvement. However, the process of gathering statistics on the given activation introduces complex implementation and notable overhead [35]. Despite the potential of quality improvement, this approach may not be an attractive solution in practice.
(a) Dynamic Quantization
(b) Ours : Temporal Dynamic Quantization
Figure 4: Overview of TDQ module. Comparison between (a) Input-dependent dynamic quantization, which requires activation statistics during both training and inference, and (b) the proposed TDQ module that enables cost-free inference based on pre-computed intervals.
Instead, we propose a novel dynamic quantization that utilizes temporal information rather than input activation. One important observation we made is that when we measured the Pearson correlation between time steps and per-tensor activation variation,  $62.1\%$  of the layers exhibited moderate temporal dependence  $(|\mathsf{r}| > 0.5)$ , and  $38.8\%$  exhibited strong temporal dependence  $(|\mathsf{r}| > 0.7)$ . This implies that many layers in diffusion models indeed possess strong temporal dependence. Although the activation distribution might change based on the input data, the overall trends remain similar within the same time frame (see Fig 2). Therefore, we can determine the optimal interval based on the differences in time steps, leading to more reliable and robust quantization results
Intuitively, we can employ a scalar quantization interval for each time step. However, this strategy falls short in capturing inter-temporal relationships and cannot be applied to scenarios where different time steps are used between training and inference. Instead, we propose using a learnable small network that predicts quantization interval by taking time step information. This strategy can evaluate activation alterations collectively across multiple timesteps, thereby making the learning process more stable and improving performance, as will be shown in Table 3.
Fig. 4 (b) presents the overview of our idea, called TDQ module. In the TDQ module, the dynamic interval  $\tilde{s}$  is generated based on the provided time step  $t$  as follows:
$$
I = \operatorname {e n c} (t), \quad \tilde {s} = f (I), \tag {6}
$$
where  $\text{enc}(\cdot)$  represents the encoding function of the time step, which will be described in Section 3.4. Here,  $I$  is the encoded feature, and  $f(\cdot)$  signifies the generator module function.
In our approach, the TDQ module is attached to each quantization operator, and the independent dynamic quantization interval is generated based on the given time step. As illustrated in the Fig. 4, the generator is implemented by simply stacking multiple linear layers with a softplus function to constrain the data range to non-negative value. Please note that all components of the generator are differentiable. As a result, during the PTQ or QAT process, the interval can be updated toward minimizing quantization error using gradient descent.
For instance, when employing a well-known QAT algorithm such as LSQ [29], the static interval can be substituted by the output of the TDQ module. The quantization function from Eq. 5 is then modified as follows:
$$
\bar {x} = \operatorname {c l i p} (\lfloor x / \tilde {s} \rceil + z, n, p), \quad \hat {x} = \tilde {s} \cdot (\bar {x} - z). \tag {7}
$$
When we use a straight-through estimator [39], the gradient can propagate through the learnable parameters of the generator. These parameters are then updated via iterative gradient descent, with the goal of minimizing the final task loss. The same pipeline can be applied to PTQ algorithms that utilize local gradient of minimizing reconstruction error [23, 24]. Consequently, the TDQ module is easily applicable to both QAT and PTQ schemes.
However, unlike input-dependent dynamic quantization, the TDQ module enables cost-free inference. After the PTQ or QAT process, the time-dependent interval can be pre-computed offline, and during inference, we can utilize the pre-computed value. In addition, the pre-computed interval provides
a strong advantage by enabling the seamless integration of TDQ module on existing frameworks without modifications.
3.4 Engineering Details
We conducted a series of experiments to enhance the stability and effectiveness of the TDQ module, and figured out that several engineering improvements play a key role in achieving reliable results. This section provides a detailed description of these improvements and their impact on the performance of the TDQ module.
Frequency Encoding of Time Step Feeding the time step directly to the generator results in inferior convergence quality. This is primarily due to the well-known low-frequency inductive bias of neural networks [40, 41]. If the time step is directly input to the generator, it tends to produce an interval that barely changes regardless of the time step. To mitigate this low-frequency bias, we use geometric Fourier encoding for time step [42, 43], as described in Eq. 6.
$$
I = \operatorname {e n c} (t) = (\sin \left(\frac {t}{t _ {\text {m a x}} ^ {0 / d}}\right), \cos \left(\frac {t}{t _ {\text {m a x}} ^ {0 / d}}\right), \sin \left(\frac {t}{t _ {\text {m a x}} ^ {2 / d}}\right), \cos \left(\frac {t}{t _ {\text {m a x}} ^ {2 / d}}\right), \dots , \sin \left(\frac {t}{t _ {\text {m a x}} ^ {d / d}}\right), \cos \left(\frac {t}{t _ {\text {m a x}} ^ {d / d}}\right)), \tag {8}
$$
where  $t$  is the current time step,  $d$  is the dimension of encoding vector, and  $I$  is the encoded vector. This encoding approach allows the TDQ module to accommodate high-frequency dynamics of time step. In this paper, we set  $t_{max}$  as 10000, empirically.
Initialization of TDQ Module Proper initialization of quantization interval is crucial, as incorrect initialization can lead to instability in QAT or PTQ processes. Existing quantization techniques only need to initialize the static step value, but we need to initialize the TDQ module's output to the desired value. To achieve this, we utilize He initialization [44] for the weights and set the bias of the TDQ module (MLP)'s last linear layer to the desired value. Given that the input to the MLP (geometric Fourier encoding) can be treated as a random variable with a mean of zero, the output of the He-initialized MLP will also have a mean of zero. Thus, we can control the mean of the MLP output to a desired value via bias adjustment. After extracting 1000 samples for the entire time step, we initialized the quantization interval to minimize the overall error and then conducted an update to allow adaptation to each time step.
In order to demonstrate the superior performance of TDQ, we conducted tests using two different models: DDIM [2], a pixel space diffusion model, and LDM [4], a latent space diffusion model. For the DDIM experiments, we utilized the CIFAR-10 dataset [45] (32x32), while for the LDM experiments, we employed the LSUN Churches dataset [46] (256x256). This allowed us to showcase the effectiveness of the proposed method in both low and high-resolution image generation scenarios. We applied PTQ and QAT to both models. However, it is important to note that while the latent diffusion model consists of a VAE and a diffusion model, we focused on quantizing the diffusion model and did not perform quantization on the VAE component.
In the absence of prior QAT studies for the diffusion model, we experimented with well-known static quantization methods, i.e., PACT [19], LSQ [29], and NIPQ [30] as baselines. Our idea is integrated on top of LSQ, by replacing the static interval as an output of TDQ module. We also provide additional experiments of TDQ integration with other methods in the Appendix. Per-layer quantization was applied to activations and weights of all convolutional and linear layers, including the activation of the attention layer. The models were trained for 200K iterations on CIFAR-10 and LSUN-churches, with respective batch sizes of 128 and 32. The learning rate schedule was consistent with the full precision model.
In PTQ, we used PTQ4DM [18], a state-of-the-art study, as a baseline for extensive analysis. For a fair comparison, we mirrored the experimental settings of PTQ4DM but modified the activation quantization operator with a TDQ module for dynamic quantization interval generation. Like PTQ4DM, we used per-channel asymmetric quantization for weight and per-tensor asymmetric quantization for activation. The weight quantization range was determined by per-channel minimum/maximum values, and the activation quantization range was trained via gradient descent to minimize blockwise
Table 1: QAT results of diffusion models. N/A represents the failure of convergence.
(CIFAR-10)	Methods	FID ↓ / IS ↑
W8A8	W4A8	W8A4	W4A4	W3A3
DDIM [2]	PACT [19]	18.43 / 7.62	22.09 / 7.55	68.58 / 5.93	51.92 / 6.41	N/A / N/A
LSQ [29]	3.87 / 9.62	4.53 / 9.38	6.2 / 9.56	7.3 / 9.45	7.63 / 9.38
NIPQ [30]	3.91 / 9.49	13.08 / 9.7	6.36 / 9.59	30.73 / 8.94	N/A / N/A
Ours	3.77 / 9.58	4.13 / 9.59	4.56 / 9.64	4.48 / 9.76	6.48 / 9.26
(Churches)	Methods	FID ↓
W8A8	W4A8	W8A4	W4A4	W3A3
LDM [4]	PACT [19]	9.20	9.94	8.59	10.35	12.95
LSQ [29]	N/A	4.92	5.08	5.06	7.21
NIPQ [30]	4.12	7.22	4.68	9.13	N/A
Ours	3.87	4.04	4.86	4.64	6.57
Table 2: PTQ results of diffusion models. FP represents the output of the full-precision checkpoint.
(CIFAR-10)	Methods	FID ↓/IS ↑
W8A8	W8A7	W8A6	FP
DDIM [2]	Min-Max	8.73 / 9	34.61 / 8.28	332.15 / 1.47
PTQ4DM [18]	6.15 / 8.84	6.2 / 8.83	22.43 / 8.77	5.59 / 8.94
Ours	5.99 / 8.85	5.8 / 8.82	5.71 / 8.84
(Churches)	Methods	FID ↓
W8A8	W8A6	W8A5	FP
LDM [4]	Min-Max	4.34	103.15	269.05
PTQ4DM [18]	3.97	4.26	7.06	4.04
Ours	3.89	4.24	4.85
reconstruction loss, as in methods like BRECQ [24]. Quantization was applied to all layers, but in the LDM PTQ experiment, the activation of the attention matrix was not quantized. Following PTQ4DM's approach, we used a calibration set of 5120 samples for PTQ, consisting of 256 images with 20 random time steps selected per image.