Auto Encoder
Variational Auto-Encoding을 이해하기 위해 기본적인 Auto-Encoding을 알아야 한다.
Auto Encoder(AE)는 데이터를 압축하고 복원하는 단순한 모델이다. Linear layer을 통해 데이터 크기를 줄이고 복원한다. Auto Encoder 구성은 다음과 같다.
- Encoder: 데이터를 압축하는 신경망 (파란 부분)
- latent variable: 데이터가 압축된 벡터
- Decoder: 데이터를 복원하는 신경망 (초록 부분)
다른 표현으로 Encoder를 Recognition model, Decoder를 Reconstruction model이라고 부른다.
class Autoencoder(nn.Module):
def __init__(self):
super(Autoencoder, self).__init__()
self.encoder = nn.Sequential(
nn.Linear(in_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, latent_dim),
nn.ReLU(),
)
self.decoder = nn.Sequential(
nn.Linear(latent_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, in_dim),
nn.Tanh(),
)이를 활용하면 이미지 노이즈를 제거할 수 있다. 전체 코드: Github.
입력을 노이즈 있는 이미지, 정답을 노이즈 없는 이미지로 두고 학습하면 노이즈를 제거하는 모델이 학습된다. 같은 맥락에서 워터마크를 제거하는 모델도 학습할 수 있다.
Variational AE 개요
Variational Auto Encoder(VAE)는 "Auto-Encoding Variational Bayes"에서 소개된 모델로, latent variable을 확률 분포에서 샘플링한다.
Encoder가 latent variable을 출력하는 대신, 평균($\mu$)과 표준편차($\sigma$)를 출력한다. 평균과 표준편차를 이용해 Gaussian 분포를 생성하고 latent variable을 샘플링한다. 즉, Gaussian 분포 $N(\mu ,\sigma^2)$에 대해 Encoder는 $\mu$와 $\sigma$를 생성하도록 학습한다. 샘플링한 latent $z$는 Decoder 입력이 된다. 조금 더 깊이 들어가보자.
확률 분포를 생성하고 샘플링하는 과정을 수식으로 표현해보자.
- $p_{\theta}(x)$: 풀려는 문제. 올바른 $x$를 생성해낼 확률.
- $p_{\theta}(x|z)$: Decoder. latent $z$로부터 $x$가 나올 확률.
- $p_{\theta}(z|x)$: Encoder. 입력 $x$로부터 latent $z$가 나올 확률.
- $q_{\phi}(z|x)$: $p_{\theta}(z|x)$의 근삿값.
먼저, Encoder는 입력 $x$가 주어졌을 때 $z$를 출력한다. 그런데 우리는 $x$에 대응하는 $z$를 알지 못한다. 따라서 $p_{\theta}(z|x)$를 구할 수 없다. 대신 Encoder를 학습시켜 $p_{\theta}$에 근사하는 $q_{\phi}$를 구한다.
다시 말해, Encoder를 학습하는 과정은 파라미터 $\phi$를 학습시켜 $q_{\phi}$가 $p_{\theta}$에 가까워지도록 한다.
Decoder는 $z$가 주어졌을 때 $x$를 출력한다. 따라서 $p_{\theta}(x|z)$로 표현할 수 있다.
참고로 $p(x)$는 Encoder + Decoder를 나타내는 식이 아니다. 다만, 정의한 문제 $p(x)=p(z)p(x|z)$를 풀기 위해 추론에 Encoder, Decoder 구조를 활용하는 것일 뿐이다.
Stochastic Gradient Variational Bayes
Loss function을 유도해보자.
$\log p_{\theta}(\mathbf{x})$는 log-likelihood로 올바른 $x$를 생성할 가능성을 나타낸다. 우리는 이 가능성을 최대로 만들어 올바른 $x$를 생성하려 한다.
아래는 Evidence Lower Bound: ELBO에 대한 식으로, 유도 과정을 생략하고 결과만 작성했다.
KL-divergence를 $\log p_{\theta}(x)$에 대해 정리하면 다음과 같다.
$$\log p_{\theta}(\mathbf{x}^{(i)}) = D_{KL} \left( q_{\phi}(\mathbf{z} | \mathbf{x}^{(i)}) \parallel p_{\theta}(\mathbf{z} | \mathbf{x}^{(i)}) \right) + \mathcal{L}(\theta, \phi; \mathbf{x}^{(i)})$$
KL-divergence 부분은 항상 양수이기 때문에 다음과 같은 부등식이 성립한다.
$$\log p_{\theta}(\mathbf{x}^{(i)}) \geq \mathcal{L}(\theta, \phi; \mathbf{x}^{(i)})$$
따라서 $\log p_{\theta}(\mathbf{x})$를 최대화하기 위해 $\mathcal{L}(\theta, \phi; \mathbf{x})$을 최대화해야 하고, 다시 말해 $- \mathcal{L}(\theta, \phi; \mathbf{x})$를 최소화해야 한다.
이 식을 다시 작성하면 다음과 같다.
$$- \mathcal{L}(\theta, \phi; \mathbf{x}^{(i)}) = D_{KL} \left( q_{\phi}(\mathbf{z} | \mathbf{x}^{(i)}) \parallel p_{\theta}(\mathbf{z}) \right) - \mathbb{E}{q{\phi}(\mathbf{z} | \mathbf{x}^{(i)})} \left[ \log p_{\theta}(\mathbf{x}^{(i)} | \mathbf{z}) \right]$$
여기서 우변은 Regularization + Reconstruction로 구성되어 있다.
- Regularization Loss: Encoder가 주어진 $x$에 대해 $z$를 잘 생성하는지
- Reconstruction Loss: Decoder가 주어진 $z$에 대해 $x$를 잘 생성하는지
정리하면, VAE의 Loss function은 Lower bound로부터 파생된다. Loss는 Encoder와 Decoder에 대한 Loss를 더한 값이다. 자세한 과정은 논문 2.2와 2.3에 기록되어 있다.
Reparameterization trick
앞서 설명했듯 VAE에서 latent $z$는 Gaussian 분포에서 샘플링한다.
평균을 $\mu$, 표준편차를 $\sigma$라 할 때,
$$z^{(i,l)}\sim q_{\phi}(z|x^{(i)})$$
$$z^{(i,l)} = \mu^{(i)} + \sigma^{(i)} \odot \epsilon^{(l)}$$
$\epsilon\sim N(0,1)$는 랜덤한 작은 값이다.
epsilon = randn_like(std)
z = mu + std * epsilonLoss Function 정의
위에서 설명했던 Loss는 일반화된 모습이었다. 구현을 위해서는 구체적인 식을 정의해야 한다.
$$p_{\theta}(z)\sim N(z;0,I)$$
$$\log q_{\phi}(z|x^{(i)})=\log N(z;\mu^{(i)},\sigma^{2(i)}I)$$
먼저, $p_{\theta}(z)$는 centered isotropic Gaussian을 따르며, $\log q_{\phi}(z|x)$도 Gaussian을 따른다고 가정한다.
$$- \mathcal{L}(\theta, \phi; \mathbf{x}^{(i)}) \simeq - \frac{1}{2} \sum_{j=1}^{J} \left( 1 + \log \left( (\sigma_{j}^{(i)})^2 \right) - (\mu_{j}^{(i)})^2 - (\sigma_{j}^{(i)})^2 \right) - \frac{1}{L} \sum_{l=1}^{L} \log p_{\theta} (\mathbf{x}^{(i)} | \mathbf{z}^{(i,l)})$$
이 식은 Gaussian 분포에 대해 Regularization Loss를 구체적으로 정의했다. 두번째 항인 Reconstruction Loss는 negative log-likelihood다. 따라서, Binary Cross Entropy로 정의할 수 있다.
def loss(x, x_reconstructed, mu, std):
# Regularization Loss
kl_div = -0.5 * sum(1 + log(std.pow(2)) - mu.pow(2) - std.pow(2))
# Reconstruction Loss
recon_loss = binary_cross_entropy(x_reconstructed, x)
return kl_div + recon_lossPytorch 구현
전체 구현은 Github: VAE에서 확인할 수 있다.
class VAE(nn.Module):
def __init__(self, input_dim, hidden_dim, latent_dim):
super(VAE, self).__init__()
self.encoder = Encoder(input_dim, hidden_dim, latent_dim)
self.decoder = Decoder(latent_dim, hidden_dim, input_dim)
def forward(self, x):
mu, logvar = self.encoder(x)
# Reparameterization trick
std = torch.exp(0.5 * logvar)
epsilon = torch.randn_like(std)
z = mu + std * epsilon
x_recon = self.decoder(z)
return x_recon, mu, logvar구현에는 표준편차 $\sigma$ 대신 $\log \sigma^2$인 logvar를 반환하도록 한다.
$\sigma$는 일반적으로 매우 작은 값으로 계산된다. 따라서 학습 과정에서 최적화가 잘 되지 않는 문제가 있다. 하지만 분산을 log 공간에 매핑시키면 값을 더 큰 범위로 변환할 수 있다. $\sigma$가 일반적으로 [0, 1] 범위를 가진다고 하면, $\log \sigma^2$는 [log(1), -inf] 범위를 가진다. 따라서 학습 과정에서 잘 최적화되는 모습을 보인다. - 출처.
참고로 $\log \sigma^2$가 음수 범위를 가지기 때문에 logvar를 출력하는 layer는 activation으로 ReLU를 사용하면 안 된다.
def loss(x, x_recon, mu, logvar):
recon_loss = nn.functional.binary_cross_entropy(x_recon, x)
kl_div = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return recon_loss + kl_div이렇게 하면 Loss function도 logvar에 대해 재정의할 수 있다.
MNIST 데이터셋을 이용해 학습하면 입력 이미지와 유사한 출력 만들어 낸다.
시각화
2차원 latent space를 시각화했다. 코드: Github
2차원 latent space, 즉 2개의 Gaussian 분포를 생성하도록 Encoder를 학습시켰다. $z$를 [-3, 3] 범위에 대해 Decoder에 입력했다. $p(z)$가 Standard Normal Distribution을 따른다고 가정했기 때문에 [-3, 3] 범위로 latent space 대부분을 시각화할 수 있다.
시각화한 이미지를 통해 샘플링된 $z$와 출력 $x$의 관계를 확인할 수 있다.