Diffusion Models for Multi-Task Generative Modeling
Authors: Changyou Chen
Han Ding
Bunyamin Sisman
Yi Xu
Ouye Xie
Benjamin Z. Yao
Son Dinh Tran
Belinda Zeng
Venue: ICLR 2024 poster (2024)
Keywords: Diffusion Models
Multi-Task Multi-Modal Learning
Generative Models
Primary area: generative models

Abstract: Diffusion-based generative modeling has been achieving state-of-the-art results on various generation tasks. Most diffusion models, however, are limited to a single-generation modeling. Can we generalize diffusion models with the ability of multi-modal generative training for more generalizable modeling? In this paper, we propose a principled way to define a diffusion model by constructing a unified multi-modal diffusion model in a common {\em diffusion space}. We define the forward diffusion process to be driven by an information aggregation from multiple types of task-data, {\it e.g.}, images for a generation task and labels for a classification task. In the reverse process, we enforce information sharing by parameterizing a shared backbone denoising network with additional modality-specific decoder heads. Such a structure can simultaneously learn to generate different types of multi-modal data with a multi-task loss, which is derived from a new multi-modal variational lower bound that generalizes the standard diffusion model. We propose several multi-modal generation settings to verify our framework, including image transition, masked-image training, joint image-label and joint image-representation generative modeling. Extensive experimental results on ImageNet indicate the effectiveness of our framework for various multi-modal generative modeling, which we believe is an important research direction worthy of more future explorations.

Published as a conference paper at ICLR 2024
DIFFUSION MODELS FOR MULTI-MODAL GENERATIVE MODELING
Changyou Chen Ouye Xie
Han Ding
Bunyamin Sisman2 Son Tran Amazon
Xu Belinda Zeng²
Benjamin Yao
University at Buffalo
ABSTRACT
Diffusion-bart generative modeling has been achieving sfate-of-the-an aar results various generation tasks. Most diffusion models, however, are limited to a single-generation modeling. Can we generalize diffusion models with the ability of multi-modal generative training for more generalizable modeling? In this paper, we propose a principled way to define a diffusion model by constructing a unified multi-modal diffusion model in a common diffiusion space. We define the forward diffusion process to be driven by an information aggregation from multiple types of task-data, e.g., images for generation task and labels for a classification task. In the reverse process, we enforce information sharing by parameterizing a shared backbone denoising network with additional modality-speeiif decoder heads, Such a structure can simultaneously learn to generate different types of multi-modal data with a multi-task loss, which is derived from a new multi-modal variational lower bound that generalizes the standard diffusion model. We propose several multi- modal generation settings to verify our framework, including image transition, masked-image training, joint image-label and joint image-representatinn generative modeling. Extensive experimental results on ImageNet indicate the effectiveness of our framework for various multi-mode generative modeling, which we believe is an important research direction worthy of more future explorations.
INTRODUCTION
The field of artificial intelligence (AI) has witnessed significant advancements in generative modeling, leading to remarkable progresses such as DALL-E (Ramesh et al.] (022) and GPT-4 (OpenAl 2023). The generative AI paradigm enables the learning of transitions from simple to complex distributions, such as from a standard Gaussian distribution to a nigh-dimensiooal image distribution. Compared to discriminative learning, generative mechanisms can arguably prioritize the overall structures of the data, offering better data fitting and potential robustness to data noise. However, while real-worr world applications often involve data of multiple types (multi-modal), including images, video, text, and labels, most existing generative models primarily focus generating a single data type or modality. Notably, the diffusion model (Sohl-Dickster
et al. 2015 Ho et al. (020), state-of sfate-of-thh-arr generative model, has been independently developed for generating image, text, audio, and label data (Dhariwal & Nichol, 2021] Li et al. 2022b Liu et al.) [023] Han et al. (022). Can we design a principled way to enable joint modeling and generating multi-mode data within the diffusion- modieo framework?
Furthermore, leveraging mudti-modal information through learning from multiple tasks and data sources has proven to be highly effective to learn generalized representations. Prominent examples include the ALBEF and BLIP models, which jointly learns from multi-modal data to match image and text ([. et al.] 202T 2022a| 2023), and the BERT model, which benefits from multi-task training such masked token prediction and consecutive sentence prediction (Devlin et al.) (019b). Can we adopt a similar setting to leverage multi-modal data and losses into the diffusion-moote framework, so to better integrate shared information among tasks for better generative modeling?
In this paper, we present our initial endeavor towards this goal by introducing the multi-mon modii diffu- sion model with multi-task learning, referred to MT-Diffusion. MT-Diffusion enables simultaneous modeling and generation of multi-modal data with a unified diffusion model. By multi-taskk we emphasize that MT-Diffusion designed 1) simultaneously generate multi-modal data (potentially

---

Published as a conference paper at ICLR 2024
Forward noising
Reverse denoising
Task-1 Data
Encoder-1
Decoder-1
Task
Scllsmddd date
Aggregation
Forward
Network Backbone
Reverse Denoising
Muliiiak lows
Tesk-2 Data
Encoder-2
Decoder-2
Task-2
loss
Diffusion space
Figure 1: Illustration of the proposed MT-Diffusion on two modalities. The diffusion process is defined in a shared diffusion space for all modality data, which are transformed from the modality- specific encoders. The forward nosing process includes a forward aggregation step that integrates information from mulii-modal data, and the reverse denosing component transforms the diffusion space back to the task-specific data spaces with learnable decoders through a mulli-task loss.
heterogeneous such as images and labels) within a unified model; and 2) seamlessly integrate multi- task learning losses into the diffusion framework in principled manner, supported by theoretical foundations. Our multi-task setting is versatile and applicable to numerous practical scenarios. It worth noting that, as an initial investigation to multi-modal diffusion models, we only focus on two modalities to demonstrate the promise of the research direction, while leaving training with more modalities as interesting future work. Particularly, we construct several practical multi-task generative learning scenarios in experiments to demonstrate the effectiveness of our framework:
Image transition:
We consider jointly modeling multi-modal data, such as images and the corresponding semantic segmentation masks, by learning to generate both in the reverse process within our WTT-Diffusion framework. We design this task synthetic experiment to qualitatively demonstrate the ability of our model on small-scaled datasets.
Masked
training]] Motivated from the previous success masked language pre- training such as BERT (Devlin et al.) 2019a) in language modeling, we propose to combine a pure generation task with masked-image generation task for generative training. We demonstrate on the ImageNet dataset that our model can be more efficient in training a generative model, and can converge to point comparable to (if not better than) the heavily tuned single-task diffusion model in terms of generative image quality. Furthermore, it can simultaneously obtain for free great imtte-restoration ability for masked image recovery.
Joint image-label generation: We jointly model images and the corresponding labels by learning generate both with our MT-Diffusion. We demonstrate on the ImageNet dataset that one can achieve better classification accuracy compared to pure supervised training.
Joint rmage--eprsssntttton generation: We also investigate simultaneously learning to generate images and representations (e.g., CLIP representations (Radford et al. (22TT) with MT-Diffusion. As this is larger-scale setting based on stable diffusion, we only provide qualitative results to demonstrate the ability of our model to generate high-quality images from text, while leaving more detailed investigations as interesting future work.
Our solution for these multi-modal generation problems diffusion model, designed to handle data from multiple modalities through both innovative algorithm and architecture designs in the diffusion forward and reverse processes. Our general idea is illustrated in Figure [ In the forward process, muuli-mmodl/multii- well-designed mechanisms (details in Section 2.2.2) so that the aggregated information can be conveniently applied to the forward noising operation of diffusion model. To deal with potentially heterogeneous data, an effective encoder architecture design is proposed to encode multi-modal data into a shared diffusion space. In the reverse process, we propose to extend the original U- Net architecture in diffusion models to simultaneously reconstruct the multi-modal data from different tasks. To this end, modality-specific decoder heads are designed be attached to the U-Net architecture to decode the diffusion latent code back to multi-moda data spaces. The forward and reverse processes are then integrated within the diffusion mechanism, leading to a loss derived from a new multi-task evidence lower bound (ELBO), as multi-task loss. Extensive experiments on the aforementioned problems are conducted to verify the effectiveness of our framework, demonstrating that our model can achieve simultaneous generation without hurting individual task performance, promising generalization of the standard diffusion model for multi-modal generative learning.
novel generalization of the standard
data are first aggregated through some
"The recent DiffMAR (Wei et al. 2023) is fundamentally different from ours. It a standard conditional diffusion model to denoise the pre-masked region; ours models both image and mask generation.
The default U-Net architecture adopted, though the Transformer (Peebles & Xie 2022 can also be used.
2

---

Published as a conference paper at ICLR 2024
2 MULTI-MODAL DIFFUSION MODELS
2.1 PRELIMINARIES ON DENIOSING DIFFUSION PROBABILITY MODELS (DDPM)
DDPM is a probability generative model that consists of a forward noising process and reverse denoising process operated on a diffusion space. The forward process gradually adds Gaussian noise into the data, which ultimately become standard Gaussian samples; and the reverse process parameterizes neural network model to reverse the forward process. Specifically, given a data sample x from the data distribution, the forward process from time 1 - 1 to time / is defined as (Zz, 1 = (NZz vI - 3r 2:-1, 3y I), where 2, represents a noisy version of the original data sample Zg time t; } is an increasing sequence converging to 1 (making x, converge to standard Gaussian sample). A reverse process is modeled by a neural network (we consider a U-Net for image data) parameterized by 0 as po (2,-1 24) (2 no (Zz, f), Eo t)). Considering all time steps f = 1, ... T, the forward and reverse processes define two joint distributions over the same set of random variables {zo.-.. ZT}. By variational principle, loss corresponding to the evidence lower bound (ELBO) can be derived to optimize the parameterized generative model 0, as
E By(co.. xx)) - pogT Σ121 log 2(2, (2(-1)
p0(21-1
2.2 MOLAI DIFFUSION MODELS
We propose the MT-Diffusion model to jointly model multi-modal data with multi- tast learning by generalizing the DDPM framework. We assume each task is associated with one data modality (the data can be the same for different tasks). For example, an unconditional image-generatin associated with image data, and an image classification task with image-label paired data. Suppose there are N modalities, where modality i is associated with task data from space X, Let x, denote one data sample from the i-th modality space, and let X A {xI,-.. ,Nx) (N) } be the union data from the N modalities. We note that our setting is quite general in the sense that the data spaces {.) can be heterogeneous, e.g., the image space versus the image-label space as from our previous example.
task is
To deal with potential heterogeneity of modality-data spaces, we propose to define MT-Diffusion in shared latent space, called diffusion space and denoted as Z. To this end, we propose to apply a mapping to project each of the original modality-data space onto the shared diffusion space, We define the mapping with encoder E, for task i, iee., E; X - Z, illustrated in Figure] [ For simplicity, we consider non- parametric or fixed-paramets parameter encoders. The specified encoder designs are detailed in Section 2.2.4 In the following, we first formally define the proposed MT-Diffusion by specifying the forward and reverse processes, as well as deriving the corresponding variaitonal lower bound, by extending the DDPM framework handle multiple data sources and multi-tast task losses.
2.2.1 FORWARD-REVERT
PROCESSES AND THE VARIATIONAL LOWER BOUND
In our design, the forward and re- verse processes will be responsible for integrating multi-modal data information and multi-task losses within the DDPM framework, re- spectively. This is implemented by first defining joint distributions over data modalities and the dif- fusion latent variables in both for- ward and reverse processes. Specif- ically, in the forward process, the noising transition from time
Modality data
Forward
Reverse
2o Figure 2: The forward (left) and reverse (right) processes of the proposed MT-Diffusion by jointly modeling a set of task data. to time t is defined to be conditioned on the modality data. To this end, we propose to define a joint distribution at time over and the diffusion latent variable 2,, conditioned on information from time t - 1, to endow the following decomposed forn
2--1 Ze 27
Zo 7/-1 7 ZT
the data X- {x,,".. xy
N
g(z, X 2(-1) = q(z Z,-1, .N)) II (xx),
(I)
where q(z, Z,, xx) represents the transition distribution of z, from time t to time
1, and }x,) denotes prior distributions of the modality data that we assume to be mutually independent for simplicity. We denote this process as forward aggregation, and the specific probability distributions will be defined in the next section. Furthermore, the reverse process is defined by simply
We assume modality data are time-independent, although is also feasible to introduce time dependency.
3

---

Published as a conference paper at ICLR 2024
reversing the forward distributions, resulting in joint distribution pa X 21) time t, where represents the reverse model parameter. Specifically, starting by sampling 27 from p(zr), we propose to decompose the reverse transition time t into the following conditional distributions: po(z;,1, X ( po 1 22) IT= Po po(x; zz). The random variable dependency and the general forward-revers processes are illustrated in Figure 2 Before specifying these distributions, we first derive an objective by matching the joint distributions of the forward and reverse processes. This results in multi-task ELBO for the proposed MT-Diffusion, based on which a final loss can be defined in Section 2.2.5
Theorem 1. The negative ELBO of MT-Diffusion endows: L = B. [Co + Li + L2 + Ca). where
Co - KL (q(zr) zo. x))l((z))
L _ Σ
KL (q(z)-1 Zo: Zt, X) (z 2,)), (2)
N
(>
C2 4
ΣΣ KL (g((x:)lloo(x, 2i)), L3 - log po(z Zo 111). 1m1
Remark 1. We can see that the prior multi-modal data distributions are within the loss term Cz. If only a single generation task is considered, the sub-loss L2 will disappeared, reducing to the standard DDPM loss. Our multi- diffusion objective defines the posterior of the transition probability q(z--1 2t, X) by conditioning on all modality data (in L1), and additionally, as formulated in Lz, parameterizes the reverse process to regularize the predicted modality-data distribution px 2c-1) so that it matches the prior modality data distribution qi (x:).
2.2.2 FORWARD AGGREGATION
The forward aggregation mainly deals with the posterior transition probability q zo: Zi, X) in Ln of equation 2 To derive explicit form, we start by specifying the forward transition probability ((z 14-1,X) X), which can consequently induce the marginal distribution q(z; Zo. X) as well as the posterior transition probability. To integrate different task information, we define the forward transition distribution as a Gaussian distribution by aggregating the task information into the mean parameter. Specifically, we define
2t + E We
(6) (x;:) ,1
at N + 1 N
qZ: Z 1, X) = N Z: V 0
(3)
N+1
where wi denotes the weight for the i-th modality representation at time t, and {a(} are weights to scale the mean and covariance of the Gaussian transition similar to DDPM. By change of notation a4 A a//(N + 1). the transition distribution can be re-written as q(z, zz-1.X) N (26 (2 +w, F(x)), - or) I)). which we will use in the following derivations and imple- mentation. With these transition distributions, multi-task information can be seamlessly incorporated into the diffusion process, which can effectively translate the reverse process with a parametric model to be defined in Section 2.2.3 Now we can derive the marginal and posterior transition distri- butions, which turn out also endow simple forms of Gaussian distributions, stated in Theorem 2 Theorem 2. Given the transition distribution equation 3 the marginal transition distribution follows
((2 20, X) N
215 V Zo
αi Ex(x)), (1 (1
210) P o. where ā₁ A II 0r and (10 is recursively defined as (6) va (wl") + +a(01) with āo
Furthermore, the posterior transition follows follows q(2=-1 Zo, 21, X) = N 1z (2r X), 1) where
Pi [2 X) Va,(1 0v.-1) Z4 +11 - / 20 EN
(1-o.)af)) Vas
(1 -
Fs Fx;));
1 1-
Σ W ( Ex(x:) and B (1 - (z=az)(11 āt-1)
N
21
(5)
Va,
V1
Ctt
1 ā₁
The posterior transition distribution equation shares a similar form as that in DDPM, with extra term of Σ w. (4) E,(x,) representing information aggregated from all tasks (thus aggregation). Note the aggregation is defined in the forward process, enabling a closed-form posterior but without losing too much modeling expressiveness compared to other complex aggregations.
Remark 2.
4

---

Published as a conference paper at ICLR 2024
Encoder E
Gradient update
Identity mapping E(X)
Zg
U-Net
Pretrained generator [X) G(X)
Aggregate
34
Net
Decoder D()
Shared
800)= (1Ne
Figure 3: Training pipeline and encoder-decoder design choices. !20 indicate three possible choices for the encoder E-): gray shaped boxes indicate stop gradients; and black dash lines mean possible connections to the encoder and decoder. "Aggregate" is implemented through equation 4
2.2.3 REVERSE PARAMETRIZATION
Based on the ELBO Theorem the reverse model is responsible for defining two sets of dis- tributions: (2 22) and po [x; Ze). The first distribution is similar to that in DDPM, and the second one is induced from decoding the diffusion latent code back to modality
leverage these distributions within unified architecture for task information sharing, spaces. To propose to parameterized the reverse model with a shared backbone network followed by N extra modality heads, each corresponding to modality. The basic structure is illustrated in Figure | Specifically, for - z,), we follow DDPM to define it as Gaussian distributions with mean and covariance denoted as pe (Z; X) and of I, respectively. Consequently, the KL-divergence L1 of equation 2 reduces to matching the mean of the two Gaussians with proper weighting scheme depending on Based on the form of the mean of g(;,-1 Zo, Zz, X) in equation 5 instead of parameteriz- ing the mean of po (2c-1 2z), we follow DDPM to parameterize the U-Net to predict the intrinsic noise z, (denoted as e). Specifically, the parametrized U-Net model eg (Zc, t) is formulated as: eo (z(,, I) co(Var Zu +ū, X ++VI- c,e e.
For the decoding distributions (x; 2,)'s in the L2 term of equation 2 the distribution forms are modality specific. We consider the following
cases
experiments:
When modality data are represented as probability vectors, e.g., labels in the classification task, we define (x; z;) as a discrete distribution, parameterized by the output of the modality head. Consequently, the KL-term in Ce is equivalent to the cross-entropy loss.
When modality data are in the form of continuous values, we define (x 7x2) as Gaussian distribution with the mean parameterized by the output of the modality head. In this case, the KL-divergence in Le reduces the MSE loss, similar to the case for pe (z,-1 (2c).
It is worth noting that different from (2(21-1 2z). the variables and 2, in ( 22) can be in different feature spaces. Thus, decoder D;(-) the form of one modality head specified above is applied project the latent code back to the modality-data space, based on which proper px(x;] 2.) is defined, as illustrated in Figure Specifically, the decoding process can be written as:
At time f Zt denoising
Diffusion
C₁
A U-Net(2,,
Task , x, A Di(c;; 0) xx;
decoding
where we use "U-Ner" to denote the output from one particular component of the U-Net, serving as the input to the decoder head (see Section 2.2.4 for more details). In other words, the reverse parameterized model consists of two parts: eo (z,, x, t) and (2 t; 0). Detailed structure designs to integrate the decoders (together with the encoder F:()) in the forward process) into the shared U-Net backbone is discussed in the next section.
2.2.4 ENCODER- DECODER DESIGNS
The encoders aim to map different task-data onto the diffusion space, and the decoders project the diffusion latent code from the shared U-Net backbone back to the task-data spaces. As the encoders associated with the forward process, propose avoid introducing extra trainable parameters in the encoders for simplicity. Furthermore, we propose to introduce trainable parameters to the decoders they parts of the parameterized reverse model. There many possible design
5

---

Published as a conference paper at ICLR 2024
choices for the encoders and decoders. Our guideline is to choose architectures to reuse existing components or some pretrained models as best as possible. Based on this principle, we recommend the following designs, with the detailed training pipeline and architectures illustrated in Figure
Encoder Design We consider three scenarios, indicated by 223 in Figure 3 1) A modality-dat space is the same as the diffusion space, e.g., both image spaces. In this case, we can define the encoder as simple mapping such as the identity mapping @ in Figure 3 2) A modality
space inhomogeneous with the diffusion space, e.g., a label space vs. an image space. In this case, we propose to either a pretrained generator (@ in Figure 3 the shared U-Net backbone (3 in Figure 3) to transfer modality-date information to the diffusion space. Particularly, for choice @ we use the cross attention mechanism in the U-Net architecture map modality datal information the diffusion space.
Decoder Design The decoders are modality and task specific. They accept outputs from one of the U-Net blocks (indicated by black-dash-line connections in Figure 3 and learn generate the original modality data. For example, in a classification task, the decoder is designed as a classifier that outputs a class label, associated with a cross-entropy loss.
2.2.5 TRAINING AND INFERENCE
Training We propose a simple training loss for our MT-Diffusion, based on the ELBO equation and the specific forward and reverse parameterization described above. In the ELBO, Lo is inde- pendent of the model parameter 0, thus can be omitted in training. By substituting the specific distributions into the BLBO and adopting the simple losy idea in DDPM (Ho et al. (2020) that ignores the weights for different timesteps, the ELBO equation 2 the proposed MT-Diffusion:
reduces to the following training loss for
L A Emse ΣΣ KL (g. (x;())ll 22)) where ze ~ g(ze| Zo, X),
(6)
and Emse A Egg DDPM; À is the weight scalar (we set it to 0.1 in our experiments). In particular, the KL terms above can endow closed forms depending on the modality-specifii (x:). For example, in a classification task with x; representing labels, both (x.) and pa (x; z,) are defined as discrete distributions, making the KL divergence equivalent to the cross entropy. We apply stochastic optimization for model learning. At each iteration, a random timestep t is first sampled. Then the corresponding 2, is sampled from the forward process with task information aggregation. We then feed 2, to the reverse model to predict the forward noise and the modality data. Finally, gradient descent is applied to update the model parameter based on the loss equation
EEsi (ze,)) - e11 log (zo 21) has the same form as the simple loss in
Inference A distinction of our model is its ability to simultaneously generating multi-modal data. We propose a generic inference procedure that can achieve both unconditional generation (with initially all missing modality data X) or conditional generation (with initially parts of X known). The basic idea is to estimate the potentially missing modality data from the corresponding heads of the reverse model outputs. The specific algorithm is summarized in Algorithm T in Appendix C
3 RELATED WORK
Diffusion-basee Models Diffusion models (Sohl-Dickstein et al. 2015 Ho et al. 2020 Song
al. 202T) have been state-of-the-as art generative models on variety of applications including image
syntheses (Dhariwal & Nichol 2021). text-to-image generation (Ramesh et al. 2021 Saharia et al.)
2022 Yu et al. [022]] Rombach al. 2022), audio generation (Kong et al. [021 Liu et al. (023))
video generation Ho et al. 2022 Harvey et al. 2022, Singer et 2023) and text generation (Austin ctt al. 2021| Li et al. 2022b; He et al. 2022] Gong et al. 2023), etc. All these models, however, only Tocus on a single generation task, in contrast to our multi-task generation.
Manipulating Diffusion Latent Spaces There have been significant efforts to manipulating latent spaces of pretrained diffusion models for a variety of downstream tasks, including text-driven image editing, inpainting, completion and etc (Gal et al.) 2022, Ruiz et al., 2022b Cohen et al.
2022 Kawar et al. 2022, Meng et al.) 2021bb Bau et al. 2021 [ Avrahami et al. 2022b a [Brr Tal et al. 2022 Lugmayr et al. 2022). There also related works to learning more discriminative latent space (Zhang et al. 2022 Preechakul et al. 2022) or manipulate a latent space for better
representations (Kwon al. 2023). These methods, however, are task-driven and do not provide

---

Published as a conference paper at ICLR 2024
Figure 4: Random samples the night2day (top, unconditional generation) and cityscape (bottom, conditional generation) datasets. The 3 pictures in each block the cityscape dataset (bottom) corre- spond to the conditional image (source), the ground-truth and the inferred target image, respectively.
Figure 5: Randomly generated examples of MT- -if-uiffuion with masked-image training. First row: image restoration from random masks; images in each block: original, masked and restored images. Second row: image restoration from half masking; each block contains two restored images to illustrate generation variance. Third row: image generation from scratch with complete masks.
theoretical foundation on the working principles; while all the tasks considered the literature can be accomplished with our MT-Diifusion framework by incorporating the corresponding task data.
Our model is also related to guided diffusion, which discussed details in Appendix D
4 EXPERIMENTS
As first work on multi-modal diffusion models, we focus on evaluating our MT-Diffusion on the 4 multi-task learning settings mentioned in the Introduction, while leaving other more complicated settings such as incorporating more tasks into training as interesting future work. We describe detailed hyperparameter settings, the four tasks and encoder-decoder designs in Appendix D In the four experiments, MT-Diffusion for image translation and joint image-vepresentariun generative modeling mostly for illustration purposes. Thus, the results are mostly qualitative, e.g., Figure demonstrates some visualizations of image-translate More details are deferred to Appendix
4.1 MT-DIFFUSION FOR MASKED-IMACE TRAINING
Task description and encoder- decoder design We propose the masked-image training, new training strategy we de- sign improve image genera- tion with MT-Diffusion. The task motivated from the masked- language pretraining paradigm in NLP models such as BERT (Devlin et al.) 2019a), which achieves significant success with multi-task training. Specifically, in addition to the standard image generation task that generates images from noise, we define an additional task, called a random inpainting task, which learns to recover from randomly masked images. Consequently, the forward process is defined start from clean-masked
Table 1: LPIPS score for masked image recovery. "Mask-m" means masking an image with m patches.
Model
Mask-5
Mask-10
Mask-15 15
Mask-20
Clean-Masked
0.311
0.414
0.461
0.491
SDEdit- Meng et al. 2022
0.400
0.466
0.497
0.513
MT MT-Diffusion
0.035
0.068
0.099
0.133
7

---

Published as a conference paper at ICLR 2024
1St
FID]
sFID|
Precision
Recallt
ADM (un-cond)
15.64
23.22
16.53
0.57
0.60
ADM (class cond)
22.69
16.35
17.49
0.59
0.62
2.06
Generation with Masked-Iman Image Training (Section
4.1
2.81
MT-Diffusion-U 23.77
26.00
25.22
0.57
0.55
MT-Diflusiun-X 34.53
9.85
15.78
0.68
0.64
Generation with Joint Image-Laber Modeling (Section 4.2
280
MT-DiffusionM 15.66
33.92
21.63
0.54
0.55
2000
2000
2000
1000
5000
BOCD
MT-Diffusion-M*
26.31
13.48
16.64
0,66
0.61
Rieetuned lberstion
MT- NT-Diffusion-F:
11.86
41.00
22.37
0.48
0.45
Figure 6: Classification accuracies vs. supervised finetuning itcrations the Im- ageNet validation data.
MT-Diffusion-E'
24.78
11.28
16.08
0.74
0.56
Table 2: Training efficiency comparisons on ImageNNet664.
image pair, and gradually add noise to the pair in the forward process. To create the masked images, we randomly sample m ~ Uniform(0--. .. 10) patches of size 16 x 16 from the original image and mask them out with zero pixel-values. These patches are placed randomly so they might overlap with each other. We adopt two choices for the noise prediction network co (-): the first only considers (z,, t) as input, denoted as MT-Diffusion-U; the other takes (z,, t),, denoted as MT-Diffusion-X. The former is more specifically designed for unconditional generation, whereas the latter is more suitable for constrained image restoration, a task we defined as inpainting a randomly masked image while keeping the unmasked region unchanged. Note current state-of-the -f-te-----e-rt diffusion models do not directly deal with this problem. Existing work for image editing and inpainting such as SDEdit (Meng et al.) 2021a) and Dreambooth (Ruiz et al. (022a) do not explicitly enforce unmask region consistency, thus their inpainting results can change the unmasked region. By considering simultane- ously generating the
image pairs with MT-Diffusion, our model can maximally learn to maintain the consistency of the unmasked regions. Similar to the image- trancttranniiion experiment, we use the identity map the encoder, and replicate the output block of the original U-Net as the additional masked-image decoder, which consists of a normalization layer and a convolution layer.
Results We implement our method based on the guided diffusion codebase (dif) on ImageNet (Deng et al.) 2009). We first demonstrate that framework can help to improve the training effi- ciency for image generation. To this end, we compare our and the ADM models (Dhariwal & Nichol [02T] be- fore convergence at 1M iteration at im- age resolution 32. For a fair compari- son, all models are trained from scratch with exactly the same hyperparameters (thus the numbers are not directly com- parable to the reported values). We adopt the same evaluation metrics the ADM under 5K samples. Quanti- tative results are shown in Table Note MT-Diffusion-U is designed to gener- ate images from complete random noise while MT-Diffusion-X needs a condi- tional masked input, which we simply define as a complete mask (all zeros) for this purpose. It observed both the two variants significantly outperform the unconditioned ADM, indicating the training efficiency and modeling effectiveness of multi-task generative learning via masked-image training. In addition, MT- MT-Diffusion-X is found to perform better than MT-Diffusion-U. We hypothesize this is because the conditional masked-ing -image information makes the training of the model easier and more effective.
Table 3: Generated image comparisons on ImageNet- 128. The subscript "g" means classifier guidance; the super- script "** on ADM indicates results from the released checkpoint; ADM without "** is the one continued train- ing on the released checkpoint. MT-Diffusion represents the MT-Diffusion-X version; superscript "f"" indicating continued finetuning on single-task image generation; and the superscript "*" indicates results from the constrained image restoration task with 10% random masking (not directly comparable with the other settings).
1st
FIDJ
sFID
Precision?
Rocallf
ADM
79.95
8.46
4.92
0.67
0.66
ADM
151.10
| 3.56
4.63
0.79
0.57
ADM
74.04
| 9.62
5.61
0.66
0.64
ADM
140.89
4.27
5.58
0.79
0.55
MT-Diffusion
78.26
8.42
5.89
0.67
0.65
MT-Diffusion,
81.06
8.06
5.83
0.68
0.65
MT-Diffusion'
84.22
7.01
5.99
0.69
0.64
MT- UT-Diffusson!
171.65
3.51
5.73
0.83
0.53
MT- MT-Diffusion;
135.35
2.15
3,86
0.72
0.68
138.21
2.02
3.84
0.73
0.69
In addition to pure image generation, one unique property of MT-Diffusion-X X is its ability to simultancously perform constrained image restoration. To this end, we randomly mask out some testing images with (5, 10, 15, 20} patches and leam to restore the masked patches. We expect
8

---

Published as a conference paper at ICLR 2024
MT-Diffusion-X X to restore masked images without changing unmasked regions. Some example results are illustrated in Figure [ which clearly demonstrate the strong ability of MT-Diffusion-X for constrained image restoration as well as generation from scratch. For quantitative evaluation, we adopt the LPIPS score (Zhang et al.] 2018) that measures the semantic similarity of the original the restored images, and compare our method with one simple baseline from Meng al. (2022), denoted as SDEdit-. The results are shown in Table []] where the row of "Clcan-Maske Masked" denotes the LPIPS scores of dlean-masked image pairs that we include for reference. It is clear that MT-Diffusion-X obtains scores closed to zero, indicating the closed similarity between restored images and original images. SDEdit,, on the other hand, obtains very high LPIPS scores, which are even higher than the Clean-Masked baseline. This is expected since SDEdit- is not specifically designed for such task.
Finally, to demonstrate the ability of our model to generate high-quality images at convergence, we compare our method with ADM for pure image generation, under a resolution of 128. We train our MT-Diffusion-X from scratch. We evaluate the ADM with two versions: the released checkpoint and continued trained version from the checkpoint with the same hyperparameters our model, for a more fair comparison. The results are shown in Table 3 evvaluated on 50K samples. Our method achieves comparable performance than ADM (if not better some metrics such the IS), while being able to perform more tasks such as the constrained image restoration demonstrated above. is also noted that the continue-trained version of ADM (started from the released checkpoint) is slightly worse than the released checkpoint, indicating the latter might have been tuned for best performance, and thus it is more fair to compare our methods with the continue-trained version of ADM.
4.2 MT-DIFFUSION FOR JOINT IMAGE-LABEL GENERATION MODELING
Task description and encoder-decoder design This is a more heterogeneous case as labels and images are in different spaces. We follow our design principle to use the diffusion U-Net as the encoder for labels via cross attention in the U-Net. The label decoder is an additional head from some layer of the U-Net. We consider two designs: 1) Add one additional MLP layer out of the middle block of the UNet to map the diffusion latent space onto the label space. This introduces minimal extra parameters into the original reverse model but might enforce some discriminative information not helpful for pure image generation. We denote this variant as MT-Diffusion-M. 2) Add pre-defined classifier at the end of the U-Net output. Since the U-Net output can be used to reconstruct the original image, this structure essentially makes the learning of image generation and classification in a sequential manner. In the experiments, we use the classifier provided in the guided diffusion codebase (dif) as the label decoder. We denote this variant as MTT-Diflusio-EE
Results We adopt the same experiment setting as the previous section. In addition to measuring the generated image quality, we also measure the classification performance using the classifier in MT-Diffusion-E. We find that for such a heterogeneous setting, continuing finetuning both the generator and classifier with single generation and classification tasks, respectively, can significantly improve single-task performance. Consequently, we adopt similar idea of classifier-free guidance to simultaneously learn pure-generation model along with the multi-task training with the shared U-Net. We denote the finetuned models as MT-DDifusionn-. reported in Table It is observed that MT-Diffusion-N. performs better than MT-Diffusion-E. which slightly under-petn perform the single-task ADM. This is expected and indicates that learning to generate heterogeneous data can trade off single-task performance. We also believe the performance gap is partly due to the un-tuned sub-optimal hyperparameter setting, With single-task finetuning, we can see that both variants outperform ADM in all metrics. To continue finetuning the classifier, we use the training and evaluation script from the codebase (dif), and compare with the pretrained classifier in classifier-guidanee
and MT-DDfflusio-E" The results are
ADM (dif). We continue finetuing the classifier from the released checkpoint as baseline. Top-1 1 and top-5 accuracies are plotted in Figure 6 It is observed that our classifier consistently outperforms the baseline, although the gap turns smaller with increasing finetuning steps.
5 CONCLUSION
We propose the multi-modal diffusion model with multi-task training, generalization of the standard diffusion model for multi-modal generation. Our model is general and flexible, which can incorporate potentially heterogeneous modality information into a unified diffusion model, compared to training on a single-task setting. We define several multi-task generative problems and test them on our proposed MT-Diffusion. Extensive experiments are performed to verify the effectiveness of our proposed framework. Interesting future works include improving the framework by better network- architectures designs and applying the method to more diverse multi-modal and multi-task settings.

---

Published as a conference paper at ICLR 2024
REFERENCES
Guided diffusion.
http://ttt////aa.ooooo//o//aaaaa futp Accessed: 2023-
05-11.
Latent Diffusion Models. hittess//gttnbb.ccm///att0///000 x 05-11.
Accessed: 2023-
Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne den Berg. Structured denoising diffusion models discrete state-spaces. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL ht.tps://openrevilt net/ ott.eee...ee/eoom2moo// wit.mmmiiidhhhiiPa..
Omri Avrahami, Ohad Fried, and Dani Lischinski. Blended latent diffusion. arXiv preprint arXiv:2206.02779, 2022a.
Omri Avrahami, Dani Lischinski, and Ohad Fried. Blended diffusion for text-driven editing of natural images. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 18208-18218, 2022b.
Zhipeng Bao, Martial Hebert, and Yu-Xiong Wang. Generative modeling for multi-task visual learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.). Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 1537-1554. PMLR, 17-23 Jul 2022. URL
mlr wtw.s//psss...e.ee....... /ttpp//wweeeaaaaooooooorerrerraa
Omer Bar-Tal Tal, Dolev Ofri-Amar, Rafail Fridman, Yoni Kasten, and Tali Dekel. Text2live: Text-driven layered image and video editing. In European Conference on Computer Vision, pp. 707-723. Springer, 2022.
David Bau, Alex Andonian, Audrey Cui, YeonHwan Park, Ali Jahanian, Aude Oliva, and Antonio Torralba. Paint by word. arXiv preprint arXiv:2/031 1095 2021.
Andreas Blattmann, Robin Rombach, Kaan Oktay, Jonas Müller, and Björ Ommer. Semi-pane parameeraeetiic neural image synthesis. In Advances in Neural Information Processing Systems, 2022.
Zitian Chen, Yikang Shen, Mingyu Ding, Zhenfang Chen, Hengshuang Zhao, Erik Leared--Mill and Chuang Gan. Mod-squad: Designing mixture of experts as modular multi-task learners. ArXiv, abs/2212.08066, 2022. URL. uttttstss:///et 263793067
semant i cscho ar org/Corpus
Jooyoung Choi, Jungbeom Lee, Chaehun Shin, Sungwon Kim, Hyunwoo J. Kim, and Sung-Hoon Yoon. Perception prioritized training of diffusion models. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11462-11471, 2022. URL https
Niv Cohen, Rinon Gal, Eli A Meirom, Gal Chechik, and Yuval Atzmon. " this my unicorn, fluffy": Personalizing frozen lasion-language representations. arXiv preprint arXhv:2204.01694 2022.
Marius Cordts, Mohamed Omran, Sebastian Ramos. Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schicle. The cityscapes dataset for semantic urban scene understanding, In Proc. the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248-255. leee, 2009.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutamova. Bert: Pre-training of deep bidirectional transformers for language understanding. ArXiv,
2019a.
10

---

Published as a conference paper at ICLR 2024
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-train Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019b. Association for Computational Linguistics. doi: 10. 10.18//3/11/11111 1423.. URLhttps https wwclannholog..org///oooooo
1423
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780- 8794, 2021.
Heshan Fernando, Han Shen, Miao Liu, Subhajit Chaudhury, Keerthiram Murugesan, and
Tianyi Chen. Mitigating gradient bias in multi-objective learning: provably convergent approach. In International Conference on Learning Representations, 2023. URL https //api. www.steiinnncoooaa.o.....ccoc
Rinon Gal, Yuval Alaluf. Yuval Atzmon, Or Patashnik, Amit H Bermano, Gal Chechik, and Daniel Cohen-Or. An image is worth one word: Personalizing text-to-image generation using textual inversion, arXiv preprint arXiv:2208.016// 0/618, 2022.
Hyojun Go, Jinyoung Kim, Yunsung Lee, Seunghyun Lee, Shinhyeok Oh, Hyeongdon Moon, and
Seungtaek Choi. Addressing negative transfer in diffusion models. ArXiv, Abss22306000344 2023.
URL https::/pppp semant escho ar. org/ Corpus 2.8 999775
Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and Lingpeng Kong. Diffuseq: Sequence to sequence text generation with diffusion models. In The Eleventh International Conference on Learning Representations, 2023. URL htttps:///ppeneeettt wwwpeewwwwwwweeeeeeeooeeeeeeeeee wtw.:neeee//tttteeeeee
FLVXs
Xizewen Han, Huangjie Zheng, and Mingyuan Zhou. CARD: Classification and regression diffusion models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL htLp:tts: httpt::://pppreeeeeeee http://////peeeeeeeeeeeeeeeeeeee formm? id=412zYEJ9d
Tiankai Hang, Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, and Bain- ing Guo. Efficient diffusion training via min-snr weighting strategy. 2023 IEEE/CVF In- ternational Conference on Computer Vision (ICCV), pp. 7407-7417,
2023. URL https
//api. www.sttttaann.oooeeooeee hwwp://wi....o.....ooooooo/22//a
William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Dietrich Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL
https: //tpr//eeeeeeeeeeeeee net ott.p//wweeee//ooomeddddd-or--. MoorumiiddORTuuuviii
Zhengfu He, Tianxiang Sun, Kuanning Wang, Xuanjing Huang, and Xipeng Qiu. Diffusion- bert: Improving generative masked language models with diffusion models. arXiv preprint 1.50092 2022.
Falk Heuer, Sven Mantowsky, Syed Saqib Bukhari, and Georg Schneider. Multitask- centeemet (men): Efficient and diverse multitask learning using an anchor free approach. 2021 IEEE/CVE International Conference on Computer Vision Workshops (ICCVW), pp. 997-1005, 2021. URL wwt/CorppssI::ao..ooooo :D:
httpps://at http://///ppl.ssenniiicsseeeo
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 6840-6851. Curran Associates, Inc., 2020. URL https::/pproeeednngg www..pp.cc/ppap.epppppppeeeeppep Balee//ADcccccBBBa00000011111100
files/papez/2022 www.sd--appppppp.p....
Jonathan Ho, Tim Salimans, Alexey A. Gritsenko, William Chan, Mohammad Norouzi, and David J. Fleet. Video diffusion models. In ICLR Workshop on Deep Generative Models for Highly Structured Data 2022. URL httthttss:/:pppreeeee
net wwtt//oouttideeeeeee R2NdD25
11

---

Published as a conference paper at ICLR 2024
Yi Hu, Jiazhi Yang, Li Chen, Keyu Li, Chonghao Sima, Xizhou Zhu, Siqi Chai, Senyao Du, Tianwei Lin, Wen Wang, Lewei Lu, Xiaosong Jia, Qiang Liu, Jifeng Dai, Yu Qiao, and Hongyang Li. Planning-orientn autonomous driving. 2023 IEEE/CVF Conference Computer Vision and Pat- tern Recognition (CVPR), pp. 17853-17862, 2022. htt.sstps:://eet semant cccho ar
#Tt//oorpuss1:2257874444
Gabriel Ilharco. Marco Tulio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi. Editing models with task arithmetic. ArXiv, abs/221 1222222 2022. URL htRps::::: /api. weemanticsccool. wwtpaaaccccccc.a.oooooooirooo.ii :ttg:/Cooous::::440000
Phillip Isola, Jun- Yan Zhu, Tinghui Zhou, and Alexci A Efros. Image-to-image translation with conditional adversarial networks. CVPR, 2017.
Junguang Jiang. Baixu Chen, Junwei Pan, Ximei Wang, Liu Dapeng, Jie Jiang, and Mingsheng Long. Forkmerge: Overcoming negative transfer in multi-task learning. ArXiv, URL https://api.
2023.
icscholla. wttp/oorpu11::25530202.
Bahjat Kawar, Shiran Zada, Oran Lang, Omer Tov, Huiwen Chang, Tali Dekel, Inbar Mosseri, and Michal Irani. Imagic: Text-based real image editing with diffusion models. arXiv preprint arXiv:2210.09272 2022.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations, 2021. URL htt.hhtps://// /ttttp::///pereviii
net/ www.mmriiioooooo d=a xFK8Ymz5J
Mingi Kwon, Jaeseok Jeong, and Youngjung Uh. Diffusion models already have semantic latent space. The Eleventh International Conference on Learning Representations, 2023. URL
ptttps://openreviie net. h
Pierre Yver Laffont, Zhile Ren, Xiaofeng Tao, Chao Qian, and James Hays. Transient attributes for high-level understanding and editing of outdoor scenes. ACM Transactions Graphics (proceedings of SIGGRAPH), 33(4), 2014.
Junnan Li, Ramprasaath R. Selvaraju, Akhilesh Deepak Gotmare, Shafiq Joty, Caiming Xiong, and Steven Hoi. Align before fuse: Vision and language representation learning with momentum distillation. In NeurlPS, 2021.
Junnan Li, Dongxu Li, Caiming Xiong, and Steven Hoi. Blip: Bootstrapping language-image pre-training for unified lisguae understanding and generation. In ICML, 2022a.
Junnan Li, Dongxu Li, Silvio Savarese, and Steven Hoi. BLIP-2: bootstrapping language-image pre-training with frozen image encoders and large language models. In ICML, 2023.
Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori Hashimoto. Diffusion- LM improves controllable text generation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances Neural Information Processing Systems, 2022b. URI. net hott:/oooum1dd3351r1r5.
Haohe Liu, Zehua Chen, Yiitan Yuan, Xinhao Mci, Xubo Liu, Danilo Mandic, Wenwu Wang. and MarkD Plumbley. Audioldm: Text-to-audio generation with latent diffusion models. ArXiv, Abs/2301.12503, 2023.
Yang Liu, Zhaowen Wang, Hailin Jin, and Ian James Wassell. Multi-task adversarial network for disentangled feature learning. 2018 JEEE/CVF Conference on Computer Vision and Pattern Recog- nition, pp. 3743-3751, 2018. URL httph::///eaat semant segaantccccooar. org/ Co CorpsIID 52203209
Alexander Long, Wei Yin, Thalaiyasingam Ajanthan, Vu Nguyen, Pulak Purkait, Ravi Garg, Alan Blair, Chunhua Shen, and Anton den Hengel. Retrieval augmented classification for long-tail visual recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6959 695996969, 2022.
12

---

Published as a conference paper at ICLR 2024
Andreas Lugmaar, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 1p..1146661 11461-11471, 2022.
Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun- Yan Zhu, and Stefano Ermon.
Sdedit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations, 2021a.
Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun- Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations. In International Conference Learning Representations, 2021b.
Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun- Yan Zhu, and Stefano Ermon. SDEdit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations, 2022.
Kevin P. Murphy. probml.
Probabilistic Machine Learning: An introduction. MIT Press, 2022. URL
OpenAI. Gpt-4 technical report. ArXiv, abs/2303.08774, 2023.
Byeongjun Park, Sangmin Woo, Hyojun Go, Lin-Young Kim, and Changick Kim. Denois- ing task routing for diffusion models. ArXiv, abs/2310.07138, 2023. URL https://api semanticscholaz: wwtpooooooIIi
William Peebles and Saining Xic. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748, 2022.
Hoang Phan, Ngoc Nguyen Tran, Trung Le, Toan Tran, Nhat Ho, and Dinh Q. Phung. Stochastic multiple target sampling gradient descent. ArXiv, abs/2206,0 0199444 2022. URL httttts:://ar wwt.or/////o..o.
gemanticscho
wttpg///rruuuuuie ID: 249395631
Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwonge and Supasorn Suwajanakorn. Diffu- sion autoencoders: Toward meaningful and decodable representation. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and llya Sutskever. Learning transferable visual models from natural language supervision. International Conference on Machine Learning, 2021.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. arXiv preprint 102.::2002.10922 2021.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. ArXiv, abs/2204.06 1225,, 2022.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Confer- ence on Computer Vision and Pattern Recognition (CVPR), pp. 10684-10695, June 2022.
Sebastian Ruder, Joachim Bingel, Isabelle Augenstein, and Anders Søgaard. Latent multi-task architecture learning. In AAAI Conference on Artificial Intelligence, 2017. URL https://app
ssmanticscho.. org/ ID:pusID:1119888
Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation. ArXiv, abs/2208. 12242, 2022a.
Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation. arXiv preprint arxiv:2208. /2242, 2022b.
13

---

Published as a conference paper at ICLR 2024
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li. Jay Whang. Emily L. Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, Seyedeh Sara Mahdavi, Raphael Gontijo Lopes, Tim Salimans, Jonathan Ho. David Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. ArXiv, abs/2205. 1487, 2022.
Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, Patrick Schramowski, Srivatsa Kundurthy, Katherine Crowson, Ludwig Schmidt, Robert Kaczmarczyk, and Jenia Jitsev. Laion-5b: An open large-scale dataset for training next generation image-text models. ArXiv, abs/2210.08402,, 2022.
Dmitry Senushkin, Nikolay Patakin, Arseny Kuznetsov, and Anton Konushin. Independent component alignment for multi-task learning. 2023 IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 20083-20093, 2023.
org/ Sxg//orpusID:2200 22:0099888200000888
Mohit Sharma, Claudio Fantacci, Yuxiang Zhou, Skanda Koppula, Nicolas Manfred Otto Heess, Jonathan Scholz, and Yusuf Aytar. Lossless adaptation of pretrained vision models for robotic manipulation. ArXiv, abs/2304.006000, 2023. URL httttt/t:://aa semant.eat....e semant.cscho" cscho ar org/
Corpus 25: 25810873
Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin, Jie An, Songyang Zhang, Qiyuan Hu, Harry Yang, Oron Ashual, Oran Gafni, Devi Parikh, Sonal Gupta, and Yaniv Taigman. Make- a-ideo- Text-to- video generation without text-video data. In The Eleventh International Conference on Learning Representations, 2023. URLhtttss: wwwpeeeeeiiieeeeeeeeeeeeoe.eeeee net/ forum wid.Jffyddnfff ff Dvgz
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 2265, Lille, France, 07-09 Jul 2015. PMLR.
Jiaming Song, Chenlin Meng, and Stefano Enmon. Denoising diffusion implicit models. In Inter- national Conference on Learning Representations, 2021. URL /ttp///peeeveeeeeeeeeeeeeeeeeeee net/
forum? 1d= 1 CCHLP
Malik Tiomoko, Hafiz Tiomoko Ali, and Romain Couillet. Deciphering and optimizing multi-task
learning: a random matrix approach. In International Conference on Learning Representations, 2021. BtRLhttps://// seep::///.aa..aaaa.....mme www.ettitnnccoaaaeeeemecccccc :
Nilesh Tripuraneni, Michael I. Jordan, and Chi Jin. On the theory of transfer learning: The importance of task diversity. ArXiv, abs/2006.1 1600, 2020. URL httptpss:/pppi..eaatccccece
org/Corpus ID: 21 99660033
Haoxiang Wang, Han Zhao, and Bo Li. Bridging multi-task learning and meta-learning: Towards
efficient training and effective adaptation. ArXiv, abss2/00.00911,, 2021. URL httpss://ppl
isemanticscholar wwtp/oopps/oo:aa......om 23:444:
Wenhai Wang, Jifeng Dai, Zhe Chen, Zhenhang Huang, Zhiqi Li, Xizhou Zhu, Xiaowei Hu, Tong Lu, Lewei Lu, Hongsheng Li, et al. Internimage: Exploring large-scale vision foundation models with deformable convolutions. arXiv preprint arXiv:2211.05778, 2022.
Chen Wei, Karttikeya Mangalam, Po- Yao Huang. Yanghao Li, Haoqi Fan, Hu Xu, Huiyu Wang. Cihang Xie, Alan Loddon Yuille, and Christoph Feichtenhofer. Diffusion models as masked autoencoders. ArXiv, abs/2304.03223, 2023.
Sen Wu, Hongyang Zhang, and Christopher Ré. Understanding and improving information transfer in multi-task learning. ArXiv, abs/2005.00944, 2020. httpltttt:://tt semanticschh a
axt/Cocpus1D::2137555
Xingqian Xu, Zhangyang Wang, Eric Zhang, Kai Wang, and Humphrey Shi. Versatile diffusion: Text, images and variations all in one diffusion model. 2023 JEEE/CVF International Conference Computer Vision (ICCV), pp. 7720-7731 2022. URL https: /api. semanticschol..
brc//oorpus1D:25352337
14

---

Published as a conference paper at ICLR 2024
Hanrong Ye and Dan Xu. T ask p rompter: S patial hannel ulti -t ask p rompting for d ense 8 cene nderstanding. 2023. URL/ttts:: /tpppaapii.a. semanticscholatt arg/CorpussD)
259905659
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, Benton C. Hutchinson, Wei Han, Zarana Parekh, Xin Li, Han Zhang, Jason Baldridge, and Yonghui Wu. Scaling autoregressive models for content-rich text-to-image generation. ArXiv, abs/2206.10789, 2022.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness deep features as perceptual metric. In CVPR, 2018.
Zijian Zhang, Zhou Zhao, and Zhijie Lin. Unsupervised representation learning from pre-trained diffusion probabilistic models. arXiv preprint arXiv:2212. 2022.
Jun- Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In Computer Vision (ICCV), 2017 IEEE International Conference on, 2017.
15

---

Published as a conference paper at ICLR 2024
A
ELBO OF THE PROPOSED MULTI-TASK DIFFUSION MODEL: THEOREM 1
We give detailed derivations of our multi-task diffusion ELBO in the following:
N
po 2z Ze I[ Po
(tt)
C = Eg
log p(z)) T
log
log
((x:)
((z Z;-1,
im
12
NV
po
(2 00(22-1
2t) [T Po (x;
Ze) q(21-1 2o. X)
Eg
logp(zz) T
log
log
(xx)
gize-1 Ze, 20, X)
g(Z, 20, X)
1=1
(>1
(zol (1 I[/= px(xx (11)
log
9(21 no; X)
p(ZT)
po [2 2z))
Bgg
log
log
9 1z
Zu; X
a 22, Zo, X)
121
(x 2x)
log
lo 121)
(xx:)
121 i=1
where we use the fact that
q(z Ze 1, x) = ((z 11-1,, 20) g(z=-1 Ze, Zo, x)q(Zz zo. x)
((2 200 x)
the second equality.
B CALCULATING THE POSTERIOR DISTRIBUTIONS: THEOREM 2
In our derivation, will frequently the following well-known property of Gaussian random
variables.
In the following, first present Lemma calculating the posterior distribution of Gaussian random variables, based on which we derive the posterior distribution of our forward process.
Lemma 3. Let 61 NN I), e2 NM12, I). Then, for Va 2 2 the random variable
A dei beq follows:
N~ (apil + buz, (a² +bb0))1)
Note in the forward aggregation, 2,-1 En wi E, (x,) = 21-1 :wm
(EN (x e
2e-1 tw?" E(X), where we define E(X) 2 Σ ب.!!! (xx)). This is equivalent the form of
considering only one extra task with modality-dait embedding E(X), egg., it suffices only consider two tasks in the proof. Consequently, in the following, we give the derivations of the forward posterior distribution with one additional task, which case we will drop the task index in x. Generalizing to N task is straightforward. To derive the marginal distribution g(2, 20, x), let 6 N (e; 0, I) for Vt. For the forward process, we have
2e = (arlzz-1 (w,-1+),E(()). + -
Vai 1//-1 21 14--2+00000 (1 + /1 - 01116-111111 + wxE(x) + VI are
Varar11 Ze-2 / + W (0 E(x) + vI 01:0--11
-orai-1 (/o12 24- tw;
/ 04-261-2
+ +Wiyaa + 1 /a E(x) + /I a,00-16
=Vo-0t-104- 10:0--14-220 2 + (w) V + W4-1 0,:4-1 + w/ 2 00:04-1000+- F(xx)
/1 -,0--104-2
16

---

Published as a conference paper at ICLR 2024
where apply Lemma Et-1 and Er into c.
3 in the third equation to consolidate the two random Gaussian variables
Let X - IT Ci d4 = Ei=1 Ii jjj and define ãy = 0. We have
āi vou (w + āt-1); and
glz Zo, x) N (zei vā, Zo + Exx, (x)
As special case, if define the forward transition distribution by letting w, 1, we will have:
ã Vo (1 + ât-1).
(8)
Lemma
Murphy 2022)), ). Define the following distributions for the prior and likelihood:
p(x) N (x; (x ply xx N (y;
Let (A A Then the posterior follows:
p(xly) N xx: A L -b) AAu) .1)
In our case in equation - E(x), L 2
we have μ A / Zo 1
1 I. A - VaiL I
1-û;-1
I I. and (1-ax)1- 1) I. Thus, the posterior q(z:- 1-
Zt. Zn, x) =
N (e 20, x), 1, where
(0, = va,(1 a Ze +(1 - va) 20 + ((1 ar)ā - ā B(x)
1- az
Va ãi-1) Zt +(1 (4)) Var- 20 + ((1 - ar) V (1 a.))ur) E(x)
1 αq
B, (1 1=)a))(1 - From the marginal distribution ((zz 20, x). we have
1 ū₁
Van 20 +ā; F(xx) + v1 04f6
+ 20
V (zt - - V1 O4.E)
(10)
Substituting equation into equation (2,, x, t) Va,(1 a(-1) +(1 (o)
9 we have
(ze -ā,E(xx) VI-
Van
1 ā-
(1-ar]ā /ar
-(1 (1-ār) E(x)
+
1 1- 81
(Va,(1 - -(t-1) + (1- ort) /ar))
2t -(1 -1x),, -(1 -11-at)
(1-ā,)ā--
Gr
1- ār
1- ae
2,
w. E(x)
Vo,
Thus, similar to the standard DDPM, we can parameterize the backward denoising process with neural network to predict the added noise except that in our case, the neural network eg would take (Ze, x, 1) as the input, i.e.,
(0 x,1) co(vor 2n +ō v oie, oie, x. Rs
17

---

Our MT- Wi--iiffusion formulation endows closed connection with the classifier guidance and classifier- free guidance mechanisms. Specifically, from equation 5 if one defines the encoder E(-) as the gradient from pretrained classifier, the posterior mean recovers the one for classifier guidance. By contrast, if one defines the encoder with the reverse U-Net, the posterior mean calculation in equation recovers the classifier-free guidance mechanism. However, an important difference the forward process, where our framework is designed to aggregate information from different encoders for multi-task learning, whereas both classifer guidance and classifer-fre guidance do not. Overall, method constitutes broader framework that can be applied to different scenarios, including generation Published as a conference paper at ICLR 2024
C INFERENCE
The inference algorithm is shown in Algorithm D set to T 1000.
If not explictly stated, the number be timesteps is
Algorithm 1 MTf-Diffusion Inference
1: 2T ~O I)
2: if' Modality data x, (Vi) not initially available then
3:
Randomly initialize modality data x;
4: end if
5:
for TT T.... do
6:
E~ N, i > 1, else e; El Vi
7:
D Get modality data encoding Get estimated noise and predicted modality data
8:
co t), U-UNNee(zz
9:
x; (Vi) not initially available then
10:
Update x; from the new X. Vi
11: end
2 Vo, (z--ai I-a 1-l (2 x)) - Ein w o +he VI-a.
b Update diffusion latent
13:
end for
14: return Zo:
D RELATED WORKS
Connections to Classifier and Classifier-Free Guidance Guided diffusion models aim to leverage prior knowledge from various guidance information for better controllable generation. For example, the classifier guidance method uses the gradient of pretrained classifier to perturb the reverse process to generate from a dlass-conditions distribution (Dhariwal & Nichol 2021). The classifier guidance simultaneously learns guidance model using the same generation network of the diffusion model (Ho & Salimans, 2022). The works that try to utilize external data such as the retrieval- augmented based methods (Blattmann et al., 2022: Long al. 2022) can also be considered special type of guidance. Although the final formulation has some connections with our method (see Appendix D). guided diffusion models essentially only handle single generation task. Our method, on the other hand, can model multiple tasks within a unified diffusion model.
image transition, masked-image pretraining, joint image-label and investigated in the experiments.
Multi-Task Learning Multi-Task Learning (MTL) paradigm in machine learning that involves training model to perform multiple tasks simultaneously, with the idea that knowledge gained from one task can help improve the performance on other related tasks. Recent development has mainly focused on multi-task learning for predictive models instead of generative models. Apart from investigating theory in multi-task learning [Wang al.) 2021 Tiomoko et al., 2021 TTipuraneni et al., 2020) Wu et al. (220)), many existing works explore different techniques to boost model performance with multi-task learning, including but not limited to architecture designs (Heuer et al..
2021 Ruder et al. 2017 Ye & Xu 2023 Sharma et al.) 2023 Chen et al.] 2022). optimization algorithms (Senushkin et 2003 Fernando et al. al.] 2023] Jiang et al.) 2023] Phan et al. 2022) and
task relationship learning (Hu et al. 2022 Ilharco et al. (022). Recent research interest has also be
expanded to applying multi-task learning in generative models (Bao et al.) 2022 Liu et al. 2018).
However, the generative models are limited to more traditional models such VAH and GAN. And
18

---

Published as a conference paper at ICLR 2024
there is limited work studying multi-task learning for diffusion models. Our work represents of the first works on integrating multi-modal generation with multi-task learning in a diffusion model, aiming to further improve generative performance and expand the scope of state-of-the-ee diffusion models.
More Recent Development on Diffusion Models There is some recent effort trying to develop multi-task diffusion models. For example, the Versatile Diffusion (Xu et al.) 2022) The Versatile diffusion focuses developing
neural architectures that make different tasks interact with each other within the single-task diffusion framework. This is different from our work in that ours not only introduces novel neural architecture but also generalizes the single-task diffusion in terms of the loss function. We believe the Versatile diffusion architecture can be incorporated into our framework for more flexible modeling.
Previous efforts have also focused on efficient training of diffusion models, e.g., P2- P2-Weighting (Choi al. 2022), Min-SNR (Hang al., 2023), ANT al. (2023), and Task Routing (Park al. 2023). Our model is orthogonal to these works with no technical overlap. Consequently, we believe there room to incorporate these techniques into our framework for further improvement, an interesting
future direction to be explored.
E DETAILED EXPERIMENTAL SETTINGS AND EXTRA RESULTS
In addition to evaluating some other datasets that will be described in the specific tasks, mainly rely on the ImageNet-IK dataset (Deng et al. 2009) with resolutions of 64 x 64 and 128 x 128, where we adopt the pre-defined training and validation splits. All experiments are conducted on a A100 GPU server consists of 8 GPUs, with a batchsize of 64, if not explicitly specified. When evaluating generation quality, we follow and adopt the popular Inception Score (IS), FID score, sFID score, Precision and Recall metrics (dif), calculated on 10K or 50K samples, where the former is for computational efficiency and latter for comparing with existing results. We note that due to our different hyperparameter settings (specified in the Appendix), some of our results are not directly comparable to some reported results in previous works. For fair comparisons, we rerun some of the baselines our settings that are consistent with our method. One additional hyperparameter of
our model the task weights equation which we set to w = 1/(11000 +) to mitigate some potentially negative influence from some heterogeneous tasks on the generated image quality when t small. We follow most of the parameter settings as in the codebases.
The training procedure is summarized in Algorithm 2
Algorithm 2 MT-Diffusion Training
1: repeat
2: Zo (zo). X (X)
Sample zo and modality data X
3: ~ Scheduler ... ,
4:
{e:} = F(X)
Get modality data encoding
5:
<~ I)
6:
2: q(z X)
b Forward aggregation via equation
7:
e, E(x:),
D Get modality data encoding
8:
c x= U-NetZZt > Noise and modality data prediction via the reverse model
9:
Take gradient descent step based on the loss equation 6
10:
until Converged
E. EXPERIMENT SETTINGS FOR IMAGE TRANSLATION WITH MT-
For the Cityscape dataset, the modality data corresponds to the semantic segmentation maps; and for the night2day dataset, the modality data corresponds to images of day time. We adopt the latent diffusion codebase from [dml and use the provided checkpoints of the VQ-VAE encoder-decode (kl-f8.pt). We use the VQ VAE encoder as the encoder for modality data; and construct an additional output head by duplicate the original output block of the U-Net structure the decoder generate
19

---

Published as a conference paper at ICLR 2024
the modality data, which consists of normalization layer, a SiLU layer, and a convolution layer. We use the default hyper-parameters for training the models for the two datasets, summarized as:
Attention resolutions: (32, 16, 9)
Diffusion steps: 1000
Learn sigma: False
Noise schedule: Linear
#channels: 320
#heads:
#res blocks: 2
Resblock updown: False
Use scale shift norm: False
Learning rate: 1.0e-4
Batch size: 32
Task description and encoder-decoder design This is more modality- homologous setting. We adopt two standard datasets, the Cityscale dataset for semantic-labels to photo translation (Cordts et al. 2016) and the night2day dataset for night-to-day photo translation (Laffont et al., 2014). We adopt the public codebase of latent diffusion model (LDM) (ldm). For the translation problem, the task data (original and translated images) are the same data space, thus we do not need to explicitly define separate encoders E-) for the modality data. Instead, we use the same pretrained image encoder in LDM map all images to the diffusion latent space. We add another head at the end of the U-Net as the decoder for target translated images.
Table 4: Classification performance on Cityscape dataset.
Model
Per-pixel acc.
Per-class acc.
Class IOU
CycleGAN Zhu et al. 2017
0.58
0.22
0.16
pix2pix [sola
2017
0.85
0.40
0.32
Internlmage-B Wang et
0.86
Single-task diffusion
0.72
0.54
0.32
MT-Diffusion
0.95
0.85
0.70
We perform image translation by generating target images conditioned on source images based on Algorithm I Some example generated image are illustrated in Figure For quantitative evaluation, we follow Zhu et al. (2017) to measure the performance in terms of per-pixel accuracy, per-class accuracy and class 100, and compare with exiting methods (Zhu et al.) [017 Isola et al.) 2017 Wang et al.) 2022). The results are shown in Table It is cleared that our model obtains the best accuracy compared to the baselines, except for the state-of-the-art Internimage- model, which is much larger image foundation model pretrained
Results
web-scaled data, thus it is not comparable. also calculate the IS and FID scores on the night2day dataset. Note prior work did not typically calculate these scores. We obtain an FID score of 37.93 and an IS of 3.94, and the IS score is even slightly better than that of the ground-truth data (3.65). As comparison, the FID and IS scores with single-task diffusion are 40.73 and 3.84, respectively.
E.2 MT-DIFFUION FOR MASKED-IMACE TRAINING
In this task, the modality data is a randomly masked version of the original images. To create randomly masked image, we random sample a coordinate (æx, y) that is within the image, then we masked out a patch (z min(æ + 16, 64)), min(y + 16, 64) by setting the corresponding pixel values to zeros. We repeat this process for m to control the ratio of masked regions. We adopt the latent diffusion codebase from dif We simply define the encoder the identity map. and define the decoder for the masked images by replicating the output block of the original U-Net, similar to the for training the models,
above Image Translation experiment. We adopt the default not specified below.
Diffusion steps: 1000
Rescale learned sigmas: False
20

---

Published as a conference paper at ICLR 2024
ALLTI
ARU
Figure 7: Random samples from scratch with MT-Diffusion by masked-image training on ImageNet- 128 x 128.
Rescale timesteps: False
Noise schedule: cosine
#channels: 192
#res blocks: 3
Learning rate: 7.0e-5
Batch size: 80
More random samples from scratch and image restoration results from both random masking and half masking are illustrated in Figure 1 and 10
21

---

Published as a conference paper at ICLR 2024
E.3 MT-DIFFUSION FOR JOINT MAGE-LABEL GENERATION MODELING
In this task, the modality data are discrete labels. We use the original U-Net as the encoder, which tasks noisy image, label and timestep as input. For simplicity, timestep to be zeros, although we believe better results can be obtained by jointly encoding with such information. For decoders, we proposes two options. one go out of the middle block of the U-Net and the other go out of the output block, as described in the main text. For the one from the middle block, we simply add one fully connected layer to define the decoder; and for the one from the output block, we adopt the pre-defined classifier from the codebase dif as the decoder. We adopt the default hyper-panne pyper-parameters for training the models, if not specified below. First, for MTT-iffussio-MM:
set the noisy image and the
Diffusion steps: 1000
Rescale learned sigmas: False
Rescale timesteps: False
Noise schedule: cosine
#channels: 192
#res blocks: 3
Learning rate: 7.0e-5
Batch size: 75
The setting for MT-Diffusion-E is the same as MT-Diffusion-M, except with some extra hyper- parameters for the pre-difined classifier:
resolution: (32,16,8)
Classifier_depth:
Classifier_pool: attention
True
ceasifier_usesscale. Shii__shff__ttttteeeeee norm: True
E.4 MT-DIFFUSION FOR JOINT IMAGE REPRESENTATION GENERATION MODELING
Task description and encoder-decoder design Finally, we apply MT-Diffusion for joint image- representation generation. The setting is similar to the image-label generation setting in Section 4.2 by replacing the label data with image representation from the CLIP model Radford et al. (2021). Similarly, we use the original U-Net as the encoder for image representations via the cross-attention mechanism. For the decoder, we append two-layer MLP the output of the middle block of the U-Net, which is expected to output image representations. The MLP project the tensor from middle block to dimension of 1024, followed by ReLU layer, and finally another layer output tensor of 1024. For MT-Diffusion for Joint Image-Representation::
Diffusion steps: 1000
Learning rate: 1.0e-5
Batch size: 2048
Results We conduct large-scale experiments based on the pretrained stable diffusion model Rom- bach et (2022 by continuing finetuning the model on the LAION dataset Schuhmann al. (2022) with our MT-Diffusion. We adopt the default hyperparameter setting as that in the codebase. Due to the sarge-scale nature, it is challenging make fair quantitative comparisons with related methods. Thus, we only show some generated examples from our method, and leave more extensive comparisons as future work. Some randomly generated examples are shown in Figure and 12 demonstrating impressive generated quality results.
We also provide visulized comparison between our MT-diffusion with the stable diffusion baseline in Figure 13 and 14 From the generated images, it appears that our method can understand the semantic meaning of the images and generate better looking images.
22

---

positive transfers in the former Published as a conference paper at ICLR 2024
E.5 DISCUSSION
Computation Efficiency Our multi-modal setting only adds small modality heads to decode back the modality space. In our two-modality setting, the additional computational overhead is minimal, amounting to approximately 10% more training 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 modalities separately in terms of both time and storage efficiency.
Extra Experiments During the rebuttal, we try to design new experiments to demonstrate 1) our model is better than pure condition model on two modalities; 2) negative transfer phenomenon in our model.
For 1), compare model with pure condition model that learns to recover images from random masked images. We run the experiments on the small CIFAR- 10 dataset. We observe that our model converge faster than the pure conditional baseline, while both converge comparable final results in terms of both and FID scores. However, we would like to emphasize that our model not only can do conditional generation, but also joint generation for multiple modals. Thus, model represents a more flexible generative model framework.
For 2), plan test model more tasks and modallities. Specifically, we plan to train our model on 5 tasks to simultancously learning to generate original images,, masked images, corresponding captions, random captions, and class labels. After spending significant efforts in implementation, we find it takes too much work to finish the experiment. In addition, we are in lack of GPU resources. Thus, unfortunately, we have to postpone this large-scale experiment. However, we wish to point out that our current results actually have the implications that more similar tasks tend have more positive transfers. For example, comparing the following two settings indicated Table 2: 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 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 setting.
23

---

Published as a conference paper at ICLR 2024
Figure 8: Random samples from scratch with MT-Diffusion by masked- training on ImageNet- 64x64
24

---

Published as a conference paper at ICLR 2024
Figure 9: Random samples for image restoration from random masking on ImageNet-64x In each block, the three images are original image, masked image and restored image, respectively.
25

---

Published as a conference paper at ICLR 2024
rtar
Figure 10: Random samples for image restoration from half masking on ImageNet-64x64, In each block, the first image the masked image, the rest three are different restored imaegs.
26

---

Published as a conference paper at ICLR 2024
Smiling wearing leather jacket, cowboy hat and kilt.
Young badger delicately sniffing yellow richly textured oil painting.
bird's- eye view of a mountain with tall trees and clean-water lake.
car is driving down curvy road with flowers blooming on the road sides.
church with beautiful landscape
dragon wearing
laptop screen showing bunch of photographs.
home buit a huge Soap bubble, windows, doors, porches, swnings, middle of SPACE, eyberpuni lights.
A dream of distant galaxy, by caspar david friedrich, matte painting trending on artstation
lone tree standing tall against starry night sky.
painting squirrel eating burger.
photo wearing sunglasses standing deck of spaceship with the Earth in the background.
Figure 11: Random samples for text-to-image generation finetuned stable diffusion v2.
27

---

Published as a conference paper at ICLR 2024
besutiul photog aadscaap background
developer office, photo, Image
hadggehog using calculator.
night filled with stars above
phito of in enoutriuutg ding lonse forex. red colored dog.
Ausi cabin als on the edge of giand lall.
There lca mear
Wilfowen dor the mecdow around the cabin Lake.
squirnel driving
A surnet ouer mourca rang, vertor image. Asterooo panting that actopss
candles
bookshelf window.
ANGR AGAG AIIRD
amazon
that reads "Amazon". thunder and lightning the mountain.
apic pointing Gancolf the Block summoning
bird.
An Anlluztratoon s neural network.
Image of animel haf moue octrous.
impr panting Canadian man An of Isteer nicig mouse through mapii trees
Otyreaw weln avk redden asticge westers desig Idarrepatt water radds cremare aeting carport
Ulose portrat bus messwoman tording Dense woodland landscage.
cel phone, painting the stplc Bombrandt
Dall buttalo
Falls.
28
Figure 12: Random samples for text-to-image generation finetuned stable diffusion v2.

---

Published as a conference paper at ICLR 2024
amazon
mazo
Transparent glass apple decorated with magical multicolored flowers inside
street sign that reads "Amazon"
brid's eye view mountain with tall trees and clean- water
A beautiful photograph of girl with Switzerland landscape
in the background with
A dragon wearing karate suit
church with beautiful landscape
A small chair sits front of table on the wooden floor. There is a bookshelf nearby the window
red colored dog
MT-Diffusin
Stable Diffusion
MT-Diffusion
Stable Diffusion
Figure 13: Visual comparisons between MT-diffusion (left) and the stable diffusion baseline (right).
29

---

Published as a conference paper at ICLR 2024
ANGR AGAG AIIRD
squirrel driving toy
An extremely angry bird
impressionist oil painting of a Canadian man riding through forest maple trees
image of animal half mouse half octopus
Smillng sloth wearing leather jacket, cowboy hat and
painting latent
A watercolor painting of chair that looks like an octopus
An illustration a slightly conscious neural network
MT-Difuusion
Stable Diffusion
MT-Difuusion
Stable Diffusion
Figure 14: Visual comparisons between MT-diffusion (left) and the stable diffusion baseline (right).
30


Reviews:
[{"review_id": "9KncLktgWx", "reviewer": "Reviewer_bBAj", "summary": "This paper introduces a novel approach to generative modeling by extending diffusion-based models to a multi-task learning framework. The proposed Multi-Task Diffusion Model (MT-Diffusion) is capable of generating multi-type data (e.g., images and their corresponding labels) within a single unified model. It integrates multi-task learning losses into the diffusion process, supported by a theoretical foundation. The authors propose and experiment with several multi-task generative settings, including image transition, masked-image training, joint image-label, and joint image-representation generation, demonstrating the framework's versatility and effectiveness on the ImageNet dataset. MT-Diffusion handles multiple data types through a shared diffusion space, with a forward process aggregating multi-task data and a reverse process using task-specific decoder heads to reconstruct data for different tasks. This approach results in a novel multi-task variational lower bound that generalizes the standard diffusion model, achieving simultaneous multi-task generation without compromising individual task performance.", "questions": "- 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.\n\n\n[2]: Choi, Jooyoung, et al. \"Perception prioritized training of diffusion models.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.\n\n[3]: Hang, Tiankai, et al. \"Efficient diffusion training via min-snr weighting strategy.\" arXiv preprint arXiv:2303.09556 (2023).\n\n[4]: Go, Hyojun, et al. \"Addressing Negative Transfer in Diffusion Models.\" arXiv preprint arXiv:2306.00354 (2023).\n\n[5]: Park, Byeongjun, et al. \"Denoising Task Routing for Diffusion Models.\" arXiv preprint arXiv:2310.07138 (2023).", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper provides a sound theoretical explanation for the utility of a multi-task loss using the Evidence Lower Bound (ELBO).\n- The idea of enabling multi-task learning for inputs of various modalities through a shared latent space is innovative.\n- Considering the connection to guided diffusion models is a thoughtful approach that takes into account the expansiveness of the research.", "weaknesses": "- The paper does not specify the extent of increased training costs resulting from the proposed methodology.\n- 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.\n- (Minor point) There is prior (possibly concurrent) work proposing a multimodal, multi-task diffusion process through a Versatile Diffusion[1] multi-flow diffusion process.\n- (Minor point) The caption of Table 2 does not provide sufficient information, making it difficult for the reader to interpret the experimental results.\n\n[1]: Xu, Xingqian, et al. \"Versatile diffusion: Text, images and variations all in one diffusion model.\" Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zDyctFvJyM", "reviewer": "Reviewer_KnqL", "summary": "This work discusses the potential of diffusion-based models in generative modeling. While current diffusion models excel in single-generation modeling, the paper explores the possibility of extending them for multi-task generative training. The authors introduce a unified multi-task diffusion model, MT-Diffusion, that operates in a shared diffusion space. This model aggregates information from multiple types of task-data and employs a shared backbone denoising network with task-specific decoder heads. The paper presents several multi-task generation settings, such as image transition, masked-image training, joint image-label, and joint image-representation generative modeling. Experimental results on ImageNet demonstrate the model's effectiveness in multi-task generative modeling.", "questions": "- Can you explain the model structure for each task and result? Especially results in Fig 4 and Fig 5.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper introduces MT-Diffusion, a novel approach to multi-task generative modeling using diffusion models.\n- The proposed model effectively aggregates information from different task-data types, enhancing its versatility.\n- Extensive experiments on ImageNet validate the model's effectiveness and potential in various multi-task generative modeling scenarios.", "weaknesses": "- Experiments are done on low resolution and small datasets, undermining its effectiveness. \n- The paper is lack of model details for each task", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ukaw9d3pDL", "reviewer": "Reviewer_oMx9", "summary": "This paper extends the diffusion process in a multi-task learning, with a shared diffusion space for all task data. The paper verifies its formulation in multiple variations of a two-task setting (though it should be relatively straightforward to generalise into many tasks), showing improved performance over standard single task learning baselines. The paper also lists several accompanied architecture designs for the proposed multi-task diffusion formulation, based on different choices of data domains.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Disclaimer: I am probably not the right person to review this paper. I have background in multi-task learning but have limited experience in diffusion models. The paper seems to more focus on diffusion models and have limited context in multi-task learning.\n\n-\tThe problem formulation is clean and straightforward. I have no problem understanding its derivation of ELBO and loss functions.\n-\tThe presented architectures consider multiple choices of data types.", "weaknesses": "I will present my concerns and weaknesses here fully based on my experience in multi-task learning.\n\n1.\tThe 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 bering used in these two papers, and from which both were based on the same assumption of using shared diffusion space as well?\n2.\tAs 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, Adasahre, 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MpwSqoeliN", "reviewer": "Reviewer_4Fko", "summary": "In this study, four different scenarios for multi-task learning are presented to enhance the performance of image generation and classification. The authors emphasize two specific settings, namely masked-image training and joint image-label generation. Through experiments conducted on ImageNet, the paper successfully demonstrates the viability of the suggested multi-task frameworks.", "questions": "- 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.\n\n- It is unclear from my understanding what the single-task learning model for ImageNet classification, depicted in Figure 6, entails. Further clarification would be appreciated.\n\n- To ensure the readiness of the paper for publication, it is crucial to address the concerns mentioned in the Weakness section. These issues should be at least partially resolved. Thank you for your attention to these matters.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "- Multi-task learning in generation, the subject of investigation in this study, is a captivating and relatively unexplored area of research.\n\n- The methods proposed in this study exhibit simplicity and effectiveness when applied to the ImageNet-64 baseline.", "weaknesses": "- 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.\n\n- The theoretical portion of the paper does not provide a clear and comprehensive explanation of the proposed multi-task models, as the focus of the paper is primarily empirical. It would be beneficial to have a more detailed architectural explanation of the models designed for various multi-task settings. If space constraints are a concern, the theoretical portion can be entirely moved to the appendix.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)