티스토리 뷰
[요약] Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images
4stack 2021. 12. 13. 01:40논문 링크: https://arxiv.org/pdf/2011.10650.pdf
서론
Generative 모델의 일종인 Variational AutoEncoder (VAE)는 주어진 data들(e.g. 얼굴 이미지들)과 prior가 주어진 latent variable들 사이의 관계를 학습함으로써 data들의 분포를 학습한다. Prior가 주어졌기도 하고, latent space와 data 사이의 관계성을 학습한다는 목표가 확실하기 때문에 저자들은 VAE가 autoregressive model과 같이 깡으로 주어진 data들의 분포를 학습해야 하는 모델들과 비교했을 때, capacity가 작은 모델로도 충분히 높은 화질의 사진들의 분포를 학습할 수 있다고 주장한다. 하지만 이전까지의 결과를 보면 PixelCNN과 같은 autoregressive model들이 VAE보다 나은 면모를 보여줬다. 그래서 저자들은 이에 의문을 제기하였고, Very Deep VAE를 통해 다음과 같은 contribution들을 달성한다.
- Layer 수에 따른 VAE의 성능에 대한 이론적 기반을 제공하고, autoregressive model이 layer를 쌓은 VAE의 특수한 형태임을 증명한다.
- 기존의 VAE들과 비교하여 훨씬 많은 layer를 안정적으로 쌓을 수 있는 학습 방법들을 제공한다.
- 실제 실험 결과를 통해 capacity와 관계없이 layer 수를 늘리는 것이 성능을 향상시키고, 모든 benchmark에서 autoregressive model을 이기는 것을 보여준다.
- 쌓았던 layer들이 실제로 유의미한 representation을 학습하였고, 이를 통해 sampling 속도가 빨라지고 scalable 해진다는 것을 보여준다.
추천 읽을거리
LVAE를 읽어보면 notation들이나 이후의 내용들을 이해하는데 도움이 될 것이다. Autoregressive model 자체에 관심이 있다면 PixelCNN를 읽어보는 것을 추천한다.
본론
VAE의 layer를 더 쌓으면 좋을까?
먼저 VAE의 layer 수를 쌓는 것이 어떤 이점을 가져오는지에 대한 분석이 필요하다. 저자들은 이와 관련하여 다음과 같은 proposition들을 증명하였다.
- Data 차원이 N일 때, N-layer VAE는 autoregressive model의 일반화이다.
- N-layer VAE는 N 차원 latent densities의 universal approximator이다.
첫 번째 proposition은 말하고자 하는 바가 명확하다. 결국 autogressive model은 여러 층을 쌓은 VAE의 특수한 형태이다. 하지만 1024x1024 사진을 생각해보면, data 차원만큼 layer를 쌓는 것은 생각보다 힘든 일이다. 그래서 저자들은 두 번째 proposition을 통해 data 차원과 layer 수가 동일할 때가 상한선이고, 그보다 적은 layer들로도 비슷한 결과를 낼 수 있다는 가능성을 보여줬다. 물론 하한선을 찾는 것은 또 다른 문제지만, 결국 layer를 쌓으면 쌓을수록 VAE의 성능이 올라간다는 것을 이론적으로 보여줬다고 볼 수 있다.
어떻게 더 많은 layer들을 안정적으로 쌓을까?
사실 VAE의 layer 수를 늘리는 것과 관련된 연구들은 많이 진행되고 있었다. 보통 Hierarchical VAE라고 지칭한다. Hierarchical VAE에 있어 한 가지 challenge는, 많이 쌓는 것은 문제 없지만 posterior collapse와 같은 문제들 때문에 data와 멀리 있는 layer들은 없느니만 못한 경우들이 자주 나타난다. 아래의 figure가 이를 보여준다. 단순히 layer들을 쌓은 VAE (초록색)의 경우, 깊은 layer들에서 posterior collapse가 나타나 유의미한 정보를 학습하지 못한 것을 정량적으로 확인할 수 있다.
저자들은 이 문제를 해결하기 위해 다음과 같은 모델 설계 방식 또는 학습 방법을 소개했다.
- 각 residual block의 마지막 convolution layer의 output을 \(\sqrt{N}\)으로 나눠줌
- upsampling 할 때 흔히 쓰던 transposed convolution 대신, nearest-neighbor upsampling을 사용
- 학습 도중 gradient norm이 매우 큰 상황이 발생하는데, gradient norm이 threshold를 넘는 경우 넘어가는 식으로 학습
각각의 장점 및 효과들은 너무 specific 해서, 궁금하면 논문에서 찾아보는 것을 권장한다. 아무튼 위와 같은 방식들을 도입하여 기존의 최대 깊이였던 30-layer의 약 2.5배에 달하는 70-layer VAE를 학습하는 데 성공하였다.
실험 결과
실험 단락에서는 1) capacity가 아닌, layer의 수를 늘렸을 때와 일반적인 benchmark에서의 성능 그리고 3) 각 layer에서 학습한 representation을 보여준다. 먼저 layer 수에 따른 성능은 다음과 같다.
결과적으로 capacity (param 수)를 고정하고 layer 수 (depth)만 늘렸음에도 불구하고, test loss가 줄어든 것을 확인할 수 있다. 이를 통해 layer 수에 따른 성능 향상이 있음을 실험적으로 알 수 있다. 그리고 일반적인 benchmark에서도 autoregressive model과 비교하여 작은 capacity를 가지고도 훨씬 나은 성능을 보여줬다. 자세한 수치는 논문을 참고하기를 바란다. 마지막으로 각 layer에서 학습한 representation는 다음과 figure를 통해 설명한다.
앞서 설명했듯이, Hierarchical VAE에서 나타나는 문제 중 하나는 깊은 layer들이 아무런 의미도 없는 representation을 학습한다는 것이다. 그리고 저자들이 제안한 VAE의 경우, upsampling을 진행하기 때문에 깊은 layer (figure 기준 왼쪽)와 얕은 layer (figure 기준 오른쪽) 각각을 구성하는 latent variable들의 수에 큰 차이가 나 (figure 기준 아래의 percentage), 얕은 layer에 유의미한 representation이 몰릴 가능성이 상당히 높다.
보통 이를 직관적으로 확인하는 방법은, 주어진 사진을 encoder에 넣어서 나온 latent variable들 중 깊은 layer에 해당하는 것들만 바꿨을 때 decoder를 통해 reconstruct 되는 이미지를 확인하는 것이다. 만약 아무런 의미도 없는 representation을 학습했다면 거의 변화가 없을 것이다. 하지만 위의 figure를 보면, 깊은 layer에서도 원본 이미지와 비교하여 reconstruct 된 이미지가 달라지는 모습을 확인할 수 있다. 실험적으로 모든 layer들이 유의미한 representation을 학습했음을 밝혔다.
이외에도 유의미한 representation을 활용하여 autoregressive과 비교하여 생성 속도나 1024x1024 사진과 같이 높은 차원의 data를 학습하는데 있어 이점을 보인다고 설명한다. 자세한 것은 논문을 참고하기를 바란다.
이론적으로나 실험적으로나 우수한 성과를 보여줬다. 단순히 layer를 쌓아서 잘됐더라 이상으로 ablation study나 theoretical justification이 잘된 논문이었다. Autoregressive model에도 빌빌대던 VAE에 새로운 가능성을 열어준 논문이라고 생각한다. 그리고 스토리 라인이나 논문의 질이 높다 생각하기에, VAE나 generative model에 관심이 있는 사람들이라면 꼭 읽어봤으면 하는 논문이다.