단계를 정렬하라: 확산 모델에서 샘플링 스케줄 최적화하기

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

April 22, 2024

저자: Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis

cs.AI

초록

디퓨전 모델(DMs)은 시각적 영역을 넘어 최첨단 생성 모델링 접근법으로 자리 잡았습니다. DMs의 주요 단점은 대규모 신경망을 통한 많은 순차적 함수 평가에 의존하는 느린 샘플링 속도입니다. DMs에서 샘플링은 샘플링 스케줄로 알려진 이산화된 노이즈 레벨 집합을 통해 미분 방정식을 푸는 것으로 볼 수 있습니다. 과거 연구들은 주로 효율적인 솔버를 도출하는 데 초점을 맞췄지만, 최적의 샘플링 스케줄을 찾는 데는 거의 관심을 기울이지 않았으며, 전체 문헌은 수작업 휴리스틱에 의존해 왔습니다. 본 연구에서는 처음으로 'Align Your Steps'라는 이름으로 고품질 출력을 위한 DMs의 샘플링 스케줄을 최적화하는 일반적이고 원칙적인 접근 방식을 제안합니다. 우리는 확률적 미적분학의 방법을 활용하여 다양한 솔버, 훈련된 DMs 및 데이터셋에 특화된 최적의 스케줄을 찾습니다. 우리는 이 새로운 접근 방식을 여러 이미지, 비디오 및 2D 토이 데이터 합성 벤치마크에서 다양한 샘플러를 사용하여 평가했으며, 최적화된 스케줄이 거의 모든 실험에서 이전의 수작업 스케줄을 능가하는 것을 관찰했습니다. 우리의 방법은 특히 적은 단계 합성 체제에서 샘플링 스케줄 최적화의 잠재력을 보여줍니다.

English

Diffusion models (DMs) have established themselves as the state-of-the-art generative modeling approach in the visual domain and beyond. A crucial drawback of DMs is their slow sampling speed, relying on many sequential function evaluations through large neural networks. Sampling from DMs can be seen as solving a differential equation through a discretized set of noise levels known as the sampling schedule. While past works primarily focused on deriving efficient solvers, little attention has been given to finding optimal sampling schedules, and the entire literature relies on hand-crafted heuristics. In this work, for the first time, we propose a general and principled approach to optimizing the sampling schedules of DMs for high-quality outputs, called Align Your Steps. We leverage methods from stochastic calculus and find optimal schedules specific to different solvers, trained DMs and datasets. We evaluate our novel approach on several image, video as well as 2D toy data synthesis benchmarks, using a variety of different samplers, and observe that our optimized schedules outperform previous hand-crafted schedules in almost all experiments. Our method demonstrates the untapped potential of sampling schedule optimization, especially in the few-step synthesis regime.