ML Math

Warmup over 1,000 steps to 3×10⁻⁴, then cosine decay to 3×10⁻⁵ by step 10,000. The rate is 0 at the start, 1.5×10⁻⁴ halfway through warmup, 3×10⁻⁴ at its peak, 1.65×10⁻⁴ halfway through the decay and 3×10⁻⁵ at the end. Warmup protects the fragile first steps; the decay lets training settle.

• A schedule tuned for one run length does not transfer; cosine decay depends on knowing the total number of steps.
• Warmup matters most with large batches and adaptive optimizers; small runs often do fine without it.
• Changing the schedule and the peak rate at once makes it impossible to tell which helped.

lr_max, lr_min, warmup, total = 3e-4, 3e-5, 1_000, 10_000

def lr(t):
    if t < warmup: return lr_max * t / warmup                     # linear warmup
    progress = (t - warmup) / (total - warmup)
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + np.cos(np.pi * progress))   # cosine decay

schedule = {t: float(f'{lr(t):.2e}') for t in (0, 500, 1_000, 5_500, 10_000)}

Assumes import numpy as np. Each snippet carries its own example numbers; the comments say which are made up or simulated.

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim, T_max=100)
scheduler = torch.optim.lr_scheduler.OneCycleLR(optim, max_lr=0.01, total_steps=1000)

# call after each epoch/step:
scheduler.step()