[논문 리뷰] Straight-Through Estimator as Projected Wasserstein Gradient Flow

Pengyu Cheng, Chang Liu외 4명|arXiv (Cornell University)|2019. 10. 05.

Advanced Neuroimaging Techniques and Applications참고 문헌 25인용 수 6

한 줄 요약

이 논문은 이산 분포 위에서의 기댓값 비용 최소화 문제에 대해, 직선 통과(ST) 추정기를 투영된 와서스타인 경로 유동(pWGF)으로 간주함으로써 ST 추정기의 이론적 기반을 구축한다. 이는 포isson과 같은 무한 지지도를 가진 분포에 대해 수렴 성능을 향상시키는 새로운 pWGF 기반 최적화 방법을 제안하며, 베르누이 및 포isson 잠재변수 모델에서 기존 방법들을 능가하고 ST의 성능을 유지한다.

ABSTRACT

The Straight-Through (ST) estimator is a widely used technique for back-propagating gradients through discrete random variables. However, this effective method lacks theoretical justification. In this paper, we show that ST can be interpreted as the simulation of the projected Wasserstein gradient flow (pWGF). Based on this understanding, a theoretical foundation is established to justify the convergence properties of ST. Further, another pWGF estimator variant is proposed, which exhibits superior performance on distributions with infinite support,e.g., Poisson distributions. Empirically, we show that ST and our proposed estimator, while applied to different types of discrete structures (including both Bernoulli and Poisson latent variables), exhibit comparable or even better performances relative to other state-of-the-art methods. Our results uncover the origin of the widespread adoption of the ST estimator and represent a helpful step towards exploring alternative gradient estimators for discrete variables.

연구 동기 및 목표

• 이산 변분 추론에서 널리 사용되지만 잘 이해되지 않는 직선 통과(ST) 기울기 추정기의 이론적 정당성을 제공하는 것.
• 이산 분포 위에서 기댓값 비용을 최소화하는 맥락에서 ST와 투영된 와서스타인 경로 유동(pWGF) 간의 관계를 체계화하는 것.
• 무한 지지도를 가진 분포(예: 포isson)에 대해 더 빠른 수렴을 달성하는 pWGF 기반의 새로운 기울기 추정기를 개발하는 것.
• 제안된 방법을 베르누이 및 포isson 잠재변수 모델에 대해 실증적으로 검증하여 최신 기준 성능과의 성능을 비교하는 것.

제안 방법

• 이 방법은 ST를 2-와서스타인 공간에서 투영된 와서스타인 경로 유동(pWGF)을 시뮬레이션하는 것으로 해석한다. 여기서 기울기는 연속적인 완화된 공간에서 계산된 후 이산 분포 가족으로 다시 투영된다.
• 각 반복 단계에서 알고리즘은 먼저 와서스타인 경로 유동을 통해 샘플을 업데이트한다: $\tilde{\bm{z}}_n = \bm{z}_n + \varepsilon \nabla f(\bm{z}_n)$, 이로써 일시적인 분포 $\tilde{\mu}_k$ 가 생성된다.
• 그 후, 분포 $\tilde{\mu}_k$ 는 2-와서스타인 거리 $W^2(\mu, \tilde{\mu}_k)$ 를 최소화함으로써 이산 가족 $\mathcal{M}$ 으로 다시 투영된다.
• 인과 분포 가정 하에, 와서스타인 거리 최소화는 각 차원의 마진 분포 거리 최소화로 줄어들며, 이는 효율적인 계산을 가능하게 한다.
• 베르누이 분포의 경우, 누적분포함수를 기반으로 한 임계값 기반 규칙을 사용해 최적의 투영을 해석적으로 유도한다.
• 무한 지지도를 가진 분포(예: 포isson)로의 확장을 위해 투영 단계를 무한 지지도를 처리할 수 있도록 조정함으로써, 더 빠른 수렴을 이룬다.

실험 결과

연구 질문

• RQ1 직선 통과 추정기는 연속 최적화 프레임워크를 통해 이론적으로 정당화될 수 있는가?
• RQ2 투영된 와서스타인 경로 유동(pWGF)은 ST 추정기의 타당하고 일반화 가능한 해석인가?
• RQ3 pWGF 기반 추정기는 특히 무한 지지도를 가진 분포에서 기존 기울기 추정기보다 성능이 뛰어나게 되는가?
• RQ4 제안된 pWGF 추정기의 성능은 베르누이 및 포isson 잠재변수 모델에서 ST 및 다른 최신 기준 방법과 비교해 어떻게 되는가?

주요 결과

• 직선 통과 추정기는 투영된 와서스타인 경로 유동(pWGF)을 시뮬레이션하는 것으로 이론적으로 정당화되며, 그 경험적 성공에 대한 엄밀한 기반을 제공한다.
• 베르누이 잠재변수 모델에서 제안된 pWGF 기반 추정기는 ST와 비교해 유사하거나 더 뛰어난 성능을 보이며, 비선형 아키텍처를 사용한 MNIST에서 테스트 ELBO가 104.6을 기록한다.
• 포isson 잠재변수 모델에서는 제안된 pWGF 변형이 ST 및 다른 기준 방법보다 더 빠른 수렴을 보이며, 이는 무한 지지도 분포를 고려한 설계 덕분이다.
• MNIST에서 모든 테스트 아키텍처에서 pWGF 방법은 Gumbel-Softmax, RELAX, REBAR를 모두 능가하며, 선형 아키텍처에서 테스트 ELBO가 129.2를 기록한다.
• 이론적 분석을 통해 인과 분포의 경우, 두 분포 간 2-와서스타인 거리가 각 차원의 제곱 마진 거리의 합과 동일하다는 것을 증명한다.
• 베르누이 분포의 경우, pWGF 하에서 최적의 투영이 누적분포함수를 포함하는 임계값 기반 규칙과 동일하다는 것이 입증되었으며, $W^2(\text{Bern}(p), \text{Bern}(q)) = |p - q|$ 를 만족한다.