직관적으로 데이터 분포는 latent space 내에서 컴팩트한 볼륨을 구성하고 무한대로 확장되지 않아야 함
하지만, Vanilla Autoencoder의 latent space는 컴팩트하지 않으며, VAE는 이를 latent space 상에서의 확률적 prior를 적용하여 개선
Derivation
먼저, latent space의 사후 확률(posterior probability)을 p(z∣x)로 정의하면, Bayes Theorem에 따라 사후 확률은 아래와 같이 표현할 수 있다.
p(z∣x)=p(x)p(x∣z)p(z)(1)
하지만 p(x)는 계산하기 매우 어렵기 때문에(p(x)=∫p(x∣z)p(z)dz 로 무수히 많은 z에 대한 적분을 계샨해야 함) 계산이 간단한 특정 확률분포(예: Gaussian) q(z)를 근사하는 방법(i.e., Variational Inference) 적용할 수 있다.
실제 사후 확률분포와 근사된 확률분포간의 KLD(Kullback-Leibler divergence)를 계산하면 아래와 같이 전개된다.
Jensen's inequality에 의하여 DKL(q(z)∣∣p(z∣x))≥0 이므로, logp(x)는 아래와 같은 lower bound를 가진다. Bayes Theorem의 정의에서 p(x)는 Evidence이기에 이를 Evidence Lower Bound(ELBO)라고도 한다.
Loss function의 첫번째 항은 reconstruction loss로 현재 샘플 z에 대한 negative log likelihood이며, 두번째 항은 KLD regularizer로 샘플링된 z에 대한 제약 조건(q(z∣x)가 p(z)와 유사해야 함)을 부여함으로써, latent space 내에서 컴팩트한 볼륨을 구성할 수 있게 한다.
q(z∣x)는 probabilistic encoder라고도 하며, p(x∣z)는 probabilistic decoder(Likelihood of generating true data sample given the latent code)라고도 한다.
KLD regularizer를 Entropy 항으로 분해해서 살펴보면, Posterior에서 샘플링된 z는 최대한 다양해야 한다는 제약 조건이 같이 포함되기에, GAN의 고질적인 문제인 mode collapse를 방지하는 효과가 있다.
하지만, KLD regularizer로 인한 posterior collapse 문제가 발생한다.
하지만, z는 샘플링된 fixed 데이터로 θ에 대한 z의 미분을 직접적으로 계산할 수 없기에, 노이즈(ϵ)를 샘플링하는 문제로 변경하여 평균과 분산에 대한 미분을 계산하는 reparameterization trick을 적용한다. (⊙는 element-wise product)
z=μ+σ⊙ϵ, where ϵ∼N(0,I)(9)
Posterior collapse
VAE의 고질적인 문제인 블러한 이미지는 복잡한 데이터 분포에서 posterior 분포가 prior 분포와 일치함으로써, latent varaible을 무시하고 학습이 진행되는 posterior collapse임.
input x에서 posterior로 흐르는 신호가 너무 미약하거나 노이즈가 많을 때, 디코더는 posterior q(z∣x)에서 샘플링되는 latent variable z를 무시하고 학습이 진행됨.
qϕ(z∣x)=pθ(z∣x)=p(z)(10)
VQ-VAE (Vector Quantized VAE)
Intuition
VAE는 continuous latent space를 학습하는 반면, VQ(Vector Quanzited; 벡터 양자화)-VAE는 discrete latent space를 학습
일반적으로 실제 세계에서 접하는 많은 데이터는 discrete representation임. (이미지는 여러 개의 discrete object로 표현되며, 언어는 discrete phonemes로 표현됨)
VQ-VAE는 K차원의 벡터를 유한한 코드 벡터 셋으로 매핑하는 방법으로 KNN과 매우 유사함.
VQ-VAE는 별도의 prior 분포를 가정하지 않고 모든 latent code에 대해 uniform한 prior를 가정함으로써, posterior가 stochastic이 아닌 deterministic임.
Derivation
e를 latent embedding space(codebook)으로 아래와 같이 정의한다. (K: the number of latent variable categories, D: embedding size)
e∈RK×D,i=1,…,K(11)
우리의 목적은 k개의 임베딩 벡터 중 Encoder output E(x)와 가장 유사한 벡터(codebook vector)를 찾고 이를 디코더 D(⋅)의 입력값으로 넣어 x를 복원하는 것으로 아래와 같이 정의할 수 있다.
zq(x)=Codebook(E(x))=ek where k=argimin∥E(x)−ei∥2(12)
Encoder output E(x)=ze(x)는 input vector x에 대한 encoder vector
k: nearest neighborhood search를 통해 K개의 임베딩 벡터 중 하나와 일치하는 코드 백터
ek: codebook에서 input과 가장 가까운 code vector
Latent vector들의 set인 codebook을 학습하기 위하여 VQ-VAE의 loss function은 3가지 loss function의 합으로 이루어져 있다.
sg(⋅): stop gradient operator로, zero gradient를 가지기에 sg(.)안의 파라메터는 업데이트되지 않는다.
VQ loss(aka. codebook loss): Encoder output과 embedding space간의 L2 에러로 input에 대응하는 codebook만을 업데이트하기 위한 용도이기 때문에, encoder output에 대한 stop gradient operator가 필요하다.
Commitment loss: VQ loss와 유사하지만 encoder의 weight를 고정함으로써 codebook을 학습하는 과정에서 동일한 신호가 다른 codebook vector로 매핑되는 경우를 줄여주는 역할을 한다.
VQ-VAE는 별도의 prior 분포를 가정하지 않고 모든 latent code에 대해 uniform한 prior를 가정함으로써, posterior가 deterministic이기에 KLD regularizer의 영향을 받지 않는다. 이를 수식으로 표현하면 다음과 같다.
q(z=ek∣x)={10for k = argmini∣∣ze(x)−ei∣∣2otherwise(14)
VQ-VAE의 학습 과정을 그림으로 표현하면 아래와 같다.
dVAE(discrete VAE)
DALL-E에서 제안된 방법으로 VQ-VAE와 유사하자만, dVAE는 벡터 양자화의 방법처럼 하나의 codebook 벡터만 선택하게 강제하지 않고 posterior 확률분포를 stochastic하게 표현할 수 있게 변경.
Discrete 샘플링을 continuous approximation 문제로 이완시키기에, loss function은 VAE의 형태와 동일
이를 위해, Encoder output은 codebook 벡터셋에 대한 범주형(categorical) 분포를 출력
범주형 분포에서 곧바로 backpropagation이 불가능하므로, 이를 Gumbel-max trick과 Gumbel-softmax Trick을 사용한 categorical reparameterization으로 해결
Derivation
g1,g2,…,gk이 Standard Gumbel 분포를 따르고 i.i.d.라고 할 때, (gk∼Gumbel(0,1)) 클래스 확률이 q(ei∣x) 인 카테고리 분포에서의 샘플 z를 다음과 같이 나타낼 수 있다.
z=codebook(argmaxi[gi+logq(ei)])(15)
(codebook(i)는 codebook 벡터셋에서 i번째 벡터를 의미하며, one-hot과 동일한 의미이다.)
이를 Gumbel-max trick이라고 하며, argmax를 softmax로 근사하여 미분 가능하게 변경한 방법이 Gumbel-softmax trick이다.
TensorFlow image preprocessing code for training the transformer (target_res = 256, channel_count = 3)
def preprocess_image(img, target_res):
h, w = tf.shape(img)[0], tf.shape(img)[1]
s_min = tf.minimum(h, w)
off_h = tf.random.uniform([], 3 * (h - s_min) // 8,
tf.maximum(3 * (h - s_min) // 8 + 1, 5 * (h - s_min) // 8),
dtype=tf.int32)
off_w = tf.random.uniform([], 3 * (w - s_min) // 8,
tf.maximum(3 * (w - s_min) // 8 + 1, 5 * (w - s_min) // 8),
dtype=tf.int32)
# Random full square crop.
img = tf.image.crop_to_bounding_box(img, off_h, off_w, s_min, s_min)
t_max = tf.minimum(s_min, round(9 / 8 * target_res))
t = tf.random.uniform([], target_res, t_max + 1, dtype=tf.int32)
img = tf.image.resize_images(img, [t, t], method=tf.image.ResizeMethod.AREA,
align_corners=True)
img = tf.cast(tf.rint(tf.clip_by_value(img, 0, 255)), tf.uint8)
# We don’t use hflip aug since the image may contain text.
return tf.image.random_crop(img, 2 * [target_res] + [channel_count])
Mixed-Precision and Distributed Training
FP32 적용시 약 50GB의 메모리가 필요하기 16GB 메모리의 V100 GPU에서는 학습 불가
FP16을 적용해도 약 24GB의 메모리가 필요하므로 Model Parallelism 필수
일반적인 mixed-precision 적용 시에는 underflow 현상으로 학습이 잘 되지 않아 여러 가지 트릭들을 사용함
Per-resblock gradient scaling: Standard loss scaling 대신 Residual Block마다 gradient scaling을 사용
16-bit precision은 반드시 필요한 경우에만 사용하며 gain, bias, embedding, unembedding, gradient, Adam momemt는 32-bit precision을 사용하고 gradient 압축을 수행하지 않음.
Data-parallel 학습 시 worker 개수가 많아질수록 underflow 빈도가 높아지기에, loss를 전체 배치 크기로 나눈 다음 gradient scale을 곱함.
Pipeline Model Parallelism도 같이 구현 (Parameter Sharding)
노드 간의 통신 시 모델 전체를 all-reduce하면 노드 간 너무나 많은 communication cost가 발생하기에 PowerSGD로 low-rank factorization을 통해 gradient를 압축해서 all-reduce communication cost를 줄임.
PowerSGD
Gain, bias, embedding, unembedding을 제외한 모든 파라메터에 대한 gradient는 PowerSGD로 gradient 압축 수행
85% compression rate 달성
Backpropagation 과정에서 gradient를 error buffer로 축적함으로써 개별적인 buffer를 할당하는 방법 대비 메모리를 절약
error buffer를 0으로 만드는 인스턴스 최소화 (mixed-precision backprop에서 생기는 nonfinite 값 혹은 체크포인트에서 모델을 재시작할 때 등)
Gram-Schmidt 대신 Householder orthogonalization을 사용함으로써 수치적인 안정성 개선
16FP 적용 시 발생하는 underflow 방지
dVAE Training
64장의 16GB V100 GPU 사용 (Batch size = 512, GPU당 batch size = 8, Total updates = 3백만)
Cosine scheduling 사용
첫 5000 updates까지 KL weight β를 0부터 6.6까지 증가
relexation temperature τ는 첫 150,000 updates까지 1에서 1/16으로 annealing
1024장의 16GB V100 GPU 사용 (Batch size = 1024, GPU당 batch size = 1, Total updates = 43만)
10% BPE dropout 적용
첫 5000 step까지 linear scheduling으로 step size를 4.5⋅10−4로 증가시키고 training loss가 감소하지 않을 때마다 step size를 절반으로 감소 (총 5번 수행하기에 최종적으로는 peak보다 32배 작은 step size를 가짐)
Illustration of VAE model with the multivariate Gaussian assumption (출처: https://lilianweng.github.io/lil-log/2018/08/12/from-autoencoder-to-beta-vae.html)
