Pruning

A human baby has far more connections in its brain than an adult does. As we grow, the unused ones are removed, and we get smarter. Networks are the same. Many weights end up close to zero, and they contribute almost nothing. So what happens if we just delete them?

• Babies have more connections
• Unused ones are removed
• Delete the weights that do nothing

How far can we go?

We train a small network. Then we copy it, and remove the weights with the smallest values. Half, then eighty percent, then ninety five percent. We won't do any retraining. Run it. Remove half of all the weights, and accuracy barely moves. Half of the network was doing almost nothing. At eighty percent, it starts to hurt, and at ninety five, it's badly damaged. But that's without retraining. In practice, you prune a little, retrain a little, and repeat. And that way, you can often remove around ninety percent with almost no loss.

import torch
from torch.nn.utils import prune
from base import load, net, fit, score

torch.manual_seed(0)
X, y, Xt, yt = load(4000)
model = fit(net(), X, y, epochs=3)
print(f"Full model: {score(model, Xt, yt):.3f}")

for amount in [0.5, 0.8, 0.95]:
    m = fit(net(), X[:1], y[:1], epochs=0)
    m.load_state_dict(model.state_dict())
    for layer in (m[1], m[3]):
        prune.l1_unstructured(layer, "weight", amount)
    print(f"{amount:.0%} of weights removed: {score(m, Xt, yt):.3f}")

Full model: 0.858
50% of weights removed: 0.842
80% of weights removed: 0.633
95% of weights removed: 0.383

import torch
from torchvision import datasets, transforms

def load(n=6000):
    t = transforms.ToTensor()
    a = datasets.MNIST("data", True, t, download=False)
    b = datasets.MNIST("data", False, t, download=False)
    return a.data[:n].float()[:, None] / 255, a.targets[:n], b.data[:2000].float()[:, None] / 255, b.targets[:2000]

def net():
    return torch.nn.Sequential(torch.nn.Flatten(), torch.nn.Linear(784, 128), torch.nn.ReLU(),
                               torch.nn.Linear(128, 10))

def fit(model, X, y, lr=1e-3, epochs=3, opt=None):
    opt = opt or torch.optim.Adam(model.parameters(), lr=lr)
    for _ in range(epochs):
        for i in torch.randperm(len(X)).split(64):
            loss = torch.nn.functional.cross_entropy(model(X[i]), y[i])
            opt.zero_grad(); loss.backward(); opt.step()
    return model

def score(model, X, y):
    return (model(X).argmax(1) == y).float().mean().item()

Pruning in practice

Be aware of one catch. Setting a weight to zero doesn't make the model faster by itself, because the maths still runs, with zeros. You need sparse hardware, or you need to prune whole channels, which really does make the network smaller. Researchers also found the lottery ticket hypothesis. Hidden within big networks are tiny ones, which would've trained just as well alone.

Summary