Paper: Diffusion Models for Multi-Task Generative Modeling (ICLR-2024)

Reviewer comment (Reproducibility):
The explanation of the implementation details (2.2.4) regarding the encoders and decoders can be perplexing, particularly in terms of how the classification label is encoded and how it is "aggregated" with other tensors.

Author response:
We acknowledge the complexity of the implementation details in Section 2.2.4. More detailed descriptions for different task settings, including how label data is encoded, are provided in Section 4 (Experimental Part) and Appendix E. Please refer to those sections for a comprehensive understanding, which we will put effort in to make it more clear in the revision.

Specifically, for label data, we utilize the original U-Net structure from the ADM model, accepting input with a noisy image and label (from an embedding layer as a conditional input). In our case, the noisy image is set as the original image, and label embeddings are jointly encoded through the U-Net.

Reviewer comment (Experiments):
While the proposed method demonstrates significant improvements on the ImageNet-64 benchmark, it lacks experiments on more widely used and challenging benchmarks, as well as comparisons with newer generation models. Additionally, some auxiliary experiments in the Appendix utilize stable diffusion, which is now commonly employed as a baseline, while the primary experiments do not present any relevant results. The absence of these experiments makes it challenging to provide sufficient justification for the superiority of the proposed complex multi-task training pipeline.

Author response:
We understand the importance of benchmarking on widely used datasets. In our main experiments, we train on both ImageNet-64 and ImageNet-128 based on the ADM codebase, providing a well-documented baseline with accepted evaluation scores for comparison. Training larger models on datasets like LAION from scratch is computationally challenging. More importantly, there are many factors such as the data quality in the training that can significantly impact the final model performance, thus making directly comparing with these models in a fair and faithful way difficult. We believe a quantitative comparison with the ADM baseline on ImageNet sufficiently demonstrates the effectiveness of our method.

For comparisons with stable diffusion, we acknowledge that it is very challenging to perform a fair and convincing comparison between our model and stable diffusion on large datasets such as the LAION dataset we used. However, to resolve your concern, we have conducted an additional visual comparison experiment to compare our model and stable diffusion. The results are plotted in two new figures, Figure 13 and Figure 14 in the Appendix. It appears that our model seem to understand the underlying semantics of the images better, and generates better looking images corresponding to the prompts in many cases. Please check our revision PDF file.

Reviewer comment (Reproducibility):
Could you please elaborate on the specific encoders employed for each task and the type of data they operate on? It would be helpful to understand the label space for each task, particularly regarding how a one-hot classification vector is encoded in the latent space.

Author response:
We acknowledge the complexity of the implementation details in Section 2.2.4. More detailed descriptions for different task settings, including how label data is encoded, are provided in Section 4 (Experimental Part) and Appendix E. Please refer to those sections for a comprehensive understanding, which we will put effort in to make it more clear in the revision.

Specifically, for label data, we utilize the original U-Net structure from the ADM model, accepting input with a noisy image and label (from an embedding layer as a conditional input). In our case, the noisy image is set as the original image, and label embeddings are jointly encoded through the U-Net.

Reviewer comment (Novelty):
The related work and experiments are all around diffusion models with very trivial baselines and simple experiments. I understand the author shape this paper as the one of the first to explore diffusion in a multi-task learning setting. But at the same time, I saw some other papers also using the multi-task learning technique in diffusion to enforce geometric constrains, particularly in the 3d shape synthesis. For example, Wonder3D (https://www.xxlong.site/Wonder3D/) and DreamCraft3D (https://mrtornado24.github.io/DreamCraft3D/) are two examples applying diffusion on both RGB and normal maps to improve multi-view/geometric consistency. I am aware both papers were released very recently and seem to be submitted to ICLR as well, I am just wondering how the proposed paper differentiates itself from the straightforward implementation of using multi-task learning in diffusion like being used in these two papers, and from which both were based on the same assumption of using shared diffusion space as well?

Author response:
For Weakness 1: We respectfully assert that our choice of baselines and experiments is non-trivial. ADM by OpenAI, a baseline model, stands as one of the state-of-the-art diffusion generative models. Additionally, conducting experiments on the ImageNet dataset, a large and computationally demanding dataset, is far from trivial. For instance, a single experiment on 8 A100 GPU server takes several days. We believe our experiments are well-justified given the computational demands of the tasks and the significance of the chosen baselines.

Thank you for pointing out these interesting references. We want to point out that an important difference between these methods and ours is that they are essentially based on the single-task diffusion model, i.e., the underlying principle is the conditional diffusion model, with additional losses to introduce some regularizations for the 3D generation. We consider at least two differences from our work: 1) the conditional generation only models a conditional distribution, whereas ours models the joint distribution of different modality data. This allows us to perform not only joint generation, but also conditional generation as the purpose of the references, i.e., ours is a more general generation framework. 2) The addition of extra losses in the reference works is mostly ad hoc, whereas in our framework, all the losses are attributed from the new multi-task ELBO, providing a more general and principled approach for multi-task diffusion generation modeling.

Reviewer comment (Experiments):
As such, I am not exactly sure how to comment and understand the performance of the proposed formulation, since the experiment setting is simple and only compare with simple single task baseline. For example, it might be more intriguing to see i) the performance / formulation without having a shared diffusion space, ii) how to compare with multi-task predictive models without any diffusion (e.g. MTAN, Adashare, Cross-Stitch, PadNet, as multi-task learning in computer vision is an active research area.), iii) is task conflict issue in multi-task learning alleviated? What are the other benefits other than improved performance?

Author response:
For Weakness 2: We acknowledge the complexity of our experimental settings, focusing on incorporating multi-task learning into diffusion models—a non-trivial task given the unique characteristics of diffusion models. We appreciate your specific questions and provide the following responses:
1. Performance/formulation without a shared diffusion space corresponds to the baseline of modeling tasks using independent single-task diffusion models, which we have compared with.

2. How to compare with multi-task predictive models without any diffusion: Since we focus on generative modeling, it is not exactly clear to us how to compare our model with these predictive models. However, we believe some ideas can be incorporated into our framework for further performance improvement, which we would like to leave as interesting future work.

3. Is task conflict issue in multi-task learning alleviated: As a first work, the task conflict issue alleviation was not investigated in this study, which we believe would need significant effort. THus, we consider it as an interesting avenue for future exploration.

4. What are the other benefits other than improved performance: We believe the most important benefit of our framework is that it enables multi-task generative modeling, i.e., one can model multiple generation tasks in a single model without hurting single-task performance. This can not only save training cost, but also alleviate the burden of maintaining different models for different tasks.

Reviewer comment (Experiments):
Experiments are done on low resolution and small datasets, undermining its effectiveness.

Author response:
We respectfully disagree with the characterization of our experiments. While the primary experiments were conducted on the ImageNet1K dataset, which is widely acknowledged as substantial in the generative modeling community, we acknowledge the existence of larger datasets such as LAION. However, due to the prohibitive cost of pretraining from scratch on these datasets, they are not considered standard for comparisons.

We conducted experiments on two resolution levels, 64x64 and 128x128, leveraging our computational resources. These experiments, although time-consuming, are deemed sufficient to demonstrate the effectiveness of our method. Additionally, qualitative experiments on the larger LAION dataset, based on the pretrained stable diffusion model, are detailed in the appendix for further insights.

Reviewer comment (Reproducibility):
The paper is lack of model details for each task.

Author response:
We appreciate your observation, and we want to clarify that all model details are provided in the experiment section and Section E of the Appendix, adhering to page limits. To enhance clarity, we are willing to paraphrase the paper and incorporate more detailed experimental settings into the main text.

Reviewer comment (Reproducibility):
Can you explain the model structure for each task and result? Especially results in Fig 4 and Fig 5.

Author response:
Detailed information on experimental settings and model structures is provided in Appendix E. The model structures closely resemble the ADM model, incorporating a shared U-Net backbone shared by all tasks, as illustrated in Figure 3. Task-specific encoder structures are introduced in our model for added specificity. For a comprehensive understanding, please refer to Appendix E, and feel free to reach out with any additional questions.

Reviewer comment (Experiments):
While significant performance improvements are shown across various metrics, including FID, the analysis lacks control of variables to confirm that these improvements truly stem from a multi-task setting. Following the previous point, it is my view that the proposed methodology likely entails considerably higher training costs and an increased number of data samples seen by the model compared to baseline learning. Therefore, it is necessary to deeply analyze whether the performance improvement is due to positive transfer resulting from multi-task learning, or merely an effect akin to data augmentation from masked samples. The absence of such analysis has influenced my evaluation towards rejection.

Author response:
For Weakness 2: Where performance improvement comes multi-task learning:
We wish to argue that the additional data augmentation in the masked image training is also a part of the multi-task setting, offering additional benefits without incurring extra overhead. We believe separating data augmentation from the multi-task setting is challenging, as that is part of the multi-task framework. 

We do not fully understand your point about positive transfer and data augmentation, because the auxiliary task only contains data augmentation. Thus, we believe the data augmentation indeed attributes to the performance gain, which can also be understood as positive transfer. We will explicitly discuss the insight in the revised manuscript.

Dear Reviewer,

We appreciate your feedback on our rebuttal. We are happy to learn that you are mostly satisfied with our rebuttal, and only have some concerns on verifying whether the performance comes from positive transfer. We are happy to explore and performance more ablation studies on this. Specifically, we have set up to run the following experiments: train our multi-task diffusion model on 5 tasks to simultaneously learning to generate original images,, masked images, corresponding captions, random captions, and class labels. Intuitively, generating original images and masked images are two closest tasks, which is expected to have more positive transfers. While the other tasks such as generating random captions are completely independent of generating images, thus we expect the performance would drop if the task of generating random captions are added. We are working on the experiments now, but please note this is a more challenging setting, thus we are not expecting to be able to produce some results before the rebuttal deadline. However, we are committed report the results in our revision.

Even though we are not able to provide the aforementioned results in the rebuttal period ending in one day, we wish to point out our current results actually have the implications that more similar tasks tend to have more positive transfers. For example, comparing the following two settings indicated in Table 2: 1) simultaneously generating images and the corresponding masked images; 2) simultaneously generating images and the corresponding labels. The former two tasks are considered closer as they are in the same data space. And from Table 2, we can clearly see that the former setting (Generation with Masked-Image Training (Section 4.1)) outperforms the latter one (Generation with Joint Image-Label Modeling (Section 4.2)), indicating there are more positive transfers in the former setting.

We hope this can address your concerns, and our explanation can provide more information on your final decision. Thank you very much for your time.

Dear Reviewer,

Thanks again for your time reviewing our paper. We just want to send out a follow-up request to please reconsider your final decision of our paper. Specifically, we believe our existing results, as pointed our above, can at least partially resolve your concern on the positive transfer and task similarity issue. Furthermore, we have setup an extra experiment to verify that. Due to the complexity of the experiment, we cannot provide results during the rebuttal period. However, we are committed to finish that in our final revision.

Furthermore, we would like to highlight that we have provided some additional experiment verifications to address the issues raised by other reviewers, which are reflected in the Appendix of the updated PDF file. As a result, all other reviewers agree to raise their scores.

Again, thanks a lot for your time and effort in reviewing our paper. We appreciate it if you can adjust your final decision.

Reviewer comment (Novelty):
(Minor point) There is prior (possibly concurrent) work proposing a multimodal, multi-task diffusion process through a Versatile Diffusion[1] multi-flow diffusion process.

Author response:
For Weakness 3: Comparison with Prior Work:
Thank you for bringing up the reference to Versatile Diffusion. After a careful review, we recognize that it is indeed a concurrent work that we did not initially notice. The Versatile diffusion mainly focuses on developing new neural architectures that make different tasks interact with each other within the single-task diffusion framework. We believe that our work is different, as it not only introduces a novel neural architecture but also generalizes the single-task diffusion in the loss function. We will enrich the discussion on the relationship between our model and Versatile Diffusion in the revision.

Reviewer comment (Presentation):
(Minor point) The caption of Table 2 does not provide sufficient information, making it difficult for the reader to interpret the experimental results.

Author response:
For Weakness 4: Table 2 Caption:
We appreciate your feedback on the clarity of the Table 2 caption. We acknowledge that additional details about different variants of our model are described in Sections 4.1 and 4.2. To enhance clarity, we will provide more comprehensive information in the title for improved interpretability in the revised manuscript.

Reviewer comment (Novelty):
How do you think the proposed off-the-shelf guidance method would integrate with previous research focused on efficient training of diffusion models, such as P2-Weighting[2], Min-SNR[3], ANT[4], and Task Routing[5]? Particularly, Min-SNR[3], ANT[4], and Task Routing[5] view the methodology of diffusion models as naturally creating a multi-task situation through various time steps, each requiring different levels of denoising. Considering this study aims to extend multi-task learning by increasing the input modality for denoising, there seems to be an overlap. I would be interested to hear your insights on this matter.

Author response:
For Questions: Relation with Previous Research:
We are grateful for the references to inspiring papers interpreting standard diffusion models from a multi-task learning perspective. While we acknowledge the relevance of these papers, we maintain that our ideas are orthogonal, with no technical overlap. We are open to exploring the incorporation of these insights into our framework for potential further improvement—a promising avenue for future work.
Thank you once again for your thoughtful review. We are committed to addressing each of these points thoroughly in the final revision, ensuring a more robust and comprehensive presentation of our work. We hope the reviewer can reconsider his/her decision.