[논문 리뷰] Generative Modeling on Manifolds Through Mixture of Riemannian Diffusion Processes

Jaehyeong Jo, Sung Ju Hwang|arXiv (Cornell University)|2023. 10. 11.

Generative Adversarial Networks and Image SynthesisComputer Science인용 수 3

한 줄 요약

이 논문은 일반 리만 다양체 위에서 시간에 따라 변화하는 확률적 과정을 종점 조건이 부여된 다리 과정의 혼합으로 구성하는 새로운 생성 모델링 프레임워크인 리만 다이퓨전 믹스처(RFM)를 제안한다. 이는 이전의 다이퓨전 모델들이 겪는 시간 역전 및 스코어 매칭 문제를 피한다. 이 방법은 훨씬 적은 훈련 중 시뮬레이션 단계로도 효율적이고 확장 가능한 훈련을 가능하게 하여, 구형, 토러스, 쌍곡면, 3D 메쉬와 같은 다양한 다양체에서 최신 기준 성능을 달성한다.

ABSTRACT

Learning the distribution of data on Riemannian manifolds is crucial for modeling data from non-Euclidean space, which is required by many applications in diverse scientific fields. Yet, existing generative models on manifolds suffer from expensive divergence computation or rely on approximations of heat kernel. These limitations restrict their applicability to simple geometries and hinder scalability to high dimensions. In this work, we introduce the Riemannian Diffusion Mixture, a principled framework for building a generative diffusion process on manifolds. Instead of following the denoising approach of previous diffusion models, we construct a diffusion process using a mixture of bridge processes derived on general manifolds without requiring heat kernel estimations. We develop a geometric understanding of the mixture process, deriving the drift as a weighted mean of tangent directions to the data points that guides the process toward the data distribution. We further propose a scalable training objective for learning the mixture process that readily applies to general manifolds. Our method achieves superior performance on diverse manifolds with dramatically reduced number of in-training simulation steps for general manifolds.

연구 동기 및 목표

• 기존의 생성 모델이 열핵 근사나 계산 비용이 큰 발산 계산에 의존하는 리만 다양체 상의 문제점을 해결한다.
• 고차원 및 일반 다양체 환경에서 스코어 기반 다이퓨전 모델과 발산이 많은 CNF 모델의 확장성과 기하학적 일관성 부족 문제를 해결한다.
• 기하학적으로 인식 가능한 원칙적인 생성 과정을 개발하여, 사전 분포에서 데이터 분포로의 이동을 리만 다리 과정의 혼합으로 직접 모델링한다.
• 시간 역전을 피하고 시뮬레이션 단계를 줄이기 위해, 일반 다양체에서 효율적인 훈련을 가능하게 하는 이중 다리 매칭 목적 함수를 도입한다.
• 특히 고차원 환경에서 복잡한 다양체 상에서 정확한 모델링을 위해 결정론적 흐름보다 확률적 과정이 필수적임을 입증한다.

제안 방법

• 시간 역전에 의존하지 않고, 고정된 종점 조건이 부여된 다리 과정의 혼합으로 생성 과정을 모델링하는 리만 다이퓨전 믹스처 프레임워크를 제안한다.
• 열핵 근사나 비가역적인 발산 계산이 필요 없는 일반 다양체 상에서의 실용적인 리만 다리 과정 가족을 도입한다.
• 전방 및 후방 다리 과정을 매칭하는 이중 다리 매칭 목적 함수를 정의하여, 임의의 리만 기하학에 적용 가능한 효율적이고 확장 가능한 훈련을 가능하게 한다.
• 다양체 기하학에 기반한 가장 가능성이 높은 종점으로 향하는 드리프트를 사용하여, 사전 분포(예: 다양체 상의 균일 분포)에서 데이터 분포로의 확률적 사상 구축.
• 연속 정규화 흐름(CNF)이 제안된 프레임워크의 특수한 경우임을 이론적으로 입증하여 그 일반성 확인.
• 고차원 추론을 가능하게 하기 위해 JAX를 사용하여 드리프트 벡터 장의 효율적인 신경망 파arameterization을 구현한다.

실험 결과

연구 질문

• RQ1 노이즈 과정의 시간 역전에 의존하지 않고도, 스코어 매칭 및 발산 계산을 피할 수 있는 리만 다양체 상의 생성 모델을 구축할 수 있는가?
• RQ2 일반 다양체 상에서 기존의 다이퓨전 및 CNF 모델에 비해 더 기하학적으로 직관적이고 확장 가능한가? 리만 다리 과정의 혼합이 그러한 대안이 될 수 있는가?
• RQ3 제안된 이중 다리 매칭 목적 함수는 이전 방법에 비해 더 적은 훈련 중 시뮬레이션 단계로 더 빠르고 정확한 훈련을 가능하게 하는가?
• RQ4 RFM의 성능은 고차원 및 비유한 기하학을 포함한 다양한 다양체에서 최신 기준 모델과 비교해 어떻게 되는가?
• RQ5 복잡한 다양체 상에서 정확한 모델링을 위해 확률적 생성 과정이 필수적인가, 아니면 결정론적 흐름으로도 충분한가?

주요 결과

• 제안된 리만 다이퓨전 믹스처(RFM)는 원형, 토러스, 쌍곡면, 3D 메쉬와 같은 다양한 다양체 상에서 합성 및 실세계 데이터셋에서 최신 기준 모델을 능가하거나 동등하게 성능을 내며, 전반적으로 뛰어난 성능을 보였다.
• 스탠포드 번이 및 스포트 더 카우와 같은 일반 폐쇄 다양체에서 RFM는 이전 CNF 모델이 요구하는 시뮬레이션 단계의 5%만으로도 뛰어난 성능을 달성했다.
• 고차원 환경에서도 효과적으로 확장되며, 발산 계산에 의존하는 스코어 기반 다이퓨전 모델보다 훨씬 빠른 훈련 시간을 기록했다.
• 이중 다리 매칭 목적 함수는 볼카노 및 스탠포드 번이 데이터셋에서 최소 15회의 시뮬레이션 단계로도 정확한 모델링을 가능하게 했으며, 단일 방향 접근 방식은 훨씬 더 많은 단계가 필요했다.
• 쌍곡면 상에서 RFM는 복잡한 다중 모달 분포(예: 랩드 가우시안 혼합)를 성공적으로 학습했고, 파oincaré 디스크 투영을 통해 정확한 밀도 추정을 시각화했다.
• 절단 실험 결과, 제안된 목적 함수의 설계 선택이 중요하며, 균일한 시간 샘플링보다 시간 스케일링 훈련이 더 뛰어난 성능을 내는 것으로 확인되었다.