[논문 리뷰] Wasserstein Dependency Measure for Representation Learning

Sherjil Ozair, Corey Lynch외 4명|arXiv (Cornell University)|2019. 03. 28.

Domain Adaptation and Few-Shot Learning참고 문헌 42인용 수 20

한 줄 요약

이 논문은 상호정보량 추정에서 KL 발산 대신 워샤르슈타인 거리(Wasserstein distance)를 사용하여 안정적인 훈련을 가능하게 하는 립시츠 연속 신경망을 활용해, 새로운 표현 학습 목표인 워샤르슈타인 의존도 측정법(Wasserstein Dependency Measure, WDM)을 제안한다. 제안된 방법인 워샤르슈타인 예측 코드화(Wasserstein Predictive Coding, WPC)는 특히 데이터의 구조가 신경망의 인덕티브 바이어스와 일치하지 않을 때 상호정보량이 높은 작업에서 대비 예측 코드화(Contrastive Predictive Coding, CPC)보다 훨씬 우수한 표현 품질을 달성한다.

ABSTRACT

Mutual information maximization has emerged as a powerful learning objective for unsupervised representation learning obtaining state-of-the-art performance in applications such as object recognition, speech recognition, and reinforcement learning. However, such approaches are fundamentally limited since a tight lower bound of mutual information requires sample size exponential in the mutual information. This limits the applicability of these approaches for prediction tasks with high mutual information, such as in video understanding or reinforcement learning. In these settings, such techniques are prone to overfit, both in theory and in practice, and capture only a few of the relevant factors of variation. This leads to incomplete representations that are not optimal for downstream tasks. In this work, we empirically demonstrate that mutual information-based representation learning approaches do fail to learn complete representations on a number of designed and real-world tasks. To mitigate these problems we introduce the Wasserstein dependency measure, which learns more complete representations by using the Wasserstein distance instead of the KL divergence in the mutual information estimator. We show that a practical approximation to this theoretically motivated solution, constructed using Lipschitz constraint techniques from the GAN literature, achieves substantially improved results on tasks where incomplete representations are a major challenge.

연구 동기 및 목표

• 상호정보량 최대화의 근본적 한계를 해결한다. 즉, 상호정보량이 클수록 날카운 하한을 확보하기 위해 필요한 표본 수가 지수적으로 증가한다.
• 영상 이해나 강화학습과 같은 고상호정보량 작업에서 상호정보량 기반 방법이 완전한 표현을 학습하지 못함을 규명한다.
• KL 기반 상호정보량 추정기의 이론적·실용적 한계를 극복하기 위해 워샤르슈타인 거리에 기반한 새로운 학습 목표를 제안한다.
• 실험적으로 WPC(워샤르슈타인 의존도 측정법의 실용적 구현)가 CPC보다 더 완전하고 강건한 표현을 학습함을 입증한다. 특히 도전적인 데이터 분포에서 성능이 뛰어나다.
• 데이터의 구조가 컨볼루션 네트워크의 인덕티브 바이어스와 일치하지 않을 때 WPC가 미니배치 크기 변화에 덜 민감하고 일반화 성능이 뛰어나다는 것을 보여준다.

제안 방법

• 상호정보량 추정에서 KL 발산을 워샤르슈타인 거리로 대체하여, 워샤르슈타인 의존도 측정법(WDM)이라 불리는 새로운 의존도 측정법을 정의한다.
• GAN 문헌에서 유래한 기법을 활용해 상호정보량 추정기에서 사용하는 신경망에 리프시츠 연속성 조건을 도입함으로써 실용적인 추정기 구축.
• 상호정보량 목표를 WDM 목표로 대체한 대비 예측 코드화(CPC)-스타일 프레임워크를 사용해 표현 모델을 훈련한다.
• 가중치 클리핑 또는 기울기 페널티를 적용해 리프시츠 제약 조건을 강제함으로써 훈련 중 안정적이고 의미 있는 기울기 갱신을 보장한다.
• 컨텍스트와 미래 표현 간의 WDM을 최대화하도록 표현 모델을 훈련함으로써, 더 많은 변동 요인을 포착하도록 유도한다.
• 다양한 고상호정보량 데이터셋(예: MultiOmniglot, CelebA, MultiviewShapes3D)을 대상으로 WPC의 성능을 CPC와 비교해 평가한다.

실험 결과

연구 질문

• RQ1 왜 상호정보량 기반 표현 학습 방법은 영상이나 강화학습과 같은 고상호정보량 환경에서 완전한 표현을 학습하지 못하는가?
• RQ2 상호정보량 추정에서 KL 발산을 워샤르슈타인 거리로 대체하면 더 강건하고 완전한 표현을 얻을 수 있는가?
• RQ3 제안된 워샤르슈타인 예측 코드화(WPC) 방법의 성능은 다양한 데이터 분포와 네트워크 아키텍처에서 대비 예측 코드화(CPC)와 비교해 어떻게 다른가?
• RQ4 리프시츠 제약 조건은 낮은 데이터 또는 고상호정보량 환경에서 표현 학습의 안정성과 일반화 능력을 어느 정도 향상시키는가?
• RQ5 WPC는 컨볼루션 네트워크의 인덕티브 바이어스와 일치하지 않는 다양한 미니배치 크기와 데이터 구조에서도 뛰어난 성능을 유지하는가?

주요 결과

• 고상호정보량(~34.43 nats)을 가진 SplitCelebA 데이터셋에서, WPC는 완전히 연결된 네트워크를 사용해 0.87의 정확도를 달성했으며, 이는 CPC의 0.85보다 뛰어나다.
• 동일한 데이터셋에서 WPC는 완전히 연결된 네트워크와 컨볼루션 네트워크를 포함한 다양한 아키텍처에서 일관된 성능을 유지했지만, CPC는 컨볼루션 네트워크에서 성능이 크게 떨어졌다.
• 컨볼루션 네트워크의 인덕티브 바이어스와 일치하지 않는 StackedMultiOmniglot에서는 SpatialMultiOmniglot보다 WPC가 CPC를 더 넓은 폭으로 앞서며, 아키텍처 불일치에 대한 강건성을 보였다.
• WPC는 미니배치 크기가 32일 때 최적의 성능을 보였고, 더 큰 배치 크기에서의 성능 향상이 거의 없었지만, CPC는 안정화를 위해 더 큰 배치 크기가 필요했다.
• MultiviewShapes3D에서 WPC는 모든 테스트된 데이터셋 및 미니배치 크기에서 CPC를 일관되게 앞서며, 다양한 데이터 분포에 대한 일반화 능력을 입증했다.
• 결과는 WDM을 사용함으로써 상호정보량 추정의 근본적 한계인 지수적 표본 복잡도를 완화함으로써 고정보량 환경에서 더 완전한 표현을 얻을 수 있음을 확인한다.