Open in Colab Download .ipynb

Lesson 35: Classic Architectures

Lessons 33-34 built and trained one small CNN. This lesson looks at how CNN architectures evolved over roughly two decades, and works through the single biggest architectural idea in that history in detail: the residual connection, and the vanishing-gradient problem it was designed to fix.

In [1]:
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

A brief tour

A handful of architectures defined the field, each solving a specific problem with the previous one:

  • LeNet-5 ("Gradient-Based Learning Applied to Document Recognition", LeCun et al., 1998★) — the template Lesson 33 already implements: conv, pool, conv, pool, then fully-connected layers. Built for digit recognition on 32x32 images.
  • AlexNet (Krizhevsky, Sutskever & Hinton, 2012★) — the same template, just much bigger (5 conv layers, ~60M parameters), trained on ImageNet (Deng et al., 2009★) using GPUs and ReLU instead of tanh/sigmoid. Its win in the 2012 ImageNet competition (top-5 error dropping from ~26% to ~15%) is usually cited as the event that restarted mainstream interest in neural networks.
  • VGG (Simonyan & Zisserman, 2014★) — replaced AlexNet's large 11x11/5x5 filters with a deep stack of small 3x3 convolutions. Two stacked 3x3 convs have the same receptive field as one 5x5 conv (Lesson 11's pyramid idea again) but fewer parameters and an extra nonlinearity in between.
  • GoogLeNet / Inception (Szegedy et al., 2015★) — went wider instead of just deeper: each "Inception module" runs several convolution sizes (1x1, 3x3, 5x5) side by side on the same input and concatenates their outputs, so the network doesn't have to commit to one filter size per layer. 1x1 convolutions are used first to cheaply shrink the channel count, keeping the wider module affordable.
  • ResNet (He et al., 2015★) — the subject of the rest of this lesson. Pushed depth from VGG's ~19 layers to 50, 101, even 152, by solving the problem that had made naively stacking more layers worse, not better.

The vanishing-gradient problem

Why did stacking more layers make plain networks worse? Backprop's chain rule multiplies a gradient by every layer's local Jacobian on its way back to the input. If those per-layer factors are consistently smaller than 1 — which happens easily with a saturating activation like sigmoid, whose derivative is at most 0.25 — a deep enough stack multiplies the gradient by a very small number many times over. The gradient reaching early layers shrinks toward zero, and those layers stop learning at all, even though nothing is mathematically wrong with the network.

Build a 30-layer, sigmoid-activated network of Linear layers and track the gradient magnitude at every depth, for two versions: a plain stack x = sigmoid(layer(x)), and a residual stack x = x + 0.3 * sigmoid(layer(x)) where each layer only has to learn a small correction added to its input, rather than replacing it outright.

In [2]:
DEPTH = 30
WIDTH = 32

def make_layers(seed):
    torch.manual_seed(seed)
    layers = nn.ModuleList([nn.Linear(WIDTH, WIDTH) for _ in range(DEPTH)])
    for layer in layers:
        nn.init.xavier_normal_(layer.weight, gain=nn.init.calculate_gain('sigmoid'))
    return layers

class PlainDeepNet(nn.Module):
    def __init__(self, seed):
        super().__init__()
        self.layers = make_layers(seed)

    def forward(self, x):
        acts = [x]
        for layer in self.layers:
            x = torch.sigmoid(layer(x))
            x.retain_grad()
            acts.append(x)
        return x, acts

class ResidualDeepNet(nn.Module):
    def __init__(self, seed):
        super().__init__()
        self.layers = make_layers(seed)

    def forward(self, x):
        acts = [x]
        for layer in self.layers:
            x = x + 0.3 * torch.sigmoid(layer(x))
            x.retain_grad()
            acts.append(x)
        return x, acts

def grad_norms(model_cls):
    torch.manual_seed(0)
    x = torch.randn(8, WIDTH, requires_grad=True)
    model = model_cls(seed=1)
    out, acts = model(x)
    out.sum().backward()
    return [a.grad.norm().item() for a in acts if a.grad is not None]

plain_norms = grad_norms(PlainDeepNet)
res_norms = grad_norms(ResidualDeepNet)

print(f'plain net:    gradient norm at layer 0 / gradient norm at layer {DEPTH-1} = {plain_norms[0] / plain_norms[-1]:.2e}')
print(f'residual net: gradient norm at layer 0 / gradient norm at layer {DEPTH-1} = {res_norms[0] / res_norms[-1]:.2e}')

plt.figure(figsize=(6, 4))
plt.semilogy(plain_norms, label='plain (sigmoid) net')
plt.semilogy(res_norms, label='residual net')
plt.xlabel('layer depth (0 = input)')
plt.ylabel('gradient norm (log scale)')
plt.title('Gradient magnitude vs. depth')
plt.legend()
plt.show()

plain net:    gradient norm at layer 0 / gradient norm at layer 29 = 2.07e-19
residual net: gradient norm at layer 0 / gradient norm at layer 29 = 1.05e+00
No description has been provided for this image

The plain network's gradient shrinks by roughly nineteen orders of magnitude between the last layer and the first — for all practical purposes, the early layers receive no learning signal at all. The residual network's gradient stays essentially flat across all 30 layers.

The reason is structural, not a matter of tuning: with a residual connection, x_{l+1} = x_l + F(x_l), so dx_{l+1}/dx_l = I + dF/dx_l. Backprop multiplies these Jacobians together across layers, but every one of them contains an identity matrix I as a direct term. That gives the gradient a path straight back to the input that never gets multiplied by a small sigmoid derivative — an unobstructed shortcut, regardless of how deep the stack gets. The plain network has no such path: every layer's Jacobian is purely dF/dx_l, so there's nothing to stop repeated multiplication from driving the product toward zero.

Batch normalization

The convolutional residual block below uses one more ingredient: batch normalization (Ioffe & Szegedy, 2015). The idea is simple — for each channel, subtract that channel's mean and divide by its standard deviation, computed across the current mini-batch (over the batch, height, and width dimensions, separately per channel), so every channel's activations always have mean 0 and standard deviation 1 going into the next layer. Two learnable parameters, a per-channel scale gamma and shift beta, are then applied on top, so the layer can still recover a different mean/scale if that's actually useful — normalization to 0/1 is a starting point the network can undo, not a hard constraint.

Deep networks without batch norm are prone to a related but distinct problem from vanishing gradients: as training updates early layers, the distribution of activations feeding into later layers keeps shifting (sometimes called "internal covariate shift"), so later layers are constantly chasing a moving target. Renormalizing at every layer keeps that distribution stable, which in practice lets much higher learning rates be used and makes deep networks noticeably easier to train — one of the reasons ResNet could push to 50+ layers where earlier architectures struggled past ~20.

In [3]:
C = 4
x_bn = torch.randn(8, C, 5, 5)

bn = nn.BatchNorm2d(C)
bn.train()
out_torch = bn(x_bn)

# from scratch: normalize each channel over (batch, height, width), then scale + shift
mean = x_bn.mean(dim=(0, 2, 3), keepdim=True)
var = x_bn.var(dim=(0, 2, 3), unbiased=False, keepdim=True)
x_norm = (x_bn - mean) / torch.sqrt(var + bn.eps)
gamma = bn.weight.view(1, C, 1, 1)
beta = bn.bias.view(1, C, 1, 1)
out_manual = gamma * x_norm + beta

print(f'max abs diff vs nn.BatchNorm2d: {(out_manual - out_torch).abs().max().item():.2e}')
print(f'per-channel mean before: {x_bn.mean(dim=(0, 2, 3)).detach().numpy().round(3)}')
print(f'per-channel std before:  {x_bn.std(dim=(0, 2, 3), unbiased=False).detach().numpy().round(3)}')
print(f'per-channel mean after:  {out_torch.mean(dim=(0, 2, 3)).detach().numpy().round(3)}')
print(f'per-channel std after:   {out_torch.std(dim=(0, 2, 3), unbiased=False).detach().numpy().round(3)}')

max abs diff vs nn.BatchNorm2d: 4.77e-07
per-channel mean before: [-0.046  0.165 -0.061 -0.165]
per-channel std before:  [0.939 1.037 0.911 0.977]
per-channel mean after:  [ 0.  0. -0.  0.]
per-channel std after:   [1. 1. 1. 1.]

A convolutional residual block

In an actual ResNet, F is a pair of small convolutions (with batch normalization and ReLU in between), and the shortcut adds the input feature map back onto their output, channel-for-channel and pixel-for-pixel:

In [4]:
class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(channels)

    def forward(self, x):
        out = torch.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        return torch.relu(x + out)  # the shortcut: add the block's input back in

block = ResidualBlock(channels=16)
x = torch.randn(4, 16, 20, 20)
y = block(x)
print(f'input shape:  {tuple(x.shape)}')
print(f'output shape: {tuple(y.shape)}  (unchanged — a residual block preserves shape, so blocks can be stacked freely)')

input shape:  (4, 16, 20, 20)
output shape: (4, 16, 20, 20)  (unchanged — a residual block preserves shape, so blocks can be stacked freely)

Exercise

  1. In PlainDeepNet/ResidualDeepNet above, change the activation from torch.sigmoid to torch.relu (ReLU's derivative is 1 for any positive input, not capped at 0.25 like sigmoid's) and rerun the gradient-norm comparison. Does the plain network's vanishing problem get better, worse, or stay about the same?
  2. Change the residual net's scale factor from 0.3 to 1.0 (i.e. x + torch.sigmoid(layer(x)) with no damping) and rerun. Does the gradient ratio change much? What does that suggest about why the residual connection works — is it the scale factor, or the + x shortcut itself?
  3. ResidualBlock above requires the input and output to have the same number of channels, since they're added directly. Real ResNets sometimes need to change channel count between blocks (e.g. 16 → 32). Sketch (in words, or in code) what the shortcut path would need to do in that case for the addition to still make sense.
← Lesson 34: Training a CNN
Lesson 36: Transfer Learning →