Pre-processing

Load in images and apply safe data augmentation. Split to a smaller train dataset that will speed up model training. The rest are validation that don't actually touch the model but can be used to assess the quality of reconstruction from the latent bottleneck at real images on the pet manifold.

import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import lpips
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import glob
import matplotlib.pyplot as plt

class DogCatDataset(Dataset):
    def __init__(self, folder_path, transform=None):
        # Collect ALL image files (both dog and cat)
        self.image_paths = glob.glob(os.path.join(folder_path, '*.jpg'))
        self.transform = transform

        # Build labels from filename: 'dog' -> 1, 'cat' -> 0
        self.labels = []
        for path in self.image_paths:
            filename = os.path.basename(path)
            if filename.startswith('dog'):
                self.labels.append(1)
            else:
                self.labels.append(0)

        # Human-readable label names (useful for plotting)
        self.label_names = {0: 'cat', 1: 'dog'}

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        label = self.labels[idx]
        return image, label

# Define standard preprocessing
train_transform = transforms.Compose([
    transforms.Resize((72, 72)),           # slightly larger than target
    transforms.RandomCrop((64, 64)),       # random 64×64 crop
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.clamp(0, 1))
])

# Validation gets no augmentation — deterministic center crop only
val_transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.clamp(0, 1))
])

dog_cat_dataset_train = DogCatDataset('dogs-vs-cats/train', transform=train_transform)
dog_cat_dataset_val   = DogCatDataset('dogs-vs-cats/train', transform=val_transform)

train_size = int(0.5 * len(dog_cat_dataset_train))
val_size   = len(dog_cat_dataset_train) - train_size

# Fix the split indices so both datasets use the same images for train/val
indices = torch.randperm(len(dog_cat_dataset_train)).tolist()
train_dataset = torch.utils.data.Subset(dog_cat_dataset_train, indices[:train_size])
val_dataset   = torch.utils.data.Subset(dog_cat_dataset_val,   indices[train_size:])

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True,  num_workers=6)
val_loader   = DataLoader(val_dataset,   batch_size=64, shuffle=False, num_workers=6)

print(f"Number of training images:   {len(train_dataset)}")
print(f"Number of validation images: {len(val_dataset)}")

Number of training images:   12500
Number of validation images: 12500

Let's look at 100 training images to see what they look like

plt.figure(figsize=(15, 15))
for i in range(100):
    image, label = train_dataset[i]
    # Convert tensor (C, H, W) to numpy (H, W, C) for plotting
    img_np = image.permute(1, 2, 0).numpy()
    
    plt.subplot(10, 10, i + 1)
    plt.imshow(img_np)
    plt.title(dog_cat_dataset_train.label_names[label])
    plt.axis('off')

plt.tight_layout()
plt.show()

Model 1: Standard VAE w/ L1 Loss

• ResBlock Encoder w/ Stride 2 Downsampling
• ResBlock Encoder w/ Bilinear Upsampling
• Pixel likelihood conditional on latent code is assumed to follow a Laplace distribution
• KL Penalty on the latent space to force compactness

Stride 2 Downsampling vs. Max Pool

In VAEs and other generative models, it is often preferred to use Stride 2 Downsampling (via convolutional layers) rather than Max Pooling for several reasons:

1. Learnable Downsampling: Unlike Max Pooling, which is a fixed operation, a stride-step convolution allows the network to learn the most optimal way to compress spatial information and reduce dimensionality.
2. Preservation of Precise Spatial Information: Max Pooling is designed for translation invariance by discarding exact spatial location in favor of presence. In reconstruction tasks (like VAEs), knowing where a feature was is crucial for the decoder to map it back correctly.
3. Differentiability: While Max Pooling is sub-differentiable, strided convolutions provide smoother gradients throughout the training process, which helps in optimizing the bottleneck latent space.
4. Information Density: Strided convolutions help maintain a more representative summary of the input signal, whereas pooling can lead to "aliasing" or loss of fine-grained detail necessary for high-quality image synthesis.

Bilinear Interpolation vs. Transposed Convolutions

In VAEs and generative models, Bilinear Interpolation followed by a standard convolution is often preferred over Transposed Convolutions for several key reasons:

1. Elimination of Checkerboard Artifacts: Transposed convolutions frequently suffer from "checkerboard artifacts" caused by uneven overlap when the kernel size is not perfectly divisible by the stride. Bilinear upsampling provides a smooth, continuous starting point for the subsequent convolution, effectively preventing these grid-like patterns.
2. More Stable Gradient Flow: By decoupling the upsampling (interpolation) from the feature learning (convolution), the optimization process becomes more stable. The model doesn't have to learn how to both expand the spatial dimensions and refine features simultaneously within a single layer.
3. Parameter Efficiency: Bilinear interpolation is a fixed, parameter-free operation. This allows the network to focus its learnable parameters on the convolutional layers that follow, leading to a more efficient use of the model's capacity.
4. Improved Reconstruction Quality: For tasks requiring high-fidelity image synthesis, starting with a smooth bilinear estimate often leads to more natural-looking textures and edges compared to the sparse initialization inherent in transposed convolutions.

L1 vs. L2 Loss

In VAEs and other generative models, L1 Loss (Mean Absolute Error) is often preferred over L2 Loss (Mean Squared Error) for the following reasons:

1. Robustness to Outliers: L2 loss penalizes large errors quadratically, which forces the model to heavily prioritize correcting outliers. L1 loss has a constant gradient, making it more robust and less likely to be dominated by a few noisy pixels.
2. Sharper Reconstructions: L2 loss tends to produce blurry images because it averages out pixel values to minimize the squared error. L1 loss does not penalize small errors as aggressively, which encourages the model to preserve sharper edges and more distinct transitions.
3. Practical Differentiability: Although L1 is not differentiable at zero, in practice, this is rarely an issue for stochastic gradient descent. The probability of an error being exactly zero is negligible, and most deep learning frameworks (like PyTorch) provide a sub-gradient (usually 0) at that point, allowing training to proceed smoothly.

This lack of differentiability doesn't matter for VAEs where the downstream task is reconstruction. This does matter a lot for regression style problems where coefficients are going to be set to zero for certain ranges of the coefficient space leading to flat gradients!

VAE Diagram

from matplotlib.patches import FancyBboxPatch
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
 
# ════════════════════════════════════════════════════════════════
# SHARED STYLE & HELPERS
# ════════════════════════════════════════════════════════════════
 
BG_COLOR    = '#1a1a2e'
ENC_COLOR   = '#16213e'; ENC_EDGE    = '#0f3460'
DEC_COLOR   = '#1a0a2e'; DEC_EDGE    = '#6a0572'
FLAT_COLOR  = '#0a2e1a'; FLAT_EDGE   = '#0d7040'
COND_COLOR  = '#2e1a0a'; COND_EDGE   = '#d4770a'
ARROW_COLOR = '#e94560'
DIM_COLOR   = '#ffd700'
OP_COLOR    = '#a0c4ff'
TITLE_COLOR = '#ffffff'
NOTE_COLOR  = '#ff9f43'
MATH_COLOR  = '#55efc4'
GRAY        = '#aaaaaa'
GOLD_EDGE   = '#e6b800'
 
 
def draw_block(ax, x, y, w, h, title, ops, in_dim, out_dim,
               face, edge, is_flat=False):
    """Draw an architecture block with shadow, title, ops list, and I/O dims."""
    ax.add_patch(FancyBboxPatch(
        (x + 0.06, y - 0.06), w, h, boxstyle="round,pad=0.06",
        facecolor='black', edgecolor='none', alpha=0.40, zorder=1))
    ax.add_patch(FancyBboxPatch(
        (x, y), w, h, boxstyle="round,pad=0.06",
        facecolor=face, edgecolor=edge, linewidth=2.5, zorder=2))
    top = y + h
    ax.text(x + w/2, top - 0.18, f"In: {in_dim}", ha='center', va='top',
            fontsize=8.5, color=DIM_COLOR, fontweight='bold', zorder=3)
    ax.text(x + w/2, top - 0.46, title, ha='center', va='top',
            fontsize=9.5, color=TITLE_COLOR, fontweight='bold', zorder=3)
    sep = top - 0.66
    ax.plot([x + 0.10, x + w - 0.10], [sep, sep], color=edge, lw=1.0, zorder=3)
    for k, op in enumerate(ops):
        ax.text(x + w/2, sep - 0.06 - k * 0.26, op, ha='center', va='top',
                fontsize=7.8, color=OP_COLOR, zorder=3,
                style='italic' if is_flat else 'normal')
    ax.plot([x + 0.10, x + w - 0.10], [y + 0.28, y + 0.28],
            color=edge, lw=1.0, zorder=3)
    ax.text(x + w/2, y + 0.22, f"Out: {out_dim}", ha='center', va='top',
            fontsize=8.5, color=DIM_COLOR, fontweight='bold', zorder=3)
 
 
def arrow(ax, x1, y1, x2, y2, color=ARROW_COLOR, lw=2.0):
    ax.annotate('', xy=(x2, y2), xytext=(x1, y1),
                arrowprops=dict(arrowstyle='->', color=color,
                                lw=lw, mutation_scale=14),
                zorder=6)
 
 
def arrow_bend(ax, x1, y1, x2, y2, rad=0.3, color=ARROW_COLOR, lw=2.0):
    ax.annotate('', xy=(x2, y2), xytext=(x1, y1),
                arrowprops=dict(arrowstyle='->', color=color, lw=lw,
                                mutation_scale=14,
                                connectionstyle=f'arc3,rad={rad}'),
                zorder=6)
 
 
def note_box(ax, x, y, w, h, text, fontsize=10,
             face='#0f3460', edge=ARROW_COLOR, text_color=DIM_COLOR):
    """Small labeled pill for inputs/outputs."""
    ax.add_patch(FancyBboxPatch(
        (x, y), w, h, boxstyle="round,pad=0.05",
        facecolor=face, edgecolor=edge, lw=1.8, zorder=3))
    ax.text(x + w/2, y + h/2, text, ha='center', va='center',
            fontsize=fontsize, color=text_color, fontweight='bold', zorder=4)
 
 
def dashed_box(ax, x, y, w, h, edge=MATH_COLOR, lw=2.0):
    ax.add_patch(FancyBboxPatch(
        (x, y), w, h, boxstyle="round,pad=0.10",
        facecolor=BG_COLOR, edgecolor=edge,
        linewidth=lw, alpha=0.90, zorder=7, linestyle='--'))
 
 
# ════════════════════════════════════════════════════════════════
# SHARED LAYOUT (encoder / decoder positions)
# ════════════════════════════════════════════════════════════════
 
BW  = 3.00;  BH  = 2.80   # encoder / decoder block size
FBW = 2.80;  FBH = 1.95   # flat (bottleneck) block size
GAP = 0.25                 # horizontal gap between encoder blocks
 
 
def layout(enc_y, dec_y):
    """Return all x/y positions given row baselines."""
    ex = [0.2 + i * (BW + GAP) for i in range(4)]
    stack_x = ex[3] + BW + 0.55
    v_gap = 0.28
    mu_y   = enc_y + (BH - (2*FBH + v_gap))/2 + FBH + v_gap
    logv_y = enc_y + (BH - (2*FBH + v_gap))/2
    lat_x  = stack_x + FBW + 0.50
    lat_mid = logv_y + (mu_y + FBH - logv_y)/2
    lat_y  = lat_mid - FBH/2
    fcd_x  = lat_x
    fcd_y  = dec_y + (BH - FBH)/2
    dx = [ex[3], ex[2], ex[1], ex[0]]
    return dict(ex=ex, dx=dx, stack_x=stack_x,
                mu_y=mu_y, logv_y=logv_y,
                lat_x=lat_x, lat_y=lat_y,
                fcd_x=fcd_x, fcd_y=fcd_y)
 
 
# ────────────────────────────────────────────────────────────────
# Encoder & decoder specs (shared across figures)
# ────────────────────────────────────────────────────────────────
 
ENC_SPECS = [
    ("Enc Block 1",
     ["Conv 3-->64, s=1", "GN * ReLU",
      "Conv 64-->128, s=2", "GN * ReLU", "ResBlock(128)"],
     "64x64x3", "32x32x128"),
    ("Enc Block 2",
     ["Conv 128-->256, s=2", "GN * ReLU", "ResBlock(256)"],
     "32x32x128", "16x16x256"),
    ("Enc Block 3",
     ["Conv 256-->512, s=2", "GN * ReLU", "ResBlock(512)"],
     "16x16x256", "8x8x512"),
    ("Enc Block 4",
     ["Conv 512-->512, s=2", "GN * ReLU", "ResBlock(512)", "Flatten()"],
     "8x8x512", "8192"),
]
 
DEC_SPECS = [
    ("Dec Block 1",
     ["Upsample x2 (4-->8)", "Conv 512-->256",
      "GN * ReLU", "ResBlock(256)"],
     "4x4x512", "8x8x256"),
    ("Dec Block 2",
     ["Upsample x2 (8-->16)", "Conv 256-->128",
      "GN * ReLU", "ResBlock(128)"],
     "8x8x256", "16x16x128"),
    ("Dec Block 3",
     ["Upsample x2 (16-->32)", "Conv 128-->64",
      "GN * ReLU", "ResBlock(64)"],
     "16x16x128", "32x32x64"),
    ("Dec Block 4",
     ["Upsample x2 (32-->64)", "Conv 64-->32",
      "GN * ReLU", "ResBlock(32)",
      "Conv 32-->3, k=1", "Sigmoid()"],
     "32x32x64", "64x64x3"),
]
 
 
def draw_encoder(ax, L, enc_y):
    for i, (t, ops, ind, outd) in enumerate(ENC_SPECS):
        draw_block(ax, L['ex'][i], enc_y, BW, BH, t, ops, ind, outd,
                   ENC_COLOR, ENC_EDGE)
    for i in range(3):
        arrow(ax, L['ex'][i]+BW, enc_y+BH/2, L['ex'][i+1], enc_y+BH/2)
 
 
def draw_bottleneck(ax, L, enc_y):
    draw_block(ax, L['stack_x'], L['mu_y'], FBW, FBH, "fc_mu",
               ["Linear(8192 --> 256)", "-->  μ(x)  ∈ R²⁵⁶"],
               "8192", "μ  (256)", FLAT_COLOR, FLAT_EDGE, is_flat=True)
    draw_block(ax, L['stack_x'], L['logv_y'], FBW, FBH, "fc_logvar",
               ["Linear(8192 --> 256)", "-->  log σ²(x)  ∈ R²⁵⁶"],
               "8192", "log σ²  (256)", FLAT_COLOR, FLAT_EDGE, is_flat=True)
    arrow(ax, L['ex'][3]+BW, enc_y + BH*0.68, L['stack_x'], L['mu_y']+FBH/2)
    arrow(ax, L['ex'][3]+BW, enc_y + BH*0.32, L['stack_x'], L['logv_y']+FBH/2)
 
 
def draw_latent(ax, L):
    draw_block(ax, L['lat_x'], L['lat_y'], FBW, FBH,
               "Reparameterize",
               ["z = μ + σ * ε", "ε ~ N(0, I)"],
               "μ,  log σ²", "z ∈ R²⁵⁶",
               FLAT_COLOR, GOLD_EDGE, is_flat=True)
    ax.add_patch(FancyBboxPatch(
        (L['lat_x'], L['lat_y']), FBW, FBH, boxstyle="round,pad=0.06",
        facecolor='none', edgecolor=GOLD_EDGE, linewidth=3.0, zorder=5))
    arrow(ax, L['stack_x']+FBW, L['mu_y']+FBH/2,
          L['lat_x'], L['lat_y']+FBH*0.72)
    arrow(ax, L['stack_x']+FBW, L['logv_y']+FBH/2,
          L['lat_x'], L['lat_y']+FBH*0.28)
 
 
def draw_fc_decode(ax, L, in_label="z  (256)", in_size=256):
    draw_block(ax, L['fcd_x'], L['fcd_y'], FBW, FBH, "fc_decode",
               [f"Linear({in_size} --> 8192)", "view(−1, 512, 4, 4)"],
               in_label, "4x4x512", FLAT_COLOR, FLAT_EDGE, is_flat=True)
 
 
def draw_decoder(ax, L, dec_y):
    for i, (t, ops, ind, outd) in enumerate(DEC_SPECS):
        draw_block(ax, L['dx'][i], dec_y, BW, BH, t, ops, ind, outd,
                   DEC_COLOR, DEC_EDGE)
    for i in range(3):
        arrow(ax, L['dx'][i], dec_y+BH/2, L['dx'][i+1]+BW, dec_y+BH/2)
    arrow(ax, L['fcd_x'], L['fcd_y']+FBH/2, L['dx'][0]+BW, dec_y+BH/2)
 
 
# ════════════════════════════════════════════════════════════════
# FIGURE 1 — BASE VAE  +  SAMPLING NOTATION
# ════════════════════════════════════════════════════════════════
 
def figure1_base_vae():
    fig, ax = plt.subplots(1, 1, figsize=(32, 15))
    ax.set_xlim(0, 32); ax.set_ylim(0, 15)
    ax.axis('off')
    fig.patch.set_facecolor(BG_COLOR); ax.set_facecolor(BG_COLOR)
 
    enc_y, dec_y = 7.8, 2.2
    L = layout(enc_y, dec_y)
 
    # Title
    ax.text(16.0, 14.55, 'Variational Autoencoder — Architecture & Sampling',
            ha='center', va='center', fontsize=16, color=TITLE_COLOR,
            fontweight='bold', zorder=5)
 
    # Encoder
    draw_encoder(ax, L, enc_y)
    note_box(ax, L['ex'][0]-0.05, enc_y+BH+0.25, BW+0.10, 0.55,
             'Input  x ∈ R⁶⁴ˣ⁶⁴ˣ³')
    arrow(ax, L['ex'][0]+BW/2, enc_y+BH+0.25,
          L['ex'][0]+BW/2, enc_y+BH+0.02)
 
    # Bottleneck
    draw_bottleneck(ax, L, enc_y)
    draw_latent(ax, L)
 
    # Latent --> fc_decode
    arrow(ax, L['lat_x']+FBW/2, L['lat_y'],
          L['lat_x']+FBW/2, L['fcd_y']+FBH)
    draw_fc_decode(ax, L)
 
    # Decoder
    draw_decoder(ax, L, dec_y)
    note_box(ax, L['dx'][3]-0.05, dec_y-0.85, BW+0.10, 0.55,
             'Reconstruction  x_hat ∈ R⁶⁴ˣ⁶⁴ˣ³',
             face='#2e0a2e')
    arrow(ax, L['dx'][3]+BW/2, dec_y,
          L['dx'][3]+BW/2, dec_y-0.28)
 
    # ── SAMPLING NOTATION (right side) ──
    sx = L['lat_x'] + FBW + 0.40
    sy = L['lat_y'] - 0.40
    sw, sh = 5.0, 2.80
    dashed_box(ax, sx, sy, sw, sh, edge=NOTE_COLOR)
 
    ax.text(sx + sw/2, sy + sh - 0.12,
            '⚡  Generation  (no encoder needed)',
            ha='center', va='top', fontsize=11, color=NOTE_COLOR,
            fontweight='bold', zorder=8)
 
    lines = [
        ("Direct sampling:", TITLE_COLOR, True),
        ("  1.  z ~ N(0, I)           sample from prior", MATH_COLOR, False),
        ("  2.  fc_decode(z)          project to 4x4x512", MATH_COLOR, False),
        ("  3.  Decoder(.) -> x_hat   upsample to image", MATH_COLOR, False),
        ("", MATH_COLOR, False),
        ("Two-stage sampling:", TITLE_COLOR, True),
        ("  1.  z2 ~ N(0, I)          sample stage-2 prior", MATH_COLOR, False),
        ("  2.  VAE2.decode(z2) -> z1  map to stage-1 space", MATH_COLOR, False),
        ("  3.  fc_decode(z1) -> Dec -> x_hat", MATH_COLOR, False),
    ]
    for k, (txt, col, bold) in enumerate(lines):
        ax.text(sx + 0.20, sy + sh - 0.48 - k * 0.24, txt,
                ha='left', va='top', fontsize=8.5, color=col,
                fontweight='bold' if bold else 'normal',
                family='monospace', zorder=8)
 
    # ── LOSS NOTATION (below decoder) ──
    lx = L['dx'][3] + BW + 0.50
    ly = dec_y - 0.85
    lw, lh = 7.0, 1.50
    dashed_box(ax, lx, ly, lw, lh, edge=MATH_COLOR)
 
    ax.text(lx + lw/2, ly + lh - 0.10,
            'Loss  (negative ELBO)',
            ha='center', va='top', fontsize=11, color=MATH_COLOR,
            fontweight='bold', zorder=8)
    loss_lines = [
        "L  =  L1(x, x_hat)  +  lam * LPIPS(x, x_hat)  +  D_KL( Q(z|x) || N(0,I) )",
        "       |--- reconstruction ---|                   |-- regularization --|",
        "  lam  determined by gradient balancing at decoder's last layer",
    ]
    for k, line in enumerate(loss_lines):
        ax.text(lx + 0.20, ly + lh - 0.45 - k * 0.28, line,
                ha='left', va='top', fontsize=8.5, color=MATH_COLOR,
                family='monospace', zorder=8)
 
    # Section labels
    enc_mid = (L['ex'][0] + L['ex'][3] + BW) / 2
    bot_mid = (L['stack_x'] + L['lat_x'] + FBW) / 2
    dec_mid = (L['dx'][3] + L['dx'][0] + BW) / 2
    for txt, xp, yp in [("◀  ENCODER  ▶", enc_mid, 11.4),
                         ("◀  BOTTLENECK  ▶", bot_mid, 11.4),
                         ("◀  DECODER  ▶", dec_mid, 1.4)]:
        ax.text(xp, yp, txt, ha='center', va='center', fontsize=10,
                color=GRAY, fontstyle='italic', zorder=5)
 
    # Legend
    legend_items = [
        mpatches.Patch(facecolor=ENC_COLOR, edgecolor=ENC_EDGE, label='Encoder Block'),
        mpatches.Patch(facecolor=FLAT_COLOR, edgecolor=FLAT_EDGE, label='Bottleneck (Linear)'),
        mpatches.Patch(facecolor=FLAT_COLOR, edgecolor=GOLD_EDGE, label='Latent Vector z'),
        mpatches.Patch(facecolor=DEC_COLOR, edgecolor=DEC_EDGE, label='Decoder Block'),
    ]
    ax.legend(handles=legend_items, loc='lower right', fontsize=9,
              facecolor=BG_COLOR, edgecolor='white',
              labelcolor='white', framealpha=0.85)
 
    plt.tight_layout()
    plt.savefig('fig1_vae_base.png', dpi=150, bbox_inches='tight',
                facecolor=fig.get_facecolor())
    plt.show()
    print("Saved: fig1_vae_base.png")

figure1_base_vae()

Saved: fig1_vae_base.png

VAE Code

import torch
import torch.nn as nn
import torch.nn.functional as F


class ResBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
            nn.GroupNorm(32, channels),
            nn.ReLU(),
            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
            nn.GroupNorm(32, channels),
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.block(x) + x)


class VAE(nn.Module):
    def __init__(self, hidden_size=256):
        super(VAE, self).__init__()
        self.hidden_size = hidden_size

        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1), # 64→64
            nn.GroupNorm(32, 32),
            nn.ReLU(),

            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),  # 64→32
            nn.GroupNorm(32, 64),
            nn.ReLU(),
            ResBlock(64),

            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),  # 32→16
            nn.GroupNorm(32, 128),
            nn.ReLU(),
            ResBlock(128),

            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),  # 16→8
            nn.GroupNorm(32, 256),
            nn.ReLU(),
            ResBlock(256),

            nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),  # 8→4
            nn.GroupNorm(32, 512),
            nn.ReLU(),
            ResBlock(512),

            nn.Flatten(),
        )

        self.fc_mu = nn.Linear(512 * 4 * 4, hidden_size)
        self.fc_logvar = nn.Linear(512 * 4 * 4, hidden_size)
        self.fc_decode = nn.Linear(hidden_size, 512 * 4 * 4)

        self.decoder = nn.Sequential(

            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),  # 4→8
            nn.Conv2d(512, 256, kernel_size=3, padding=1),
            nn.GroupNorm(32, 256),
            nn.ReLU(),
            ResBlock(256),

            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),  # 8→16
            nn.Conv2d(256, 128, kernel_size=3, padding=1),
            nn.GroupNorm(32, 128),
            nn.ReLU(),
            ResBlock(128),

            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),  # 16→32
            nn.Conv2d(128, 64, kernel_size=3, padding=1),
            nn.GroupNorm(32, 64),
            nn.ReLU(),
            ResBlock(64),

            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),  # 32→64
            nn.Conv2d(64, 32, kernel_size=3, padding=1),
            nn.GroupNorm(32, 32),
            nn.ReLU(),
            ResBlock(32),

            nn.Conv2d(32, 3, kernel_size=1, stride=1),  # final projection
            nn.Sigmoid(),
        )

    def get_last_decoder_layer(self):
        """Return the weight of the final Conv2d before Sigmoid for adaptive weight calc."""
        # decoder[-2] is the Conv2d(32, 3, 1), decoder[-1] is Sigmoid
        return self.decoder[-2].weight

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

    def forward(self, x):
        encoded = self.encoder(x)
        mu = self.fc_mu(encoded)
        logvar = self.fc_logvar(encoded)
        logvar = torch.clamp(logvar, min=-10, max=10)
        z = self.reparameterize(mu, logvar)
        z_projected = self.fc_decode(z).view(-1, 512, 4, 4)
        reconstruction = self.decoder(z_projected)
        return reconstruction, mu, logvar

Loss Function

Two pieces for an image $\mathbf X$ and it's reconstruction after passing through the bottleneck $\hat{\mathbf X}$. Assuming $M$ pixels in the image and $K$ dimensions at the bottleneck:

• Negative log-likelihood = $\sum \limits_{m = 1}^M | x_m - \hat{x}_m|$

• KL Divergence w/ $K$ dim standard normal prior = $\sum \limits_{k = 1}^K -\frac{1}{2} \left[1 + \log \sigma^2_k - \mu_k^2 - \sigma^2_k \right]$

ELBO derivation says just add 'em together and minimize!

def generator_loss(reconstruction, x, mu, logvar, beta=1.0, free_bits=0.1):

    #n_pixels = x.shape[1] * x.shape[2] * x.shape[3]

    # L1 at ELBO scale (sum over pixels, mean over batch)
    # This is the natural negative likelihood loss!  Must be this scale for the KL
    # Term to Balance

    l1_loss = F.l1_loss(reconstruction, x, reduction='sum') / x.shape[0]

    nll_loss = l1_loss

    # KL divergence (sum over latent dims, mean over batch)
    # Normal Posterior vs. Normal Prior
    # Independent, so final solution at log scale is sum
    KLD_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
    # Any KL loss by dim less than free_bits is coded as free_bits
    # Removes incentive at the margins to exactly match prior
    KLD = torch.sum(torch.clamp(KLD_per_dim, min=free_bits)) / x.shape[0]

    total = nll_loss + (KLD * beta)

    return total, nll_loss.item(), KLD.item()

Model Setup

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
hidden_size = 256

model = VAE(hidden_size=hidden_size).to(device)

print(model)

VAE(
  (encoder): Sequential(
    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): GroupNorm(32, 32, eps=1e-05, affine=True)
    (2): ReLU()
    (3): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (4): GroupNorm(32, 64, eps=1e-05, affine=True)
    (5): ReLU()
    (6): ResBlock(
      (block): Sequential(
        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 64, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 64, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (7): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (8): GroupNorm(32, 128, eps=1e-05, affine=True)
    (9): ReLU()
    (10): ResBlock(
      (block): Sequential(
        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 128, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 128, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (11): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (12): GroupNorm(32, 256, eps=1e-05, affine=True)
    (13): ReLU()
    (14): ResBlock(
      (block): Sequential(
        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 256, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 256, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (15): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (16): GroupNorm(32, 512, eps=1e-05, affine=True)
    (17): ReLU()
    (18): ResBlock(
      (block): Sequential(
        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 512, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 512, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (19): Flatten(start_dim=1, end_dim=-1)
  )
  (fc_mu): Linear(in_features=8192, out_features=256, bias=True)
  (fc_logvar): Linear(in_features=8192, out_features=256, bias=True)
  (fc_decode): Linear(in_features=256, out_features=8192, bias=True)
  (decoder): Sequential(
    (0): Upsample(scale_factor=2.0, mode='bilinear')
    (1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (2): GroupNorm(32, 256, eps=1e-05, affine=True)
    (3): ReLU()
    (4): ResBlock(
      (block): Sequential(
        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 256, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 256, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (5): Upsample(scale_factor=2.0, mode='bilinear')
    (6): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (7): GroupNorm(32, 128, eps=1e-05, affine=True)
    (8): ReLU()
    (9): ResBlock(
      (block): Sequential(
        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 128, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 128, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (10): Upsample(scale_factor=2.0, mode='bilinear')
    (11): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (12): GroupNorm(32, 64, eps=1e-05, affine=True)
    (13): ReLU()
    (14): ResBlock(
      (block): Sequential(
        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 64, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 64, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (15): Upsample(scale_factor=2.0, mode='bilinear')
    (16): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (17): GroupNorm(32, 32, eps=1e-05, affine=True)
    (18): ReLU()
    (19): ResBlock(
      (block): Sequential(
        (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 32, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 32, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (20): Conv2d(32, 3, kernel_size=(1, 1), stride=(1, 1))
    (21): Sigmoid()
  )
)

AdamW optimizer. Learning rate is set using a cosine annealing procedure with warm restarts every 50x iterations

from torch import optim

# Optimization
# AdamW with default weight decay
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
#Long burn
num_epochs = 500


# Cosine w/ Warm Restarts
scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
    optimizer,
    T_0=50,          # First cycle is 50 epochs
    T_mult=2,        # Double the length of each subsequent cycle
    eta_min=1e-6     # Let it decay very close to zero
)

Slowly turn on the KL annealing to prevent immediate collapse. 5 iterations with no regularization, 10 slowly ramping from 0 KL penalty to the correct KL penalty!

## KL Annealing

import numpy as np

def get_beta(epoch, warmup_start=10, warmup_end=20, beta_max=1.0):
    if epoch <= warmup_start:
        return 0.0
    elif epoch <= warmup_end:
        t = epoch - warmup_start
        warmup_steps = warmup_end - warmup_start
        return beta_max * 0.5 * (1 - np.cos(np.pi * t / warmup_steps))
    else:
        return beta_max

A function to allow us to visualize evolution of image reconstructions and to show the decoded form of generated images from the latent space:

def visualize_results(model, epoch, temperature=0.7):
    model.eval()
    with torch.no_grad():
        train_batch, _ = next(iter(train_loader))
        train_batch = train_batch[:5].to(device)
        train_recon, _, _ = model(train_batch)

        val_batch, _ = next(iter(val_loader))
        val_batch = val_batch[:5].to(device)
        val_recon, _, _ = model(val_batch)

        z_random = torch.randn(10, hidden_size).to(device) * temperature
        z_projected = model.fc_decode(z_random).view(-1, 512, 4, 4)
        generated = model.decoder(z_projected)

        fig, axes = plt.subplots(3, 10, figsize=(20, 6))

        for i in range(5):
            axes[0, i * 2].imshow(train_batch[i].cpu().permute(1, 2, 0))
            axes[0, i * 2].set_title("Train Ori")
            axes[0, i * 2 + 1].imshow(train_recon[i].cpu().permute(1, 2, 0))
            axes[0, i * 2 + 1].set_title("Train Recon")

            axes[1, i * 2].imshow(val_batch[i].cpu().permute(1, 2, 0))
            axes[1, i * 2].set_title("Val Ori")
            axes[1, i * 2 + 1].imshow(val_recon[i].cpu().permute(1, 2, 0))
            axes[1, i * 2 + 1].set_title("Val Recon")

        for i in range(10):
            axes[2, i].imshow(generated[i].cpu().permute(1, 2, 0))
            axes[2, i].set_title(f"Gen {i + 1}")

        for ax in axes.flatten():
            ax.axis("off")

        plt.suptitle(f"Epoch {epoch}")
        plt.tight_layout()
        plt.show()

Now, our training loop (with some additional niceties that will make for better visualization):

from tqdm import tqdm
import numpy as np
import torch
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import ipywidgets as widgets
from IPython.display import display
import base64, io
from PIL import Image as PILImage

# =============================================================================
# DASHBOARD SETUP  — Plotly + ipywidgets (survives nbconvert → HTML)
# =============================================================================

history = {'epoch': [], 'total': [], 'nll': [], 'kld': []}
snapshot_store = {}   # epoch → list of 30 PIL images

# ── Loss figure (Plotly FigureWidget — updates in-place) ──────────
loss_fig = make_subplots(
    rows=1, cols=3,
    subplot_titles=('Total ELBO Loss', 'NLL  (Reconstruction)', 'KL Divergence'),
)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#e94560', width=2),
                               name='Total'), row=1, col=1)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#a0c4ff', width=2),
                               name='NLL'),   row=1, col=2)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#ffd700', width=2),
                               name='KLD'),   row=1, col=3)
loss_fig.update_layout(
    height=320, paper_bgcolor='#1a1a2e', plot_bgcolor='#0d0d1a',
    font=dict(color='white'), showlegend=False,
    margin=dict(t=40, b=30, l=40, r=20),
)
for i in range(1, 4):
    loss_fig.update_xaxes(gridcolor='#222244', title_text='Epoch',
                           row=1, col=i)
    loss_fig.update_yaxes(gridcolor='#222244', row=1, col=i)

loss_widget = go.FigureWidget(loss_fig)


def _update_loss_plots():
    ep  = history['epoch']
    with loss_widget.batch_update():
        loss_widget.data[0].x = ep;  loss_widget.data[0].y = history['total']
        loss_widget.data[1].x = ep;  loss_widget.data[1].y = history['nll']
        loss_widget.data[2].x = ep;  loss_widget.data[2].y = history['kld']


# ── Snapshot viewer (ipywidgets — buttons + image grid) ───────────
def _tensor_to_b64(img_np):
    """Convert H×W×3 float32 numpy array → base64 PNG string."""
    arr = (np.clip(img_np, 0, 1) * 255).astype(np.uint8)
    buf = io.BytesIO()
    PILImage.fromarray(arr).save(buf, format='PNG')
    return 'data:image/png;base64,' + base64.b64encode(buf.getvalue()).decode()

ROW_LABELS = ['Train  Orig / Recon', 'Val  Orig / Recon', 'Generated']

# 30 image widgets arranged in 3 rows × 10 cols
img_widgets = [[widgets.Image(format='png', width=80, height=80)
                for _ in range(10)] for _ in range(3)]

row_boxes = []
for r in range(3):
    label = widgets.Label(
        value=ROW_LABELS[r],
        layout=widgets.Layout(width='130px', display='flex',
                              align_items='center'))
    row_box = widgets.HBox(
        [label] + img_widgets[r],
        layout=widgets.Layout(align_items='center'))
    row_boxes.append(row_box)

snap_label = widgets.Label(
    value='Snapshot — no data yet',
    layout=widgets.Layout(margin='4px 0 4px 0'))
snap_label.style.font_size = '13px'

snap_idx = [0]

btn_prev = widgets.Button(description='◀  Prev',
                          button_style='',
                          layout=widgets.Layout(width='100px'))
btn_next = widgets.Button(description='Next  ▶',
                          button_style='',
                          layout=widgets.Layout(width='100px'))
for b in (btn_prev, btn_next):
    b.style.button_color = '#16213e'
    b.style.text_color   = 'white'

epoch_dropdown = widgets.Dropdown(
    options=[],
    description='Epoch:',
    layout=widgets.Layout(width='150px'))


def _render_snapshot(epoch):
    imgs = snapshot_store.get(epoch)
    if imgs is None:
        return
    for r in range(3):
        for c in range(10):
            b64 = _tensor_to_b64(imgs[r * 10 + c])
            # FigureWidget images need raw bytes
            raw = base64.b64decode(b64.split(',')[1])
            img_widgets[r][c].value = raw
    snap_label.value = f'Snapshot — Epoch {epoch}'


def _on_prev(_b):
    keys = sorted(snapshot_store.keys())
    if not keys: return
    snap_idx[0] = max(0, snap_idx[0] - 1)
    epoch_dropdown.value = keys[snap_idx[0]]


def _on_next(_b):
    keys = sorted(snapshot_store.keys())
    if not keys: return
    snap_idx[0] = min(len(keys) - 1, snap_idx[0] + 1)
    epoch_dropdown.value = keys[snap_idx[0]]


def _on_dropdown_change(change):
    if change['name'] == 'value' and change['new'] is not None:
        keys = sorted(snapshot_store.keys())
        snap_idx[0] = keys.index(change['new'])
        _render_snapshot(change['new'])


btn_prev.on_click(_on_prev)
btn_next.on_click(_on_next)
epoch_dropdown.observe(_on_dropdown_change)

nav_bar      = widgets.HBox([btn_prev, epoch_dropdown, btn_next])
snap_viewer  = widgets.VBox([snap_label, nav_bar] + row_boxes)
dashboard    = widgets.VBox([loss_widget, snap_viewer])
display(dashboard)


def _capture_snapshot(epoch):
    model.eval()
    imgs = []
    with torch.no_grad():
        tb, _ = next(iter(train_loader))
        tb = tb[:5].to(device);  tr, _, _ = model(tb)
        for i in range(5):
            imgs.append(tb[i].cpu().permute(1,2,0).numpy())
            imgs.append(tr[i].cpu().permute(1,2,0).numpy())

        vb, _ = next(iter(val_loader))
        vb = vb[:5].to(device);  vr, _, _ = model(vb)
        for i in range(5):
            imgs.append(vb[i].cpu().permute(1,2,0).numpy())
            imgs.append(vr[i].cpu().permute(1,2,0).numpy())

        z_rand = torch.randn(10, hidden_size).to(device) * 0.7
        gen = model.decoder(model.fc_decode(z_rand).view(-1, 512, 4, 4))
        for i in range(10):
            imgs.append(gen[i].cpu().permute(1,2,0).numpy())

    snapshot_store[epoch] = imgs
    # Update dropdown options
    epoch_dropdown.options = sorted(snapshot_store.keys())
    snap_idx[0] = len(snapshot_store) - 1
    epoch_dropdown.value = epoch
    model.train()

# =============================================================================
# TRAINING LOOP  — plain VAE, single optimizer
# =============================================================================

best_avg_loss = float('inf')
global_step   = 0

for epoch in range(1, num_epochs + 1):
    model.train()
    train_loss = 0.0
    total_nll  = 0.0
    total_kld  = 0.0

    beta = get_beta(epoch)
    #beta = 1
    pbar = tqdm(train_loader,
                desc=f"Epoch {epoch:3d}/{num_epochs}", leave=False)

    for batch_idx, (data, _) in enumerate(pbar):
        data = data.to(device)

        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)

        loss, nll, kld = generator_loss(
            reconstruction = recon_batch,
            x              = data,
            mu             = mu,
            logvar         = logvar,
            beta           = beta,
            free_bits      = 0.1,
        )

        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)
        optimizer.step()

        train_loss += loss.item()
        total_nll  += nll
        total_kld  += kld
        global_step += 1

        pbar.set_postfix({
            'Loss': f'{train_loss / (batch_idx+1):.1f}',
            'NLL':  f'{total_nll  / (batch_idx+1):.1f}',
            'KLD':  f'{total_kld  / (batch_idx+1):.1f}',
            'β':    f'{beta:.3f}',
        })

    # ── Epoch averages ────────────────────────────────────────────
    n          = len(train_loader)
    avg_loss   = train_loss / n
    avg_nll    = total_nll  / n
    avg_kld    = total_kld  / n
    current_lr = optimizer.param_groups[0]['lr']

    # ── Checkpoint on best training ELBO ─────────────────────────
    if avg_loss < best_avg_loss:
        best_avg_loss = avg_loss
        torch.save({
            "epoch":            epoch,
            "global_step":      global_step,
            "model_state_dict": model.state_dict(),
            "opt_state_dict":   optimizer.state_dict(),
            "best_avg_loss":    best_avg_loss,
            "hidden_size":      model.hidden_size,
        }, "vae_best.pth")

    scheduler.step()

    # ── Live dashboard update ─────────────────────────────────────
    if epoch >= 20:
        history['epoch'].append(epoch)
        history['total'].append(avg_loss)
        history['nll']  .append(avg_nll)
        history['kld']  .append(avg_kld)
        _update_loss_plots()

    # ── Snapshot every 10 epochs ──────────────────────────────────
    if epoch % 10 == 0:
        _capture_snapshot(epoch)

VBox(children=(FigureWidget({
    'data': [{'line': {'color': '#e94560', 'width': 2},
              'mode': 'l…

Model 2: VAE + Perceptual Loss

Our reconstructions and generated images seem right - they look like dogs and cats. But, it's hard to actually say if they are or aren't because the images are incredibly blurry!

This is a problem with our choice of loss function - through the ELBO only requires us to minimize NLL + KL, minimizing the NLL leads to behavior that really penalizes shifts from the original position. As an example, we want to construct our latent space in such a way that cat in returns a cat out. Since our loss function is per-pixel loss, a cat outline shifted just one pixel to the right drastically increases the loss! To the human eye, though, a cat that is one pixel wider than it should be is unimportant - a cat is a cat is a cat.

As with CNNs and transitioning to U-Nets, images have two properties - what and where. Our VAE as is does a decent job of handling both, but the loss structure leads the VAE to prioritize where more than what!

An easy solution is to try to measure what and add a little where. This can be accomplished by combining the per-pixel loss with perceptual loss.

Perceptual loss leverages the convolutional backbone of a CNN trained for classification and says that two identical images should have the same feature maps throughout the convolutional backbone. For a 64 x 64 x 3 image and a 64 x 64 x 3 reconstruction, we should see that the feature map for this image given in the first layer of a CNN backbone (32 height x 32 width x 64 channels, for example) should be the same for both images if they are identical.

As CNN backbones are constructed to induce hierarchical feature arrangements - local patterns early and global patterns late - the difference between the feature maps for a true image and its reconstruction at each layer contains different levels of information for what and where. Something like the absolute value of the differences of these feature maps would provide more information about what is actually different between the images, regardless of where.

Let PL denote a measure of perceptual loss. Then, we can define a loss function to minimize as:

$$ \text{PL} + \text{NLL} + \text{KLD} $$

The most common way to measure perceptual loss is using LPIPS (Learned Perceptual Image Patch Similarity; Zhang et al. 2018). LPIPS uses a pretrained VGG ImageNet backbone and passes both an image and its reconstruction through and captures all intermediate representations. Then,

1. Channel-wise normalization: At each layer, the feature activations are normalized to unit length along the channel dimension. This ensures that all channels contribute equally regardless of their raw activation magnitudes — without this step, a channel with activations in the range [0, 100] would dominate the loss compared to a channel in [0, 1], regardless of which is more perceptually meaningful.

2. Learned per-channel weighting: Each channel is multiplied by a scalar weight that was trained on human perceptual judgments. Zhang et al. collected a dataset (BAPPS) where humans were shown image triplets and asked "which of these two patches is more similar to the reference?" The learned weights encode which VGG channels humans actually care about — amplifying channels that detect perceptually important features (edges, textures, object parts) and suppressing channels that detect things humans are insensitive to.

3. Squared difference, averaged spatially: At each layer $l$, the squared difference of the weighted, normalized features is averaged over spatial positions:

$$d_l = \frac{1}{H_l W_l} \sum_{h,w} \left\| w_l \odot \left( \hat{x}_l^{h,w} - \hat{y}_l^{h,w} \right) \right\|_2^2$$

4. Sum across layers: The final LPIPS score sums across all extracted layers:

$$\text{LPIPS}(x, \hat{x}) = \sum_l d_l$$

The critical insight is that VGG features capture spatial relationships between pixels. Early layers detect edges and textures. Later layers detect object parts and compositions. Two images can have low LPIPS loss even if they differ pixel-by-pixel — as long as they share the same visual structure. A cat shifted one pixel to the right has nearly identical VGG features to the original cat, so the perceptual loss is close to zero even though the per-pixel loss is large.

This frees the decoder from the tyranny of pixel-perfect alignment. Instead of hedging its bets and producing a blurry average, it can commit to sharp, coherent outputs — because the loss function now rewards perceptual similarity rather than pixel identity.

The tradeoff: LPIPS has no probabilistic interpretation. The ELBO derivation gives us NLL + KLD, and bolting LPIPS onto that is theoretically unprincipled — we're adding a term that doesn't correspond to any likelihood model. And this can create problems! If the perceptual loss on its natural scale is much bigger than the NLL, then the reconstructions will prioritize structure over location. If the opposite is true, then the reconstructions will prioritize location over structure. And if these two things are really big or small and not appropriately weighted, then the KL divergence could completely dominate the optimization procedure (posterior collapse to the prior) or get completely lost (shift to deterministic autoencoders). The best solution is to find a principled way to balance these terms through gradient balancing.

Gradient Balancing

The core insight is that loss magnitudes don't matter — gradient magnitudes do. The optimizer doesn't see loss values directly. It sees gradients. Two loss terms can have identical values but wildly different gradient norms at the decoder's weights, meaning one term dominates the parameter update while the other has no practical influence.

So instead of comparing loss values and guessing a weight, we compare gradient norms at a specific point in the network — typically the decoder's last layer, since that's where all reconstruction signals converge before producing the output image.

Given two loss terms $\mathcal{L}_A$ and $\mathcal{L}_B$ that we want to balance, we compute:

$$\lambda = \frac{\left\| \nabla_{w} \mathcal{L}_A \right\|}{\left\| \nabla_{w} \mathcal{L}_B \right\| + \epsilon}$$

where $w$ are the weights of the decoder's last layer and $\epsilon$ is a small constant for numerical stability. Then we use $\lambda$ to scale $\mathcal{L}_B$:

$$\mathcal{L}_{\text{total}} = \mathcal{L}_A + \lambda \cdot \mathcal{L}_B + \mathcal{L}_{\text{KLD}}$$

When $\lambda$ is large, it means $\mathcal{L}_A$ produces much larger gradients than $\mathcal{L}_B$, so $\mathcal{L}_B$ needs to be amplified to compete. When $\lambda$ is small, $\mathcal{L}_B$ is already strong enough relative to $\mathcal{L}_A$.

Why this works: $\lambda$ is recomputed at every training step, so it adapts automatically as the loss landscape changes during training. Early in training when reconstructions are poor, the gradient norms might favor one term heavily. Later when reconstructions improve, the balance shifts. A fixed weight can't track this — gradient balancing can.

Why the last decoder layer? This is the final point where the network can influence the output image. If a loss term has zero gradient at this layer, it has zero influence on what pixels are produced, regardless of how large the loss value is. Measuring gradients here tells us exactly how much each term is actually steering the decoder's output.

In PyTorch, the computation uses torch.autograd.grad to peek at the gradients without performing a full backward pass:

def compute_adaptive_weight(loss_a, loss_b, last_layer_weight):
    grads_a = torch.autograd.grad(loss_a, last_layer_weight, retain_graph=True)[0]
    grads_b = torch.autograd.grad(loss_b, last_layer_weight, retain_graph=True)[0]
    weight = torch.norm(grads_a) / (torch.norm(grads_b) + 1e-4)
    weight = torch.clamp(weight, 0.0, 1e4).detach()
    return weight

Note the .detach() — $\lambda$ is treated as a constant for the backward pass. We don't want to differentiate through the balancing computation itself; we just want it to set the scale.

For our VAE with perceptual loss, we set $\mathcal{L}_A$ = NLL (the ELBO-grounded term we trust) and $\mathcal{L}_B$ = LPIPS (the term with no natural scale). This makes NLL the anchor and dynamically scales LPIPS to have proportional gradient influence at the decoder's output layer. The KLD term is left unscaled because it operates on the encoder's output (mu and logvar), not the decoder — its gradients flow through a completely different path and are already at the correct ELBO scale relative to NLL.

Let's start building this improved VAE! We'll use the exact same VAE architecture as before. The only thing that changes is the Loss computation.

First, we need to create the LPIPS module that will allow us to compute the perceptual loss:

import lpips

class LPIPSLoss(nn.Module):
    """LPIPS perceptual loss. Expects [0, 1] inputs, converts to [-1, 1] internally."""
    def __init__(self):
        super().__init__()
        self.lpips = lpips.LPIPS(net='vgg')
        for param in self.parameters():
            param.requires_grad = False

    def forward(self, reconstruction, target):
        return self.lpips(reconstruction * 2 - 1, target * 2 - 1).mean()

Next, we need to set up the adaptive weight function:

def compute_adaptive_weight(loss_a, loss_b, last_layer_weight):
    grads_a = torch.autograd.grad(loss_a, last_layer_weight, retain_graph=True)[0]
    grads_b = torch.autograd.grad(loss_b, last_layer_weight, retain_graph=True)[0]
    weight = torch.norm(grads_a) / (torch.norm(grads_b) + 1e-4)
    weight = torch.clamp(weight, 0.0, 1e4).detach()
    return weight

Then, we need to set up our generator loss function:

def perceptual_generator_loss(reconstruction, x, mu, logvar, perceptual_loss_fn, last_layer_weight, beta=1.0, free_bits=0.1):
    # NLL Loss (L1 at ELBO scale)
    nll_loss = F.l1_loss(reconstruction, x, reduction='sum') / x.shape[0]

    # Perceptual Loss via LPIPS
    p_loss = perceptual_loss_fn(reconstruction, x)

    # KL Divergence
    KLD_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
    KLD = torch.sum(torch.clamp(KLD_per_dim, min=free_bits)) / x.shape[0]

    # Gradient balancing: compute adaptive weight to scale perceptual loss to NLL scale
    d_weight = compute_adaptive_weight(nll_loss, p_loss, last_layer_weight)

    wp_loss = d_weight * p_loss

    # Total ELBO with balanced perceptual term
    total = nll_loss + (wp_loss) + (beta * KLD)

    return total, nll_loss.item(), KLD.item(), wp_loss.item(), d_weight.item()

As before, we need to instantiate our models (VAE and LPIPS) and start up a new optimizer.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
hidden_size = 256

model = VAE(hidden_size=hidden_size).to(device)

print(model)

VAE(
  (encoder): Sequential(
    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): GroupNorm(32, 32, eps=1e-05, affine=True)
    (2): ReLU()
    (3): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (4): GroupNorm(32, 64, eps=1e-05, affine=True)
    (5): ReLU()
    (6): ResBlock(
      (block): Sequential(
        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 64, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 64, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (7): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (8): GroupNorm(32, 128, eps=1e-05, affine=True)
    (9): ReLU()
    (10): ResBlock(
      (block): Sequential(
        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 128, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 128, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (11): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (12): GroupNorm(32, 256, eps=1e-05, affine=True)
    (13): ReLU()
    (14): ResBlock(
      (block): Sequential(
        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 256, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 256, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (15): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (16): GroupNorm(32, 512, eps=1e-05, affine=True)
    (17): ReLU()
    (18): ResBlock(
      (block): Sequential(
        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 512, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 512, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (19): Flatten(start_dim=1, end_dim=-1)
  )
  (fc_mu): Linear(in_features=8192, out_features=256, bias=True)
  (fc_logvar): Linear(in_features=8192, out_features=256, bias=True)
  (fc_decode): Linear(in_features=256, out_features=8192, bias=True)
  (decoder): Sequential(
    (0): Upsample(scale_factor=2.0, mode='bilinear')
    (1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (2): GroupNorm(32, 256, eps=1e-05, affine=True)
    (3): ReLU()
    (4): ResBlock(
      (block): Sequential(
        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 256, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 256, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (5): Upsample(scale_factor=2.0, mode='bilinear')
    (6): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (7): GroupNorm(32, 128, eps=1e-05, affine=True)
    (8): ReLU()
    (9): ResBlock(
      (block): Sequential(
        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 128, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 128, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (10): Upsample(scale_factor=2.0, mode='bilinear')
    (11): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (12): GroupNorm(32, 64, eps=1e-05, affine=True)
    (13): ReLU()
    (14): ResBlock(
      (block): Sequential(
        (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 64, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 64, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (15): Upsample(scale_factor=2.0, mode='bilinear')
    (16): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (17): GroupNorm(32, 32, eps=1e-05, affine=True)
    (18): ReLU()
    (19): ResBlock(
      (block): Sequential(
        (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): GroupNorm(32, 32, eps=1e-05, affine=True)
        (2): ReLU()
        (3): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): GroupNorm(32, 32, eps=1e-05, affine=True)
      )
      (relu): ReLU()
    )
    (20): Conv2d(32, 3, kernel_size=(1, 1), stride=(1, 1))
    (21): Sigmoid()
  )
)

perceptual_loss_fn = LPIPSLoss().to(device)

print(perceptual_loss_fn)

Setting up [LPIPS] perceptual loss: trunk [vgg], v[0.1], spatial [off]

/home/kmcalist/.local/lib/python3.10/site-packages/torchvision/models/_utils.py:208: UserWarning:

The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.

/home/kmcalist/.local/lib/python3.10/site-packages/torchvision/models/_utils.py:223: UserWarning:

Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=VGG16_Weights.IMAGENET1K_V1`. You can also use `weights=VGG16_Weights.DEFAULT` to get the most up-to-date weights.

Loading model from: /home/kmcalist/.local/lib/python3.10/site-packages/lpips/weights/v0.1/vgg.pth
LPIPSLoss(
  (lpips): LPIPS(
    (scaling_layer): ScalingLayer()
    (net): vgg16(
      (slice1): Sequential(
        (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): ReLU(inplace=True)
        (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (3): ReLU(inplace=True)
      )
      (slice2): Sequential(
        (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (6): ReLU(inplace=True)
        (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (8): ReLU(inplace=True)
      )
      (slice3): Sequential(
        (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (11): ReLU(inplace=True)
        (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (13): ReLU(inplace=True)
        (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (15): ReLU(inplace=True)
      )
      (slice4): Sequential(
        (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (18): ReLU(inplace=True)
        (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (20): ReLU(inplace=True)
        (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (22): ReLU(inplace=True)
      )
      (slice5): Sequential(
        (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (25): ReLU(inplace=True)
        (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (27): ReLU(inplace=True)
        (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (29): ReLU(inplace=True)
      )
    )
    (lin0): NetLinLayer(
      (model): Sequential(
        (0): Dropout(p=0.5, inplace=False)
        (1): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      )
    )
    (lin1): NetLinLayer(
      (model): Sequential(
        (0): Dropout(p=0.5, inplace=False)
        (1): Conv2d(128, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      )
    )
    (lin2): NetLinLayer(
      (model): Sequential(
        (0): Dropout(p=0.5, inplace=False)
        (1): Conv2d(256, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      )
    )
    (lin3): NetLinLayer(
      (model): Sequential(
        (0): Dropout(p=0.5, inplace=False)
        (1): Conv2d(512, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      )
    )
    (lin4): NetLinLayer(
      (model): Sequential(
        (0): Dropout(p=0.5, inplace=False)
        (1): Conv2d(512, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      )
    )
    (lins): ModuleList(
      (0): NetLinLayer(
        (model): Sequential(
          (0): Dropout(p=0.5, inplace=False)
          (1): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
      )
      (1): NetLinLayer(
        (model): Sequential(
          (0): Dropout(p=0.5, inplace=False)
          (1): Conv2d(128, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
      )
      (2): NetLinLayer(
        (model): Sequential(
          (0): Dropout(p=0.5, inplace=False)
          (1): Conv2d(256, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
      )
      (3-4): 2 x NetLinLayer(
        (model): Sequential(
          (0): Dropout(p=0.5, inplace=False)
          (1): Conv2d(512, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
      )
    )
  )
)

from torch import optim

# Optimization
# AdamW with default weight decay
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
#Long burn
num_epochs = 500


# Cosine w/ Warm Restarts
# scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
#     optimizer,
#     T_0=50,          # First cycle is 50 epochs
#     T_mult=2,        # Double the length of each subsequent cycle
#     eta_min=1e-6     # Let it decay very close to zero
# )

Now, let's train the thing!

from tqdm import tqdm
import numpy as np
import torch
from plotly.subplots import make_subplots
import ipywidgets as widgets
from IPython.display import display
import base64, io
from PIL import Image as PILImage

import plotly.graph_objects as go

# =============================================================================
# DASHBOARD SETUP  — Plotly + ipywidgets (survives nbconvert → HTML)
# =============================================================================

history = {'epoch': [], 'total': [], 'nll': [], 'kld': [], 'perceptual': [], 'p_weight': []}
snapshot_store = {}   # epoch → list of 30 PIL images

# ── Loss figure (Plotly FigureWidget — updates in-place) ──────────
loss_fig = make_subplots(
    rows=1, cols=5,
    subplot_titles=('Total ELBO Loss', 'NLL  (Reconstruction)', 'KL Divergence',
                    'Perceptual Loss', 'Perceptual Weight'),
)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#e94560', width=2),
                               name='Total'), row=1, col=1)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#a0c4ff', width=2),
                               name='NLL'),   row=1, col=2)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#ffd700', width=2),
                               name='KLD'),   row=1, col=3)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#90ee90', width=2),
                               name='Perceptual'), row=1, col=4)
loss_fig.add_trace(go.Scatter(x=[], y=[], mode='lines',
                               line=dict(color='#ff9f40', width=2),
                               name='P-Weight'), row=1, col=5)
loss_fig.update_layout(
    height=320, paper_bgcolor='#1a1a2e', plot_bgcolor='#0d0d1a',
    font=dict(color='white'), showlegend=False,
    margin=dict(t=40, b=30, l=40, r=20),
)
for i in range(1, 6):
    loss_fig.update_xaxes(gridcolor='#222244', title_text='Epoch',
                           row=1, col=i)
    loss_fig.update_yaxes(gridcolor='#222244', row=1, col=i)

loss_widget = go.FigureWidget(loss_fig)


def _update_loss_plots():
    ep  = history['epoch']
    with loss_widget.batch_update():
        loss_widget.data[0].x = ep;  loss_widget.data[0].y = history['total']
        loss_widget.data[1].x = ep;  loss_widget.data[1].y = history['nll']
        loss_widget.data[2].x = ep;  loss_widget.data[2].y = history['kld']
        loss_widget.data[3].x = ep;  loss_widget.data[3].y = history['perceptual']
        loss_widget.data[4].x = ep;  loss_widget.data[4].y = history['p_weight']


# ── Snapshot viewer (ipywidgets — buttons + image grid) ───────────
def _tensor_to_b64(img_np):
    """Convert H×W×3 float32 numpy array → base64 PNG string."""
    arr = (np.clip(img_np, 0, 1) * 255).astype(np.uint8)
    buf = io.BytesIO()
    PILImage.fromarray(arr).save(buf, format='PNG')
    return 'data:image/png;base64,' + base64.b64encode(buf.getvalue()).decode()

ROW_LABELS = ['Train  Orig / Recon', 'Val  Orig / Recon', 'Generated']

# 30 image widgets arranged in 3 rows × 10 cols
img_widgets = [[widgets.Image(format='png', width=80, height=80)
                for _ in range(10)] for _ in range(3)]

row_boxes = []
for r in range(3):
    label = widgets.Label(
        value=ROW_LABELS[r],
        layout=widgets.Layout(width='130px', display='flex',
                              align_items='center'))
    row_box = widgets.HBox(
        [label] + img_widgets[r],
        layout=widgets.Layout(align_items='center'))
    row_boxes.append(row_box)

snap_label = widgets.Label(
    value='Snapshot — no data yet',
    layout=widgets.Layout(margin='4px 0 4px 0'))
snap_label.style.font_size = '13px'

snap_idx = [0]

btn_prev = widgets.Button(description='◀  Prev',
                          button_style='',
                          layout=widgets.Layout(width='100px'))
btn_next = widgets.Button(description='Next  ▶',
                          button_style='',
                          layout=widgets.Layout(width='100px'))
for b in (btn_prev, btn_next):
    b.style.button_color = '#16213e'
    b.style.text_color   = 'white'

epoch_dropdown = widgets.Dropdown(
    options=[],
    description='Epoch:',
    layout=widgets.Layout(width='150px'))


def _render_snapshot(epoch):
    imgs = snapshot_store.get(epoch)
    if imgs is None:
        return
    for r in range(3):
        for c in range(10):
            b64 = _tensor_to_b64(imgs[r * 10 + c])
            raw = base64.b64decode(b64.split(',')[1])
            img_widgets[r][c].value = raw
    snap_label.value = f'Snapshot — Epoch {epoch}'


def _on_prev(_b):
    keys = sorted(snapshot_store.keys())
    if not keys: return
    snap_idx[0] = max(0, snap_idx[0] - 1)
    epoch_dropdown.value = keys[snap_idx[0]]


def _on_next(_b):
    keys = sorted(snapshot_store.keys())
    if not keys: return
    snap_idx[0] = min(len(keys) - 1, snap_idx[0] + 1)
    epoch_dropdown.value = keys[snap_idx[0]]


def _on_dropdown_change(change):
    if change['name'] == 'value' and change['new'] is not None:
        keys = sorted(snapshot_store.keys())
        snap_idx[0] = keys.index(change['new'])
        _render_snapshot(change['new'])


btn_prev.on_click(_on_prev)
btn_next.on_click(_on_next)
epoch_dropdown.observe(_on_dropdown_change)

nav_bar      = widgets.HBox([btn_prev, epoch_dropdown, btn_next])
snap_viewer  = widgets.VBox([snap_label, nav_bar] + row_boxes)
dashboard    = widgets.VBox([loss_widget, snap_viewer])
display(dashboard)


def _capture_snapshot(epoch):
    model.eval()
    imgs = []
    with torch.no_grad():
        tb, _ = next(iter(train_loader))
        tb = tb[:5].to(device);  tr, _, _ = model(tb)
        for i in range(5):
            imgs.append(tb[i].cpu().permute(1,2,0).numpy())
            imgs.append(tr[i].cpu().permute(1,2,0).numpy())

        vb, _ = next(iter(val_loader))
        vb = vb[:5].to(device);  vr, _, _ = model(vb)
        for i in range(5):
            imgs.append(vb[i].cpu().permute(1,2,0).numpy())
            imgs.append(vr[i].cpu().permute(1,2,0).numpy())

        z_rand = torch.randn(10, hidden_size).to(device) * 0.7
        gen = model.decoder(model.fc_decode(z_rand).view(-1, 512, 4, 4))
        for i in range(10):
            imgs.append(gen[i].cpu().permute(1,2,0).numpy())

    snapshot_store[epoch] = imgs
    epoch_dropdown.options = sorted(snapshot_store.keys())
    snap_idx[0] = len(snapshot_store) - 1
    epoch_dropdown.value = epoch
    model.train()

# =============================================================================
# TRAINING LOOP  — VAE + Perceptual Loss
# =============================================================================

best_avg_loss = float('inf')
global_step   = 0

for epoch in range(1, num_epochs + 1):
    model.train()
    train_loss    = 0.0
    total_nll     = 0.0
    total_kld     = 0.0
    total_p_loss  = 0.0
    total_p_weight = 0.0

    #beta = 1.0
    beta = get_beta(epoch, warmup_start=5, warmup_end=15)
    pbar = tqdm(train_loader,
                desc=f"Epoch {epoch:3d}/{num_epochs}", leave=False)

    for batch_idx, (data, _) in enumerate(pbar):
        data = data.to(device)

        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)

        last_layer_weight = model.get_last_decoder_layer()

        loss, nll, kld, p_loss, p_weight = perceptual_generator_loss(
            reconstruction     = recon_batch,
            x                  = data,
            mu                 = mu,
            logvar             = logvar,
            perceptual_loss_fn = perceptual_loss_fn,
            last_layer_weight  = last_layer_weight,
            beta               = beta,
            free_bits          = 0.1,
        )

        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)
        optimizer.step()

        train_loss     += loss.item()
        total_nll      += nll
        total_kld      += kld
        total_p_loss   += p_loss
        total_p_weight += p_weight
        global_step    += 1

        pbar.set_postfix({
            'Loss':   f'{train_loss    / (batch_idx+1):.1f}',
            'NLL':    f'{total_nll     / (batch_idx+1):.1f}',
            'KLD':    f'{total_kld     / (batch_idx+1):.1f}',
            'PL':     f'{total_p_loss  / (batch_idx+1):.4f}',
            'PW':     f'{total_p_weight/ (batch_idx+1):.2f}',
            'β':      f'{beta:.3f}',
        })

    # ── Epoch averages ────────────────────────────────────────────
    n              = len(train_loader)
    avg_loss       = train_loss    / n
    avg_nll        = total_nll     / n
    avg_kld        = total_kld     / n
    avg_p_loss     = total_p_loss  / n
    avg_p_weight   = total_p_weight / n
    current_lr     = optimizer.param_groups[0]['lr']

    # ── Checkpoint on best training ELBO ─────────────────────────
    if avg_loss < best_avg_loss:
        best_avg_loss = avg_loss
        torch.save({
            "epoch":            epoch,
            "global_step":      global_step,
            "model_state_dict": model.state_dict(),
            "opt_state_dict":   optimizer.state_dict(),
            "best_avg_loss":    best_avg_loss,
            "hidden_size":      model.hidden_size,
        }, "vae_perceptual_best.pth")

    #scheduler.step()

    # ── Live dashboard update ─────────────────────────────────────
    if epoch >= 20:
        history['epoch']     .append(epoch)
        history['total']     .append(avg_loss)
        history['nll']       .append(avg_nll)
        history['kld']       .append(avg_kld)
        history['perceptual'].append(avg_p_loss)
        history['p_weight']  .append(avg_p_weight)
        _update_loss_plots()

    # ── Snapshot every 10 epochs ──────────────────────────────────
    if epoch % 10 == 0:
        _capture_snapshot(epoch)

VBox(children=(FigureWidget({
    'data': [{'line': {'color': '#e94560', 'width': 2},
              'mode': 'l…

An Improvement: The Two Stage Sampler

Above we can see that our reconstructions are decent now! Not blurry. Just... off a little. Facial features are weird, edges are a little weird. There's still a little blur.

That all said, the reconstructions are recognizable, but the generated images look like dog and cat parts thrown together!

Our theory is that each $\mathbf X$ is a draw of the random vector $\mathbf z$ from $P(Z)$. $P(Z)$ dictates the arrangment of the 256 latent values that are coherent - the VAE is learning a mapping/projection of the data space into the cat/dog space and vice-versa.

The problem is that we don't know $P(Z)$! We've learned $P(z | \mathbf X)$, but not the marginal distribution. With a large enough training set, we expect that $P(\mathbf z | \mathbf X)$ represents the true latent distribution - correlations and all - since we've seen a lot of $\mathbf X$ values and mapped them to optimal $\mathbf z$ values.

But, do we know the structure of $P(\mathbf z | \mathbf X)$? Can we sample from it?

By construction:

$$ P(\mathbf z | \mathbf X) = P(\mathbf X | f_\theta(\mathbf z)) \mathcal N_K(\mathbf z | \mathbf 0 , \mathbf I) $$

So, we expect that the posterior is close-ish to a normal distribution centered on 0 and with unit variance.

What if we take draws from this distribution and pass them through the decoder?

import matplotlib.pyplot as plt

model.eval()
with torch.no_grad():
    # Draw 100 random samples from the normal prior N(0, I)
    z_random = torch.randn(100, hidden_size).to(device)
    
    # Project and decode
    z_projected = model.fc_decode(z_random).view(-1, 512, 4, 4)
    generated_images = model.decoder(z_projected)

# Plot a 10x10 grid of the generated images
plt.figure(figsize=(15, 15))
for i in range(100):
    # Convert from (C, H, W) to (H, W, C) for plotting
    img_np = generated_images[i].cpu().permute(1, 2, 0).numpy()
    
    plt.subplot(10, 10, i + 1)
    plt.imshow(img_np)
    plt.axis('off')

plt.tight_layout()
plt.show()