Why Deep Stacks Fail
An optimisation problem mistaken for overfitting, and its one-line fix
Week 5 ended with a stack that stopped improving as it got deeper. That is worth taking seriously, because the obvious explanation is wrong and the right one led to the single most useful architectural idea of the last decade.
The degradation problem
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms
MEAN, STD = (0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)
tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize(MEAN, STD)])
train_all = datasets.CIFAR10('data', train=True, download=True, transform=tf)
test_all = datasets.CIFAR10('data', train=False, download=True, transform=tf)
CLASSES = train_all.classes
train_set = Subset(train_all, range(4000))
val_set = Subset(test_all, range(2000))
def fit(model, epochs=4, lr=0.05, train=None, seed=0, report_train=False):
torch.manual_seed(seed)
loader = DataLoader(train if train is not None else train_set,
batch_size=128, shuffle=True)
val = DataLoader(val_set, batch_size=512)
opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9,
weight_decay=5e-4)
sched = torch.optim.lr_scheduler.OneCycleLR(
opt, max_lr=lr, total_steps=epochs * len(loader))
loss_fn = nn.CrossEntropyLoss()
for _ in range(epochs):
model.train()
for xb, yb in loader:
opt.zero_grad()
loss_fn(model(xb), yb).backward()
opt.step()
sched.step()
model.eval()
def score(dl):
right = seen = 0
with torch.no_grad():
for xb, yb in dl:
right += (model(xb).argmax(1) == yb).sum().item()
seen += yb.numel()
return right / seen
if report_train:
return score(DataLoader(train if train is not None else train_set,
batch_size=512)), score(val)
return score(val)
def plain(depth, norm):
torch.manual_seed(0)
def block(cin, cout):
layers = [nn.Conv2d(cin, cout, 3, padding=1, bias=not norm)]
if norm:
layers.append(nn.BatchNorm2d(cout))
return layers + [nn.ReLU()]
layers = block(3, 32)
for _ in range(depth):
layers += block(32, 32)
layers += [nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(32, 10)]
return nn.Sequential(*layers)
print('%18s %8s %12s %12s' % ('', 'depth', 'train acc', 'val acc'))
for norm in [False, True]:
for depth in [2, 6, 12]:
tr, va = fit(plain(depth, norm), epochs=4, report_train=True)
print('%18s %8d %12.4f %12.4f'
% ('with batchnorm' if norm else 'no normalisation',
depth, tr, va))
no normalisation 2 0.2120 0.2220
no normalisation 6 0.1050 0.0975
no normalisation 12 0.1055 0.0990
with batchnorm 2 0.3465 0.3395
with batchnorm 6 0.3740 0.3675
with batchnorm 12 0.3827 0.3665
Read the top three rows first. Without normalisation, adding depth does not merely fail to help: the six and twelve layer stacks collapse to 0.10, which on ten classes is chance. They learned nothing at all, and they are strictly more powerful than the two layer version that learned something.
Look at the training accuracy, not the validation accuracy
If the deeper model were overfitting, its training accuracy would be higher and its validation accuracy lower. It is not: the deeper network is worse at fitting the data it was trained on, despite having strictly more capacity. Any function the shallow network computes, the deep one could compute by making its extra layers do nothing at all.
So it is not failing for lack of capacity. It is failing because gradient descent cannot find the good solution, which is an optimisation problem rather than a statistical one.
The bottom three rows show batch normalisation rescuing it, and that is the historical order. Normalisation made stacks of this depth trainable, networks then got deeper, and the same wall reappeared. Residual connections are the general fix, and they are what allow hundreds of layers rather than tens.
The residual block
from torch import nn
class ResidualBlock(nn.Module):
def __init__(self, channels):
super().__init__()
self.body = nn.Sequential(
nn.Conv2d(channels, channels, 3, padding=1, bias=False),
nn.BatchNorm2d(channels), nn.ReLU(),
nn.Conv2d(channels, channels, 3, padding=1, bias=False),
nn.BatchNorm2d(channels))
self.act = nn.ReLU()
def forward(self, x):
return self.act(x + self.body(x)) # the whole idea
torch.manual_seed(0)
block = ResidualBlock(16)
x = torch.randn(2, 16, 8, 8)
print('in ', tuple(x.shape), ' out', tuple(block(x).shape))
# with the body zeroed, the block is the identity
with torch.no_grad():
for prm in block.body.parameters():
prm.zero_()
print('body zeroed, output equals relu(input):',
bool(torch.allclose(block(x), torch.relu(x))))
body zeroed, output equals relu(input): True
What it does to the gradient
from torch import nn
torch.manual_seed(0)
def gradient_at_input(residual, depth=20):
x = torch.randn(4, 16, requires_grad=True)
h = x
layers = [nn.Sequential(nn.Linear(16, 16), nn.ReLU())
for _ in range(depth)]
for layer in layers:
h = h + layer(h) if residual else layer(h)
h.sum().backward()
return x.grad.norm().item()
torch.manual_seed(0)
print('plain stack of 20 gradient norm at the input %.3e'
% gradient_at_input(False))
torch.manual_seed(0)
print('residual stack of 20 gradient norm at the input %.3e'
% gradient_at_input(True))
residual stack of 20 gradient norm at the input 4.735e+02
The addition gives the gradient a route back to the early layers that does not pass through any weights at all. Differentiating x + f(x) gives 1 + f'(x), and that 1 is a path along which the gradient arrives undiminished no matter how many blocks it has travelled through.