[논문 리뷰] Optimizing DDPM Sampling with Shortcut Fine-Tuning

Ying Fan, Kangwook Lee|arXiv (Cornell University)|2023. 01. 31.

Machine Learning in Healthcare인용 수 4

한 줄 요약

이 논문은 DDPM 샘플링을 최적화하기 위해 적분 확률 거리 측도(IPM)를 직접 최소화하는 강화학습 기반의 방법인 슈트컷 퍼리티파이닝(SFT)을 제안한다. 이는 역방향 확산 과정을 건너뛰며, 샘플링을 정책 그래เดียน트 문제로 설정함으로써 10단계의 샘플링으로도 SOTA 수준의 샘플 품질을 달성한다. CIFAR-10과 CelebA에서 각각 FID 점수 2.28과 2.01을 기록하며, 1000단계의 전체 단계 DDPM을 초월한다.

ABSTRACT

In this study, we propose Shortcut Fine-Tuning (SFT), a new approach for addressing the challenge of fast sampling of pretrained Denoising Diffusion Probabilistic Models (DDPMs). SFT advocates for the fine-tuning of DDPM samplers through the direct minimization of Integral Probability Metrics (IPM), instead of learning the backward diffusion process. This enables samplers to discover an alternative and more efficient sampling shortcut, deviating from the backward diffusion process. Inspired by a control perspective, we propose a new algorithm SFT-PG: Shortcut Fine-Tuning with Policy Gradient, and prove that under certain assumptions, gradient descent of diffusion models with respect to IPM is equivalent to performing policy gradient. To our best knowledge, this is the first attempt to utilize reinforcement learning (RL) methods to train diffusion models. Through empirical evaluation, we demonstrate that our fine-tuning method can further enhance existing fast DDPM samplers, resulting in sample quality comparable to or even surpassing that of the full-step model across various datasets.

연구 동기 및 목표

• 고품질 생성을 위해 많은 단계(T ≈ 1000)가 필요한 표준 DDPM 샘플링의 비효율성을 해결하기 위해.
• 학습된 역방향 확산 과정을 따르지 않고도 샘플링을 향상시킬 수 있는지 탐색하기 위해.
• 역방향 과정을 모방하는 대신 IPM 최소화를 통해 직접적으로 샘플 품질을 최적화하는 방법을 개발하기 위해.
• DDPM에 대한 IPM에 대한 그래디언트 디센트와 정책 그래데이언트 간의 이론적 연결을 수립하여, 확산 모델의 RL 기반 학습을 가능하게 하기 위해.

제안 방법

• DDPM 샘플링을 강화학습 문제로 설정하는 새로운 알고리즘인 SFT-PG를 제안하며, 여기서 생성기(DDPM 샘플러)는 정책 그래데이언트를 통해 최적화된다.
• DDPM 샘플러에 대한 IPM에 대한 그래디언트 디센트가 수학적으로 확률적 정책 그래데이언트와 동일함을 보여주어, RL 스타일의 학습이 가능하다.
• 학습 중 단계적인 향상 조건을 이론적으로 해석할 수 있도록 도와주는 IPM의 서브티튜트 함수를 도입한다.
• 정책 그래데이언트 학습의 안정성을 높이기 위해 크리틱-생성기 번갈아 학습 방식을 사용하며, 그래디언트 페널티와 베이스라인 정규화를 적용한다.
• 고정된 분산 스케줄과 FastDPM에서 유도된 서브샘플링 스케줄을 초기화로 사용하며, 오직 평균 예측 헤드만을 피니튜닝한다.
• 그래디언트 클리핑과 크리틱 업데이트 당 다중 생성기 단계를 적용하여 학습 안정성과 성능을 향상시킨다.

Figure 1 : Image denoising is similar to a closed-loop control system: finding paths from pure noise to natural images.

실험 결과

연구 질문

• RQ1 역방향 확산 과정을 따르지 않고, 대신 다른 샘플링 경로를 발견함으로써 DDPM 샘플링 품질을 향상시킬 수 있는가?
• RQ2 IPM에 대해 DDPM 샘플러에 대한 그래디언트 디센트는 정책 그래데이언트와 동일한가? 이는 확산 모델의 RL 기반 학습을 가능하게 하는가?
• RQ3 IPM의 서브티튜트 함수는 피니튜닝 중 안정적이고 단조로운 향상을 이끌 수 있는가?
• RQ4 직접적인 IPM 최소화는 FastDPM 및 Analytic-DPM과 같은 유사한 빠른 샘플링 방법보다 더 적은 단계로 더 높은 샘플 품질을 달성할 수 있는가?
• RQ5 SFT-PG를 사용한 피니튜닝은 오직 T′=10개의 샘플링 단계로도 전체 단계 DDPM(T=1000)과 유사하거나 더 뛰어난 성능을 달성할 수 있는가?

주요 결과

• SFT-PG는 10단계 샘플링으로 CIFAR-10에서 FID 2.28, CelebA에서 FID 2.01을 기록하며, 전체 단계 DDPM(FID 3.03 및 3.26)을 초월한다.
• 스위스롤 데이터셋에서 SFT-PG는 T=100일 때 FID 2.36에서 T′=10일 때 FID 0.64로 감소시켜 매우 적은 단계로도 뚜렷한 향상을 보였다.
• 기존의 빠른 샘플러를 향상시켰다: SFT-PG는 CIFAR-10에서 T′=10일 때 FastDPM(FID 29.43), Analytic-DPM(FID 22.94), SN-DPM(FID 16.33)보다 더 낮은 FID를 기록했다.
• 5개의 생성기 단계와 γ=1.0로 학습할 경우, MNIST에서 FID 0.82를 기록하여 최고 성능을 달성했으며, 이는 그래디언트 클리핑과 단계 수가 성능에 영향을 준다는 것을 보여준다.
• 베이스라인 정규화를 적용한 SFT-PG B는 토이 데이터셋에서 SFT와 SFT-PG보다 略적으로 더 높은 최종 성능을 기록하여 안정성 향상을 확인했다.
• 이론적 분석을 통해 IPM 기반의 DDPM에 대한 그래디언트 디센트가 정책 그래데이언트와 동일하다는 것을 확인하였으며, 이는 확산 모델 학습에 RL 방법을 적용한 최초의 사례이다.

Figure 2 : A visual illustration of the key idea of Shortcut Fine-Tuning (SFT). DDPMs aim at learning the backward diffusion model, but this approach is limited to a small number of steps. We propose the idea of not following the backward process and exploring other unexplored paths that can lead to