DeeDiff: Dynamic Uncertainty-Aware Early Exiting for Accelerating Diffusion Model Generation

Thank you very much for your detailed and constructive review.

Response to weakness 1:

Thank you for providing more related recent work on our topic. We will cite this work in our revised version.

Indeed, the ICML-23 workshop[1] paper shares similar ideas with us. However, there are a few key points that set us apart from them. First of all, the exiting strategy in [1] is static and pre-defined by hand and the model requires finetune under each strategy, which in fact introduces more cost and instability to achieve acceleration. Instead, our DeeDiff learns dynamic exiting strategy and the performance and efficiency trade-off are simply controlled by the threshold. Second, our method achieves better performance with higher efficiency compared with [1]. For example, [1] obtains 4.98 FID with 32.09% acceleration ratio while our method gains 3.9 FID with 46.2% acceleration ratio.

Response to weakness 2:

We apologize for the lack of detailed descriptions. We will provide more details for BERTxiT and CALM in the revised version. In short, BERTxiT utilizes a learning strategy to extend early exiting to BERT models and apply average layer-wise loss ($L = \frac{1}{N} \sum_{i=1}^N L_i$, N is the number of layers) to train the network while CALM apply decay layer-wise loss ($L = \frac{i}{N} \sum_{i=1}^N L_i$). Furthermore, CALM uses the similarity of adjacent layers and confidence to decide to exit and calibrates local early exits from global constraints. In our experiments, we follow their training strategy and we only apply similarity to decide exiting for CALM since confidence-based exiting is hard to be applied to diffusion models. During training, we chose the best evaluation epoch of BERTxiT and CALM for a fair comparison.

As for the comparison with S-Pruning[2], the GFLOPs we reported are not synonymous with Multiply-Accumulate Operations (MACs). Typically, 1 MACs roughly equals 2 GFLOPs which means that our backbone has slightly less computation cost than S-Pruning (11.97 vs 12.2). We apologize for our fault in the sampling step. We would like to provide the performance of our method on DDIM 100 sampling steps as follows:

Response to weakness 3:

The answer is partly yes. At the beginning of training, the loss will focus more on early stages. However, with further optimization, the loss will gradually focus on the latter steps, achieving globally optimal.

Thank you very much for pointing out our minor mistake. We will fix it in our revised version.

Response to weakness 4:

We apologize for the faults of citation and thank reviewer for your detailed comments. We will update the citation in the revised version.

[1] Moon, Taehong, et al. "Early Exiting for Accelerated Inference in Diffusion Models." ICML 2023 Workshop on Structured Probabilistic Inference & Generative Modeling. 2023.

[2] Fang, Gongfan, Xinyin Ma, and Xinchao Wang. "Structural Pruning for Diffusion Models." arXiv preprint arXiv:2305.10924 (2023).

Thank you very much for your detailed and constructive review.

Response to weakness 1:

The running time is unstable since it is influenced by hardware environments such as memory and I/O. Therefore, we provide GFLOPs to show the theoretical computation cost for fair comparison. The extra computation cost inducted by the UEM is negligible compared with the acceleration brought by our methods.

We apologize for the lack of text-to-image generated results and text-image alignment. We will add images and alignment metrics in the revised version.

Response to weakness 2:

We apologize for the unclearness of our notations. The grayness refers to the uncertainty value, where the grayer points mean higher uncertainty. The uncertainty map is generated from UEM. During inference, we will compute the average uncertainty at the image level and utilize it as a proxy to make exiting decisions.

In notations, $g_i$ is in fact the output layer used in U-ViT to reshape transformer tokens back to images, which are the final output for generation. $w_t$ and $b_t$ are the main learned parameters of UEM. Therefore, $g_i$ is not fixed during training. Also, the training objective is to minimize the MSE of $u_{i,t}$ and $\hat{u}_{i,t}$ since we hope to optimize $w_t$ and $b_t$.

Thank you a lot for your detailed and constructive review.

Response to weakness 1:

Our method does not significantly increase the training cost. On the one hand, UEM is a basic linear layer that entails minimal computation costs and few parameters. On the other hand, our model utilizes the pretrained weight upon the baseline model which brings benefits to model convergence. We will provide the exact training cost compared with other methods in the revised version. Furthermore, the cost added in inference is negligible compared with the acceleration brought by our methods. Therefore, our method is also able to transfer to large diffusion models such as stable diffusion.

Response to weakness 2:

Thank you very much for your detailed advice. We provide the performance of the models that exit with fixed layers as follows:

In the table, 'Exit at 7 layer' shares similar GFLOPs with our dynamic method. The table shows that our dynamic method gains better performance than the static exiting method, demonstrating the effectiveness of our proposed dynamic exiting strategy.

Thank you very much for your detailed and constructive review.

Response to weakness 1:

UEM should not be dismissed as a mere pruning technique. UEM is able to learn the uncertainty per sample. We sincerely invite reviewers to refer to Figure 4. The gray images are generated by UEM, where the grayer points mean higher uncertainty. It shows that UEM can learn unique features for each sample such as the face structure shown in Figure 4. Moreover, in order to distinguish our method from simple pruning, we will provide the statistics of samples using different average layers in our revised version (Please check the appendix).

Response to weakness 2 and 3:

We would like to respectfully clarify the reviewer's misunderstanding of UEM and the uncertainty-aware layer-wise loss. Our proposed UEM and loss do not alter the structure of the base model during training. The experiments of ours w/o EE follow the same experimental settings (number of layers) as the baseline model. Furthermore, previous works that apply simple layer-wise loss generally meet performance drops after training because of hard optimization. For example, in MuE[1], please compare the results of the first line in Table 1 with the first line in Table 2. The performance drops slightly with simple layer-wise loss while our proposed uncertainty-aware layer-wise loss brings extra benefits to performance on CIFAR-10 and CelebA. We also report the ours w/o EE performance on COCO and ImageNet.

Our method achieves superior performance compared to other early-exiting techniques, without exiting during inference. We believe that the extent of additional benefits our method can provide is dependent on the size of the datasets.

Summary response to weakness 1, 2, 3:

While we are happy to further discuss the added benefits of our method, we would like to emphasize its primary contribution - that DeeDiff achieves the best efficiency while minimizing performance loss. The additional benefits provide us with a deeper understanding of diffusion models, which we will discuss more in the revised version.

Response to weakness 4:

We apologize for the lack of detailed descriptions. We will provide more details for BERTxiT and CALM in the revised version. In short, BERTxiT utilizes a learning strategy to extend early exiting to BERT models and apply average layer-wise loss ($L = \frac{1}{N} \sum_{i=1}^N L_i$, N is the number of layers) to train the network while CALM apply decay layer-wise loss ($L = \frac{i}{N} \sum_{i=1}^N L_i$). Furthermore, CALM uses the similarity of adjacent layers and confidence to decide to exit and calibrates local early exits from global constraints. In our experiments, we follow their training strategy and we only apply similarity to decide exiting for CALM since confidence-based exiting is hard to be applied to diffusion models. During training, we chose the best evaluation epoch of BERTxiT and CALM for a fair comparison.

Response to weakness 5:

Our method is not limited to Transformer architecture. We would like to reviewer to check Table 4 in the appendix. We provide the performance and efficiency of our method applied in CNN-based diffusion models.

Response to weakness 6:

We apologize for the lack of clarity for notations. $L_{i,t}$ is the output features of the $i$-th layer at $t$ timestep. In U-ViT[2], there is an output layer that maps the tokens to noise images, namely $g_i$. We would like to invite the reviewer to refer to figure 1 in U-ViT[2].

Response to weakness 7:

Thank you for your generous suggestion. We report the performance of the small model here:

The small model is trained with 7 layers from scratch. All settings stay unchanged compared with baseline models. Also, we choose the best evaluation epoch for a fair comparison.

[1] Tang S, Wang Y, Kong Z, et al. You Need Multiple Exiting: Dynamic Early Exiting for Accelerating Unified Vision Language Model[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023: 10781-10791.

[2] Bao F, Nie S, Xue K, et al. All are worth words: A vit backbone for diffusion models[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023: 22669-22679.

Dear reviewers and ACs,

The revised version is available now. Thank you all for your support and review. The updated contents are marked as red words. Since the page limit, most updated experiments are added in appendix. We sincerely invite all reviewers to check our revised version. We hope to have more discussions with reviewers. Thanks again!