PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Diffusion Training via Min-SNR Weighting Strategy

Tiankai Hang, Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, Baining Guo

arXiv:2303.09556v3cs.CV

TL;DR

Diffusion training converges slowly partly because optimization directions conflict across timesteps, complicating effective experimentation. Min-SNR-γ treats timesteps as multi-task objectives and uses loss weighting to balance them, accelerating convergence and achieving a record ImageNet-256 FID score.

  • Problem

    Slow convergence in diffusion training is linked to conflicting optimization directions across timesteps, which challenge effective experimentation.

  • Method

    Min-SNR-γ treats each timestep’s denoising process as an individual task and uses a predefined global loss-weighting strategy to mitigate gradient conflicts.

  • Results

    3.4× faster convergence and a record FID score of 2.06 on the ImageNet 256×256 benchmark were achieved with the proposed weighting strategy.

  • Takeaways & Limitations

    The strategy effectively balances different timesteps and improves diffusion training across prediction targets and network architectures.

  • Takeaways & Limitations

    The analysis assumes the network predicts the noiseless state x0, while noting that different prediction objectives can be transformed into one another.

Abstract

from arXiv · show

Denoising diffusion models have been a mainstream approach for image generation, however, training these models often suffers from slow convergence. In this paper, we discovered that the slow convergence is partly due to conflicting optimization directions between timesteps. To address this issue, we treat the diffusion training as a multi-task learning problem, and introduce a simple yet effective approach referred to as Min-SNR-$γ$. This method adapts loss weights of timesteps based on clamped signal-to-noise ratios, which effectively balances the conflicts among timesteps. Our results demonstrate a significant improvement in converging speed, 3.4$\times$ faster than previous weighting strategies. It is also more effective, achieving a new record FID score of 2.06 on the ImageNet $256\times256$ benchmark using smaller architectures than that employed in previous state-of-the-art. The code is available at https://github.com/TiankaiHang/Min-SNR-Diffusion-Training.

1. Introduction

The paper attributes slow diffusion-training convergence to conflicting gradients across timesteps and proposes Min-SNR-γ, a loss-weighting strategy that balances these conflicts. The strategy accelerates convergence and achieves a new ImageNet 256×256 FID record.

  • Results: 3.4 times faster convergence is achieved than the baseline with a non-conflicting weighting strategy.Figure 1 attributes the speedup to the weighting strategy and reports superior performance.
  • Problem: Conflicting optimization directions across timesteps can slow diffusion training because shared model weights receive incompatible gradient signals.Optimizing denoising for one noise level can harm reconstruction at other noise levels.
  • Method: Min-SNR-γ treats denoising at each timestep as an individual task in a multi-task learning formulation.The method assigns timestep loss weights according to task difficulty.
  • Method: Clamped signal-to-noise ratios are used as loss weights to alleviate conflicting gradients across timesteps.This is a predefined global step-wise weighting strategy rather than run-time adaptive weighting.
  • Results: The strategy produces Pareto objective values almost as low as the optimal step-wise loss weights and better than other existing weighting strategies.This result indicates alleviation of the gradient-conflicting issue.
  • Results: FID 2.06 establishes a new record on the ImageNet 256×256 benchmark.The strategy also improves models with other prediction targets and network architectures.

2. Related Works

The related work positions diffusion models as strong generative models across modalities and connects their training challenges to multi-task learning, where jointly learned tasks can experience negative transfer.

  • Denoising Diffusion Models: Diffusion models are strong generative models because they can model complex distributions and produce high-fidelity, diverse images.Their applications include text-to-video, 3D avatar, image translation, image manipulation, music generation, and drug discovery.
  • Multi-task Learning: Multi-task learning jointly learns related tasks so knowledge from one task can benefit other tasks.Negative transfer is a central challenge because joint training can hurt learning instead of helping it.

3. Method

This section formulates diffusion training as a multi-task optimization problem over timesteps and introduces Min-SNR-γ weighting to reduce conflicts between those tasks. The strategy uses clamped SNR weights and is closest to Pareto optimality among the compared predefined weighting schemes.

  • 3.1. Preliminary: Diffusion models use forward Gaussian noising and reverse Gaussian denoising processes, with shared denoising parameters often reused across all timesteps.The model can predict noise, the noiseless state x0, or velocity, with equivalent objectives obtainable through loss-weight modifications.
  • 3.1. Preliminary: Finetuning selected timestep ranges can benefit nearby timesteps but harm timesteps farther away, indicating conflicting optimization effects.The experiment clusters timesteps into bins, finetunes one bin at a time, and evaluates losses in the other bins.
  • 3.2. Diffusion Training as Multi-Task Learning: Each timestep is treated as an individual task with loss L_t(θ), and the objective seeks an update that decreases every timestep’s loss without sacrificing another task.When no such direction remains, the optimization reaches a Pareto-stationary state; regularization prevents any timestep from receiving an excessively small weight.
  • 3.3. Pareto optimality of diffusion models: Iterative Frank-Wolfe and UGD weighting require extra optimization at every training iteration and produce unstable weights when gradients are estimated from limited samples.Increasing sample counts can stabilize the weights but raises computation cost.
  • 3.4. Min-SNR-γ Loss Weight Strategy: Min-SNR-γ assigns each timestep the weight w_t = min{SNR(t), γ}, preventing excessive focus on small noise levels.The compared alternatives include constant weighting, SNR weighting, Max-SNR-γ, and UGD optimization weighting.
  • 3.4. Min-SNR-γ Loss Weight Strategy: Min-SNR-γ is closest to the Pareto optimum among the compared fixed weighting strategies, supporting simultaneous optimization across timesteps.UGD achieves the lowest objective value, while Min-SNR-γ is reported as the closest to the optimum and is intended to improve convergence.

4. Experiments

The experiments evaluate Min-SNR-γ across weighting designs, timestep ranges, prediction targets, architectures, and image-generation benchmarks. Across these settings, it converges faster, lowers training loss, and improves FID, with strong results on CelebA and ImageNet.

  • Experimental setup: The experiments cover unconditional and conditional image generation on CelebA and ImageNet, using ViT and UNet backbones.ImageNet is evaluated at 64 × 64 and 256 × 256 resolutions.
  • Weighting designs: Min-SNR-γ uses wt = min{SNR(t), γ}, with γ = 5 as the default value for predicting x0.The comparison includes constant, SNR, truncated-SNR, and Min-SNR-γ weighting.
  • Convergence analysis: 3.4× speedup is achieved in reaching an FID score of 10 compared with other weighting methods on ImageNet 256×256 using ViT-B.The SNR weighting strategy performed worst, which may reflect disproportionate focus on less noisy stages.
  • Convergence analysis: Min-SNR-γ achieves lower unweighted training loss across all analyzed timestep ranges, whereas constant and SNR weighting favor different noise-intensity regimes.The timestep ranges are [0, 100), [200, 300), [600, 700), and [800, 900).
  • Prediction targets: Min-SNR-γ converges faster for noise ϵ and velocity v prediction, while constant or Max-SNR-γ settings can diverge when predicting noise.The results indicate that balancing timestep loss weights is independent of re-parameterization.
  • Architectures and benchmarks: The strategy remains effective across architectures: UNet experiments show faster convergence and better FID, while CelebA ViT-Small reaches FID 2.14.On ImageNet 256×256, ViT-XL reaches FID 2.08 after 2.1M iterations and FID 2.06 after longer training with Min-SNR-5.

5. Conclusion

The paper attributes slow diffusion-training convergence to conflicting optimization directions across timesteps and addresses this with Min-SNR-γ loss weighting. Experiments report faster training and state-of-the-art FID on ImageNet-256.

  • Conclusion: Conflicting optimization directions between timesteps may cause slow convergence in diffusion training.The conclusion frames this conflict as the central training issue addressed by the paper.
  • Conclusion: Min-SNR-γ treats diffusion training as multi-task learning and balances timestep contributions with a novel weighting strategy.Each timestep is regarded as an individual task.
  • Conclusion: The method boosts diffusion training several times faster and achieves the state-of-the-art FID score on ImageNet-256.The conclusion reports these outcomes at the paper level without specifying additional benchmark values.

A. Proof for Theorem 1

The proof formulates diffusion training as simultaneous optimization of timestep losses and derives an update direction that does not increase any task loss. Regularization stabilizes the direction search, while a nonzero solution yields Pareto improvement or indicates Pareto stationarity.

  • Pareto optimality is reached when changing the parameter point would increase at least one task loss, preventing one task from being sacrificed for another.
  • A shared update direction δ must be nonzero and satisfy the paper’s inequalities for every timestep loss.
  • If a gradient vector u has ⟨g_t, u⟩ ≥ 0 for every task, then −u is an update direction that lowers each task loss.
  • Maximizing the minimum gradient inner product converts the simultaneous-improvement requirement into an optimization problem over u.
  • A regularization term prevents the unconstrained update vector from becoming unbounded and destabilizing training.
  • The regularized solution satisfies ⟨g_t, u⟩ ≥ 0 for all tasks; if no nonzero solution exists, training has reached Pareto stationarity.

C.1. Architecture Settings

The experiments use ViT and UNet architectures across CelebA and ImageNet resolutions, with model configurations adjusted for benchmark comparisons and controlled ablations.

  • ViT-Small is used for CelebA 64 × 64, while ViT-Base is the default backbone for ablations.
  • A 21-layer ViT-Large is used for ImageNet 64 × 64 to enable a relatively fair comparison with U-ViT.
  • ViT-XL matches DiT’s depth, hidden size, and patch size for ImageNet 256 × 256 comparison.
  • The UNet ablation adjusts ADM-based settings to keep parameters and FLOPs close to ViT-B.
  • 296M-parameter ADM architecture training for 900K iterations reaches an FID score of 2.11 on ImageNet 64 × 64.
  • For ImageNet 256 × 256, the 395M LDM setting operates in a 32 × 32 × 4 latent space.

C.2. Training Settings

Training uses AdamW with optimizer coefficients selected by backbone, and sampling settings vary across ablations and ImageNet resolutions.

  • AdamW is the default optimizer, with (β1, β2) set to (0.9, 0.999) for UNet and (0.99, 0.99) for ViT.
  • EDM’s Heun sampler is used unless otherwise specified, with 30 steps for ablations, 20 for ImageNet 64 × 64, and 50 for ImageNet 256 × 256.

D.1. Ablation Study on Pixel Space

The pixel-space ablation evaluates Min-SNR-γ with x0 and ϵ prediction targets using a ViT-B model and conditional guidance on ImageNet 64 × 64.

  • The ablation trains a ViT-B diffusion model for 800K iterations with batch size 512 on ImageNet 64 × 64 pixel space.
  • Both x0 and ϵ prediction targets use the Min-SNR-γ loss weight with γ = 5.
  • A pretrained ADM noisy classifier at 64 × 64 provides conditional guidance.

D.1.1 Min-SNR-γ on EDM

Min-SNR-γ is applied to EDM and is reported to accelerate convergence, with ImageNet 256 × 256 results reaching FID 2.06. The evaluation uses the official ImageNet-64 training setting but a shorter training duration than EDM.

  • Min-SNR-γ helps EDM converge faster, as shown by the changing FID curve during training.The curve tracks FID against training images using Heun sampling with 18 steps (NFE=35).
  • 2.06 FID is achieved on ImageNet 256 × 256 with the proposed model.The result is presented as a visual result in Figure 13.
  • The comparison uses DPM Solver to sample 50k images for FID calculation with classifier guidance in the ImageNet 64 × 64 ablation.
  • The EDM evaluation keeps the official ImageNet-64 batch size and optimizer settings.
  • The EDM experiment was not trained as long as the approximately 2k-epoch EDM ImageNet training because of compute limitations.

D.2. Visual Results on Different Datasets

The paper reports additional randomly synthesized samples across CelebA and conditional ImageNet at 64 × 64, using both UNet and ViT backbones, plus ImageNet 256 × 256 results.

  • The additional visual results include randomly synthesized CelebA 64 × 64 and conditional ImageNet 64 × 64 samples without cherry-picking.
  • 1.60 FID accompanies the CelebA 64 × 64 samples generated with a UNet backbone.
  • 2.28 FID accompanies the ImageNet 64 × 64 samples generated with a ViT backbone.
  • 2.14 FID accompanies the ImageNet 64 × 64 samples generated with a UNet backbone.
  • 2.06 FID accompanies the ImageNet 256 × 256 samples generated with a ViT backbone.
Loading 2303.09556v3…