TimeStep Master: Asymmetrical Mixture of Timestep LoRA Experts for Versatile and Efficient Diffusion Models in Vision

4. Question: The ablation study shows that the increase of r does not necessarily increase performance.

All experiments presented in our paper were conducted under the constraint of maintaining equal training costs. When we remove this constraint and train each timestep expert with n=8 and 4000 training steps, the results are shown in the table below. The experimental results demonstrate that the model's performance indeed increases with the growth of r.

5. Question: Can this framework be applied to more tasks?

We have verified the effectiveness of TSM on three classic architectures (U-Net, DiT, MMDiT), two modalities (image, video), and three tasks (domain adaptation, post-pretraining, model distillation). Due to time constraints, we will complete the comparative experiments of TSM on more downstream tasks in future versions of the paper, such as image and video restoration, image and video segmentation, etc.

Thank you for your valuable reviews, we will answer your questions one by one regarding these weaknesses.

1. Question: The two-stage TSM approach relies on a relatively complex design and may not generalize easily to tasks outside the evaluated benchmarks.

We have verified the effectiveness of TSM on three classic architectures (U-Net, DiT, MMDiT), two modalities (image, video), and three tasks (domain adaptation, post-pretraining, model distillation). We think the generalization ability and versatility of our TSM have been fully verified. Furthermore, most of our hyperparameter selections are directly followed by the baseline [1, 2]. For additional hyperparameters regarding the selection of core LoRA and context LoRA, there is actually the simplest selection method, which is our We have obtained SOTA settings in multiple experiments, that is, n=8 for core LoRA and n=1 for context LoRA. In addition, no additional transcendent parameters are introduced, so it will be easy to transfer our method to other fields.

[1] Kaiyi Huang, Kaiyue Sun, Enze Xie, Zhenguo Li, and Xihui Liu. T2i-compbench: A comprehensive benchmark for open-world compositional text-to-image generation. Advances in Neural Information Processing Systems, 36:78723–78747, 2023.

[2] Tianwei Yin, Michael Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and William T Freeman. Improved distribution matching distillation for fast image synthesis. arXiv preprint arXiv:2405.14867, 2024a.

2. Question: Extra parameters limit the scalability of the model. It also inherits the inference procedure of the diffusion process, which is also computationally burdensome.

Because the amount of LoRA parameters we introduced is very small, as shown in the table below, the increased time and memory consumption of our method compared to the origin model is indeed acceptable. Therefore, TSM does not actually affect model scalability.

We have supplemented the following table with a comparison of the time and memory usage required to run the original model and our method on a single A100, where the batch size is set to 1.

3. Question: Compared to Vanilla LoRA, some improvements, such as in color and shape metrics shown in Table 1, are relatively modest.

In Table 1, compared to vanilla LoRA, TSM demonstrates notable improvements across all three fundamental architectures (UNet, DiT, and MMDiT). However, the enhancement brought by TSM to SD3 is relatively less pronounced than those observed in SD1.5 and PixArt-α. This is primarily because SD3's performance on T2I-Compbench has already approached saturation levels. Therefore, the fact that TSM achieves approximately 1% improvement in attributes such as color and shape is actually quite remarkable and commendable. Furthermore, as shown in Tables 2 and 3, the advantages of TSM over vanilla LoRA are demonstrated even more comprehensively.

3. Question: The prompt misalignment problem is not exclusive to LoRA-tuned models and may not be solely caused by the different noise levels at various timesteps.

Thank you for the analysis. Perhaps there are various reasons for the degradation of model capabilities caused by LoRA, such as the low-rank characteristics of the matrix, insufficient training, etc. However, TSM can largely solve the degradation problem caused by LoRA under the premise of approximate training overhead. Therefore, we believe that TSM will replace LoRA as the preferred solution for low-cost fine-tuning of diffusion models.

4. Question: It might be sufficient to fine-tune only the UNet parameters without modifying the text encoder.

As shown in Table 1, 2, 3, 4, 5, 6, our comparison experiments between vanilla LoRA and TSM, the modules for adding LoRA are all the same, where LoRA is consistently added to the text encoder. Therefore, our ablation experiments are fair and can fully prove that our TSM is significantly better than vanilla LoRA. Besides, most of our hyperparameter selections (including whether to train a text encoder or not) are directly followed by the baselines [1, 2].

[1] Kaiyi Huang, Kaiyue Sun, Enze Xie, Zhenguo Li, and Xihui Liu. T2i-compbench: A comprehensive benchmark for open-world compositional text-to-image generation. Advances in Neural Information Processing Systems, 36:78723–78747, 2023.

[2] Tianwei Yin, Michael Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and William T Freeman. Improved distribution matching distillation for fast image synthesis. arXiv preprint arXiv:2405.14867, 2024a.

5. Question: It’s unclear whether the gating-based combination method explored the potential degradation in model performance.

Table 9 in our article shows our ablation experiments on router architecture. Besides, the experimental results in Tables 1, 2, and 3 all show that even such a simple router architecture is enough to make our method better than existing methods.

6. Question: Ablation studies on LoRA combination methods (beyond just the gating inputs) would be useful. A related work is a Mixture of LoRA Experts.

In the table below, we further add a comparison of MoE LoRA trained from scratch (even though we are working at the same time) with TSM. It can be seen that our TSM still achieves the best results.

7. Question: Its practical use in certain scenarios remains to be seen.

We have verified the effectiveness of TSM on three classic architectures (U-Net, DiT, MMDiT), two modalities (image, video), and three tasks (domain adaptation, post-pretraining, model distillation). We think the practicality and versatility of our TSM have been fully verified.

8. Question: In Fig. 4, the "in the style of Sci-Fi" prompt does not seem to be well followed, and in Fig. 5, some images (e.g., the black-and-white image) show poor understanding of the prompt.

The Sci-Fi actually denotes Science-Fiction, and the test prompts are selected from EvalCrafter benchmark. To be noted, in the example shown on the left side of Fig. 4, the result generated by VC2 lacks a crucial main character, while the result from VC2 with LoRA is missing a bird. The absence of the main subject represents a relatively more significant error. In the example presented in Fig. 5 (the black-and-white image), SD1.5 has generated numerous peculiar elements, and the image produced by SD1.5 with LoRA even includes green vegetation, which evidently deviates substantially from the input prompt. Compared to other methods, our method has greatly improved the problem of mismatch between text and generated contents. Although there may still be some small mismatches, it does improve a lot.

9. Question: Does this only result in an additional 1M parameters? Could the authors clarify more specifically where the 1M increase in training parameters comes from? Is it solely due to the added gating mechanism?

Thank you for the suggestion. We added the analysis of parameters. As shown in the table below, because the structure of the router is very simple, the additional training parameters introduced in the second stage are only 1M and the increased parameters of TSM compared to the origin model are acceptable.

Thank you for your valuable reviews, we will answer your questions one by one regarding these weaknesses.

1. Question: The paper does not clearly compare its method with non-LoRA approaches that address prompt misalignment issues.

Thanks for the reminder. These methods perform additional training-free operations on the attention module, but their improvement in the essential capabilities of the model is limited. We conducted experiments on SD2 and compared the results with other methods (including training-free methods) as follows. The results demonstrate that TSM still has a large performance advantage compared to training-free methods.

[1] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bj¨orn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10684–10695, 2022a.

[2] N. Liu, S. Li, Y. Du, A. Torralba, and J. B. Tenenbaum, “Compositional visual generation with composable diffusion models,” in ECCV, 2022.

[3] W. Feng, X. He, T.-J. Fu, V. Jampani, A. Akula, P. Narayana, S. Basu, X. E. Wang, and W. Y. Wang, “Training-free structured diffusion guidance for compositional text-to-image synthesis,” in ICLR, 2023.

[4] H. Chefer, Y. Alaluf, Y. Vinker, L. Wolf, and D. Cohen-Or, “Attend-andexcite: Attention-based semantic guidance for text-to-image diffusion models,” in ACM Trans. Graph., 2023.

[5] Kaiyi Huang, Kaiyue Sun, Enze Xie, Zhenguo Li, and Xihui Liu. T2i-compbench: A comprehensive benchmark for open-world compositional text-to-image generation. Advances in Neural Information Processing Systems, 36:78723–78747, 2023.

2. Question: A comparison of computational overhead and training parameters would provide more insight into the advantages of the proposed method.

Because the amount of LoRA parameters we introduced is very small, as shown in the table below, the increased time and memory consumption of our method compared to the origin model is acceptable.

We have supplemented the following table with a comparison of the time and memory usage required to run the original model and our method on a single A100, where the batch size is set to 1.

4. Question: Could an evaluation based on FID be added?

Thank you for the suggestion. We added the results of our FID experiment on PixArt are shown in the following table. PixArt-α's result is measured by ourselves since the model they published is not the model used to measure FID in their paper.

5. Question: Would further increasing the number of experts theoretically enhance the results even more?

Intuitively, it seems this way, but our experiments have shown that this is not the case. The premise that must be ensured in our experiments is that with the same training cost, as the timesteps are further subdivided and the number of experts increases, the number of steps each expert is trained on decreases. The experiments in Table 6 indicate that when n≤4, the model's performance improves as the number of experts increases. However, if n continues to increase, the model's performance may even decline. We supplemented the experimental result of n=16 as follows, which further verified our conjecture.

Thank you for your valuable reviews, we will answer your questions one by one regarding these weaknesses.

1. Question: This results in higher memory usage and greater latency during inference.

Thank you for your concern. We have added the analysis of the inference memory and time usage. For training, because the amount of LoRA parameters we introduced is very small, as shown in the table below, the increased time and memory consumption of our method compared to the origin model is acceptable.

For inference, we have supplemented the following table with a comparison of the time and memory usage required to run the original model and our method on a single A100, where the batch size is set to 1.

2. Question: Although the assembling stage can be understood from the flowchart, the explanation of Equation (7) may be somewhat confusing, especially regarding the subscript notation of ϵ, which lacks clarity.

Thank you for your detailed reminder. As we stated in sections 139-140 of the article, ϵ is the noise sampled from a Gaussian distribution. We are sorry for the misunderstanding.

3. Question: The results for SD2 with the proposed method are missing. What could be the reason for this?

In our paper, we chose a UNet-based model (SD1.5), a DiT-based model (PixArt-α) and a MMDiT-based model (SD3) as our pretrained model. The reason we chose SD1.5 instead of SD2 is that SD1.5 seems to be more widely circulated on major platforms. The results show that by using SD1.5+TSM significantly exceeds the SD1.5 and its various variants. To further verify our observation, we added the experimental results of TSM on SD2 in the following table. It can be seen that we achieved the most advanced results on SD2 as well.

[1] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bj¨orn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10684–10695, 2022a.

[2] N. Liu, S. Li, Y. Du, A. Torralba, and J. B. Tenenbaum, “Compositional visual generation with composable diffusion models,” in ECCV, 2022.

[3] W. Feng, X. He, T.-J. Fu, V. Jampani, A. Akula, P. Narayana, S. Basu, X. E. Wang, and W. Y. Wang, “Training-free structured diffusion guidance for compositional text-to-image synthesis,” in ICLR, 2023.

[4] H. Chefer, Y. Alaluf, Y. Vinker, L. Wolf, and D. Cohen-Or, “Attend-andexcite: Attention-based semantic guidance for text-to-image diffusion models,” in ACM Trans. Graph., 2023.

[5] Kaiyi Huang, Kaiyue Sun, Enze Xie, Zhenguo Li, and Xihui Liu. T2i-compbench: A comprehensive benchmark for open-world compositional text-to-image generation. Advances in Neural Information Processing Systems, 36:78723–78747, 2023.

4. Question: It is not clear why such a choice is made if this is the real rank of the target matrix. In addition, it is hard to justify why different blocks share the same rank.

The rank of LoRA is not the real rank of the parameter matrix. Instead, it is used to fine-tune the parameter matrix by approximating it to a low-rank matrix. Our setting of r directly follows the rank used in the previous work [1, 2], which shows that the generalization of our method is superior. In addition, in the table below, we added the ablation experiment between r and model performance when n=8. It is obvious that, as r increases, the performance of the model also increases significantly. For the sake of simplicity, using the same LoRA in different blocks allows us to quickly start training a new model. If the optimal parameters are searched for each block, the calculation cost will be relatively large.

[1] Kaiyi Huang, Kaiyue Sun, Enze Xie, Zhenguo Li, and Xihui Liu. T2i-compbench: A comprehensive benchmark for open-world compositional text-to-image generation. Advances in Neural Information Processing Systems, 36:78723–78747, 2023.

[2] Tianwei Yin, Michael Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and William T Freeman. Improved distribution matching distillation for fast image synthesis. arXiv preprint arXiv:2405.14867, 2024a.

Thank you for your valuable reviews, we will answer your questions one by one regarding these weaknesses.

1. Question：The reason for using different scales of intervals in the assembling stage is not clearly explained.

According to our exploration during the experiment, under different timestep conditions, the variances between the features in the diffusion model are very different, and according to the lines 072-075 of our paper, previous work [1, 2] has also analyzed that there are substantial variance discrepancies between the model's intermediate features under different timestep conditions. Therefore, we believe it would be challenging to fit such widely divergent distributions using identical parameters. This reasoning led us to propose our TSM.

[1] Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Qinsheng Zhang, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, et al. ediff-i: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01324, 2022.

[2] Ma, Xinyin and Fang, Gongfan and Wang, Xinchao. DeepCache: Accelerating Diffusion Models for Free . CVPR, 2024.

2. Question: The paper lacks a comparison between the asymmetrical MoE (Mixture of Experts) and the standard MoE.

In the following table, we add the comparison between MoE LoRA trained from scratch and TSM, where TSM still achieves the best results. Furthermore, in Table 7 of our paper, on the premise of inheriting the well-trained timestep LoRA experts from TSM 1-stage, we compared asymmetric MoE and symmetric MoE under three model architectures. Experiments proved that asymmetric MoE is better than Common MoE.

3. Question: In Figure 3, if only a single scale with n=8 intervals were used, would the vanilla MoE already be capable of learning how to combine the context experts with the core expert through top-2 gating?

Thank you for your thoughts on our paper and your suggestions for improvements. In fact, we have done this experiment before, and the results are shown in the table below. Our conclusion is that if LoRA trained under the same interval is used as experts in MoE, the performance of the model cannot continue to improve. We speculate that it is because LoRA trained under the same interval can only specialize in the timestep condition within the training range. For example, LoRA trained under large noise levels does not help LoRA trained under small noise levels. Therefore, we use LoRA trained under different intervals as MoE experts in TSM. A large number of experimental results in the paper also prove that our asymmetric design is effective.