VDVAE: Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images

https://arxiv.org/pdf/2011.10650.pdf

 

Introduction

VAE는 prior가 주어지고, latent space와 data 사이의 관계를 학습한다 확실한 목표를 가진 model이다.

그래서 이 논문의 저자는 VAE가 Autoregressive Model보다 좊은 화질의 사진들의 분포를 잘 학습할 수 있다고 주장한다. 그래서 아래와 같은 contrivution을 달성하는데

1. Layer 층 수에 따른 VAE의 성능에 대한 이론을 세우고, autoregressive model이 layer를 쌓은 VAE의 특수한 형태임을 증명
2. 안정적으로 훨씬 많은 layer를 쌓을 수 있는 방법 제공
3. 실험을 통해 capacity와 상관 없이 alyer 수를 늘리는 것이 주는 성능 향상의 영향과 모든 benchmark에서 autoregressive model을 이기는 것을 보여줌
4. layer들이 유의미한 representation을 학습하였으며 이 덕에 sampling 속도가 빨리지고 scalable 해짐을 보여줌

Hierarchical VAE

LVAE: Ladder Variational Autoencoders (2016)

• Latent variable Z를 여러개의 layer로 쪼개어 z0, z1, ... 로 표현

BIVA: A Very Deep Hierarchy of Latent Variables for Generative Modeling (2019)

• LVAE에서 더 deep 하게 확장
• hierarchical latent variable 사이에 skip-connection 적용

NVAE: A Deep Hierarchical Variational Autoencoder (2020)

• VDVAE가 나오기 전 SOTA
• BIVA + IAF

VDVAE

Encoder는 기존의 LVAE, NVAE와 비슷하다.

Generator가 FLow에 의존하지 않고 Res Block 방식으로 구현되었다.

이로 인해 속도가 더 빨라지고 더 깊은 layer를 쌓을 수 있었다.

Autoregressive Model (PixelCNN)

Autoregressive Model의 가장 큰 장점이 해상도가 매우 높은 realistic한 output을 생성해낼 수 있다는 점이다.

하지만 그럼에도 단점이 존재하는데 첫번째로는 1)시간이 오래걸리는 점이며 두번째로는 2)데이터에서 관측되는 패턴만 생성이 된다는 점이다. 

 

이 논문의 저자는 Hierarchical VAE가 Autoregressive 방식보다 우월하다고 입증한다.

여기서 중요한 점이 Layer의 깊이가 Deep해야한다는 점이다.

 

VAE가 깊어지면 왜 좋을까

VAE의 layer가 깊어지는 게 어떤 이점을 가져오는지 분석해보자.

1. Data 차원이 N일 때, N-layer VAE는 autoregressive model의 일반화이다.
2. N-layer VAE는 N차원 latent density의 universal approximator이다.

1번이 의미하는 것은 Autoregressive Model이 결국은 VAE를 깊게 쌓은 것의 한 형태라는 의미이다.

하지만 깊게 쌓는 것 자체가 어렵기 때문에 저자는 2번을 통해 data 차원가 layer 수가 같을 때까지가 상한선이고 이보다 적은 layer들도 비슷한 결과를 낼 수 있다는 것을 보여줬다.

즉, VAE의 Layer를 deep하게 쌓을수록 좋은 성능을 보일 수 있다는 것을 증명한 것이다.

 

그렇다면 어떻게 해야 VAE의 Layer를 안정적으로 Deep하게 쌓을 수 있을까?

이 부분은 앞서 보여준 것 처럼 Hierarchical VAE를 통해 찾아볼 수 있다.

하지만 Hierarchical VAE에서도 한계가 존재하는데, 깊어진 layer들이 별 의미가 없는 경우가 나타난다는 점이다.

위의 그래프에서 보이는 것처럼 그냥 layer를 쌓으면 초록색 그래프처럼 posterior collapse가 일어나 의미있는 정보를 학습하지 못하게 된다.

여기서 posterior collapse란 approximate posterior가 prior를 그대로 mimic하며, model은 latent variable을 무시한 상태에서 학습이 진행된다는 의미이다. 다시 말해서 latent variable이 유의미하게 학습되지 않는다는 소리이다. 이렇게 되면 VAE는 그저 쓰레기 모델이 되어버린다...

그래서 이를 해결하기 위해 저자는 아래와 같은 방법을 소개한다.

• residual block 마지막에 ouput을 root(N)으로 나눠줌
• upsampling을 할 때 transposed convolution 대신 nearest-neighbor upsampling을 사용
• 학습 도중 gradient norm이 매우 큰 상황이 발생할 때 gradient norm이 threshold를 넘으면 넘어가는 식으로 학습

이 결과 70 layer VAE 학습에 성공하였다.