Paper: Denoising Diffusion Step-aware Models (ICLR-2024)

Reviewer comment (Reproducibility):
The approach relies on "slimmable networks" and "slimmable counterparts" of standard convolution and normalization. However, these aren't discussed in much detail. I think the paper could benefit from some additional background on what slimmable networks are and how they work, and more details on the particular slimmable network architecture used for this work. (In particular, how many slimmable switches were used, and how do they fit into the U-Net architecture?)

Author response:
**Background of slimmable networks**

Our implementation of slimmable networks draws inspiration from the US-Net [15]. These networks are characterized by their ability to operate at various widths, providing a flexible and universal solution for network scalability. During the training of slimmable networks, we focus on optimizing the smallest, largest, and a selection of randomly sampled middle-sized sub-networks. This approach implicitly enhances the performance of all potential networks within the supernet.

In practical terms, our slimmable UNet adheres to the structure of ADM, but with a significant modification: all convolution layers are replaced by slimmable convolutions. These specialized convolutions are capable of adaptively processing tensors with a varying number of input channels. To accommodate a broader range of sub-networks, we adjusted the group number in the group normalization layer from 32 to 16. Additionally, we chose to omit the Batch Normalization (BN) calibration stage, as proposed in [11], since our diffusion UNet exclusively utilizes GroupNorm. Regarding the sizes of the sub-networks, we offer seven different options, corresponding to 2/8, 3/8, 4/8, 5/8, 6/8, 7/8, and 8/8 of the original ADM's width. This results in a total strategy space of 7^(num_timesteps).

Reviewer comment (Reproducibility):
The objective for the evolutionary search algorithm is also presented at a fairly high level and could use more details. The authors mention using NSGA-II to "balance conflicting objectives", but it's not clear to me what objectives were used. Algorithm 2 suggests that the search objective was a linear combination of FID and FLOPS, but the details of the linear combination weights are not specified.

Author response:
**Details of the search algorithms**

For the search parameters, the process encompasses a total of 10 iterations, with each iteration involving a population of 50, and maintaining a mutation rate of 0.001. The initial generation crucially includes a mix of uniform non-step-aware strategies and some random strategies. This specific approach to initialization and mutation has been empirically found to facilitate easier convergence of the search algorithm. Furthermore, a weight parameter is incorporated, which multiplies the GFLOPs to strike a balance between image quality and computational efficiency. For CIFAR-10, we set the weight parameter to 0.1 to favor higher image quality, while for CelebA, the FLOPs weight parameter is adjusted to 0.25. These parameters were manually selected to ensure that there are no compromises in generation quality.

Reviewer comment (Reproducibility):
Could you provide more details on the slimmable network architecture and on the configuration of the search algorithm, as I discuss in [W1] and [W2]? For the search algorithm, what was the value of \( w_M \) used, and how critical is this choice? (Or, are the FID and FLOPs measurements automatically balanced by NGSA-II somehow?)

Author response:
**Background of slimmable networks**

Our implementation of slimmable networks draws inspiration from the US-Net [15]. These networks are characterized by their ability to operate at various widths, providing a flexible and universal solution for network scalability. During the training of slimmable networks, we focus on optimizing the smallest, largest, and a selection of randomly sampled middle-sized sub-networks. This approach implicitly enhances the performance of all potential networks within the supernet.

In practical terms, our slimmable UNet adheres to the structure of ADM, but with a significant modification: all convolution layers are replaced by slimmable convolutions. These specialized convolutions are capable of adaptively processing tensors with a varying number of input channels. To accommodate a broader range of sub-networks, we adjusted the group number in the group normalization layer from 32 to 16. Additionally, we chose to omit the Batch Normalization (BN) calibration stage, as proposed in [11], since our diffusion UNet exclusively utilizes GroupNorm. Regarding the sizes of the sub-networks, we offer seven different options, corresponding to 2/8, 3/8, 4/8, 5/8, 6/8, 7/8, and 8/8 of the original ADM's width. This results in a total strategy space of 7^(num_timesteps).

**Details of the search algorithms**

For the search parameters, the process encompasses a total of 10 iterations, with each iteration involving a population of 50, and maintaining a mutation rate of 0.001. The initial generation crucially includes a mix of uniform non-step-aware strategies and some random strategies. This specific approach to initialization and mutation has been empirically found to facilitate easier convergence of the search algorithm. Furthermore, a weight parameter is incorporated, which multiplies the GFLOPs to strike a balance between image quality and computational efficiency. For CIFAR-10, we set the weight parameter to 0.1 to favor higher image quality, while for CelebA, the FLOPs weight parameter is adjusted to 0.25. These parameters were manually selected to ensure that there are no compromises in generation quality.

Reviewer comment (Presentation):
Figure 4a and Figure A are quite interesting, but they also look somewhat blurry. I can sort of see some fuzzy "bands" at different points of the trajectory; are these actual changes in the slimmed model size or are these some sort of compression artifact in the figure image? Also, what is the resolution of the X axis here, e.g. where are the boundaries between different steps? It's hard to understand exactly what the plot is showing, and I think this plot would be more readable if it were presented as a line graph (instead of as a heatmap), since its only 1 dimensional.

Author response:
**Elaboration on the search result**

For visual propose, I smoothed the color binds to present a general trend instead of the details. The fuzzy area are some mutation points of the search result. They are actual changes in model size. I think your suggestion of drawing line graph is very useful. Therefore, I plot them in our modified version in the **Appendix Section D**.

Thanks for your remind. In this revised version, we have added experiments, discussion, and detailed content to the original paper. The newly added and modified content is highlighted in red for easy identification. Additionally, we have included the line graph in the main content of our paper for clearer illustration.

Regarding the observed fluctuations in our search results, we attribute these to uncertain timesteps of the diffusion process. In multiple CIFAR-10 experiment runs, the fluctuation area consistently appears within the 400-600 step range, and the internal patterns of these fluctuation areas show little similarity across different runs. Despite this, the overall trend in all the searched strategies demonstrates a similar behavior of increased reliance on larger models as the timestep approaches T→0.

Based on these observations, we hypothesize that: 1) there are some uncertain timesteps that are challenging for our search algorithm to accurately pinpoint, and 2) our current search strategy may not have reached an 'optimal' point. Nevertheless, this does not largely weaken our method, since our results have clearly demonstrated that our approach is significantly more effective than the original non-step-aware diffusion. Moreover, pursuing an extremely optimal result would entail substantially higher computational costs, we believe that our current strategy has achieved a favorable balance between efficiency and performance.

Reviewer comment (Experiments):
The paper lacks experiments on higher resolution data. Can the authors provide results on $\geq 128$ resolution images?

Author response:
Hi. Thanks for your comments. We would like to update our empirical results on CelebA-HQ 128x128 and conditional CIFAR-10. 

Our experiments with CelebA-HQ at a resolution of 128x128 have demonstrated the efficacy of our DDSM with higher resolution data. In this study, we trained a new slimmable supernet for 128x128 resolution data from scratch. We then applied the strategy initially searched on the 64x64 CelebA-HQ dataset to this supernet. The results were promising: DDSM achieved a 61% reduction in FLOPs while preserving the quality of generation. 

Similarly, our experiments on conditional CIFAR-10 showcased DDSM's adaptability to guided diffusion. Here, we constructed a novel slimmable supernet, incorporating class labels as conditions. During inference, we leveraged classifier guidance. By applying the strategy devised for the unconditional CIFAR-10 to this supernet, we maintained the same level of FLOPs efficiency and preserved the FID score.

These experiments underscore the robustness and versatility of our DDSM. They also show the good transferability of DDSM across different settings of the same dataset. To save computational cost, one could search on a small proxy dataset (for example, low-resolution data) and then apply the result to high-resolution data. Moreover, it's worth noting that the performance of DDSM on each dataset could potentially be further enhanced by restarting the search process individually for them.

|  | CelebA-HQ 128x128 |  | conditional CIFAR-10 |  |
| --- | --- | --- | --- | --- |
|  | ADM | +DDSM | ADM | +DDSM |
| FID | 7.53 | 7.71 | 2.48 | 2.52 |
| GLOPs | 194.00 | 76.18 | 12.14 | 6.20 |

Reviewer comment (Presentation):
Figure 3 is unclear. What do networks with mixed colors (blue and green) mean?

Author response:
We apologize for any confusion caused and have now provided a more detailed description in the caption of Figure 3 for clarity. In this figure, the color green represents the activated weights within the network, whereas blue represents the inactive weights. The varying proportions of blue and green in different networks indicate networks of various sizes. The supernet, which has all of its weights activated, is depicted entirely in green.

During the training phase, our focus is on the supernet, which we train by optimizing its various sub-networks. This step is crucial for establishing a foundation of diverse network sizes and capabilities. In the subsequent searching stage, we creatively combine these sub-networks into different step-aware strategies. The objective here is to explore and identify the most optimal strategy, one that effectively balances performance and efficiency.

Reviewer comment (Experiments):
Is DDSM compatible with guided diffusion? Can the authors provide some demonstrations?

Author response:
Hi. Thanks for your comments. We would like to update our empirical results on CelebA-HQ 128x128 and conditional CIFAR-10. 

Our experiments with CelebA-HQ at a resolution of 128x128 have demonstrated the efficacy of our DDSM with higher resolution data. In this study, we trained a new slimmable supernet for 128x128 resolution data from scratch. We then applied the strategy initially searched on the 64x64 CelebA-HQ dataset to this supernet. The results were promising: DDSM achieved a 61% reduction in FLOPs while preserving the quality of generation. 

Similarly, our experiments on conditional CIFAR-10 showcased DDSM's adaptability to guided diffusion. Here, we constructed a novel slimmable supernet, incorporating class labels as conditions. During inference, we leveraged classifier guidance. By applying the strategy devised for the unconditional CIFAR-10 to this supernet, we maintained the same level of FLOPs efficiency and preserved the FID score.

These experiments underscore the robustness and versatility of our DDSM. They also show the good transferability of DDSM across different settings of the same dataset. To save computational cost, one could search on a small proxy dataset (for example, low-resolution data) and then apply the result to high-resolution data. Moreover, it's worth noting that the performance of DDSM on each dataset could potentially be further enhanced by restarting the search process individually for them.

|  | CelebA-HQ 128x128 |  | conditional CIFAR-10 |  |
| --- | --- | --- | --- | --- |
|  | ADM | +DDSM | ADM | +DDSM |
| FID | 7.53 | 7.71 | 2.48 | 2.52 |
| GLOPs | 194.00 | 76.18 | 12.14 | 6.20 |

Reviewer comment (Novelty):
Using a different network at different timestep has been explored before, such as in e-diff-i.

Author response:
**Discussion on EDiff-I**

While eDiff-I [6] adopts multiple expert models at each step to improve image quality, it does not enhance efficiency, as all these expert models are uniformly sized. Our method distinguishes itself by utilizing models of varying sizes for acceleration purposes. Moreover, unlike eDiff-I, which manually assigns models to each step, we introduce an evolutionary search approach. This innovation not only automates model allocation but also optimizes the process, potentially yielding more effective outcomes.

[6] eDiff-I: Text-to-Image Diffusion Models with an Ensemble of Expert Denoisers, Balaji et al..

Reviewer comment (Experiments):
The compatibility claims with DDIM and latent diffusion are fairly cursory. More detailed experiments showing accelerated performance combining these methods could better showcase modularity.

Author response:
**Details of combining DDIM and Latent Diffusion**

For the DDIM experiment, we adopted the weights from the 1000-step ADM model and applied the DDIM sampling schedule. As indicated in Table 4, our method significantly boosts DDIM's speed by 48%, 45%, 53%, and 62% for 10, 50, 100, and 1000 steps, respectively. This suggests that while additional steps enhance generation quality, they also introduce excess computational load. Our approach proves increasingly beneficial as the number of steps increases.

In our latent diffusion experiment, we employed the AutoEncoderKL of SD1.4 to transform images into latent vectors, upon which our DDSM was trained. To adapt to the reduced spatial size, we modified the U-Net downsampling from 4 to 2. Our DDSM achieved a 60% acceleration in latent diffusion methods.

Reviewer comment (Experiments):
The evolutionary search itself requires non-trivial compute resources. Analysis of the search costs and scalability could be insightful.

Author response:
**Analysis of the evolutionary search**

A detailed discussion of search costs can be found in **Section B** of our appendix. In deep learning, models like StableDiffusion and Imagen are usually trained once and then used repeatedly. Though these models require substantial training time, they efficiently generate a large number of images upon completion. Our DDSM aims to reduce the inference costs in diffusion models by employing a step-aware pruning strategy, achieving up to 76% faster performance. Despite DDSM incurring higher initial training and search costs (about 2-3 times the training cost of standard models, with search costs similar to training a conventional model), these are one-time expenses. The specific time costs are detailed in Table B.

Reviewer comment (Experiments):
You claim compatibility with methods like DDIM and latent diffusion, but details are limited. Could you provide in-depth quantitative experiments demonstrating accelerated performance when combining DDSM with these existing diffusion acceleration techniques?

Author response:
**Details of combining DDIM and Latent Diffusion**

For the DDIM experiment, we adopted the weights from the 1000-step ADM model and applied the DDIM sampling schedule. As indicated in Table 4, our method significantly boosts DDIM's speed by 48%, 45%, 53%, and 62% for 10, 50, 100, and 1000 steps, respectively. This suggests that while additional steps enhance generation quality, they also introduce excess computational load. Our approach proves increasingly beneficial as the number of steps increases.

In our latent diffusion experiment, we employed the AutoEncoderKL of SD1.4 to transform images into latent vectors, upon which our DDSM was trained. To adapt to the reduced spatial size, we modified the U-Net downsampling from 4 to 2. Our DDSM achieved a 60% acceleration in latent diffusion methods.

Reviewer comment (Evaluation):
The search cost and scalability of the evolutionary algorithm is unclear. Could you analyze the computational requirements of the search procedure and discuss how it scales with factors like dataset size?

Author response:
**Analysis of the evolutionary search**

A detailed discussion of search costs can be found in **Section B** of our appendix. In deep learning, models like StableDiffusion and Imagen are usually trained once and then used repeatedly. Though these models require substantial training time, they efficiently generate a large number of images upon completion. Our DDSM aims to reduce the inference costs in diffusion models by employing a step-aware pruning strategy, achieving up to 76% faster performance. Despite DDSM incurring higher initial training and search costs (about 2-3 times the training cost of standard models, with search costs similar to training a conventional model), these are one-time expenses. The specific time costs are detailed in Table B.