Adversarial Score identity Distillation: Rapidly Surpassing the Teacher in One Step

We would like to thank the reviewer for providing a clear summary of the improvements of SiDA over previous methods. Regarding the weaknesses pointed out, we would like to provide the following clarifications:

W1: We will include a comprehensive comparison on ImageNet 512x512 with a wide range of baselines, including DiT and MAR, in our revision. For instance, using a 481M model with 256 autoregressive steps, MAR (Li et al., 2024) achieves an FID of 2.74 without CFG and using 100 diffusion NFEs, and an FID of 1.73 with CFG and using 100*2 diffusion NFEs. SiD$^2$A on EDM2-S (280M) already outperforms this, achieving an FID of 1.669 without CFG and with only a single generation step. Could you clarify which paper you are referring to for RDM?

W2: We will add ImageNet 512x512 results and detailed comparisons with existing methods in the revised paper. Please refer to our general response for the current results on this point.

W3: The term “One-Step” is defined differently in OSGAN and SiDA. In SiDA, we retain the alternating training between the discriminator (the encoder of the fake score network) and generator, referring to the single forward pass of the generator as “One-Step” generation. OSGAN, however, modifies alternating training into “One-Step” training, eliminating the need for alternation. We will clarify this distinction in the revision.

W4: We commit to open-sourcing both the PyTorch code and model checkpoints. These are not provided at the moment to prevent any accidental violations of ICLR's anonymity requirements.

Q1: A gap exists between the pretrained teacher's estimated score and the true score. The addition of an adversarial loss helps bridge this gap, guiding the generator toward a closer approximation of the true data distribution.

Q2: In SiD, the generator is guided by a biased teacher network, while in SiDA, this bias can be corrected with the help of an adversarial loss that leverages real data. Even when using as little as 5% of the real data, this real-data-based guidance remains highly effective. Please refer to our response to Reviewer kYRQ's Q1 for further details.

We thank the reviewer for providing a clear summary of our paper.

Presentation Clarification

We would like to clarify our statement that “the proposed weight honors the existing weighting mechanisms.” In the training of the generator and fake score networks in SiDA, we adopt the weighting strategies from SiD, which have been proven to be highly effective and stable. These strategies assign weights to diffused images based on their signal-to-noise ratios. When incorporating the adversarial loss, our goal was to minimize disruption to these established weighting mechanisms while leveraging the adversarial loss to correct the bias of the pretrained teacher.

To achieve this, we average the spatial outputs of the encoder across all locations in a batch, scale this average by a constant, and add it to the unweighted SiD loss for the generator and the unweighted $x_0$-prediction denoising loss for the fake score network. This approach has proven effective across four datasets using EDM without dataset-specific tuning and six model sizes (ranging from 125M to 1.5B parameters) when distilling EDM2 on ImageNet 512x512 without model-specific tuning.

The key takeaway in Eqs. (7–10) is that SiDA’s weighting mechanism remains unaffected by the introduction of the batch-averaged adversarial loss. Additionally, Eq. (7) differs from PatchGAN in several significant ways: (1) the input is noise-corrupted and depends on $t$; (2) it is reweighted in Eq. (8) using data-specific weights; and (3) the logits of the discriminator are derived from the output of the encoder component of the diffusion U-Net after channel-wise pooling.

Novelty

As discussed in Related Work, this paper is motivated by the success of Truncated Diffusion, Diffusion-GAN, and SiD. The novel contribution lies in how we combine the SiD loss with the Diffusion-GAN loss without introducing additional parameters beyond those in SiD while preserving its weighting mechanism.

Existing adversarial diffusion distillation methods often rely heavily on their adversarial components for good performance. For instance, Adversarial Diffusion Distillation (ADD) by Sauer et al. (2023) barely outperforms using adversarial loss alone and performs poorly without it. In contrast, SiDA builds upon the already strong baseline of SiD, with adversarial loss added specifically to correct the bias of the pretrained teacher network. This novel insight into correcting the pretrained teacher’s bias using adversarial loss adds significant value to the field and highlights the importance of this work, making it a compelling contribution worthy of publication.

Evaluation and Performance

We will expand our evaluation to include distillation results for EDM2 across six model sizes (125M to 1.5B parameters), pretrained on ImageNet 512x512. Additionally, we will compare SiDA to over 25 baselines that report results on ImageNet 512x512. The results show that SiDA-distilled EDM2 achieves lower FIDs despite being smaller, significantly faster, and not relying on CFG.

Regarding text-to-image experiments, we have observed promising results with SiDA. However, we believe these results merit their own publication, focusing not only on achieving lower FID but also on improving image-text alignment, which is beyond the scope of this paper.

Efficiency

We will provide further comparisons demonstrating that SiDA is only about 10% slower per iteration than SiD. Despite this slight overhead, SiDA achieves significantly faster FID reduction and surpasses the teacher in less time, making it approximately an order of magnitude faster than SiD when considering its efficiency in improving performance.

We appreciate the reviewer's positive feedback on our paper.

• In response to the identified weaknesses, we provide the following clarifications:

  • W1: We’d like to clarify that data-free distillation, which does not require real or teacher-synthesized images, is a distinctive feature of methods like Diff-Instruct and SwiftBrush, which optimize KL divergence under diffusion, and SiD, which optimizes Fisher divergence under diffusion. In scenarios lacking real data, SiD stands as the state-of-the-art choice. Most other distillation methods, including consistency distillation, CTM, and DMD, necessitate access to real or teacher-synthesized data. SiDA not only surpasses the teacher model but also outperforms both data-free and non-data-free distillation methods.

  • W2: We will include results on distilling EDM2 pretrained on ImageNet 512×512 in the revised paper.

  • W3: The revised paper will detail our approach to model parameter selection. Given that SiDA demonstrates robustness to these settings, we believe the provided guidelines are sufficient. However, we are open to conducting additional ablation studies if the reviewer deems them necessary for a deeper understanding of our method.

  • W4: A schematic illustration of the SiDA algorithm will be added to the paper.

• Addressing the reviewer's specific questions:

  • Q1: The VAE latents for the ImageNet (512×512) dataset total approximately 157 GB. In our experiments with SiDA for EDM2-XS, we tested utilizing only 5% of the images (7.8 GB). The Appendix includes a comparison of the FID trajectories between using 5% and the full dataset. The results indicate that, with just 5% of the data, the SiDA framework achieves an initial convergence rate closely following that of the full dataset. However, it reaches a plateau earlier and converges to a higher FID—worse than using the full data but still clearly outperforming the teacher. This finding actually suggests potential applications of SiDA in domain adaptation, which we plan to explore in future research.

  • Q2: We have implemented the suggested changes; please refer to our response in W2.

  • Q3: Please see below for a separate response.

  • Q4: The revised paper includes a description of our parameter selection process. The approach involves adjusting parameters so that the SiD loss and adversarial loss components have comparable scales for both the generator and fake score networks, maintaining values roughly between 1 and 10,000 to prevent underflow or overflow in fp16. Notably, setting the adversarial loss coefficient too low results in SiDA's performance aligning closely with SiD, as the adversarial loss becomes dominated by the SiD loss component.

  • Q5: In our ablation study on CIFAR10 (unconditional), we observed that the current SiDA implementation with $\alpha=1.2$ can diverge. This issue can be mitigated by extending the initial warmup period using the SiD loss alone before introducing the adversarial loss. For SiD$^2$A, which is essentially SiDA with the generator initialized from a much better starting point, the model performs well at $\alpha=1.2$ and remains stable even when $\alpha$ is increased up to 1.5.

We trust these clarifications address your concerns and welcome any further discussion.

We appreciate the reviewer’s recognition of SiDA’s SOTA performance.

1. Can SiDA be used to distill DiT-based models?

   Response: Our research primarily focuses on distilling U-Net-based EDM models, which are well-recognized and benchmarked in diffusion-based generation and distillation. We have now also added results for EDM2, which shows a significantly lower FID on ImageNet 512 x 512 compared to all existing DiT-based models, as detailed in Table 2 of the EDM2 paper and our forthcoming revision.

   Despite this, we acknowledge the versatility of DiT-based architectures. SiDA is not inherently limited to U-Net backbones and shows potential for adapting to transformer-based architectures. For example, in the context of transformer-based diffusion models, we could divide the architecture into “encoder” and “decoder” components. Insights from studies like REPA (Yu, Sihyun, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, and Saining Xie. "Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think." arXiv preprint arXiv:2410.06940 (2024).) suggest that aligning self-supervised features from the encoder, particularly those from middle layers (6th to 8th in a 24-layer model), with a diffusion model could enhance DiT/SiT model training.

   While integrating SiDA with transformer-based models is intriguing, this exploration is beyond the scope of our current research. We plan to pursue this direction in future projects and invite researchers with the necessary computing resources to adapt our SiDA code, which will be open-sourced, to distill DiT-based models and contribute their findings to the community.

2. Can SiDA be applied to high-resolution models?

   Response: Yes, SiDA excels in high-resolution settings, achieving remarkably low FIDs on ImageNet 512x512 without CFG and through one-step generation. We have documented our current results for distilling EDM2 and will provide further updates after revising our paper based on the reviewers' feedback.

3. Why is $L^1(\theta)$ still necessary with an additional discriminator, i.e., why does SiDA still require tuning $\alpha$?

   Response: Omitting $L^1(\theta)$, or setting $\alpha=0$, was feasible in SiD but resulted in unsatisfactory FID scores. With the inclusion of adversarial loss, this configuration in both SiDA and SiD$^2$A now yields satisfactory outcomes, yet still underperforms compared to the recommended settings of $\alpha=1.0$ for SiDA and $\alpha=1.0$ or $1.2$ for SiD$^2$A. This indicates that, despite the support from adversarial loss, maintaining $\alpha>0$ is important to correct the approximation bias in the SiD gradient effectively.

We have completed the revision of the manuscript, incorporating class-conditional random generations of ImageNet 512x512 using the same set of random noise in Figures 16–34. Additionally, we have added Figures 8-10 to illustrate the rapid improvement of SiD, SiDA, and SiD$^2$A in visual quality as the iterations progress.

Dear Reviewers,

Thank you for your valuable feedback, which has guided our revisions to the manuscript. We have incorporated all the revisions outlined in our responses to each review. We will soon further update the paper to include visualizations that demonstrate how image generation rapidly improves its visual quality during the distillation of EDM2 on ImageNet 512x512. Additionally, we will showcase randomly generated images from specific class labels—if there are particular classes you would like to see, please let us know.

We invite you to review our updated paper and our responses. We appreciate any further feedback you may provide.

Thank you,

Authors of SiDA