🎓 A1Learner
🔍
0 of 53 lessons completed

PyTorch Residual Connections

A residual connection is a shortcut that adds a layer's input back to that layer's output. The addition gives the gradient a direct path from the end of a stack back to its start.

A residual layer adds its input back

One layer runs on a small input, and the input joins the result by addition. The gap between the two sides rounds to zero.

torch.manual_seed(0)
layer = nn.Sequential(nn.Linear(8, 8), nn.ReLU(), nn.Linear(8, 8))
xr = torch.randn(3, 8)
outr = xr + layer(xr)
gap = ((outr - xr) - layer(xr)).abs().max().item()
print("Matches within 0.000001:", gap < 1e-6)

Gradient norms at depth four

A gradient norm is the overall size of a gradient tensor. The norm comes from the square root of the summed squares of its entries.

def first_grad_norm(use_residual, depth):
    torch.manual_seed(0)
    layers = nn.ModuleList([nn.Sequential(nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 32)) for _ in range(depth)])
    torch.manual_seed(1)
    xx = torch.randn(16, 32)
    target = torch.randn(16, 32)
    h = xx
    for lyr in layers:
        h = h + lyr(h) if use_residual else lyr(h)
    loss = ((h - target) ** 2).mean()
    loss.backward()
    return layers[0][0].weight.grad.norm().item(), h

pf, ph = first_grad_norm(False, 4)
rf, rh = first_grad_norm(True, 4)
print(f"Plain first-layer gradient norm at depth 4: {pf:.4f}")
print(f"Residual first-layer gradient norm at depth 4: {rf:.4f}")

The plain stack also shrinks its signal

Output spread is the standard deviation of the stack's output values. The plain stack's spread shrinks while the residual stack keeps its spread.

print("Plain output spread at depth 4:", round(ph.std().item(), 4))
print("Residual output spread at depth 4:", round(rh.std().item(), 4))

Example

import torch
from torch import nn


def build_stack(depth, seed):
    torch.manual_seed(seed)
    return nn.ModuleList(
        [nn.Sequential(nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 32)) for _ in range(depth)]
    )


def first_grad_norm(use_residual, depth=8):
    layers = build_stack(depth, 0)
    torch.manual_seed(1)
    x = torch.randn(16, 32)
    target = torch.randn(16, 32)
    h = x
    for layer in layers:
        if use_residual:
            h = h + layer(h)
        else:
            h = layer(h)
    loss = ((h - target) ** 2).mean()
    loss.backward()
    first = layers[0][0].weight.grad.norm().item()
    last = layers[-1][0].weight.grad.norm().item()
    return first, last


plain_first, plain_last = first_grad_norm(False)
resid_first, resid_last = first_grad_norm(True)
print("Depth: 8 layers")
print(f"Plain stack, first layer gradient norm: {plain_first:.6f}")
print(f"Plain stack, last layer gradient norm: {plain_last:.4f}")
print(f"Residual stack, first layer gradient norm: {resid_first:.4f}")
print(f"Residual stack, last layer gradient norm: {resid_last:.4f}")

Key points

  • Use x + layer(x) to keep the input inside the output of every layer
  • Gradient norm: the overall size of a gradient tensor, from the square root of its summed squares
  • A plain stack at depth 8 gave a first-layer gradient norm of 0.000004 in the closing example
  • The residual stack at the same depth gave a first-layer gradient norm of 0.5434
  • The Intermediate lesson Checking Gradient Flow first watched gradient norms layer by layer

In this example, two stacks of eight layers train on the same seeded batch. Each layer is a widening-free pair of Linear layers with a ReLU between them. The plain stack passes each output straight into the next layer. The residual stack adds each layer's input back to its output. One backward pass runs on each stack. The plain stack's first layer receives a gradient norm of 0.000004, while its last layer receives 0.0255. The residual stack's first layer receives 0.5434, close to its last layer's 0.5474. The signal that teaches the early layers survives the trip in the residual stack.

Note

A residual connection adds no parameters, because the addition itself holds no weights

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does a residual connection add to a layer's output?