Last Updated on 2024-02-27 by Clay
介紹
變分自動編碼器（Variational AutoEncoder, VAE) 是自動編碼器（AutoEncoder, AE）的進階變體，架構與原本的自動編碼器相似，同樣都是由編碼器（Encoder）和解碼器（Decoder）所組成。
但是 VAE 有幾個與本來的 AutoEncoder 不同的特點：
- 使用潛在空間（latent space）向量替換編碼（encode）
原本的 AutoEncoder 擁有一個確定的編碼過程，將輸入的向量通過編碼器的神經網路轉換成一組潛在特徵向量，而 VAE 則是透過生成『平均值』和『變異數』兩組向量，並讓向量中的元素兩兩配對生成一組常態分佈的向量。
請參考 VAE 的架構圖與下方公式。
其中，⊙ 為元素對元素的乘法。並且 ϵ 是從標準常態分佈 N(0, 1) 抽樣來的。並且在文獻中，通常會看到用平方根的標準差 σ 表示，這也是為了讓保持為正數。
- 重參數化（reparameterization trick）
如果我們直接讓 Encoder 的任務變成不接受任何輸入（或僅有noise為輸入）直接輸出常態分佈（帶有均值和標準差），我們將面臨一個問題 —— 這些隨機產生的常態分佈畢竟是隨機的，我們將難以進行 backward propagation 來更新模型權重，梯度會卡在 decoder 那一層。
而重參數化換了個角度去生成常態分佈採樣：我們透過 encoder 產生的，是均值與方差，以此來生成固定的一組組的常態分佈採樣。之所以是一組組不是一組，是因為我們可能產生的均值有 mean1, mean2, …, meanN，並且方差有 var1, var2, …, varN，所以我們總共有 N 組常態分佈採樣。
但若是讓模型直接透過神經網路生成這些組成常態分佈採樣的參數，我們基本上做的就是普通的神經網路輸出，組成常態分佈採樣的隨機化部分壓根沒經過神經網路，自然就可以做梯度下降來更新模型權重了。
另外，由於我們最後要讓 VAE 生成假資料，是透過常態分佈的輸入給 decoder 並讓 decoder 解碼出假資料。若是我們在訓練時非得讓 encoder 生成常態分佈，即意味著我們需要額外給定一個輸入（可能是噪音或圖片）給 encoder —— 當然在某些情況下我們可能會需要這樣綁定生成特徵的能力。
- 引入 KL 散度當作 loss function 的一部分
這是為了讓我們重參數化的輸出更貼其標準常態分佈。之所以選擇 KL 散度有以下原因：
- 數學性質：因為 KL 散度與最大化似然律有關，所以在機率建模中是一個合理的評估指標
- 非對稱性：P 和 Q 之間的 KL 散度和 Q 和 P 之間的 KL 散度不同
- 高效計算：對標準常態分佈來說，KL 散度可以直接計算
而 loss function 的另外一個部分，就是跟原本的 AutoEncoder 一樣的重建損失（reconstruction loss），我們可以直接透過 MSE 來計算輸入與最終輸出之間的落差，以此來改進。
所以整個 VAE 模型等於需要同時學習兩件事情：一個就是中間的隨機常態分佈要跟標準常態分佈相像、而通過編碼器、解碼器解碼出來的最終輸出需要與輸入越像越好。
VAE 其他特性
當然，除了以上重點之外，VAE 也還具備以下有趣的特色：
模型架構
VAE 可以使用任意深度學習神經網路層來實現，比如全連接層、CNN 和 RNN。實現的自由度非常高。但值得注意的是，在寫 loss function 時，有個乘以 KL 散度的 Beta 係數，該 Beta 係數的值越大，意味著模型越是關注中間 latent space 的常態分佈採樣需要越接近標準常態分佈越好，而 Beta 值越小，則是重建原始圖片的重要性越高。
這個 Beta 值並不好設，我反覆調整了幾次，這才讓模型在兩邊的 loss 達到平衡 —— 在測試的過程中，更是找到了 Beta-VAE 的架構，可以把 Beta 參數也變成模型學習的一部分，自動化地去設定。
生成能力
訓練結束後，VAE 的 decoder 部份可以被單獨提取之來，並使用我們從標準常態分佈中的採樣資料作為輸入，解碼出新的、從未見過的資料。
潛在空間（latent space）性質
VAE 的潛在空間有著良好的表現能力，比方說相似的資料會在潛在空間中叢聚在一起。
變體與發展
VAE 還有著許多的變體，比如 conditional VAE、Beta-VAE、disentangled VAE… 等等。這些都是為了解決特定的問題而被研究出的變體。
當然，VAE 也有其限制（尤其是最基礎的版本），比方說難以訓練或是生成的圖像較為模糊等等。
程式
以下來介紹我實作的程式部分。
首先是模型的架構。
from typing importTupleimport torch
class VAE(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
code_dim = 2# Encoder
self.encoder = torch.nn.Sequential(
torch.nn.Linear(784, 128),
torch.nn.ReLU(),
)
# Mean & Variance
self.mean_layer = torch.nn.Linear(128, code_dim)
self.log_variance_layer = torch.nn.Linear(128, code_dim)
# Decoder
self.decoder = torch.nn.Sequential(
torch.nn.Linear(code_dim, 128),
torch.nn.ReLU(),
torch.nn.Linear(128, 784),
torch.nn.Sigmoid(),
)
def reparameterization(self, means: torch.Tensor, log_variances: torch.Tensor) -> torch.Tensor:
std = torch.exp(0.5 * log_variances)
eps = torch.randn_like(std)
return means + eps * std
def forward(self, inputs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
# Encoder
hiddens = self.encoder(inputs)
# Mean & Variance
means = self.mean_layer(hiddens)
log_variances = self.log_variance_layer(hiddens)
# Get gaussion distribution
gaussian_distribution_codes = self.reparameterization(means=means, log_variances=log_variances)
# Decoder
decoded = self.decoder(gaussian_distribution_codes)
return means, log_variances, gaussian_distribution_codes, decoded訓練的程式，需要把剛才的 models.py 中定義好的模型架構 import 進來才行。
# coding: utf-8import torch
import torchvision
from models import VAE
def main() -> None:
# Settings
device = torch.device("cuda:0"if torch.cuda.is_available() else"cpu")
epochs = 100
batch_size = 128
lr = 2e-3# DataLoader
train_dataset = torchvision.datasets.MNIST(
root="../../data/MNIST/",
train=True,
download=True,
transform=torchvision.transforms.ToTensor(),
)
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
# Model
model = VAE().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
loss_function = torch.nn.MSELoss()
# Trainfor epoch inrange(1, epochs+1):
for inputs, labels in train_dataloader:
inputs = inputs.view(-1, 784).to(device)
# Forward
means, log_vars, codes, decoded = model(inputs)
# Backward
optimizer.zero_grad()
# Reconstruction loss
reconstruction_loss = loss_function(decoded, inputs)
# KL divergence loss
KL_divergence = torch.mean(-0.5 * torch.sum(1 + log_vars - means ** 2 - log_vars.exp(), dim=1), dim=0)
# Total loss
beta = 0.001
loss = reconstruction_loss + beta * KL_divergence
loss.backward()
optimizer.step()
# Show progressprint(f"[{epoch}/{epochs}] Loss: {loss.item()}")
# Save
torch.save(model, "VAE.pth")
if __name__ == "__main__":
main()訓練結束後，我們也可以可以看到視覺化的結果。
# coding: utf-8import torch
import torchvision
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from models import VAE
def main() -> None:
# Load model
model = torch.load("VAE.pth")
model.to("cuda:0").eval()
print(model)
# DataLoader
test_dataset = torchvision.datasets.MNIST(
root="../../data/MNIST/",
train=False,
download=True,
transform=torchvision.transforms.ToTensor(),
)
test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=16, shuffle=False)
# Plot
axis_x = []
axis_y = []
answers = []
with torch.no_grad():
for data in test_dataloader:
inputs = data[0].view(-1, 28*28).to("cuda:0")
answers += data[1].tolist()
means, log_vars, code, outputs = model(inputs)
axis_x += code[:, 0].tolist()
axis_y += code[:, 1].tolist()
# Use a colormap
colors = [cm.rainbow(i/9) for i in answers]
scatter = plt.scatter(axis_x, axis_y, c=colors)
cbar = plt.colorbar(scatter)
cbar.set_ticks([i/9for i inrange(10)])
cbar.set_ticklabels([str(i) for i inrange(10)])
plt.show()
plt.savefig("VAE.png")
if __name__ == "__main__":
main()Output:
比起原本的 AutoEncoder，相同標籤的資料顯然更能表示群聚在一起的潛在空間特性了。
[Machine Learning] AutoEncoder 基本介紹 (附 PyTorch 程式碼)
我們也能透過 VAE，來畫出最具代表性的由潛在空間生成的圖像。因為這能呈現出潛在空間的連續性，相似的點會對應到相似的圖像。
這次的 VAE 實作算是一個很好玩的體驗，我以為我對 PyTorch 以及深度學習的概念已經足夠熟悉了，尤其是在今年做完 [論文閱讀] The Forward-Forward Algorithm: Some Preliminary Investigation 的實現後 —— 但實際上要不參考他人的實作，自己從零開始刻起還是跌跌撞撞地遇到許多問題，包括 loss function 的公式部分一開始我寫錯了好幾次。
研究所時有個同學似乎能花一天看論文，然後花了一天就能完全復現 VAE 的結果，現在想想真的是相當厲害，真的是該跟人家多學習學習才是。