[논문 리뷰] Patch Diffusion: Faster and More Data-Efficient Training of Diffusion Models

Zhendong Wang, Yifan Jiang외 6명|arXiv (Cornell University)|2023. 04. 25.

Advanced Neuroimaging Techniques and Applications인용 수 32

한 줄 요약

패치 디퓨전은 좌표 조건 입력과 임의의 패치 크기 스케줄링을 갖춘 패치-단위 조건부 스코어 매칭을 도입하여 확산 모델 학습을 가속하고 데이터 효율성을 향상시키며, 샘플 품질을 유지하거나 개선하는 한편 최소 2x 빠른 학습을 달성한다.

ABSTRACT

Diffusion models are powerful, but they require a lot of time and data to train. We propose Patch Diffusion, a generic patch-wise training framework, to significantly reduce the training time costs while improving data efficiency, which thus helps democratize diffusion model training to broader users. At the core of our innovations is a new conditional score function at the patch level, where the patch location in the original image is included as additional coordinate channels, while the patch size is randomized and diversified throughout training to encode the cross-region dependency at multiple scales. Sampling with our method is as easy as in the original diffusion model. Through Patch Diffusion, we could achieve $\mathbf{\ge 2 imes}$ faster training, while maintaining comparable or better generation quality. Patch Diffusion meanwhile improves the performance of diffusion models trained on relatively small datasets, $e.g.$, as few as 5,000 images to train from scratch. We achieve outstanding FID scores in line with state-of-the-art benchmarks: 1.77 on CelebA-64$ imes$64, 1.93 on AFHQv2-Wild-64$ imes$64, and 2.72 on ImageNet-256$ imes$256. We share our code and pre-trained models at https://github.com/Zhendong-Wang/Patch-Diffusion.

연구 동기 및 목표

• 확산 모델 학습의 시간 및 데이터 요구를 줄여 민주화하는 것을 목표로 한다.
• UNet 백본 및 샘플링 방법에 독립적인 패치 단위 학습 프레임워크를 제안한다.
• 교차 영역 의존성을 포착하기 위해 패치 수준의 좌표 조건화 및 패치 크기 스케줄링을 개발한다.
• 패치 디퓨전이 학습 시간을 절반으로 단축하면서도 생성 품질을 유지하거나 향상시키는지, 특히 작은 데이터세트에서 확인한다.

제안 방법

• 무작위 이미지 패치에서 denoising score matching으로 학습된 패치-와이즈 디노이저 D_theta(x; sigma_t)를 정의한다.
• 패치 위치 i, j와 패치 크기 s에 조건화하여 작은 패치와 큰 패치의 혼합에서 샘플링된 패치 크기를 사용한다.
• 입력에 두 개의 좌표 채널을 추가하여 패치 위치를 인코딩하는 픽셀 수준 좌표 시스템을 도입한다.
• 학습 중 패치 크기를 섞기 위한 패치 크기 스케줄링(확률적 또는 점진적)을 사용하고, 제어 가능한 비율 p로 전체 크기 이미지도 포함한다.
• 각 역단계에서 전체 이미지 좌표를 제공하여 표준 확산 모델의 샘플링 방식은 유지한다.
• 평가 시 CFG를 옵션으로 두고 EDM-DDPM++ 및 Latent Diffusion 백본 위에서 Patch Diffusion을 시연한다.

실험 결과

연구 질문

• RQ1 좌표 및 크기 조건화가 적용된 패치-와이즈 스코어 매칭이 확산 모델의 데이터 효율성과 학습 시간을 개선하는가?
• RQ2 패치 크기 스케줄링(확률적 vs 점진적)이 생성 품질과 학습 비용에 어떤 영향을 미치는가?
• RQ3 패치 기반 학습을 사용하면서 explicit 패치 병합 없이 표준 확산 모델처럼 샘플링이 가능할까?
• RQ4 패치 디퓨전이 작은 데이터세트와 사전 학습된 확산 모델의 미세 조정에 유리한가?

주요 결과

• 패치 디퓨전은 벤치마크 데이터셋에서 생성 품질을 유지하거나 향상시키면서 최소 2x 빠른 학습 속도를 달성한다.
• 좌표 조건화와 소형/대형 패치의 혼합 사용이 지역 간 일관성과 FID 점수를 개선한다.
• 학습 중 전체 크기 이미지 비율 p가 약 0.5인 지점에서 효율성과 품질의 균형이 잘 맞는다.
• 패치 디퓨전은 데이터 효율성을 높여 비교적 작은 데이터셋(약 5k 이미지 수준)에서도 경쟁력 있는 결과를 가능하게 한다.
• 대형 사전 학습 모델(예: ControlNet)을 Patch Diffusion으로 미세 조정하면 학습 효율이 두 배로 늘어나도 성능 저하가 없다.
• 제한된 크기의 데이터세트에서 Patch Diffusion은 FID에서 기준선보다 일관되게 우수하며 약 2x의 학습 시간 감소를 달성한다.