[논문 리뷰] Accelerating Parallel Sampling of Diffusion Models

Zhiwei Tang, Jiasheng Tang외 3명|arXiv (Cornell University)|2024. 02. 15.

Advanced Mathematical Modeling in Engineering인용 수 4

한 줄 요약

이 논문은 디퓨전 모델을 위한 훈련 불필요한, 보편적인 병렬 샘플링 알고리즘인 ParaTAA를 제안한다. 이 알고리즘은 자동재귀적 샘플링 과정을 고정점 반복을 통해 풀이하는 삼각형 비선형 방정식으로 재구성한다. 삼각형 앤더슨 가속화와 적응형 윈도우링을 활용하여 ParaTAA는 추론 단계를 4배에서 14배 감소시켜 순차적 샘플링과 거의 동일한 이미지 품질을 달성한다. 예를 들어, 스테이블 디퓨전에서 100단계 DDIM 결과를 단 7단계만으로 재현한다.

ABSTRACT

Diffusion models have emerged as state-of-the-art generative models for image generation. However, sampling from diffusion models is usually time-consuming due to the inherent autoregressive nature of their sampling process. In this work, we propose a novel approach that accelerates the sampling of diffusion models by parallelizing the autoregressive process. Specifically, we reformulate the sampling process as solving a system of triangular nonlinear equations through fixed-point iteration. With this innovative formulation, we explore several systematic techniques to further reduce the iteration steps required by the solving process. Applying these techniques, we introduce ParaTAA, a universal and training-free parallel sampling algorithm that can leverage extra computational and memory resources to increase the sampling speed. Our experiments demonstrate that ParaTAA can decrease the inference steps required by common sequential sampling algorithms such as DDIM and DDPM by a factor of 4$\sim$14 times. Notably, when applying ParaTAA with 100 steps DDIM for Stable Diffusion, a widely-used text-to-image diffusion model, it can produce the same images as the sequential sampling in only 7 inference steps. The code is available at https://github.com/TZW1998/ParaTAA-Diffusion.

연구 동기 및 목표

• 높은 품질의 생성에도 불구하고 느린 추론 속도로 인해 실시간 배포가 제한되는 자동재귀적 디퓨전 샘플링의 문제를 해결한다.
• 모델의 미세조정이나 distillation 없이도 본질적으로 순차적인 샘플링 과정을 병렬화할 수 있는 방법을 개발한다.
• 순차적 방법(예: DDIM, DDPM)과 비교해 추론 단계를 극적으로 줄이면서도 이미지 품질과 다양성을 유지한다.
• 모델의 무게를 변경하지 않고도 추가적인 계산 및 메모리 자원을 활용해 추론 속도를 높일 수 있도록 한다.
• 디퓨전 모델에서 병렬 샘플링의 수렴 분석 및 최적화를 위한 이론적 기반 프레임워크를 제공한다.

제안 방법

• 고정점 반복(FP)을 사용해 디퓨전 샘플링 과정을 삼각형 비선형 방정식의 체계로 재구성함으로써 중간 단계의 병렬 계산을 가능하게 한다.
• 과거 반복값을 삼각형 구조로 활용하는 삼각형 앤더슨 가속화(TAA)를 도입해 고정점 반복의 수렴 속도를 가속화한다.
• 수렴 속도와 계산 비용 간의 트레이드오프를 제어하기 위해 적응형 윈도우링을 적용하여 사용할 과거 반복값의 수를 동적으로 조절한다.
• 유사한 프롬프트에서 이전에 수행한 샘플링 경로를 이용해 수렴을 부트스트랩함으로써 유사한 프롬프트에 대해 단계 수를 크게 줄인다.
• 모델 재학습 없이도 기존 순차적 샘플러(예: DDIM, DDPM)와 통합하기 위해 비선형 방정식의 계수 구조를 수정한다.
• 모델 가중치를 변경하지 않고도 16비트 정밀도와 효율적인 GPU 메모리 액세스를 최적화해 벽시계 시간을 감소시킨다.

(a) DDIM 100 steps

실험 결과

연구 질문

• RQ1 디퓨전 모델의 자동재귀적 샘플링 과정이 이미지 품질을 유지하면서 효과적으로 병렬화될 수 있는가?
• RQ2 고정점 반복이 디퓨전 샘플링을 위한 비선형 방정식 체계를 푸는 데 있어 수렴 속도를 체계적으로 최적화할 수 있는가?
• RQ3 윈도우 크기와 초기화 전략이 병렬 샘플링에서 수렴 속도와 추론 단계 감소에 미치는 영향은 어떠한가?
• RQ4 모델의 미세조정 없이도 추가적인 계산 자원을 얼마나 효과적으로 활용해 추론 단계를 줄일 수 있는가?
• RQ5 제안된 방법이 순차적 샘플링과 거의 동일한 결과를 내면서도 수십 배에서 수백 배의 속도 향상을 달성할 수 있는가?

주요 결과

• ParaTAA는 DDIM 및 DDPM과 같은 순차적 샘플링 방법이 요구하는 추론 단계를 4배에서 14배로 감소시킨다.
• 스테이블 디퓨전에서 100단계 DDIM을 사용할 경우, ParaTAA는 단 7단계의 추론으로 동일한 이미지 품질을 달성하여 14배의 속도 향상을 보였다.
• 윈도우 크기를 10으로 설정했을 때, ParaTAA는 스테이블 디퓨전에서 DDIM의 추론 단계를 100에서 25로 줄여 4배의 속도 향상을 달성했다.
• 윈도우 크기를 10에서 20으로 두 배로 늘였을 때도 단계 수는 25에서 21로만 감소하여 특정 윈도우 크기 이상에서는 수익 감소 현상이 나타남을 보여주었다.
• 유사한 프롬프트에서 이전에 수행한 경로를 초기화로 사용할 경우, 수렴 시간이 최대 3배까지 단축되었으며, 3단계 내로도 고품질 이미지를 생성할 수 있었다.
• 이 방법은 순차적 샘플링과 유사한 이미지 무결성과 다양성을 유지하며, 수렴성과 해의 유일성에 대한 이론적 보장도 제공한다.

(b) DDPM 100 steps