[논문 리뷰] Soft Truncation: A Universal Training Technique of Score-based Diffusion Model for High Precision Score Estimation

Dong‐Jun Kim, Seungjae Shin외 3명|arXiv (Cornell University)|2021. 06. 10.

Advanced Neuroimaging Techniques and Applications인용 수 10

한 줄 요약

이 논문은 고정된 절단 하이퍼파rameter 대신 랜덤 변수를 도입함으로써 모든 확산 시간 동안 개선된 스코어 추정을 가능하게 하는 보편적인 학습 기법인 소프트 절단(Soft Truncation)을 제안한다. 각 학습 스텝에서 랜덤한 최소 확산 시간을 샘플링함으로써, CIFAR-10, CelebA, CelebA-HQ, STL-10에서 음의 로그우도(NLL)와 프레셰 인ception 거리(FID) 모두에서 최신 기술 수준의 성능을 달성하며, 오랫동안 지속된 밀도 추정과 샘플 품질 간의 상충관계를 해결한다.

ABSTRACT

Recent advances in diffusion models bring state-of-the-art performance on image generation tasks. However, empirical results from previous research in diffusion models imply an inverse correlation between density estimation and sample generation performances. This paper investigates with sufficient empirical evidence that such inverse correlation happens because density estimation is significantly contributed by small diffusion time, whereas sample generation mainly depends on large diffusion time. However, training a score network well across the entire diffusion time is demanding because the loss scale is significantly imbalanced at each diffusion time. For successful training, therefore, we introduce Soft Truncation, a universally applicable training technique for diffusion models, that softens the fixed and static truncation hyperparameter into a random variable. In experiments, Soft Truncation achieves state-of-the-art performance on CIFAR-10, CelebA, CelebA-HQ 256x256, and STL-10 datasets.

연구 동기 및 목표

• 스코어 기반 확산 모델에서 밀도 추정(NLL)과 샘플 생성 품질(FID) 간의 반대 상관관계를 해결하기 위해.
• 특히 NLL가 가장 민감한 저노이즈 영역에서 전체 확산 시간 범위에 걸쳐 스코어 네트워크 학습을 향상시키기 위해.
• FID 또는 NLL에 최적화된 별도의 모델 설정이 필요 없도록 통합된 학습 기법을 도입함으로써 이를 제거하기 위해.
• 구조적 변경 없이 가능하고 이론적으로 탄탄하며, likelihood 및 생성 성능를 모두 향상시키는 보편적으로 적용 가능한 방법을 제공하기 위해.

제안 방법

• 고정된 ε 대신 랜덤한 절단 변수 τ를 도입하여, 각 미니배치마다 재샘플링되며 스코어 추정을 위한 최소 확산 시간을 정의한다.
• 스코어 네트워크를 [τ, T] 간격에서만 학습시키며, τ 이하의 스코어는 무시함으로써, 무작위 샘플링을 통한 전체 시간 범위의 커버리지 확보.
• 이 방법은 확산 시간에 대한 가중치 함수를 암묵적으로 유도하며, 손실 내 일반 기대값과 동치되며, 이는 NLL와 FID를 균형 있게 조절하는 데 성공한 이유를 설명한다.
• 소프트 절단은 어떤 스코어 기반 확산 모델이나 SDE 설정과도 호환되며, 모델 아키텍처나 추론 과정에 변화가 필요 없다.
• 학습 목표는 확산 모델에 특화된 MLE의 일반화된 형태인 최대 흐트러진 우도 추정(MPLE)으로 정의된다.
• 이 기법은 VE, VP, ODE 기반 설정을 포함한 다양한 데이터셋과 SDE에서 실험적으로 검증되었다.

실험 결과

연구 질문

• RQ1 단일 학습 기법이 확산 모델에서 동시에 밀도 추정(NLL)과 샘플 품질(FID)을 향상시킬 수 있는가?
• RQ2 실제로 고정된 절단 하이퍼파rameter ε가 NLL와 FID 간의 상충관계를 초래하는 이유는 무엇인가?
• RQ3 학습 중 최소 확산 시간을 랜덤화하면 전체 시간 범위에 걸쳐 더 나은 스코어 추정이 가능한가?
• RQ4 소프트 절단은 이론적으로 가중 우도 목표 함수와 동치인가, 만약 그렇다면 가중치는 어떻게 유도되는가?
• RQ5 소프트 절단은 다양한 모델 아키텍처, SDE 설정 및 데이터셋에 일반화되는가?

주요 결과

• CIFAR-10에서 소프트 절단은 VP SDE, DDPM++를 사용해 FID 3.45를 달성하였으며, FID에 유리한 모델 수준에 도달하면서도 NLL 3.03을 유지해 NLL에 유리한 모델 수준과 유사한 성능을 보였다.
• CelebA-HQ 256에서 소프트 절단은 FID 3.96과 NLL 3.01을 기록하여, FID 최적화 및 NLL 최적화를 위한 기준 모델을 모두 능가했다.
• 절단 분석에서 소프트 절단이 τ ∼ P₁일 경우, 고정된 ε=10⁻⁵에 비해 CIFAR-10(RVE, UNCSN++)에서 FID를 2.74점 감소시켰으며, 낮은 NLL를 유지했다.
• 고정된 절단 대비 소프트 절단은 CIFAR-10에서 최대 FID 2.74점 감소, CelebA에서 1.26점 감소를 기록했으며, NLL에 악영향을 주지 않았다.
• P₁(구간 [0.1, 1.0]에서 균일 분포)을 사용한 소프트 절단이 최적의 균형을 이룩했으며, CIFAR-10(VP, DDPM++)에서 FID를 6.70에서 3.96으로 감소시켰고, NLL는 안정적으로 유지되었다.
• 이 기법은 아키텍처에 관계없이 강건하다: CIFAR-10에서는 FID를 1.26~2.74점 향상시키며, CelebA에서는 0.7~1.2점 향상되었으며, 일관된 NLL 성능 유지를 보였다.