Learn AI Series):Exercise 1: Complete ViT for CIFAR-10 vs CNN baseline.
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import time
class PatchEmbedding(nn.Module):
def __init__(self, img_size=32, patch_size=4, in_channels=3, d_model=128):
super().__init__()
self.n_patches = (img_size // patch_size) ** 2
self.proj = nn.Conv2d(in_channels, d_model,
kernel_size=patch_size, stride=patch_size)
def forward(self, x):
x = self.proj(x)
return x.flatten(2).transpose(1, 2)
class TransformerBlock(nn.Module):
def __init__(self, d_model, n_heads, d_ff):
super().__init__()
self.norm1 = nn.LayerNorm(d_model)
self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
self.norm2 = nn.LayerNorm(d_model)
self.ff = nn.Sequential(
nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))
def forward(self, x):
x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
x = x + self.ff(self.norm2(x))
return x
class SmallViT(nn.Module):
def __init__(self, img_size=32, patch_size=4, d_model=128,
n_heads=4, n_layers=4, n_classes=10, d_ff=512):
super().__init__()
self.patch_embed = PatchEmbedding(img_size, patch_size, 3, d_model)
n_patches = self.patch_embed.n_patches
self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
self.pos_embed = nn.Parameter(
torch.randn(1, n_patches + 1, d_model) * 0.02)
self.layers = nn.ModuleList([
TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)])
self.norm = nn.LayerNorm(d_model)
self.head = nn.Linear(d_model, n_classes)
def forward(self, x):
batch = x.size(0)
patches = self.patch_embed(x)
cls = self.cls_token.expand(batch, -1, -1)
x = torch.cat([cls, patches], dim=1) + self.pos_embed
for layer in self.layers:
x = layer(x)
return self.head(self.norm(x[:, 0]))
# CNN baseline
cnn = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(128, 256, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1),
nn.Flatten(), nn.Linear(256, 10)
)
transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(32, padding=4),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
test_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_data = datasets.CIFAR10('.', train=True, download=True, transform=transform)
test_data = datasets.CIFAR10('.', train=False, transform=test_transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
test_loader = DataLoader(test_data, batch_size=256)
def train_model(model, name, epochs=20):
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(epochs):
model.train()
total_loss = 0
for imgs, labels in train_loader:
logits = model(imgs)
loss = nn.CrossEntropyLoss()(logits, labels)
opt.zero_grad()
loss.backward()
opt.step()
total_loss += loss.item()
model.eval()
correct, total = 0, 0
with torch.no_grad():
for imgs, labels in test_loader:
preds = model(imgs).argmax(-1)
correct += (preds == labels).sum().item()
total += labels.size(0)
acc = correct / total
if epoch % 5 == 0 or epoch == epochs - 1:
print(f"{name} epoch {epoch:>2d}: loss={total_loss/len(train_loader):.3f}, "
f"test_acc={acc:.1%}")
vit = SmallViT()
print(f"ViT params: {sum(p.numel() for p in vit.parameters()):,}")
print(f"CNN params: {sum(p.numel() for p in cnn.parameters()):,}\n")
train_model(cnn, "CNN", epochs=20)
print()
train_model(vit, "ViT", epochs=20)
The CNN trains faster in the first few epochs because of its built-in spatial biases -- convolutions are a strong prior for image data. With data augmentation (random flip + random crop), the ViT closes the gap by epoch 15-20, but on CIFAR-10's modest 50K images it typically doesn't surpass the CNN. On larger datasets, the crossover point favors the ViT.
Exercise 2: Patch embedding visualization and position similarity.
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
transform = transforms.Compose([transforms.ToTensor()])
cifar = datasets.CIFAR10('.', train=True, download=True, transform=transform)
img, label = cifar[0] # single image: (3, 32, 32)
print(f"Image shape: {img.shape}, label: {label}")
pe = PatchEmbedding(img_size=32, patch_size=4, d_model=128)
patches = pe(img.unsqueeze(0))
print(f"Patch embeddings: {patches.shape}") # (1, 64, 128)
print(f"Number of patches: {patches.shape[1]}")
print(f"Each patch: 4x4x3 = {4*4*3} pixels -> {patches.shape[2]}-dim vector")
# Verify patch extraction matches raw pixels
n_per_side = 32 // 4 # 8
for row in range(2):
for col in range(2):
patch_idx = row * n_per_side + col
raw_patch = img[:, row*4:(row+1)*4, col*4:(col+1)*4]
print(f"Patch ({row},{col}) idx={patch_idx}: "
f"raw pixel range [{raw_patch.min():.3f}, {raw_patch.max():.3f}]")
# Position embedding similarity (randomly initialized ViT)
vit = SmallViT(img_size=32, patch_size=4, d_model=128,
n_heads=4, n_layers=4, n_classes=10, d_ff=512)
pos = vit.pos_embed.detach().squeeze(0) # (65, 128) -- 64 patches + CLS
sim = F.cosine_similarity(pos.unsqueeze(0), pos.unsqueeze(1), dim=-1)
print(f"\nPosition similarity matrix: {sim.shape}") # (65, 65)
# Top similar pairs (excluding self-similarity)
sim_no_diag = sim.clone()
sim_no_diag.fill_diagonal_(float('-inf'))
flat_idx = sim_no_diag.view(-1).topk(6).indices
for idx in flat_idx:
i, j = idx.item() // 65, idx.item() % 65
print(f" Positions {i} and {j}: similarity = {sim[i, j]:.4f}")
print("\nBefore training, similarities are roughly random.")
print("After training, nearby patches would show high similarity.")
Before training the position embeddings are randomly initialized, so the similarity matrix shows no spatial pattern. After training on real image data, the model discovers that neighboring patches should have similar position embeddings -- it learns the 2D grid structure purely from the optimization objective. This is what we discussed in the episode about 1D position embeddings learning 2D spatial structure.
Exercise 3: Window attention vs full attention -- cost and speed comparison.
import torch
import torch.nn as nn
import math
import time
class FullAttention(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
self.n_heads = n_heads
self.d_k = d_model // n_heads
self.qkv = nn.Linear(d_model, d_model * 3)
self.proj = nn.Linear(d_model, d_model)
def forward(self, x):
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.n_heads, self.d_k)
qkv = qkv.permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)
attn = attn.softmax(dim=-1)
out = (attn @ v).transpose(1, 2).reshape(B, N, C)
return self.proj(out)
class WindowAttention(nn.Module):
def __init__(self, d_model, n_heads, window_size=7):
super().__init__()
self.window_size = window_size
self.n_heads = n_heads
self.d_k = d_model // n_heads
self.qkv = nn.Linear(d_model, d_model * 3)
self.proj = nn.Linear(d_model, d_model)
def forward(self, x, H, W):
B, N, C = x.shape
ws = self.window_size
x = x.view(B, H, W, C)
pad_h = (ws - H % ws) % ws
pad_w = (ws - W % ws) % ws
if pad_h > 0 or pad_w > 0:
x = nn.functional.pad(x, (0, 0, 0, pad_w, 0, pad_h))
Hp, Wp = x.shape[1], x.shape[2]
nH, nW = Hp // ws, Wp // ws
windows = x.view(B, nH, ws, nW, ws, C)
windows = windows.permute(0, 1, 3, 2, 4, 5).contiguous()
windows = windows.view(-1, ws * ws, C)
qkv = self.qkv(windows).reshape(-1, ws*ws, 3, self.n_heads, self.d_k)
qkv = qkv.permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)
attn = attn.softmax(dim=-1)
out = (attn @ v).transpose(1, 2).reshape(-1, ws*ws, C)
out = self.proj(out)
out = out.view(B, nH, nW, ws, ws, C)
out = out.permute(0, 1, 3, 2, 4, 5).contiguous()
out = out.view(B, Hp, Wp, C)
if pad_h > 0 or pad_w > 0:
out = out[:, :H, :W, :]
return out.view(B, H * W, C)
d_model = 96
full_attn = FullAttention(d_model, n_heads=3)
win_attn = WindowAttention(d_model, n_heads=3, window_size=7)
full_attn.eval()
win_attn.eval()
H, W = 56, 56
x = torch.randn(1, H * W, d_model)
# Theoretical cost
N = H * W # 3136
w = 7 * 7 # 49
n_windows = (H // 7) * (W // 7) # 64
full_flops = N * N * d_model
win_flops = n_windows * w * w * d_model
print(f"Full attention: O({N}^2 * {d_model}) = {full_flops:,} ops")
print(f"Window attention: O({n_windows} * {w}^2 * {d_model}) = {win_flops:,} ops")
print(f"Theoretical speedup: {full_flops / win_flops:.1f}x\n")
# Wall-clock timing
with torch.no_grad():
t0 = time.perf_counter()
for _ in range(10):
_ = full_attn(x)
t_full = (time.perf_counter() - t0) / 10
t0 = time.perf_counter()
for _ in range(10):
_ = win_attn(x, H, W)
t_win = (time.perf_counter() - t0) / 10
print(f"Full attention: {t_full*1000:.1f}ms")
print(f"Window attention: {t_win*1000:.1f}ms")
print(f"Measured speedup: {t_full / t_win:.1f}x")
The theoretical speedup is dramatic -- full attention on a 56x56 feature map is O(3136^2) while 7x7 window attention is O(64 * 49^2), roughly 64x cheaper. The measured speedup will be lower than theoretical (memory overhead, window reshaping operations, CUDA kernel launch costs on GPU), but still substantial. This is why the Swin Transformer can handle high-resolution inputs that would be impractical with standard ViT's global attention.
Here we go! For the past several episodes, everything we've built has been discriminative. We've trained models that take an input and predict a label or a value. A CNN takes an image and says "cat" or "dog." A transformer takes a sequence and predicts the next token. Even our ViT from last episode was a classifier -- it chops an image into patches, processes them with self-attention, and outputs a class prediction. All of these are asking the question: given this input, what category does it belong to?
Today we flip the script entirely. Instead of classifying existing data, we're going to create new data. Generative models learn the underlying distribution of a training set and then sample from it -- producing new images, new music, new text that look like they could have been in the training set but never were. And the most fascinating way to train a generative model? Make two neural networks fight each other.
This is the Generative Adversarial Network -- the GAN -- and it was introduced by Ian Goodfellow in 2014. The idea is almost philosophically elegant: one network tries to create convincing fakes, and another network tries to catch them. As the detector gets better, the forger gets better. When the game reaches equilibrium, the forger produces samples indistinguishable from real data.
GANs were the dominant generative approach for images from roughly 2014 to 2022, and understanding them is essential background for the diffusion models and other generative techniques we'll cover later in the series. The adversarial training principle also shows up in robustness training, domain adaptation, and quite some other areas of modern ML ;-)
The analogy everyone uses (and honestly, it's a good one): a counterfeiter and a detective. The counterfeiter creates fake banknotes. The detective examines banknotes and decides whether each is real or fake. If the detective catches a fake, the counterfeiter learns what gave it away and improves. If the detective is fooled, the counterfeiter knows that technique works.
Over time, both get better. The counterfeiter produces increasingly realistic fakes. The detective develops increasingly subtle detection methods. If training goes well, the counterfeiter eventually produces fakes so good that the detective can't do better than random guessing -- 50/50.
In neural network terms:
import torch
import torch.nn as nn
class Generator(nn.Module):
def __init__(self, z_dim=64, out_dim=784):
super().__init__()
self.net = nn.Sequential(
nn.Linear(z_dim, 128), nn.ReLU(),
nn.Linear(128, 256), nn.ReLU(),
nn.Linear(256, out_dim), nn.Tanh()
)
def forward(self, z):
return self.net(z)
class Discriminator(nn.Module):
def __init__(self, in_dim=784):
super().__init__()
self.net = nn.Sequential(
nn.Linear(in_dim, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 1), nn.Sigmoid()
)
def forward(self, x):
return self.net(x)
G = Generator(z_dim=64, out_dim=784)
D = Discriminator(in_dim=784)
z = torch.randn(4, 64)
fake = G(z)
score = D(fake)
print(f"Noise: {z.shape}")
print(f"Generated: {fake.shape}") # (4, 784) -- flattened 28x28 images
print(f"D scores: {score.shape}") # (4, 1) -- probabilities
print(f"D thinks real: {score.detach().mean():.4f}")
Two details worth noting. The generator uses Tanh to output values in [-1, 1] (matching normalized image data). The discriminator uses LeakyReLU in stead of ReLU -- this prevents "dead neurons" that can cause the discriminator to stop providing useful gradients to the generator. We covered the dead neuron problem back in episode #40 when discussing training challenges; LeakyReLU allows a small negative slope (0.2) so gradients always flow, even for negative inputs.
The training objective is a minimax game. The discriminator wants to maximize its ability to classify real vs fake. The generator wants to minimize the discriminator's ability -- it wants the discriminator to think generated samples are real.
Mathematically: min_G max_D E[log D(x)] + E[log(1 - D(G(z)))]
The first term: for real data x, the discriminator wants D(x) close to 1 (high log probability). The second term: for fake data G(z), the discriminator wants D(G(z)) close to 0, so log(1 - 0) = 0 (maximum). But the generator wants D(G(z)) close to 1, so log(1 - 1) approaches negative infinity (minimum).
In practice, training alternates between two steps:
import torch.optim as optim
z_dim = 64
G = Generator(z_dim, 784)
D = Discriminator(784)
opt_G = optim.Adam(G.parameters(), lr=2e-4, betas=(0.5, 0.999))
opt_D = optim.Adam(D.parameters(), lr=2e-4, betas=(0.5, 0.999))
criterion = nn.BCELoss()
# Simulated training step (real data would come from DataLoader)
batch_size = 64
real_data = torch.randn(batch_size, 784) * 0.5 # placeholder for real images
# Step 1: Train Discriminator
z = torch.randn(batch_size, z_dim)
fake_data = G(z).detach() # detach: don't backprop through G
real_labels = torch.ones(batch_size, 1)
fake_labels = torch.zeros(batch_size, 1)
d_loss_real = criterion(D(real_data), real_labels)
d_loss_fake = criterion(D(fake_data), fake_labels)
d_loss = d_loss_real + d_loss_fake
opt_D.zero_grad()
d_loss.backward()
opt_D.step()
# Step 2: Train Generator
z = torch.randn(batch_size, z_dim)
fake_data = G(z)
g_loss = criterion(D(fake_data), real_labels) # G wants D to say "real"
opt_G.zero_grad()
g_loss.backward()
opt_G.step()
print(f"D loss: {d_loss.item():.4f}, G loss: {g_loss.item():.4f}")
The .detach() on the generator's output in step 1 is critical. When training the discriminator, we don't want gradients flowing back into the generator -- we're only updating D's parameters. In step 2, we do want gradients flowing through D and into G (but we don't update D's parameters -- only G's).
Notice the generator's loss uses real_labels as the target. The generator is trying to fool the discriminator -- it wants D to output 1 (real) for generated samples. This is sometimes called the "non-saturating" loss because log(D(G(z))) provides stronger gradients when D(G(z)) is small (early in training) compared to the original minimax log(1 - D(G(z))) formulation. Stronger gradients early on means the generator actually learns something before the discriminator becomes too dominant.
Having said that, let me show you a more complete training loop with actual MNIST data to make this concrete:
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# MNIST setup
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)) # scale to [-1, 1]
])
mnist = datasets.MNIST('.', train=True, download=True, transform=transform)
loader = DataLoader(mnist, batch_size=64, shuffle=True)
z_dim = 64
G = Generator(z_dim, 784)
D = Discriminator(784)
opt_G = optim.Adam(G.parameters(), lr=2e-4, betas=(0.5, 0.999))
opt_D = optim.Adam(D.parameters(), lr=2e-4, betas=(0.5, 0.999))
criterion = nn.BCELoss()
for epoch in range(5):
d_total, g_total, n = 0, 0, 0
for real_imgs, _ in loader:
batch = real_imgs.size(0)
real_flat = real_imgs.view(batch, -1)
real_lbl = torch.ones(batch, 1) * 0.9 # label smoothing
fake_lbl = torch.zeros(batch, 1)
# Train D
z = torch.randn(batch, z_dim)
fake = G(z).detach()
d_loss = criterion(D(real_flat), real_lbl) + criterion(D(fake), fake_lbl)
opt_D.zero_grad()
d_loss.backward()
opt_D.step()
# Train G
z = torch.randn(batch, z_dim)
fake = G(z)
g_loss = criterion(D(fake), torch.ones(batch, 1))
opt_G.zero_grad()
g_loss.backward()
opt_G.step()
d_total += d_loss.item()
g_total += g_loss.item()
n += 1
print(f"Epoch {epoch}: D_loss={d_total/n:.3f}, G_loss={g_total/n:.3f}")
After a few epochs, the generator starts producing things that vaguely resemble digits. After many more epochs with a proper DCGAN architecture (coming up next), the results get dramatically better. But even with this simple MLP-based GAN on MNIST, the principle works -- the generator learns to produce 28x28 images from pure noise, guided entirely by the discriminator's feedback.
GAN training is notoriously unstable. The generator and discriminator are in a dynamic equilibrium, and quite some things can go wrong:
Oscillation: D gets too strong, G can't make progress. D's loss drops to near zero (it classifies everything correctly), which means G gets no useful gradient signal. Then G overfits to a specific trick, D adjusts, and the cycle repeats without convergence.
Mode collapse: the generator discovers that producing one specific output (or a few specific outputs) consistently fools the discriminator. In stead of generating diverse faces, it generates the same face over and over. The generator has "collapsed" to a few modes of the data distribution in stead of covering the full distribution. Think about why this happens -- if the generator finds one image that D consistently rates as "real," the gradient signal says "keep doing this." The generator has no incentive to explore other outputs because exploration risks worse scores.
Vanishing gradients: if D is too confident (outputs very close to 0 or 1), the gradients through the sigmoid are tiny, and G learns very slowly. This is the same saturation problem we saw in episode #40 with sigmoid activations.
Non-convergence: unlike supervised learning where the loss monotonically decreases (at least on the training set), GAN training is a game -- the losses oscillate as G and D take turns improving. There's no simple metric that tells you "training is going well" until you visually inspect the generated samples. This is genuinly one of the most frustrating aspects of working with GANs in practice.
Practical tricks that help:
# Tricks that make GAN training more stable
# 1. Label smoothing: 0.9 instead of 1.0 for "real"
real_labels = torch.ones(batch_size, 1) * 0.9
# 2. Adam with beta1=0.5 (standard GAN recipe)
opt = optim.Adam(params, lr=2e-4, betas=(0.5, 0.999))
# 3. LeakyReLU in discriminator
nn.LeakyReLU(0.2) # NOT nn.ReLU()
# 4. Batch normalization (but NOT in D's output layer)
nn.BatchNorm1d(256)
# 5. Train D more than G when needed
for _ in range(2): # 2 D steps per G step
# ... train discriminator ...
# ... train generator once ...
# 6. Spectral normalization on D's weights
nn.utils.spectral_norm(nn.Linear(256, 128))
These tricks were discovered through years of painful experimentation by the community. They're not derived from theory -- they're empirical recipes that make the training game more stable. The GAN literature is full of papers proposing new loss functions, normalization schemes, and architectural constraints specifically to tame training instability.
The original GAN used fully connected layers, which doesn't scale well to images. In 2016, Radford et al. proposed DCGAN (Deep Convolutional GAN) with specific architectural guidelines that made GAN training much more stable on image data:
class DCGANGenerator(nn.Module):
def __init__(self, z_dim=100, channels=1):
super().__init__()
self.net = nn.Sequential(
# z_dim -> 256 x 7 x 7
nn.ConvTranspose2d(z_dim, 256, 7, 1, 0, bias=False),
nn.BatchNorm2d(256), nn.ReLU(),
# 256x7x7 -> 128x14x14
nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
nn.BatchNorm2d(128), nn.ReLU(),
# 128x14x14 -> channels x 28x28
nn.ConvTranspose2d(128, channels, 4, 2, 1, bias=False),
nn.Tanh()
)
def forward(self, z):
return self.net(z.view(-1, z.size(1), 1, 1))
class DCGANDiscriminator(nn.Module):
def __init__(self, channels=1):
super().__init__()
self.net = nn.Sequential(
# channels x 28x28 -> 64x14x14
nn.Conv2d(channels, 64, 4, 2, 1, bias=False),
nn.LeakyReLU(0.2),
# 64x14x14 -> 128x7x7
nn.Conv2d(64, 128, 4, 2, 1, bias=False),
nn.BatchNorm2d(128), nn.LeakyReLU(0.2),
# 128x7x7 -> 1
nn.Conv2d(128, 1, 7, 1, 0, bias=False),
nn.Sigmoid()
)
def forward(self, x):
return self.net(x).view(-1, 1)
G = DCGANGenerator(z_dim=100)
D = DCGANDiscriminator()
z = torch.randn(4, 100)
fake_imgs = G(z)
scores = D(fake_imgs)
print(f"Generated images: {fake_imgs.shape}") # (4, 1, 28, 28)
print(f"Discriminator scores: {scores.shape}") # (4, 1)
print(f"G params: {sum(p.numel() for p in G.parameters()):,}")
print(f"D params: {sum(p.numel() for p in D.parameters()):,}")
The DCGAN guidelines that became standard practice:
ConvTranspose2d (transposed convolution, sometimes misleadingly called "deconvolution") is the generator's key operation. If you remember from episodes #45-46, a regular convolution with stride 2 downsamples -- it takes a 14x14 feature map and produces a 7x7 one. A transposed convolution does the opposite: it upsamples -- a 7x7 feature map becomes 14x14, then 28x28. The generator starts from a random vector and progressively "paints" an image at increasing resolution, layer by layer. The discriminator does the reverse -- it takes a full-resolution image and progressively compresses it down to a single real/fake probability.
Generating tiny 28x28 MNIST digits is one thing. Generating photorealistic 1024x1024 faces is something else entirely. The problem: training a GAN on high-resolution images from the start is extremely unstable. The discriminator easily distinguishes blurry early-stage generator outputs from crisp real images, providing no useful gradient signal.
Progressive growing (Karras et al., 2017, NVIDIA) solved this by starting training at low resolution (4x4) and gradually adding layers to both G and D during training:
Each resolution stage starts easy -- the generator just needs to get the rough structure right -- and refines progressivly. By the time the model reaches 1024x1024, it has already learned faces at all lower resolutions, so the high-resolution details are refinements on an already solid foundation.
StyleGAN (Karras et al., 2019, also NVIDIA) built on progressive growing with a radically different generator architecture. In stead of feeding the noise vector directly into the first layer, StyleGAN introduces a mapping network that transforms the noise z into an intermediate latent vector w through 8 fully-connected layers, and then injects w into each layer through adaptive instance normalization (AdaIN). Different layers control different aspects of the image:
class MappingNetwork(nn.Module):
"""StyleGAN's mapping network: z -> w through 8 FC layers."""
def __init__(self, z_dim=512, w_dim=512, n_layers=8):
super().__init__()
layers = []
for i in range(n_layers):
in_d = z_dim if i == 0 else w_dim
layers.extend([nn.Linear(in_d, w_dim), nn.LeakyReLU(0.2)])
self.net = nn.Sequential(*layers)
def forward(self, z):
return self.net(z)
class AdaIN(nn.Module):
"""Adaptive Instance Normalization -- injects style at each layer."""
def __init__(self, channels, w_dim=512):
super().__init__()
self.norm = nn.InstanceNorm2d(channels)
self.style = nn.Linear(w_dim, channels * 2) # scale + shift
def forward(self, x, w):
style = self.style(w).unsqueeze(2).unsqueeze(3)
gamma, beta = style.chunk(2, dim=1)
return gamma * self.norm(x) + beta
mapping = MappingNetwork(z_dim=128, w_dim=128)
adain = AdaIN(channels=64, w_dim=128)
z = torch.randn(2, 128)
w = mapping(z)
feature_map = torch.randn(2, 64, 16, 16)
styled = adain(feature_map, w)
print(f"z: {z.shape} -> w: {w.shape}")
print(f"Feature map: {feature_map.shape} -> Styled: {styled.shape}")
The key insight: different layers of the generator control different levels of detail:
This separation means you can mix styles from different images: take the coarse structure from one face and the fine details from another. Feed one w vector to the early layers and a different w vector to the late layers, and you get a controllable blend. StyleGAN2 (2020) and StyleGAN3 (2021) further improved quality and eliminated artifacts. By 2021, StyleGAN-generated faces were virtually indistinguishable from real photographs to the human eye.
One of the most important theoretical improvements to GAN training came from Arjovsky et al. in 2017: the Wasserstein GAN (WGAN). The idea: replace the binary cross-entropy loss with the Wasserstein distance (also called Earth Mover's Distance) between the real and generated distributions.
The original GAN minimizes the Jensen-Shannon divergence between distributions. The problem: when the real and generated distributions don't overlap (which is common early in training -- the generated images are obviously fake), the JS divergence is constant and provides zero gradient. The Wasserstein distance, on the other hand, always provides a meaningful gradient even when distributions don't overlap.
class WGANCritic(nn.Module):
"""WGAN uses a 'critic' (no sigmoid) instead of a discriminator."""
def __init__(self, in_dim=784):
super().__init__()
self.net = nn.Sequential(
nn.Linear(in_dim, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 1) # NO sigmoid -- raw score
)
def forward(self, x):
return self.net(x)
# WGAN training step (simplified)
critic = WGANCritic()
generator = Generator(z_dim=64)
opt_C = optim.RMSprop(critic.parameters(), lr=5e-5)
opt_G = optim.RMSprop(generator.parameters(), lr=5e-5)
batch_size = 64
real_data = torch.randn(batch_size, 784) * 0.5
# Train critic (5 steps per G step is typical for WGAN)
for _ in range(5):
z = torch.randn(batch_size, 64)
fake = generator(z).detach()
c_loss = -(critic(real_data).mean() - critic(fake).mean())
opt_C.zero_grad()
c_loss.backward()
opt_C.step()
# Weight clipping (original WGAN -- WGAN-GP uses gradient penalty instead)
for p in critic.parameters():
p.data.clamp_(-0.01, 0.01)
# Train generator
z = torch.randn(batch_size, 64)
fake = generator(z)
g_loss = -critic(fake).mean()
opt_G.zero_grad()
g_loss.backward()
opt_G.step()
print(f"Critic loss: {c_loss.item():.4f}, G loss: {g_loss.item():.4f}")
Notice three differences from the standard GAN: (1) no sigmoid on the critic (it outputs an unbounded score, not a probability), (2) the critic is trained multiple times per generator step (typically 5), and (3) weight clipping enforces the Lipschitz constraint that the Wasserstein distance requires. WGAN-GP (Gradient Penalty, Gulrajani et al.) later replaced weight clipping with a gradient penalty term, which works much better in practice.
WGANs are significantly more stable to train -- the critic loss actually correlates with sample quality (unlike standard GAN losses), so you can monitor training progress numerically in stead of just staring at generated images ;-)
Since roughly 2022, diffusion models (which we'll cover later in this series) have largely replaced GANs as the dominant generative approach for images. Stable Diffusion, DALL-E 2, Midjourney -- all diffusion-based.
Why the shift? Diffusion models are easier to train (no adversarial instability), produce higher diversity (less mode collapse), and handle text-conditioning more naturally. GANs can be faster at inference (one forward pass vs. many denoising steps) and are still used in specific applications -- real-time video synthesis, super-resolution, image editing, face manipulation -- but for general-purpose image generation, diffusion won.
Understanding GANs still matters though. The adversarial principle appears throughout modern AI: adversarial training for robustness (training a classifier on adversarial examples to make it harder to fool), discriminator heads in other architectures (used for domain adaptation, style transfer), and the game-theoretic thinking that GANs introduced. And many production systems still run GAN-based components -- particularly where inference speed matters, since a GAN generator produces an image in a single forward pass while a diffusion model needs 20-50 denoising steps.
Generative models that produce photorealistic human faces raise serious ethical questions, and I think it's important we address them directly.
Deepfakes: GAN-generated face-swaps can put anyone's face into any video. This enables fraud, harassment, non-consensual content, and political manipulation. The technology for creating convincing deepfakes has become accessible to anyone with a decent GPU and an afternoon to spare -- and that's a real problem.
Identity fraud: synthetic faces can be used to create fake social media profiles, bypass identity verification systems, and manufacture fake personas at scale. Every time you see "This Person Does Not Exist" generating a perfect face, remember that the same technology can generate thousands of fake LinkedIn profiles.
Consent: StyleGAN was trained on datasets of real people's faces (FFHQ -- Flickr-Faces-HQ, scraped from Flickr). Those people didn't consent to having their likeness used to train a model that generates synthetic faces. This raises fundamental questions about data rights that the field is still grappling with.
Detection: the arms race between generation and detection is ongoing. Techniques exist (frequency analysis, inconsistency detection, forensic classifiers), but they're always playing catch-up with improving generators. No detection method has remained reliable for more than a year or two before generators learned to avoid the artifacts being detected.
The technology itself is neutral -- the same GAN that enables deepfakes also enables medical image synthesis (generating rare pathology examples for training diagnostic models), privacy-preserving data sharing (generating synthetic patient records that preserve statistical properties without exposing real individuals), and creative tools. But responsible development requires acknowledging the dual-use potential and investing in detection and attribution alongside generation.
Exercise 1: Build a complete GAN training pipeline on MNIST and visualize the results. Create a Generator (z_dim=64, layers: 64->256->512->784, ReLU + Tanh output) and Discriminator (784->512->256->1, LeakyReLU 0.2 + Sigmoid). Train for 50 epochs on MNIST with Adam (lr=2e-4, betas=(0.5, 0.999)), batch size 64, and label smoothing (0.9 for real labels). Every 10 epochs, generate 16 samples from fixed noise vectors and save the 784-dim outputs reshaped to 28x28. Print the discriminator's average score on real data and fake data at each epoch -- track how these evolve. By epoch 50, the generated images should look recognizably like digits (if blurry).
Exercise 2: Implement mode collapse detection. Train the same GAN from Exercise 1, but this time after every epoch, generate 100 samples and compute the pairwise cosine similarity between all generated images (flattened to vectors). The average pairwise similarity tells you about diversity: if all images look the same (mode collapse), similarity is very high (close to 1.0). If images are diverse, similarity is lower. Plot (or print) this metric across 50 training epochs. Also compute the standard deviation of pixel values across the 100 generated images -- mode collapse should show low variance. Compare two training runs: one with the standard training recipe, and one where you intentionally cause mode collapse by training G for 5 steps per D step (the opposite of what WGAN recommends).
Exercise 3: Build a DCGAN for MNIST and compare against the MLP-based GAN. Implement the DCGANGenerator and DCGANDiscriminator from this episode. Train both the MLP GAN (Exercise 1) and the DCGAN for 20 epochs on MNIST with identical hyperparameters (Adam, lr=2e-4, betas=(0.5, 0.999), batch size 64). After training, generate 100 samples from each model. For each model, compute: (a) the average discriminator score on generated samples (higher = more realistic to D), (b) the pixel-level mean and standard deviation of generated images (should roughly match MNIST statistics: mean ~0.13, std ~0.31 for raw [0,1] values), and (c) the diversity metric from Exercise 2. The DCGAN should produce sharper, more structured images because the convolutional architecture has the right inductive biases for spatial data.
@scipio