[논문 리뷰] Flow Matching in Latent Space

Quan Dao, Hao Phung외 2명|arXiv (Cornell University)|2023. 07. 17.

Generative Adversarial Networks and Image Synthesis인용 수 5

한 줄 요약

이 논문은 사전 훈련된 오토인코더의 잠재공간에서 플로우 매칭을 적용함으로써 고해상도 이미지 생성의 계산 효율성과 확장성을 향상시키는 새로운 프레임워크인 잠재 플로우 매칭(LFM)을 제안한다. 압축된 잠재 표현에서 플로우 매칭을 훈련하고, 조건부 생성을 위해 분류기 자유 가이던스를 통합함으로써, 이미지 복원, 의미에서 이미지 생성, 클래스 조건부 생성 등의 작업에서 최신 기술 수준의 성능을 달성하며, CelebA-HQ 256에서 FID 점수 4.09와 26.3을 기록했고, 빠른 추론과 낮은 훈련 비용을 유지한다.

ABSTRACT

Flow matching is a recent framework to train generative models that exhibits impressive empirical performance while being relatively easier to train compared with diffusion-based models. Despite its advantageous properties, prior methods still face the challenges of expensive computing and a large number of function evaluations of off-the-shelf solvers in the pixel space. Furthermore, although latent-based generative methods have shown great success in recent years, this particular model type remains underexplored in this area. In this work, we propose to apply flow matching in the latent spaces of pretrained autoencoders, which offers improved computational efficiency and scalability for high-resolution image synthesis. This enables flow-matching training on constrained computational resources while maintaining their quality and flexibility. Additionally, our work stands as a pioneering contribution in the integration of various conditions into flow matching for conditional generation tasks, including label-conditioned image generation, image inpainting, and semantic-to-image generation. Through extensive experiments, our approach demonstrates its effectiveness in both quantitative and qualitative results on various datasets, such as CelebA-HQ, FFHQ, LSUN Church & Bedroom, and ImageNet. We also provide a theoretical control of the Wasserstein-2 distance between the reconstructed latent flow distribution and true data distribution, showing it is upper-bounded by the latent flow matching objective. Our code will be available at https://github.com/VinAIResearch/LFM.git.

연구 동기 및 목표

• 픽셀 공간 플로우 매칭 모델의 높은 계산 비용과 느린 추론 문제를 해결하기 위해.
• 사전 훈련된 오토인코더의 잠재공간으로 플로우 매칭을 이전하여 효율적이고 확장 가능한 고해상도 이미지 합성 구현을 위해.
• 클래스 조건부 이미지 생성, 복원, 의미에서 이미지 생성을 포함한 다양한 조건부 생성 작업으로 플로우 매칭을 확장하기 위해.
• 조건부 생성 품질 향상을 위해 플로우 매칭에 분류기 자유 가이던스를 도입하기 위해.
• 잠재공간에서 재구성된 분포와 진짜 데이터 분포 사이의 워셔스타인-2 거리에 대한 이론적 상한을 제공하기 위해.

제안 방법

• 사전 훈련된 오토인코더의 잠재공간에 플로우 매칭을 적용하여, 뱀머프 표현을 생성 모델링을 위한 잠재 다양체로 사용한다.
• 잠재 코드, 조건(예: 클래스 레이블, 마스크, 또는 이미지), 시간을 입력으로 받는 조건부 속도장 네트워크를 설계하여 플로우 역학을 예측한다.
• 조건부 및 무조건적 플로우 양측을 훈련하여 조건부 흐름을 유연하게 추론할 수 있도록, 속도장 네트워크에 분류기 자유 가이던스를 통합한다.
• 잠재공간 내에서 효율적이고 정확한 추론을 위해 적응형 오일러 적분기(예: 룬게-쿠타)를 사용한다.
• 플로우 매칭 목표함수를 사용하여, 잠재 플로우 매칭 네트워크와 조건부 임베딩 헤드(예: 레이아웃에서 이미지로)를 함께 훈련한다.
• 이론적 분석을 통해 재구성된 분포와 진짜 데이터 분포 사이의 워셔스타인-2 거리는 잠재 플로우 매칭 목표함수에 의해 상한이 있음을 확인한다.

실험 결과

연구 질문

• RQ1 잠재공간에서의 플로우 매칭이 고해상도 이미지 품질을 유지하면서 계산 비용을 크게 줄일 수 있는가?
• RQ2 잠재 플로우 매칭은 복원 및 의미에서 이미지 생성과 같은 다양한 조건부 이미지 생성 작업에 얼마나 잘 일반화되는가?
• RQ3 분류기 자유 가이던스는 플로우 매칭에 효과적으로 적용되어 조건부 생성 성능을 향상시킬 수 있는가?
• RQ4 잠재 플로우 매칭 목표함수와 진짜 데이터 분포 사이의 워셔스타인-2 거리 간의 이론적 관계는 어떠한가?
• RQ5 오토인코더 아키텍처의 선택이 잠재 플로우 매칭의 성능에 어떤 영향을 미치는가?

주요 결과

• LFM는 CelebA-HQ 256 복원 벤치마크에서 FID 점수 4.09를 기록하여 대부분의 기존 방법을 능가하며, MAT의 SOTA 점수 2.94에 가까워졌다.
• 의미에서 이미지 생성 작업에서 LFM는 CelebA-HQ 256에서 FID 점수 26.3을 기록하여 단순한 아키텍처와 특수 설계 없이도 뛰어난 성능을 보였다.
• 적응형 오일러 적분기를 사용해 단 50단계로도 빠른 추론이 가능하여 경쟁적인 추론 속도를 달성했다.
• LFM에서의 분류기 자유 가이던스는 동일한 설정에서 FID를 16.71에서 8.58로 크게 향상시켜 성능 향상을 이끌었다.
• 이론적 분석을 통해 재구성된 분포와 진짜 데이터 분포 사이의 워셔스타인-2 거리는 잠재 플로우 매칭 목표함수에 의해 상한이 있음을 확인하여, 방법의 통계적 일致성을 검증했다.
• LDM보다 더 작은 모델을 사용했음에도 불구하고, 분류기 자유 가이던스를 통한 성능 향상이 뚜렷했으며, 이는 더 큰 아키텍처로의 확장 잠재력이 있음을 시사한다.