Diffusion Models for Multi-Task Generative Modeling

Thank you for your thoughtful review and valuable feedback. We appreciate the opportunity to address your concerns:

For Weakness 1: Implementation Details (2.2.4): 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.

For Weakness 2: Theoretical Portion: Response: Thanks for the suggestions. We respectfully disagree that our work is primarily empirical. Indeed, there is no theory on how to design a diffusion model that can handel multi-modal multi-task data by incorporating different task information in the forward process like our model. We consider this as one key contribution of our work. We believe it is necessary to present the underlying theoretical perspective of our model, which will provide guidance to design specific model architectures for multi-task generative modeling. However, we agree that the presentation can be simplified and emphasizes more on the model design part. We will be happy to do a thorough revision of our paper to reflect the suggestions from the reviewer.

For Weakness 3: Experiments on Widely Used Benchmarks: 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 question on what the single-task learning model for ImageNet classification, depicted in Figure 6, entails. Response: The single-task learning method in Figure 6 is the simple U-Net trained based on pure supervised learning, which is the most direct way to do classification.

I acknowledge the authors' efforts in providing clarification regarding weakness 1 and weakness 2.

However, I still find the writing in the current version of the paper to be unsatisfactory because of the limited space. The description of the implementation details of the method in section 2.2.4 and the experiment section remain unclear. Additionally, there are several typos and format issues in the reference section.

Furthermore, for weakness 3, the absence of a comparison with the prevalent latent diffusion raises doubts about the effectiveness of the proposed method.

We appreciate your thoughtful feedback on our work and would like to address the concerns raised in your review. We are committed to addressing these points in our revised manuscript, providing additional clarity and context where necessary.

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.

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.

Dear Reviewer,

As the deadline for rebuttal is approaching, we just want to update with you the CIFAR-10 experiment we are running. We have run the experiment on CIFAR-10 with MT-diffusion for masked-image training for around 8 hours on a single A100 GPU, as well as the conditional diffusion model that generate images from random masked images. The experiments are still not converging. In order to catch up with the rebuttal deadline, we have used the intermediate checkpoints (at the same running time) to generate 10K images, and calculate the FID and IS scores. The results are listed as follows:

Please note since the algorithm has not converged, these results are not comparable to current state of the art. However, we believe our final results can catch up quickly. More importantly, the current results indicate our model by modeling the joint distribution performs better than the simple conditional diffusion model. We will continue the experiment and add the final results to the final revision.

We appreciate your willingness to adjust your scores. Given our current and existing results, as well as our above rebuttals, can you please consider doing so? Thank you very much.

Thank you for your thoughtful evaluation and constructive feedback. We appreciate the opportunity to address your concerns:

For Weakness 1: Experiments on Low Resolution and Small Datasets: 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.

For Weakness 2: Lack of Model Details for Each Task: 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.

For the Question on Model Structure: 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.

Thank you for your thoughtful review and valuable feedback. We appreciate the careful consideration given to our work. Below, we address each of the raised points and provide additional clarification.

For Weakness 1: Increased Training Costs: We appreciate your concern regarding potential increased training costs. In our two-task setting, the additional computational overhead is minimal, amounting to approximately 10% more time per iteration than a pure single-task diffusion model. Importantly, our model remains significantly more efficient than training two individual diffusion models for the two tasks separately in terms of both time and storage efficiency. We'd like to highlight that our method not only handles multi-task learning effectively but also converges faster, as indicated in Table 2. We will incorporate this discussion into the final revision for enhanced clarity.

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.

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.

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.

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.

The concerns I have raised have been largely addressed, and I appreciate the efforts made by the authors in this regard. However, my concern regarding Weakness 2 remains unaddressed. I acknowledge that it may be challenging to completely disentangle the effects of data augmentation from multi-task learning. Nonetheless, it is imperative to support, through experimental or theoretical analysis, whether the proposed multi-task scenarios actually synergized to enhance performance. For instance, assessing task affinity to verify positive transfer between tasks would significantly bolster the validation of this research's utility.

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.