[논문 리뷰] Maximum Likelihood Training of Implicit Nonlinear Diffusion Models

Dongjun Kim, Byeonghu Na외 4명|arXiv (Cornell University)|2022. 05. 27.

Generative Adversarial Networks and Image Synthesis인용 수 10

한 줄 요약

이 논문은 데이터에 적응하는 비선형 확산 과정을 학습하기 위해 잠재공간에서 정규화 플로우를 학습하여 데이터 공간에서의 비선형 드리프트와 볼라티리티를 암묵적으로 모델링하는 암묵적 비선형 확산 모델(INDM)을 제안한다. INDM는 CelebA에서 FID 1.75의 최고 성능을 기록하며 최대가능도 추정(MLE)을 통해 훈련되며, 샘플 품질과 훈련 안정성 측면에서 선형 기준 모델인 DDPM++을 능가한다.

ABSTRACT

Whereas diverse variations of diffusion models exist, extending the linear diffusion into a nonlinear diffusion process is investigated by very few works. The nonlinearity effect has been hardly understood, but intuitively, there would be promising diffusion patterns to efficiently train the generative distribution towards the data distribution. This paper introduces a data-adaptive nonlinear diffusion process for score-based diffusion models. The proposed Implicit Nonlinear Diffusion Model (INDM) learns by combining a normalizing flow and a diffusion process. Specifically, INDM implicitly constructs a nonlinear diffusion on the extit{data space} by leveraging a linear diffusion on the extit{latent space} through a flow network. This flow network is key to forming a nonlinear diffusion, as the nonlinearity depends on the flow network. This flexible nonlinearity improves the learning curve of INDM to nearly Maximum Likelihood Estimation (MLE) against the non-MLE curve of DDPM++, which turns out to be an inflexible version of INDM with the flow fixed as an identity mapping. Also, the discretization of INDM shows the sampling robustness. In experiments, INDM achieves the state-of-the-art FID of 1.75 on CelebA. We release our code at https://github.com/byeonghu-na/INDM.

연구 동기 및 목표

• 기존의 확산 모델이 고정된 선형 확산 과정에 의존함으로써 표현 능력과 훈련 효율성에 제약을 받는 문제를 해결하고자 한다.
• 데이터 공간에서 학습 가능한 비선형 전진 SDE를 학습함으로써, 잠재공간에서의 가역적 정규화 플로우를 통해 비선형성과 볼라티리티를 모델링할 수 있는지 탐구하고자 한다.
• 유연한 전방 SDE의 드리프트 및 볼라티리티 계수를 플로우 기반 변환을 통해 암묵적으로 모델링함으로써, 확산 모델에 대해 최대가능도 추정(MLE) 훈련을 가능하게 하고자 한다.
• 다양한 스텝 크기에서의 비선형 SDE 이산화 과정이 안정성을 유지하도록 보장함으로써 샘플링의 안정성을 향상시키고자 한다.
• 적절히 훈련된 비선형 확산이 선형 확산 기준 모델에 비해 가능도 및 샘플 품질 지표에서 뛰어난 성능을 보일 수 있음을 입증하고자 한다.

제안 방법

• INDM은 잠재공간에서의 선형 확산을 비선형 확산 과정으로 변환하기 위해 학습 가능한 정규화 플로우를 적용함으로써 데이터 공간에서 비선형 확산 과정을 구성한다.
• 정규화 플로우는 가역적이며 미분 가능하므로 전진 및 역방향 과정을 통해 역전파가 가능하며, 이는 비선형 SDE를 훈련하는 데 필수적이다.
• 데이터 공간의 전방 SDE는 플로우 매핑을 통해 암묵적으로 정의된다: $\mathbf{x}_t = \mathbf{f}_{\bm{\phi}}(\mathbf{z}_t)$, 여기서 $\mathbf{z}_t$는 선형 SDE를 따른다.
• 모델은 진짜 데이터 스코어를 근사하는 노이즈 제거 스코어 손실을 최소화하여 스코어 네트워크 $\mathbf{s}_{\bm{\theta}}$를 훈련함으로써 변분 경계를 통해 MLE 훈련을 수행한다.
• 손실이 음의 로그가능도의 변분 하한에 해당하도록 하기 위해 가중치 함수 $\lambda(t) = g^2(t)$를 사용한다.
• 샘플링은 사전 분포 $\mathbf{x}_T \sim \mathcal{N}(0, I)$에서 시작하여 시간 역행 SDE를 풀어 수행되며, 이는 학습된 스코어 네트워크와 플로우 기반 드리프트 보정을 통해 유도된다.

실험 결과

연구 질문

• RQ1 데이터 공간에서 학습 가능한 데이터 적응형 비선형 확산 과정이 고정된 선형 확산에 비해 생성 모델링 성능을 향상시킬 수 있는가?
• RQ2 잠재공간의 정규화 플로우를 통한 전방 SDE의 암묵적 모델링이 표준 선형 확산에 비해 드리프트 및 볼라티리티 계수의 훈련 효율성을 높일 수 있는가?
• RQ3 제안된 방법은 확산 모델에서 최대가능도 추정(MLE) 훈련을 달성할 수 있으며, 가능도 및 샘플 품질 측면에서 스코어 기반 훈련과 비교해 어떻게 성능을 내는가?
• RQ4 비선형 확산 모델의 이산화 안정성은 DDPM++과 같은 선형 확산 모델에 비해 어떻게 비교되는가?
• RQ5 플로우 네트워크에 의해 도입된 비선형성은 CIFAR-10 및 CelebA와 같은 벤치마크 데이터셋에서 FID 및 NLL 지표에 얼마나 기여하는가?

주요 결과

• INDM은 CelebA $64 \times 64$에서 기존의 모든 확산 모델, 특히 DDPM++ 및 NCSN++를 능가하는 최고 성능인 FID 1.75를 달성한다.
• CIFAR-10에서는 NLL가 2.97(수정 후)이며, NELBO와 NLL 간 격차가 0.02로, 거의 최적의 가능도 추정 성능을 보인다.
• INDM은 CIFAR-10(VE)에서 FID를 2.28로, VP에서 2.90으로 낮추었으며, 이는 DDPM++(FID 2.64)과 NCSN++(FID 2.20)을 모두 능가하는 가능도 및 샘플 품질 성능을 보였다.
• 모델은 샘플링 이산화에 대해 강력한 내성적 안정성을 보이며, 선형 기준 모델과 달리 굵은 시간 스텝에서도 높은 품질의 생성을 유지한다.
• INDM는 거의 MLE 성능을 달성하였으며, CIFAR-10에서 NELBO와 NLL 간 격차가 단지 0.02로, 효과적인 가능도 최적화를 보여준다.
• 절단 실험 결과, 플로우 네트워크를 통한 비선형 확산이 필수적임을 확인하였다: 플로우를 항등함수로 고정(즉, DDPM++ 유사)할 경우 훈련 곡선이著しく 열 劣화되고 FID가 높아진다.