[논문 리뷰] SA-Solver: Stochastic Adams Solver for Fast Sampling of Diffusion Models

Shuchen Xue, Mingyang Yi외 5명|arXiv (Cornell University)|2023. 09. 10.

Bayesian Methods and Mixture Models인용 수 6

한 줄 요약

이 논문은 확산 확률 모델에서 빠르고 고품질의 샘플링을 위한 새로운 스토하스틱 아담스 방법인 SA-Solver를 제안한다. 변동성을 통제하는 확산 SDE의 가족을 도입하고 효율적인 선형 다단계 SDE 해법기를 활용하여, 기존의 결정론적 및 스토하스틱 샘플러들과 비교해도 더 적은 기능 평가 횟수(NFEs)로도 다양한 벤치마크에서 최고 수준의 FID 점수를 달성한다. 특히 낮은 NFE 설정과 부분적으로 훈련된 모델에서 뛰어난 성능을 보이며 슈퍼리아이어스를 초월한다.

ABSTRACT

Diffusion Probabilistic Models (DPMs) have achieved considerable success in generation tasks. As sampling from DPMs is equivalent to solving diffusion SDE or ODE which is time-consuming, numerous fast sampling methods built upon improved differential equation solvers are proposed. The majority of such techniques consider solving the diffusion ODE due to its superior efficiency. However, stochastic sampling could offer additional advantages in generating diverse and high-quality data. In this work, we engage in a comprehensive analysis of stochastic sampling from two aspects: variance-controlled diffusion SDE and linear multi-step SDE solver. Based on our analysis, we propose extit{SA-Solver}, which is an improved efficient stochastic Adams method for solving diffusion SDE to generate data with high quality. Our experiments show that extit{SA-Solver} achieves: 1) improved or comparable performance compared with the existing state-of-the-art (SOTA) sampling methods for few-step sampling; 2) SOTA FID on substantial benchmark datasets under a suitable number of function evaluations (NFEs). Code is available at https://github.com/scxue/SA-Solver.

연구 동기 및 목표

• 확산 모델을 위한 기존의 훈련 없음 샘플링 방법의 비효율성과 품질 제한 문제를 해결하기 위해.
• 결정론적 해법기와 비교해 더 높은 품질과 다양성을 가진 데이터 생성 잠재력을 스토하스틱 샘플링이 지닌다.
• 조정 가능한 노이즈 스케일을 가진 새로운 노이즈 제어 확산 SDE의 가족을 설계하여 샘플링 효율성을 향상시키기 위해.
• 기존 ODE 해법기의 성능을 뛰어넘거나 이를 충족하는 효율적인 선형 다단계 SDE 해법기를 개발하기 위해.
• 추가 모델 훈련 없이도 낮은 NFE 설정에서 최고 수준의 생성 품질을 달성하기 위해.

제안 방법

• 원래의 것과 동일한 변동 분포를 가지지만 조절 가능한 노이즈 스케일(변동성)을 가진 확산 SDE의 가족을 제안하며, 조정 가능한 매개변수 τ를 통해 제어한다.
• 확산 과정에 특화된 지수 가중 적분을 통합하여 [29]에서 제안한 스토하스틱 아담스 방법을 변형하여 유도된 SDE를 해결한다.
• 제안된 스토하스틱 아담스 해법기의 이론적 수렴 차수 분석을 도입하여 수치적 안정성과 정확성을 보장한다.
• 효율성을 향상시키고 샘플링에 필요한 기능 평가 횟수(NFEs)를 줄이기 위해 선형 다단계 프레임워크를 활용한다.
• 다양성과 품질의 균형을 맞추기 위해 노이즈 스케일 τ를 校정하여 생성 과정에서 제어 가능한 트레이드오프를 가능하게 한다.
• 표준 DPMs와 동일한 역과정을 사용하지만, 제안된 SA-Solver를 통한 향상된 수치적 통합을 적용한다.

실험 결과

연구 질문

• RQ1 노이즈 분산을 통제하는 스토하스틱 샘플링이 결정론적 해법기와 비교해 확산 모델의 생성 품질을 향상시킬 수 있는가?
• RQ2 기존 ODE 해법기의 성능을 뛰어넘거나 이를 충족하는 효율적인 선형 다단계 SDE 해법기를 어떻게 설계할 수 있는가?
• RQ3 다양한 NFE 수준에서 노이즈 스케일(τ)과 샘플링 품질(FID로 측정) 사이의 최적 트레이드오프는 무엇인가?
• RQ4 확산 모델이 완전히 훈련되지 않은 상태에서도 제안된 방법이 뛰어난 성능을 유지하는가? 이는 모델 용량에 대한 강건성을 시사하는가?
• RQ5 훈련 없이도 스토하스틱 해법기를 사용해 최고 수준의 FID 점수를 달성할 수 있는가? 이는 최소한의 기능 평가 횟수로도 가능한가?

주요 결과

• SA-Solver는 모든 벤치마크 데이터셋(CIFAR-10, LSUN-Bedroom, ImageNet)에서 적절한 기능 평가 횟수(NFEs)로 최고 수준의 FID 점수를 기록하며, 결정론적 및 스토하스틱 베이스라인을 모두 압도한다.
• 극도로 낮은 NFE(예: 5–10) 설정에서도 SA-Solver는 DDIM, DPM-Solver, UniPC와 같은 기존 최고 수준의 방법들과 비교해 FID 점수를 유사하거나 뛰어나게 기록한다.
• LSUN-Bedroom 256x256 데이터셋에서 SA-Solver는 NFE=60일 때 최소 FID 3.06을 기록하며, EDM(SDE) 및 UniPC를 포함한 모든 베이스라인을 초월한다.
• 분류기 없는 가이던스를 사용한 ImageNet 512x512 DiT 모델에서, NFE=60일 때 SA-Solver는 DDIM 및 UniPC보다 더 다양하고 세밀한 샘플을 생성한다.
• 모델이 완전히 훈련되지 않은 상태에서도 SA-Solver는 결정론적 샘플러보다 뛰어난 성능을 유지하며, 모델 용량에 대한 강건성을 시사한다.
• 제거 실험 결과, 노이즈 스케일 τ를 증가시킬수록(최대 τ=1.0까지) 낮은 NFE에서 샘플 품질이 향상되며, 대부분의 벤치마크에서 τ=0.8–1.0에서 최적 성능을 기록한다.