[논문 리뷰] SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization

Yuhta Takida, Takashi Shibuya외 8명|arXiv (Cornell University)|2022. 05. 16.

Domain Adaptation and Few-Shot Learning인용 수 13

한 줄 요약

이 논문은 정지 기울기나 코드북 재설정과 같은 히우리스틱에 의존하지 않고 코드북 활용도를 향상시키는 스토하스틱 양자화와 자가 자동 조절 학습을 갖춘 변분 오토인코더인 SQ-VAE를 제안한다. 훈련 가능한 스토하스틱 복원 및 양자화를 도입함으로써 SQ-VAE는 훈련 과정에서 자연스럽게 결정론적 양자화로 수렴하게 되어, 시각 및 음성 작업에서 VQ-VAE와 표준 VAE보다 더 나은 재구성 및 생성 성능을 달성한다.

ABSTRACT

One noted issue of vector-quantized variational autoencoder (VQ-VAE) is that the learned discrete representation uses only a fraction of the full capacity of the codebook, also known as codebook collapse. We hypothesize that the training scheme of VQ-VAE, which involves some carefully designed heuristics, underlies this issue. In this paper, we propose a new training scheme that extends the standard VAE via novel stochastic dequantization and quantization, called stochastically quantized variational autoencoder (SQ-VAE). In SQ-VAE, we observe a trend that the quantization is stochastic at the initial stage of the training but gradually converges toward a deterministic quantization, which we call self-annealing. Our experiments show that SQ-VAE improves codebook utilization without using common heuristics. Furthermore, we empirically show that SQ-VAE is superior to VAE and VQ-VAE in vision- and speech-related tasks.

연구 동기 및 목표

• 정지 기울기나 코드북 재설정과 같은 히우리스틱에 의존하지 않고도 코드북의 일부 벡터만 사용되는 VQ-VAE의 코드북 붕괴 문제를 해결한다.
• 정지 기울기, EMA, 또는 코드북 재설정에 의존하지 않는 표준 변분 베이즈 프레임워크 내에서 작동하는 학습 프레임워크를 개발한다.
• 스토하스틱성의 감소를 통해 시간이 지남에 따라 자가 자동 조절 과정을 통해 더 큰 코드북 활용도를 달성한다.
• 범주형 데이터에 적합한 변형을 포함하여 다양한 데이터 유형으로 일반화할 수 있는 방법을 설계한다. 특히, 베르누이-미제르 분포를 사용한 특화된 변형을 제공한다.
• 초기 설정 조정 없이도 보조 기법 없이도 시각 및 음성 벤치마크에서 재구성 및 생성 성능이 뛰어나다.

제안 방법

• 연속 잠재 변수를 코드북 벡터로 매핑하는 훈련 가능한 확률 분포를 사용하는 스토하스틱 복원 과정을 도입한다.
• 코드북 벡터 위의 훈련 가능한 후행 범주형 분포를 사용한 스토하스틱 양자화를 적용함으로써 이산 선택을 통한 기울기 흐름을 가능하게 한다.
• 스토하스틱성의 감소가 훈련 과정에서 점차 줄어들어 자연스럽게 결정론적 양자화로 수렴하는 자가 자동 조절 메커니즘을 사용한다.
• 표준 역전파를 통해 최적화할 수 있는 증거 하한 경계(ELBO)로 목적 함수를 제안하며, 직선 추정 기법이 필요 없다.
• 두 가지 변형을 제안한다: 일반 데이터를 위한 가우시안 SQ-VAE와 범주형 데이터를 위한 vMF SQ-VAE이며, 후자는 이산 분포에서 성능 향상을 위해 베르누이-미제르 분포를 사용한다.
• 스토하스틱성의 온도를 유일한 핵심 초기 설정으로 사용하여, Gumbel-Softmax 또는 Concrete 근사 기법을 따르며 모델을 훈련시킨다.

실험 결과

연구 질문

• RQ1 VAE 프레임워크 내에서 스토하스틱 양자화가 결정론적 VQ-VAE에 비해 코드북 활용도를 향상시키는가?
• RQ2 SQ-VAE의 자가 자동 조절 과정이 히우리스틱 개입 없이 더 나은 표현 학습과 재구성 성능을 이끌어내는가?
• RQ3 SQ-VAE는 표준 VAE 및 VQ-VAE에 비해 시각 및 음성 작업에서 어떻게 성능을 내는가?
• RQ4 vMF SQ-VAE는 이미지 픽셀과 같은 이산 분포를 가진 범주형 데이터에서 가우시안 SQ-VAE를 능가하는가?
• RQ5 코드북 크기를 늘릴 경우, VQ-VAE에서 자주 발생하는 붕괴 현상과는 달리 SQ-VAE에서는 성능 향상이 이루어지는가?

주요 결과

• SQ-VAE는 정지 기울기, EMA, 코드북 재설정을 사용하지 않아도 VQ-VAE보다 높은 코드북 활용도를 달성하여 코드북의 효과적인 탐색이 이루어졌음을 시사한다.
• SQ-VAE의 자가 자동 조절 과정은 양자화의 스토하스틱성을 점차 감소시켜 안정적인 수렴과 향상된 표현 학습을 이끌어낸다.
• MNIST와 Gray-CelebA에서 SQ-VAE는 VAE 및 VQ-VAE보다 더 낮은 재구성 손실을 기록하며 더 나은 재구성 성능를 보였다.
• CelebA-HQ 256×256에서 SQ-VAE는 고해상도 재구성 및 생성 샘플을 생성하여 대규모 데이터셋에 대한 확장성을 입증했다.
• vMF SQ-VAE는 가우시안 SQ-VAE보다 범주형 데이터에서 뚜렷한 성능 향상을 보였으며, 이는 이산 분포에 대해 베르누이-미제르 분포의 효과성을 입증한다.
• SQ-VAE에서는 코드북 크기를 늘릴수록 성능 향상이 이루어지지만, VQ-VAE에서는 더 큰 코드북이 자주 붕괴를 유도하고 성능 향상이 없어지는 것과는 대조된다.