[논문 리뷰] Variational inference via Wasserstein gradient flows

Marc Lambert, Sinho Chewi외 3명|arXiv (Cornell University)|2022. 05. 31.

Markov Chains and Monte Carlo Methods인용 수 12

한 줄 요약

이 논문은 가우시안 측도의 Bures–Wasserstein 공간에서 워셔스타인 기울기 유동을 사용하는 새로운 변분 추론 방법을 제안하며, 복잡한 사후 분포의 원칙적인 근사화를 가능하게 한다. 이 방법은 로그-볼록성 조건 하에서 강력한 이론적 보장을 제공하며, ODE 기반 최적화를 통해 혼합 가우시안을 사용한 다중모달 타겟의 효과적인 근사화를 보여준다.

ABSTRACT

Along with Markov chain Monte Carlo (MCMC) methods, variational inference (VI) has emerged as a central computational approach to large-scale Bayesian inference. Rather than sampling from the true posterior $π$, VI aims at producing a simple but effective approximation $\hat π$ to $π$ for which summary statistics are easy to compute. However, unlike the well-studied MCMC methodology, algorithmic guarantees for VI are still relatively less well-understood. In this work, we propose principled methods for VI, in which $\hat π$ is taken to be a Gaussian or a mixture of Gaussians, which rest upon the theory of gradient flows on the Bures--Wasserstein space of Gaussian measures. Akin to MCMC, it comes with strong theoretical guarantees when $π$ is log-concave.

연구 동기 및 목표

• 타겟 사후 분포가 로그-볼록일 경우, 특히 가우시안 및 혼합 가우시안 근사화에 대해 기존 변분 추론(VI)의 이론적 보장 부족 문제를 해결한다.
• 확률 측도의 Bures–Wasserstein 공간에서 기울기 유동에 기반한 체계적인 VI 프레임워크를 개발한다.
• 시간에 따라 변화하는 사후 근사 분포를 ODE로 모델링하여 변분 칼만 필터링 아이디어를 연속 시간 동역학계로 확장한다.
• 헤시안 행렬 계산이 필요 없이 평균 및 공분산 매개변수를 학습하기 위한 확장 가능한 최적화 기반 알고리즘을 제공한다.
• 혼합 가우시안 근사화를 통해 평균장 VI를 초월하여 상관관계와 다중모달 구조를 포착한다.

제안 방법

• 제안 분포와 진짜 사후 분포 간의 KL 발산을 최소화하는 변분 추론을, 가우시안 측도의 Bures–Wasserstein 기하학에서 기울기 유동을 통해 수식화한다.
• 이름표 없이 적분을 통한 부분적 통합과 워셔스타인 기울기 유동 이론을 활용해 각 가우시안 성분의 평균 및 공분산 매개변수의 연속 시간 ODE를 유도한다.
• K개의 가우시안 성분을 갖는 입자 기반 표현을 사용하며, 각 성분의 평균과 공분산의 하우스홀더 분해 요소가 결합된 ODE에 의해 진화한다.
• 4차 룬게-쿠타 방법을 사용해 ODE를 수치적으로 통합하며, 상태 벡터는 모든 평균 및 벡터화된 하우스홀더 매개변수를 통합한다.
• 공분산 행렬의 정의성 조건을 유지하기 위해 공분산을 직접이 아닌 하한 삼각 하우스홀더 분해 요소를 진화시킨다.
• 수치적 통합을 위해 세타 포ints를 사용하는 쿠바처 규칙을 활용해 ODE 내 기대값을 계산한다.

실험 결과

연구 질문

• RQ1 Bures–Wasserstein 공간에서의 워셔스타인 기울기 유동은 기존 변분 추론에 비해 이론적으로 타당하고 확장 가능한 가우시안 및 혼합 가우시안 근사화의 대안이 될 수 있는가?
• RQ2 타겟 사후 분포가 로그-볼록일 경우, 라운지엔-디퓨전과 유사한 강력한 수렴 보장을 달성할 수 있는가?
• RQ3 특히 가우시안 성분 수를 늘릴 경우, 이 방법은 다중모달 사후 분포를 얼마나 잘 근사할 수 있는가?
• RQ4 알고리즘은 초기화에 얼마나 민감한가? 실무에서 열악한 초기화를 어떻게 완화할 수 있는가?
• RQ5 헤시안 정보가 필요 없이도, 복잡한 사후 분포의 구조(예: 상관관계 및 다중 모드)를 효과적으로 포착할 수 있는가?

주요 결과

• 제안된 방법은 로그-볼록 타겟 하에서 이론적 수렴 보장을 달성하며, MCMC 기법(예: 라운지엔 디퓨전)에서 알려진 강력한 비점근 분석을 변분 추론으로 확장한다.
• 이중 모달 타겟에 대해 20개의 가우시안으로 구성된 혼합 모델이 두 모드를 성공적으로 포착하였으며, 등고선도는 진짜 사후 분포와 밀도 있는 일치를 보였다.
• KL 발산은 약 30개의 통합 단계(스텝 크기가 0.1일 경우 약 300단계) 후에 안정화되어 최적화 과정의 수렴을 나타낸다.
• 가우시안 입자의 수를 늘일수록 근사 품질이 향상되며, 4~6개의 비균일 가중치를 갖는 모드를 가진 타겟에서 더 낮은 KL 발산과 더 정확한 밀도 추정이 가능해졌다.
• 초기화가 수렴에 큰 영향을 미친다: 한 모드 근처에 초기화된 입자는 그 모드 근처에 머물러 있으며, 모드에서 동일한 거리에 있는 입자는 공분산이 증가하는 평균 방향으로 수렴한다. 이는 다수의 초기 입자를 사용할 필요성을 강조한다.
• 이론적 통합을 통해 헤시안 계산을 피하는 Hessian-무관 업데이트를 도출함으로써, 효율적이고 안정적인 최적화를 실현하였다.