SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 7: Behavior of multi-task learning methods through SNR-based interval clustering across training iterations. A similar trend as in Fig. 3 is observed.
(b) Gradient weights of NashMTL
(c) Loss weight of UW
their cosine similarities calculated by every 10K iterations are averaged. Finally, we can plot the average cosine similarity against the timestep axis as in Fig. 1. For plotting them against the log-SNR axis, the values of the axis were adjusted, and the empty parts were filled with linear interpolation.
For ADM and LDM, the pairwise cosine similarity between gradients is calculated during 1M training iterations and 400K training iterations, respectively.
Negative Transfer Analysis To calculate the negative transfer gap in Eq. 3, we need to additionally train the model on denoising tasks within specific timestep interval  $[t_1, t_2]$ . Since we plot five intervals [1, 200], [201, 400], [401, 600], [601, 800], and [801, 1000], we trained the model on denoising tasks for each interval. Each model is trained for 600K iterations in ADM and 300K iterations in LDM on the FFHQ dataset. For the model trained on entire denoising tasks, we used the trained model the same as in Section 5.1. ADM is trained on 1M iterations and LDM is trained on 400K iterations. All of these models are trained with the same batch size and learning rate as experiments in Section 5.1 (See Appendix E).
DDIM 50-step sampler [67] was used for the generation. FID is calculated with Clean-FID [51] by setting the entire 70K FFHQ dataset as reference images. Since the official code of Clean-FID supports FID calculation with statistics from these reference images, we used it and reported FID with 10k generated images.
C Implementation Details for MTL methods
We describe how MTL methods are applied in Section 4.2. To be more self-contained, we hereby present implementation details for MTL methods. For the implementation of MTL methods, we used the official code of LibMTL [39] $^5$ . NashMTL [46] supports practical speed-up by updating gradient weights  $\alpha$  every few iterations, not every iteration. We utilize this by updating  $\alpha$  every 25 training iterations.
D Additional Experimental Results
We present additional experimental results to supplement the empirical findings presented in Section 5. In Section D.1, we provide visualizations of the behavior of MTL methods with other clustering methods that were not covered in Section 5.3. Furthermore, we examine the impact of our hyperparameter, the number of clusters  $k$ , in Section D.2. To validate the effectiveness of interval clustering compared to other clustering methods, we present additional results in Section D.3. In Section D.4, we delve deeper into comparing the performance of stronger MTL baselines such as Linear Scalarization (LS) [80, 35] and Random Loss Weighting (RLW) [40] with our proposed approach.
4https://github.com/GaParmar/clean-fid
<sup>5</sup>https://github.com/median-research-group/LibMTL
(b) Gradient weights of NashMTL
Figure 8: Behavior of multi-task learning methods through gradient-based interval clustering across training iterations. A similar trend as in Fig. 3 is observed.
Table 5: FID-10K scores of the LDM trained using a combination of UW and PCgrad methods on the FFHQ dataset while varying the value of  $k$ . Notably, integrating MTL methods with two clusters significantly improves FID scores. Increasing  $k$  from 2 to 5 also enhances FID scores, but further increasing  $k$  from 5 to 8 shows similar results.
Clustering	Vanilla	Number of clusters (k)
k=2	k=5	k=8
Timestep	10.56	9.563	9.151	9.083
SNR	9.606	9.410	9.367
Gradient	9.634	9.033	9.145
D.1 Visualization for the Behavior of MTL Methods with Other Clustering Methods
Due to space constraints in our main paper, we were unable to include the behavior analysis of MTL methods for SNR-based and gradient-based interval clustering. However, we present these results in Fig. 7 and 8, which show similar trends to the observations depicted in Fig. 4. These findings suggest valuable insights into the behavior of MTL methods, regardless of the clustering objectives.
Firstly, we observed a notable increase in the occurrence of conflicting gradients as the timestep difference between tasks increased. This observation suggests that the temporal distance between denoising tasks plays a crucial role in determining the frequency of conflicting gradients.
Secondly, we noted that both loss and gradient balancing methods assign higher weights to task clusters with higher levels of noise. This finding indicates that these methods allocate more importance to the noisier tasks.
D.2 Analysis: The Number of Interval Clusters
To understand the impacts of the number of clusters  $k$ , we conducted experiments by varying  $k$  with 2, 5, and 8. We trained a model for timestep-based, SNR-based, and gradient-based clustering with each  $k$ , resulting in nine trained models. For MTL methods, we used combined methods with UW [29] and PCgrad [83] as in Section 5.4. All training configurations such as learning rate and training iterations are the same as in Section 5.1. We evaluate 10K generated samples from the DDIM 50-step sampler [67] for all methods with the FID score [51, 18].
Table 5 shows the results. Notably, we made an intriguing observation regarding the integration of MTL methods with only two clusters, which resulted in a noteworthy enhancement in FID scores. Additionally, we found that increasing the number of clusters, denoted as  $k$ , from 2 to 5 also exhibited a positive impact on improving FID scores. However, our findings indicated that further increasing  $k$  from 5 to 8 did not yield significant improvements and resulted in similar outcomes. From these results, we conjecture that increasing the number of clusters to greater than five has no significant effect.
Table 7: The results of Random Loss Weighting (RLW) and Linear Scalarization (LS) on the FFHQ dataset in ADM architecture.
Clustering	Method	FID	Precision	Recall
-	Vanilla	24.95	0.5427	0.3996
Timestep	RLW	38.06	0.4634	0.3293
LS	25.34	0.5443	0.3868
SNR	RLW	35.13	0.4675	0.3404
LS	25.69	0.5369	0.3843
Gradient	RLW	36.19	0.4643	0.3392
LS	26.12	0.5120	0.3878
D.3 Comparison Interval Clustering with Task Grouping Method
To show the effectiveness of interval clustering methods for denoising task grouping in diffusion models, we compare high-order approximation (HOA)-based grouping methods [72, 12].
For grouping  $N$ -tasks in deep neural networks, the early attempt [72] established a two-stage procedure: (1) compute MTL performance gain for all task combinations and (2) search best groups for maximizing MTL performance gain across the groups. However, performing (1) requires huge computation since MTL performance gain should be measured for all  $2^{N} - 1$  combinations. Therefore, they reduce computation by HOA, which utilizes MTL gains on only pairwise task combinations. Also, the HOA scheme is inherited by the following work, task affinity grouping [12], which uses their defined task affinity score instead of MTL gains. Different from these works, our interval clustering aims to group the tasks with interval constraints.
For a fair comparison, we use a pairwise gradient similarity averaged across training iterations between denoising tasks for the objective of HOA-based grouping and interval clustering. In this case, the HOA-based grouping becomes cosine similarity grouping used in [12], and interval clustering becomes gradient-based clustering in our method. However, for HOA-based grouping, a solution of brute force searching with branch-and-bound-like algorithm [72, 12] requires computational complexity of  $O(2^{N})$ . It incurs enormous costs in diffusion with many denoising tasks. Therefore, we use a beam-search scheme in [68]. We set the number of clusters as 5 for both methods.
Table 6: Comparison interval clustering and high order approximation-based task grouping. DDIM-50 step sampler is used.
Clustering	FID-10k
HOA	9.873
Interval	9.033
We apply the combined method with UW [29] and PCgrad [83] as in Section 5.3 for the resulting clusters from both HOA-based grouping and interval clustering. We trained the model on the FFHQ dataset [27] and used LDM architecture [56]. All training configurations are the same as in Section 5.1. For evaluation metrics, we use FID and its configurations are the same as in Section 5.1.
Table 6 shows the results, indicating that the interval clustering outperforms HOA-based task grouping.
D.4 Comparison to Random Loss Weighting and Linear Scalarization
Linear Scalarization (LS) [80, 35] and Random Loss Weighting (RLW) [40] can serve as strong baselines for MTL methods. Therefore, validating the superiority of our method compared to theirs can emphasize the necessity of applying sophisticated MTL methods such as UW, PCgrad, and NashMTL. Accordingly, we provide the results of comparative experiments for LS and RLW on the FFHQ dataset using ADM architecture in Table 7. We note that all experimental configuration is the same as in vanilla training in Section 5.1.
As shown in the results, LS achieves slightly worse performance than vanilla training, which suggests that simply re-framing the diffusion training task as an MTL task and applying LS is not enough. Also, RLW achieves much worse performance compared to vanilla training. It appears that the randomness introduced by loss weighting interferes with diffusion training. These results indicate that sophisticated MTL methods are indeed responsible for significant performance gain.
E Detailed Experimental Settings in Section 5
In this section, we describe the details of experimental settings in Section 5. For validating the effectiveness in both pixel-space and latent-space diffusion models in unconditional generation, we used ADM [8] and LDM [56] as same in our observational study (refer to details of architecture in Appendix B).
E.1 Detailed Settings of Comparative Evaluation and Analysis (Section 5.1 and 5.3)
A single A100 GPU is used for experiments in Section 5.1 and 5.3.
**Setup for Unconditional Generation** We trained the models on FFHQ [27] and CelebA-HQ [26] datasets. All training was performed with AdamW optimizer [43] with the learning rate as  $1e-4$  or  $2e-5$ , and better results were reported. For ADM, we trained 1M iteration with batch size 8 for the FFHQ dataset and trained 400K iterations with batch size 16 for the CelebA-HQ dataset. For LDM, we trained 400K iterations with batch size 30 for both FFHQ and CelebA-HQ datasets. We generate 10K samples with a DDIM-50 step sampler and measure FID [18], Precision [36], and Recall [36] scores. For all evaluation metrics, we use all training data as reference data. FID is calculated with clean-FID [51], and Precision and Recall are computed with publicly available code. All analyses are conducted above trained models.
Setup for Class-Conditional Generation We trained the DiT-S/2 [52] on ImageNet dataset [7]. All training was performed with the AdamW optimizer [43] with the learning rate of  $1e - 4$  or  $2e - 5$ , and better results were reported. As in DiT [52], we applied the classifier-free guidance [19] and trained 800K iterations with a batch size of 50. All samples are generated by a DDPM 250-step sampler. For evaluation metrics, we follow the evaluation protocol in ADM [8], by using their evaluation code<sup>7</sup>. We used the cosine schedule [47] for noise scheduling and SD-XL VAE [53] for our VAE.
E.2 Detailed Settings of Comparison to Loss Weighting Methods (Section 5.2)
We trained the DiT-L/2 [52] on ImageNet dataset [7]. All training was performed with the AdamW optimizer [43] with the learning rate of  $1e - 4$ . As in DiT [52], we applied the classifier-free guidance [19] and trained 400K iterations with a batch size of 256. All samples are generated by a DDPM 250-step sampler and classifier-guidance scale of 1.5. We used the cosine schedule [47] for noise scheduling. For experiments, we used 8 A100 GPUs.
E.3 Detailed Settings of Combining MTL Methods with Sophisticated Training Objectives (Section 5.4)
We trained three different models: vanilla LDM, vanilla LDM with P2 [6], and vanilla LDM with P2, PCgrad [83], and UW [29] applied simultaneously. All training configurations are the same in Section 5.1 but we use 500K iterations. We generate 50K samples for evaluation with a DDIM 200-step sampler and evaluate FID.
F Qualitative Results
In this section, we provide qualitative comparison results, which were omitted from the main paper due to space constraints. In Figure 9, 10, 11 and 12, we visualize the generated images by all models that are used for results in Table 1. As shown in the results, we can observe that incorporating MTL methods for diffusion training can improve the quality of generated images. One noteworthy observation is that UW [29] tends to generate higher-quality images compared to NashMTL [46] and PCGrad [83]. This finding aligns with the results observed in Table 1.
Moreover, we plot the randomly selected samples from 50K generated data in Fig. 13. Despite being randomly selected, the majority of the generated images exhibit remarkable fidelity.
$^{6}$ https://github.com/youngjung/improved-precision-and-recall-metric-pytorch $^{7}$ https://github.com/openai/guided-diffusion/tree/main/evaluations
Figure 9: Qualitative comparison of ADM trained on the FFHQ dataset.
Figure 10: Qualitative comparison of LDM trained on the FFHQ dataset.
G Dynamic Programming Algorithm for Interval Clustering
In this section, we introduce the algorithm for optimizing the interval cluster and the implementation details. The optimal solution of interval clustering can be found using dynamic programming for a  $L_{cluster}$  function [2, 49, 76]. The sub-problem is then defined as finding the minimum cost of clustering  $\mathcal{X}_{1,i} = \{1,\dots ,i\}$  into  $m$  clusters. By saving the minimum cost of clustering  $\mathcal{X}_{1,i} = \{1,\dots ,i\}$  into  $m$  clusters to the matrix  $D[i,m]$ , the value in  $D[T,k]$  represents the minimum clustering costs for the original problem in Eq. 4. For some timestep  $m\leq j\leq i$ ,  $D[j - 1,m - 1]$  must contain the minimum costs for clustering  $\mathcal{X}_{1,j - 1}$  into  $(m - 1)$  clusters [49, 76]. This establishes the optimal substructure for dynamic programming, which leads to the recurrence equation as follows:
$$
D [ i, m ] = \min  _ {m \leq j \leq i} \left\{D [ j - 1, m - 1 ] + L _ {\text {c l u s t e r}} \left(\mathcal {X} _ {j, i}\right) \right\}, \quad 1 \leq i \leq T, \quad 1 \leq m \leq k. \tag {6}
$$
To obtain the optimal intervals  $l_{1},\ldots ,l_{k}$ , we use  $S[i,m]$  to record the argmin solution of Eq. 6. Then, we backtrack the solution in  $O(k)$  time from  $S[T,k]$  by assigning  $l_{m} = S[l_{m + 1} - 1,m]$  from  $m = k$  to  $m = 1$  by initializing  $l_{k} = S[T,k]$ .
Interval clustering with SNR-based or gradient-based objectives can produce unbalanced sizes of each interval, which causes unbalanced allocation of task clusters due to randomly sampled timestep  $t$ . Therefore, we add constraints on the size of each cluster to avoid seriously unbalanced task clusters. To add constraints on the size of each cluster  $n_i = |I_i| = r_i - l_i + 1$  for  $i = 1, \dots, k$ , we define the lower and upper bounds of it as  $m_I$  and  $M_I$  with  $m_I \leq \frac{n_i}{k} \leq M_I$ . In Eq. 6, the  $m$ -th cluster (i.e.,  $\mathcal{X}_{j,i}$ ) size  $n_m$  must range from  $m_I$  to  $M_I$ , yielding  $i + 1 - M_I \leq j \leq i + 1 - m_I$ . Furthermore, to
Figure 11: Qualitative comparison of ADM trained on the CelebA-HQ dataset.
Figure 12: Qualitative comparison of LDM trained on the CelebA-HQ dataset.
satisfy the  $(m - 1)$ -clusters constraint,  $1 + (m - 1)m_{I} \leq j$ . Finally, Eq. 6 with constraints on the size of the cluster is derived as follows:
$$
D[i,m] = \min_{\substack{\max \{1 + (m - 1)m_{I},i + 1 - M_{I}\} \leq j\\ j\leq i + 1 - m_{I}}}\left\{D[j - 1,m - 1] + L_{cluster}(\mathcal{X}_{j,i})\right\} ,1\leq i\leq T,1\leq m\leq k. \tag{7}
$$
Specifically, we assign  $\left\lfloor \frac{T}{2k} \right\rfloor$  and  $\left\lceil \frac{3T}{2k} \right\rceil$  to  $m_I$  and  $M_I$ , respectively.
H Broader Impacts
Revisiting Diffusion Models through Multi-Task Learning Our work revisits diffusion model training from a Multi-Task Learning aspect. We show that negative transfer still occurs in diffusion models and addressing it with MTL methods can improve the diffusion models. Starting from our work, a better understanding of the multi-task learning characteristics in diffusion models can lead to further advancements in diffusion models.
Negative Societal Impacts Generative models, including diffusion models, have the potential to impact privacy in various ways. For instance, in the context of DeepFake applications, where generative models are used to create realistic synthetic media, the training data plays a critical role in shaping the model's behavior.
When the training data is biased or contains problematic content, the generative model can inherit these biases and potentially generate harmful or misleading outputs. This highlights the importance
Figure 13: Randomly selected images from generated images of LDM with combined methods of UW, PCgrad, and P2 on the FFHQ dataset. DDIM 250-step sampler is used.
of carefully selecting and curating the training data for generative models, particularly when privacy and ethical considerations are at stake.
I Limitations
Our work has two limitations that can be regarded as future works. Firstly, we have not yet completely resolved the issue of negative transfer in the training of diffusion models as shown in Fig. 5. This indicates that learning entire denoising tasks still causes degradation in certain denoising tasks. By successfully addressing this degradation and enabling the model to harmoniously learn entire denoising tasks, we anticipate significant improvements in the performance of the diffusion model.
Secondly, our study does not delve into the architectural design aspects of multi-task learning methods. While our focus lies on model-agnostic approaches in MTL, it is worthwhile to explore the possibilities of designing appropriate architectures within an MTL framework. Previous works in diffusion models utilize timestep and noise level as input, which can be considered as using task embeddings scheme [73]. By revisiting these aspects, the architecture of the diffusion model can be further advanced in future works.