SlowGuess/NeurIPS_2023_OCR huggingface dataset:
After dividing the denoising tasks into task clusters via interval clustering, we apply multi-task learning methods to the resulting task clusters. As mentioned in Section 2, previous multi-task learning works have tracked down the following causes for negative transfer: (1) conflicting gradient, (2) difference in gradient magnitude, and (3) imbalanced loss scale. In this work, we leverage one representative method that tackles each of the causes mentioned above, namely, (1) PCgrad [83], (2) NashMTL [46], and (3) Uncertainty Weighting [29].
For each training step in diffusion modeling, we compute the noise prediction loss  $L^l$  for the  $l$ -th data within the minibatch. As shown in Eq 2, calculating  $L^l$  involves sampling the timestep  $t^l$ , in which case  $L^l$  is a loss incurred on the denoising task  $\mathcal{D}^{t^l}$ . We may then assign  $L^i$  to the appropriate task cluster by considering the corresponding timestep. Subsequently, we may group up the losses as  $\{L_{I_i}\}_{i = 1,\dots,k}$ , where  $L_{I_i}$  is the loss for the  $i$ -th task cluster. (More details in Appendix C)
1. PCgrad [83] In each iteration, PCgrad projects the gradient of a task onto the normal plane of the gradient of another task when there is a conflict between their gradients. Specifically, PCgrad first calculates the per-interval gradient  $\nabla_{\theta}L_{I_i}$ . Then, if the other interval gradient  $\nabla_{\theta}L_{I_j}$  for  $i \neq j$  has negative cosine similarity with  $\nabla_{\theta}L_{I_i}$ , it projects  $\nabla_{\theta}L_{I_i}$  onto the normal plane of  $\nabla_{\theta}L_{I_j}$ . PCgrad repeats this process with all of the other interval gradients for all interval gradients, resulting in a projected gradient per interval. Finally, model parameters are updated with the summation of projected gradients.
2. NashMTL [46] In NashMTL, the aggregation of per-task gradients is treated as a bargaining game. It aims to update model parameters with weighted summed gradients  $\Delta \theta = \sum_{i = i}^{k}\alpha_{i}\nabla_{\theta}L_{I_{i}}$  by obtaining the Nash bargaining solution to determine  $\alpha_{i}$ , where  $\Delta \theta$  is in the ball of radius  $\epsilon$  centered zero,  $B_{\epsilon}$ . They define the utility function for each player as  $u_{i} = \langle \nabla_{\theta}L_{I_{i}},\Delta \theta \rangle$ , then the unique Nash bargaining solution can be obtained by  $\arg \max_{\Delta \theta \in B_{\epsilon}}\sum_{i}\log (u_{i})$ . By denoting  $G$  as matrix whose columns contain the gradients  $\nabla_{\theta}L_{I_i}$ ,  $\alpha \in \mathbb{R}_{+}^{k}$  is the solution to  $G^{\mathrm{T}}G\alpha = 1 / \alpha$  where  $1 / \alpha$  is the element-wise reciprocal. To avoid the optimization to obtain  $\alpha$  for each iteration, they update  $\alpha$  once every few iterations.
3. Uncertainty Weighting (UW) [29] UW uses task-dependent (homoscedastic) uncertainty to weight task cluster losses. By utilizing observation noise parameter  $\sigma_{i}$  for  $i$ -th task clusters, the total loss function is  $\sum_{i} L_{I_i} / \sigma_i^2 + \log(\sigma_i)$ . As the noise parameter for the  $i$ -th task clusters loss  $\sigma_i$  increases, the weight of  $L_{I_i}$  decreases, and vice versa. The  $\sigma_i$  is discouraged from increasing too much by regularizing with  $\log(\sigma_i)$ .
In this section, we demonstrate the efficacy of our proposed method by addressing the negative transfer issue in diffusion training. First, we provide the comparative evaluation in Section 5.1, where
Table 1: Quantitative comparison to vanilla training (Vanilla) on the unconditional generation. Integration of MTL methods using interval clustering consistently improves FID scores and generally enhances precision compared to vanilla training.
Model	Clustering	Method	Dataset
FFHQ [27]	CelebA-HQ [26]
FID (↓)	Precision (↑)	Recall (↑)	FID (↓)	Precision (↑)	Recall (↑)
ADM [8, 6]	Vanilla	24.95	0.5427	0.3996	22.27	0.5651	0.4328
Timestep	PCgrad [83]	22.29	0.5566	0.4027	21.31	0.5610	0.4238
NashMTL [46]	21.45	0.5510	0.4193	20.58	0.5724	0.4303
UW [29]	20.78	0.5995	0.3881	17.74	0.6323	0.4023
SNR	PCgrad [83]	20.60	0.5743	0.4026	20.47	0.5608	0.4298
NashMTL [46]	23.09	0.5581	0.3971	20.11	0.5733	0.4388
UW [29]	20.19	0.6297	0.3635	18.54	0.6060	0.4092
Gradient	PCgrad [83]	23.07	0.5526	0.3962	20.43	0.5777	0.4348
NashMTL [46]	22.36	0.5507	0.4126	21.18	0.5682	0.4369
UW [29]	21.38	0.5961	0.3685	18.23	0.6011	0.4130
LDM [56]	Vanila	10.56	0.7198	0.4766	10.61	0.7049	0.4732
Timestep	PCgrad [83]	9.599	0.7349	0.4845	9.817	0.7076	0.4951
NashMTL [46]	9.400	0.7296	0.4877	9.247	0.7119	0.4945
UW [29]	9.386	0.7489	0.4811	9.220	0.7181	0.4939
SNR	PCgrad [83]	9.715	0.7262	0.4889	9.498	0.7071	0.5024
NashMTL [46]	10.33	0.7242	0.4710	9.429	0.7062	0.4883
UW [29]	9.734	0.7494	0.4797	9.030	0.7202	0.4938
Gradient	PCgrad [83]	9.189	0.7359	0.4904	10.31	0.6954	0.4927
NashMTL [46]	9.294	0.7234	0.4962	9.740	0.7051	0.5067
UW [29]	9.439	0.7499	0.4855	9.414	0.7199	0.4952
our method can boost the quality of generated samples significantly. Next, we compare previous loss weighting methods for diffusion models to UW with interval clustering in Section 5.2, verifying our superior effectiveness to existing methods. Then, we analyze the behavior of adopted MTL methods, which serve to explain the effectiveness of our method in Section 5.3. Finally, we demonstrate that our method can be readily combined with more sophisticated training objectives to boost performance even further in Section 5.4. Extensive information on all our experiments can be found in Appendix E.
5.1 Comparative Evaluation
Experimental Setup Here, we demonstrate that incorporating MTL methods into diffusion training improves the performance of diffusion models. For comparison, we consider unconditional and class-conditional image generation. For unconditional image generation, we used FFHQ [27] and CelebA-HQ [26] datasets, where all images were resized to  $256 \times 256$ . For class-conditional image generation experiments, we employed the ImageNet dataset [7], also resized to  $256 \times 256$  resolution.
For architecture, we adopt widely recognized architectures for image generation. Specifically, we use the lightweight ADM [6, 8] and LDM [56] for unconditional image generation, while employing DiT-S/2 [52] with classifier-free guidance [19] for class-conditional image generation. We train the model using our method: We consider every possible pair of (1) interval clustering (timestep-, SNR-, and gradient-based) and (2) MTL method (PCgrad, NashMTL, and Uncertainty Weighting (UW)), and report the results. We used  $k = 5$  in interval clustering throughout experiments.
For evaluation metrics, we use FID [18] and precision [36] for measuring sample quality, and recall [36] for assessing sample diversity and distribution coverage. IS [61] is additionally used for the evaluation metric in the class-conditional image generation setting. Finally, for sample generation, we use DDIM [67] sampler with 50 steps for unconditional generation and DDPM 250 steps for class conditional generation, and all evaluation metrics are calculated using  $10\mathrm{k}$  generated samples.
Comparison in Unconditional Generation As seen in Table 1 our method significantly improves performance upon conventionally trained diffusion models (denoted vanilla in the table). In particular, there is an improvement in FID in all cases, and an improvement in precision scores in all but two cases, which highlights the efficacy of our method. Also, given strong results for both pixel- and latent-space models, we can reasonably infer that our method is generally applicable.
We also observe the distinct characteristics of each multi-task learning method considered. Uncertainty Weighting tends to achieve higher improvements in sample quality compared to PCgrad and NashMTL. Indeed, UW achieves superior FID and Precision for ADM, while excelling in Precision for LDM.
Figure 3: Quantitative comparison to vanilla training (Vanilla) on ImageNet  $256 \times 256$  dataset with DiT-S/2 architecture and classifier-free guidance. Integration of MTL methods using interval clustering consistently improves FID, IS, and Precision compared to vanilla training.
Table 2: Comparison between MinSNR and ANT-UW. DiT-L/2 is trained on ImageNet.
Method	FID	IS	Precision	Recall
Vanilla	12.59	134.60	0.73	0.49
MinSNR	9.58	179.98	0.78	0.47
ANT-UW	6.17	203.45	0.82	0.47
Table 3: GPU memory usage and runtime comparison on FFHQ dataset in LDM architecture.
Method	GPU memory usage (GB)	# Iterations / Sec
Vanilla	34.126	2.108
PCgrad	28.160	1.523
NashMTL	38.914	2.011
UW	34.350	2.103
However, UW sacrifices distribution coverage in exchange for sample quality, resulting in lower Recall compared to other methods. Meanwhile, NashMTL scores higher in recall and lower in precision compared to other methods, suggesting it has better distribution coverage while sacrificing sample quality. Finally, PCgrad tends to show a balanced performance in terms of precision and recall. We further look into behaviors of different MTL methods in Section 5.3.
Due to space constraints, we provide a comprehensive collection of generated samples in Appendix F. In summary, diffusion models trained with our method produce more realistic and high-fidelity images compared to conventionally trained diffusion models.
Comparison in Class-Conditional Generation We illustrate the results of quantitative comparison on class-conditional generation in Fig. 3. The results show that our methods outperform vanilla training in FID, IS, and Precision. In particular, UW and Nash-MTL significantly boost these metrics, showing superior improvement in generation quality. These results further support the generalizability of MTL methods through the interval clustering on class-conditional generation and the transformer-based diffusion model.
5.2 Comparison to Loss Weighting Methods
Since UW is a loss weighting method, validating the superiority of UW with interval clustering compared to previous loss weighting methods such as P2 [6] and MinSNR [17] highlights the effectiveness of our method. We name UW by incorporating interval clustering as Addressing Negative Transfer (ANT)-UW. We trained DiT-L/2 with MinSNR and UW with  $k = 5$  on the ImageNet across 400K iterations, using a batch size of 240. All methods are trained by AdamW optimizer [43] with a learning rate of  $1e - 4$ . Table 2 shows that ANT-UW dramatically outperforms MinSNR, emphasizing the effectiveness of our method. An essential note is that the computational cost of ANT-UW remains remarkably similar to vanilla training as shown in Section 5.3, ensuring that our enhanced performance does not come at the expense of computational efficiency. Additionally, we refer to the results in [50], showing that our ANT-UW outperforms P2 and MinSNR when DI-T-L/2 is trained on the FFHQ dataset.
5.3 Analysis
To provide a better understanding of our method, we present various analysis results here. Specifically, we compare the memory and runtime of MTL methods, analyze the behavior of MTL methods adopted, provide a convergence analysis, and assess the extent to which negative transfer has been addressed.
(a) Average conflict in PCgrad
(b) Gradient weights of NashMTL
(c) Loss weight of UW
Figure 4: Behavior of multi-task learning methods across training iterations. (a): With increasing timestep difference, gradient conflicts between task clusters become more frequent in PCgrad. (b) and (c): Both UW and NashMTL allocate higher weights to task clusters that handle noisier inputs.
Figure 5: Convergence analysis on FFHQ dataset. Compared to baselines, all methods exhibit fast convergence and achieve good final performance.
Memory and Runtime Comparison We first compared the memory usage and runtime between MTL methods and vanilla training for a deeper understanding of their cost. We conducted measurements of memory usage and runtime with  $k = 5$  on the FFHQ dataset using the LDM architecture and timestep-based clustering, and the results are shown in Table 3. PCgrad has a slower speed of 1.523 iterations/second compared to vanilla training, but its GPU memory usage is lower due to the partitioning of minibatch samples. Meanwhile, NashMTL has a runtime of 2.011 iterations/second. Even though NashMTL uses more GPU memory, it has a better runtime than PCgrad because it computes per-interval gradients occasionally. Concurrently, UW shows similar runtime and GPU memory usage as vanilla training, which is attributed to its use of weighted loss and a single backpropagation process.
Behavior of MTL Methods We analyze the behavior of different multi-task learning methods during training. For PCgrad, we calculate the average number of gradient conflicts between task clusters per iteration. For UW, we visualize the weights allocated to the task cluster losses over training iterations. Finally, for NashMTL, we visualize the weights allocated to per-task-cluster gradients over training iterations. We used LDM trained on FFHQ for our experiments. Although we only report results for time-based interval clustering for conciseness, we note that MTL methods exhibit similar behavior across different clustering methods. Results obtained using other clustering methods can be found in Appendix D.1.
The resulting visualizations are provided in Fig. 4. As depicted in Fig. 4a, the task pair that shows the most gradient conflicts is  $I_{1}$  and  $I_{5}$ , namely, task clusters apart in timesteps. This result supports our hypothesis that temporally distant denoising tasks may be conflicting, and as seen in Section 5.1, PCgrad seems to mitigate this issue. Also, as depicted in Fig. 4b and 4b, both UW and NashMTL tend to allocate higher weights to task clusters that handle noisier inputs, namely,  $I_{4}, I_{5}$ . This result suggests that handling noisier inputs may be a difficult task that is underrepresented in conventional diffusion training.
Faster Convergence In Fig. 5, we plot the trajectory of the FID score over training iterations, as observed while training on FFHQ. We can observe that all our methods enjoy faster convergence and better final performance compared to the conventionally trained model. Notably, for pixel space
Figure 6: Negative transfer gap (NTG) comparison on the FFHQ dataset. Integration of MTL methods tends to improve the negative transfer gap. Methods that fail to improve NTG in areas where the baseline records low NTG tend to achieve lesser improvements in the baseline.
diffusion (ADM), UW converges extremely rapidly, while beating the vanilla method by a large margin. Overall, these results show that our method may not only make diffusion training more effective but also more efficient.
Reduced Negative Transfer Gap We now demonstrate that our proposed method indeed mitigates the negative transfer gap we observed in Section 3.2. We used the same procedure introduced in Section 3.2 to calculate the negative transfer gap for all methods considered, for the FFHQ dataset.
As shown in Fig. 6 our methods improve upon negative transfer gaps. Specifically, for tasks that exhibit severe negative transfer gaps in the baseline (e.g., [601, 800], [801, 1000] for ADM, and [401, 600], [601, 800] for LDM), our methods mitigate the negative transfer gap for most cases, even practically removing it in certain cases. Another interesting result to note is that models less effective in reducing negative transfer (NashMTL-SNR for LDM and PCgrad-Grad for ADM) indeed show worse FID scores, which supports our hypothesis that resolving negative transfer leads to performance gain. We also note that even the worst-performing methods still beat the vanilla model.
5.4 Combining MTL Methods with Sophisticated Training Objectives
Finally, we show that our method is readily applicable on top of more sophisticated training objectives proposed in the literature. Specifically, we train an LDM by applying both UW and PCgrad on top of the P2 objective [6] and evaluate the performance on the FFHQ dataset. We chose UW and PCgrad based on a previous finding that combining the two methods leads to performance gain [41]. Also, we chose the gradient-based clustering method due to its effectiveness for LDM on FFHQ. As seen in Table 4, when combined with P2, our method improves the FID from 7.21 to 5.84.
Table 4: Combining our method with P2 on the FFHQ dataset. DDIM 200-step sampler is used.
Type	Method	FID-50k
GAN	PGan [63]	3.39
AR	VQGAN [11]	9.6
Diffusion (LDM)	D2C [65]	13.04
Vanilla	9.1
P2	7.21
P2 + Ours	5.84
In this work, we studied the problem of better training diffusion models, with the distinction of reducing negative transfer between denoising tasks in a multi-task learning perspective. Our key contribution is to enable the application of existing multi-task learning techniques, such as PCgrad and NashMTL, that were challenging to implement due to the increasing computation costs associated with the number of tasks, by clustering the denoising tasks based on their various task affinity scores. Our experiments validated that the proposed method effectively mitigated negative transfer and improved image generation quality. Overall, our findings contribute to advancing diffusion models. Starting from our work, we believe that addressing and overcoming negative transfer can be the future direction to improve diffusion models.
[79] Sen Wu, Hongyang R Zhang, and Christopher Ré. Understanding and improving information transfer in multi-task learning. In International Conference on Learning Representations, 2019. 17
[80] Derrick Xin, Behrooz Ghorbani, Justin Gilmer, Ankush Garg, and Orhan Firat. Do current multi-task optimization methods in deep learning even help? Advances in Neural Information Processing Systems, 35:13597-13609, 2022. 18, 20
[81] Yilun Xu, Shangyuan Tong, and Tommi S. Jaakkola. Stable target field for reduced variance score estimation in diffusion models. In The Eleventh International Conference on Learning Representations, 2023. 4
[82] Xingyi Yang, Daquan Zhou, Jiashi Feng, and Xinchao Wang. Diffusion probabilistic model made slim. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 22552-22562, 2023. 2
[83] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. Advances in Neural Information Processing Systems, 33:5824-5836, 2020. 2, 3, 4, 6, 7, 19, 20, 21
A Relation to Negative Transfer Gap in Previous Literature 17
B Detailed Experimental Settings for Observational Study 17
C Implementation Details for MTL methods 18
D Additional Experimental Results 18
D.1 Visualization for the Behavior of MTL Methods with Other Clustering Methods 19
D.2 Analysis: The Number of Interval Clusters 19
D.3 Comparison Interval Clustering with Task Grouping Method 20
D.4 Comparison to Random Loss Weighting and Linear Scalarization 20
E Detailed Experimental Settings in Section 5 21
E.1 Detailed Settings of Comparative Evaluation and Analysis (Section 5.1 and 5.3) 21
E.2 Detailed Settings of Comparison to Loss Weighting Methods (Section 5.2) 21
E.3 Detailed Settings of Combining MTL Methods with Sophisticated Training Objectives (Section 5.4) 21
F Qualitative Results 21
G Dynamic Programming Algorithm for Interval Clustering 22
H Broader Impacts 23
I Limitations 24
A Relation to Negative Transfer Gap in Previous Literature
Previous works on transfer and multi-task learning have explored measuring the negative transfer [25, 79, 77]. For the source task  $\mathcal{T}_{src}$  and the target task  $\mathcal{T}_{tgt}$ , the negative transfer can be defined as the phenomenon that the source task negatively transfer to the target task. Denote the model trained on both source and target task as  $\Theta (\mathcal{T}_{tgt},\mathcal{T}_{src})$  and the model only trained on the target task as  $\Theta (\mathcal{T}_{tgt})$ . With performance measure  $P$  for the model on  $\mathcal{T}_{tgt}$ , negative transfer can be quantified by utilizing negative transfer gap (NTG):
$$
N T G \left(\mathcal {T} _ {t g t}, \mathcal {T} _ {s r c}\right) = P \left(\Theta \left(\mathcal {T} _ {g t}\right)\right) - P \left(\Theta \left(\mathcal {T} _ {t g t}, \mathcal {T} _ {s r c}\right)\right). \tag {5}
$$
For  $P$ , higher is better,  $NTG > 0$  indicates that negative transfer occurs, showing that additionally training on  $\mathcal{T}_{src}$  negatively affects the learning of  $\mathcal{T}_{tgt}$ .
In our study of negative transfer in diffusion models, the target task involves denoising tasks within a specific timestep interval as  $\mathcal{T}_{tgt} = \mathcal{D}^{[t_1,t_2]}$ , while the source task comprises the remaining denoising tasks as  $\mathcal{T}_{src} = \mathcal{D}^{[1,T]} \setminus \mathcal{D}^{[t_1,t_2]}$ .
However, since a model trained only a subset of entire denoising tasks cannot generate samples properly, we cannot utilize the sample quality metrics (e.g. FID [18]) for  $P$  to measure  $P(\Theta(\mathcal{T}_{gt}))$  in Eq. 5 for arbitrary timestep intervals. This is a different point from a typical MTL setting, where the performance of each task can be measured.
Alternatively, we redefine  $NTG$  with the difference in sample quality resulting from denoising by different models,  $\Theta(\mathcal{T}_{tgt})$  and  $\Theta(\mathcal{T}_{tgt},\mathcal{T}_{src})$ , in the  $[t_1,t_2]$  interval. During the sampling procedure with a model trained on entire denoising tasks, we use  $\Theta(\mathcal{T}_{tgt})$  or  $\Theta(\mathcal{T}_{tgt},\mathcal{T}_{src})$  in  $[t_1,t_2]$ . Denote the resulting samples with  $\Theta(\mathcal{T}_{tgt},\mathcal{T}_{src})$  as  $\{\tilde{x}_0\}$  and the resulting samples with  $\Theta(\mathcal{T}_{tgt})$  as  $\{\tilde{x}_0^{[t_1,t_2]}\}$ . Then, by comparing the quality of these samples as Eq. 3, we can measure how much the denoising of  $[t_1,t_2]$  degrades in terms of sampling quality.
Furthermore, the success of multi-expert denoisers in prior studies [37, 13, 1] suggests the potential existence of negative transfer. By distinctly separating parameters for denoising tasks, they might mitigate this negative transfer, leading to enhanced performance in their generation.
B Detailed Experimental Settings for Observational Study
In this section, we provide the details on experimental settings in Section 3. The training details and the architectures used are the same as those in Section 5. All experiments are conducted with a single A100 GPU and with FFHQ dataset [27].
For the pixel-space diffusion model, we use the lightweight ADM as same in [6]. It inherits the architecture of ADM [8], but it uses fewer base channels, fewer residual blocks, and a self-attention with a single resolution. Specifically, the model uses one residual block per resolution with 128 base channels and  $16 \times 16$  self-attention with 64 head channels. A linear schedule with  $T = 1000$  is used for diffusion scheduling. We referenced the training scripts in the official code<sup>2</sup> for implementation.
For the latent-space diffusion model, we use the LDM architecture as the same settings for FFHQ experiments in [56]. Specifically, an LDM-4-VQ encoder and decoder are used, in which the resolution of latent vectors is reduced by four times compared to the original images and has a vector quantization layer with 8092 codewords. The denoising model has 224 base channels with multipliers for each resolution as 1, 2, 3, 4 and has two residual blocks per resolution. Self-attention with 32 head channels is used for 32, 16, and 8 resolutions. For diffusion scheduling, the linear schedule with  $T = 1000$  is used. We conducted experiments with the official code<sup>3</sup>. In general, we utilized the pre-trained weights provided by LDM. However, if our retraining results demonstrated superior performance, we reported them.
Task Affinity Analysis To measure the task affinity score between denoising tasks, we first calculate  $\nabla_{\theta}L_{t}$  for  $t = 1,\dots ,T$  every 10K iterations during training. The gradient is calculated with 1000 samples in the training dataset. Then, the pairwise cosine similarity of the gradient is computed and
$^{2}$ https://github.com/jychoi118/P2-weighting
<sup>3</sup>https://github.com/CompVis/latent-diffusion