Efficient Diffusion Training via Min-SNR Weighting Strategy

  • 2023-03-16 18:59:56
  • Tiankai Hang, Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, Baining Guo
  • 5

Abstract

Denoising diffusion models have been a mainstream approach for imagegeneration, however, training these models often suffers from slow convergence.In this paper, we discovered that the slow convergence is partly due toconflicting optimization directions between timesteps. To address this issue,we treat the diffusion training as a multi-task learning problem, and introducea simple yet effective approach referred to as Min-SNR-$\gamma$. This methodadapts loss weights of timesteps based on clamped signal-to-noise ratios, whicheffectively balances the conflicts among timesteps. Our results demonstrate asignificant improvement in converging speed, 3.4$\times$ faster than previousweighting strategies. It is also more effective, achieving a new record FIDscore of 2.06 on the ImageNet $256\times256$ benchmark using smallerarchitectures than that employed in previous state-of-the-art.

 

Quick Read (beta)

loading the full paper ...