Teori dan Implementasi Variational Autoencoder (VAE) dengan PyTorch untuk Generative Modeling

Mengenal Arsitektur Variational Autoencoder dan Perbedaannya dengan Autoencoder Standar

Autoencoder standar bekerja dengan mengompresi input menjadi representasi terkompresi di latent space melalui enkoder, kemudian mendekodenya kembali menjadi rekonstruksi. Tujuan utamanya adalah mempelajari representasi efisien dari data input. Namun, autoencoder standar memiliki keterbatasan fundamental: latent space yang dihasilkan tidak terstruktur dan tidak kontinu. Ini berarti kita tidak bisa mengambil titik acak di latent space dan mengharapkan output yang masuk akal karena model hanya mempelajari pemetaan diskrit dari input ke titik-titik tertentu.

Variational Autoencoder (VAE) mengatasi masalah ini dengan pendekatan probabilistik. Alih-alih memetakan input ke sebuah titik tunggal di latent space, enkoder VAE memetakan input ke parameter distribusi probabilitas — yaitu mean (μ) dan variance (σ²). Lebih spesifik, enkoder menghasilkan dua vektor: mu dan log_var, yang merepresentasikan mean dan log-variance dari distribusi Gaussian di latent space. Dekoder kemudian mengambil sample dari distribusi ini untuk merekonstruksi input asli.

Perbedaan kunci inilah yang membuat VAE bersifat generative. Karena latent space VAE terstruktur sebagai distribusi kontinu, kita bisa mengambil sample dari area mana pun di ruang tersebut dan dekoder akan menghasilkan output yang koheren. Autoencoder biasa menghasilkan titik-titik terisolasi, sedangkan VAE menghasilkan region kontinu yang bisa di-sample secara bermakna.

Gambar: Arsitektur dasar VAE — enkoder memetakan input ke parameter distribusi latent (μ, σ), kemudian sample z diambil dari distribusi tersebut dan didekode menjadi rekonstruksi. — Sumber: [Wikimedia Commons](https://commons.wikimedia.org/wiki/File:VAE_Basic.png) (CC BY-SA 4.0 — EugenioTL)

Memahami Peran Reconstruction Loss dan KL Divergence dalam VAE

Loss function VAE terdiri dari dua komponen yang saling melengkapi: reconstruction loss dan KL divergence. Keduanya bekerja bersama untuk mencapai dua tujuan yang terkadang bertentangan — menghasilkan rekonstruksi akurat sekaligus menjaga latent space tetap teratur.

Reconstruction loss mengukur seberapa mirip output yang direkonstruksi dengan input asli. Pada kasus data gambar seperti MNIST, kita umumnya menggunakan Binary Cross-Entropy (BCE) karena nilai piksel dinormalisasi ke rentang [0,1]. Komponen ini memastikan bahwa informasi dari input tidak hilang selama proses encoding-decoding.

KL divergence berperan sebagai regularisasi yang mendorong distribusi latent q(z|x) agar mendekati prior distribution p(z) yang kita tetapkan — biasanya standard normal distribution N(0,1). Regularisasi ini memaksa latent space menjadi terstruktur dan kontinu. Tanpa KL divergence, model akan cenderung mengabaikan regularisasi dan menghasilkan latent space yang tidak teratur, mirip dengan autoencoder biasa.

Kedua komponen ini menciptakan trade-off: reconstruction loss mendorong fidelity tinggi, sementara KL divergence mendorong regularitas latent space. Bobot relatif keduanya dikontrol melalui hyperparameter β pada β-VAE. Rumusan loss function VAE dapat dituliskan sebagai:

$$ \mathcal{L} = -\mathbb{E}_{z \sim q(z|x)}[\log p(x|z)] + \beta \cdot D_{KL}(q(z|x) \| p(z)) $$

Suku pertama adalah reconstruction loss (negative log-likelihood), dan suku kedua adalah KL divergence yang mengukur divergensi antara distribusi posterior yang dipelajari dengan prior.

Reparameterization Trick dan Perannya dalam Pelatihan VAE

Proses sampling dari distribusi Gaussian z ~ N(μ, σ²) memiliki satu masalah dalam konteks neural network: operasi sampling bersifat stochastic dan tidak differentiable. Ini berarti gradient tidak bisa mengalir melalui node sampling saat backpropagation, sehingga parameter μ dan σ tidak bisa di-update.

Reparameterization trick menyelesaikan masalah ini dengan memisahkan komponen deterministic dan stochastic. Alih-alih langsung mengambil sample dari N(μ, σ²), kita mengambil sample ε dari N(0,1) yang bersifat eksternal terhadap parameter model, lalu mentransformasikannya:

def reparameterize(mu, log_var):
    std = torch.exp(0.5 * log_var)
    eps = torch.randn_like(std)
    return mu + eps * std

Dengan trik ini, ε diperlakukan sebagai input noise eksternal, sementara μ dan σ tetap menjadi bagian dari computational graph. Gradient dapat mengalir melalui μ dan σ tanpa terhalang oleh operasi sampling. Inilah yang memungkinkan VAE dilatih end-to-end menggunakan stochastic gradient descent.

Gambar: Reparameterization Trick — sampling dari distribusi posterior dipecah menjadi komponen deterministic (μ + σ) dan komponen stochastic eksternal (ε ~ N(0,1)), sehingga gradient dapat mengalir melalui parameter μ dan σ saat backpropagation. — Sumber: [Wikimedia Commons](https://commons.wikimedia.org/wiki/File:Reparameterization_Trick.png) (CC BY-SA 4.0 — EugenioTL)

Implementasi Enkoder dan Dekoder VAE dengan PyTorch

Implementasi VAE untuk dataset MNIST membutuhkan tiga class utama: Encoder, Decoder, dan VAE yang menggabungkan keduanya. Enkoder menerima gambar grayscale 28×28 dan menghasilkan parameter distribusi latent, sementara dekoder merekonstruksi gambar dari latent vector.

import torch
import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)

    def forward(self, x):
        h = torch.relu(self.fc1(x))
        mu = self.fc_mu(h)
        log_var = self.fc_logvar(h)
        return mu, log_var

class Decoder(nn.Module):
    def __init__(self, latent_dim=20, hidden_dim=400, output_dim=784):
        super().__init__()
        self.fc1 = nn.Linear(latent_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, z):
        h = torch.relu(self.fc1(z))
        return torch.sigmoid(self.fc2(h))

class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
        super().__init__()
        self.encoder = Encoder(input_dim, hidden_dim, latent_dim)
        self.decoder = Decoder(latent_dim, hidden_dim, input_dim)

    def reparameterize(self, mu, log_var):
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        mu, log_var = self.encoder(x)
        z = self.reparameterize(mu, log_var)
        return self.decoder(z), mu, log_var

Output:

VAE model created: 652,824 parameters

Arsitektur ini menggunakan fully-connected layers dengan hidden_dim=400 dan latent_dim=20. Enkoder memproses input 784 dimensi (28×28 yang di-flatten) menjadi dua vektor 20 dimensi — mu dan log_var. Reparameterization trick terjadi di method forward class VAE, menghasilkan z yang kemudian di-pass ke dekoder. Dekoder mengembalikan output 784 dimensi dengan aktivasi sigmoid karena nilai piksel dinormalisasi.

Gambar: Arsitektur VAE dengan reparameterization — diagram menunjukkan aliran data dari input x ke enkoder yang menghasilkan μ dan σ, lalu melalui reparameterization trick (z = μ + σ ⊙ ε) menuju dekoder yang menghasilkan rekonstruksi x'. — Sumber: [Wikimedia Commons](https://commons.wikimedia.org/wiki/File:Reparameterized_Variational_Autoencoder.png) (CC BY-SA 4.0 — EugenioTL)

Melatih VAE pada Dataset MNIST dan Memvisualisasikan Hasil

Proses training VAE pada MNIST melibatkan forward pass, kalkulasi loss sebagai kombinasi BCE dan KL divergence, kemudian backpropagation untuk update parameter. Loss function perlu diimplementasikan secara eksplisit.

from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import torch.optim as optim

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.view(-1))
])

train_loader = DataLoader(
    datasets.MNIST('./data', train=True, download=True, transform=transform),
    batch_size=128, shuffle=True
)

def loss_function(x_recon, x, mu, log_var):
    bce = nn.functional.binary_cross_entropy(x_recon, x, reduction='sum')
    kld = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
    return (bce + kld) / x.size(0)

vae = VAE()
optimizer = optim.Adam(vae.parameters(), lr=1e-3)

for epoch in range(20):
    vae.train()
    train_loss = 0
    for batch_idx, (data, _) in enumerate(train_loader):
        optimizer.zero_grad()
        recon_batch, mu, log_var = vae(data)
        loss = loss_function(recon_batch, data, mu, log_var)
        loss.backward()
        train_loss += loss.item()
        optimizer.step()
    print(f'Epoch {epoch+1}, Loss: {train_loss / len(train_loader):.4f}')

Output:

Epoch 1, Loss: 163.6289
Epoch 2, Loss: 121.2987
Epoch 3, Loss: 114.4573
Epoch 4, Loss: 111.5716
Epoch 5, Loss: 109.7974
Epoch 6, Loss: 108.6668
Epoch 7, Loss: 107.8710
Epoch 8, Loss: 107.1789
Epoch 9, Loss: 106.6471
Epoch 10, Loss: 106.2007
Epoch 11, Loss: 105.8686
Epoch 12, Loss: 105.5331
Epoch 13, Loss: 105.2519
Epoch 14, Loss: 105.0616
Epoch 15, Loss: 104.7857
Epoch 16, Loss: 104.5901
Epoch 17, Loss: 104.4028
Epoch 18, Loss: 104.2085
Epoch 19, Loss: 104.1089
Epoch 20, Loss: 103.9075

Training loop di atas berjalan selama 20 epoch dengan Adam optimizer. Setiap batch melalui proses: forward pass menghasilkan rekonstruksi dan parameter latent, loss function menggabungkan BCE untuk reconstruction quality dan KLD untuk latent regularization, lalu backpropagation meng-update seluruh parameter. Setelah training selesai, kita bisa membandingkan gambar asli dengan hasil rekonstruksi untuk memvalidasi kualitas model.

Menghasilkan Gambar Baru dari Latent Space yang Sudah Terlatih

Salah satu kemampuan utama VAE adalah menghasilkan gambar sintetis yang belum pernah dilihat model. Setelah training, dekoder dapat digunakan secara mandiri — kita cukup memberikan random latent vector dan dekoder akan menghasilkan gambar baru.

import matplotlib.pyplot as plt

vae.eval()
with torch.no_grad():
    z = torch.randn(64, 20)
    samples = vae.decoder(z).view(-1, 28, 28).cpu()

fig, axes = plt.subplots(8, 8, figsize=(8, 8))
for i, ax in enumerate(axes.flat):
    ax.imshow(samples[i], cmap='gray')
    ax.axis('off')
plt.tight_layout()
plt.show()

Output:

Kita mengambil 64 random vector dari N(0,1) dan menjalankannya melalui dekoder. Hasilnya adalah grid 8×8 gambar digit yang terlihat realistis meskipun model belum pernah melihat kombinasi latent tersebut. Eksperimen lebih lanjut bisa dilakukan dengan interpolasi linear antara dua latent vector untuk mengamati transisi mulus antar digit, menunjukkan bahwa latent space VAE benar-benar kontinu dan terstruktur.

Best Practices dan Keterbatasan VAE dalam Praktik

VAE memiliki keunggulan dalam menghasilkan latent space yang terstruktur dan kontinu, namun ada trade-off yang perlu dipahami. Gambar yang dihasilkan VAE cenderung tampak blurry dibandingkan dengan output dari Generative Adversarial Network (GAN). Ini terjadi karena VAE mengoptimalkan likelihood dari seluruh distribusi data (coverage), sementara GAN fokus pada kualitas visual (sharpness).

Untuk meningkatkan kualitas hasil, beberapa variasi bisa dipertimbangkan. β-VAE menambahkan hyperparameter β pada KL divergence term untuk mengontrol keseimbangan antara reconstruction fidelity dan latent disentanglement. Meningkatkan kapasitas dekoder dengan layer convolutional juga dapat membantu menghasilkan gambar yang lebih tajam.

Gunakan VAE ketika prioritas adalah structured latent representation dan coverage data yang luas — misalnya untuk anomaly detection, data imputation, atau sebagai komponen dalam sistem yang lebih besar. Jika target utama adalah realisme visual, alternatif seperti GAN, VQ-VAE, atau Conditional VAE mungkin lebih sesuai.

Eksplorasi generative model tidak berhenti di VAE. Di kursus Deep Learning Rumah Coding, kami membahas generative model secara lebih mendalam — termasuk GAN, Diffusion Model, dan aplikasi praktisnya dalam computer vision. Kunjungi halaman kursus kami untuk melihat silabus lengkap dan memulai perjalanan di bidang generative AI.