[논문 리뷰] Efficient Diffusion Training via Min-SNR Weighting Strategy

Tiankai Hang, Shuyang Gu외 6명|arXiv (Cornell University)|2023. 03. 16.

Domain Adaptation and Few-Shot Learning인용 수 5

한 줄 요약

이 논문은 분산 훈련을 다중 작업 학습 문제로 간주함으로써 시간대별 가중치를 신호 대 잡음비(SNR) 기반으로 할당하는 간단하면서도 효과적인 손실 가중 전략인 Min-SNR-γ를 제안한다. 이는 수렴 속도를 크게 향상시키며, 더 작은 모델을 사용하여 ImageNet 256×256에서 FID 점수 2.06의 새로운 SOTA 성능을 달성한다.

ABSTRACT

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$ imes$ faster than previous weighting strategies. It is also more effective, achieving a new record FID score of 2.06 on the ImageNet $256 imes256$ 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.

연구 동기 및 목표

• 시간대 간 최적화 방향의 갈등으로 인해 노이즈 제거 분산 모델 훈련이 느리게 진행되는 문제를 해결하기 위해.
• 각 시간대를 개별 작업으로 간주하는 다중 작업 학습 문제로 분산 훈련을 재정의하기 위해.
• 적응형 파레토 최적화의 불안정성과 계산 비용을 피하면서도 안정적이고 효율적이며 확장 가능한 손실 가중 전략을 개발하기 위해.
• 아키텍처 수정 없이도 훈련 효율성과 생성 품질을 향상시키고, ImageNet 벤치마크에서 최신 기술 수준의 성능을 달성하기 위해.

제안 방법

• 다중 작업 학습 프레임워크에서 각 노이즈 제거 시간대를 개별 작업으로 간주하기 위해.
• 클램프된 신호 대 잡음비(SNR)를 가중치 요소로 사용하는 사전 정의된 전역 손실 가중 전략인 Min-SNR-γ를 도입하기 위해.
• 극단적인 가중치를 방지하기 위해 SNR에 클램프 함수를 적용하여 훈련 안정성 향상과 기울기 갈등 감소를 위해.
• 클램프 강도를 제어하는 하이퍼파라미터 γ를 사용하여 다양한 훈련 설정에서의 강인성 확보하기 위해.
• 실행 시간 적응형 가중치를 피하기 위해 고정된 단계 기반 손실 가중 일정을 사용하여 효율성과 안정성 향상시키기 위해.
• 아키텍처 수정 없이 표준 아키텍처(UNet, ViT)를 사용하여 새로운 손실 가중 전략으로 모델 훈련하기 위해.

실험 결과

연구 질문

• RQ1 시간대 간 공유된 모델 가중치가 존재함에도 불구하고 분산 훈련이 왜 느리게 수렴하는가?
• RQ2 시간대 간 충돌하는 최적화 방향을 체계적인 손실 가중 전략을 통해 완화할 수 있는가?
• RQ3 적응형 다중 작업 학습 방법(예: 파레토 최적화)과 비교해 전역적이고 사전 정의된 손실 가중 전략은 분산 훈련에서 어떻게 성능을 내는가?
• RQ4 간단한 SNR 기반 가중 전략이 수렴 속도와 FID 점수에서 기존 베이스라인을 능가할 수 있는가?
• RQ5 Min-SNR-γ는 다양한 아키텍처와 생성 목표(예: ε 대 x₀ 예측)에 일반화되는가?

주요 결과

• Min-SNR-γ는 이전의 가중 전략보다 3.4배 빠른 수렴 속도를 달성하여 훈련을 크게 가속화한다.
• 이 방법은 ImageNet 256×256 벤치마크에서 FID 점수 2.06을 기록하며 새로운 SOTA 성능을 달성한다.
• 더 작은 모델(예: 395M 파라미터)을 사용해도 ImageNet 256×256에서 FID 2.81을 달성하며, DiT-XL-2와 같은 더 큰 모델을 능가한다.
• Min-SNR-5를 사용해 훈련한 ViT-XL 모델는 단지 210만 번의 반복만에 FID 2.08을 달성하여 DiT 대비 3.3배 더 빠르게 수렴한다.
• 이 방법은 UNet과 ViT 백본 모두에서 성능 향상을 이끌었으며, ε 및 x₀ 예측 목표 모두에 효과적으로 작용한다.
• 제거 실험 결과, 하이퍼파라미터 γ에 대해 강인하며 다양한 설정에서 성능 저하가 최소한이었다.