Denoising Task Routing for Diffusion Models
Authors: Byeongjun Park
Sangmin Woo
Hyojun Go
Jin-Young Kim
Changick Kim
Venue: ICLR 2024 poster (2024)
Keywords: Diffusion Model Architecture
Multi-Task Learning (MTL)
Diffusion Models
Primary area: generative models
TL;DR: Simple add-on strategy improves diffusion model architectures by explicitly routing denoising tasks in diffusion models.

Abstract: Diffusion models generate highly realistic images by learning a multi-step denoising process, naturally embodying the principles of multi-task learning (MTL). Despite the inherent connection between diffusion models and MTL, there remains an unexplored area in designing neural architectures that explicitly incorporate MTL into the framework of diffusion models. In this paper, we present Denoising Task Routing (DTR), a simple add-on strategy for existing diffusion model architectures to establish distinct information pathways for individual tasks within a single architecture by selectively activating subsets of channels in the model. What makes DTR particularly compelling is its seamless integration of prior knowledge of denoising tasks into the framework: (1) Task Affinity: DTR activates similar channels for tasks at adjacent timesteps and shifts activated channels as sliding windows through timesteps, capitalizing on the inherent strong affinity between tasks at adjacent timesteps. (2) Task Weights: During the early stages (higher timesteps) of the denoising process, DTR assigns a greater number of task-specific channels, leveraging the insight that diffusion models prioritize reconstructing global structure and perceptually rich contents in earlier stages, and focus on simple noise removal in later stages. Our experiments reveal that DTR not only consistently boosts diffusion models' performance across different evaluation protocols without adding extra parameters but also accelerates training convergence. Finally, we show the complementarity between our architectural approach and existing MTL optimization techniques, providing a more complete view of MTL in the context of diffusion training. Significantly, by leveraging this complementarity, we attain matched performance of DiT-XL using the smaller DiT-L with a reduction in training iterations from 7M to 2M. Our project page is available at https://byeongjun-park.github.io/DTR/

Published as a conference paper at ICLR 2024
DENOISING TASK ROUTING FOR DIFFUSION MODELS
Byeongjun Park i Sangmin Woo Hyojun Go of
Jin- Vin-Young Kim 91 Changick Kim
KAIST Twelve Labs (t: Equal contribution,
*: Corresponding author)
{pb13810, smwoo95, changick Bragicckkkkiin @kaist. .a..stac.... .r (gohyojun15, seago0828 (emoo0a.8)ggma....... com
ABSTRACT
Diffusion models generate highly realistic images by learning a multi-step denois- ing process, naturally embodying the principles of multi-task learning (MTL). Despite the inherent connection between diffusion models and MTL, there re- mains an unexplored area in designing neural architectures that explicitly incor- porate MTL into the framework of diffusion models. In this paper, we present Denoising Task Routing (DTR), a simple add-on strategy for existing diffusion model architectures to establish distinct information pathways for individual tasks within a single architecture by selectively activating subsets of channels in the model. What makes DTR particularly compelling is its seamless integration of prior knowledge of denoising tasks into the framework: (1) Task Affinity: DTR activates similar channels for tasks at adjacent timesteps and shifts activated chan- nels as sliding windows through timesteps, capitalizing on the inherent strong affinity between tasks at adjacent timesteps. (2) Task Weights: During the early stages (higher timesteps) of the denoising process, DTR assigns a greater number of task-specific channels, leveraging the insight that diffusion models prioritize reconstructing global structure and perceptually rich contents in earlier stages, and foeus simple noise removal in later stages. Our experiments reveal that DTR not only consistently boosts diffusion models' performance across different evaluation protocols without adding extra parameters but also accelerates train- ing convergence. Finally, we show the complementarity between our architectural approach and existing MTL optimization techniques, providing a more complete view of MTL in the context of diffusion training. Significantly, by leveraging this complementarity, we attain matched performance of DiT-XL using the smaller DiT-L with reduction in training iterations from 7M to 2M. Our project page is available at https:/byeonnjjun park. www.kk..tuu..iooo. github.io/DTT/
INTRODUCTION
Diffuuson models (Ho et al., 2020; Sohl-Dickstein et al., 2015; Song et al., 2021) have made signif- icant strides in generative modeling across various domains, including image (Dhariwal & Nichol, 2021; Rombach et al., 2022), video (Harvey et al., 2022), 3D (Woo et al., 2023), audio (Kong et al., 2020) and natural language (Li et al., 2022). In particular, they have demonstrated their versatility across a broad spectrum of image generation scenarios such as unconditional (Ho et al., 2020; Song al., 2020), class-conditional (Dhariwal & Nichol, 2021; Nichol & Dhariwal, 2021), and text-to- image generation (Nichol et al., 2021; Ramesh et al., 2022; Saharia et al., 2022).
Diffusion models are designed to learn denoising tasks across various noise levels by reversing the forward process that distorts data towards predefined noise distribution. Recent studies (Hang et al., 2023; Go et al., 2023a) have shed light on the multi-task learning (MTL) (Caruana, 1997) aspect inherent in diffusion models, where a single neural network handles multiple denoising tasks. They particularly focus on enhancing the optimization MTL in diffusion models, employing techniques such as loss weighting (Hang al., 2023) and task clustering (Go et al., 2023a), aiming address the issue of negative transfer - phenomenon that arises when shared parameters are between conflicting tasks. While these efforts demonstrate the promise of viewing diffusion models MTL, there remains an unexplored avenue for designing neural architectures from an MTL perspective within the context of diffusion models.

---

Published as a conference paper at ICLR 2024
One common practice in diffusion models is to condition the models with timesteps (or noise lev- els) through differentiable operation (Ho et al., 2020; Karras et al., 2022; Rombach et al., 2022), prompting the model's behavior by adjusting representation of model according to timesteps (or noise levels). This can be seen as an implicil way of incorporating MTL aspects into the architec- tural design. However, we argue that this may not fully address negative transfer, as it places the entire burden of task adaptability solely on implicit signals.
In this paper, we take step beyond implicit conditioning and explicitly tackle multiple denois- ing tasks by making a simple modification to existing diffusion model architectures. Specifically, we draw inspiration from prior works on task routing (Strezoski et al., 2019; Ding et al., 2023), which enables the establishment of distinct information pathways for individual tasks within single model. The distinct information pathways are implemented through task-specific channel masking, making task routing effectively handle numerous tasks (Strezoski et al., 2019). However, we observe that a naive random routing approach (Strezoski et al., 2019), which allocates random pathways for each task, does not take account into the inter-task relationship between denoising tasks in diffusion models, resulting in a detrimental impact on performance.
To tackle this challenge, we present the Denoising Task Routing (DTR), simple add-on strategy for existing diffusion model architectures. DTR enhances them by establishing task-specific pathways that integrate prior knowledge of diffusion-denoising tasks, such as: (1) Task Affinity: Considering strong task affinity between adjacent timesteps (Go et al., 2023a), DTR activates similar channels for tasks at adjacent timesteps by sliding windows over channels throughout the timesteps and ac- tivating channels within the window. (2) Task Weights: Inspired by the observation that diffusion models prioritize reconstructing the global structure and perceptually rich contents in the early stages (higher timesteps) (Choi et al., 2022), DTR allocates an increased number of task-specific channels to denoising tasks higher timesteps.
Building upon this foundation, DTR offers notable advantages: (1) Simple Implementation: DTR
can be integrated with minimal lines of code, streamlining its adoption. (2) Elevated Performance:
DTR significantly elevates the quality of the generated samples. (3) Accelerated Convergence:
DTR enhances the convergence speed of existing diffusion models. (4) Efficiency: DTR achieves
these without extra parameters and incurs only negligible computational cost for channel masking.
Finally, we conduct experiments across various image generation tasks, such as unconditional, class- conditional, and text-to-image generation, with FFHQ (Karras et al., 2019), ImageNet (Deng et al., 2009), and MS-COCS dataset (Lin et al., 2014), respectively. By incorporating our proposed DTR into two prominent architectures, DiT (Peebles & Xie, 2022) and ADM (Dhariwal & Nichol, 2021), we observe a significant enhancement in the quality of generated images, thereby validating the benefits of our DTR. Moreover, we demonstrate the seamless compatibility of MTL optimization techniques tailored for diffusion models with our MTL architectural design for DTR. Significantly, attain similar DiT-XL's performance using the smaller DiT-L with reduction in training itera- tions from 7M to 2M, showcasing the efficiency and effectiveness of our approach.
RELATED WORK
Diffusion model architecture. Advancements in diffusion model architecture center on integrating well-established architectural components into the framework of diffusion models. Earlier works use a UNet-based architecture (Ronneberger et al., 2015) and propose several improvements. For exam- ple, DDPM (Ho et al., 2020) uses group normalization (Wu & He, 2018) and self-attention ( (Vaswani et al., 2017), IDDPM (Nichol & Dhariwal, 2021) uses multi-head self-attention, Song et al. (2021) proposes to scale skip connections by 1/v2, and ADM (Dhariwal Nichol, 2021) proposes the adaptive group normalization. Recently, several works propose transformer-baseee architectures for diffusion models instead of UNet, including GiT (Yang et al., 2022), U-ViT (Bao et al., 2023), RIN (Jabri et al., 2022), DiT (Peebles & Xie, 2022) and MDT (Gao et al., 2023). Unlike these works, our objective is to incorporate the MTL aspects into architectural design. Specifically, we propose a simple add-on strategy to improve existing diffusion models with task routing, and validate our method upon both representative UNet and Transformer based architectures, ADM and DIT.
Multi-task learning (MTL). MTL (Caruana, 1997; Sener & Koltun, 2018) aims to improve ef- ficiency and prediction accuracy across multiple tasks by sharing parameters and learning them simultaneously. This approach stands in contrast to training separate models for each task, allowing
2

---

Published as a conference paper at ICLR 2024
the model to leverage inductive knowledge transfer among related tasks. However, MTL encounters challenges that conflicting tasks exist, leading to a phenomenon known as negative transfer (Ruder, 2017), where knowledge learned in one task negatively impacts the performance of another.
To address the negative transfer, previous research explores optimization strategies and architectural designs. Optimization strategies focus on mitigating two main problems: (1) conflicting gradients and (2) unbalanced losses or gradients. Conflicting gradients between tasks can cancel each other thus resulting in suboptimal updates. To mitigate this, Yu et al. (2020) project gradient onto the normal plane of conflicting gradient and Chen et al. (2020) stochastically drop elements in gradi- ents. Imbalanced learning, where tasks with larger losses or gradients dominate the training, is also addressed through loss balancing (Kendall et al., 2018 88 and gradient balancing (Navon et al., 2022).
In terms of MTL architectures, researchers develop both implicit and explicit methods. Implicit methods guide the model to learn multiple tasks with task embeddings, avoiding the extensive task- specific modifications (Sun et al., 2021; Zhang et al., 2018; Popovic et al., 2021; Pilault et al., 2021). Explicit methods, on the other hand, embed task-specific behaviors directly into the architecture through task-specific branches (Long et al., 2017; Vandenhende et al., 2019), task-specific mod- ules (Liu et al., 2019; Maninis et al., 2019), feature fusion across multiple network branches (Gao et al., 2019; Misra et al., 2016), and task routing mechanisms (Strezoski et al., 2019: Pfeiffer et al., 2023). Task routing, in particular, demonstrates its scalability while requiring minimal additional pa- rameters, making it suitable for handling a large number of tasks. Therefore, in our work, we adopt task routing to enhance explicit MTL design within existing diffusion model architectures. Further- more, in contrast to prior research on task routing (Strezoski et al., 2019; Pfeiffer et al., 2023; Pascal et al., 2021), it noteworthy that our proposed method introduces novel approach that incorporates priors for inter-task relationships without the need for extra parameters.
MTL contexts in diffusion models. Recent studies (Hang et al., 2023; Go et al., 2023a) revisit diffusion models as a form of MTL, where single neural network simultancously learns multi- ple denoising tasks with various noise levels. They observe negative transfer between denoising tasks and seek to enhance diffusion models by addressing the issue from an MTL optimization per- spective. However, there remains limited exploration of architectural improvements from an MTL architectural perspective. To bridge this gap, our work proposes architectural enhancements within the framework of MTL for diffusion models.
Conditioning the model with timestep (Ho et al., 2020) or noise level (Song et al., 2021) can be perceived as an implicit method of incorporating MTL aspects into architectural design. For in- stance. DDPM (Ho et al., 2020) adds the Transformer sinusoidal position embedding (Vaswani et al., 2017) into each residual block, which is widely adopted for various diffusion models including LDM (Rombach al., 2022), ADM (Dhariwal & Nichol, 2021), DiT (Peebles & Xie, 2022) and EDM (Karras et al., 2022). However, we argue that relying solely on implicit signals is insufficient for effectively mitigating negative transfer. Our goal in this paper is to explicitly incorporate prior knowledge of denoising tasks into the existing diffusion model architectures with task routing.
PRELIMINARY
Diffusion models. Diffusion models (Dhariwal & Nichol, 2021; Song et al., 2020) stochastically transform an original data 2o into latent, often following a Gaussian distribution, by iteratively adding noise- - the forward process. To make diffusion models generative, they need to learn to reverse the perturbed data back to its original distribution p(æ)) the reverse process. The for- ward process can be conceptualized as fixed-length Markov chain comprising T discrete steps. At each timestep f. along this chain, represented as Tt, the data undergoes transformation based on a conditional distribution ((x))T/20)) Specifically, ( is modeled as a Gaussian distribu- tion ((zz;; Vazzo, (1 ā, [), where ã; represents noise schedule parameter, and x; denotes the noisy version of the input 2o at time t. The reverse process recovers the original data by modeling p(xx-1/x;). which approximates the distribution q( l This equips the model to effectively "undo" the diffusion steps and reconstruct the original data from the noisy observations. To achieve this, many diffusion models commonly use the training strategy of DDPM (Ho al., 2020), which aims to optimize a noise prediction network Eo Ex 8) by minimizing a simple objective ΣT Ca with respect to f, where La is defined as:
Le == V. ||E ())
3

---

Published as a conference paper at ICLR 2024
Task routing. Task routing (Strezoski al., 2019; Ding et al., 2023) is proposed to explicitly establish task-specific pathways within a single neural network. In practice, task routing employs a C-dimensional task-specific binary mask m E {0, 1] associated with the task D. Formally, the task routing is implemented by task-specific channel masking at the 1-th layer, given the input z' and transformation function F", can be expressed as:
2 = mD O (2'),
(2)
where denotes channel-wise multiplication. We note that this operation performs a conditional feature-wise transformation, allowing the neural network to create task task-specific subnetworks within a single model. By explicitly separating in-model data flows, the neural network builds its own ben- eficial sharing practices, effectively addressing negative transfer issues that may arise from sharing channels between conflicting tasks. One significant advantage of task routing lies in its scalability. does not significantly increase the number of parameters computational complexity with the addi- tion of tasks. As demonstrated in (Strezoski et al., 2019), task routing exhibits excellent scalability, proving its effectiveness even for scenarios with hundreds of tasks.
4 METHODOLOGY
In this section, introduce Denoising Task Routing (DTR), a straightforward add-on strategy on existing diffusion model architectures to enhance the learning of multiple denoising tasks. We first describe our DTR, focusing on the integration of the task routing framework (Strezoski et al., 2019; Ding et al., 2023) into the diffusion model framework in Sec. 4.1. Next, we consider a naive random routing method and discuss its limitations on handling prior knowledge of denoising tasks in Sec. 4.2. Finally, we present a detailed description of DTR in Sec. 4.3, which explicitly considers the prior knowledge of denoising tasks in the routing mask creation.
4.1 TASK ROUTING FOR DIFFUSION MODELS
We conceptualize diffusion training as a form of MTL, where each task corresponds to the denoising task D, at a specific timestep E {1.....3 T) learned by L, in Eq. (1). Typically, T often surpasses 1000, resulting in thousands of denoising tasks being jointly optimized in single model.
Many diffusion models employ a multi-layered residual block structure in their architecture (Rombach et al., 2022; Dhariwal & Nichol, 2021; Peebles & Xie, 2022; Song et al., 2021). These models commonly adopt the practice of ini- tializing each residual block as the identity function. For example, ADM (Dhariwal & Nichol, 2021) initializes the final convolutional layer of every residual block as zero. On the other hand, DiT-based models (Peebles & Xie, 2022: Gao et al., 2023) utilize -- in their transformer blocks for initializing them as identity functions. To eas- ily integrate these practices into task routing, we apply task routing at the level of residual blocks, emphasizing block-wise task routing as the foundational element of our method. By denoting the 1-th block as Block' and its in- put as 2 E HxxC C our denoising task routing repre- sented as:
DTR
1-1
Block!
Routing Mask
Active Task Dy
mps
Figure 1: The overview of DTR. DTR makes explicit task-specific pathways by channel masking.
_+1 z + Block' (mp, @ x'),,
(3)
where mD, E 1) denotes task-specific binary mask for D,-
Figure provides a concise overview of our DTR scheme, which is adaptable to general residual block structures including ADM (Dhariwal & Nichol, 2021) and DiT (Peebles & Xie, 2022). In both the inference and training stages, the activated mask is set according to the input timestep of the noise prediction network ég. Through this approach, we explicitly establish task-specific pathways within a single noise prediction network ep. Detailed implementations for incorporating DTR in ADM and DiT architectures can be found in Appendix A.
4

---

Published as a conference paper at ICLR 2024
4.2 A NAIVE RANDOM TASK ROUTING APPROACH (R-TR)
The remaining part of designing task routing is to establish task-specific routes by defining task- specific routing mask m. To establish task-specific routes, we first consider setting m., as ran- dom masks (Strezoski et al., 2019), which actívates a predefined portion of random channels for each task. For each denoising task Dr. we randomly sample Cg [&C]] channel indices from the set {1, (I.....C) C}. where 0 < _ - 1. Then, the routing mask mDD is configured to assign a value of one to the randomly sampled channel indices and value of zero to others.
Here, activation ratio 3 determines the trade-off between task-specific units versus task -gnneal units within the model architecture. When = 1, the model is the same as the original model without task routing, where all units are shared among tasks. As 3 decreases from 1, the model allocates fewer units for sharing, thereby enhancing task-specificity.
However, employing random masking for diffusion models might overlook the inter-task relation- ships between denoising tasks. A recent work (Go et al., 2023a) shows that affinity between denois- ing tasks increases as the proximity of their timestep increases, suggesting that sharing units between tasks with closer timesteps beneficial. Despite this, the expected number of shared channels re- mains constant for pairs of distinct tasks (as we observed in Appendix B), implying that random masking inherently cannot consider timestep proximity. Additionally, we empirically validate this in Sec. 5 and results be found in Table 1. To leverage the prior knowledge of the diffusion task, we design a more tailored mask in the next section.
4.3 MASK CREATION OF DENOISING TASK ROUTING (DTR)
To adequately reflect the specific characteristics of diffusion de- noising tasks, we propose a novel masking strategy grounded in recent findings in the field: (1) Task Affinity: Denoising tasks at adjacent timesteps have a higher task affinity than those at distant timesteps (Go et al., 2023a).
(2) Task Weights: Previous stud- ies have shown improvements in diffusion training by assign- ing higher weights to denoising tasks at higher timesteps com- pared to lower timesteps (Hang et al., 2023; Choi et al., 2022; Go et al., 2023a). This aligns with the observation that diffu- sion models primarily learn perceptually rich content at higher timesteps, whereas they focus on straightforward noise removal at lower timesteps (Choi et al., 2022).
T
Random
DTR < c)
C.
T
DTR (a 1)
DTR (e 1)
To integrate the concept of (1) Task Affinity, we employ a slid- ing window of size C3 within the mask, activating channels within its boundaries. As we increase timesteps from 1 to T, the sliding window gradually shifts. This ensures that denoising tasks at neighboring timesteps engage similar sets of channels, while those at distant timesteps reduce channel sharing. The un- derlying principle here is that sharing channels between tasks having higher task affinity proves (Fifty et al., 2021). To incorporate (2) Task Weights,
Figure 2: Routing masks in ran- dom routing and DTR with vary- ing a (3 is fixed to 0.8). The ac- tivated and deactivated channels are colored in yellow and purple, respectively.
beneficial for training, as demonstrated we use an additional parameter, a. This modulates the shifting ratio of the sliding window across timesteps, manipulating the amount of allocation of task-dedicated channels to each timestep.
To incorporate the above two concepts, the available start index of activated channels in fO.... C CB),, and we quantized this start index according to timestep as (+)², enabling modulation of shifting ratio of the sliding window. Formally, the masks are initialized as:
1, if -G) (41) <e 5 [(C Cp) (4) Cp.
M.D.C 0, otherwise.
(4)
A conceptual visualization of our masking strategy shown in Fig. 2. When 0 < a < 1, it allocates more task-dedicated channels to smaller timesteps. At = 1, task-dedicated channels are evenly distributed across all timesteps. Lastly, for a > 1, more desk-dedic channels are assigned to higher timesteps. aligning with intent give more weight to the higher timesteps. However, setting a to too large values causes the initial sliding window shift to occur at very large timesteps, in turn, leads to a situation where many tasks no longer have dask-dedie channels.
5

---

Published as a conference paper at ICLR 2024
Latent-Levvl Diffusion Model. Imuge Resolution x
Pixel-Levvl Diffusion Model. Imuge Resolution 64:664
FFHQ
ImageNet
MS-COCO
FFHQ
ImageNet
Model
Model
FIDL
FID
1St
Prect
Rect
FIDL
FID,
FID]
1S1
Prect
DIT
10.99
12.59
134.60
0.73
0.49
7.70
ADM
3.12
6.34
72.05
0.85
0.46
DIT
DIT + DIR
7.32
8.90
156.48
0.77
0.51
7.04
ADM + DIR
2.65
4,91
82.44
0.85
0.46
Table 1: Comparative results. We evaluate unconditional image generation on FFHQ, class-
conditional image generation on ImageNet, and text-conditional image generation MS-COCY We set the activation ratio / to 0.8 for both R-TR and DTR. Note that our DTR achieves substantial performance improvements without additional parameters significant computational costs.
EXPERIMENTAL RESULTS
In this section, we present experimental results to validate the effectiveness of our method. To be- gin, outline our experimental setups in Sec. 5.1. Then, we provide the results of a comparative evaluation in Sec. 5.2, showing that our method significantly improves FID, IS, Precision, and Re- call metrics compared the baseline. Finally. we delve into comprehensive analysis of DTR in Sec. 5.3, dissecting its performance across multiple dimensions.
5.1
EXPERIMENTAL. SETUP
Due to space constraints, we provide a concise overview of our experimental setups here. More extensive information regarding all our experimental settings can be found in Appendix C.
Evaluation protocols. To assess the effectiveness of our method, we conducted compre- hensive evaluation across three image-generation tasks: 1) Unconditional generation: utilized FFHQ (Karras et al., 2019), which contains 70K training images of human faces. 2) Class- conditional generation: used ImageNet (Deng al., 2009), which contains 1,281,167 training images from 1K different classes. 3) Text-to-Image Image generation: we used MS-COCO (Lin et al., 2014), which contains 82,783 training images and 40,504 validation images, each annotated with descriptive captions.
Evaluation metries. To evaluate the quality of the generated samples, we used FID (Heusel al., 2017), IS (Salimans et al., 2016), and Precision/Recall (Kynkäänniemi et al., 2019), Specifically, FID is used for sample quality in unconditional and rext-to-image generation. Then, FID, IS, and Precision are used for sample quality measure and Recall is used for diversity measure in class- conditional generation.
Models. To verify the broad applicability of our method, we utilized two representative architec- tures: UNet-based ADM (Dhariwal & Nichol, 2021) and Transformer-basee DIT (Peebles Xie, 2022). For text-to-image generation, used CLIP text encoder (Radford et al., 2021) to transform textual descriptions into a sequence of embeddings for the condition diffusion models.
5.2 COMPARATIVE EVALUATION
Quantitative results. We quantitatively validate our approach on well-established architectures, egg., DiT and ADM. The results are presented in Table I. Firstly, we observe that naive random routing (R-TR) leads performance degradation. This occurs because the R-TR approach lacks the capability to incorporate prior knowledge of the diffusion model (Task Affinity) specific to de- noising tasks, as relies on random instantiation of routing masks. In contrast, DTR incorporates the prior knowledge of denoising tasks in diffusion models, mentioned as Task Affinity and Task Weights in Sec. 4.3. Therefore, through its straightforward design, our DTR consistently demon- strates significant performance enhancements across all metrics for three datasets when compared to the model without DTR. Note that DTR achieves substantial performance improvements with no extra parameters and with negligible computational overhead for multiplications of channel masks.
Compatibility of DTR with MTL loss weighting techniques. In Table 2, we show that our DTR, an MTL architectural approach for diffusion models, is compatible with MTL loss weighting techniques specifically designed for diffusion models (Go et al., 2023a; Hang et al., 2023) as well improved loss weighting method (Choi al., 2022), both in class-conditional and unconditional

---

Published as a conference paper at ICLR 2024
Class- Class-Conditioaa ImageNet 256x256.
DIT-L/2
DIT-L// DTR
Loss Weight Type
FID.,
1S1
Prect
Rect
FID_
1S1
Prect
Rect
Vanilla
12.59
134.60
0.73
0.49
8.90
156.48
0.77
0.51
Min-SNR (Hang al., 2023)
9.58
179.98
0.78
0.47
8.24
186.02
0.79
0.50
ANT-UW (Go al., 2023a)
5.85
206.68
0.84
0.46
4.61
208.76
0.84
0.48
Unconditional FFHQ 256
Loss Weight Type
FID-10K
Vanilla
SNR (Hang al., 2023)
ANT-UW (Go al., 2023a)
(Choi et al., 2022)
D/T-B/2
12.93
9.73
9.30
10.08
DIT-B/2 DIR
8.82
8.93
8.81
8.83
Table Compatibility of DTR with MTL. loss weighting methods. In a class-conditional gener- ation, utilizing both DTR and loss weighting techniques significantly boosts performance, showing their complementarity. In unconditional generation, employing only DTR nearly matches the best performance, which underscores the effectiveness of DTR standalone solution.
VanlliaDTR
MIR
Mn-SNR+DTR
ANT-UW
ANT-UN-DTR
30.00
350
0.70
25/0
300
0.00
20.00
250
200
0.50
15:00
0.79
1000
150
180
0.02
0.40
2:00
0.0
15
2.0
1.0
1.3
23
23
15
20
2.
1.5
155
20
25
Scale
Scale
Scala
Figure 3: Compatibility of DTR and MTL loss weighting methods w.r.t. guidance scale. DTR robustly boosts the performance across various guidance scales for all metrics.
generation scenarios. Here, we use the DiT architecture as our baseline model. Initially, we observe that applying loss weighting techniques yields superior performance compared to using no such techniques. In a coassioondiiinal generation, the simultancous use of both DTR and loss weighting techniques consistently boosts performance, implying that these two approaches complement each other effectively. Furthermore, in Fig. 3, we provide insights into performance variations across different guidance scales. The results demonstrate the robustness of our DTR across guidance scales, as DTR generally enhances FID, IS, and Precision. unconditional generation, DTR essentially takes on the role of loss weighting techniques, reducing the necessity of additional loss weighting techniques. As a result, employing only DTR leads to performance levels that are nearly equivalent to those achieved with the combination of DTR and loss weighting techniques. This underscores the effectiveness of our DTR approach as standalone solution for unconditional generation tasks.
Qualitative results. Due space limitations, we present comprehensive set of generated exam- ples in Appendix for qualitative comparison. To summarize our findings, our method for training diffusion models produces images that exhibit improved realism and fidelity when compared dif- fusion models without DTR.
More Training Itera- tions. We have explored the impact of longer training on model performance. Our extensive training of DÍT-L/2 + DTR with ANT-UW (Go al., 2023a) for million iterations signifi- cantly enhanced FID scores to 2.33 on Ima- geNet 256x256. Table 3 shows our method, de- spite fewer parameters and iterations, outper- forms vanilla DiT-XL/2 and rivals DiT-XL after 7 million iterations. This underscores approach's 's efficiency, demonstrating dramatic improve- ments by integrating MTL into diffusion model architecture and optimization.
Further Results
Methed
FID]
Train Iters
Params
Flops (G)
DIT-XL/2 DIT-XL/2
2.27
7M
675M
118.64
2.55
2.35M
675M
118.64
DIT-L/2
2.33
2M
458M
80.73
Table More Training Iterations. Although DTR used smaller parameters, DTR shows similar performance compared to the larger model trained more iterations.
5.3 ANALYSIS
Mask instantiation strategy, Given that the routing mask of DTR is instantiated by two hyper- parameters, and B, we conduct ablation studies to assess the impact of varying them in Fig.

---

Published as a conference paper at ICLR 2024
Bassire (1-2.))
cm (d cm (1-0.d)
Laconditional FFHO 256:226
Clus-Conditionl ImageNet 256x:26
1-77% (pw2.:0
D-77 (p-2
Model
PiD!
Model
I55
Prect
Reef
(3
10.99
DIF (8
12.59
134,60
0.73
0.49
14
7.78
0.5
11.55
146.05
0149
7.78
10.31
149.23
0.49
13
12
INT-22 DTR
7.77
7.64
INEL22 DTR
2.0
3.0
8,89
9.39
154,05
0.77
0.77
0.49
0.49
4.0
7.32
4.0
8.90
156.45
0.77
0.51
11
5.0
7.40
5.0
9.66
156.47
0.77
0.50
10
Text-Conditional MS-C0C0 256 256
Clas-Condiiional ImageNet 64::6
Medell
Model
FIN]
1S1
Pre11
Reef
DNT-B/2
7.70
ADM 1)
0.34
72.05
0.46
8:25
5.30
74.74
0.45
2.0
INTB22 DTR
7.04
ADM DTR
5.48
74.75
035
0.46
0.46
0
5.0
6,93
5.0
5.22
78.12
0.86
0.45
Figure 4: 0x B ablation. We use DiT-B/2 FFHQ
Table 4: Mask instantiation ablation. We set B = 0.8 for DTR
due to its stable performance, as shown in Fig. 4. In general, a =4
256x256.
yields the best results.
Cuss-Condiiionll ImageNet 256x 256.
DTT-4
50-522 OTR
DTTLG- AiE 40--000-0
Model
FID]
1St
Prect
Rect
S
DTTLG
DIT.S/2
44.28
32.31
0.41
0.53
D/T.S/2 DTR
37.43
38.97
FDD33
DIT-B/2
27.96
64,72
0.57
0,52
DIT-B/2 DTR
16.58
87.94
DIT-1/2
12.59
134.60
0.73
0.49
10
DIT-L/2 DTR
156.48
0.51
Training Stops (K)
Training Staps [K]
Table 5: Impact of DTR w.r.t. model size. Note that DTT achieves consistent improvements across model sizes.
(a) Model Size
(b) Loss Weight Type
Figure- 5: Convergence comparison on ImageNet.
DTR accelerates faster FID-IOK improvement.
To provide a clear context, we set a baseline (DiT-B/2 without any task routing) and include R-TR (DīT-B/2 with random routing) for comparison. We initially observe that setting 3 to 0.8 rather than 0.6 leads to superior performance for both DTR and R-TR and notably, B = 0.8 exhibits robust behavior w.r.t. variations in a. Consequently, we fix B at 0.8.
To delve deeper into the impact of a on performance, we report the results by changing a on each dataset in Table 4. Increasing a corresponds allocating more channels tasks higher timesteps. Our results indicate that a 4 yields the best performance across almost all datasets and evalua- tion metrics. Note that increasing a leads to significant performance improvements up to a certain threshold (a = 4), beyond which performance begins to degrade. This suggests that allocating a moderately larger capacity to tasks higher timesteps is beneficial for overall performance. This suggests that our design principle for DTR, Task Weights - allocating moderately larger capacity to tasks at higher timesteps, is beneficial for overall performance. From the results, we opt to fix / at 0.8 and a at 4 for further evaluations.
Impact of DTR with respect to model size. In Table 5, we present the results of a controlled scal- ing study of DiT on the ImageNet dataset, focusing on how DTR affects the performance according to the DiT model sizes (S, B, L). Initially, we observe considerable performance improvements as scale up the model. Importantly, applying DTR further enhanced performance across all model sizes, with larger models benefiting more. It is hypothesized because as the model size increases, the total number of channels increases, thus more task-dedicated channels can be allocated.
Convergence speed. In Fig. 5, we present a comparative analysis of convergence speed, focusing on the impact of training with and without DTR. First, we investigate training convergence by in- tegrating DTR into DiT models of varying sizes (S, B, L). As shown in the results, the addition of DTR leads significant speed boost regardless of the model size. In particular, training DiT-B/2 without DTR takes roughly 400K training iterations to reach an FID score of 31, whereas, with DTR, achieves the same result in only 200K iterations, effectively doubling the speed. Addition- ally, we explore the synergy between DTR and MTL loss weighting methods (ANT-UW (Go et al., 2023a) and Min-SNR (Hang et al., 2023)), which can boost the convergence of DIT, in the context of DiT-L/2. While using MTL loss weighting methods alone provides certain degree of convergence acceleration, integrating DTR can further enhance convergence speed. Moreover, DTR mitigates the
8

---

Published as a conference paper at ICLR 2024
Shallow
Depth
Deep
(T) (0) vanilla
Figure 6: Comparison of CKA representation similarity on FFHQ dataset. We show how inter- task representation similarity changes when applying task routing in 12 DiT blocks. The horizontal and vertical axes in each plot represent timesteps (from & to f = T). We assess three configu- rations: (a) DiT alone vs. (b) DiT with R-TR vs. (c) DiT with DTR. R-TR generally increases CKA similarity, whereas DTR decreases it. Brighter/darker color represents higher/lower similarity.
saturation issue that emerges at 300K in DīT-L/2 ANT-UW, leading to a more stable convergence and accelerated learning. Through this, we confirm that explicitly handling negative transfer with DTR can significantly improve training dynamics.
Representation analysis via CKA. In Fig. we employ Centered Kernel Alignment (CKA) (Ko- mblith et al., 2019) to visualize the similarity between intra-model representations. Specifically, examine how similar or different the model representations are at different timesteps within each DiT block to gain insights into the model's behavior. CKA quantifies this similarity, where higher score indicates that the model behaves similarly across different timesteps, while lower score im- plies that the model's behavior varies significantly across different timesteps.
We make noteworthy observations by comparing three scenarios: DiT model without task routing vs, DiT with random routing vs. DIT with DTR: (1) Upon introducing random routing to the DiT model, we notice an overall increase in CKA compared to the baseline. This implies that with random routing, the model's representations remain more similar across various timesteps. (2) With DTR, we observe a distinct pattern in the CKA scores, where there are high scores at lower timesteps and low scores at higher timesteps. We can interpret that at higher timesteps, the model primarily focuses on learning discriminative features that are relevant to specific timesteps, whereas at lower timesteps, the model tends exhibit similar behavior across different timesteps. This aligns with our design principle, Task Weights. (3) In the later blocks with DTR, there is notable highlight on diagonal elements. This suggests that the model takes into account Task Affinity, reflecting the model's is ability to make its behavior more similar for adjacent timesteps.
Comparison to multi-experts strategy. We compare DiT-L/2 equipped with DTR against multi-experts model (DiT-B/2 x 4), with each expert specializing in certain timesteps, e.g., [4]..... /4, T]. Here, we show that DTR outperforms the multi-experts denoiser method (Balaji et al., 2022; Lee et al., 2023). For detailed results, please refer to Appendix D.
DISCUSSION
In this work, we have proposed DTR, a simple add-on strategy for diffusion models that estab- lishes task-specific pathways within single model while embedding prior knowledge of denoising tasks into the model through explicit architectural modifications. Our experimental findings clearly indicate that DTR represents significant leap forward compared current diffusion model archi- tectures. which rely solely on implicit signals. Importantly, this improvement is achieved only with a negligible computational cost and without introducing additional parameters. Our work conveys two important messages: (1) We found that relying solely on implicit signals for enhancing task adapt- ability of diffusion models, e.g., conditioning on timesteps or noise levels, proves insufficient to mitigate negative transfer between denoising tasks. (2) By explicitly addressing the issue of negative transfer and incorporating prior knowledge of denoising tasks into diffusion model architectures, work shows promise in enhancing their performance across various generative tasks. To the best of our knowledge, our study is the first to advance diffusion model architecture from an MTL perspective. We hope our work will inspire further investigations in this direction.

---

Published as a conference paper at ICLR 2024
7 ETHICS STATEMENT
Generative models, such as diffusion models, have the potential to exert profound societal influence, with particular implications for deep fake applications and the handling of biased data. A critical focus is on the potential amplification of misinformation and the erosion of trust in visual media. In addition, when generative models are trained on datasets with biased or deliberately manipulated content, there is the unintended consequence of inadvertently reinforeing and exacerbating social biases, thereby facilitating the spread of deceptive information and the manipulation of public per- ception. We will encourage the research community discuss ideas prevent these unintended consequences.
REPRODUCIBILITY STATEMENT
We present details implementation and experimental setups in main manuscripts and Ap- pendix. To further future works from our work, we release our experimental codes and checkpoints https://gthhb..ommyyyoogguuuuu
REFERENCES
Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Karsten Kreis, Miika Ait- tala, Timo Aila, Samuli Laine, Bryan Catanzaro, al. ediffi: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01322 2022. 9, 18
Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22669-22679, 2023.
Rich Caruana. Multitask learning. Machine learning, 28:41-75, 1997. 1,2
Zhao Chen, Jiquan Ngiam, Yanping Huang, Thang Luong, Henrik Kretzschmar, Yuning Chai, and Dragomir Anguelov. Just pick sign: Optimizing deep multitask models with gradient sign dropout. Advances in Neural Information Processing Systems, 33:2039-2050, 2020. 3
Jooyoung Choi, Jungbeom Lee, Chaehun Shin, Sungwon Kim, Hyunwoo Kim, and Sungroh Yoon. Perception prioritized training of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11472-11488 11481, 2022. 2, 5,
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248 255.25.. leee, 2009. 2, 6, 17
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in neural information processing systems, 34:8780-8799 8794, 2021. 1, 2, 4, 6,
Chuntao Ding, Zhichao Lu, Shangguang Wang. Ran Cheng, and Vishnu Naresh Boddeti. Mitigating task interference in multi-task learning via explicit task routing with non-learnable primitives. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7756-7765, 2023. 2, 4
Chris Fifty, Ehsan Amid, Zhe Zhao, Tianhe Yu, Rohan Anil, and Chelsea Finn. Efficiently identify- ing task groupings for multi-task learning. Advances in Neural Information Processing Systems, 34:27503-2751 2021.
Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is strong image synthesizer. arXiv preprint arXiv:2303. 10000...43300 14389, 2023. 2,
Yuan Gao, Jiayi Ma, Mingbo Zhao, Wei Liu. and Alan L Yuille. Nddr-enn: Layerwise feature fusing in multi-task cnns by neural discriminative dimensionality reduction. Proceedings of the IEEF/CV conference computer vision and pattern recognition, pp. 3205-3214, 2019.
10

---

Published as a conference paper at ICLR 2024
Hyojun Go, Jinyoung Kim, Yunsung Lee, Seunghyun Lee, Shinhyeok Oh, Hyeongdon Moon, and Seungtaek Choi. Addressing negative transfer in diffusion models. In Thirty-sevente Conference on Neural Information Processing Systems, 2023a. 1, 5, 17. 19
Hyojun Go, Yunsung Lee, Jin- Young Kim, Seunghyun Lee, Myeongho Jeong, Hyun Seung Lee, and Seungtack Choi. Towards practical -aug-aad diffusion models. In Proceedings of the JEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 1962-1971, 2023b. 18
Tiankai Hang, Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, and Baining Guo. Efficient diffusion training via min-snr weighting strategy. arXiv preprint arXiv:2303. 09556, 2023. 3, 5, 6, 8,
William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. Advances in Neural Information Processing Systems, 35:
27953-27965, 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016. 15
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. Advances in neural information processing systems, 30, 2017. 6, 17
Jonathan Ho and Tim Salimans. 220.:22207.12998, 2022. 17
Classifier-free
diffusion guidance.
arXiv preprint
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840- 6851, 2020. 2, 3, 17
Allan Jabri, David Fleet, and Ting Chen. Scalable adaptive computation for iterative generation.
arXiv preprint arXiv:22 1972, 2022. 2
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401-4410, 2019. 6,
Tero Karras, Miia Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in Neural Information Processing Systems, 053265655-26777 2022. 2,
Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In Proceedings of the IEEE conference computer vision and patlern recognition, pp. 7482-7491, 2018.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprini arXiv:2009.0976/ 2020.
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In International conference on machine learning, pp. 3519- 3529. PMLR. 2019. 9
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in Neural Information Processing Systems, 32, 2019. 6, 17
Yunsung Lee, Jin- Young Kim, Hyojun Go. Myeongho Jeong, Shinhyeok Oh, and Seungtack Choi. auliiearchieectue multi multi-expert diffusion models. arXiv preprint arXiv:2306.0493 04990, 2023. 9, 18
Xiang Li. John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion- Im improves controllable text generation. Advances in Neural Information Processing Systems, 55:4328-4343 2022.
11

---

Published as a conference paper at ICLR 2024
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer EscionEECV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Pro- ceedings, Part V 13, pp. 740-755. Springer, 2014. 2, 6, 17
Shikun Liu, Edward Johns, and Andrew J Davison. End-to-end multi-task learning with attention. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp- 1871-1880, 2019. 3
Mingsheng Long, Zhangjic Cao, Jianmin Wang, and Philip S Yu. Learning multiple tasks with multilinear relationship networks. Advances in neural information processing systems, 30, 2017.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019. 17
Kevis-Koking Maninis, Ilija Radosavovic, and lasonas Kokkinos. Attentive single-tasking of multi- ple tasks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recogni- tion, pp. 1851-1860, 2019.
Ishan Misra, Abhinav Shrivastava, Abhinav Gupta, and Martial Hebert. Cross-stitch networks for multi-task learning. In Proceedings the IEEE conference on computer vision and pattern recog- nition, pp. 3994-4003, 2016.
Aviv Navon, Aviv Shamsian, Idan Achituve, Haggai Maron, Kenji Kawaguchi, Gal Chechik, and Ethan Fetaya. Multi-task learning as a bargaining game. In International Conference on Machine Learming, pp. (6428-16446. PMLR, 2022. 3
Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint 1rXiv:2112.10744 2021.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffuuson probabilistic models. In International Conference on Machine Learning, pp. 8162-8171, PMLR, 2021. 1, 17
Lucas Pascal, Pietro Michiardi, Xavier Bost, Benoit Huet, and Maria Zuluaga. Maximum roaming multi-task learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 9331-9341, 2021. 3, 19
William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748, 2022. 2, 3, 4, 6, 15
Jonas Pfeiffer, Sebastian Ruder, Ivan Vulié, and Edoardo Maria Ponti. Modular deep learning. arXiv preprint arXiv:2302. //529, 2023. 3
Jonathan Pilault, Amine El hattami, and Christopher Pal. Conditionally adaptive multi-tast learn- ing: Improving transfer learning in NLP using fewer parameters & less data. In International Conference on Learning Representations, 2021.
Nikola Popovic, Danda Pani Paudel, Thomas Probst, Guolei Sun, and Luc Van Gool. Composite- tasking: Understanding images by spatial composition of tasks. Proceedings of the IEEF/CVF Conference Computer Vision and Pattern Recognition, pp. 6870-6880, 2021.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, al. Learning transferable visual models from natural language supervision. In International conference on machine learning. pp-
8748-8763. PMLR, 2021. 6, 18
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 1(2)3, 2022.
12

---

Published as a conference paper at ICLR 2024
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Con- ference Computer Vision and Pattern Recognition, pp. 10684-10695, 2022. 1, 2,4,
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical CnrrventionMMCCAA.
image segmentation. Medical Image Computing and Computer-Assistet 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18. pp. 234-241. Springer, 2015. 2
Sebastian Ruder. An overview of multi-task learning in deep neural networks. arXiv preprint arXiv: 2017. 3
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Bureu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Informa- lion Processing Systems, 45:36479-364999 2022.
Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016. 6, 17
Ozan Sener and Vladlen Koltun. Multi-task learning as multi-objective optimization. Advances in neural information processing systems, 31, 2018. 2
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamies. In International conference on machine learn- ing. pp. 2256-2265. PMLR, 2015.
Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02500 2020. 1.3
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 2021. 1, 3,
Gjorgji Strezoski, Nanne van Noord, and Mareel Worring. Many task learning with task routing. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 1384-1384, 2019. 2, 3, 5,
Guolei Sun, Thomas Probst, Danda Pani Paudel, Nikola Popović, Menelaos Kanakis, Jagruti Patel, Dengxin Dai, and Lue Van Gool. Task switching network for multi-task learning. In Proceedings of the IEEE/CVF international conference computer vision, pp. 8291-8300 2021. 3
Simon Vandenhende, Stamatios Georgoulis, Bert De Brabandere, and Luc Van Gool. Branched multi-task networks: deciding what layers share. arXiv preprint arXiv: 02920, 2019. 3
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. 2, 3, 15, 18
Sangmin Woo, Byeongjun Park, Hyojun Go, Jin-Young Kim, and Changick Kim. Harmonyview: Harmonizing consistency and diversity in one-image-to-3d. arXiv preprint arXiv:2312.15980 2023.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp. 3-19, 2018. 2, 15
Xiulong Yang, Sheng-Min Shih, Yinlin Fu, Xiaoting Zhao, and Shihao Ji. Your vit is secretly a hybrid diseriminnaivve-gneratvee diffusion model. arXiv preprint arXiv:2208.077770 2022. 2
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. Advances Neural Information Processing Systems, 33:5824-5836. 2020. 3
Yu Zhang, Ying Wei, and Qiang Yang. Learning to multitask. Advances in Neural Information Processing Systems, 31, 2018.
13

---

Published as a conference paper at ICLR 2024
APPENDIX
CONTENTS
A Implementation Details on Denoising Task Routing
15
A.1 Implementation Details on ADM and DiT
15
A.2 Pseudocode
16
The Average Portion of Shared Channels Random Masking Strategy
17
Detailed Experimental Setup in Section 5
17
D Comparison to Multi-Experts Strategy
18
Comparison computational complexity.
19
F Potential Alternatives of Sliding Window
19
G Limitations and Future Works
19
Qualitative Results
20
H.1 Qualitative Results for Comparative Evaluation
20
H.2 Qualitative Results from DiT-L/2 with DTR and ANT-UW
20
14

---

Published as a conference paper at ICLR 2024
Scale
Scale
Poistwse Feadforned
Feedforward
Pointrie
Feedforward
Scale
Seale: Shif
Scall, shirt
Scals, Shini
Layer
Layer Noea
Layer Noen
Eayar
Seale
Scale
Scale
Mulisiliaal
Hood Self Antentiioo
Multi-lead Self-Antzntion
Self- Aieentinn
Mtlli-liaul
Self Anetton
Scale, Shin
Scale, Shuft
Scale, Shift
Layer Noce
Layer Noen
Layer Norm
ML.7
Impor Tokers
Conditioiig
Ixpuc Toicrns
Irpat Tokors
Imper Tokcons
DIT Block
Reformulated DIT Block
DTR on Reformulated DIT Block
DTR on DIT Block
Figure 7: DiT block with DTR. @ represents an element-wise multiplication. We only show the conditioning block in the DiT block (leftmost), as all blocks use the same conditioning block.
A IMPLEMENTATION DETAILS ON DENOISING TASK ROUTING
In this section, we present the implementation details on Denoisng Task Routing (DTR). Firstly, in Sec. A.I, describe the details of implementation for incorporating DTR in ADM (Dhariwal & Nichol, 2021) and DiT (Peebles Xie, 2022) architectures. To provide further details, we illustrate pseudocode for the task routing mechanism and the routing mask instantiation in Sec. A.2.
A.1 IMPLEMENTATION DETAILS ADM AND DIT
ADM (Dhariwal & Nichol, 2021) We apply DTR on two types of residual blocks, an Atten- tion (Vaswani al., 2017) block, and ResNet (He et al., 2016) block. Due to the potential effect of the channel masking on local rumning mean and variance, DTR is positioned right after the group normalization layer (Wu He, 2018) in both types of residual blocks. Then, can easily apply DTR since ADM uses the same configured residual block in Eq. (3).
DiT (Peebles & Xie, 2022) DiT introduces rero-initialized adaptive layer normalization (adaLN- Zero) in transformer blocks, where the normalization parameters are regressed from the timestep em- beddings and conditions. While DiT also utilizes residual connections within the transformer block, it does not directly adhere to the formulation outlined in Eq. (3). Consequently, we reformulate the DiT block, introducing the necessary modifications integrate DTR.
The original (-th DiT block DYTBlock' outputs 21+1 given the input z' as:
[+1 DiTBlock' (=²) (z' + Attn' (z²)) MLP' Attn'(z!)),
(5)
where Attn' and MT. represent a multi-head self-attention and pointwise feedforward layer in I-th DiT block, both including adaLN-Zero layers. We note that this can be regarded residual block by defining the block as:
Block' (2') = Attn' + MLP (z' + Attn'(z')).
(6)
The left two blocks in Fig. 7 show overview of how reformulate the DiT block. We then apply DTR on the reformulated DIT block by using Eq. (3) and Eq. (6), which is expanded as:
_+1 = + Attn (mDD @ 2²) + ((MD, @ 2') Atnn' Att''mp, @ 2²)).
(7)
Here, using the existing DIT block in Eq. (5), can be expressed
_+1 (1 mp,) @ z' DiTBlock (mp, ④ 2²).
Interestingly, this can be implemented simply by applying DTR to the original DiT block and then incorporating skip connection from the output with complementary routing mask at the end.
15

---

Published as a conference paper at ICLR 2024
Pseudo Code 1 [NumPy-like] Random Masking (Left) vs. DTR Masking (Right)
C: numbec channels
E: chancels
3 beta:
activation zatio
alpha:
channel shifting parazetes
bets:
activation ratio
def andionrraaking(t c, betal:
initialize
dtr_naskinngf,
nunber activated channels pam activ int (beta G]
nask ap.xecon (R, $ nunber of activated channels oun activ int (bata c]
10 fill the mask with onee
10
11 nask(;, activc
11
nunber of desctivated channels
randonly chuffle columes of mask[;, np.randon, perddom.ermtaatioooooo
12
activ
linearly
pointa
sp..iiepace
15
t apply scaling factor to 1inear points
16
As alpha
17
calculate the chennel offeet for timestaps offeet
18
.o.roudd))
for t in
start offset[t]
22
end offset[t] nun_ activ nask(t, rend]
23
Pseudo Code 2 [Simplified] ADM block (Left) vs.
ADM block + DTR (Right)
input representation routing
(x):
apply normalization, 5110
def forward is, mask]:
5
S11U (n)
5
apply normalization, $15U 2:100/nor((1]
6
pe-intrrpolatiio
7
(ipeanple
apply
apply
conv (h)
t up-intarpolation cone
apply normalisation, 8110, conv
10
R_R CoDV (wpsanplpl (nL#))
11
conv (8130 (h)))
11
E apply cony
representation (upsanple (x))
12
apply nornalizatios, 81IU,
COLV (Dozn
add original regresentation return conv (xpsanple RLE
16
Pseudo Code 3 [Simplified] DiT block (Left) DiT block DTR (Right)
e: input ropresentation
E: input ropresentation
2 mask: routing nask
ded (e):
apply attention, connection
fosward
attention
apply routing
7 apely norn, mip, skip connection
mask
nip (soz22 (e))
$ apply nozn, attantion. skip connection
roturn 8
8
aLS attention (noml [)
10 apply nozn, slp, connection
9
add original ropresentation complement return (1-maak) x_x
A.2 PSEUDOCODE
Our DTR is easy to implement yet highly effective. Adding just a few lines of code can lead to significant performance boost. This be observed in the pseudocode examples. These code snippets illustrate the concept of random masking and the implementation of masking using the DTR (see Pseudo Code 1). To provide further clarity, also offer pseudocode for simplified version of
16

---

Published as a conference paper at ICLR 2024
the ADM ResBlock and the DiT block, both extended with DTR functionality (see Pseudo Code and Pseudo Code 3, respectively).
B THE AVERAGE PORTION OF SHARED CHANNELS IN RANDOM MASKING STRATEGY
To verify that the random masking does not take into account the relationships between tasks, we derive the expected value (X) of the shared channel, where is a random variable representing the number of shared channels for two tasks Di, and Dt,- For ease of understanding, two tasks D;, and Dr, as and j. Intuitively, when i = j. the all channels are shared, yielding R(X) C3. In the case of + without loss of generality, we first sample the channel indices set R(1) = (i1.... &c, } C (1,-.,C) C'} for task i. When selecting R(j), the probability P(k) of selecting l: shared channels from R(i) and selecting the rest from others is (G2) (C-CC) (c.-R)/(C (Cg) Finally, the expectation value of is derived as follows:
abbreviate
Cp. CD (P(K), where P(k)= (C,
if j. otherwise.
B(X)
(9)
Note that the expectation value of two distinct tasks remains consistent, which indicates that the randomly initialized routing mask falls short of representing the inter-task relationship as it assumes that all denoising tasks are equally related. As extensively studied in previous studies (Go et al., 2023a), Task Affinity is one of the prior knowledge in the field of diffusion models. This underlines why the use of random routing methods leads a noticeable degradation in performance.
C DETAILED EXPERIMENTAL SETUP IN SECTION
Training details. We employed the AdamW optimizer (Loshchilov & Hutter, 2019) with fixed learning rate of le-4. No weight decay was applied during training. A batch size of 256 was used and a horizontal flip was applied the training data. We utilized classifier-free guidance (Ho & Salimans, 2022) with a guidance scale set to 1.5 in conditional generation settings such as text-to- image generation and class-conditional image generation. For the FFHQ dataset (Karras et al., 2019), we trained for 100k iterations and evaluated model performance on 50K samples. On the ImageNet dataset (Deng et al., 2009), we trained for 400K iterations and evaluated models using 50K samples. In experiments on MS-COCO dataset (Lin et al., 2014), we trained for 400K iterations and evaluated model performance on 50K samples.
The diffusion timestep T was set to 1,000 for training and DDPM 250-step (Ho et al., 2020) for sample generation. We used a cosine scheduling strategy (Nichol & Dhariwal, 2021) 1) and applied an exponential moving average (EMA) to the model's parameters with decay of 0.9999 to enhance stability. All the models were trained on 8 NVIDIA A100 GPUs. We implemented the task routing on the official code of DIT and ADM
For implementing loss weighting methods such as Min-SNR (Hang et al., 2023), ANT-UW (Go al., 2023a), and P2 (Choi al., 2022), we used the officially released code of Min-SNR and P2d for implementing their weighting method. However, since ANT does not release their code, we re-implemented ANT-UW with timestep-based clustering when the number of clusters is 8.
Evaluation metrics. We evaluated diffusion models using FID (Heusel et al., 2017), IS (Salimans al., 2016), and Precision/Recall (Kynkäänniemi et al., 2019). Lower FID indicates a closer distri- bution match between generated and real data, suggesting higher quality and diversity of generated samples. Higher IS implies that the generated data are of higher quality and diversity. Precision measures whether generated images fall within the estimated manifold of real images, while Recall measures the reverse. Higher Precision and Recall reflect better alignment between the generated
17

---

Published as a conference paper at ICLR 2024
Caas-Conditional ImageNet 256x256.
Method
# of Parameters
FID
1St
Prect
Reef
Vanilla (DiT-L/2)
458M
12.59
134.60
0.73
0.49
Multi-experts (DiT-B/2 x DTR (DiT-L/2)
521M
9.48
149.79
0.75
0.51
458M
8.90
156.48
0.77
0.51
Table 6: Comparison between DTR and multi-expert strategy. Although DTR used smaller pa- rameters, DTR outperforms the multi-expert strategy in terms of FID, IS, and Precision.
and real data distribution. We followed the evaluation protocol of ADM and used codebase". Unless otherwise stated, FID is calculated with 50K generated samples.
Additional architectural details for DēT. Since the official DiT code only offers the implementation for class-conditional genera- tion, we have extended it to include implementations for uncondi- tional and rext-conditional generation.
For the unconditional generation, we set the number of classes to one following the recommendation of the authors DIT 6, For the text-conditional generation, we utilize text tokens from CLIP (Rad- ford et al., 2021) text encoder to condition the diffusion model. Figure 8 briefly shows the implemented DiT block for the text- conditional generation. Given that we utilize conditions as a se- quence of text tokens, opposed to single token in the class- conditional generation, the parameters udeLN-Zero are solely re- gressed from the timestep embeddings. To condition text tokens, we incorporate a multi-head cross-attention layer within the DiT block. This layer follows the same structural design the multi-head self- attention, with text tokens serving as keys and values in the cross- attention layer (Vaswani et al., 2017). Note that optimizing unconditional and text-conditional DiT blocks is beyond the scope of our current focus, leaving opportunities for further improvement.
Ps.7i
sat Homaes
Figure DiT block for text- to-image generation.
For all DiT experiments, we employ a VAE encoder/decoder from Stable Diffusion' 2 to obtain the latent feature of input images. This VAE maps 256 x 256 x 3 images into a compact latent repre- sentation with 32 32 dimension.
D COMPARISON TO MULTI-EXPERTI STRATEGY
Although our works focus on effectively building single neural network for diffusion models, comparing our DTR to works using multiple neural networks (Balaji al., 2022; Go al., 2023b; Lee et al., 2023) for diffusion models further support the effectiveness of our method.
Regarding this, we compare multi-experts and DTR, when using a similar number of parameters. For constructing multi-experts, we used four DřT-B/2 models, and each model trained on a specific area of the four parts of timesteps {1,... T). We trained each model with 200K iterations with a learning rate of le-4 and batch size of 256. Each DiT-B/2 model has 1130.3M parameters, total used parameters for multi-experts are 521.2M parameters. For comparison above multi-experts with DTR, we used the DiT-L//2 model which has 458M parameters and used in experiments in Sec. 5.
Table shows the results of the comparison between DTR and multi-experts strategy on the Ima- geNet 256x256 dataset. As shown in the results, both the multi-experts strategy and DTR outperform vanilla training. Notably, our DTR outperforms the multi-experts strategy in terms of FID, IS, and Precision. This result implies that explicitly handling negative transfer in single model with DTR can outperform wewaumeter-spprraae models for covering denoising tasks.
com/ http:://hw.....oeeeooooee/eee///
"https://hugging http://////ooonngoooooooooaaaaae
18

---

Published as a conference paper at ICLR 2024
E COMPARISON OF COMPUTATIONAL COMPLEXITY.
Despite not requiring additional parameters, DTR incurs minimal computational cost for channel masking. To clar- ify this cost, we report floating point operations (FLOPs) and average training iterations executed per second across different model sizes (S, B, L) of DiT in Tab. 7. The results show a negligible increase in GFLOPs and a correspond- ing decrease in average training speed. This supports the computational efficiency of our DTR, demonstrating that it requires only marginal computation from its adoption ex- isting models.
Method
GFLOPs
Iters/scc
DiT-S/2
6.06
8.19
DiT-S/2 DTR
6.06
8.14
DiT-B/2
23.01
6.90
DrT-B/2 DTR
23.02
6.87
DiT-L/2
80.71
3.72
DiT-L/2 DTR
80.73
3.71
Table Computation comparison.
F POTENTIAL ALTERNATIVES OF SLIDING WINDOW
Masking Strategy Type
DiT-B/2
Vanilla
MaxRoaming (Pascal et al., 2021)
ERCDT
CDTR
DTR
FID4
10.99
39.90
10.13
9.61
7.32
Table 8: Masking Strategy Alternatives. DTR outperforms several masking strategy alternatives, which fall short of adequately incorporating task weights or task affinity.
Here, to further support the effectiveness of DTR, compare more extensive baselines. First, choose MaxRoaming (Pascal et al., 2021) which utilizes the optimization strategy on randomly ini- tialized channel masks. Second, we employ tinestep-baser clustering Go et al. (2023a) on DTR for validating the effects of our masking strategy on fine-grained denoising tasks compared to clustering these tasks. We used k 8 of cluster size for timestep-based clustering, and initialized masks re- garding each cluster one task. We denote this as CDTR. Third, we explicitly route with clustered denoising tasks (ERCDT), where half of the channels are shared across all denoising tasks, while the remaining channels are segmented and activated for specific tasks. Comparing ERCDT with DTR can also validate the effects of whether clustering is applied or not. We trained DiT-B/2 on the FFHQ dataset using each strategy, and we present the results in Table 8.
The results show that our DTR significantly outperforms all masking strategy alternatives. The al- ternatives show suboptimal performance due to their failure to incorporate the diffusion prior to task weight and task affinity. For MaxRoaming, we suggest that this phenomenon due to the detrimen- tal effects of introducing randomness into the prior. As illustrated by ANT (Go al., 2023a), the randomness causes negative impacts on performance, and Randomness in MaxRoaming also causes alternatives, CDTR and ERCDT, the primary rea-
this performance degradation. For the other son for this discrepancy lies in the inability of alternative methods to adequately capture and reflect nuanced, proximal relationships among denoising tasks inherent in the clustering approach. For ex- ample. the denoising task at / = 1 is considered nearly equivalent to tasks at ( = 5 or / 100 within the context of k 8 clusters, failing intervals. Furthermore, while tasks at t = 124 and = 125 belong to the same cluster, timesteps t 125 and 126 fall into different clusters, not effectively reflecting the one-timestep difference between them. This limitation impedes the performance of ERCDT and CDTR compared to DTR, reinforeing the effectiveness of our proposed method. It is noteworthy that both CDTR and ERCIT outperform vanilla training, indicating that task routing, even with discrete representations through task clustering, enhances performance by incorporating relationships among denoising tasks.
recognize the higher affinity between tasks at closer time
G LIMITATIONS AND FUTURE WORKS
In this work, we have proposed fixed task-specific masks that incorporate the prior knowledge of de- noising tasks in diffusion models. Although showed that these fixed masks achieve dramatic performance improvements, task-specific masks are not changed and optimized through training pro- cedures. Despite the immutability of masks having advantages in training speed and computation,
19

---

Published as a conference paper at ICLR 2024
further optimization can be more beneficial. By utilizing well-known methods such reinforcement learning and evolutionary algorithms, the masks can be more optimized than our DTR masks and these can be future work from work. Additionally, starting from work, another future study could be to architecturally consider resource partitioning among multiple denoising tasks.
H QUALITATIVE RESULTS
H.1 QUALITATIVE RESULTS FOR COMPARATIVE EVALUATION
Qualitative Comparison on FFHQ Dataset Figure shows the qualitative comparison of results on unconditional facial image generation between baseline, R-TR, and DTR. Our proposed method has better performance in generating realistic images.
Qualitative Comparison on ImageNet Dataset ation, we show the generated results from baseline, R-TR, and DTR. As illustrated in Fig. 10, our method outperforms others.
For the comparison of conditional image gener-
Qualitative Comparison on MS-COCO Dataset To further verify the effectiveness of the pro- posed method, we compare the qualitative results of the Text-to-mmage generation task between baselines, R-TR, and DTR in Fig. 11.
H.2 QUALITATIVE RESULTS FROM DIT-L/2 WITH DTR AND ANT-UW
Figures 12, 13, 14, 15 illustrates the generated images by DiT-L with DTR and ANT-UW trained on 400K iterations. As shown in the results, highly realistic images generated by DTR and ANT-UW despite the model being only trained on 400K iterations with a batch size of 256.
Baxeline
R-TR
DTR
Figure 9: Qualitative comparison between baseline, random routing (R-TR), and denoising
task routing (DTR) FFHQ dataset.
20

---

Published as a conference paper at ICLR 2024
Baseline
R-TR
DTR
Ringlet
Basenji
Psimacus crithacus
Polecal
Police van
Dining table
Organ
Figure 10: Qualitative comparison between baseline, random routing (R-TR), and denoising task routing (DTR) on ImageNet dataset.
"A bed with brightt geaen and nebchiing
"A black atove top oven witting ciicchen.
"People are teakkinng deva covered
"Tvo horsex are grazing on the rolling hills"
"Century ald boilcings acand neaz stzeeta in eity."
"Contury old bui buildings atand near atroets in city."
Boselnn
R-TR
DTR
Figure 11: Qualitative comparison between baseline, random routing (R-TR), and denoising task routing (DTR) MS-COCO dataset.
21

---

Published as a conference paper at ICLR 2024
Figure 12: Uncurated 256x256 DIT-L/2 samples. Classifier-free guidanzce scale = 2.0. Class label = "golden retriever" (207)
22

---

Published as a conference paper at ICLR 2024
Figure 13: Uncurated 256x256 DIT-L/2 samples. Classifier-free guidance scale = 2.0. Class label = "panda" (388)
23

---

Published as a conference paper at ICLR 2024
Figure 14: Uncurated 256x256 DIT-L/2 samples. Classifier-free guidance scale = 4.0. Class label = "cliff drop-off" (972)
24

---

Published as a conference paper at ICLR 2024
Figure 15: Uncurated 256x256 DIT-L/2 samples. Classifier-free guidance scale = 2.0. Class label = "lake shore" (975)
25


Reviews:
[{"review_id": "67iJkjCxQQ", "reviewer": "Reviewer_8VoW", "summary": "This paper introduces Denoising Task Routing (DTR), an add-on strategy for diffusion models that incorporates multi-task learning (MTL). The proposed channel masking strategy effectively boosts performance without introducing any extra parameters. The experiments demonstrate consistent improvement across evaluation protocols.", "questions": "See Weaknesses.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. The proposed routing mask strategy is interesting as it leverages the task similarity between adjacent timesteps. \n\n2. The experiment conducted in this study is comprehensive and demonstrates significant performance improvement.", "weaknesses": "1. The idea of considering diffusion models as multi-task learning has previously been proposed by Hang et al. (2023) and Go et al. (2023a). The proposed masking strategy in this work is a simple modification of TR (Strezoski et al., 2019). Its novelty is limited.\n\n2. It lacks an ablation study to evaluate the necessity of the proposed masking strategy. Ding et al. (2023) propose to divide channels into shared channels and task-specific channels. Assigning each time-step cluster (Go et al., 2023a) to the respective task-specific channels can serve as an important baseline.", "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": "jyJlvZpKo8", "reviewer": "Reviewer_VZ4w", "summary": "This paper proposes Denoising Task Routing (DTR), a simple add-on strategy for existing diffusion model architectures to establish distinct information pathways for individual tasks within a single architecture by selectively activating subsets of channels in the model. The authors incorporate two prior knowledge aspects of diffusion-denoising tasks—task affinity and task weights—into the model architecture design to mitigate the negative transfer phenomenon. The paper provides empirical results on several image generation tasks and a qualitative analysis to validate the effectiveness of DTR.", "questions": "* Would DTR achieve better performance by reducing the overlap channels at higher timesteps? Given the authors' assertion in Figure 6 that \"at higher timesteps, the model primarily focuses on learning discriminative features that are relevant to specific timesteps, whereas at lower timesteps, the model tends to exhibit similar behavior across different timesteps,\" it appears that denoising tasks at higher timesteps have less correlation. Would assigning these tasks entirely distinct channels be beneficial?\n* Is there a typographical error in Equation (4)? Should the first $t$ be replaced with $t-1$?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper effectively addresses the negative transfer phenomena by establishing task-specific pathways for multiple denoising tasks. The concept of integrating key prior knowledge in diffusion and task routing is well-presented and could potentially influence future work on architecture design in diffusion models.\n* The implementation, although simple, is effective and yields significant performance gains on multiple benchmarks.\n* The paper is structured well, making it easy to understand and follow.", "weaknesses": "* The empirical analysis could be more comprehensive in decoupling the contributions of task weights and task affinity. As I understand, the results in Figure 4 only ablate the significance of the synergy of the two priors. To study the direct contribution of **Task Weights**, it would be helpful to compare `DTR with random routing but task-dedicated allocation channels` with `Random Task Routing (R-TR)`. Similarly, to study the contribution of **Task Affinity**, a comparison between `DTR with task-unified allocation channels but sliding window channels` and `R-TR` would be useful.\n* The paper does not adequately explain the sensitivity of DTR to different masking strategies. The authors should elaborate on why they chose Equation (4) as the masking strategy and discuss potential alternatives.", "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": "LpzqvR1hSR", "reviewer": "Reviewer_wtsR", "summary": "In this paper, diffusion training is cast as multi-task learning, where each task corresponds to the denoising task at a specific timestep. The authors present Denoising Task Routing (DTR), a simple add-on strategy for existing diffusion model architectures to establish distinct information pathways for individual tasks within a single architecture by selectively activating subsets of channels in the model. Besides, the channel partitioning considers task affinity and task weights in diffusion models. Extensive experiments demonstrate the effectiveness and efficiency of the proposed method.", "questions": "Please address my concerns in the \"Weaknesses\" section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. This paper proposes a simple add-on strategy for existing diffusion model architectures, which is simple yet effective, without introducing additional parameters, and contributes to accelerating convergence during training.\n2. Extensive experiments demonstrate the effectiveness and efficiency of the proposed method.\n3. The paper is well-written and easy to follow.", "weaknesses": "1. Some advanced routing methods [1, 2] improve the random routing by considering the inter-task relationship. Hence, it is better to discuss and compare the proposed method with them.\n\n2. In Figure 9, the images generated by the baseline (the first row) look very strange and both R-TR and DTR methods alleviate it (the second and third rows). So why the random routing method can work well? In particular, in the fifth case/column, the image generated by R-TR looks better than the one generated by DTR. Why?\n\n[1] Pascal et al. Maximum Roaming Multi-Task Learning. AAAI 2021.\n\n[2] Ding et al. Mitigating Task Interference in Multi-Task Learning via Explicit Task Routing with Non-Learnable Primitives. CVPR 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": {}}]

Author rebuttal:

Decision: Accept (poster)