ResNet

In twenty fifteen, researchers noticed something strange. A fifty six layer network did worse than a twenty layer network, even on its training data. That made no sense. The deeper network could simply copy the shallow one, and pass the signal through the extra layers unchanged. So why couldn't it learn that? It turned out that learning to do nothing is surprisingly hard.

• Deeper networks did worse
• Even on training data
• Doing nothing is hard to learn

The skip connection

The answer was almost comically simple. Add the input directly to the output of the layers. Now, these layers don't need to learn the whole answer. They only need to learn the change, the small correction on top. If there's nothing to add, they just output zero, which is easy. And the gradient gets a clean highway straight back to the early layers. A plus sign. That's all it was. It won every major vision contest that year.

A residual block

Here's the block in code. Two convolutions with batch norm, just as usual. The whole trick is on the last line. X plus y. Run it. The shape going in is the same as the shape coming out, which must be true for the addition to work. So you can stack these as many times as you like.

import torch
from torch import nn

class Block(nn.Module):
    def __init__(self, ch):
        super().__init__()
        self.conv1 = nn.Conv2d(ch, ch, 3, padding=1)
        self.conv2 = nn.Conv2d(ch, ch, 3, padding=1)
        self.bn1, self.bn2 = nn.BatchNorm2d(ch), nn.BatchNorm2d(ch)

    def forward(self, x):
        y = torch.relu(self.bn1(self.conv1(x)))
        y = self.bn2(self.conv2(y))
        return torch.relu(x + y)

block = Block(16)
print("Shape in: ", [1, 16, 32, 32])
print("Shape out:", list(block(torch.rand(1, 16, 32, 32)).shape))

Shape in:  [1, 16, 32, 32]
Shape out: [1, 16, 32, 32]

Proving it works

Let's prove it. We build two very deep networks, each with thirty blocks. They're the same in every way, except that one has skip connections. We train each one on the same small problem. Run it. The plain network is stuck, and barely learns at all, because its signals die in the depth. With skips, the very same network trains. One plus sign is the whole difference. You'll find this idea again inside every transformer.

import torch
from torch import nn

def build(skip, depth=30):
    class Layer(nn.Module):
        def __init__(self):
            super().__init__()
            self.f = nn.Sequential(nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 32))
        def forward(self, x):
            return torch.relu(x + self.f(x)) if skip else torch.relu(self.f(x))
    return nn.Sequential(nn.Linear(4, 32), *[Layer() for _ in range(depth)],
                         nn.Linear(32, 1))

torch.manual_seed(0)
X = torch.rand(256, 4)
y = (X[:, 0] * X[:, 1] + X[:, 2:].sum(1, keepdim=True).ravel()).unsqueeze(1)

for name, skip in [("Plain", False), ("With skips", True)]:
    torch.manual_seed(1)
    net = build(skip)
    opt = torch.optim.Adam(net.parameters(), lr=3e-4)
    for _ in range(300):
        loss = ((net(X) - y) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    print(f"{name:<11} 30-block network, final loss {loss.item():.4f}")

Plain       30-block network, final loss 0.2110
With skips  30-block network, final loss 0.0001

Summary