[논문 리뷰] Multistep Consistency Models

Jonathan Heek, Emiel Hoogeboom외 1명|arXiv (Cornell University)|2024. 03. 11.

Complex Systems and Decision MakingDecision Sciences인용 수 3

한 줄 요약

이 논문은 일괄적 일致성 모델과 TRACT의 통합인 다단계 일치 모델을 소개한다. 이는 단일 단계 일치성과 완전한 확산 샘플링 사이를 보간하는 모델이다. 공유된 파라미터를 가진 2~8단계에 걸쳐 세그먼트 단위로 일치성 모델을 훈련시킴으로써, 8단계에서 이미지 품질 기준 최고 성능(FID 1.4, ImageNet64; FID 2.1, ImageNet128)을 달성하였으며, 표준 확산 모델의 품질을 유지하면서도 고속 샘플링을 구현하였다.

ABSTRACT

Diffusion models are relatively easy to train but require many steps to generate samples. Consistency models are far more difficult to train, but generate samples in a single step. In this paper we propose Multistep Consistency Models: A unification between Consistency Models (Song et al., 2023) and TRACT (Berthelot et al., 2023) that can interpolate between a consistency model and a diffusion model: a trade-off between sampling speed and sampling quality. Specifically, a 1-step consistency model is a conventional consistency model whereas a $\infty$-step consistency model is a diffusion model. Multistep Consistency Models work really well in practice. By increasing the sample budget from a single step to 2-8 steps, we can train models more easily that generate higher quality samples, while retaining much of the sampling speed benefits. Notable results are 1.4 FID on Imagenet 64 in 8 step and 2.1 FID on Imagenet128 in 8 steps with consistency distillation, using simple losses without adversarial training. We also show that our method scales to a text-to-image diffusion model, generating samples that are close to the quality of the original model.

연구 동기 및 목표

• 표준 확산 모델과 저단계 일치 모델 간의 성능 격차를 해소하기 위해.
• 다단계 추론을 통해 샘플링 속도와 생성 품질 사이의 부드러운 트레이드오���을 가능하게 하기 위해.
• TRACT과 일치성 디스틸레이션을 일반화함으로써 일치성 모델링의 훈련 안정성과 샘플 품질을 향상시키기 위해.
• 통합 오차를 보정하고 저단계 생성에서의 흐림을 줄이기 위한 결정론적 샘플러(aDDIM)를 개발하기 위해.
• ImageNet128과 텍스트-이미지 생성과 같은 복잡한 벤치마크에서 뛰어난 성능을 보여주기 위해.

제안 방법

• 방법은 확산 과정을 사전 정의된 세그먼트로 나누며, 각 세그먼트에 대해 별도의 일치성 모델을 훈련하지만, 모든 세그먼트 간에 파라미터를 공유한다.
• 일치성 훈련(CT)과 일치성 디스틸레이션(CD)을 다단계로 일반화하여, 더 부드럽고 학습 가능성이 높은 노이즈 제거 궤적을 가능하게 한다.
• 일치성 모델링에서 유도된 단계 스케줄 안내와 동기화된 드롭아웃을 도입하여 훈련 안정성을 향상시킨다.
• 새로운 결정론적 샘플러인 조정된 DDIM(aDDIM)는 노이즈 예측을 확장하여 통합 오차를 보정하고, 저단계 샘플링에서의 흐림을 줄인다.
• 이 프레임워크는 1단계 일치 모델과 무한 단계의 전체 확산 모델 사이를 보간할 수 있으며, 8단계 버전은 표준 확산 모델 성능을 그대로 재현한다.
• 디스틸레이션을 위해, 이 방법은 사전 훈련된 확산 모델에서 고품질의 16단계 일치 모델을 aDDIM을 사용해 훈련한다.

실험 결과

연구 질문

• RQ1 일치성 모델은 단일 단계 샘플링을 초월하여, 성능 향상과 함께 속도 이점을 유지할 수 있는가?
• RQ2 공유된 파라미터를 가진 다단계 일치성 훈련은 단일 단계 일치성 또는 표준 확산보다 더 뛰어난 성능을 낼 수 있는가?
• RQ3 aDDIM과 같은 결정론적 샘플러는 ImageNet128과 같은 복잡한 데이터셋에서 소수의 단계로 경쟁 가능한 FID 점수를 달성할 수 있는가?
• RQ4 단계 스케줄 안내는 저단계 다단계 일치 모델의 성능에 어떤 영향을 미치는가?
• RQ5 다단계 일치 모델은 최소한의 샘플링 단계로 텍스트-이미지 생성에서 표준 확산 모델의 성능을 재현할 수 있는가?

주요 결과

• 제안된 다단계 일치 모델은 8단계 샘플링만으로 ImageNet64에서 SOTA FID 1.4, ImageNet128에서 FID 2.1을 달성하였다.
• 4단계 샘플링으로도 ImageNet64에서 FID 1.6, ImageNet128에서 FID 2.3을 달성하여 저비용 추론에서도 뛰어난 성능을 보였다.
• 단지 8단계만으로도 표준 확산 모델 성능(예: 100단계 DDIM)을 그대로 재현하여 품질 격차를 효과적으로 해소하였다.
• 단계 스케줄 안내가 저단계 모델의 성능을 크게 향상시켰으며, 단계 수가 증가함에 따라 수익 감소 효과가 나타났다.
• aDDIM 샘플러 덕분에 ImageNet128에서도 경쟁 가능한 FID 점수를 확보하였으며, 이는 이전에 신뢰할 수 있는 결정론적 저단계 기반 벤치마크가 없던 분야였다.
• 텍스트-이미지 생성에서 16단계 디스틸레이션 모델은 동일한 난수 시드를 사용할 경우 원본 100단계 DDIM 모델과 거의 구분되지 않는 샘플을 생성하였다.