[논문 리뷰] Parallel Sampling of Diffusion Models

Andy Y. Shih, Suneel Belkhale외 3명|arXiv (Cornell University)|2023. 05. 25.

Advanced Neuroimaging Techniques and Applications인용 수 9

한 줄 요약

ParaDiGMS는 Picard 반복을 통해 역방향 확산 샘플링을 병렬화하여 확산 모델 샘플링을 가속하며, 로봇 공정 및 이미지 생성 작업 전반에서 품질 저하 없이 2–4x 속도 향상을 달성한다.

ABSTRACT

Diffusion models are powerful generative models but suffer from slow sampling, often taking 1000 sequential denoising steps for one sample. As a result, considerable efforts have been directed toward reducing the number of denoising steps, but these methods hurt sample quality. Instead of reducing the number of denoising steps (trading quality for speed), in this paper we explore an orthogonal approach: can we run the denoising steps in parallel (trading compute for speed)? In spite of the sequential nature of the denoising steps, we show that surprisingly it is possible to parallelize sampling via Picard iterations, by guessing the solution of future denoising steps and iteratively refining until convergence. With this insight, we present ParaDiGMS, a novel method to accelerate the sampling of pretrained diffusion models by denoising multiple steps in parallel. ParaDiGMS is the first diffusion sampling method that enables trading compute for speed and is even compatible with existing fast sampling techniques such as DDIM and DPMSolver. Using ParaDiGMS, we improve sampling speed by 2-4x across a range of robotics and image generation models, giving state-of-the-art sampling speeds of 0.2s on 100-step DiffusionPolicy and 14.6s on 1000-step StableDiffusion-v2 with no measurable degradation of task reward, FID score, or CLIP score.

연구 동기 및 목표

• 더 빠른 샘플링으로 샘플 품질을 저하시키지 않으면서 계산을 속도에 맞춰 교환하려는 동기 부여.
• 사전 학습된 확산 모델 및 기존의 빠른 샘플러와 호환되는 일반적인 병렬 샘플링 프레임워크를 개발.
• 로봇 정책과 이미지 생성 모델 전반에 걸쳐 속도 향상을 시연하면서 보상, FID, CLIP와 같은 지표를 보존.
• 이전 빠른 샘샘 기법들(DDIM, DPMSolver)과 분류기-무 Guidance와의 호환성 및 일반화 가능성 확인.

제안 방법

• 역시간 확산 미분방정식을 Picard 업데이트를 사용하여 반대로 풀고 전체 디노이징 궤적의 추측을 반복적으로 정제한다.
• Picard 업데이트를 이산화하여 시간 단계별로 drifts s(x_i, t_i)의 병렬 계산과 그 뒤의 프리픽스 합 누적을 가능하게 한다.
• GPU 메모리에 맞추기 위한 슬라이딩 윈도우 미니배치 전략과 총변이(distance) 보장을 보장하는 중단 기준을 도입한다.
• 역 SDE로 확장할 때 고정된 순방향 노이즈 실현을 결정론적 ODE로 흡수한다.
• ParaDiGMS 내에서 Euler 또는 고차 솔버를 대체하고 ParaDDPM, ParaDDIM, ParaDPMSolver와의 통합으로 호환성을 입증한다.
• 분포 간 거리의 공차를 순차적 기준선과 비교하여 허용 오차 기반 중단 기준을 사용한다.

실험 결과

연구 질문

• RQ1 확산 모델 샘플링을 디노이징 단계를 병렬화하여 샘플 품질을 희생하지 않고도 가속화할 수 있는가?
• RQ2 Picard 반복 기반 병렬 샘플링과 기존의 빠른 샘플러(DDIM, DPMSolver)의 속도와 품질을 비교하면 어떠한가?
• RQ3 접근 방식이 분류기-무 가이드 및 선행 속도 향상과 호환되는가, 그리고 이미지 및 로봇 정책 작업 모두에 일반화될 수 있는가?

주요 결과

• ParaDiGMS는 로봇 및 이미지 생성 모델 전반에서 약 2–4x의 속도 향상을 달성하면서 작업 보상, FID, CLIP 점수의 측정 가능한 저하 없이 작동한다.
• DiffusionPolicy에서 100단 샘플링 시 샘플당 소요 시간이 0.74초에서 0.2초로 감소(예시: 100단 액션 생성 작업).
• StableDiffusion-v2에서 1000단 샘플링 시 샘플당 시간이 50.0초에서 14.6초로 감소하되 품질 저하가 없다.
• ParaDiGMS는 DDPM, DDIM, DPMSolver와 호환되어 작업 전반에서 결합 속도 향상을 가능하게 한다(ParaDDPM, ParaDDIM, ParaDPMSolver).
• StableDiffusion-v2의 잠재 공간 확산에서 멀티 GPU 구성 시 순수 속도 향상이 3배를 넘게 나타나며 배치 창 크기와 하드웨어에 따라 속도 향상이 달라진다.
• ParaDiGMS는 COCO 캡션에 대한 CLIP 점수로 측정한 샘플 품질을 여러 구성에서 유지한다.