[논문 리뷰] Generative Modeling with Optimal Transport Maps

Litu Rout, Alexander Korotin외 1명|arXiv (Cornell University)|2021. 10. 06.

Generative Adversarial Networks and Image Synthesis참고 문헌 56인용 수 5

한 줄 요약

이 논문은 고차원의 환경 공간(예: 이미지)에서 직접 최적 운반(OT) 맵을 사용하는 새로운 엔드 투 엔드 방법을 제안한다. 이는 잠재 공간 근사치를 회피하며, 워샤르슈타인-2 거리에 대한 최소-최대 최적화를 설정하고 오차 한계를 증명함으로써, 이미지 생성 및 양방향이 아닌 이미지 복원(예: 노이즈 제거, 색상화, 메쉬 복구)에서 최신 기술 수준(SOTA) 성능을 달성한다. 이는 이전의 오토에코더 기반 잠재 공간을 사용하는 OT 기반 모델보다 뛰어나다.

ABSTRACT

With the discovery of Wasserstein GANs, Optimal Transport (OT) has become a powerful tool for large-scale generative modeling tasks. In these tasks, OT cost is typically used as the loss for training GANs. In contrast to this approach, we show that the OT map itself can be used as a generative model, providing comparable performance. Previous analogous approaches consider OT maps as generative models only in the latent spaces due to their poor performance in the original high-dimensional ambient space. In contrast, we apply OT maps directly in the ambient space, e.g., a space of high-dimensional images. First, we derive a min-max optimization algorithm to efficiently compute OT maps for the quadratic cost (Wasserstein-2 distance). Next, we extend the approach to the case when the input and output distributions are located in the spaces of different dimensions and derive error bounds for the computed OT map. We evaluate the algorithm on image generation and unpaired image restoration tasks. In particular, we consider denoising, colorization, and inpainting, where the optimality of the restoration map is a desired attribute, since the output (restored) image is expected to be close to the input (degraded) one.

연구 동기 및 목표

• 고차원 데이터에서 복원 정확도가 떨어지는 오토에코더 기반 잠재 공간에 의존하는 이전 최적 운반 기반 생성 모델의 한계를 해결한다.
• 생성 모델링을 위해 원래의 환경 공간(예: 이미지의 픽셀 공간)에서 최적 운반 맵을 직접, 엔드 투 엔드로 계산하는 방법을 개발한다.
• 입력-출력 공간의 차원이 동일하거나 다를 경우 모두에 대해 계산된 OT 맵에 대한 이론적 오차 한계를 수립한다.
• 대규모 시각 작업(예: 이미지 생성 및 양방향이 아닌 이미지 복원)에서 OT 맵을 생성 모델로써 실용적으로 활용할 수 있음을 보여준다.
• 입력 속성(예: 콘텐츠, 구조)을 유지하는 안정적이고 효율적이며 이론적으로 탄탄한 OT 기반 생성 모델링 프레임워크를 제공한다. 이는 복원 작업에 매우 중요하다.

제안 방법

• 동일한 차원의 환경 공간 내에서 원천 분포와 목표 분포 간의 제곱 비용(워샤르슈타인-2 거리)에 대한 최적 운반 맵을 계산하기 위한 최소-최대 최적화 문제를 수립한다.
• 운반 맵 $ G $ 와 이중 잠재 함수 $ ho $ 를 위한 신경망 아키텍처를 사용하며, 칸토로비치-루빈슈타인 이중 공식을 활용한 GAN 유사 목적 함수로 학습한다.
• 다른 차원의 공간에 있는 분포를 처리하기 위해 투영 기반 정규화와 수정된 이중 최적화 목적 함수를 도입하여 프레임워크를 확장한다.
• 계산된 OT 맵에 대한 이론적 오차 한계를 도출하며, 이는 이중 갭과 이중 잠재 함수의 스무스함에서 유도되며, 약한 조건 하에서도 진짜 OT 맵으로 수렴함을 보장한다.
• 배치 정규화와 잔여 블록을 사용하여 안정성과 성능을 높이기 위해 운반 맵 $ G $ 와 잠재 함수 $ ho $ 를 함께 확률적 경사 하강법으로 학습한다.
• 기울기 페널티와 스펙트럼 정규화를 포함한 WGAN-QC 유사 학습 설정을 채택하여 학습을 안정화하고 샘플 품질을 향상시킨다.

실험 결과

연구 질문

• RQ1 고차원 환경 공간(예: 이미지 픽셀)에서 오토에코더 기반 잠재 공간에 의존하지 않고, 최적 운반 맵을 직접 효과적으로 생성 모델로 사용할 수 있는가?
• RQ2 딥 신경망을 사용하여 제곱 비용(워샤르슈타인-2 거리)에 대해 환경 공간에서 최적 운반 맵을 효율적으로 계산할 수 있는가?
• RQ3 원천 및 목표 분포가 서로 다른 차원의 공간에 있을 경우, 계산된 OT 맵에 대해 어떤 이론적 보장(예: 오차 한계)을 제공할 수 있는가?
• RQ4 제안된 환경 공간 기반 OT 맵이 이미지 생성 및 양방향이 아닌 이미지 복원 작업에서 기존의 OT 기반 생성 모델보다 뛰어나게 성능을 발휘하는가?
• RQ5 최적성은 핵심 조건인 복원 작업(예: 노이즈 제거, 색상화)에서 OT 맵이 입력 속성(예: 콘텐츠, 구조)을 유지할 수 있는가?

주요 결과

• 제안된 방법은 CIFAR-10(11.2) 및 CelebA(14.8)에서 최신 기술 수준(FID) 점수를 기록하며, 이전의 OT 기반 모델을 능가하고 최신 기술 수준의 GAN과 경쟁 가능하다.
• CelebA $128 imes128$ 이미지 생성 벤치마크에서 FID 점수 14.8을 달성하여 고해상도 이미지 합성에서 뛰어난 샘플 품질과 다양성을 입증한다.
• 양방향이 아닌 이미지 복원 작업(예: 노이즈 제거, 색상화, 메쉬 복구)에서 OT 맵 기반 모델은 CycleGAN 및 기타 베이스라인 대비 뛰어난 시각적 품질과 구조적 정확도를 보이며, LPIPS는 낮고 FID는 높다.
• 동차 및 이종 차원 설정 모두에서 계산된 OT 맵에 대한 이론적 오차 한계가 확립되었으며, 약한 정규성 조건 하에서 진짜 최적 맵으로 수렴함을 보여준다.
• 계산 복잡도가 OT 기반 GANs(예: WGAN-QC)와 유사하여, ImageNet 및 CelebA에서 대규모 학습을 수행할 수 있으며 추가 오버헤드 없이 실행 가능하다.
• 제거 분석 결과, 잠재 공간 OT 방법보다 환경 공간에서 직접 OT 맵을 사용할 경우 성능이 더 뛰어나며, 특히 오토에코더 복원 오차가 품질을 떨어뜨리는 고차원 이미지 작업에서 두드러진다.