VAE(Variational AutoEncoder)란 무엇인가?
VAE는 현재 딥러닝 분야에서 매우 활발하게 연구되고 있는 deep generative model 중 하나로, 다른 대표적인 모델로 GAN(Generative Adversarial Network)이 있다. GAN은 랜덤한 노이즈로부터 새로운 샘플 데이터를 생성해내는 과정에 집중하기에 데이터의 분포 자체에는 관심이 없는 반면, VAE는 데이터의 특징 및 분포를 학습하고 이를 잠재 변수(latent variable)로 인코딩하는 과정을 포함한다는 차이점이 있다.
VAE는 원본 데이터에 담겨있는 특징들을 잘 표현하는 잠재 변수의 분포를 찾는 인코더 네트워크와, 잠재 변수 분포에서 샘플링한 값으로부터 원본과 유사한 이미지를 생성해내는 디코더 네트워크로 구분된다. 만약 전체 네트워크 학습이 잘 이루어진다면 샘플링한 잠재 변수로부터 디코딩을 거쳐 새로운 데이터를 생성할 수 있다. VAE의 근간을 이루는 수학적 백그라운드를 알고 싶다면 이 링크를 참고하자.
개요
고해상도의 이미지를 얻기 위해 낮은 해상도의 이미지를 생성하는 네트워크에서 시작하여 점차 해상도를 높여나가는 점진적인(progressive) 학습 기법을 GAN에 적용하는 연구는 잘 알려져 있지만, 잠재 변수 관점에서 점진적인 학습을 접목한 연구는 비교적 덜 조명받았다. 물론 VAE에서의 점진적(또는 순차적) 학습을 다룬 선행 연구들도 있으나[2][3], 이들 연구는 대부분 VAE로부터 더 높은 해상도의 이미지를 생성하기 위해 네트워크 용량을 점차 키워나가는 접근법을 다루고 있다.
반면, 이 논문에서는 원본 데이터에 담긴 계층적인 표현들을 추출 및 구분(disentangle)하는 데 초점을 맞춘다. 계층적인 표현들을 추출하기 위해 네트워크 깊이에 따라 VAE 네트워크의 용량을 증가시키며, 이렇게 얻어진 계층적 표현들은 서로 다른 추상화 레벨로써 이미지 생성에 사용된다. 뿐만 아니라, 잠재 변수의 형태로 얻어진 추상적 표현들이 서로 섞이지 않고 잘 구분되도록 점진적인 학습을 도입하였다. 그 결과 학습 초기에는 원본 데이터의 가장 추상적인 표현부터 학습하며, 모델의 깊이가 깊어질수록 더 낮은 레벨의 추상적인 표현까지 학습한다. 이러한 방법론을 적용했을 때 각각의 계층적인 표현들이 서로 잘 구분됨을 여러 메트릭들을 기준으로 실험적으로 비교하였다.
pro-VLAE 구조
본 논문에서 제안하는 pro-VLAE는 일반적인 VAE와 크게 두 가지 차이점이 존재한다. 하나는 VAE에 계층적 구조를 접목시켰다는 것이고, 다른 하나는 이렇게 만들어진 VAE를 높은 추상화 레벨부터 시작하여 점진적으로 학습을 수행한다는 것이다.
- 계층적 구조 도입
기존의 VAE가 단 하나의 인코더와 디코더 네트워크 쌍을 가지고 있는 반면, pro-VLAE에서는 여러 추상화 레벨을 갖는 계층적 구조를 표현하기 위해 여러 개의 인코더 및 디코더 네트워크 블록을 쌓는다. 이로 인해 잠재 변수 z 또한 서로 다른 추상화 레벨을 갖는 z = {z1, x2, …, zL}로 분해된다. 이러한 계층적 구조는 이전에 제시된 VLAE[4]와 유사한데, 논문 저자들이 새로 제안한 모델을 pro-VLAE이라 명명한 것 또한 이런 이유에서다.
2. 계층적 구조 VAE의 점진적 학습
이렇게 정의된 계층적 구조 VAE는 학습 단계에서 한 번에 모든 인코더 및 디코더 네트워크를 학습시키지 않는다. 대신 가장 높은 추상화 레벨(l = L)에서 시작하여, 최종적으로 가장 낮은 추상화 레벨(l = 1)에 이르기까지 잠재 변수 z_l을 점진적으로 학습시킨다.
어떻게 학습시킬 것인가
안타깝게도, 새로 설계한 pro-VLAE 구조만으로는 학습이 성공적으로 이루어지지 않았다고 한다. 그렇기에 논문 저자들은 pro-VLAE 학습에 적합한 두 가지 새로운 전략을 제시 및 적용했다.
1. fade-in 기법
첫 번째 전략은 매 학습 스텝마다 새로운 레벨의 컴포넌트가 추가될 때, 이를 바로 적용하지 않고 점차 가중치를 늘려나가는 ‘fade-in’ 기법이다. 기존에 학습되던 네트워크에 새로운 컴포넌트를 직접적으로 추가하면 그래디언트에 갑작스런 변화가 야기되는데, 이는 VAE에서 잠재 변수 분포의 분산이 폭발(explode)하는 등 학습에 악영향을 끼친다. 이를 방지하기 위해 fade-in 계수 α를 도입하여, 인코더와 디코더의 새로운 컴포넌트를 추가할 때 처음에는 α=0의 가중치에서 시작하여 α=1까지 점차 증가시킨다.
2. 새로운 KL penalty 추가
두 번째는 기존의 목적함수에 새로운 항을 도입하는 것이다. pro-VLAE의 목적함수 L은 기존 VAE의 목적함수 L_pro에 새로 정의된 목적함수 L_pre-trained을 합친 구조로 이루어져 있다.
L_pro는 기존에 제시된 β-VAE[5]의 목적함수와 유사한 구조다. L_pro는 다시 Reconstruction loss와 KL divergence 항으로 나뉘며, 일반 ELBO(Evidence Lower Bound)에서 두 분포가 서로 다른 정도를 표현하는 KL divergence에 β만큼의 가중치를 준다. VAE에서 ELBO와 같은 목적함수가 어떠한 수학적 의미를 갖는지 궁금하다면 이 링크를 참고하자. 이때 s번째 점진적 스텝에서 학습 네트워크에 포함된 잠재 변수 분포 {z_L-s, …, z_L}만이 목적함수 계산에 사용됨을 유의하자.
pro-VLAE의 목적함수에는 기존 목적함수에 L_pre-trained 항이 추가된다. 이는 점진적 학습의 s번째 스텝에서 아직 네트워크에 포함되지 않은 잠재 변수 분포 {z_1, z_2, …, z_L-s-1}에 대해서도 KL penalty를 부여한다는 의미다. 이때 γ라는 계수는 0.5라는 작은 값으로 설정하여, 낮은 레벨의 잠재 변수가 네트워크 학습 과정에 추가되기 전부터 일종의 정규화(regularization) 역할을 수행할 수 있도록 돕는다.
성능 평가
저자들은 dSprites (Matthey et al. (2017)), 3DShapes (Burgess & Kim (2018)), MNIST (LeCun et al. (1998)), CelebA (Liu et al. (2015)) 등 4개의 벤치마크 데이터셋으로 pro-VLAE의 성능을 평가했다. 아래 그림은 서로 다른 4개의 모델(β-VAE, VLAE, teacher-student VAE, pro-VLAE)에 대하여 데이터 특성을 얼마나 잘 구분해내는가를 평가한 것으로, x축은 reconstruction factor를, y축은 왼쪽부터 각각 factor metric/MIG/MIG-sup을 나타낸다. 값이 높을수록 데이터의 특성을 각 계층별로 잘 구분해냄을 의미하며, pro-VLAE가 다른 모델 대비 높은 성능을 보임을 확인할 수 있다.
벤치마크로 사용된 3DShapes는 총 6가지 요소를 포함하고 있다. 아래 그림에서, pro-VLAE는 추상화 레벨이 가장 높은 z3는 물체 색깔/벽 색깔/바닥 색깔의 3요소를, z2는 물체의 크기/방향/모양의 3요소를 각 계층별로 잘 분리하여 표현한다. 특히 z1에서는 더 이상 구분해낼 요소가 없기에 어떠한 특성도 표현하지 않음을 알 수 있다. 점진적인 학습이 적용되지 않는 VLAE에서는 총 9만큼의 잠재 차원에 걸쳐 일부 요소들이 서로 섞여 표현됨을 확인할 수 있으며, 세 종류의 메트릭에 대해서도 pro-VLAE가 더 높은 값을 가진다.
아래 그림은 CelebA 데이터셋을 대상으로 각 계층의 특성을 시각화한 것이다. 가장 높은 추상화 레벨의 잠재 변수는 성별과 인종 등의 특성들을 잘 구분해내며, 가장 낮은 레벨에서는 아이섀도우 여부처럼 보다 덜 추상적인, 즉 구체적인 정보를 내재하고 있음을 확인할 수 있다.
결론
본 포스팅에서는 VAE를 통해 원본 데이터에 담긴 여러 추상화 레벨의 특성들을 보다 잘 추출 및 구분하는 연구를 살펴보았다. pro-VLAE는 일반적인 VAE에 계층적 표현 방법론과 점진적 학습 과정을 접목하였으며, 추가로 fade-in 전략과 pre-trained KL penalty를 도입하여 성능을 개선하였다. 그렇게 함으로써 서로 다른 추상화 레벨에서의 특성들이 얽히지 않고 잘 구분됨을 실험적으로 보였다.
이처럼 데이터의 특성들이 서로 겹치지 않도록 구분해낼 수 있다면, 주어진 잠재 공간 중 데이터의 고유한 특성들을 모두 표현해낼 수 있는 최소한의 잠재 공간만이 사용되며 나머지는 쓰이지 않을 것이다. 이로부터 원본 데이터의 특성을 모두 포함하는 샘플 데이터를 생성하는 데 필요한 최소한의 잠재 공간 크기를 알 수 있고, 이를 활용한 VAE 네트워크 용량의 최적화 또한 가능할 것이라 생각된다. 논문 저자들이 소스코드를 깃허브에 공개해 놓았으니, pro-VLAE 구현체에 관심이 있는 독자라면 직접 소스코드를 확인해봐도 좋을 것이다.
레퍼런스
[1] https://arxiv.org/pdf/2002.10549
[2] Gregor, Karol, et al. “Draw: A recurrent neural network for image generation.” arXiv preprint arXiv:1502.04623 (2015).
[3] Lezama, José. “Overcoming the Disentanglement vs Reconstruction Trade-off via Jacobian Supervision.” ICLR (Poster). 2019.
[4] Zhao, Shengjia, Jiaming Song, and Stefano Ermon. “Learning hierarchical features from deep generative models.” International Conference on Machine Learning. 2017.
[5] Higgins, Irina, et al. “beta-vae: Learning basic visual concepts with a constrained variational framework.” (2016).
