Addressing Negative Transfer in Diffusion Models

TL;DR

Diffusion-model training across multiple noise levels constitutes a large-scale multi-task learning problem over denoising tasks $\{\mathcal{D}^t\}_{t=1}^T$ with loss $L_{simple}$. The authors identify two key findings: high task affinity for nearby denoising tasks that declines with greater timesteps/SNR gaps (O1) and the presence of negative transfer during joint diffusion training (O2). To scale MTL, they introduce interval clustering to group denoising tasks into contiguous timesteps and apply MTL methods (PCgrad, NashMTL, Uncertainty Weighting) within clusters, guided by a dynamic-programming interval clustering framework with timesteps, SNR, and gradient affinity costs. Extensive experiments on FFHQ, CelebA-HQ, and ImageNet across pixel- and latent-space diffusion models demonstrate reduced negative transfer, improved generation quality, and faster convergence when combining interval clustering with MTL, with additive gains when paired with sophisticated objectives like P2. The work offers a practical, scalable approach to harness MTL for diffusion models and suggests promising future directions in task-embedding design and architecture-aware MTL strategies.

Abstract

Diffusion-based generative models have achieved remarkable success in various domains. It trains a shared model on denoising tasks that encompass 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.

Figures (13)

• 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.
• 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.
• 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.
• 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.