Variational AutoEncoder

from torchvision import datasets
import matplotlib.pyplot as plt
import numpy as np
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

train = datasets.MNIST("./_mnist", train=True, download=True)
test = datasets.MNIST("./_mnist", train=False, download=True)
print(train.data.shape)
print(test.data.shape)

print(train.data[0])
print("Original shape:", train.data[0].shape)
print("Label:", train.targets[0].item())

plt.figure(figsize=(3, 3))
r = 0
plt.imshow(train.data[r], cmap="gray")
plt.title(f"Raw MNIST image, label = {train.targets[r].item()}")
plt.axis("off")
plt.show()

#note 28 * 28 = 784
Xtr = (train.data.float() / 255.0).view(-1, 784)
ytr = train.targets

Xte = (test.data.float() / 255.0).view(-1, 784)
yte = test.targets

print("train:", Xtr.shape, "test:", Xte.shape)

class MNIST_VAE(nn.Module):
    def __init__(self, x_dim=784, z_dim=20, h=400):
        super().__init__()
        self.enc1       = nn.Linear(x_dim, h)
        self.enc_mu     = nn.Linear(h, z_dim)
        self.enc_logvar = nn.Linear(h, z_dim)
        self.dec1       = nn.Linear(z_dim, h)
        self.dec2       = nn.Linear(h, x_dim)

    def encode(self, x):
        h = F.relu(self.enc1(x))
        return self.enc_mu(h), self.enc_logvar(h)

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

    def decode(self, z):
        h = F.relu(self.dec1(z))
        return torch.sigmoid(self.dec2(h))         # Bernoulli mean

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

def vae_loss_bernoulli(x, x_rec, mu, logvar):
    bce = F.binary_cross_entropy(x_rec, x, reduction="none").sum(dim=1)
    kl  = 0.5 * (mu.pow(2) + logvar.exp() - 1 - logvar).sum(dim=1)
    return bce.mean(), kl.mean()

# --- train MNIST VAE on CPU ---
mnist_model = MNIST_VAE()

opt = torch.optim.Adam(mnist_model.parameters(), lr=1e-3)

dl_tr = DataLoader(TensorDataset(Xtr), batch_size=128, shuffle=True)

EPOCHS = 40

vae_history = []
vae_rec_history = []
vae_kl_history = []

for epoch in range(EPOCHS):
    rec_e, kl_e, n = 0.0, 0.0, 0

    for (xb,) in dl_tr:
        x_rec, mu, logvar = mnist_model(xb)

        rec, kl = vae_loss_bernoulli(xb, x_rec, mu, logvar)
        loss = rec + kl

        opt.zero_grad()
        loss.backward()
        opt.step()

        rec_e += rec.item() * xb.size(0)
        kl_e  += kl.item()  * xb.size(0)
        n     += xb.size(0)

    avg_rec = rec_e / n
    avg_kl = kl_e / n
    avg_neg_elbo = (rec_e + kl_e) / n

    vae_rec_history.append(avg_rec)
    vae_kl_history.append(avg_kl)
    vae_history.append(avg_neg_elbo)

    print(
        f"epoch {epoch+1:2d}  "
        f"rec={avg_rec:.2f}  "
        f"kl={avg_kl:.2f}  "
        f"-ELBO={avg_neg_elbo:.2f}"
    )

mnist_model.eval()

with torch.no_grad():
    # Randomly choose 8 test images
    idx = torch.randperm(Xte.shape[0])[:8]
    xb = Xte[idx]
    labels = yte[idx]

    # Decoder output: \hat{p}_i
    x_rec, _, _ = mnist_model(xb)

    # Bernoulli sample: \tilde{y}_i ~ Bernoulli(\hat{y}_i)
    x_sample = torch.bernoulli(x_rec)

fig, ax = plt.subplots(3, 8, figsize=(10, 4))

for j in range(8):
    ax[0, j].imshow(xb[j].view(28, 28), cmap="gray")
    ax[0, j].set_title(f"input: {labels[j].item()}")
    ax[0, j].axis("off")

    ax[1, j].imshow(x_rec[j].view(28, 28), cmap="gray")
    ax[1, j].axis("off")

    ax[2, j].imshow(x_sample[j].view(28, 28), cmap="gray")
    ax[2, j].axis("off")

ax[1, 0].set_title("reconstruction", loc="left")
ax[2, 0].set_title("Bernoulli sample", loc="left")

plt.tight_layout()
plt.show()

with torch.no_grad():
    z = torch.randn(36, 20)
    x_gen = mnist_model.decode(z)

fig, ax = plt.subplots(6, 6, figsize=(7, 7))
for i, a in enumerate(ax.flat):
    a.imshow(x_gen[i].view(28, 28), cmap="gray")
    a.axis("off")

plt.suptitle("New Samples from the VAE prior", y=0.92)
plt.show()

mnist_model.eval()

with torch.no_grad():
    num1 = 6
    num2 = 9

    idx_num1 = (yte == num1).nonzero(as_tuple=True)[0]
    idx_num2 = (yte == num2).nonzero(as_tuple=True)[0]

    i_num1 = idx_num1[torch.randint(len(idx_num1), (1,))].item()
    i_num2 = idx_num2[torch.randint(len(idx_num2), (1,))].item()
    
    x1, x2 = Xte[i_num1], Xte[i_num2]

    mu1, _ = mnist_model.encode(x1.unsqueeze(0))
    mu2, _ = mnist_model.encode(x2.unsqueeze(0))

    ts = torch.linspace(0, 1, 10).view(-1, 1)
    zs = (1 - ts) * mu1 + ts * mu2
    xs = mnist_model.decode(zs)

fig, ax = plt.subplots(1, 10, figsize=(12, 1.5))
for i, a in enumerate(ax):
    a.imshow(xs[i].view(28, 28), cmap="gray")
    a.axis("off")

plt.suptitle("Linear interpolation in latent space", y=1.05)
plt.show()

mnist_model.eval()

with torch.no_grad():
    num1 = 6
    num2 = 9

    idx_num1 = (yte == num1).nonzero(as_tuple=True)[0]
    idx_num2 = (yte == num2).nonzero(as_tuple=True)[0]

    i_num1 = idx_num1[torch.randint(len(idx_num1), (1,))].item()
    i_num2 = idx_num2[torch.randint(len(idx_num2), (1,))].item()
    
    y1, y2 = Xte[i_num1], Xte[i_num2]

    # Latent-space interpolation
    mu1, _ = mnist_model.encode(y1.unsqueeze(0))
    mu2, _ = mnist_model.encode(y2.unsqueeze(0))

    ts = torch.linspace(0, 1, 10).view(-1, 1)

    zs = (1 - ts) * mu1 + ts * mu2
    ys_latent = mnist_model.decode(zs)

    # Naive pixel-space interpolation
    ys_naive = (1 - ts) * y1 + ts * y2

fig, ax = plt.subplots(2, 10, figsize=(12, 3))

for i in range(10):
    ax[0, i].imshow(ys_latent[i].view(28, 28), cmap="gray")
    ax[0, i].axis("off")

    ax[1, i].imshow(ys_naive[i].view(28, 28), cmap="gray")
    ax[1, i].axis("off")

ax[0, 0].set_ylabel("latent", rotation=0, labelpad=25)
ax[1, 0].set_ylabel("pixel", rotation=0, labelpad=25)

plt.suptitle("Latent-space interpolation vs naive pixel-space interpolation", y=1.02)
plt.tight_layout()
plt.show()

# ------------------------------------------------------------
# Denoising MNIST images using the fitted VAE
# ------------------------------------------------------------

mnist_model.eval()

# Number of test images to show
n_show = 10

# Probability of flipping each pixel
flip_prob = 0.02

# Take n_show random test images
idx = torch.randperm(Xte.shape[0])[:n_show]
x_clean = Xte[idx]

# Make random 0/1 flip mask
flip_mask = torch.rand_like(x_clean) < flip_prob

# Since MNIST pixels are grayscale in [0,1], first binarize for clean flipping
x_binary = (x_clean > 0.5).float()

# Flip selected pixels: 0 -> 1 and 1 -> 0
x_noisy = x_binary.clone()
x_noisy[flip_mask] = 1.0 - x_noisy[flip_mask]

with torch.no_grad():
    # Encode noisy image and decode it
    # This gives the VAE reconstruction/denoised image
    x_denoised, mu, logvar = mnist_model(x_noisy)

fig, ax = plt.subplots(3, n_show, figsize=(1.3 * n_show, 4.0))

for j in range(n_show):
    ax[0, j].imshow(x_clean[j].view(28, 28), cmap="gray")
    ax[0, j].axis("off")

    ax[1, j].imshow(x_noisy[j].view(28, 28), cmap="gray")
    ax[1, j].axis("off")

    ax[2, j].imshow(x_denoised[j].view(28, 28), cmap="gray")
    ax[2, j].axis("off")

ax[0, 0].set_title("original", loc="left")
ax[1, 0].set_title("noisy", loc="left")
ax[2, 0].set_title("VAE denoised", loc="left")

plt.tight_layout()
plt.show()

from sklearn.manifold import TSNE

# ------------------------------------------------------------
# 1. Encode test images into latent means
# ------------------------------------------------------------

mnist_model.eval()

n_plot = 3000   # t-SNE can be slow, so use a subset of test images

X_sub = Xte[:n_plot]
y_sub = yte[:n_plot].numpy()


with torch.no_grad():
    mu_z, logvar_z = mnist_model.encode(X_sub)

# Convert latent means to NumPy
Z_mu = mu_z.numpy()

print("Latent mean shape:", Z_mu.shape)
# Should be (n_plot, 20) if z_dim = 20


# ------------------------------------------------------------
# 2. Run t-SNE on the latent means
# ------------------------------------------------------------

tsne = TSNE(
    n_components=2,
    perplexity=30,
    learning_rate="auto",
    init="pca",
    random_state=0
)

Z_tsne = tsne.fit_transform(Z_mu)


# ------------------------------------------------------------
# 3. Plot t-SNE projection colored by digit labels
# ------------------------------------------------------------

plt.figure(figsize=(6, 5))

sc = plt.scatter(
    Z_tsne[:, 0],
    Z_tsne[:, 1],
    c=y_sub,
    s=6,
    cmap="tab10",
    alpha=0.75
)

plt.colorbar(sc, label="digit label")
plt.xlabel("t-SNE 1")
plt.ylabel("t-SNE 2")
plt.title("t-SNE of VAE latent means $\\mu_\\phi(x)$")
plt.grid(alpha=0.3)
plt.show()

class DirectBernoulliDecoder(nn.Module):
    def __init__(self, x_dim=784, z_dim=20, h=400):
        super().__init__()
        self.z_dim = z_dim
        self.dec1 = nn.Linear(z_dim, h)
        self.dec2 = nn.Linear(h, x_dim)

    def logits(self, z):
        h = F.relu(self.dec1(z))
        return self.dec2(h)

    def decode(self, z):
        return torch.sigmoid(self.logits(z))


def log_bernoulli_prob_from_logits(x, logits):
    """
    Compute log p(x | z) for every pair in a batch of x's and z's.

    x:      shape (B, D)
    logits: shape (M, D), where pi_theta(z_m) = sigmoid(logits[m])

    returns:
        log_probs with shape (B, M), where entry (i, m) is log p_theta(x_i | z_m).
    """
    # -binary_cross_entropy_with_logits gives
    # sum_j x_j log sigmoid(logit_j) + (1-x_j) log(1-sigmoid(logit_j)).
    x_expanded = x[:, None, :]           # (B, 1, D)
    logits_expanded = logits[None, :, :] # (1, M, D)

    log_prob = -F.binary_cross_entropy_with_logits(
        logits_expanded.expand(x.shape[0], -1, -1),
        x_expanded.expand(-1, logits.shape[0], -1),
        reduction="none"
    ).sum(dim=2)

    return log_prob


def mc_marginal_loglik(x, model, M=16, z_fixed=None):
    """
    Monte Carlo approximation to log p_theta(x_i), for each row x_i.

    If z_fixed is supplied, it is reused. Reusing z_fixed gives a deterministic
    approximate objective. Drawing new z each time gives an unbiased estimate
    of the likelihood integral but a noisier training objective.
    """
    if z_fixed is None:
        z = torch.randn(M, model.z_dim, device=x.device)
    else:
        z = z_fixed.to(x.device)

    logits = model.logits(z)                  # (M, 784)
    log_p_x_given_z = log_bernoulli_prob_from_logits(x, logits)  # (B, M)

    # log mean_m exp(log p(x | z_m))
    return torch.logsumexp(log_p_x_given_z, dim=1) - math.log(z.shape[0])


def train_direct_decoder(
    X_train,
    z_dim=20,
    h=400,
    M=16,
    batch_size=256,
    epochs=5,
    lr=1e-3,
    fixed_mc_samples=True,
    random_seed=0
):
    """
    Maximize the Monte Carlo approximation to the marginal likelihood.

    Important: this is usually much harder than VAE training. For classroom
    speed, start with epochs=5 and M=16. Larger M gives a better likelihood
    approximation but is slower.
    """
    torch.manual_seed(random_seed)

    model = DirectBernoulliDecoder(x_dim=784, z_dim=z_dim, h=h)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    dl = DataLoader(TensorDataset(X_train), batch_size=batch_size, shuffle=True)

    z_fixed = None
    if fixed_mc_samples:
        z_fixed = torch.randn(M, z_dim)

    history = []

    for epoch in range(epochs):
        total_nll, n = 0.0, 0

        for (xb,) in dl:
            loglik_i = mc_marginal_loglik(xb, model, M=M, z_fixed=z_fixed)
            nll = -loglik_i.mean()

            opt.zero_grad()
            nll.backward()
            opt.step()

            total_nll += nll.item() * xb.size(0)
            n += xb.size(0)

        avg_nll = total_nll / n
        history.append(avg_nll)
        print(f"direct epoch {epoch+1:2d}  MC marginal NLL={avg_nll:.2f}")

    return model, history


# Fit the direct model to the same training data Xtr.
# This can be slower than VAE training because every image is compared with M latent samples.
direct_model, direct_history = train_direct_decoder(
    Xtr,
    z_dim=20,
    h=400,
    M=64,
    batch_size=256,
    epochs=40,
    lr=1e-3,
    fixed_mc_samples=True,
    random_seed=0
)

# ------------------------------------------------------------
# Visual comparison of generated samples from both fitted decoders
# ------------------------------------------------------------

mnist_model.eval()
direct_model.eval()

with torch.no_grad():
    z = torch.randn(8, 20)
    x_gen_vae = mnist_model.decode(z)
    x_gen_direct = direct_model.decode(z)

fig, ax = plt.subplots(2, 8, figsize=(10, 2.8))

for j in range(8):
    ax[0, j].imshow(x_gen_vae[j].view(28, 28), cmap="gray")
    ax[0, j].axis("off")

    ax[1, j].imshow(x_gen_direct[j].view(28, 28), cmap="gray")
    ax[1, j].axis("off")

ax[0, 0].set_title("VAE generated samples", loc="left")
ax[1, 0].set_title("Direct-fit generated samples", loc="left")

plt.tight_layout()
plt.show()

# ------------------------------------------------------------
# Reconstruction comparison with pixel-flip noise
# ------------------------------------------------------------
# The direct model has no encoder, so it cannot reconstruct by encoding x.
# To reconstruct an image under the direct model, we search over candidate z's
# and choose the z that gives the largest p_theta(x | z).
#
# This is deliberately slower and less elegant than the VAE encoder.

def direct_reconstruct_by_prior_search(model, x, M=20000):
    """
    For each image x_i, draw M prior candidates z_m and pick
    argmax_m p_theta(x_i | z_m). This gives a crude reconstruction
    for the direct model.
    """
    model.eval()

    with torch.no_grad():
        z = torch.randn(M, model.z_dim)
        logits = model.logits(z)
        log_probs = log_bernoulli_prob_from_logits(x, logits)  # (B, M)
        best = log_probs.argmax(dim=1)
        x_rec = torch.sigmoid(logits[best])

    return x_rec


# ------------------------------------------------------------
# Pixel-flip noise function
# ------------------------------------------------------------

def add_pixel_flip_noise(x, f=0.01):
    """
    Flip each pixel independently with probability f.

    If x_ij is close to 0, it becomes close to 1.
    If x_ij is close to 1, it becomes close to 0.
    """
    flip_mask = torch.bernoulli(f * torch.ones_like(x))
    x_noisy = torch.where(flip_mask == 1, 1.0 - x, x)
    return x_noisy


n_show = 8
f = 0.02

xb = Xte[:n_show]

# Add pixel-flip noise before reconstruction
xb_noisy = add_pixel_flip_noise(xb, f=f)

with torch.no_grad():
    x_rec_vae, _, _ = mnist_model(xb_noisy)

x_rec_direct = direct_reconstruct_by_prior_search(direct_model, xb_noisy, M=20000)


fig, ax = plt.subplots(4, n_show, figsize=(10, 5.0))

for j in range(n_show):
    ax[0, j].imshow(xb[j].view(28, 28), cmap="gray")
    ax[0, j].axis("off")

    ax[1, j].imshow(xb_noisy[j].view(28, 28), cmap="gray")
    ax[1, j].axis("off")

    ax[2, j].imshow(x_rec_vae[j].view(28, 28), cmap="gray")
    ax[2, j].axis("off")

    ax[3, j].imshow(x_rec_direct[j].view(28, 28), cmap="gray")
    ax[3, j].axis("off")

ax[0, 0].set_title("original input", loc="left")
ax[1, 0].set_title(f"noisy input, f={f}", loc="left")
ax[2, 0].set_title("VAE reconstruction", loc="left")
ax[3, 0].set_title("Direct model search", loc="left")

plt.tight_layout()
plt.show()