Paper: Denoising Task Routing for Diffusion Models (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
### **W1: Comparison with other routing methods**
Thank you for providing us with further methods.

**Max-roaming [1]**: We sincerely appreciate your recommendation. Comparing our method to max-roaming can further support the effectiveness of our method. We applied max-roaming to DiT-B/2 and trained it on the FFHQ dataset. The below table shows the comparative results between DTR and max-roaming.

| Method      | FID      |
|-------------|----------|
| DTR         | **7.32** |
| Max-Roaming | 39.90    |

Consequently, their approach shows suboptimal performance due to its failure to incorporate the diffusion prior to task weight and task affinity. We suggest that this phenomenon is due to the detrimental effects of introducing randomness into the prior. Indeed, as illustrated by Go et al. [3], the application of random loss weighting has been observed to result in lower performance compared to the vanilla scenario.

**ETR-NLP [2]**: utilizes slightly different routing with channel masking in DTR and TR and max-roaming [1]. Instead of task-specific masks in channel mask routing, task-specific layers (convolution layer) are employed for calculating representations in task-specific channels. Indeed, these task-specific layers require additional parameters, and they grow as the number of tasks increases, resulting in a huge cost in diffusion where there are a lot of denoising tasks (e.g. $T$ =1000). Meanwhile, our DTR does not require additional parameters for task routing. Also, their official code has not been released yet. In this regard, we did not conduct a direct comparison to ETR-NLP.
Nevertheless, we tried to implement ETR-NLP in block-wise routing during the discussion phase, however, we failed to converge ETR-NLP despite trying various hyperparameters for training. We suspect that unaddressed elements in the scope of their paper might play a role in this outcome, and we will conduct experiments upon their code release.

Reviewer comment (Experiments):
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.

Author response:
### **W1: Studies for validating 1) task affinity and 2) task weights.**

We agree that the contribution of both prior 1) task affinity and 2)  task weights should be analyzed in our work. In this regard, we would like to respectfully emphasize that the contributions of the performance for both priors are already validated in Figure 4. 

1. *Task Affinity*: the random routing method (R-TR) routes the tasks without specific consideration of task affinity. As detailed in Section 4.2, the expected number of shared channels remains constant for any pairs of distinct denoising tasks. Contrarily, our Denoising Task Routing (DTR) utilizes a sliding window strategy, shifting activated channels by the timestep from 1 to $T$. This approach ensures that tasks at neighboring timesteps share similar channel sets, thereby embedding task affinity into the routing process. For instance, when $\alpha=1$ in Eq. (4), DTR shifts channels linearly, focusing solely on task affinity without considering task weights, as seen in Figure 2. The performance superiority of DTR ($\alpha=1$) over R-TR, as shown in Figure 4, serves as an empirical validation of the benefits of incorporating task affinity in DTR.
2. *Task Weights*: By adjusting $\alpha$, DTR also accounts for task weights alongside task affinity. When $\alpha > 1$, a higher number of dedicated channels are allocated to later timesteps, in line with our intention to assign greater weight to these timesteps. We compared the performance of DTR with $\alpha=1$ and DTR with $\alpha>1$, finding that the latter generally surpasses the former, as evidenced in Figure 4 and Table 4. This comparison substantiates the impact of considering task weights in our model.

In summary, our results validate the efficacy of DTR in the context of both task affinity alone and in combination with task weights. We regret to inform you that we are currently unable to fully comprehend your queries regarding 'DTR with random routing but task-dedicated allocation channels with Random Task Routing (R-TR)' and 'DTR with task-unified allocation channels but sliding window channels and R-TR'. We kindly request further clarification on these points to provide a more thorough response. We eagerly anticipate a productive and enlightening discussion on these matters.

Reviewer comment (Experiments):
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.

Author response:
### **W2-1: Rationale for choosing Equation (4).**
Thank you for pointing this out. As mentioned above in W1, our DTR firstly incorporates task affinity with activated channels as a sliding window according to timesteps, then modulates the shifting ratio of the sliding window across timesteps for consideration of task weights. In this regard, we implemented the idea as Eq. (4), However, we acknowledge that the conceptual explanation of Eq. (4) had been lacking in the manuscript. Through Eq. (4), each mask has $C_{\beta}$ sequential activated channels, while it is shifted as like a sliding window. In this situation, the available start index of activated channels is in {$0, \dots ,C − Cβ$}, and we quantized this start index according to timestep as $(\frac{t}{T})^{\alpha}$, enabling modulation of shifting ratio of the sliding window.

### **W2-2: Sensitivity of DTR masking strategies.**
We would like to respectfully emphasize that we have validated the sensitivity of DTR masking strategies with regard to the hyperparameters $\alpha$ and $\beta$ in Figure 4 and Table 4. 

### **W2-3: Discussion on potential alternatives.**
Thank you for pointing this out. We agree that discussion on potential alternatives for masking strategies can promote future work by starting from our work. In our work, task-specific masks are not changed and optimized through training procedures. By utilizing well-known methods such as reinforcement learning and evolutionary algorithms, the masks can be more optimized than our DTR masks. Thanks again for pointing this out, we will add this discussion to our manuscripts.

Reviewer comment (Experiments):
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?

Author response:
### **Q1-1: Would DTR achieve better performance by reducing the overlap channels at higher timesteps?**
We validated that reducing overlapped channels at higher timesteps improves the performance in Figure 4 and Table 4. By increasing $\alpha$ until adequate level, overlapped channels at higher timestep reduce. As shown in the results, the FID score is improved, achieving better performance.

### **Q1-2  Would assign these tasks entirely distinct channels be beneficial?**
In the typical configuration of the Diffusion model, where $T=1000$, employing distinct channels for each task becomes impractical. Assuming 16 channels per task, this would necessitate 16,000 channels, surpassing the 1024 channels of DiT-L, which we are currently handling. Therefore, we emphasize the impracticality of allocating distinct channels for each task.
Furthermore, claiming completely distinct channels implies disregarding the inter-task relationships in denoising tasks. As demonstrated by Go et al. [1], there are correlations between tasks at higher timesteps. We believe that leveraging these correlations and efficiently sharing channels can enhance parameter efficiency, as exemplified in our handling of the limited 1024 channels in DiT-L.

Reviewer comment (Writing):
Is there a typographical error in Equation (4)? Should the first $t$ be replaced with $t-1$?

Author response:
### **Q2: Typo**
We sincerely appreciate you pointing out typos. We will fix it.

Reviewer comment (Novelty):
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.

Author response:
Thank you for raising this important point, and we sincerely apologize for any confusion caused.
While it is acknowledged that the diffusion model has been established in prior research as a multi-task learning framework for denoising tasks, it is noteworthy that the emphasis in these studies was primarily on optimization, with less consideration given to the architectural aspect. In the field of multi-task learning, it is widely accepted that both optimization and architecture play a pivotal role. We have shown that this architectural approach is the first point (in diffusion) and is also orthogonal to their methodology.

In addition, Task Routing (TR) differs from our approach in several key respects. First, we have successfully eliminated the need for additional parameters when applying Task Routing (TR) to existing architectures. Second, our research emphasizes the importance of considering inter-task relationships and warns against a naive application that could lead to performance degradation. This underscores a critical message: performance improvements in the diffusion model are achievable by accounting for prior knowledge of denoising tasks.

In light of these distinctions, we note that our contribution to task routing is significant and remains a valid addition when viewed through the lens of incorporating diffusion as a form of multi-task learning.

Reviewer comment (Experiments):
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.

Author response:
**Comparison with ETR-NLP (Ding et al. (2023) [2])**

ETR-NLP divides channels into shared and task-specific channels, however, it utilizes slightly different routing with channel masking in DTR and TR and max-roaming [1]. Instead of task-specific masks in channel mask routing, task-specific layers (convolution layer) are employed for calculating representations in task-specific channels. Indeed, these task-specific layers require additional parameters, and they grow as the number of tasks increases, resulting in a huge cost in diffusion where there are a lot of denoising tasks (e.g. $T$=1000). Meanwhile, our DTR does not require additional parameters for task routing. Also, their official code has not been released yet. In this regard, we did not conduct a direct comparison to ETR-NLP.

Nevertheless, we tried to implement ETR-NLP in block-wise routing during the discussion phase, however, we failed to converge ETR-NLP despite trying various hyperparameters for training. We suspect that unaddressed elements in the scope of their paper might play a role in this outcome, and we will conduct experiments upon their code release.

**Comparison with the timestep cluster-based task-specific channels (Go et al, 2023a [3])**

Thank you for suggesting a valuable recommendation. Comparing our DTR to task-specific channels which are allocated according to timestep clusters can show differences between routing with entire tasks and clustered tasks. We utilized timestep-based clustering [3] with ($k=8$) and assigned channel masks by regarding each cluster as one task and trained DiT on the FFHQ dataset. The below table shows the results, and our DTR dramatically outperforms the timestep cluster-based DTR, showing that routing entire tasks is more beneficial than routing clustered tasks. Thank you for improving our work, and we will add this to the manuscript. 


| Method                     | FID      |
|----------------------------|----------|
| DTR                        | **7.32** |
| DTR w timestep-cluster [3] | 9.61     |