Training: Optimization and Regularization

The loss is the single number that training minimizes. Choosing the wrong one silently cripples your model.

• Mean Squared Error (MSE): $\mathcal{L} = \frac{1}{N}\sum_i (y_i - \hat{y}_i)^2$. Use for regression. Penalizes large errors quadratically — outliers dominate.
• Binary Cross-Entropy: $\mathcal{L} = -\frac{1}{N}\sum_i [y_i \log \hat{p}_i + (1-y_i)\log(1-\hat{p}_i)]$. Use for binary classification with a sigmoid output.
• Categorical Cross-Entropy: $\mathcal{L} = -\frac{1}{N}\sum_i \sum_k y_{ik} \log \hat{p}_{ik}$. Use for multiclass with softmax output.

Cross-entropy is preferred over MSE for classification because it is derived from maximum likelihood, produces sharper gradients near decision boundaries, and avoids the flat gradient saturation problem of MSE + sigmoid.

All neural network optimizers are variants of gradient descent:

$\theta \leftarrow \theta - \eta \nabla_\theta \mathcal{L}$

where $\eta > 0$ is the learning rate and $\nabla_\theta \mathcal{L}$ is the gradient of the loss with respect to parameters $\theta$.

Full-batch gradient descent computes the gradient over the entire dataset — exact but slow for large data. Stochastic gradient descent (SGD) uses a single example per step — noisy but fast. Mini-batch SGD (the standard) uses a batch of $B$ examples, balancing signal quality and compute:

for epoch in range(num_epochs):
    for X_batch, y_batch in dataloader:  # batch size B
        loss = compute_loss(model(X_batch), y_batch)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

The noise in SGD is actually beneficial: it acts as an implicit regularizer and helps escape sharp local minima.

Vanilla SGD oscillates in ravines and crawls on flat plateaus. Modern optimizers fix this.

Adam update ($\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$ are standard defaults):

$m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t, \quad v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2$ $\hat{m}_t = m_t/(1-\beta_1^t), \quad \hat{v}_t = v_t/(1-\beta_2^t)$ $\theta_t = \theta_{t-1} - \eta\, \hat{m}_t / (\sqrt{\hat{v}_t} + \epsilon)$

Adam is robust to learning-rate choice and converges fast. SGD + momentum often wins at convergence quality with careful tuning — many large-scale vision papers still use it.

The learning rate $\eta$ controls how large each step is. Get it wrong and nothing else matters.

• Too large: loss oscillates or diverges. Weights overshoot the minimum.
• Too small: training is painfully slow and may stall in flat regions.
• Just right: loss decreases smoothly and plateaus at a good value.

Learning rate schedules help:

• Step decay: halve $\eta$ every $k$ epochs.
• Cosine annealing: $\eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min})(1 + \cos(\pi t/T))$.
• Warmup: start tiny, ramp up over first few epochs, then decay. Critical for large batch training and transformers.

# PyTorch cosine annealing
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=num_epochs, eta_min=1e-6
)
for epoch in range(num_epochs):
    train_one_epoch()
    scheduler.step()

Every training run lives somewhere on this spectrum:

• Underfitting (high bias): training loss is high. Model too simple, or training stopped too early, or learning rate too small.
• Overfitting (high variance): training loss is low, validation loss is high. Model memorized training data instead of generalizing.
• Good fit: both training and validation loss are low and close.

Diagnose by watching the validation loss curve. If it stops falling and starts rising while training loss keeps falling — you're overfitting.

Four levers to fight overfitting:

1. More data (most effective).
2. Simpler model (fewer parameters).
3. Regularization (keep the model but constrain it).
4. Early stopping (stop before the gap grows).

The bias-variance decomposition: expected test error = (bias)$^2$ + variance + irreducible noise. Regularization trades a small increase in bias for a large drop in variance.

L2 regularization (weight decay) adds a penalty to the loss:

$\mathcal{L}_{\text{reg}} = \mathcal{L} + \frac{\lambda}{2} \|\theta\|_2^2$

The gradient update becomes $\theta \leftarrow \theta(1 - \eta\lambda) - \eta \nabla \mathcal{L}$. Weights shrink toward zero each step. In Adam, weight decay is applied separately from the adaptive scaling (AdamW) — this distinction matters in practice and AdamW is the preferred implementation.

Dropout randomly zeroes each neuron's output with probability $p$ during training:

import torch.nn as nn
# During training, 50% of activations zeroed; outputs scaled by 1/(1-p)
self.dropout = nn.Dropout(p=0.5)
x = self.dropout(self.relu(self.linear(x)))

At inference, dropout is disabled and all neurons contribute. Effect: the network can't rely on any single neuron — it must learn distributed, redundant representations. This reduces co-adaptation and acts like an ensemble of $2^n$ thinned networks.

Batch normalization (BatchNorm) normalizes the pre-activations within each mini-batch, then applies a learned scale $\gamma$ and shift $\beta$:

$\hat{z}_i = \frac{z_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}, \qquad y_i = \gamma \hat{z}_i + \beta$

where $\mu_B$ and $\sigma_B^2$ are the batch mean and variance. Benefits:

• Allows higher learning rates without divergence.
• Reduces sensitivity to weight initialization.
• Acts as a mild regularizer (the batch statistics add noise).

At inference, $\mu_B$ and $\sigma_B^2$ are replaced by running statistics accumulated during training. Placement: typically after the linear layer and before the activation (Linear → BN → ReLU). BatchNorm has become almost universal in CNNs; transformers use Layer Norm instead (normalizes over the feature dimension, not the batch — independent of batch size).

Early stopping: monitor validation loss after each epoch. Stop training when validation loss stops improving for $k$ consecutive epochs (patience $k$), and restore the best checkpoint.

best_val_loss = float('inf')
patience, epochs_no_improve = 10, 0
for epoch in range(max_epochs):
    train_loss = train_one_epoch()
    val_loss = evaluate(val_loader)
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), 'best.pt')
        epochs_no_improve = 0
    else:
        epochs_no_improve += 1
        if epochs_no_improve == patience:
            break
model.load_state_dict(torch.load('best.pt'))

A production-grade training loop stacks all of this: AdamW optimizer, cosine LR schedule with warmup, dropout and weight decay in the model, BatchNorm for stable training, early stopping to avoid overfitting. Each component fixes a specific failure mode — none is optional on real problems.

Things that silently wreck training runs, ranked by how often they bite:

• Forgot optimizer.zero_grad(): gradients accumulate across batches, giving wrong updates.
• Wrong loss for the task: MSE on a classification head (should be cross-entropy), or cross-entropy with logits passed through sigmoid first (double-applies sigmoid — use BCEWithLogitsLoss).
• Data not normalized: raw pixel values [0, 255] cause wildly different gradient scales vs. [-1, 1]. Always normalize inputs to zero mean and unit variance.
• Learning rate too high with Adam: Adam is sensitive — start with $\eta = 3 \times 10^{-4}$ and tune from there.
• Using weight decay with vanilla Adam instead of AdamW: in vanilla Adam, weight decay gets absorbed into the adaptive scaling and is much weaker than intended.
• Evaluating with dropout active: always call model.eval() before validation and model.train() before training.