[논문 리뷰] Fast Sampling of Diffusion Models via Operator Learning

Hongkai Zheng, Weili Nie외 3명|arXiv (Cornell University)|2022. 11. 24.

Advanced Neuroimaging Techniques and Applications인용 수 17

한 줄 요약

이 논문은 단일 forward pass를 통해 확산 모델의 빠르고 병렬적인 샘플링을 위한 신규 신경 연산자 기반 방법인 DSNO를 제안한다. 확산 모델의 백본에 푸리에 도메인 시간 컨볼루션 레이어를 통합함으로써 DSNO는 연속 시간 역확산 과정의 궤적을 학습하며, 한 번의 모델 평가로 CIFAR-10에서 FID 점수 3.78, ImageNet-64에서 7.83의 최신 기술 수준 성능을 달성한다.

ABSTRACT

Diffusion models have found widespread adoption in various areas. However, their sampling process is slow because it requires hundreds to thousands of network evaluations to emulate a continuous process defined by differential equations. In this work, we use neural operators, an efficient method to solve the probability flow differential equations, to accelerate the sampling process of diffusion models. Compared to other fast sampling methods that have a sequential nature, we are the first to propose a parallel decoding method that generates images with only one model forward pass. We propose diffusion model sampling with neural operator (DSNO) that maps the initial condition, i.e., Gaussian distribution, to the continuous-time solution trajectory of the reverse diffusion process. To model the temporal correlations along the trajectory, we introduce temporal convolution layers that are parameterized in the Fourier space into the given diffusion model backbone. We show our method achieves state-of-the-art FID of 3.78 for CIFAR-10 and 7.83 for ImageNet-64 in the one-model-evaluation setting.

연구 동기 및 목표

• 시간 민감한 응용 프로그램을 위한 실시간 추론을 가능하게 하기 위해 확산 모델 샘플링을 가속화한다.
• 기존의 빠른 샘플링 방법이 다수의 forward pass 가 필요로 하는 순차적 처리의 한계를 극복한다.
• 노이즈에서 데이터로의 연속 시간 솔루션 궤적을 매핑하는 미분 가능하고 연속적인 해법을 개발한다.
• 추론 시간을 단일 forward pass로 극적으로 단축하면서도 높은 샘플 품질을 유지한다.
• 파rameter 효율적인 푸리에 공간 시간 컨볼루션을 사용하여 시간 모델링을 확산 모델에 통합한다.

제안 방법

• 초기 가우시안 노이즈에서 연속 시간 역확산 과정의 솔루션 궤적을 매핑하는 신경 연산자인 DSNO를 제안한다.
• 궤적을 따라 시간적 상관관계를 모델링하기 위해 푸리에 도메인에서 파arameter화된 시간 컨볼루션 레이어를 도입한다.
• 푸리에 기반 연산을 통해 모든 시간 스텝의 출력을 동시에 계산함으로써 순차적 계산을 피하고 병렬 디코딩을 가능하게 한다.
• 기존의 U-Net 백본에 시간 컨볼루션 블록을 통합하여 표준 확산 모델 아키텍처와의 호환성을 유지한다.
• 샘플의 정밀도를 향상시키기 위해 L1 및 시각적( perceptual, LPIPS) 손실의 조합을 사용해 모델을 훈련시킨다.
• 일반화 및 FID 성능 향상을 위해 2차 시간 이산화와 다중 해상도 감시를 사용한다.

실험 결과

연구 질문

• RQ1 신경 연산자를 사용해 확산 모델의 연속 시간 역과정을 효과적으로 모델링하여 빠른 샘플링을 가능하게 할 수 있는가?
• RQ2 푸리에 도메인에서의 시간 컨볼루션 레이어가 확산 궤적의 장거리 시간 의존성을 효율적으로 포착할 수 있는가?
• RQ3 푸리에 기반 연산을 통한 병렬 디코딩이 샘플 품질을 훼손하지 않고 단일 스텝 생성을 가능하게 하는가?
• RQ4 제안된 방법이 단일 forward pass로 최신 기술 수준의 FID 점수를 달성할 수 있는가?
• RQ5 손실 함수의 선택(예: L1 대비 LPIPS)이 최종 샘플 품질과 FID 점수에 어떤 영향을 미치는가?

주요 결과

• DSNO는 한 번의 모델 평가 설정에서 CIFAR-10에서 최신 기술 수준의 FID 점수 3.78을 달성한다.
• DSNO는 단일 모델 평가로 ImageNet-64에서 최신 기술 수준의 FID 점수 7.83을 달성한다.
• VGG 기반 LPIPS 손실의 사용은 표준 L1 손실 대비 FID 점수 향상에 유의미한 기여를 한다.
• 시간 해상도를 4에서 8로 증가시키면 FID가 약간 향상되지만, 이로 인한 이득의 감소율을 고려하면 4가 효율성 측면에서 충분하다.
• 제안된 시간 컨볼루션 블록은 모델 크기를 10%만 증가시키면서도 병렬 디코딩과 높은 성능을 가능하게 한다.
• 이 방법은 진정으로 병렬 디코딩을 가능하게 하여, 순차적 솔버와 달리 모든 시간 스텝의 출력을 단일 forward pass 내에서 생성한다.