SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hyojun Go $^{1*}$  JinYoung Kim $^{1*}$  Yunsung Lee $^{2*}$  Seunghyun Lee $^{3*}$  Shinhyeok Oh $^{3}$  Hyeongdon Moon $^{4}$  Seungtaek Choi $^{5\dagger}$
Twelvelabs<sup>1</sup> Wrtn Technologies<sup>2</sup> Riidd<sup>3</sup> EPFL<sup>4</sup> Yanolja<sup>5</sup> {gohyojun15, seago0828}@gmail.com<sup>1</sup>, sung@wrtn.io<sup>2</sup>, {seunghyun.lee shinhyeok.oh}@riiid.co<sup>3</sup>, hyeongdonmoon@epfl.ch<sup>4</sup>, seungtaek.choi@yanolja.com<sup>5</sup>
Diffusion-based generative models have achieved remarkable success in various domains. It trains a shared model on denoising tasks that encompasses different noise levels simultaneously, representing a form of multi-task learning (MTL). However, analyzing and improving diffusion models from an MTL perspective remains under-explored. In particular, MTL can sometimes lead to the well-known phenomenon of negative transfer, which results in the performance degradation of certain tasks due to conflicts between tasks. In this paper, we first aim to analyze diffusion training from an MTL standpoint, presenting two key observations: (O1) the task affinity between denoising tasks diminishes as the gap between noise levels widens, and (O2) negative transfer can arise even in diffusion training. Building upon these observations, we aim to enhance diffusion training by mitigating negative transfer. To achieve this, we propose leveraging existing MTL methods, but the presence of a huge number of denoising tasks makes this computationally expensive to calculate the necessary per-task loss or gradient. To address this challenge, we propose clustering the denoising tasks into small task clusters and applying MTL methods to them. Specifically, based on (O2), we employ interval clustering to enforce temporal proximity among denoising tasks within clusters. We show that interval clustering can be solved using dynamic programming, utilizing signal-to-noise ratio, timestep, and task affinity for clustering objectives. Through this, our approach addresses the issue of negative transfer in diffusion models by allowing for efficient computation of MTL methods. We validate the efficacy of proposed clustering and its integration with MTL methods through various experiments, demonstrating 1) improved generation quality and 2) faster training convergence of diffusion models. Our project page is available at https://gohyojun15.github.io/ANT_diffusion/.
Diffusion-based generative models [20, 66, 71] have accomplished remarkable achievements in various generative tasks, including image [8], video [21, 23], 3D shape [44, 54], and text generation [38]. In particular, they have shown excellent performance and flexibility in a wide range of image generation settings, including unconditional [28, 47], class-conditional [22], and text-conditional image generation [1, 48, 55]. Consequently, improving diffusion models has garnered significant interest.
The framework of diffusion models [20, 66, 71] comprises gradually corrupting the data towards a given noise distribution and its subsequent reverse process. A model is optimized by minimizing the weighted sum of denoising score-matching losses across various noise levels [20, 69] for learning the reverse process. This can be interpreted as diffusion training aiming to train a single shared model to
$^{*}$ Co-first author  $^{1,2,4,5}$ Work done while at Riiid  $\dagger$  Corresponding author
denoising its input across various noise levels. Therefore, diffusion training is inherently multi-task learning (MTL) in nature, where each noise level represents a distinct denoising task.
However, analyzing and improving diffusion models from an MTL perspective remains underexplored. In particular, sharing one model between tasks may lead to competition between conflicting tasks, resulting in a phenomenon known as negative transfer [24, 25, 57, 78], leading to poorer performance compared to learning individual tasks with separate models. Negative transfer has been a critical issue in MTL research, and related works have demonstrated that the performance of multi-task models can be improved by remediating negative transfer [24, 25, 57, 78, 83]. Considering these, we argue that negative transfer should be investigated in diffusion models, and if present, addressing it is a potential direction for improving diffusion models.
In this paper, we characterize how multi-task diffusion model is, and whether there exists negative transfer in denoising tasks. In particular, (O1) we first observe that task affinity [12, 78] between two denoising tasks is negatively correlated with the difference in noise levels, indicating that they may be less conflict as the noise levels become more similar [78]. This suggests that adjacent denoising tasks should be considered more harmonious tasks than non-adjacent tasks in terms of noise levels.
Next, (O2) we observe the presence of negative transfer from diffusion model training. During sampling within a specific timestep interval, utilizing a model trained exclusively on denoising tasks within that interval generates higher-quality samples compared to a model trained on all denoising tasks simultaneously. This finding implies that simultaneously learning all denoising tasks can cause degraded denoising within a specific time interval, indicating the occurrence of negative transfer.
Based on these observations, we focus on improving diffusion models by addressing negative transfer. To achieve this, we first propose to leverage the existing multi-task learning techniques, such as dealing with issues of conflicting gradients [5, 83], differences in gradient magnitudes [42, 46, 64], and imbalanced loss scales [4, 16, 29]. However, unlike previous MTL studies that typically focused on small sets of tasks, the presence of a large number of denoising tasks ( $\approx$  thousands) in diffusion models makes it computationally expensive since MTL methods generally require calculating per-task loss or gradient in each iteration [4, 5, 16, 24, 29, 42, 46, 64, 78, 83].
To address this, we propose a strategy that first clusters the entire denoising tasks and then applies multi-task learning methods to the resulting clusters. Specifically, inspired by (O1), we formulate the interval clustering problem which groups denoising tasks by pairwise disjoint timestep intervals. Based on the interval clustering, we propose timesteps, signal-to-noise ratios, and task affinity score-based interval clustering and show that these can be clustered by dynamic programming as [2, 76, 49]. Through our strategy, we can address the issue of negative transfer in diffusion models by allowing for efficient computation of multi-task learning methods.
We evaluated our proposed methods through extensive experiments on widely-recognized datasets: FFHQ [27], CelebA-HQ [26], and ImageNet [7]. For a comprehensive analysis, we employed various models, including Ablated Diffusion Model (ADM) [8], Latent Diffusion Model (LDM) [56], and Diffusion Transformer (DiT) [52]. These models represent diverse diffusion architectures spanning pixel-space, latent-space, and transformer-based paradigms. Our results underscore a significant enhancement in image generation quality, attributed to a marked reduction in negative transfer. This affirms the merits of our clustering proposition and its synergistic integration with MTL techniques.
Diffusion Models Diffusion models [20, 66, 71] are a family of generative models that generate samples from noise via a learned denoising process. Diffusion models beat other likelihood-based models, such as autoregressive models [62, 75], flow models [9, 10], and variational autoencoders [32] in terms of sample quality, and sometimes outperform GANs [14] in certain cases [8]. Moreover, pre-trained diffusion models can be easily applied to downstream image synthesis tasks such as image editing [30, 45] and plug-and-play generation [13, 15]. From these advantages, several works have applied diffusion models for various domains [3, 23, 38, 44, 54] and large-scale models [48, 56, 58].
Several studies have focused on improving diffusion models in various aspects, such as architecture [1, 8, 28, 52, 82], sampling speed [33, 60, 67], and training objectives [6, 17, 31, 70, 74]. Among these, the most closely related studies are improving training objectives, as we aim to enhance optimization between denoising tasks from the perspective of multi-task learning (MTL). Several works [31, 70, 74]
redesign training objectives to improve likelihood estimation. However, these objectives may lead to sample quality degradation and training instability and require additional techniques such as importance sampling [70, 74] and sophisticated parameterization [31] to be successfully applied. On the other hand, P2 [6] proposes a weighted training objective that prioritizes denoising tasks for certain noise levels, where the model is expected to learn perceptually rich features. Similar to P2, we aim to improve the sample quality of diffusion models from an MTL perspective, and we will show that our method is also beneficial to P2.
As a concurrent work, MinSNR [17] shares a common insight with us that diffusion training is essentially multi-task learning. However, their observation lacks a direct connection to negative transfer in terms of sample quality. They address the instability and inefficiency of multi-task learning optimization in diffusion models, mainly due to a large number of denoising tasks. In contrast, our work delves deeper into exploring negative transfer and task affinity, and we propose the application of MTL methods through task clustering to overcome the identified challenges in MinSNR.
Multi-Task Learning Multi-Task Learning (MTL) is an approach that trains a single model to perform multiple tasks simultaneously [57]. Although sharing parameters between tasks can reduce the overall number of parameters, it may also result in a negative transfer, causing performance degradation because of conflicting tasks during training procedure [24, 25, 57, 78].
Prior works have tracked down three causes of negative transfer: (1) conflicting gradient, (2) the difference in gradient magnitude, and (3) imbalanced loss scale. First, Conflicting gradients among different tasks may negate each other, resulting in poorer updates for a subset of, or even for all tasks. PCgrad [83] and Graddrop [5] mitigate this by projecting conflicting parts of gradients and dropping elements of gradients based on the degree of conflict, respectively. Second, tasks with larger gradients may dominate tasks with smaller gradients due to differences in gradient magnitude across tasks. Different optimization schemes have been proposed to equalize gradient magnitudes, including MGDA-UB [64], IMTL-G [42], and NashMTL [46]. Similarly, imbalanced loss scales may cause tasks with smaller losses to be dominated by those with larger losses. To balance task losses, uncertainty [29], task difficulty [16], and gradient norm [4] is exploited.
Adapting MTL methods and negative transfer formulation to diffusion models is challenging since these techniques are typically designed for scenarios with a small number of tasks and easily measurable individual task performance. Our goal is to address this challenge and demonstrate that observing negative transfer in diffusion models and mitigating it can improve them.
3 Preliminaries and Observation
We first provide the necessary background information on diffusion models and their multi-task nature. Next, we conduct analyses that yield two important observations: (O1) task affinity between two tasks is negatively correlated with the difference in noise levels, and (O2) negative transfer indeed exists in diffusion training, i.e., the model is overburdened with different, potentially conflicting tasks.
3.1 Preliminaries
Diffusion model [20, 66, 71] consists of two processes: a forward process and a reverse process. The forward process  $q$  gradually injects noise into a datapoint  $\mathbf{x}_0$  to obtain noisy latents  $\{\mathbf{x}_1,\dots ,\mathbf{x}_T\}$  as:
$$
\begin{array}{l} q (\mathbf {x} _ {t} | \mathbf {x} _ {0}) = \mathcal {N} (\mathbf {x} _ {t} | a _ {t} \mathbf {x} _ {0}, \sigma_ {t} ^ {2} \mathbf {I}), \quad q (\mathbf {x} _ {t} | \mathbf {x} _ {s}) = \mathcal {N} (\mathbf {x} _ {t} | \alpha_ {t | s} \mathbf {x} _ {s}, (\sigma_ {t} ^ {2} - \alpha_ {t | s} ^ {2} \sigma_ {s} ^ {2}) \mathbf {I}), \quad 1 \leq s <   t \leq T (1) \\ \text {w h e r e} \alpha_ {t}, \sigma_ {t} \text {c h a r a c t e r i z e t h e s i g n a l - t o - n o i s e r a t i o n S N R} (t) = \alpha_ {t} ^ {2} / \sigma_ {t} ^ {2}, \text {a n d} \alpha_ {t | s} = \alpha_ {t} / \alpha_ {s}. \text {H e r e ,} \\ \operatorname {S N R} (t) \text {d e c r e a s e s i n t , s u c h t h a t b y t h e d e s i g n a t e d f i n a l t i m e s t e p t = T , q} (\mathbf {x} _ {T}) \approx \mathcal {N} (\mathbf {0}, \mathbf {I}). \end{array}
$$
The reverse process is a parameterized model trained to restore the original data from data corrupted during the forward process. The widely adopted training scheme uses a simple noise-prediction objective [8, 20, 34, 56, 59] that trains the model to predict the noise component  $\epsilon$  of the latent  $\mathbf{x}_t = \alpha_t\mathbf{x}_0 + \sigma \epsilon$ ,  $\epsilon \sim \mathcal{N}(\mathbf{0},\mathbf{I})$ . More formally, the objective is as follows:
$$
L _ {s i m p l e} = \mathbb {E} _ {t, x _ {0}, \epsilon} [ L _ {t} ], \quad \text {w h e r e} L _ {t} = | | \epsilon - \epsilon_ {\theta} (\mathbf {x} _ {t}, t) | | _ {2} ^ {2}. \tag {2}
$$
Let us denote by  $\mathcal{D}^t$  the denoising task at timestep  $t$  trained by minimizing the loss  $L_{t}$  (Eq. 2). Then, since a diffusion model jointly learns multiple denoising tasks  $\{D_t\}_{t = 1,\dots ,T}$  using a single shared model  $\epsilon_{\theta}$ , it can be regarded as a multi-task learner. Also, we denote by  $\mathcal{D}^{[t_1,t_2]}$  the set of tasks  $\{\mathcal{D}^{t_1},\mathcal{D}^{t_1 + 1},\ldots ,\mathcal{D}^{t_2}\}$  henceforth.
(a) ADM (timestep  $t$ )
(b) ADM (log-SNR)
(c) LDM (timestep  $t$ )
(d) LDM (log-SNR)
Figure 1: Task affinity scores plotted against timestep and log-SNR axes in ADM and LDM. As the timestep and SNR differences decrease, task affinity increases, implying more aligned gradient directions between denoising tasks and reduced negative impact on their joint training.
(a) ADM
Figure 2: Negative transfer gap (NTG) with FID score of ADM and LDM for denoising tasks  $\mathcal{D}^{[\cdot,\cdot]}$ . If NTG is negative,  $\mathcal{D}^{[\cdot,\cdot]}$ -trained model outperforms the entire denoising tasks-trained model in terms of denoising latent  $\{\mathbf{x}_t\}_{t\in [\cdot,\cdot]}$ , showing the occurrence of negative transfer. Negative transfer occurs in both ADM and LDM.
(b) LDM
3.2 Observation
By considering diffusion training as a form of multi-task learning, we can analyze how the diffusion model learns the denoising task. We experimentally analyze diffusion models with two concepts in multi-task learning: 1) Task affinity [72, 12]: measuring which combinations of denoising tasks may yield a more positive impact on performance. 2) Negative transfer [68, 24, 25, 57, 78, 83]: degradation in denoising tasks caused by multi-task learning. We use a lightweight ADM [8] used in [6] and LDM [56] with FFHQ  $256 \times 256$  dataset [27] for analyze diffusion models trained on both pixel and latent space.
(O1) Task Affinity Analysis We first analyze how the denoising tasks  $\mathcal{D}^{[1,T]}$  relate to each other by measuring task affinities [72, 12]. In particular, we adopt the gradient direction-based task affinity score [78]: for two given tasks  $\mathcal{D}^i$  and  $\mathcal{D}^j$ , we calculate the pairwise cosine similarity between gradients from each task loss, i.e.,  $\nabla_{\theta}L_{i}$  and  $\nabla_{\theta}L_{j}$ , then average the similarities across training iterations. Task affinity score assumes that cooperative (conflicting) tasks produce similar (conflicting) gradient directions, and it has been to correlate with the MTL model's overall performance [78]. Although there have been attempts to divide diffusion model phases using signal-to-noise ratio [6] and a trace of covariance of training targets [81], we are the first to provide an explicit and fine-grained analysis of task affinities among denoising tasks.
In Fig. 1, we visualize the task affinity scores among denoising tasks, for both ADM and LDM, with both timestep and log-SNR as axes. As can be seen in Fig. 1, task affinity between two tasks  $\mathcal{D}^i$ ,  $\mathcal{D}^j$  is high for neighboring tasks, i.e.,  $i \approx j$ , and decreases smoothly as the difference in SNRs (or timesteps) increases. This suggests that tasks sharing temporal/noise-level proximity can be cooperatively learned without significant conflict. Also, this result hints at the possibility that denoising tasks for vastly different SNRs (distant in timesteps) may potentially be conflicting.
(O2) Negative Transfer Analysis Next, we show that there exist negative transfers among different denoising tasks  $\mathcal{D}^{[1,T]}$ . Negative transfer refers to a multi-task learner's performance degradation due to task conflicts, and it can be identified by observing the performance gap between a multi-task learner and specific-task learners. For ease of observation, we group up tasks by intervals, based on
the observation (O1) that more neighboring tasks in timesteps have higher task affinity. Specifically, we investigate whether the task group  $\mathcal{D}^{[t_1,t_2]}$  suffers negative impacts from the remaining tasks.
To quantify the negative transfer, we follow the procedure: First, we generate samples  $\{\tilde{\mathbf{x}}_0\}$  using a model trained on all denoising tasks  $\mathcal{D}^{[1,T]}$ . Next, we repeat the same sampling procedure, except we replace the model with a model trained on  $\mathcal{D}^{[t_1,t_2]}$  for the latent  $\{\mathbf{x}_t\}_{t\in [t_1,t_2]}$ ; We denote the resulting samples by  $\{\tilde{\mathbf{x}}_0^{[t_1,t_2]}\}$ . If  $\{\tilde{\mathbf{x}}_0^{[t_1,t_2]}\}$  exhibits superior quality compared to  $\{\tilde{\mathbf{x}}_0\}$ , it indicates that the model trained solely on  $\mathcal{D}^{[t_1,t_2]}$  performs better in denoising the latent  $\{\mathbf{x}_t\}_{t\in [t_1,t_2]}$  than the model trained on the entire denoising task. This suggests that  $\mathcal{D}^{[t_1,t_2]}$  suffers from negative transfer by learning other tasks. More formally, given a performance metric  $P$ , FID [18] in this paper, we define the negative transfer gap:
$$
N T G \left(\mathcal {D} ^ {\left[ t _ {1}, t _ {2} \right]}\right) := P \left(\left\{\tilde {\mathbf {x}} _ {0} ^ {\left[ t _ {1}, t _ {2} \right]} \right\}\right) - P \left(\left\{\tilde {\mathbf {x}} _ {0} \right\}\right), \tag {3}
$$
where  $NTG < 0$  indicates that negative transfer occurs. The relationship between the negative transfer gap in previous literature and our negative transfer gap is described in Appendix A.
We visualize the negative transfers among denoising tasks for both lightweight ADM [6, 8] and LDM [56] in Fig. 2. The results indicate that negative transfer occurs in three out of the five considered task groups for both models. Notably, negative transfers often have a significant impact, such as a 7.56 increase in FID for ADM in the worst case. Therefore, we hypothesize that there is room for improving the performance of diffusion models by mitigating negative transfer, which motivates us to leverage well-designed MTL methods for diffusion training.
4 Methodology
In Section 3.2, we make two observations: (O1) Denoising tasks with a larger difference in  $t$  and  $\mathrm{SNR}(t)$  exhibit lower task affinity, (O2) Negative transfer occurs in diffusion training. Inspired by these observations, we aim to remediate the negative transfer in diffusion by leveraging MTL methods. Although MTL methods are reported effective when there are only a few tasks, they are impractical for diffusion models with a large number of denoising tasks since they require computing per-task gradients or loss at each iteration. In this section, to deal with challenges, we propose a strategy that first groups the denoising tasks as task clusters and then applies the multi-task learning methods by regarding each task cluster as one distinct task.
4.1 Interval Clustering
Here, we first introduce a scheme that groups all denoising tasks  $\mathcal{D}^{[1,T]}$  into a small number of task clusters. This is a necessary step for applying well-established MTL methods, for they usually involve computationally expensive subroutines such as computing per-task gradients or loss in each training iteration. Our key idea is to enforce temporal proximity of denoising tasks within task clusters, given our observation (O1) that task affinity is higher for tasks closer in timesteps. Therefore, we assign tasks in pairwise disjoint time intervals.
To obtain the disjoint time intervals, we leverage an interval clustering algorithm [2, 49] that optimizes for various clustering costs. In our case, interval clustering assigns diffusion timesteps  $\mathcal{X} = \{1,\dots ,T\}$  to  $k$  contiguous intervals  $I_{1},\ldots ,I_{k}$ , with  $\coprod_{i = 1}^{k}I_{i}\cap \mathcal{X} = \mathcal{X}$ , where  $\coprod$  denotes disjoint union. Let  $I_{i} = [l_{i},r_{i}]$ ,  $l_{i}\leq r_{i}$  for  $i = 1,\dots ,k$ , then we have  $l_{1} = 1$ , and  $r_i = l_{i + 1} - 1$  ( $i < k$  and  $r_k = T$ ). The interval clustering problem is defined as:
$$
\min  _ {l _ {1} = 1 <   l _ {2} <   \dots <   l _ {k}} \sum_ {i = 1} ^ {n} L _ {\text {c l u s t e r}} \left(I _ {i} \cap \mathcal {X}\right), \tag {4}
$$
where  $L_{cluster}$  denotes the cluster cost.
Generally, it is known that an interval clustering problem of  $n$  data points with  $k$  intervals can be solved via dynamic programming in  $O(n^{2}k\omega (n))$  [49], where  $\omega (n)$  is the time required to calculate the one-cluster cost for  $L_{cluster}(\mathcal{X})$ . If the size of each cluster is too small, it is challenging to learn the corresponding task cluster, so we add constraints on the cluster size for dynamic programming. More details regarding the dynamic programming algorithm can be found in Appendix G.
It remains to design the clustering cost function  $L_{cluster}$  to optimize for. We present three clustering cost functions: timestep-based, SNR-based, and gradient-based.
1. Timestep-based Clustering Cost Intuitively, one simple clustering cost is based on timesteps. We use the absolute timestep difference for the clustering objective by setting  $L_{cluster}(I_i \cap \mathcal{X}) = \sum_{t=l_i}^{r_i} ||t_{center}^i - t||_1^1$  in Eq. 4 where  $t_{center}^i$  denotes the center of interval  $I_i$ . The resulting intervals divide up the timesteps into  $k$  uniform intervals.
2. SNR-based Clustering Cost Another useful metric to characterize a denoising task is its signal-to-noise ratio (SNR). Indeed, it has been previously observed that a denoising task encounters perceptually different noisy inputs depending on its SNR [6]. Also, we already observed that denoising tasks with similar SNRs show high task affinity scores (see Section 3.2). Based on this, we use the absolute log-SNR difference for clustering cost. We define the clustering cost as  $L_{\text{cluster}}(I_i \cap \mathcal{X}) = \sum_{t=l_i}^{r_i} ||\log \operatorname{SNR}(t_{\text{center}}^i) - \log \operatorname{SNR}(t)||_1^1$ .
3. Gradient-based Clustering Cost Finally, we consider the gradient direction-based task affinity scores (see Section 3.2 for a definition) for clustering cost. Task affinity scores have been used as a metric to group cooperative tasks [78]. Based on a similar intuition, we design a clustering cost as follows:  $L_{cluster}(I_i \cap \mathcal{X}) = -\sum_{t=l_i}^{r_i} \mathrm{TAS}(t_{center}^i, t)$  where  $\mathrm{TAS}(\cdot)$  is the gradient-based task affinity score. While leveraging more fine-grained information regarding task affinities, this cost function requires computing and storing gradients throughout training.
4.2 Incorporating MTL Methods into Diffusion Model Training