Training Neural Networks
Pick the optimizer, init, normalization, regularization and schedule. Then debug the run when the loss misbehaves.
- Concepts and terms
- Tasks and commands
- Annotated code
- Diagrams
- Defaults and numbers
- Troubleshooting
- Questions and tips
One AdamW step#
From Lesson 03.06
The training step#
From Lesson 03.06
Training loop with AdamW and cosine#
optimizer = optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)for epoch in range(100): optimizer.zero_grad()clear old gradients output = model(torch.randn(32, 784))forward pass loss = torch.nn.functional.cross_entropy(output, torch.randint(0, 10, (32,))) loss.backward()fill .grad by chain rule torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step()update the weights scheduler.step()after optimizer.step()Warmup then cosine, from the HuggingFace library
scheduler = get_cosine_schedule_with_warmup( optimizer, num_warmup_steps=2000, num_training_steps=100000)From Lesson 03.06, Lesson 03.09
Key formulas#
- Chain rule
- \frac{dy}{dx} = f'(g(x)) \cdot g'(x)
- Adam moments
- m_t = \beta_1 m_{t-1} + (1 - \beta_1) \cdot \text{gradient} and v_t = \beta_2 v_{t-1} + (1 - \beta_2) \cdot \text{gradient}^2
- Bias correction
- \hat{m} = \frac{m_t}{1 - \beta_1^t} and \hat{v} = \frac{v_t}{1 - \beta_2^t}
- AdamW update
- w = w - \text{lr} \cdot \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon} - \text{lr} \cdot \lambda w
- Cosine decay
- \text{lr}(t) = \text{lr}_{\min} + 0.5 \cdot (\text{lr}_{\max} - \text{lr}_{\min}) \cdot \left(1 + \cos\left(\frac{\pi t}{T}\right)\right)
- Linear warmup
- \text{lr}(t) = \text{lr}_{\max} \cdot \frac{t}{\text{warmup\_steps}} \quad \text{for } t < \text{warmup\_steps}
- Init variance
- Kaiming \mathrm{Var}(w) = \frac{2}{\text{fan}_{\text{in}}}, Xavier \mathrm{Var}(w) = \frac{2}{\text{fan}_{\text{in}} + \text{fan}_{\text{out}}}
- Inverted dropout
- \text{output} = \text{activation}(z) \cdot \text{mask} / (1 - p) during training
From Lesson 03.03, Lesson 03.06, Lesson 03.07, Lesson 03.08, Lesson 03.09
Pick by situation#
| Situation | Choice |
|---|---|
| Transformer or LLM | AdamW, lr=1e-4, wd=0.01-0.1 |
| CNN or ResNet | SGD + Momentum, lr=0.1, momentum=0.9 |
| GAN | Adam, lr=2e-4, beta1=0.5 |
| Fine-tuning | AdamW, lr=2e-5, wd=0.01 |
| Not sure yet | AdamW, lr=3e-4, wd=0.01 |
| Sigmoid or tanh layers | Xavier init |
| ReLU, GELU or Swish layers | Kaiming init |
| Transformer residual layers | scale by 1/sqrt(2N), N layers |
From Lesson 03.06, Lesson 03.08
Init, norm and regularization in code#
| Task | Code |
|---|---|
| Kaiming init for ReLU | nn.init.kaiming_normal_(layer.weight, |
| Xavier init | nn.init.xavier_uniform_(layer.weight) |
| BatchNorm after a Linear | nn.BatchNorm1d(256) |
| LayerNorm in a transformer | nn.LayerNorm(d_model) |
| Dropout in an MLP | nn.Dropout(0.3) |
| Cosine decay to a floor | CosineAnnealingLR(optimizer, |
| Find the first NaN or Inf | with torch.autograd.detect_anomaly(): |
| Measure the gradient norm | torch.nn.utils.clip_grad_norm_(model.parameters(), |
Optimizer defaults#
| Setting | Default | Note |
|---|---|---|
| Adam lr | 0.001 | change lr first |
| beta1 | 0.9 | almost never change |
| beta2 | 0.999 | change second if needed |
| epsilon | 1e-8 | almost never change |
| AdamW weight_decay | 0.01 | not scaled by the adaptive rate |
| SGD momentum | 0.9 | about the last 10 gradients |
| SGD weight decay on CNNs | 1e-4 | with step or cosine schedule |
| Gradient clip max_norm | 1.0 | clip by global norm |
From Lesson 03.06
Regularization and schedule defaults#
| Setting | Default | Note |
|---|---|---|
| Dropout, transformers | 0.1 | LayerNorm models |
| Dropout, MLPs | 0.5 | half the neurons off |
| BatchNorm momentum | 0.1 | 90% old plus 10% new |
| BatchNorm eps | 1e-5 | added to the variance |
| Warmup share | 3-5% | of total steps |
| Cosine floor eta_min | 1e-5 | lowest lr |
| Step decay factor | 0.1 | at fixed epochs |
| Log clamp eps | 1e-7 | keeps log(p) finite |
From Lesson 03.07, Lesson 03.09, Lesson 03.13
Troubleshooting#
| Symptom | Cause | Fix |
|---|---|---|
| Loss stuck at -log(1/num_classes) | model predicts a uniform distribution | check data pipeline and labels |
| Loss NaN after a few steps | learning rate too high | reduce lr by 10x |
| Loss NaN at once | log(0) or division by zero | clamp predictions to [eps, 1-eps] |
| Loss oscillates | no optimizer.zero_grad() | call it before loss.backward() |
| Test accuracy varies between runs | no model.eval() | add model.eval() and torch.no_grad() |
| Train acc high, test acc low | overfitting | dropout, weight decay, more data |
| Gradients all zero | dead ReLUs or detached graph | LeakyReLU, check .requires_grad |
| Early layers stop learning | sigmoid derivative is at most 0.25 | ReLU or GELU with Kaiming init |
From Lesson 03.13, Lesson 03.03, Lesson 03.08
Glossary#
- SGD
- stochastic gradient descent on a mini-batch.
- AdamW
- Adam with decoupled weight decay.
- BatchNorm
- normalizes each feature across the mini-batch.
- LayerNorm
- normalizes each sample across its features.
- RMSNorm
- LayerNorm without the mean subtraction.
- 1cycle
- lr goes up, then down, in one cycle.
- SGDR
- cosine annealing with warm restarts.
- LR range test
- raise lr during a short run and find where loss diverges.
Questions#
- Why AdamW and not Adam plus L2?
- Adam scales the L2 term by its adaptive rate. AdamW shrinks every weight by the same proportion.
- Why warmup?
- Adam moments start at zero. Large early steps follow bad statistics.
- Why bias correction?
- At step 1 with beta1 0.9, m is 0.1 times the gradient.
- Why LayerNorm in transformers?
- It does not depend on batch size or sequence length.
- Why Kaiming for ReLU?
- ReLU zeros half the activations, so the variance is 2/fan_in.
Tips#
- Run the overfit-one-batch test first. Loss must go to nearly zero.
- Try 3 learning rates 10x apart before you blame anything else.
- If you tune one hyperparameter, tune the learning rate.
- Print tensor shapes after every operation while you debug.
- Check that gradient sizes are similar across layers.
- Use
.to(device)on the model and the data.
From Lesson 03.13, Lesson 03.06