[논문 리뷰] TRACT: Denoising Diffusion Models with Transitive Closure Time-Distillation

David Berthelot, Arnaud Autef외 7명|arXiv (Cornell University)|2023. 03. 07.

Advanced Neuroimaging Techniques and Applications인용 수 12

한 줄 요약

TRACT는 전이 관계를 이용한 증류를 통해 단일 단계 및 소수 단계 확산 샘플링을 크게 개선하여 아키텍처를 변경하지 않고 CIFAR-10 및 64×64 ImageNet에서 1단계 DDIM의 최첨단 FID 점수를 달성합니다.

ABSTRACT

Denoising Diffusion models have demonstrated their proficiency for generative sampling. However, generating good samples often requires many iterations. Consequently, techniques such as binary time-distillation (BTD) have been proposed to reduce the number of network calls for a fixed architecture. In this paper, we introduce TRAnsitive Closure Time-distillation (TRACT), a new method that extends BTD. For single step diffusion,TRACT improves FID by up to 2.4x on the same architecture, and achieves new single-step Denoising Diffusion Implicit Models (DDIM) state-of-the-art FID (7.4 for ImageNet64, 3.8 for CIFAR10). Finally we tease apart the method through extended ablations. The PyTorch implementation will be released soon.

연구 동기 및 목표

• 확산 모델의 추론 비용을 줄이고 아키텍처 변경 없이 단일 또는 소수 단계 샘플링을 가능하게 한다.
• 객체가 수축되는 BTD의 한계(목표 저하 및 SWA 비호환성 등)를 확인한다.
• 질을 유지하면서 적은 단계로 학습하기 위해 자기 교사 EMA를 사용하여 시간 간격을 다단 증류하는 TRACT를 제안한다.
• TRACT가 CIFAR-10 및 64×64 ImageNet에서 1–2단계로 최첨단 또는 경쟁력 있는 FID를 달성하는지 확인하고 변형을 분석한다.

제안 방법

• 이진 시간 증류(BTD)를 전이적 닫힘 시간 증류(TRACT)로 확장하여 증류 단계 수를 log2(T)에서 소수 상수(1–2)로 줄인다.
• t에서 t'로 학생을 증류하도록 학습시키고, t' < t를 사용하여 전이적 닫힘을 수행하기 위해 자기 교사 EMA를 활용한다(식 6–9).
• 다단 점프의 타깃 생성을 위해 학생 가중치의 자기 교사 EMA를 사용한다(알고리즘 1).
• RK 및 DDIM-VE 교사와 함께 VE/EDM 설정에 TRACT를 적용하고 대응하는 타깃 및 손실을 도출한다(식 11–15).
• 목표 저하를 완화하기 위해 증류 단계를 제한하고 EMA 및 추론 시 EMA를 활용한 자기 교사 학습으로 안정성을 높인다.
• 그룹 기반 증류, 손실 가중치 및 EMA 업데이트를 포함한 학습 세부 정보를 제공한다(부록 참조).

Figure 1 : Transitive Closure Distillation of a group $\{t_{i},\ldots,t_{j}\}$ .

실험 결과

연구 질문

• RQ1 TRACT가 아키텍처 변경 없이 1–2 단계 추론으로 고품질 샘플을 달성할 수 있는가?
• RQ2 증류 단계를 줄이는 것이 목표 저하를 완화하고 효과적인 SWA를 가능하게 하는가?
• RQ3 CIFAR-10 및 64×64 ImageNet에서 VE/EDM 교사 및 대안 샘플러(RK, DDIM-VE)와 함께 TRACT의 성능은 어떤가?

주요 결과

방법	NFEs	FID	매개변수
TRACT-EDM-256M ∗	1	3.78 ± 0.01	56M
TRACT-96M ∗	1	4.17 ± 0.03	56M
TRACT-256M	1	4.45 ± 0.05	60M
BTD-96M [44]	-	9.12	60M
TRACT-96M	2	3.32 ± 0.02	60M
TRACT-EDM-256M ∗	2	3.55 ± 0.01	56M
TRACT-EDM-96M ∗	2	3.75 ± 0.02	56M
BTD-96M [44]	-	4.51	60M
TRACT-96M	1	7.43 ± 0.07	296M
TRACT-EDM-96M ∗	1	7.52 ± 0.05	296M

• 1단계 TRACT는 CIFAR-10 FID를 9.1(BTD)에서 4.5로 개선하며 같은 아키텍처와 예산으로 1단계 설정을 달성합니다.
• 1단계 TRACT는 EDM 교사를 사용한 1단계 샘플링에서 64×64 ImageNet의 FID를 7.4로 달성하여 BTD 기반선 대비 향상됩니다.
• 2단계 TRACT는 32단계 교사에서 1단계로 증류한 CIFAR-10에서 3.32 FID를 달성하고, 64×64 ImageNet에서 단일 단계 증류로 7.43 FID를 달성합니다.
• TRACT-EDM-256M은 CIFAR-10에서 1 NFE로 3.78±0.01 FID를 달성하였고; TRACT-EDM-96M은 CIFAR-10에서 1 NFE로 3.75±0.02 FID를 달성합니다(표 1 및 관련 텍스트).
• 64×64 ImageNet에서 TRACT-96M은 1 NFE로 7.43±0.07 FID를 달성하고, 동일 설정의 TRACT-EDM-96M은 7.52±0.05를 달성합니다(표 2 및 관련 텍스트).
• 변형은 EMA 기반 자기 교차 학습으로 2단계 일정에서 최상의 성능을 보이고, 더 많은 단계는 목표 저하로 인해 성능이 저하됩니다.

Figure 5 : 1-step FID for 2-phases $T:1024\to 32\to 1$ TRACT distilled models. Each curve maps to a different way to set the inference time EMA momentum $\mu$ across training lengths. Dashed lines correspond to fixing a $\mu$ value, solid lines correspond to fixing $\epsilon=\mu^{N}$ .