[논문 리뷰] Robust Training of Vector Quantized Bottleneck Models

Adrian Łańcucki, Jan Chorowski외 7명|arXiv (Cornell University)|2020. 05. 18.

Speech Recognition and Synthesis참고 문헌 35인용 수 8

한 줄 요약

이 논문은 코드북 학습을 온라인 k-means 클러스터링으로 간주함으로써 VQ-VAE(벡터 양자화 변동형 오토인코더)의 향상된 훈련 기법을 제안한다. 코드북 학습률을 증가시키고, 양자화 이전에 배치 정규화를 적용하며, 주기적으로 데이터에 의존하는 코드북 재초기화를 수행함으로써 더 안정적인 훈련, 더 넓은 코드북 활용도, 그리고 음성, 이미지, 단위 탐지 작업 전반에서 향상된 성능을 달성한다.

ABSTRACT

In this paper we demonstrate methods for reliable and efficient training of discrete representation using Vector-Quantized Variational Auto-Encoder models (VQ-VAEs). Discrete latent variable models have been shown to learn nontrivial representations of speech, applicable to unsupervised voice conversion and reaching state-of-the-art performance on unit discovery tasks. For unsupervised representation learning, they became viable alternatives to continuous latent variable models such as the Variational Auto-Encoder (VAE). However, training deep discrete variable models is challenging, due to the inherent non-differentiability of the discretization operation. In this paper we focus on VQ-VAE, a state-of-the-art discrete bottleneck model shown to perform on par with its continuous counterparts. It quantizes encoder outputs with on-line $k$-means clustering. We show that the codebook learning can suffer from poor initialization and non-stationarity of clustered encoder outputs. We demonstrate that these can be successfully overcome by increasing the learning rate for the codebook and periodic date-dependent codeword re-initialization. As a result, we achieve more robust training across different tasks, and significantly increase the usage of latent codewords even for large codebooks. This has practical benefit, for instance, in unsupervised representation learning, where large codebooks may lead to disentanglement of latent representations.

연구 동기 및 목표

• 비정상적인 인코더 출력과 민감한 코드북 초기화로 인해 발생하는 VQ-VAE 훈련의 불안정성과 수렴 불량 문제를 해결한다.
• 대규모 VQ-VAE 모델에서 코드북 활용도를 향상시켜 더 나은 분리성과 표현 품질을 가능하게 한다.
• 코드북 학습을 온라인 k-means 클러스터링으로 간주함으로써 깊이 있는 이산 블로킹 모델의 훈련 동역학을 안정화시킨다.
• 이미지 생성, 음성 인식, 비지도 단위 탐지와 같은 다양한 작업에서 아키텍처 및 최적화 수정 사항의 영향을 평가한다.
• 다양한 훈련 향상 기법을 조합할 경우 개별 수정 사항보다 훨씬 우수하고 일관된 성능을 달성할 수 있음을 입증한다.

제안 방법

• 코드북 학습을 온라인 k-means 클러스터링으로 간주하고, 알려진 k-means 최적화 전략을 활용한다.
• 코드북 학습률을 인코더보다 높게 설정하여 수렴 속도를 가속화하고 코드북 안정성을 향상시킨다.
• 양자화 레이어 이전에 배치 정규화를 적용하여 잠재 특징의 크기를 균형 잡히게 하고 기울기 흐름을 향상시킨다.
• 최근의 인코더 출력을 오프라인 k-means를 사용해 클러스터링하여 주기적으로 데이터에 의존하는 코드북 재초기화를 수행한다.
• 직접 전파를 위해 스트레이트스러우트 추정기를 사용하여 이산 양자화 연산을 통해 전파를 유지하면서 전체적으로 미분 가능성을 확보한다.
• 증가된 코드북 학습률, 배치 정규화, 재초기화를 모두 조합하여 최적의 훈련 성능을 달성한다.

실험 결과

연구 질문

• RQ1 코드북 초기화 스케일은 VQ-VAE 훈련의 안정성과 수렴에 어떤 영향을 미치는가?
• RQ2 주기적인 데이터에 의존하는 코드북 재초기화는 훈련 중 인코더 활성화의 비정상성 문제를 어느 정도 완화하는가?
• RQ3 코드북과 인코더에 별도의 학습률를 설정할 경우 훈련 동역학과 코드북 활용도에 어떤 영향을 미치는가?
• RQ4 양자화 이전의 배치 정규화는 표현 품질과 훈련 안정성에 어떤 영향을 미치는가?
• RQ5 이중 훈련 향상 기법의 조합은 이미지, 음성, 단위 탐지 작업에서의 최종 성능에 어떤 영향을 미치는가?

주요 결과

• 주기적인 데이터에 의존하는 코드북 재초기화는 특히 다른 향상 기법과 조합되었을 때 수렴 속도를 크게 향상시키고 초기 BPD 값을 감소시킨다.
• 코드북 학습률을 증가시키면 수렴 속도가 빨라지고 훨씬 더 안정적인 훈련이 가능해지며, 특히 큰 코드북일 경우 두드러진다.
• 양자화 이전의 배치 정규화는 CIFAR-10의 모든 설정에서 BPD(비트/차원)에 소량이지만 일관된 향상을 가져온다.
• 증가된 코드북 학습률, 배치 정규화, 주기적 재초기화의 조합은 모든 평가된 작업에서 최고의 성능과 가장 강력한 훈련 안정성을 달성한다.
• 코드북 재초기화는 코드북 붕괴를 방지하고, 큰 코드북일지라도 잠재 코드워드의 사용도를 높여 표현의 분리성을 향상시킨다.
• 제안된 훈련 전략은 음성 인식과 단위 탐지와 같은 지도 및 비지도 작업에서 최종 성능을 향상시키며, BPD와 클러스터링 지표에서 일관된 성과 향상을 이룬다.