[논문 리뷰] Matryoshka Diffusion Models

Jiatao Gu, Shuangfei Zhai외 3명|arXiv (Cornell University)|2023. 10. 23.

Generative Adversarial Networks and Image Synthesis인용 수 10

한 줄 요약

MDM은 NestedUNet를 사용하여 다중 해상도 입력을 함께 노이즈 제거하는 엔드 투 엔드 고해상도 확산 프레임워크이며, 점진적 학습을 통해 최대 1024^2 해상도까지 도달합니다.

ABSTRACT

Diffusion models are the de facto approach for generating high-quality images and videos, but learning high-dimensional models remains a formidable task due to computational and optimization challenges. Existing methods often resort to training cascaded models in pixel space or using a downsampled latent space of a separately trained auto-encoder. In this paper, we introduce Matryoshka Diffusion Models(MDM), an end-to-end framework for high-resolution image and video synthesis. We propose a diffusion process that denoises inputs at multiple resolutions jointly and uses a NestedUNet architecture where features and parameters for small-scale inputs are nested within those of large scales. In addition, MDM enables a progressive training schedule from lower to higher resolutions, which leads to significant improvements in optimization for high-resolution generation. We demonstrate the effectiveness of our approach on various benchmarks, including class-conditioned image generation, high-resolution text-to-image, and text-to-video applications. Remarkably, we can train a single pixel-space model at resolutions of up to 1024x1024 pixels, demonstrating strong zero-shot generalization using the CC12M dataset, which contains only 12 million images. Our code is released at https://github.com/apple/ml-mdm

연구 동기 및 목표

• 연쇄적 또는 잠재 확산 파이프라인 없이 고해상도 이미지/비디오 생성을 촉진한다.
• 확장 공간에서 다중 해상도 확산 프로세스를 도입한다.
• 해상도 간 계산을 공유하기 위한 NestedUNet 아키텍처를 제안한다.
• 점진적으로 더 높은 해상도를 도입하기 위한 진행적 학습 스케줄을 개발한다.
• 클래스 조건화 이미지 생성, 텍스트-투-이미지 및 텍스트-투-비디오 작업에서 효과를 입증한다.

제안 방법

• 표준 확산을 다운샘플링 연산 D^r(.)로 연결된 여러 해상도 z^r_t를 가진 확장 공간으로 확장한다.
• 단일 노이즈 제거 함수에서 해상도 간 매개변수와 계산을 공유하기 위해 NestedUNet를 도입한다.
• 해상도 간 손실에 가중치를 두는 다중 해상도 노이즈 제거 목표(Eq. 3)로 학습한다(ω_t^r).
• 낮은 해상도에서 시작하여 점진적으로 더 높은 해상도를 추가하는 진행형 학습(multi-phase training)을 도입한다.
• 최종 해상도가 서로 다른 샘플들을 동시 학습하기 위한 혼합 해상도 학습을 가능하게 한다.
• 테스트 시 모든 해상도가 병렬로 생성되는 병렬 추론을 제공한다.

실험 결과

연구 질문

• RQ1 확장 공간에서의 다중 해상도 확산이 고해상도 생성의 수렴 속도와 최종 품질에 어떤 영향을 미치는가?
• RQ2 엔드-투-엔드 고해상도 확산이 연쇄적 및 잠재 확산 방법과 비슷하거나 이를 능가하면서 학습/추론 파이프라인을 단순화할 수 있는가?
• RQ3 이미지와 비디오 생성에서 진행형 학습과 중첩 수준 깊이가 효율성과 품질에 미치는 영향은 무엇인가?
• RQ4 다중 해상도 확산 모델이 비교적 적은 데이터로 텍스트-투-이미지 및 텍스트-투-비디오 작업에 잘 일반화되는가?

주요 결과

모델	FID (ImageNet 256x256)	FID (MS-COCO 256x256)	비고
ADM (Nichol & Dhariwal, 2021)	10.94	-	ImageNet 256x256 (FID)
CDM (Ho et al., 2022b)	4.88	-	ImageNet 256x256 (FID)
LDM-4 (Rombach et al., 2022)	10.56	-	ImageNet 256x256 (FID)
LDM-4* (Rombach et al., 2022)	3.60	-	ImageNet 256x256 (FID)
Ours (cfg=1)	8.92	-	MDM, ImageNet 256x256 (FID)
Ours (cfg=1.2)*	6.62	-	MDM, ImageNet 256x256 (CFG)
LDM-8 (Rombach et al., 2022)	-	23.31	MS-COCO 256x256 (FID)
LDM-8* (Rombach et al., 2022)	-	12.63	MS-COCO 256x256 (FID)
Dalle-2* (Ramesh et al., 2022)	-	10.39	MS-COCO 256x256 (FID)
IMAGEN* (Saharia et al., 2021)	-	7.27	MS-COCO 256x256 (FID)
Ours (cfg=1)	-	18.35	MDM, CC12M 256x256 (FID)
Ours (cfg=1.35)*	-	13.43	MDM, CC12M 256x256 (CFG)

• MDM은 연쇄적 또는 잠재 확산 없이도 고해상도(최대 1024^2)에서 더 빠른 수렴과 높은 품질을 달성한다.
• 다중 해상도 손실과 진행형 학습이 학습 효율과 최종 결과를 크게 향상시킨다.
• 동등한 설정에서 더 많은 매개변수와 추론 스텝에도 불구하고 Cascaded DM은 MDM보다 성능이 낮다.
• MDM은 CC12M(12M 이미지)를 사용한 텍스트-투-이미지에서 1024^2 해상도에 대해 강력한 제로샷 능력을 보여준다.
• 아블레이션은 다중 해상도 손실과 진행형 학습이 모두 유익하며, 중첩 레벨을 높이면 수렴이 개선되되 비용은 최소화된다.
• MDM은 ImageNet 및 MS-COCO FID 지표에서 여러 베이스라인과 경쟁적이거나 우수하고, WebVid-10M에서 고해상도 텍스트-투-비디오 생성을 가능하게 한다.