Training, Regularization, Optimization

The Training Loop

Training a model means iteratively adjusting parameters θ to minimize a loss function L(θ) on the training data. The fundamental update rule is gradient descent:

θ ← θ − α ∇_θ L(θ)

where α is the learning rate and ∇_θ L is the gradient of the loss with respect to all parameters.

Gradient Descent Variants

Batch Gradient Descent

Use ALL training examples to compute the gradient:

∇L = (1/n) Σᵢ ∇L(xᵢ, yᵢ; θ)

Pros: stable gradient estimate, guaranteed convergence. Cons: O(n) per update — impractical for large n.

Stochastic Gradient Descent (SGD)

Use ONE example at a time:

θ ← θ − α ∇L(xᵢ, yᵢ; θ)

Pros: one update per example, escapes local minima via noise. Cons: noisy, bouncy loss curve, poor hardware utilization.

Mini-Batch SGD

Use a small batch of B examples (typically 32–512):

θ ← θ − α (1/B) Σᵢ∈batch ∇L(xᵢ, yᵢ; θ)

Pros: good hardware utilization (GPU vectorization), moderate noise, practical speed.

This is the default in deep learning.

# PyTorch training loop (mini-batch SGD)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

for epoch in range(n_epochs):
    for batch_x, batch_y in dataloader:   # dataloader shuffles each epoch
        optimizer.zero_grad()
        y_pred = model(batch_x)
        loss = criterion(y_pred, batch_y)
        loss.backward()     # backpropagation: compute gradients
        optimizer.step()    # update parameters

Advanced Optimizers

Momentum

Accumulates an exponentially decaying moving average of past gradients:

v ← β v + (1−β) ∇L θ ← θ − α v

β = 0.9 is a common default. Momentum dampens oscillations and accelerates progress in the relevant direction.

RMSProp

Adapts learning rate per parameter by dividing by the root mean square of recent gradients:

s ← ρ s + (1−ρ) (∇L)² θ ← θ − α ∇L / (√s + ε)

Adam (Adaptive Moment Estimation)

Combines momentum (first moment) and RMSProp (second moment):

m ← β₁ m + (1−β₁) ∇L # first moment v ← β₂ v + (1−β₂) (∇L)² # second moment m̂ = m/(1−β₁ᵗ) # bias correction v̂ = v/(1−β₂ᵗ) # bias correction θ ← θ − α m̂ / (√v̂ + ε)

Defaults: α=1e-3, β₁=0.9, β₂=0.999, ε=1e-8.

Adam is the most popular optimizer for deep learning — fast convergence, works well without tuning.

AdamW

Adam with decoupled weight decay (L2 penalty applied directly to weights, not the adaptive term):

θ ← θ − α m̂/(√v̂+ε) − λ θ

Recommended over Adam for Transformers and large models where regularization matters.

import torch.optim as optim

optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, nesterov=True)

Learning Rate

The most important hyperparameter. Too large: loss diverges. Too small: slow convergence.

Learning Rate Schedule

from torch.optim.lr_scheduler import (
    CosineAnnealingLR, OneCycleLR, ReduceLROnPlateau, LinearLR
)

scheduler = CosineAnnealingLR(optimizer, T_max=n_epochs, eta_min=1e-6)
# or
scheduler = OneCycleLR(optimizer, max_lr=1e-3, steps_per_epoch=len(loader),
                        epochs=n_epochs)
# After each epoch or step:
scheduler.step()

Learning Rate Finder

Run a short experiment sweeping LR exponentially from low to high, plot loss vs. LR, pick the LR just before the loss stops falling:

# Using fastai's lr_find or pytorch-lightning's lr_finder
from torch_lr_finder import LRFinder
finder = LRFinder(model, optimizer, criterion)
finder.range_test(train_loader, end_lr=100, num_iter=100)
finder.plot()   # pick LR at steepest descent

Regularization

Regularization reduces overfitting by discouraging the model from fitting the training data too closely.

L2 Regularization (Weight Decay)

Add λ Σθ² to the loss. Penalizes large weights, encourages smooth solutions.

L1 Regularization

Add λ Σ|θ| to the loss. Induces sparsity — drives irrelevant weights to zero.

Dropout

During training, randomly set each neuron's output to zero with probability p:

self.dropout = nn.Dropout(p=0.5)
# In forward:
x = self.dropout(x)   # active only during training; nn.Dropout handles eval mode

At test time, all neurons are active, but their outputs are scaled by (1−p) — equivalent to ensembling exponentially many subnetworks.

Dropout rates: 0.5 for FC layers, 0.1–0.3 for convolutional/transformer layers.

Batch Normalization

Normalize the output of each layer across the mini-batch to have mean 0 and std 1, then scale and shift with learnable parameters γ, β:

BN(x) = γ · (x − μ_B)/√(σ²_B + ε) + β

Effects: allows higher learning rates, reduces sensitivity to initialization, has slight regularizing effect, accelerates training. Works best with batch size ≥ 32.

self.bn = nn.BatchNorm1d(256)   # for FC layers
self.bn = nn.BatchNorm2d(64)    # for conv layers (per channel)

Layer Normalization

Normalize across the feature dimension instead of batch dimension. Used in Transformers (works at batch size 1).

self.ln = nn.LayerNorm(d_model)

Spectral Normalization

Constrain the weight matrix by its largest singular value. Used in GANs for training stability.

Data Augmentation

Expand the effective training set by applying label-preserving transformations:

from torchvision import transforms

train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
])

Early Stopping

Monitor validation loss; stop training when it stops improving:

best_val_loss = float("inf")
patience, wait = 10, 0

for epoch in range(max_epochs):
    train_one_epoch(model, optimizer, train_loader)
    val_loss = evaluate(model, val_loader)

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        wait = 0
        torch.save(model.state_dict(), "best_model.pt")
    else:
        wait += 1
        if wait >= patience:
            print(f"Early stopping at epoch {epoch}")
            break

model.load_state_dict(torch.load("best_model.pt"))

Backpropagation

Backpropagation is the algorithm for computing gradients of the loss with respect to all parameters in a neural network, using the chain rule:

∂L/∂θₗ = ∂L/∂aₗ · ∂aₗ/∂θₗ

where aₗ is the pre-activation of layer l. Forward pass computes activations; backward pass computes gradients layer by layer from output to input.

Vanishing Gradients

In deep networks, gradients can shrink exponentially as they propagate backward (especially with sigmoid/tanh activations):

∂L/∂θ₁ = ∂L/∂aₙ · Π (∂aₗ/∂aₗ₋₁)

If each factor < 1, the product → 0 for deep L. Solutions: - Use ReLU activations (gradient = 1 when active) - Residual connections (skip connections carry gradient directly) - Batch normalization (prevents saturation) - Careful initialization (Xavier, He)

Exploding Gradients

Gradients can grow exponentially, especially in RNNs. Solutions: - Gradient clipping: if ‖∇L‖ > max_norm, scale down: ∇L ← ∇L · max_norm/‖∇L‖

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Weight Initialization

Poor initialization can cause vanishing/exploding activations before training even starts.

# PyTorch default is Kaiming uniform for Conv/Linear
nn.init.kaiming_normal_(layer.weight, mode="fan_out", nonlinearity="relu")
nn.init.xavier_uniform_(layer.weight)
nn.init.zeros_(layer.bias)

Hyperparameter Tuning

from sklearn.model_selection import RandomizedSearchCV
from scipy.stats import loguniform

param_dist = {
    "learning_rate": loguniform(1e-4, 1e-1),
    "n_estimators": [100, 200, 500],
    "max_depth": [3, 5, 7, 9],
    "subsample": [0.6, 0.8, 1.0]
}
search = RandomizedSearchCV(model, param_dist, n_iter=50, cv=5, n_jobs=-1)
search.fit(X_train, y_train)
print(search.best_params_)

Mixed Precision Training

Use 16-bit floats (FP16) for most operations, keep 32-bit (FP32) for critical accumulations. Speeds up training 2–4× on modern GPUs with tensor cores:

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
with autocast():
    output = model(x)
    loss = criterion(output, y)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

Gradient Accumulation

Simulate large batch sizes when GPU memory is limited:

accumulation_steps = 8   # simulate batch_size * 8

optimizer.zero_grad()
for i, (x, y) in enumerate(loader):
    output = model(x)
    loss = criterion(output, y) / accumulation_steps
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()