SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yanjing  $\mathbf{L}\mathbf{i}^{1\dagger \dagger}$ , Sheng  $\mathbf{X}\mathbf{u}^{1\dagger}$ , Xianbin  $\mathbf{C}\mathbf{a}\mathbf{o}^{1*}$ , Xiao  $\mathbf{S}\mathbf{u}^{2*}$ , Baochang Zhang $^{1,3,4}$
<sup>1</sup>Beihang University
$^{3}$ Zhongguancun Laboratory
4 Nanchang Institute of Technology
{yanjingli, shengxu}@buaa.edu.cn
Denoising diffusion generative models are capable of generating high-quality data, but suffers from the computation-costly generation process, due to a iterative noise estimation using full-precision networks. As an intuitive solution, quantization can significantly reduce the computational and memory consumption by low-bit parameters and operations. However, low-bit noise estimation networks in diffusion models (DMs) remain unexplored yet and perform much worse than the full-precision counterparts as observed in our experimental studies. In this paper, we first identify that the bottlenecks of low-bit quantized DMs come from a large distribution oscillation on activations and accumulated quantization error caused by the multi-step denoising process. To address these issues, we first develop a Timestep-aware Quantization (TaQ) method and a Noise-estimating Mimicking (NeM) scheme for low-bit quantized DMs (Q-DM) to effectively eliminate such oscillation and accumulated error respectively, leading to well-performed low-bit DMs. In this way, we propose an efficient Q-DM to calculate low-bit DMs by considering both training and inference process in the same framework. We evaluate our methods on popular DDPM and DDIM models. Extensive experimental results show that our method achieves a much better performance than the prior arts. For example, the 4-bit Q-DM theoretically accelerates the 1000-step DDPM by  $7.8\times$  and achieves a FID score of 5.17, on the unconditional CIFAR-10 dataset.
Denoising diffusion models, also known as score-based generative models [10, 33, 35], have recently shown remarkable success in various generative tasks such as images [10, 35, 22], audio [21], video [31], and graphs [23]. These models have also demonstrated flexibility in downstream tasks, making them attractive for tasks such as super-resolution [26, 7] and image-to-image translation [29]. Compared to Generative Adversarial Networks (GANs) [8], historically considered state-of-the-art, diffusion models have proven to be superior in terms of quality and diversity in most of these tasks and applications. The process of diffusion models involves gradually transforming real data into Gaussian noise, which is then reversed via a denoising process to generate real data [10, 40]. However, such denoising process is time-consuming and involves iterating a neural network for noise estimation over thousands of timesteps, despite producing a significant amount of images. Therefore, researchers are actively working on accelerating this generation process to reduce its long iterative process and high inference cost for sample generation. To achieve this, one pipeline is to
† Equal contribution. * Corresponding author.
$\ddagger$  This work was done during her internship at Shanghai Artificial Intelligence Laboratory.
(a)  $t = 100$
(b)  $t = 400$
Figure 1: Studies on the activation distribution w.r.t. time-step. Per (output) channel activation ranges of the first attention block in diffusion model on different timestep. The boxplot visualizes key statistical measures for each channel, including the minimum and maximum values, the 2nd and 3rd quartiles, and the median.
(c)  $t = 700$
focus on sample trajectory learning, to develop faster sampling strategies [28, 22, 1]. While the other pipeline directly compresses and accelerates the noise estimation networks based on network quantization technology [30], which is particularly suitable for AI chips because of the low-bit parameters and operations. Prior post-training quantization (PTQ) methods [30, 19, 17] on diffusion models (DMs) or other neural networks directly compute quantized parameters based on pre-trained full-precision models, which constrains the model performance to a sub-optimized level without fine-tuning. Furthermore, quantizing DMs based on PTQ methods to ultra-low bits (e.g., 4 bits or lower) is ineffective and suffers from a significant performance reduction.
Differently, quantization-aware training (QAT) [16, 18] methods perform quantization during back propagation and generally achieve a less performance drop with a higher compression rate than PTQ. For instance, QAT has been shown to be effective for CNNs [5, 18] and ViTs [16, 18] and BERT [24]. However, QAT methods for low-bit quantization of diffusion models remain largely unexplored. Therefore, we first build a low-bit quantized DM baseline, a straightforward yet effective solution based on common techniques [5]. Our experimental studies reveal that the severe performance drop of low-bit quantized DMs, such as PTQ [30] and baseline [5], lies in the activation distribution oscillation and quantization error accumulation caused by the denoising process.
As shown in Fig. 1, the output distribution of the noise estimation network at each time step can differ significantly, resulting in activation distribution oscillation. Particularly, the distribution of activation in a specific layer varies significantly across different timesteps during training. We also observe that errors between full-precision activations and quantized activations gradually accumulate across timesteps during the sampling process (inference), making it harder to produce well-performed quantized DMs.
Drawing on the aforementioned insights, we propose a Timestep-aware Quantization (TaQ) method to address the oscillating distribution issue. By smoothing out these fluctuations and introducing more precise scaling factors into activations, we effectively enhance the performance of the low-bit quantized DMs. We further design a new training scheme for quantized DMs, dubbed Noise-estimating Mimicking (NeM), which can reduce the accumulated errors and promote the performance of quantized DMs based on the knowledge of full-precision counterparts. In this way, we achieve a new QAT method for low-bit quantized DM (Q-DM) via incorporating all the explorations (see the overview in Fig. 2). Overall, the contributions of this paper can be summarized as follows:
- To the best of our knowledge, we proposed the first QAT method towards efficient low-bit DMs, dubbed Q-DM, by fully considering both training and inference process in the same framework.
- We introduce a Timestep-aware Quantization (TaQ) method to mitigate activation distribution oscillation caused by the random-sampled timestep in the training process. We develop a Noise-estimating Mimicking (NeM) scheme to reduce accumulated errors, by which the Q-DMs are able to achieve comparable performance as the full-precision counterparts.
Figure 2: Overview of the proposed Q-DM framework. We introduce the timestep-aware quantization in an architecture perspective and a noise imitation training scheme incorporated in the optimization process. From left to right, we respectively show the detailed architecture of single Q-AttnBlock in Q-DM and the training framework of Q-DM.
- Extensive experiments on the CIFAR-10 and ImageNet datasets show that our Q-DM outperforms the baseline and 8-bit PTQ method by a large margin, and achieves comparable performances as the full-precision counterparts with a considerable acceleration rate.
Network Quantization. Quantizing neural networks (QNNs) often possess low-bit  $(1\sim 4$  -bit) weights and activations to accelerate the model inference and save the memory usage. Specifically, ternary weights are introduced to reduce the quantization error in TWN [15]. DoReFa-Net [41] exploits convolution kernels with low bit-width parameters and gradients to accelerate both the training and inference. TTQ [42] uses two full-precision scaling coefficients to quantize the weights to ternary values. Zhuang et al. [43] present a  $2\sim 4$  -bit quantization scheme using a two-stage approach to alternately quantize the weights and activations, which provides an optimal trade-off among memory, efficiency, and performance. Jung et al. [12] parameterize the quantization intervals and obtain their optimal values by directly minimizing the task loss of the network and also the accuracy degeneration with further bit-width reduction. ZeroQ [2] supports both uniform and mixed-precision quantization by optimizing for a distilled dataset, which is engineered to match the statistics of batch normalization across different layers of the network. Xie et al. [39] introduces transfer learning into network quantization to obtain an accurate low-precision model by utilizing the Kullback-Leibler (KL) divergence. PwLQ [6] enables accurate approximation for tensor values that have bell-shaped distributions with long tails and finds the entire range by minimizing the quantization error.
Diffusion Model. The high cost of denoising through networks and the long iterative process make it difficult to implement diffusion models widely. To accelerate diffusion probabilistic models (DMs) [10], previous research has focused on finding shorter sampling trajectories while maintaining DM performance. Wavegrad [3] introduces grid search, which finds an effective trajectory with only six timesteps, but this approach cannot be generalized for longer trajectories due to its exponentially growing time complexity. Watson et al. [38] model the trajectory searching as a dynamic programming problem. Song et al. [34] construct non-Markovian diffusion processes that lead to the same training objective, but whose reverse process can be much faster to sample from. For DMs with continuous timesteps, Song et al. [33, 35] have formulated the DM in the form of an ordinary differential equation (ODE) and improved sampling efficiency by using faster ODE solvers. Jolicoeur-Martineau et al. [11] have introduced an advanced SDE solver to accelerate the reverse process via an adaptively larger sampling rate. Analytic-dpm [1] has estimated variance and KL divergence using the Monte Carlo method and a pretrained score-based model with derived analytic forms that are simplified from the score-function. In addition to those training-free methods, Luhman & Luhman [20] have
compressed the reverse denoising process into a single-step model, while San-Roman et al. [28] has dynamically adjusted the trajectory during inference. However, implementing these methods requires additional training after obtaining a pretrained DM, which makes them less desirable in most situations. In summary, all these DM acceleration methods can be categorized as finding effective sampling trajectories.
Unlike prior works, we demonstrate that diffusion models can be accelerated by compressing the network in each noise estimating iteration, which is orthogonal with the fast sampling methods mentioned above. To the best of our knowledge, this is the first study to explore low-bit quantized diffusion models in a quantization-aware training (QAT) manner.
3 Background and Challenge
Forward process. Let  $\mathbf{x}_0$  be a sample from the data distribution  $\mathbf{x}_0 \sim q(\mathbf{x})$ . A forward diffusion process adds Gaussian noise to the sample for  $T$  times, resulting in a sequence of noisy samples  $\mathbf{x}_1, \dots, \mathbf{x}_T$  as:
$$
q \left(\mathbf {x} _ {t} \mid \mathbf {x} _ {t - 1}\right) = \mathcal {N} \left(\mathbf {x} _ {t}; \sqrt {1 - \beta_ {t}} \mathbf {x} _ {t - 1}, \beta_ {T} \mathbf {I}\right), \tag {1}
$$
where  $\beta_{t} \in (0,1)$  is the variance schedule and controls the strength of the Gaussian noise in each step. The forward diffusion process satisfies the Markov property since each step relies solely on the preceding step. Additionally, as the number of steps increases towards infinity ( $T \to \infty$ ), the final state  $\mathbf{x}_T$  converges to an isotropic Gaussian distribution. A notable property of the forward process is that it admits sampling  $\mathbf{x}_t$  at an arbitrary timestep  $t$  in closed form as:
$$
q \left(\mathbf {x} _ {t} \mid \mathbf {x} _ {0}\right) = \mathcal {N} \left(\mathbf {x} _ {t}; \sqrt {\bar {\alpha} _ {t}} \mathbf {x} _ {0}, (1 - \bar {\alpha} _ {t}) \mathbf {I}\right). \tag {2}
$$
Reverse process. To generate a sample from a Gaussian noise input  $\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$  using diffusion models, the forward process is reversed. However, since the actual reverse conditional distribution  $q(\mathbf{x}_{t-1}|\mathbf{x}_t)$  is unknown, diffusion models use a learned conditional distribution  $p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)$  that approximates the real reverse conditional distribution with a Gaussian distribution. This approximation is expressed as:
$$
p _ {\theta} \left(\mathbf {x} _ {t - 1} \mid \mathbf {x} _ {t}\right) = \mathcal {N} \left(\mathbf {x} _ {t - 1}; \tilde {\boldsymbol {\mu}} _ {\theta , t} \left(\mathbf {x} _ {t}\right), \tilde {\beta} _ {t} \mathbf {I}\right). \tag {3}
$$
By using the re-parameterization trick presented in [10], it becomes possible to derive the mean  $\tilde{\pmb{\mu}}_{\theta ,t}(\mathbf{x}_t)$  and  $\tilde{\beta}_t\mathbf{I}$  as follows:
$$
\tilde {\boldsymbol {\mu}} _ {\theta , t} \left(\mathbf {x} _ {t}\right) = \frac {1}{\sqrt {\alpha_ {t}}} \left(\mathbf {x} _ {t} - \frac {1 - \alpha_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \boldsymbol {\epsilon} _ {\theta , t}\right), \tag {4}
$$
$$
\tilde {\beta} _ {t} = \frac {1 - \bar {\alpha} _ {t - 1}}{1 - \bar {\alpha} _ {t}} \cdot \beta_ {t},
$$
where  $\alpha_{t} = 1 - \beta_{t}$ ,  $\bar{\alpha}_{t} = \prod_{i=1}^{t} \alpha_{i}$  and  $\epsilon_{\theta}$  is a function approximator intended to predict  $\epsilon$  from  $\mathbf{x}_{t}$  [10].
Training. At training time, the goal of optimization is to minimize the negative log-likelihood, i.e.,  $-\log p_{\theta}(\mathbf{x}_0)$ . With variational inference, a lower bound of it could be found, denoted as  $L_{\mathrm{VLB}}$ :
$$
L _ {\mathrm {V L B}} = \mathbb {E} _ {q \left(\mathbf {x} _ {0: T}\right)} \left[ \log \frac {q \left(\mathbf {x} _ {1 : T} \mid \mathbf {x} _ {0}\right)}{p _ {\theta} \left(\mathbf {x} _ {0 : T}\right)} \right] \geq - \log p _ {\theta} \left(\mathbf {x} _ {0}\right). \tag {5}
$$
It is found in [10] that using a simplified loss function to  $L_{\mathrm{VLB}}$  often obtains better performance:
$$
L _ {\text {s i m p l e}} = \mathbb {E} _ {t, \mathbf {x} _ {0}, \boldsymbol {\epsilon} _ {t}} \left[ \| \boldsymbol {\epsilon} _ {t} - \boldsymbol {\epsilon} _ {\theta} \left(\sqrt {\bar {\alpha} _ {t}} \mathbf {x} _ {0} + \sqrt {1 - \bar {\alpha} _ {t}} \boldsymbol {\epsilon} _ {t}, t\right) \| ^ {2} \right]. \tag {6}
$$
Sampling. At inference time, a Gaussian noise tensor  $\mathbf{x}_T$  is sampled and is denoised by repeatedly sampling the reverse distribution  $p_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t)$ .  $\tilde{\mu}_{\theta ,1}(\mathbf{x}_1)$  is taken as the final generation result, with no noise added in the final denoising step.
Figure 3: Input activation distribution of the first Q-AttnBlock in diffusion model on different timestep with a model trained on CIFAR10 [13] by DDPM.
Figure 4: Distance between the outputs of the full-precision model and different bit-width baseline models trained on CIFAR-10 [13] by DDIM with 100 sampling steps.
3.2 Quantization
Given an  $N$ -layer CNN model, we denote its weight set as  $\mathbf{W} = \{\mathbf{w}^n\}_{n=1}^N$  and input feature map set as  $\mathbf{A} = \{\mathbf{a}_{\mathrm{in}}^n\}_{n=1}^N$ . The  $\mathbf{w}^n \in \mathbb{R}^{C_{\mathrm{out}}^n \times C_{\mathrm{in}}^n \times K^n \times K^n}$  and  $\mathbf{a}_{\mathrm{in}}^n \in \mathbb{R}^{C_{\mathrm{in}}^n \times W_{\mathrm{in}}^n \times H_{\mathrm{in}}^n}$  are the convolutional weight and the input feature map in the  $n$ -th layer, where  $C_{\mathrm{in}}^n, C_{\mathrm{out}}^n$  and  $K^n$  respectively stand for input channel number, output channel number and the kernel size. Also,  $W_{\mathrm{in}}^n$  and  $H_{\mathrm{in}}^n$  are the width and height of the feature maps. Then, the convolutional outputs  $\mathbf{a}_{\mathrm{out}}^n$  can be technically formulated as:
$$
\mathbf {a} _ {\text {o u t}} ^ {n} = \mathbf {w} ^ {n} \otimes \mathbf {a} _ {\text {i n}} ^ {n}, \tag {7}
$$
where  $\otimes$  represents the convolution operation. Herein, we omit the non-linear function for simplicity. Quantized neural network intends to represent  $\mathbf{w}^n$  and  $\mathbf{a}^n$  in a low-bit format such that the float-point convolutional outputs can be approximated as:
$$
\hat {\mathbf {w}} ^ {n} = \boldsymbol {s} ^ {w ^ {n}} \circ Q (\mathbf {w} ^ {n}) = \boldsymbol {s} ^ {w ^ {n}} \circ \lfloor \operatorname {c l i p} \left(\mathbf {w} ^ {n} / \boldsymbol {s} ^ {w ^ {n}}, - 2 ^ {b - 1}, 2 ^ {b - 1} - 1\right) \rceil
$$
$$
\hat {\mathbf {a}} _ {\mathrm {i n}} ^ {n} = s ^ {a _ {\mathrm {i n}} ^ {n}} \cdot Q \left(\mathbf {a} _ {\mathrm {i n}} ^ {n}\right) = s ^ {a _ {\mathrm {i n}} ^ {n}} \cdot \left\lfloor \operatorname {c l i p} \left(\mathbf {a} _ {\mathrm {i n}} ^ {n} / s ^ {a _ {\mathrm {i n}} ^ {n}}, - 2 ^ {b - 1}, 2 ^ {b - 1} - 1\right) \right] \tag {8}
$$
$$
\mathbf {a} _ {\text {o u t}} ^ {n} = \hat {\mathbf {a}} _ {\text {i n}} ^ {n} \otimes \hat {\mathbf {w}} ^ {n} \approx s ^ {a _ {\text {i n}} ^ {n}} \cdot \boldsymbol {s} ^ {w ^ {n}} \circ [ Q (\mathbf {w} ^ {n}) \odot Q (\mathbf {a} _ {\text {i n}} ^ {n}) ],
$$
where  $\circ$  denotes the channel-wise multiplication,  $\odot$  denotes the efficient GEMM operations, and  $s^{w^n} = \{s_1^{w^n}, s_2^{w^n}, \dots, s_{C_{\mathrm{out}}^n}^{w^n}\} \in \mathbb{R}_+^{C_{\mathrm{out}}^n}$  is known as the channel-wise scaling factor vector [25] to mitigate the output gap between Eq. (7) and its approximation of Eq. (8). Meanwhile, we use the layer-wise quantization for input activations and the scaling factor of activations  $s_{\mathrm{in}}^{a_{\mathrm{in}}^n} \in \mathbb{R}_+$  is a scalar.
3.3 Challenge Analysis
Here we identify two major challenges on low-bit DMs, specific to the multi-step inference process and random-sampled-step training process of diffusion models. Namely, we investigate on the distribution oscillation of the activations, and the accumulated quantization error resulted from the multi-step denoising process.
Activation distribution oscillation. To understand the distribution change of diffusion models, we investigate the activation distribution, w.r.t. timestep in the training process. Theoretically, if the distribution changes w.r.t. timestep, it would be difficult to implement previous QAT methods. We analyze the overall activation distributions of the noise estimation network, as shown in Fig. 3. We can observe that at different timesteps, the corresponding activation distributions have large discrepancies, e.g., Fig. 3(a) v.s. Fig. 3(b), which makes previous QAT methods [16] in-applicable for multi-timestep models, i.e., diffusion models.
Quantization error accumulation. Quantization of a noise estimation network introduces disturbances to the weights and activations, resulting in errors in each layer's output. Previous studies [4]
have found that these errors tend to accumulate across layers, making it more challenging to quantize deeper neural networks. In the case of diffusion models (DMs), at each time step  $t$ , the input of the model  $(\mathbf{x}_{t-1})$  is obtained from the model's output at the previous time step  $t$ , i.e.,  $\mathbf{x}_t$ . As depicted in Fig. 4, the MSE distance, representing the quantization error of low-bit quantized DMs, exhibits a noticeable growth along with the decrease of sampling timestep. This implies that as the denoising process moves towards later timestep, the accumulation of quantization errors becomes more prominent.
4 The Proposed Q-DM
4.1 Timestep-aware Quantization
To tackle the distribution oscillation in the training process, we first introduce the quantized attention block that efficiently takes into account the timestep. This structure allows for the numerical analysis of activation ranges across different timesteps and mitigates distribution oscillation of low-bit quantized DMs. We recall the quantization in the attention block based on Eq. (8), which is formulated as:
$$
\hat {\mathbf {a}} _ {q} (\mathbf {x} _ {t}, t) = s ^ {a _ {q} (\mathbf {x} _ {t}, t)} \cdot Q (\mathbf {a} _ {q} (\mathbf {x} _ {t}, t)), \quad \hat {\mathbf {a}} _ {k} (\mathbf {x} _ {t}, t) = s ^ {a _ {k} (\mathbf {x} _ {t}, t)} \cdot Q (\mathbf {a} _ {k} (\mathbf {x} _ {t}, t))
$$
$$
\mathbf {A} \left(\mathbf {x} _ {t}, t\right) = \operatorname {s o f t m a x} \left[ \left(\hat {\mathbf {a}} _ {q} \left(\mathbf {x} _ {t}, t\right) \cdot \hat {\mathbf {a}} _ {k} \left(\mathbf {x} _ {t}, t\right) ^ {\top}\right) / \sqrt {d} \right], \tag {9}
$$
$$
\hat {\mathbf {A}} \left(\mathbf {x} _ {t}, t\right) = s ^ {A \left(\mathbf {x} _ {t}, t\right)} \cdot Q \left(\mathbf {A} \left(\mathbf {x} _ {t}, t\right)\right),
$$
$$
\mathbf {a} _ {\mathrm {o u t}} (\mathbf {x} _ {t}, t) = \hat {\mathbf {A}} (\mathbf {x} _ {t}, t) \cdot \hat {\mathbf {a}} _ {v} (\mathbf {x} _ {t}, t) ^ {\top},
$$
where  $\mathbf{A}$  is the attention score.
In the  $i$ -th mini-batch, the timestep is represented as  $\{t_1,\dots ,t_{b_i}\}$ , where  $b_{i}$  is the batch size of the  $i$ -th batch. We denote  $i\in \{1,\dots ,B\}$ , and  $B$  is the number of batch. Therefore, we calculate the timestep-aware distribution divergence for the query activation  $\mathbf{a}_q$  as:
$$
\gamma_ {q; t} = \sum_ {i = 1} ^ {B} \frac {1}{b _ {i}} \sum_ {j = 1} ^ {b _ {i}} \mathbf {a} _ {q} \left(\mathbf {x} _ {t _ {j}}, t _ {j}\right),
$$
$$
\sigma_ {q; t} ^ {2} = \sum_ {i = 1} ^ {B} \frac {1}{b _ {i}} \sum_ {j = 1} ^ {b _ {i}} \left[ \mathbf {a} _ {q} \left(\mathbf {x} _ {t _ {j}}, t _ {j}\right) - \gamma_ {q; t} \right] ^ {2}, \tag {10}
$$
where  $\gamma_{q;t}$  and  $\sigma_{q;t}^2$  are statistical mean and variance of query activation  $\mathbf{a}_q$ . And the calculation of the key activation  $\mathbf{a}_k$  is likewise.
Based on such statistical results, the query and key activations in each specific timestep are smoothed as:
$$
\tilde {\mathbf {a}} _ {q} \left(\mathbf {x} _ {t}, t\right) = \left[ \mathbf {a} _ {q} \left(\mathbf {x} _ {t}, t\right) - \gamma_ {q; t} \right] / \sqrt {\sigma_ {q ; t} ^ {2} + \psi} \tag {11}
$$
$$
\tilde {\mathbf {a}} _ {k} (\mathbf {x} _ {t}, t) = [ \mathbf {a} _ {k} (\mathbf {x} _ {t}, t) - \gamma_ {k; t} ] / \sqrt {\sigma_ {k ; t} ^ {2} + \psi},
$$
where  $\psi$  is constant to avoid 0 denominator. With the above timestep-aware smoothing process, we formulate our timestep-aware quantization as:
$$
\hat {\mathbf {a}} _ {q} \left(\mathbf {x} _ {t}, t\right) = s ^ {a _ {q} \left(\mathbf {x} _ {t}, t\right)} \cdot \operatorname {T a Q} \left(\mathbf {a} _ {q} \left(\mathbf {x} _ {t}, t\right)\right), \quad \hat {\mathbf {a}} _ {k} \left(\mathbf {x} _ {t}, t\right) = s ^ {a _ {k} \left(\mathbf {x} _ {t}, t\right)} \cdot \operatorname {T a Q} \left(\mathbf {a} _ {k} \left(\mathbf {x} _ {t}, t\right)\right)
$$
$$
\mathbf {A} \left(\mathbf {x} _ {t}, t\right) = \operatorname {s o f t m a x} \left[ \left(\hat {\mathbf {a}} _ {q} \left(\mathbf {x} _ {t}, t\right) \cdot \hat {\mathbf {a}} _ {k} \left(\mathbf {x} _ {t}, t\right) ^ {\top}\right) / \sqrt {d} \right], \tag {12}
$$
$$
\hat {\mathbf {A}} (\mathbf {x} _ {t}, t) = s ^ {A (\mathbf {x} _ {t}, t)} \cdot \operatorname {T a Q} (\mathbf {A} (\mathbf {x} _ {t}, t)),
$$
in which  $\mathrm{TaQ}(*) = \lfloor \mathrm{clip}([* - \gamma_{*;t}] / [s^{*}\cdot \sqrt{\sigma_{*;t}^{2} + \psi}], - 2^{b - 1},2^{b - 1} - 1)]$ . The smoothed activations are less sensitive to the random sampled timestep in the trianing process and the timestep-aware quantization, to some extent, dismisses the distribution oscillation phenomenon.