[논문 리뷰] Variational Wasserstein gradient flow

Jiaojiao Fan, Qinsheng Zhang외 2명|arXiv (Cornell University)|2021. 12. 04.

Advanced Neuroimaging Techniques and Applications인용 수 8

한 줄 요약

이 논문은 밀도에 의존하는 목적 함수를 매개변수 함수에 대한 변분 형식으로 대체함으로써 고차원 경험 분포에서 효율적이고 확장 가능한 계산을 가능하게 하는 변분 워샤르스키 경량 흐름 방법을 제안한다. f-발산의 변분 형태와 입력 볼록 신경망을 활용함으로써, 고비용의 밀도 추정 및 헤시안 행렬식 근사가 필요 없이, CIFAR10에서 최고 수준의 FID 점수(23.1)를 달성하며 안정적인 학습과 기존 GAN보다 향상된 샘플 품질을 확보한다.

ABSTRACT

Wasserstein gradient flow has emerged as a promising approach to solve optimization problems over the space of probability distributions. A recent trend is to use the well-known JKO scheme in combination with input convex neural networks to numerically implement the proximal step. The most challenging step, in this setup, is to evaluate functions involving density explicitly, such as entropy, in terms of samples. This paper builds on the recent works with a slight but crucial difference: we propose to utilize a variational formulation of the objective function formulated as maximization over a parametric class of functions. Theoretically, the proposed variational formulation allows the construction of gradient flows directly for empirical distributions with a well-defined and meaningful objective function. Computationally, this approach replaces the computationally expensive step in existing methods, to handle objective functions involving density, with inner loop updates that only require a small batch of samples and scale well with the dimension. The performance and scalability of the proposed method are illustrated with the aid of several numerical experiments involving high-dimensional synthetic and real datasets.

연구 동기 및 목표

• 밀도 추정 및 헤시안 행렬 계산으로 인해 고차원 공간에서 워샤르스키 경량 흐름의 계산 불가능성 문제를 해결하기 위해.
• 경험 분포 상에서 워샤르스키 경량 흐름을 수치적으로 안정적이고 확장 가능한 방법으로 계산하기 위해.
• 밀도에 의존하는 목적 함수를 표본 기반 평가가 가능한 변분 형식으로 대체하기 위해.
• 기본 GAN과 비교해 학습 안정성을 향상시키고 모드 붕괴를 방지하기 위해.
• 합성 및 실제 데이터 세트(예: MNIST 및 CIFAR10 포함)에서 고차원 데이터에 대한 본 방법의 효과성을 입증하기 위해.

제안 방법

• f-발산 기반의 변분 표현을 사용해 목적 함수를 정의하고, 매개변수화된 함수의 클래스에 대해 최대화한다.
• 입력 볼록 신경망(ICNN)을 사용해 볼록 함수의 기울기로 최적의 운반 지도를 매개변수화한다.
• 작은 배치의 샘플에 대한 내부 반복 업데이트를 통해 명시적 밀도 평가 또는 헤시안 행렬식 근사가 필요 없도록 한다.
• 스토하스틱 최적화를 통해 JKO 체계를 구현하며, 각 시간 단계는 정규화된 워샤르스키 거리와 변분 목적 함수를 최소화한다.
• 초기 분포(예: 정규분포)를 목표 분포로 점진적으로 변화시키기 위해 연속된 운반 지도의 시퀀스를 사용한다.
• 학습된 운반 지도를 따라 샘플링 및 밀도 평가를 가능하게 한다.

실험 결과

연구 질문

• RQ1 고차원 설정에서 워샤르스키 경량 흐름의 목적 함수에 대한 변분 형식이 명시적 밀도 추정이 필요 없도록 할 수 있는가?
• RQ2 제안된 방법은 기존 JKO 기반 접근법과 비교해 계산 효율성 및 확장성 측면에서 어떻게 성능을 내는가?
• RQ3 변분 형식이 순간 일치 및 적분 확률 거리 측정 기반 임bedding 부등식과 같은 기하학적 및 통계적 성질을 유지하는가?
• RQ4 모드 붕괴 없이 CIFAR10과 같은 이미지 생성 벤치마크에서 최고 수준의 샘플 품질을 달성할 수 있는가?
• RQ5 단계 크기 및 흐름 단계 수에 따라 방법의 성능은 어떻게 변화하는가?

주요 결과

• 제안된 방법은 CIFAR10에서 23.1의 Fréchet Inception Distance(FID) 점수를 달성하여 GAN 및 기타 플로우 기반 생성 모델을 능가한다.
• a=5.0의 단계 크기로 학습할 경우 안정적인 수렴을 보이며 모드 붕괴를 방지하지만, 더 큰 값은 불안정성과 열악한 샘플 품질로 이어진다.
• 헤시안 행렬식 근사를 요구하는 이전의 JKO 기반 방법과 비교해 뛰어난 확장성과 계산 효율성을 보였다.
• 변분 목적 함수는 순간 일치 성질과 적분 확률 거리 측정 기반 임베딩 부등식을 만족하여 이론적 타당성을 뒷받침한다.
• FID 점수는 시간 단계에 따라 수렴함을 보이며, 흐름 단계 수가 증가함에 따라 안정적인 학습 역학을 유지함을 시사한다.
• CIFAR10에서 경쟁력 있는 Inception Score(7.48 ± 0.12)를 달성하여 AE-OT-GAN 및 OTM과 같은 최고 수준의 모델과 유사한 성능을 보였다.