class SquareRootScheduler:
def __init__(self, lr=0.1):
self.lr = lr
def __call__(self, epoch):
return self.lr * pow(epoch + 1.0, -0.5)
scheduler = SquareRootScheduler(lr=0.3)
d2l.plot(d2l.arange(num_epochs), [scheduler(t) for t in range(num_epochs)],
xlabel='epoch', ylabel='learning rate')Baseline: constant \eta = 0.3
train loss 0.064, train acc 0.977, test acc 0.892Both failure modes on one plot: loss noisy to the end (riding the noise floor), and test accuracy stalls while train accuracy climbs — overfitting.
Square-root decay
\eta_t = \eta_0 (t+1)^{-1/2} — the convex-optimal rate from 9.3. A scheduler is just a callable:
Square-root decay: training
train loss 0.095, train acc 0.968, test acc 0.901Smoother and quieter — but test accuracy lands below the constant baseline. Timid at both ends: gives up the high early rate within ~3 epochs, yet ends with the largest tail \eta of any decay here. Shape matters.
Multiplicative and piecewise constant
\eta_t = \max(\eta_{\min}, \eta_0\, \alpha^t) — aggressive, floor as safety net:
class FactorScheduler:
def __init__(self, factor=1, stop_factor_lr=1e-7, base_lr=0.1):
self.factor = factor
self.stop_factor_lr = stop_factor_lr
self.base_lr = base_lr
def __call__(self, epoch):
return max(self.stop_factor_lr, self.base_lr * self.factor ** epoch)
scheduler = FactorScheduler(factor=0.9, stop_factor_lr=0.01, base_lr=0.3)
d2l.plot(d2l.arange(50), [scheduler(t) for t in range(50)],
xlabel='epoch', ylabel='learning rate')
Piecewise constant: ride each noise floor until progress stalls, then cut the rate — the ImageNet-era staircase:
class MultiFactorScheduler:
def __init__(self, milestones, factor, base_lr):
self.milestones = milestones
self.factor = factor
self.base_lr = base_lr
def __call__(self, epoch):
lr = self.base_lr
for milestone in self.milestones:
if epoch >= milestone:
lr *= self.factor
return lr
scheduler = MultiFactorScheduler(milestones=[15, 25], factor=0.5, base_lr=0.3)
d2l.plot(d2l.arange(num_epochs), [scheduler(t) for t in range(num_epochs)],
xlabel='epoch', ylabel='learning rate')
Cosine decay
\eta_t = \eta_T + \frac{\eta_0 - \eta_T}{2}(1 + \cos(\pi t / T))
One parameter, no milestones, no kinks (Loshchilov & Hutter, 2016).
train loss 0.094, train acc 0.968, test acc 0.908The stay-high, decay-hard shapes (multiplicative, piecewise, cosine) beat the baseline and are too close to call from single runs. Cosine won on convenience, not measured superiority.
Warmup: dead vs. alive
Even a BatchNorm net has a ceiling. Cold start at \eta = 7.5 (25× the baseline) — collapsed to chance within one epoch:
train loss 2.314, train acc 0.099, test acc 0.100Same rate through a 5-epoch ramp — trains to 80–90%:
train loss 0.215, train acc 0.918, test acc 0.895Warmup + cosine
The default recipe of the late 2010s, and still strong:
train loss 0.048, train acc 0.987, test acc 0.906
WSD: the loss cliff
train loss 0.032, train acc 0.993, test acc 0.908Plateau loss sits above cosine’s, then drops in a cliff when the decay begins. Same final range as cosine at the same budget.
Branching off the plateau
Train warmup + stable only (no horizon committed), keep going as long as you like, then clone and decay whenever you want a finished model:
net_branch = copy.deepcopy(net_plateau)
decay = WSDScheduler(max_update=6, decay_steps=6, base_lr=0.3, final_lr=0.01)
train(net_branch, train_iter, test_iter, 6, loss,
torch.optim.SGD(net_branch.parameters(), lr), device, decay,
animator=board, epoch_offset=30)
board.fig
train loss 0.016, train acc 0.998, test acc 0.909The branch matches the full cosine run’s range — with the horizon chosen after the fact. One long run + cheap branched decays = models at many budgets (Hägele et al., 2024).
Plain SGD carries no state, so cloning parameters sufficed. With momentum or Adam, branch the optimizer state too.