[논문 리뷰] OMS-DPM: Optimizing the Model Schedule for Diffusion Probabilistic Models

Enshu Liu, Xuefei Ning외 3명|arXiv (Cornell University)|2023. 06. 15.

Advanced Neural Network Applications인용 수 6

한 줄 요약

OMS-DPM은 노이즈 제거 단계 동안 다양한 사전 훈련된 모델을 동적으로 선택하여 확산 확률 모델의 모델 스케줄을 최적화하는 새로운 방법을 제안한다. 이로 인해 생성 품질과 속도가 크게 향상된다. 예측기 기반의 진화적 탐색을 통해 스탠리 디퓨전보다 2배 빠른 샘플링을 달성하면서도, CIFAR-10, CelebA, ImageNet 및 LSUN-Church에서 최신 기준 FID 점수를 유지한다.

ABSTRACT

Diffusion probabilistic models (DPMs) are a new class of generative models that have achieved state-of-the-art generation quality in various domains. Despite the promise, one major drawback of DPMs is the slow generation speed due to the large number of neural network evaluations required in the generation process. In this paper, we reveal an overlooked dimension -- model schedule -- for optimizing the trade-off between generation quality and speed. More specifically, we observe that small models, though having worse generation quality when used alone, could outperform large models in certain generation steps. Therefore, unlike the traditional way of using a single model, using different models in different generation steps in a carefully designed \emph{model schedule} could potentially improve generation quality and speed \emph{simultaneously}. We design OMS-DPM, a predictor-based search algorithm, to optimize the model schedule given an arbitrary generation time budget and a set of pre-trained models. We demonstrate that OMS-DPM can find model schedules that improve generation quality and speed than prior state-of-the-art methods across CIFAR-10, CelebA, ImageNet, and LSUN datasets. When applied to the public checkpoints of the Stable Diffusion model, we are able to accelerate the sampling by 2$ imes$ while maintaining the generation quality.

연구 동기 및 목표

• 확산 확률 모델(DPMs)의 샘플링 속도가 느려서 생성당 수백~수천 번의 신경망 평가가 필요로 하는 문제를 해결한다.
• 생성 품질과 속도를 향상시키기 위해 다양한 노이즈 제거 단계에서 다른 모델을 사용하는 것—모델 스케줄이라는 간과된 최적화 차원을 식별하고 활용한다.
• 재훈련 없이도 사전 훈련된 모델만을 사용하여 임의의 시간 예산 하에서 최적의 모델 스케줄을 찾는 방법을 개발한다.
• Stable Diffusion와 같은 공개 체크포인트를 포함하여 다양한 데이터셋과 사전 훈련된 모델 간에 효율적이고 일반화 가능한 스케줄링을 가능하게 한다.
• 노이즈 제거 단계에서의 모델 이질성(다양성)을 활용하여 기존 방법보다 더 나은 품질-속도 트레이드오프를 달성한다.

제안 방법

• 작은 양의 훈련 데이터를 사용하여 모델 스케줄에서 생성 품질을 예측할 수 있는 예측기 기반의 탐색 알고리즘인 OMS-DPM을 제안한다.
• 작은 모델 스케줄 서브셋에 대해 성능 예측기를 훈련시어, 미리 보지 않은 스케줄로의 일반화를 가능하게 하여 전체 평가의 비용을 줄인다.
• 시간 예산 내에서 지수적으로 증가하는 모델 스케줄 탐색 공간을 효율적으로 탐색하기 위해 진화 알고리즘을 사용한다.
• DDIM 단계 건너뛰기 및 DPM-Solver 순서 선택과 같은 기존 가속 기법을 스케줄 탐색의 하이퍼파rameter로 간주함으로써 탄력적인 통합을 지원한다.
• 예측 성능에 기반해 각 노이즈 제거 단계에 크기와 속도가 다른 사전 훈련된 모델을 다르게 할당하도록 설계한다.
• 각 데이터셋당 소수의 데이터 포인트로 예측기를 훈련시켜, 최소한의 재훈련으로 새로운 데이터셋에 대해 재사용 가능한 이식성을 확보한다.

실험 결과

연구 질문

• RQ1 다양한 노이즈 제거 단계에 다른 사전 훈련된 모델을 동적으로 할당하면 DPM에서 생성 품질과 샘플링 속도를 동시에 향상시킬 수 있는가?
• RQ2 모델 스케줄 최적화는 오직 노이즈 스케줄, 이산화 또는 솔버 설계에 집중하는 전통적 가속 기법과 비교해 어떻게 다를까?
• RQ3 제한된 시간 예산 하에서 모델 크기와 배치(초기 vs. 후기 단계)는 최종 생성 품질에 어떤 영향을 미치는가?
• RQ4 경량 예측기가 다양한 모델 스케줄과 데이터셋으로 일반화되어, 전체 평가 없이도 효율적인 탐색을 가능하게 할 수 있는가?
• RQ5 최적의 모델 스케줄의 구조는 다양한 데이터셋과 샘플링 예산에 따라 어떻게 달라지는가?

주요 결과

• OMS-DPM은 스탠리 디퓨전 모델의 샘플링 속도를 2배로 높였으며, FID 점수는 유지하거나 향상시켰다. 이는 기준 방법보다 뛰어난 성능을 보였다.
• CIFAR-10에서 1400ms 지연 예산 하에서 FID는 기준값 3.56에서 3.19로 감소하여 품질-속도 트레이드오프의 뚜렷한 향상을 보였다.
• 엄격한 시간 예산 하에서는 오직 가장 작고 빠른 모델들(예: 2~3개 모델)만 사용하는 것이 최적의 전략이었으며, 더 큰 모델은 단계가 너무 적을 경우 오차를 유발하기 때문이다.
• LSUN-Church에서는 최적 전략이 노이즈 제거 과정의 초기에 더 큰 모델을 사용하는 반면, ImageNet-64와 CIFAR-10에서는 더 큰 모델이 후기 단계에서 더 효과적이었으며, 이는 데이터셋에 따라 다른 스케줄링 패턴을 의미한다.
• 100개의 데이터 포인트만으로도 훈련된 예측기는 잘 일반화되어, 최소한의 평가 비용으로 효율적인 탐색을 가능하게 하였다.
• DPM-Solver의 경우, 엄격한 예산 하에서는 1차 및 2차 솔버가 선호되며, 충분한 시간이 확보된 경우에만 3차 솔버가 사용된다. 이는 최적 스케줄에서의 적응형 솔버 선택을 보여준다.