[논문 리뷰] Vector Quantized Wasserstein Auto-Encoder

Tung-Long Vuong, Trung Le외 5명|arXiv (Cornell University)|2023. 02. 12.

Generative Adversarial Networks and Image Synthesis인용 수 5

한 줄 요약

이 논문은 이산 코드워드 시퀀스의 분포와 데이터 분포 간의 워셔스타인 거리 최소화를 통해 깊이 있는 이산 잠복 표현을 학습하는 새로운 방법인 벡터 양자화 워셔스타인 오토에코더(VQ-WAE)를 제안한다. 워셔스타인 거리의 이론적 성질을 활용하고 목적함수를 학습 가능한 형태로 재구성함으로써, VQ-VAE 및 SQ-VAE 변종에 비해 더 나은 코드북 활용도, 향상된 이미지 복원 및 생성 능력, 더 나은 분리되고 군집화된 잠복 표현을 달성한다.

ABSTRACT

Learning deep discrete latent presentations offers a promise of better symbolic and summarized abstractions that are more useful to subsequent downstream tasks. Inspired by the seminal Vector Quantized Variational Auto-Encoder (VQ-VAE), most of work in learning deep discrete representations has mainly focused on improving the original VQ-VAE form and none of them has studied learning deep discrete representations from the generative viewpoint. In this work, we study learning deep discrete representations from the generative viewpoint. Specifically, we endow discrete distributions over sequences of codewords and learn a deterministic decoder that transports the distribution over the sequences of codewords to the data distribution via minimizing a WS distance between them. We develop further theories to connect it with the clustering viewpoint of WS distance, allowing us to have a better and more controllable clustering solution. Finally, we empirically evaluate our method on several well-known benchmarks, where it achieves better qualitative and quantitative performances than the other VQ-VAE variants in terms of the codebook utilization and image reconstruction/generation.

연구 동기 및 목표

• 기존 VQ-VAE 변종의 코드북 활용도 및 표현 품질에 대한 한계를 해결하기 위해.
• 최적 운반 이론을 활용해 생성적 관점에서 깊이 있는 이산 표현을 학습하기 위해.
• 워셔스타인 거리와 벡터 양자화를 연결하는 이론적으로 탄탄한 프레임워크를 개발하여 더 나은 군집화와 분리도를 달성하기 위해.
• 코드북 내 코드워드의 더 균형 잡히고 의미 있는 사용을 보장함으로써 학습된 표현의 품질을 향상시키기 위해.
• 워셔스타인 거리 최소화가 복원, 생성 및 후속 표현 성능 향상에 기여함을 입증하기 위해.

제안 방법

• 디코더가 코드워드 시퀀스의 이산 분포를 데이터 분포로 매핑하는 생성 모델을 설정하고, 워셔스타인 거리 최소화를 통해 학습한다.
• 재구성 항목과 잠복 표현 및 코드워드 분포 간의 워셔스타인 거리를 포함하는 등가 목적함수를 유도한다.
• 워셔스타인 거리와 군집화를 연결하는 이론적 프레임워크를 도입하여 잠복 공간의 구조적 제어를 가능하게 한다.
• 잠복 공간 내 최근접 이웃 할당을 통해 정의된 벡터 양자화를 사용하는 결정론적 인코더 및 디코더를 적용한다.
• 학습을 안정화하고 백프로파게이션 동안 코드북을 고정하면서 인코더를 통해 기울기 흐름을 유지하기 위해 정지 기울기 연산을 적용한다.
• 이미지 생성을 가능하게 하기 위해 간단한 이산 잠복 시퀀스에 대한 사전 분포(예: 자기회귀 모델)를 사용한다.

Figure 1: Reconstruction results for the FFHQ dataset.

실험 결과

연구 질문

• RQ1 VAE 유사 프레임워크에서 워셔스타인 거리를 효과적으로 활용하여 이산 표현 학습을 향상시킬 수 있는가?
• RQ2 잠복 코드 분포와 데이터 분포 간의 워셔스타인 거리 최소화가 더 나은 코드북 활용도 및 표현 품질을 이끌어내는가?
• RQ3 워셔스타인 거리의 이론적 성질을 활용해 더 구조적이고 군집화된 잠복 표현을 달성할 수 있는가?
• RQ4 이미지 복원, 생성 및 후속 분류 성능 측면에서 VQ-WAE는 VQ-VAE 및 SQ-VAE에 비해 어떻게 비교되는가?
• RQ5 제안된 프레임워크는 개선된 최적화를 통해 더 다양하고 의미 있는 코드워드를 자연스럽게 유도할 수 있는가?

주요 결과

• VQ-WAE는 모든 데이터셋에서 단일 레이어 분류 정확도가 가장 높았으며, CIFAR10에서는 50.19%, MNIST에서는 95.62%, SVHN에서는 38.38%를 기록하여 VQ-VAE 및 SQ-VAE를 능가했다.
• 이미지 생성 성능에서 VQ-WAE는 가장 낮은 FID 스코어를 기록했으며, CIFAR10에서는 무조건적 기준으로 87.73, 조건부 기준으로 88.51을 기록했고, MNIST에서는 무조건적 기준으로 8.21, 조건부 기준으로 3.88를 기록하여 뛰어난 이미지 품질을 입증했다.
• VQ-WAE에서는 코드북 활용도가 크게 향상되어, 더 많은 코드워드가 활성화되고 VQ-VAE 및 SQ-VAE에 비해 死코드워드 또는 미사용 코드워드가 훨씬 적었다.
• 잠복 공간 시각화 결과 VQ-WAE 표현은 각 클래스별로 잘 분리되고 조밀한 군집을 형성한 반면, VQ-VAE 및 SQ-VAE 표현은 산만하게 흩어져 있고 클래스 간에 겹쳐져 있었다.
• 절단 실험을 통해 VQ-WAE는 코드북을 최적의 방식으로 활용하여 더 다양한 의미 있는 코드워드와 더 잘 구조화된 잠복 다각형을 생성함을 확인했다.
• 이론적 분석을 통해 워셔스타인 거리 공식화에서 최적의 결합(OP, 정리 2.3)이 코드워드 시퀀스에 대한 공동 분포 학습 경로를 시사하며, 이는 향후 생성 모델링 향상에 기여할 수 있다.

Figure 2: Latent distribution over the codebook on test-set.