Ensembling Diffusion Models via Adaptive Feature Aggregation

Q2: The paper claim that SDv1.5 failed to understand the number in the prompt, but this phenomenon is hard to reproduce.

A2: The inability to reproduce the results stems from the fact that the models we used are not SDv1.5 but fine-tuned versions, including EpicRealism (ER), MajicMixRealistic (MMR), and RealisticVision (RV). Due to additional fine-tuning with private datasets, these models may exhibit a loss of certain generation capabilities, such as object counting.

[1] SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis

[2] https://blackforestlabs.ai/announcing-black-forest-labs/

[3] https://laion.ai/blog/laion-coco/

Questions

Q1: The detailed time consumption during inference is not provided for comparison.

A1: Based on your comment, we compare the efficiency between our AFA and the baselines, when considering 2, 3, and 6 base models.

Efficiency Comparison on 2 base models:

Efficiency Comparison on 3 base models:

Efficiency Comparison on 6 base models:

As shown in the above tables, when using the same 50 inference steps, ensembling methods (e.g., MagicFusion and AFA) do not have an efficiency advantage over a single base model or merging methods. The number of model parameters, TFLOPs, and inference time all increase linearly with the number of base models.

However, thanks to AFA’s high tolerance for fewer inference steps, it can achieve similar performance with reduced steps. Therefore, under fewer inference steps, AFA achieves efficiency in terms of TFLOPs and inference time comparable to that of a single base models and merging methods.

IR metrics of AFA with varying the number of base models and inference steps:

The details can be found in Appendix F and H of the revised manuscript.

W3: AFA are trained on JourneyDB, but baselines are not fine-tuned on the same dataset.

R3: All the compared baselines are training-free, meaning they cannot be fine-tuned using a text-to-image dataset. However, since JourneyDB is a high-quality dataset, particularly in terms of image quality, we wondered whether AFA would perform well on a dataset with lower image quality. To investigate this, we conducted an experiment by training AFA using LAION-COCO [3], which has lower image quality compared to JourneyDB.

Quantitative comparison in Group I and II trained on JourneyDB and LAION-COCO (showing the average results of Groups I and II):

As shown in the above results, the model trained on LAION-COCO performs similarly to the one trained on JourneyDB, indicating that AFA is not sensitive to the quality of the training dataset. Its consistent performance across datasets of varying quality highlights its robustness and generalization ability.

Regardless of whether the training dataset is high-quality or low-quality, AFA effectively leverages the features of the base models to deliver stable and reliable results.

More details can be found in Appendix M of the revised manuscript.

W2: Experimental details about hardware used, memory consumption and quantitative analysis on time consumption are missing.

R2: Based on your comment, we add details of our used hardware and the memory consumption in Appendix N of the revised manuscript. Specially, We train our AFA in an environment equipped with 8 NVIDIA V100 GPUs, each with 32GB of memory. When ensembling two base models, the GPU memory consumption is about 11GB under FP16 precision and a batch size of 1. When ensembling three base models, the memory consumption increases to 14GB under the same settings. For six base models, the memory consumption reaches about 30GB. Due to the frozen parameters of the base models, training our AFA demands relatively modest GPU memory, even when ensembling a larger number of base models.

To address your concern regarding the efficiency analysis, we compare the efficiency of our AFA when ensembling 2, 3, and 6 base models.

Efficiency Comparison on 2 base models:

Efficiency Comparison on 3 base models:

Efficiency Comparison on 6 base models:

As shown in the above tables, when using the same 50 inference steps, our AFA does not offer an efficiency advantage, as the inference time significantly higher compared to a single base model. Furthermore, the number of model parameters, TFLOPs, and inference time all increase linearly with the number of base models.

However, thanks to AFA’s high tolerance for fewer inference steps, it can achieve similar performance with reduced steps. Therefore, under fewer inference steps, AFA achieves efficiency in terms of TFLOPs and inference time comparable to that of a single base models and merging methods.

IR metrics of AFA with varying the number of base models and inference steps:

The details can be found in Appendix F and H of the revised manuscript.

Thanks for your valuable comments. We will address the concerns below.

Weaknesses

W1: The method fails when considering ensembling different diffusion model architectures.

R1: Based on your comment, we evaluate the generality of our AFA on other diffusion architectures. We select three SDXL [1] models and two FLUX .1 [dev] [2] models from CivitAI.

SDXL consists two U-Nets: the denoising U-Net and the refining U-Net. Both U-Nets shares a similar architecture of SDv1.5, but contain larger blocks. Consistent with the method used for ensembling SDv1.5 models, we employ the SABW module to aggregate features from each block.

Quantitative Comparison of ensembling SDXL models:

FLUX .1 [dev] is a DiT-based model. To ensemble models with this architecture, we also utilize the SABW module to aggregate features from each DiT Transformer block.

Quantitative Comparison of ensembling FLUX .1 [dev] models:

The performance of the ensembled model surpasses that of the individual base models, demonstrating the generality of our AFA on different model architectures. However, the observed improvement is relatively small, likely because the base models already exhibit strong performance, and ensembling leads to only marginal gains. Additionally, due to the large size of the base models, ensembling multiple such models may be less practical.

It is worth noting that our AFA cannot be applied to ensembling models with different architectures, as the misalignment of block features makes block-wise aggregation challenging. We view ensembling models with diverse architectures as a promising direction for future research.

The details can be found in Appendix I of the revised manuscript.

It is weird that SDXL can achieve a 7.3 AES score and a 35.5 HPS score in the Draw Bench Prompts. Actually, I have tested SDXL based on HPSv2 benchmark prompts (including 3200 prompts), which resulted in a 27.7 HPS score and a 6.0 AES score. Do the test prompts cause the differences? What are the AES score and HPS score of the used journeyDB data?

Thanks for your valuable comments. We will address the concerns below.

Weaknesses

W1: AFA, especially SABW, introduces a large degree of additional complexity.

R1: SABW does not significantly increase computational complexity. While SABW introduces additional parameters, the increase is minimal—only about 50M parameters, which is just 1/16 of the base model’s parameter count of approximately 800M. Furthermore, due to the high tolerance for fewer inference steps, AFA can ensemble multiple base models and achieve superior generation performance with fewer steps compared to a single base model. In other words, the computational complexity of AFA remains comparable to that of a single base model when generating images.

Questions

Q1: Did AFA underperform in any specific cases or with certain prompts?

A1: Yes. Specifically, when all the base diffusion models produce poor results for a given prompt, AFA’s performance becomes limited. Since AFA depends on aggregating features from these base models, if they all struggle with a particularly challenging prompt, the aggregated output is unlikely to show significant improvement.

W2: The comparison of computational efficiency is missing.

R2: Based on your comment, we compare the efficiency between our AFA and the baselines, when considering 2, 3, and 6 base models.

Efficiency Comparison on 2 base models:

Efficiency Comparison on 3 base models:

Efficiency Comparison on 6 base models:

As shown in the above tables, when using the same 50 inference steps, ensembling methods (e.g., MagicFusion and AFA) do not have an efficiency advantage over a single base model or merging methods. The number of model parameters, TFLOPs, and inference time all increase linearly with the number of base models.

However, thanks to AFA’s high tolerance for fewer inference steps, it can achieve similar performance with reduced steps. Therefore, under fewer inference steps, AFA achieves efficiency in terms of TFLOPs and inference time comparable to that of a single base models and merging methods.

IR metrics of AFA with varying the number of base models and inference steps:

The details can be found in Appendix F and H of the revised manuscript.

Q2 (1): During the training process of AFA, it appears that intermediate inference features of all base diffusion models at different timesteps are required. Wouldn’t this process be quite time-consuming?

A2 (1): During training, only the parameters of the SABW modules are updated, while the parameters of the base models remain frozen. In each training step, a single timestep is sampled, and AFA learns to aggregate features specific to this timestep. Overall, our AFA has a time-efficient training process.

Q2 (2): Could the authors provide more information on the training time efficiency of AFA compared to other baselines?

A2 (2): The baselines, including the merging-based and MagicFusion, are the training-free methods, making a direct comparison of training efficiency with our AFA infeasible. Specifically, when ensembling three base models to denoise a 512×512 image, the training FLOPs of our AFA amount to 6.42T.

[1] SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis

[2] https://blackforestlabs.ai/announcing-black-forest-labs/

Questions

Q1: Could the authors provide further explanation as to why you chose to randomly split the six diffusion models into two groups of three, rather than integrating all six models directly?

A1: To evaluate the effectiveness of our AFA across different model groups, we randomly selected two groups, each containing three models. The choice of three models for ensembling was arbitrary, and our AFA is fully capable of handling more than three models. To assess the scalability of our AFA, we applied it to ensemble a larger number of base models.

IR metrics of AFA with varying the number of base models and inference steps:

As shown in the results above, ensembling more base models leads more performance improvements. Additionally, increasing the number of the ensembled base models enhances the tolerance for fewer inference steps.

More details can be found in Appendix F of the revised manuscript.

W3: Comparing AFA's results to those of the baselines under the same reduced-step inference would better validate AFA's practical time efficiency.

R3: Based on your suggestion, we compare our AFA with the baseline methods under fewer inference steps.

IR metrics comparison with varying inference steps:

As demonstrated by the results, our AFA not only surpasses the baseline methods at higher inference steps but also achieves significantly better performance with fewer inference steps. This highlights AFA's superior tolerance to the fewer inference steps compared to the baselines, resulting in improved inference efficiency. More details can be seen in Appendix G of the revised manuscript.

W2: A study on the impact of the number of integrated diffusion models would better validate the effectiveness of SABW.

R2: Based on your suggestion, we apply AFA to ensemble a larger number of base models.

Based on the results, ensembling more base models leads more performance improvements. Additionally, increasing the number of the ensembled base models enhances the tolerance for fewer inference steps. For example, when ensembling 4 or 5 base models, only 20 inference steps are sufficient, and with 6 base models, just 10 inference steps are sufficient. Further details can be found in Appendix F of the revised manuscript.

Thanks for your valuable comments. We will address the concerns below.

Weaknesses

W1: It appears that SABW is only suitable for integrating diffusion models with identical U-Net architectures, which may limit its flexibility and applicability.

R1: Based on your comment, we evaluate the generality of our AFA on other diffusion architectures. We select three SDXL [1] models and two FLUX .1 [dev] [2] models from CivitAI.

SDXL consists two U-Nets: the denoising U-Net and the refining U-Net. Both U-Nets shares a similar architecture of SDv1.5, but contain larger blocks. Consistent with the method used for ensembling SDv1.5 models, we employ the SABW module to aggregate features from each block.

Quantitative Comparison of ensembling SDXL models:

FLUX .1 [dev] is a DiT-based model. To ensemble models with this architecture, we also utilize the SABW module to aggregate features from each DiT Transformer block.

Quantitative Comparison of ensembling FLUX .1 [dev] models:

The performance of the ensembled model surpasses that of the individual base models, demonstrating the generality of our AFA on different model architectures. However, the observed improvement is relatively small, likely because the base models already exhibit strong performance, and ensembling leads to only marginal gains. Additionally, due to the large size of the base models, ensembling multiple such models may be less practical.

It is worth noting that our AFA cannot be applied to ensembling models with different architectures, as the misalignment of block features makes block-wise aggregation challenging. We view ensembling models with diverse architectures as a promising direction for future research.

The details can be found in Appendix I of the revised manuscript.

Dear Reviewer TNJb,

We have carefully prepared a response to address your concerns. Could you kindly take a moment to review it and let us know if it resolves the issues you raised? If you have any additional questions or suggestions, we would be happy to address them.

Thank you for your time and consideration.

The Authors

Questions

Q1: How does the method handle the situation that base models have correlated features or that base model bias is generated by imbalanced data?

A1: The training goal of our AFA is to achieve better denoising capability by learning a spatial attention map to linearly aggregate the features from the base models. During inference, given a specific prompt, the SABW module adaptively assigns weights to aggregate the features of the base models, striving to achieve best denoising capability.

Regardless of whether the base models exhibit correlated features or biases, the SABW module assigns higher weights to features that contribute positively to the denoising process and lower weights to those that may hinder the denoising process.

Q2: How is the method’s performance when facing particularly noisy input data?

A2: According to the principles of diffusion models, the image generation process begins by sampling a fully noisy latent representation that follows the standard Gaussian distribution. Consequently, during inference, our method only needs to handle noise conforming to this distribution. While different initial noise samples may introduce some diversity, the generated images will remain robust and faithfully adhere to the context of the given prompt.

Q3: What modifications should be considered when applying AFA to integrate models with different architectures?

A3: Due to the misalignment of block features, our AFA cannot be directly applied to ensemble models with different architectures. Achieving this would likely require an additional feature alignment method. For instance, inspired by REPA [3], a contrastive learning objective could be employed to align features from models with different architectures. We believe this direction holds significant potential and warrants further research.

Q4: Are there any potential optimization strategies to reduce the extra computational complexity brought by AFA to every single inference?

A4: In Appendix E of the original manuscript, we explore another MoE-based method, which applies the SABW module to determine the next block to be executed among the base models. While this method has the potential to further reduce the computational complexity of each inference step, we identified several disadvantages:

• Poor Performance. As shown in Table 4 and 5 of Appendix E, the performance of the MoE methods is inferior to that of our AFA. Additionally, the MoE methods do not support fewer inference steps.
• Comparable Training Computation Efficiency. Since all experts are still required to participate in training, the training TFLOPs of the MoE methods are comparable to that of our AFA.
• Inability to Perform Parallel Inference. Since the classifier-free guidance (CFG, Section 3.5) activates two different routes, the parallel inference cannot be conducted, which will increase the inference time. Furthermore, when multiple images need to be generated simultaneously, the MoE-based methods are inefficient, because the activated route for each image must be processed individually.

Exploring other alternative ensembling strategies to reduce the computational complexity of each inference step is a promising direction for further research.

Q5: What if the contributions of different models have conflict impacts on the result significantly?

A5: Given a prompt, if the base models have significantly conflict impacts on the result, it indicates that some base models contribute positively for the denoising process, while some have a negative impact. In such cases, the well-trained SABW module assigns higher weights to the features of the positive base models and lower weights to those of the negative base models.

[1] SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis

[2] https://blackforestlabs.ai/announcing-black-forest-labs/

[3] Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think

W2: As the number of base models increases, the method's reliance on multiple base models may result in overfitting if, say, the base models have correlated features.

R2: Based on your comment, we design an experiment to evaluate the robustness of our AFA against highly correlated features.

Specifically, we selected one high-quality base model and one low-quality base model. Using AFA, we ensemble the high-quality model with several copied low-quality models to evaluate whether the ensembled model will perform toward that of the low-quality model due to the dominance of the correlated features from the low-quality model.

Based on the results, although the performance the ensembled model slightly declines when ensembling with large number of low-quality models, it still outperforms the high-quality model. This demonstrates that our AFA exhibits strong robustness to highly correlated features.

Details can be found in Appendix L of the revised manuscript.

Thanks for your valuable comments. We will address the concerns below.

Weaknesses

W1: There are concerns about AFA’s capability to ensemble a broader range of model types or architectures.

R1: Based on your comment, we evaluate the generality of our AFA on other diffusion architectures. We select three SDXL [1] models and two FLUX .1 [dev] [2] models from CivitAI.

SDXL consists two U-Nets: the denoising U-Net and the refining U-Net. Both U-Nets shares a similar architecture of SDv1.5, but contain larger blocks. Consistent with the method used for ensembling SDv1.5 models, we employ the SABW module to aggregate features from each block.

Quantitative Comparison of ensembling SDXL models:

FLUX .1 [dev] is a DiT-based model. To ensemble models with this architecture, we also utilize the SABW module to aggregate features from each DiT Transformer block.

Quantitative Comparison of ensembling FLUX .1 [dev] models:

The performance of the ensembled model surpasses that of the individual base models, demonstrating the generality of our AFA on different model architectures. However, the observed improvement is relatively small, likely because the base models already exhibit strong performance, and ensembling leads to only marginal gains. Additionally, due to the large size of the base models, ensembling multiple such models may be less practical.

The details can be found in Appendix I of the revised manuscript.

Dear Reviewer x8ta,

We have carefully prepared a response to address your concerns. Could you kindly take a moment to review it and let us know if it resolves the issues you raised? If you have any additional questions or suggestions, we would be happy to address them.

Thank you for your time and consideration.

The Authors

Q2: How does AFA handle the diffusion models with different fine-tuned text encoders?

A2: Each ensembled diffusion model uses its own fine-tuned text encoder to encode the prompt. The resulting text embeddings are then injected into its U-Net blocks via the cross-attention mechanism. Since AFA aggregates the output features from the U-Net blocks, the fine-tuned text encoders remain unaffected by the feature aggregation process.

Q3: What is meant by "generating 4 images"?

A3: "Generating 4 images" refers to the process where, for each prompt, the models generate 4 images from different initialized noises.

[1] SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis

[2] https://blackforestlabs.ai/announcing-black-forest-labs/

Questions

Q1: How AFA generalize to the models with new diffusion frameworks, such as Stable Diffusion XL and DiT-based models?

A1: Based on your comment, we evaluate the generality of our AFA on other diffusion architectures. We select three SDXL [1] models and two FLUX .1 [dev] [2] models from CivitAI.

SDXL consists two U-Nets: the denoising U-Net and the refining U-Net. Both U-Nets shares a similar architecture of SDv1.5, but contain larger blocks. Consistent with the method used for ensembling SDv1.5 models, we employ the SABW module to aggregate features from each block.

Quantitative Comparison of ensembling SDXL models:

FLUX .1 [dev] is a DiT-based model. To ensemble models with this architecture, we also utilize the SABW module to aggregate features from each DiT Transformer block.

Quantitative Comparison of ensembling FLUX .1 [dev] models:

The performance of the ensembled model surpasses that of the individual base models, demonstrating the generality of our AFA on different model architectures. However, the observed improvement is relatively small, likely because the base models already exhibit strong performance, and ensembling leads to only marginal gains. Additionally, due to the large size of the base models, ensembling multiple such models may be less practical.

It is worth noting that our AFA cannot be applied to ensembling models with different architectures, as the misalignment of block features makes block-wise aggregation challenging. We view ensembling models with diverse architectures as a promising direction for future research.

The details can be found in Appendix I of the revised manuscript.

Thanks for your valuable comments. We will address the concerns below.

Weaknesses

W1: Some of the presentations need polishing.

R1: Thank you for your suggestion. We have refined the text (line 083-085), corrected the mistakes (line 354), and made Figure 2 easier to read in the revised manuscript.

W2: The authors should highlight the difference between the spatial attention of SABW and the typical multi-head self-attention. The authors should explain why the special architecture of SABW is necessary in comparison to some variant of self-attention.

R2: Thanks for your insightful suggestions. We have highlighted the difference and explained the necessity in the revised manuscript (line 250-258).

Specifically, the typical multi-head self-attention mechanism simply uses three linear layers to project the input features into three separate spaces(i.e., query, key, and value), and then computes the attention map across all the projected features. Our spatial-attention mechanism, based on an important experimental insight from Figure 1, takes into account the divergent denoising capabilities of multiple diffusion models, which are influenced by the prompts and denoising steps, and change dynamically across spatial locations. To account for this, our projection method incorporates the input prompt and the current denoising step. Moreover, it is sufficient to compute the attention map independently for each spatial location. As a result, better generation performance can be achieved by enabling multiple models to collaborate at different spatial locations under the conditions of the prompt and the denoising step.

W3: This study focuses on improving general generation capability, not on combining the styles or concepts of multiple models through ensembling/merging. The authors should highlight this distinction for the readers.

R3: Thank you for your insightful suggestion. We have emphasized this point in the revised manuscript (line 139 and 161, i.e., the footnote).

We appreciate Reviewer Ynbz for pointing out the issue of weird results for SDXL. Upon re-examination, we discovered that the reported AES, PS, and HPSv2 scores for the SDXL and FLUX architectures were incorrectly calculated due to a coding error (torch.max() was used instead of torch.mean()). As a result, the values represent the maximum scores rather than the average. We are very sorry for this oversight.

The revised average results for SDXL and FLUX are as follows. Thankfully, the conclusions remain unchanged: the performance of the ensembled model surpasses that of the individual base models. However, the observed improvement is relatively modest.

We have updated our manuscript accordingly in Appendix I.

Quantitative Comparison of ensembling SDXL models (Average):

Quantitative Comparison of ensembling FLUX .1 [dev] models (Average):

Additionally, we now also include the maximum and minimum values for completeness, as follows.

Quantitative Comparison of ensembling SDXL models (Maximum):

Quantitative Comparison of ensembling SDXL models (Minimum):

Quantitative Comparison of ensembling FLUX .1 [dev] models (Maximum):

Quantitative Comparison of ensembling FLUX .1 [dev] models (Minimum):

Dear Reviewer DhvC,

We have carefully prepared a response to address your concerns. Could you kindly take a moment to review it and let us know if it resolves the issues you raised? If you have any additional questions or suggestions, we would be happy to address them.

Thank you for your time and consideration.

The Authors

Q2: Other ensemble strategies can be compared, such as model distillation or more sophisticated attention mechanisms.

A2: Based on your suggestion, we attempted to further distill the ensembled model into significantly fewer steps. Specifically, we explore using LCM [1] to achieve this. In our experiments, we kept the parameters of the denoisers frozen while training only the parameters of the SABW modules. However, this approach is unsuccessful. A possible reason for this failure is that the frozen parameters may have impeded the distillation process. The details can be found in Appendix J of the revised manuscript.

Experimental results for distilling into fewer inference steps using LCM:

Q3: Could the authors extend their evaluations to diverse datasets to testify its generality?

A3: Based on your question, we evaluate AFA on three additional datasets: DiffusionDB [2], JourneyDB [3], and LAION-COCO [4]. The results are as follows. Please note that we present the averaged results of Group I and II. The experimental settings, the full experimental results, and comparisons with the baselines can be found in Appendix K of the revised manuscript. These results demonstrate that AFA consistently outperforms both the base models, highlighting the robustness and effectiveness of AFA.

Quantitative comparison on DiffusionDB:

Quantitative comparison on JourneyDB:

Quantitative comparison on LAION-COCO:

[1] Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference

[2] DiffusionDB: A Large-Scale Prompt Gallery Dataset for Text-to-Image Generative Models

[3] JourneyDB: A Benchmark for Generative Image Understanding

[4] https://laion.ai/blog/laion-coco/

W4: The paper does not address the scalability of the AFA method when ensembling a larger number of models, which could become computationally prohibitive.

R4: Based on your comment, we apply AFA to ensemble a larger number of base models.

Based on the results, ensembling more base models leads more performance improvements. Additionally, increasing the number of the ensembled base models enhances the tolerance for fewer inference steps. For example, when ensembling 4 or 5 base models, only 20 inference steps are sufficient, and with 6 base models, just 10 inference steps are sufficient. Further details can be found in Appendix F of the revised manuscript.

W2: The performance of AFA is contingent on the quality of the base models.

R2: The performance of the ensembled model being limited by the quality of the base models is a broad and inherent limitation in model ensembling. A straightforward way to mitigate this issue is to select high-quality base models, which, fortunately, is not cost-prohibitive.

W3: This paper does not extensively explore alternative ensemble strategies that could potentially offer further improvements.

R3: In Appendix E of the original manuscript, we explore another MoE-based method, which applies the SABW module to determine the next block to be executed among the base models. While this method has the potential to further improve performance and reduce FLOPs, we identified several disadvantages:

• Poor Performance. As shown in Table 4 and 5 of Appendix E, the performance of the MoE methods is inferior to that of our AFA. Additionally, the MoE methods do not support fewer inference steps.
• Comparable Training Computation Efficiency. Since all experts are still required to participate in training, the training TFLOPs of the MoE methods are comparable to that of our AFA.
• Inability to Perform Parallel Inference. Since the classifier-free guidance (CFG, Section 3.5) activates two different routes, the parallel inference cannot be conducted, which will increase the inference time. Furthermore, when multiple images need to be generated simultaneously, the MoE-based methods are inefficient, because the activated route for each image must be processed individually.

Exploring other alternative ensembling strategies is a promising direction for further research.

Thanks for your valuable comments. We will address the concerns below.

Weaknesses

W1: AFA's efficiency is lower for individual inference steps compared to merging-based methods. However, AFA is tolerance for fewer inference steps can offset this initial inefficiency.

R1: Compared to merging-based methods, our AFA does not offer advantages in a single inference step. However, since image generation requires multiple inference steps, a fair comparison should evaluate efficiency of the entire generation process.

Based on the experimental results, our AFA achieves comparable performance with fewer inference steps.

IR metrics of AFA with varying the number of base models and inference steps:

Building on this, we compare the efficiency of the merging-based methods and our AFA using 2, 3, and 6 base models.

Efficiency Comparison on 2 base models:

Efficiency Comparison on 3 base models:

Efficiency Comparison on 6 base models:

Under fewer inference steps, AFA demonstrate efficiency in terms of TFLOPs and inference time comparable to that of a single base model.

The details can be found in Appendix F and H of the revised manuscript.

Questions

Q1: Has any research optimizing the SABW module to reduce its computational overhead been done? potentially through model pruning, quantization, or more efficient network architectures.

A1: In Appendix E the original manuscript, we explore the MoE-based method that uses the SABW module to determine the next block to be executed among the base models. This method can be viewed as a form of model pruning. While this method has the potential to reduce computational overhead, it suffers from poor performance.

Since the computational overhead of SABW is lower than that of the U-Net denoisers, applying quantization and designing more efficient architectures for SABW could not significantly improve the computational overhead of AFA. Nonetheless, we acknowledge that optimizing of the efficiency of AFA remains necessary and is a worthwhile direction for further research.

Dear Reviewer WgNL,

We have carefully prepared a response to address your concerns. Could you kindly take a moment to review it and let us know if it resolves the issues you raised? If you have any additional questions or suggestions, we would be happy to address them.

Thank you for your time and consideration.

The Authors

Thanks for your valuable comments. We will address the concerns below.

Weaknesses

W1: It lacks comparison with the MoE diffusion approaches.

R1: Since the denoising experts in ERNIE-ViLG 2.0 [1], eDiff-I [2], and MEME [3] are specifically trained for particular timestep intervals, they are not compatible with the model ensembling setting, where the denoisers from the base models remain frozen. As an alternative, in Appendix E of the original manuscript, we briefly compared AFA with an intuitive Block-Level MoE method. To follow the methods of the papers that you mentioned, we additionally compare with the Denoiser-Level MoE method in the Appendix E of the revised manuscript. Specifically, we apply our SABW module to select the experts using the Gumbel softmax and the reparameterization trick. Since only one base model is activated during inferencing, the MoE methods may be more efficiency for inference.

However, we have identified several drawbacks of the MoE-based methods:

• Poor Performance. As shown in Table 4 and 5 of Appendix E, the performance of the MoE methods is inferior to that of our AFA. Additionally, the MoE methods do not support fewer inference steps.
• Comparable Training Computation Efficiency. Since all experts are still required to participate in training, the training TFLOPs of the MoE methods are comparable to that of our AFA.
• Inability to Perform Parallel Inference. Since the classifier-free guidance (CFG, Section 3.5) activates two different routes, the parallel inference cannot be conducted, which will increase the inference time. Furthermore, when multiple images need to be generated simultaneously, the MoE-based methods are inefficient, because the activated route for each image must be processed individually.

Quantitative comparison with 50 inference steps (showing the average results of Groups I and II):

Quantitative comparison with 20 inference steps:

Efficiency comparison:

Detailed results and analysis are provided in Appendix E of the revised manuscript. Additionally, we have included introductions of the papers you mentioned in Appendix E of the revised manuscript.

Q2: (1) How does the spatial attention mechanism in SABW interact with different types of prompts? (2) Are there certain prompt patterns where timestep-specific expert selection might be more appropriate than spatial aggregation?

A2: (1) All prompts interact with the SABW module through the cross-attention mechanism within its Transformer layer. (2) In our response to Weakness #1, we demonstrate that, under most prompt patterns, our AFA (i.e., spatial-aware feature-aggregation method) statistically outperforms the MoE methods (i.e., timestep-specific expert-selection methods). Identifying the specific prompt patterns where MoE may outperform AFA warrants further investigation.

[1] ERNIE-ViLG 2.0: Improving Text-to-Image Diffusion Model with Knowledge-Enhanced Mixture-of-Denoising-Experts

[2] eDiff-I: Text-to-Image Diffusion Models with an Ensemble of Expert Denoisers

[3] Multi-Architecture Multi-Expert Diffusion Models

[4] SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis

[5] https://blackforestlabs.ai/announcing-black-forest-labs/

[6] Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference

Questions

Q1: Have you explored knowledge distillation or other compression techniques to reduce the inference overhead of AFA?

A1: Based on your suggestion, we attempted to further distill the ensembled model into significantly fewer steps. Specifically, we explore using LCM [6] to achieve this. In our experiments, we kept the parameters of the denoisers frozen while training only the parameters of the SABW modules. However, this approach is unsuccessful. A possible reason for this failure is that the frozen parameters may have impeded the distillation process. Exploring additional techniques to reduce the inference overhead of AFA remains a promising direction for future research. The details can be found in Appendix J of the revised manuscript.

Experimental results for distilling into fewer inference steps using LCM:

W5: Lack of analysis on the trade-off between the granularity of adaptation (spatial/temporal) and computational cost.

R5: In our response to Weakness #1, we pointed out that while the MoE methods (i.e., timestep-specific expert-selection methods) outperform our AFA (i.e., spatial-aware feature-aggregation method) in inference TFLOPs, they offer no advantages in terms of model parameters, training TFLOPs, or inference time.

W4: There is uncertainty regarding the generalizability across different model architectures.

R4: Based on your comment, we evaluate the generality of our AFA on other diffusion architectures. We select three SDXL [4] models and two FLUX .1 [dev] [5] models from CivitAI.

SDXL consists two U-Nets: the denoising U-Net and the refining U-Net. Both U-Nets shares a similar architecture of SDv1.5, but contain larger blocks. Consistent with the method used for ensembling SDv1.5 models, we employ the SABW module to aggregate features from each block.

Quantitative Comparison of ensembling SDXL models:

FLUX .1 [dev] is a DiT-based model. To ensemble models with this architecture, we also utilize the SABW module to aggregate features from each DiT Transformer block.

Quantitative Comparison of ensembling FLUX .1 [dev] models:

The performance of the ensembled model surpasses that of the individual base models, demonstrating the generality of our AFA on different model architectures. However, the observed improvement is relatively small, likely because the base models already exhibit strong performance, and ensembling leads to only marginal gains. Additionally, due to the large size of the base models, ensembling multiple such models may be less practical.

It is worth noting that our AFA cannot be applied to ensembling models with different architectures, as the misalignment of block features makes block-wise aggregation challenging. We view ensembling models with diverse architectures as a promising direction for future research.

The details can be found in Appendix I of the revised manuscript.

W3: The limitations associated with the computational costs should be clearly explained and guidelines for practical application should be provided.

R3: Based on your suggestion, we compare the efficiency of merging methods and ensembling methods when considering 2, 3, and 6 base models.

Efficiency Comparison on 2 base models:

Efficiency Comparison on 3 base models:

Efficiency Comparison on 6 base models:

As shown in the above tables, when using the same 50 inference steps, ensembling methods (e.g., MagicFusion and AFA) do not have an efficiency advantage over a single base model or merging methods. The number of model parameters, TFLOPs, and inference time all increase linearly with the number of base models.

However, thanks to AFA’s high tolerance for fewer inference steps, it can achieve similar performance with reduced steps, while the baseline methods fail to maintain performance.

IR metrics of AFA with varying the number of base models and inference steps:

IR metrics comparison with varying inference steps:

Therefore, under fewer inference steps, AFA achieves efficiency in terms of TFLOPs and inference time comparable to that of a single base models and merging methods. For practical applications, we recommend using fewer inference steps when ensembling more base models to enhance inference efficiency.

The details can be found in Appendix F, G, and H of the revised manuscript.

W2: Missing theoretical analysis on when and why spatial-aware feature aggregation outperforms timestep-specific expert selection.

R2: In our response for Weakness #1, the MoE methods (i.e., the timestep-specific expert selection method) underperform our AFA (i.e., the spatial-aware feature aggregation method). In fact, the MoE method can be considered a specific case of our AFA. If, at each inference step, the learned weights from our AFA assign a value of 1 to all the features of a base model and 0 to others, this scenario effectively becomes an expert selection along the temporal axis. However, unlike the MoE methods, our AFA can flexibly assign weights to features from all base models, enabling a more refined and adaptive feature aggregation. This flexibility is crucial for achieving superior performance.