Accelerate Quantization Aware Training for Diffusion Models with Difficulty-aware Time Allocation

Response

Thank you for the valuable comments and suggestions. We respond to your concerns point by point as follows and are glad to discuss with you.

Q1. The performance of the drop predictor.

Our experiments prove our proposed drop predictor performs well, which can effectively accelerate the QAT training. The principle behind this is that the component is trained under the guidance of the difficulty level, which guarantees a lower drop probability for the difficult step and a higher probability for the easy step. Our reliable designs for training objectives ensure the drop predictor performs well. Also, Figure A2 in the appendix shows the behaviours of the drop predictor in the training process, which shows it performs well.

The uniform drop is an ablation study setting which doesn’t use the drop predictor but sets the drop probability to a constant. So it is not the default setting, which is just used to prove the effectiveness of our proposed components.

Q2. The hyperparameters.

The most important hyperparameter of our framework is M which controls the average drop probability. Actually, we have explored the influence of different M in Table 3. Table. 3 shows that our method is not sensitive to M and can consistently gain faster convergence compared to the baseline. Other hyperparameters such as learning rate is chosen just as the common practice in ML research.

Hope our response can address your concerns. Please feel free to communicate with us if you have any questions.

Response

Thank you for the constructive comments and suggestions. We respond to your concerns point by point as follows and are glad to discuss with you.

W1 & W2. The range of the probability $p_t$ and the clarification about the selection of M.

The reviewer might be concerned about if $p_t$ can exceed 1, and the selection about M. We think the reviewer might misunderstand some details of our method. As shown in Equation 8, we minimize $L_\text{cons }=\left(\sum_{t=1}^Tp_t-M\right)^2$, so it means that the average drop probability of timesteps equals to $M/T$. Therefore, M is set to the product of a constant and T, and neither large M nor large T has any effect on our framework. For example, if we set M=0.6T as in the paper, the average of $p_t$ should equal to 0.6, which is a valid probability belonging to [0,1]. Moreover, we also use a simple clip operation to ensure each probability $p_t$ lies in a valid range, i.e., [0,1]. According to our experiments (Figure A2), these probabilities are indeed in the legal range, even without the clip operation.

Besides, the ablation study in Table 3 proves that our method is not sensitive to M and can be consistently effective, and M=0.6T performs well in all our experimental settings.

W3 About the baselines.

We used EfficientDM as the main baseline because we could not find any other baselines. EfficientDM extended QAT to diffusion models for the first time, but at present no one else proposed new diffusion model QAT methods. We believe this is not because the diffusion QAT is not important, but because the technologies ( LSQ + Knowledge Distillation ) used by EfficientDM are mature and perform well. Therefore, we believe validating our method based on EfficientDM can demonstrate the significance and superiority of our method.

Q1

Please refer to W1 to W3

Q2. How to determine the convergence iteration.

Sorry for the confusion. We clarify that our procedure to determine the convergence iteration is that: For both baseline and our method, we test the model for every 100 iterations. We train the model until the performance deteriorates twice in a row. For example, if the FID scores of three consecutive evaluations are 10.91, 10.92 and 11.00, we will stop training, and choose the best iteration so far as the convergence iteration. We make this clear in the revised paper, marked by blue (Line 347).

Hope our response can address your concerns. Please feel free to communicate with us if you have any questions.

Thank you for your response.

I see your method use $p_t$ to select step in the training course. However, such a selection method is heuristic. The selection of one step seems to be independent with other steps, which is weird.

I also read comments from other reviewers and understand that this paper seems to focus on speeding up training. In this perspective, the authors need to train a model from scratch to provide a rigorous evaluation. In the current version, the authors only finetuned a pretrained stable diffusion. Such an evaluation only provides some behaviors of their method.

Since their method is mostly heuristic and the current paper is unclear about the main focus (due to lack of comprehensive evaluation), I maintain my current score.

Thank you for your reply! We think there may be three points that need clarification.

1. Our task cannot train from scratch. The reviewer recommends us to train the model from scratch, however, our task (Quantization-aware Training, QAT) must train the quantization parameters and model weights from a pretrained checkpoint[1,2]. We aim to speed up QAT rather than the usual diffusion training.
2. We use $p_t$ to select steps, where $p_t$ is the output of the predictor taking t as the input. [3] have demonstrated this kind of predictor can capture temporal information well. The selection of one step is not independent with other steps, because the predictor the trained with various t and can capture temporal relations between steps.
3. Our method is not heuristic. As presented in the paper, $p_t$ is dynamically optimized according to the difficulty coefficient $w_t$, where $w_t$ is calculated based on the task loss values in the current sliding window (Eq.7). The overall framework is adaptive and dynamic, rather than simply heuristic.

References

[1] Learned step size quantization. ICLR’19

[2] EfficientDm: Efficient quantization-aware fine-tuning of low-bit diffusion models. ICLR’23

[3] Temporal dynamic quantization for diffusion models. NIPS’23

Response

Thank you for the constructive comments and suggestions. We respond to your concerns point by point as follows and are glad to discuss with you.

W1. Clarification about our contributions.

We’d like to clarify that our contributions are two-fold:(i) We propose to automatically and dynamically identify the difficulties of training with QAT. (2) We propose a drop mechanism to achieve fast QAT training based on the dynamic difficulty. Although previous works[1,3] (the reviewer mentioned) and [8,9] (which we have compared with in the section 4.4) have explored accelerating diffusion training based on difficulty, they are either static or heuristic, while our method can dynamically and automatically measure the difficulty of different stages in the training process. Dynamically measuring the difficulty is important, because the difficulty varies with the current training state. Figure A2 can prove this, where the drop probability of timestep=201 decreases at first and then increases. Moreover, actually we have done experiments that compare our method to previous works (i.e., P2-weighting [8] and ANT-UW [9]) that identify the difficulty, as shown in Table 4. We also add a comparison with the suppression sampling in [3] (implement it in the QAT framework). We summarize the comparison to these difficulty-aware methods [3,8,9], emphasizing the superiority and contribution of our method.

W2. Comparison to fast sampling methods[4,5,6,7].

The reviewer hopes us compare our method to these fast sampling methods[4,5,6,7], however, we are very confused about why it is required. Our task is to accelerate the training of QAR, but [4,5,6,7] aim to accelerate the sampling process and don’t involve training. We guess the reviewer may hope us to compare to fast training methods, which we present in W1. Besides, our method is easy to implement and not heavy, which only requires less than 40 lines of code.

W3. The rationale behind allocating more training time.

• After identifying the difficulty, just increasing the learning rate will result in unfairness, so it is not reasonable. Specifically, in deep learning, we usually don't know the best learning rate unless we try it one by one. We can’t ensure assigning a higher learning rate to some timesteps can make them better. So we set the learning rate fixed, which is also used by previous works that the reviewer mentioned [1,3].
• The rationale behind allocating more training time. In contrast, we choose to allocate more training time to difficult timesteps. Naturally, difficult timesteps require more training time compared to easy ones. To speed up the overall convergence, our method avoids wasting training time on easy timesteps that don’t need so much time to reach convergence. For example, assuming that there are 3 different difficulty levels which require training time of 1, 2, 4 hours to achieve convergence. The baseline needs 4x3=12 hours totally (allocating the same training time for each) , while ideally our method only needs 1+2+4=7 hours by adaptive time allocation. To visualize this, Figure A5 in the revised paper compares the training time allocation between our method and the baseline.

W4. Experiments with more models.

In the appendix (Section A.6), we have reported the results with the model DiT-XL/2. In the rebuttal, we add the experiments for EDM[10]. We validate our method on FFHQ dataset using the EDM model from its official repository. The results with bitwidth W4/A8 are presented in the following table.

The reduction ratio of EDM is lower than SD because the EDM model is relatively smaller, which means the costs of drop predictor are relatively larger. But we believe we can improve it by using a smaller predictor.

Hope our response can address your concerns. Please feel free to communicate with us if you have any questions.

Response

Thank you for the constructive comments and suggestions. We respond to your concerns point by point as follows and are glad to discuss with you.

W1. Details about the Coefficient of Variation.

The Coefficient of Variation for timestep $t$ is calculated based on the latest sliding window of task losses at timestep t, as shown in Eq.(7). The task loss is defined in Eq.(5), which is the MSE loss between the original noise prediction and the quantized noise prediction.

In the experiments, we set the window size to 5, as wrote in the Line 346. Experiments show it works well for various models, sampling strategies, and quantization settings.

W2. Analysis versus uniform allocation, and evaluating the effectiveness of the timestep drop mechanism independently.

• Analysis versus uniform allocation. Firstly, we clarify the difference between our difficulty-aware time allocation and the uniform allocation. Specifically, difficult timesteps require more training time compared to easy ones. To speed up the overall convergence, our method avoids wasting training time on easy timesteps that don’t need so much time to reach convergence . For example, assuming that there are 3 different difficulty levels which require training time of 1, 2, 4 hours to achieve convergence. The uniform allocation needs 4x3=12 hours totally (allocating the same training time for each) , while ideally our method only needs 1+2+4=7 hours by adaptive time allocation. Moreover, to visualize this, we use Figure A5 in the revised paper to compare the training time allocation between our method and the uniform allocation, which proves the superiority of our method.
• Evaluating the effectiveness of the timestep drop mechanism independently. In the ablation study (Section 4.3), our adaptive drop mechanism outperforms the uniform drop and heuristic drop, which we believe can prove the effectiveness and rationality of our proposed drop mechanism.

W3. Expand the description of difficulty.

In the original paper, we define difficult timesteps as those that are oscillatory and difficult to converge, while easy timesteps are smooth and easy to converge. To quantitatively measure the difficulty level, we propose a difficulty coefficient $w_t$ , as defined in Eq.(7). Therefore, timesteps with higher $w_t$ are difficult, and those with lower $w_t$ are easy.

W4 & Q2. The analysis of the computational efficiency gains achieved through the proposed method. How the timestep drop mechanism affects the overall convergence speed.

We think this question can be answered with W2's response. Our method avoids wasting training time on easy timesteps that don’t need so much time to reach convergence, so it can gain efficiency. Moreover, we recommend the reviewer to refer to Figure A5, which visualizes how our method saves training time by time allocation versus the baseline.

Q1. Details about the frequency and method of updating the difficulty metrics during training.

To ensure that the difficulty assessment of timesteps remains accurate throughout the training process, our method adopts a dynamic updating mechanism as presented in the paper. Specifically, “dynamic” means that after every back-propagation, we add the task loss gained from this to the sliding window, and remove the oldest loss value. Then the $w_t$ is calculated using the losses in the updated sliding windows. We simply use a FIFO queue to implement the process. We clarify this in the revised paper, marked by orange (Line 291).

Q3. Insights on potential trade-offs in model quality.

We may not understand what the reviewer means by “tradeoff”. We’d appreciate it if you could make it clear. By far, We don't seem to observe an obvious trade-off in our framework.

Q4. Findings on how the choice of bit-width impacts both the training time and the final model performance.

Frankly, according to our experiments, we find no obvious correlation between the bit width and the final model performance. But bit-width seems to affect the overall training time. Low bit-width makes the training more difficult and therefore tends to lengthen the training time.

Hope our response can address your concerns. Please be free to communicate with us if have any questions.