[논문 리뷰] Neural Wasserstein Gradient Flows for Maximum Mean Discrepancies with Riesz Kernels

Fabian Altekrüger, Johannes Hertrich외 1명|arXiv (Cornell University)|2023. 01. 27.

Generative Adversarial Networks and Image Synthesis인용 수 5

한 줄 요약

이 논문은 비연속 리에즈 커널을 가진 최대 평균 차이(MMD) 기능에 대해 신경망 기반 워셔스타인 경사 하강 흐름을 제안한다. 전이 및 속도 계획을 매핑이 아닌 생성 신경망을 사용해 근사함으로써 특이 측도와 절대 연속 측도를 모두 처리할 수 있다. 이 방법은 정방향 및 역방향 스킴이 해석적 해에 수렴함을 보이며, 비연속 커널을 가진 MMD 흐름과 바리센터에 적용된다.

ABSTRACT

Wasserstein gradient flows of maximum mean discrepancy (MMD) functionals with non-smooth Riesz kernels show a rich structure as singular measures can become absolutely continuous ones and conversely. In this paper we contribute to the understanding of such flows. We propose to approximate the backward scheme of Jordan, Kinderlehrer and Otto for computing such Wasserstein gradient flows as well as a forward scheme for so-called Wasserstein steepest descent flows by neural networks (NNs). Since we cannot restrict ourselves to absolutely continuous measures, we have to deal with transport plans and velocity plans instead of usual transport maps and velocity fields. Indeed, we approximate the disintegration of both plans by generative NNs which are learned with respect to appropriate loss functions. In order to evaluate the quality of both neural schemes, we benchmark them on the interaction energy. Here we provide analytic formulas for Wasserstein schemes starting at a Dirac measure and show their convergence as the time step size tends to zero. Finally, we illustrate our neural MMD flows by numerical examples.

연구 동기 및 목표

• 비연속 리에즈 커널을 가진 MMD 기능의 워셔스타인 경사 하강 흐름을 계산하기 위한 신경망 기반 방법을 개발하는 것. 이는 특이 측도와 절대 연속 측도 간의 전이를 가능하게 한다.
• 기존의 JKO 스킴과 입자 흐름의 한계를 극복하기 위해, 전이 매핑에 의존하지 않고 생성 네트워크를 사용해 전이 및 속도 계획을 근사함으로써, 이를 통해 전이 및 속도 계획을 근사하는 것.
• 대상 측도가 해석적으로 알려져 있지 않고 커널이 비볼록이며 비연속일 경우에도 MMD 흐름과 바리센터를 계산할 수 있도록 하는 것.
• 절대 연속 측도를 넘어서는 MMD 기반 생성 모델링을 위한 수치적으로 안정적이고 확장 가능한 프레임워크를 제공하는 것.
• 비연속 측도 진화를 포함한 복잡한 흐름에 대한 수렴 분석과 수치 실험을 통해 신경 스킴을 검증하는 것.

제안 방법

• 리에즈 커널을 가진 MMD 기능의 전이 계획의 분해를 학습하는 생성 신경망을 사용해 조르단-킨들러러-오토(JKO) 역방향 스킴을 근사한다.
• 워셔스타인 최강하강 흐름 기반의 정방향 스킴을 사용하며, 이 역시 생성 네트워크를 통해 속도 계획을 근사함으로써 특이 초기 측도에서의 진화를 가능하게 한다.
• MMD의 상호작용 에너지 및 위치 에너지 항을 이중 적분으로 모델링하여 표본 기반 근사를 가능하게 하며, 효율적인 기울기 계산을 가능하게 한다.
• 워셔스타인 거리와 리에즈 커널의 구조를 고려한 손실 함수를 사용해 신경망을 훈련함으로써, 기저 최적 수송 역학과의 일致성을 확보한다.
• MMD에 리에즈 커널을 적용할 경우 단일 차원 정렬을 통해 슬라이싱 표현이 가능하므로, 기울기 계산을 빠르게 할 수 있으며 이는 훈련 효율성을 향상시킨다.
• 이 스킴을 사용해 MMD 바리센터를 계산하고, 비연속적이며 비볼록 에너지 경관을 가진 복잡한 목표 밀도에서 샘플링을 수행한다.

실험 결과

연구 질문

• RQ1 비연속 리에즈 커널을 가진 MMD 기능에 대해 신경망이 워셔스타인 경사 하강 흐름을 효과적으로 근사할 수 있는가? 특히 측도가 특이 측도에서 절대 연속 측도로 전이되는 경우에 대해.
• RQ2 초기 측도가 특이일 경우, 신경 역방향 및 정방향 스킴이 고전적 입자 흐름과 비교해 수렴성과 안정성 면에서 어떻게 다른가?
• RQ3 측도의 절대 연속성을 가정하지 않고도 생성 네트워크가 전이 및 속도 계획을 얼마나 잘 모델링할 수 있는가?
• RQ4 시간 간격이 0으로 수렴할 때, 제안된 방법이 MMD 흐름의 해석적 해에 수렴하는가?
• RQ5 대상 측도가 서로 다른 부분다양체에 지지되어 있을 경우, 신경 스킴이 MMD 바리센터를 성공적으로 계산할 수 있는가?

주요 결과

• 시간 간격이 0으로 수렴할 때 신경 역방향 및 정방향 스킴이 리에즈 커널을 가진 MMD 흐름의 해석적 해에 수렴함을 확인하여 이론적 일致성을 입증하였다.
• 이 방법은 측도 전이를 성공적으로 처리한다: 딜라 측도는 차원에 따라 절대 연속 측도 또는 특이하지 않은 비딜라 측도로 진화하며, 이는 고전적 입자 흐름이 모델링할 수 없는 현상이다.
• 수치 실험 결과, 초기 측도가 특이일 경우 둘 다 신경 스킴이 고전적 입자 흐름보다 우수한 성능을 보이며, 비연속 역학을 더 잘 포착한다.
• 리에즈 커널의 사용 덕분에 슬라이싱을 통한 기울기 계산이 효율적이며, 이는 신경망의 신속한 훈련과 안정적인 최적화를 가능하게 한다.
• 신경 스킴은 원과 정사각형 경계에 지지된 측도 간의 MMD 바리센터를 성공적으로 계산하였으며, 이론적 바리센터에 정확하게 수렴하였다.
• MNIST 예제에서 초기 측도를 특이 측도가 아닌 균일 측도로 설정했을 경우, 모든 방법에서 유사한 성능를 보였으며, 이는 신경 스킴이 초기 측도 유형에 대해 강건함을 시사한다.