[논문 리뷰] Consistent Diffusion Models: Mitigating Sampling Drift by Learning to be Consistent

Giannis Daras, Yuval Dagan외 2명|arXiv (Cornell University)|2023. 02. 17.

Model Reduction and Neural Networks인용 수 7

한 줄 요약

이 논문은 샘플링 드리프트를 완화하기 위해 일관성 속성을 강제하는 새로운 학습 목표인 일관성 확산 모델(CDM)을 제안한다: 모델의 예측이 자신의 생성 샘플에 대해 시간 단계에 걸쳐 일관되게 유지된다. 이 일관성 손실을 약한 노이즈 제거 스코어 매칭 목표와 조합함으로써, 드리프트되지 않은 데이터만으로도 CIFAR-10에서 SOTA FID 점수 5.45를 달성하고 AFHQ 및 FFHQ에서도 향상된 성능을 보였다.

ABSTRACT

Imperfect score-matching leads to a shift between the training and the sampling distribution of diffusion models. Due to the recursive nature of the generation process, errors in previous steps yield sampling iterates that drift away from the training distribution. Yet, the standard training objective via Denoising Score Matching (DSM) is only designed to optimize over non-drifted data. To train on drifted data, we propose to enforce a \emph{consistency} property which states that predictions of the model on its own generated data are consistent across time. Theoretically, we show that if the score is learned perfectly on some non-drifted points (via DSM) and if the consistency property is enforced everywhere, then the score is learned accurately everywhere. Empirically we show that our novel training objective yields state-of-the-art results for conditional and unconditional generation in CIFAR-10 and baseline improvements in AFHQ and FFHQ. We open-source our code and models: https://github.com/giannisdaras/cdm

연구 동기 및 목표

• 역방향 샘플링 중 오차가 누적되고 스코어 함수 학습이 불완전해지면서 발생하는 샘플링 드리프트 문제를 해결하기 위해.
• 훈련 중에 볼 수 없었던 분포 외의 점들(예: 드리프트된 샘플)에 대한 학습된 스코어 함수의 일반화를 향상시키기 위해.
• 훈련 중에 진짜 스코어 함수나 타겟 데이터 분포에 대한 액세스가 필요 없도록 하는 학습 목표를 개발하기 위해.
• 이론적이고 실증적으로, 시간 단계에 걸쳐 일관성을 강제하면 더 정확하고 안정적인 샘플링이 가능하다는 것을 입증하기 위해.
• 샘플링 과정을 수정하지 않고도 이미지 생성 벤치마크에서 SOTA 성능을 달성하기 위해.

제안 방법

• 일관성 속성 도입: 시간 t에서 모델이 시간 t에 생성된 샘플에 대해 내놓는 노이즈 제거 예측은, 동일한 샘플에 대해 시간 t에서의 동일한 예측과 일치해야 한다.
• 모델의 예측이 서로 다른 시간 단계에서 자신의 생성 경로를 따라 일관성이 떨어지지 않도록 하는 일관성 손실을 제안한다.
• 약한 노이즈 제거 스코어 매칭(이하 DSM) 목표와 일관성 손실을 조합하며, 이는 오직 드리프트되지 않은 데이터(즉, 노이즈로 오염된 깨끗한 데이터)에서만 훈련된다.
• DSM 손실과 일관성 손실을 동시에 사용하여 모델을 엔드 투 엔드로 훈련함으로써, 추론 시 드리프트된 샘플에도 일반화할 수 있도록 한다.
• 일관성 손실을 통해, 명시적으로 최적화되지 않은 시간 단계에 대해서도 전체 생성 경로에 걸쳐 스코어 함수를 암묵적으로 정규화한다.
• 모델이 명시적으로 훈련된 시간 단계를 초월해 일반화하는지 검증하기 위해 샘플링 중 조기 정지 기법을 아블레이션으로 활용한다.

(a) Consistency Property Testing on AFHQ. The plot illustrates how the Consistency Loss, $L^{2}_{t,t^{\prime},x_{t}}$ , behaves for $t^{\prime}=0$ , as $t$ changes.

실험 결과

연구 질문

• RQ1 시간 단계에 걸쳐 모델 예측의 일관성을 강제하면 샘플링 드리프트가 완화되는가?
• RQ2 약한 DSM 지도 하에 일관성 기반 정규화가 드리프트된 데이터에 대한 스코어 함수의 일반화를 향상시키는가?
• RQ3 명시적인 드리프트된 데이터 훈련 없이도 또는 샘플링 과정의 수정 없이도 모델이 SOTA 성능을 달성할 수 있는가?
• RQ4 일관성 속성과 학습된 스코어 함수의 정확성 사이의 이론적 관계는 무엇인가?
• RQ5 일관성 손실은 명시적으로 훈련되지 않은 시간 단계로의 일반화 능력에 어떤 영향을 미치는가?

주요 결과

• 제안된 CDM 방법은 CIFAR-10에서 SOTA FID 점수 5.45를 달성하였으며, 베이스라인 EDM 모델(FID 5.81)을 능가했다.
• 아블레이션 연구 결과, 특정 시간 단계에서 DSM 손실을 제거하면 성능이 크게 하락함(FID 6.59), 반면 조기 정지 샘플링은 열악한 결과(FID 14.52)를 보였다.
• 일관성 손실 덕분에 모델은 명시적으로 훈련되지 않은 시간 단계로도 일반화할 수 있었으며, DSM에 사용된 시간 단계를 초월해 샘플링을 계속해도 강력한 성능을 유지함을 입증했다.
• 이론적 분석 결과, 한 시간 단계에서 스코어가 완벽하게 학습되었고 일관성 속성이 전역적으로 유지된다면, 스코어는 전역적으로 완벽하게 학습된다.
• 실증 결과는 AFHQ 및 FFHQ에서도 일관된 향상이 있었음을 보여주며, 이는 CIFAR-10 외 다양한 데이터셋으로의 일반화 가능성을 시사한다.
• 훈련 시간은 약 1.5배 증가하여, 추론 효율성 향상의 여지가 있음을 시사한다.

(b) Consistency Property Testing on AFHQ. The plot illustrates how the Consistency Loss, $L^{2}_{t,t^{\prime},x_{t}}$ , behaves for $t=0$ , as $t^{\prime}$ changes.