Unit 04.01: Model capacity and hidden layers
How wide should the hidden layer be? There is no formula, but there is a measurement — and the number to watch is not the one most people look at.
Read the gap, not the training score
Capacity is roughly how complex a function the model can represent. Too little and it cannot fit the pattern at all. Too much and it fits noise in the training data.
The task here is a circle, which no linear boundary can separate. Compare four capacities and watch both columns:
import torch
from torch import nn
torch.manual_seed(0)
X = torch.rand(400, 2) * 2 - 1
y = ((X[:, 0] ** 2 + X[:, 1] ** 2) < 0.5).float().unsqueeze(1) # a circle
Xtr, ytr, Xva, yva = X[:300], y[:300], X[300:], y[300:]
def evaluate(hidden):
torch.manual_seed(0)
layers = [nn.Linear(2, hidden), nn.ReLU(), nn.Linear(hidden, 1)] if hidden else [nn.Linear(2, 1)]
model = nn.Sequential(*layers)
opt = torch.optim.Adam(model.parameters(), lr=0.05)
for _ in range(400):
opt.zero_grad()
nn.BCEWithLogitsLoss()(model(Xtr), ytr).backward()
opt.step()
with torch.no_grad():
acc = lambda a, b: ((torch.sigmoid(model(a)) > 0.5).float() == b).float().mean().item()
return acc(Xtr, ytr), acc(Xva, yva), sum(p.numel() for p in model.parameters())
print(f"{'hidden':>7} {'params':>7} {'train':>7} {'val':>7}")
for hidden in (0, 4, 16, 64):
tr, va, params = evaluate(hidden)
print(f"{hidden:>7} {params:>7} {tr:>7.3f} {va:>7.3f}")
# No hidden layer cannot represent a circle at all. More width helps until it
# stops helping -- watch the gap between train and val, not train alone.
With no hidden layer the model cannot represent a circle at all — train and validation are both poor, which is underfitting. As width grows, training accuracy climbs steadily. Validation climbs too, then flattens.
That flattening is the signal. Once validation stops improving, extra width is buying memorisation.
The mistake this prevents
Choosing width by watching training accuracy. It improves almost without limit, so it always suggests "bigger". The validation column is the one that tells you when to stop.
Takeaway
Increase capacity while validation improves. Stop when only training improves.
