[논문 리뷰] Parallel Streaming Wasserstein Barycenters

Matthew Staib, Sebastian Claici외 2명|arXiv (Cornell University)|2017. 05. 21.

Markov Chains and Monte Carlo Methods참고 문헌 12인용 수 26

한 줄 요약

이 논문은 반정적 이분할 바리센터를 기반으로 한 확률적 경사 하강법을 사용하여 임의의 확률 분포의 워셔스타인 바리센터를 통신 효율적이고 병렬 스트리밍 방식으로 계산하는 알고리즘을 제안한다. 이는 연속적이고 비정상적인 입력 측정값에 대해서도 확장 가능하고 실시간으로 바리센터를 추정할 수 있도록 하며, 이론적 수렴 보장과 함께 미세 격자(예: $n \approx 10^6$) 및 대규모 베이지안 추론 작업에서 기존 방법보다 뛰어난 경험적 성능을 보인다.

ABSTRACT

Efficiently aggregating data from different sources is a challenging problem, particularly when samples from each source are distributed differently. These differences can be inherent to the inference task or present for other reasons: sensors in a sensor network may be placed far apart, affecting their individual measurements. Conversely, it is computationally advantageous to split Bayesian inference tasks across subsets of data, but data need not be identically distributed across subsets. One principled way to fuse probability distributions is via the lens of optimal transport: the Wasserstein barycenter is a single distribution that summarizes a collection of input measures while respecting their geometry. However, computing the barycenter scales poorly and requires discretization of all input distributions and the barycenter itself. Improving on this situation, we present a scalable, communication-efficient, parallel algorithm for computing the Wasserstein barycenter of arbitrary distributions. Our algorithm can operate directly on continuous input distributions and is optimized for streaming data. Our method is even robust to nonstationary input distributions and produces a barycenter estimate that tracks the input measures over time. The algorithm is semi-discrete, needing to discretize only the barycenter estimate. To the best of our knowledge, we also provide the first bounds on the quality of the approximate barycenter as the discretization becomes finer. Finally, we demonstrate the practical effectiveness of our method, both in tracking moving distributions on a sphere, as well as in a large-scale Bayesian inference task.

연구 동기 및 목표

• 다양한 원천에서 온 동일하지 않은 분포이거나 연속적인 확률 측정값을 효율적이고 통신 비용이 적은 방식으로 확장 가능하게 집계하는 문제를 해결하기 위해.
• 시간이 지남에 따라 변화하는 비정상적인 입력 분포에 적응하는 실시간 스트리밍 방식의 워셔스타인 바리센터 계산을 가능하게 하기 위해.
• 바리센터 근사의 정확도에 대한 이론적 수렴 경계를, 바리센터 지지점 수가 증가함에 따라 제공하기 위해.
• 특히 미세 이산화에 대해 기존 방법의 확장성 한계를 극복하기 위해, 대규모 선형 프로그래밍이나 정규화된 최적 운반 이론과 같은 방법의 한계를 극복하기 위해.
• 고정밀도 바리센터 추정이 필수적인 대규모 베이지안 추론 및 센서 융합 응용 분야를 지원하기 위해.

제안 방법

• 알고리즘은 $n$개의 지지점이 있는 이산 바리센터를 반복적으로 갱신하는 확률적 경사 하강법(SGD)을 사용한다. 이 경우 바리센터만 이산화되어 있으며, 반정적 이분할 접근 방식을 취한다.
• 각 워커 머신은 병렬로 입력 측정값의 부분집합을 처리하고, 입력 분포의 표본 액세스를 통해 확률적 기울기를 계산한다.
• 각 반복에서 워커 간 통신은 단일 정수 값만 필요로 하여 분산 환경에서 높은 통신 효율성을 확보한다.
• 핵심 최적화 문제는 최적 운반의 이중 공식을 활용하여 바리센터 이중 잠재변수에 대한 볼록 최대화 문제로 공식화된다.
• 알고리즘은 비정상적인 분포에 대해 강건하며, 변화하는 입력 측정값에 대응해 바리센터 추정치를 지속적으로 갱신한다.
• 바리센터 지지점 수 $n$이 증가함에 따라 근사 오차에 대한 이론적 수렴이 입증된다.

실험 결과

연구 질문

• RQ1 연속적인 확률 분포의 워셔스타인 바리센터를 스트리밍 환경에서 확장 가능하고 통신 효율적이며 병렬적으로 계산할 수 있는 알고리즘을 설계할 수 있는가?
• RQ2 바리센터 근사의 정확도는 바리센터 지지점 수 $n$에 따라 어떻게 달라지며, 이에 대한 이론적 수렴 경계를 설정할 수 있는가?
• RQ3 매 단계에서 대규모 최적화 문제를 다시 풀지 않고도 비정상적인 입력 분포에서 동적으로 바리센터 추정치를 유지할 수 있는가?
• RQ4 미세 격자에서 기존의 선형 프로그래밍이나 정규화된 최적 운반 이론과 비교할 때 정확도와 확장성 측면에서 제안된 방법의 성능은 어떠한가?
• RQ5 대규모 베이지안 추론 작업에서 부분 사후 분포를 포함할 때 이 방법의 실용적 영향은 무엇인가?

주요 결과

• 바리센터 지지점 수 $n \approx 10^4$일 때, 317초 후에 진짜 사후분포와 약 26의 워셔스타인 거리를 기록하며, 유사한 격자에서 선형 프로그래밍 방법보다 뛰어난 성능을 보였다.
• $n \approx 10^6$일 경우, 16스레드 노드당 메모리 사용량이 2GB 미만으로 유지되며 고정밀도를 확보했고, 선형 프로그래밍 접근 방식은 $n=480$일 때조차 메모리 제약으로 실패했다.
• $n \approx 10^4$에서 다양한 스텝 사이즈에 대해 선형 프로그래밍보다 더 나은 근사를 제공했으며, 조기 종료를 통해 훨씬 뛰어난 결과를 얻을 수 있었다.
• 이론적 경계는 $n$이 증가할수록 근사 오차가 감소함을 보여주며, 일반적인 경우에서 반정적 이분할 바리센터 추정에 대해 알려진 바 없는 첫 번째 수렴 보장이다.
• 알고리즘은 구상 위에서 움직이는 분포를 실시간으로 추적할 수 있으며, 베이지안 추론에서 워셔스타인 평균 부분 사후분포(WASP)의 정확도를 크게 향상시켰다.
• 메모리 사용량이 극히 적고 높은 병렬 효율성을 확보하여 $n \approx 10^6$개 지지점까지 확장 가능하며, 대규모 응용 분야에서의 실용성을 입증했다.