PyTorch University Path · Core lesson

Optimization, initialization, and generalization

Prerequisites: gradients, a training loop, cross-entropy, and train/validation/test splits. The examples continue the models in logistic regression and MLP classification.

Download the supervised CPU runner. It supports --optimizer sgd|adam, explicit learning rate and weight decay, deterministic generated data, and optional checkpoint writing.

Quick map

1. Decide whether the failure is capacity, optimization, or generalization.
2. Check shapes, finite values, and gradients.
3. Overfit one tiny batch.
4. Compare SGD and Adam with controlled settings.
5. Add regularization only when the evidence shows overfitting.

Separate three different problems

Use one example throughout: a classifier has training loss 1.2 and validation loss 1.3. If both stay flat, suspect optimization. If training falls to 0.1 while validation rises to 1.8, suspect overfitting. If both are low but the test set fails, inspect the split and data match.

Suppose training accuracy is low. The model may be too limited, the optimizer may not find good parameters, or the data may not contain enough useful signal.

Suppose training accuracy is high but validation accuracy is low. The model may have memorized details that do not repeat in new examples.

These are not the same failure.

Optimization means finding parameter values that reduce training loss. Generalization means performing well on independent examples from the intended use.

Regularization means adding constraints or noise that can reduce dependence on fragile training details. A regularizer can improve generalization while making training loss slightly worse.

underfitting: training performance is poor and validation performance is poor
overfitting: training performance is strong but validation performance is much worse
healthy fit: both are useful, with a believable validation gap

Do not diagnose from one final number. Plot or print training and validation loss across epochs.

Inspect a baseline and class-level errors. Check data leakage before celebrating unusually good validation results.

SGD follows the current gradient

Calculate one update: with parameter 2.0, gradient 0.6, and learning rate 0.1, SGD gives 2.0 - 0.1 × 0.6 = 1.94. Predict this direction before comparing optimizers.

Stochastic gradient descent, or SGD, subtracts a learning-rate-scaled gradient from each parameter. “Stochastic” means the gradient usually comes from a sampled mini-batch rather than the complete dataset.

parameter_next = parameter_now - learning_rate * gradient

The learning rate controls step size. If it is far too small, loss changes slowly and training appears stuck.

If it is too large, updates jump across useful regions; loss can oscillate, rise, or become nan. Momentum keeps a moving direction from recent gradients.

It can speed movement through a long shallow valley and reduce side-to-side bouncing.

velocity_next = momentum * velocity_now + gradient
parameter_next = parameter_now - learning_rate * velocity_next

The runner uses momentum 0.9 for SGD. That does not make 0.9 universally correct.

It is a reasonable starting point for the small teaching tasks. A controlled learning-rate experiment might compare 0.008, 0.08, and 0.8 while holding seed, model, batches, and epochs fixed.

Adam gives parameters adaptive step scales

Adam keeps moving averages of each parameter's gradient and squared gradient. The first average estimates direction.

The second estimates recent gradient magnitude. Their ratio gives each parameter an adaptive scale.

Adam also corrects the averages early in training because they start at zero.

first_average = beta1 * old_first + (1 - beta1) * gradient
second_average = beta2 * old_second + (1 - beta2) * gradient_squared

update is proportional to corrected_first / (sqrt(corrected_second) + epsilon)

Adam often reaches a useful result with less learning-rate tuning than plain SGD. It is not automatically better on every task.

Optimizers follow different paths and can produce different generalization. Compare them with the same splits and a fair search over learning rates.

Do not compare one carefully tuned Adam run against one arbitrary SGD run.

# Two valid experiments; compare held-out results, not optimizer reputation.
torch.optim.SGD(model.parameters(), lr=0.08, momentum=0.9)
torch.optim.Adam(model.parameters(), lr=0.01)

Call optimizer.zero_grad(set_to_none=True) before each backward pass because PyTorch accumulates gradients. Then call loss.backward() and optimizer.step(). Reversing or omitting these operations changes the algorithm.

Initialization sets the starting signal scale

Every trainable layer needs initial values before the first update. If every hidden unit starts with identical weights, those units receive identical gradients and remain duplicates.

Random initialization breaks this symmetry. The random scale also matters.

Very large weights can create huge activations or saturated nonlinearities. Very small weights can shrink signals and gradients through many layers.

PyTorch's linear and convolution layers use sensible default initializers based on their input size. ReLU networks often use Kaiming-style initialization, which aims to preserve activation variance through layers.

Tanh networks often use Xavier-style initialization. You usually should begin with framework defaults, verify signal and gradient behavior, and customize only with a reason.

fan_in = number of inputs feeding one unit

Kaiming scale for ReLU is related to sqrt(2 / fan_in)
Xavier scale balances fan_in and fan_out

Initialization does not replace normalization, careful architecture, or learning-rate choice. It only gives training a reasonable starting point. To debug initialization, print activation means and standard deviations for one batch, confirm logits are finite, and inspect whether gradients are zero or enormous after one backward pass.

Weight decay discourages large parameter values

Weight decay gently pulls weights toward zero during optimization. For ordinary SGD, it closely matches adding an L2 penalty to the objective. A coefficient controls the strength.

regularized_objective = data_loss + lambda * sum(weight_squared)

The idea is not that small weights are always correct. It is that among solutions that fit training data, a smoother lower-magnitude solution may depend less on accidental details.

Too much decay causes underfitting. The runner defaults to 1e-4, and the CLI lets you test alternatives with --weight-decay.

Biases and normalization parameters are often excluded from weight decay in larger projects, while ordinary model weights receive it. The teaching runner applies one setting to all parameters for clarity.

AdamW separates weight decay from Adam's adaptive gradient update and is common in transformer training. That distinction matters when reproducing a published recipe.

Dropout trains many noisy subnetworks

During training, dropout randomly sets some activations to zero and scales the survivors. This stops the network from relying on exactly the same hidden path every time. During evaluation, dropout is disabled and the complete network is used.

model = nn.Sequential(
    nn.Linear(input_size, 64),
    nn.ReLU(),
    nn.Dropout(p=0.2),
    nn.Linear(64, class_count),
)

model.train()  # dropout is random and active
model.eval()   # dropout is deterministic and inactive

The probability p=0.2 means a hidden activation has a 20% chance of being dropped during training. More dropout is not always more protection.

On a tiny or already difficult model, strong dropout can prevent learning. The supplied supervised runner intentionally omits dropout so the core tasks remain small and deterministic; add it as an experiment after the baseline works.

Dropout is also a common source of confusing evaluation. torch.no_grad() does not disable dropout.

Only model.eval() changes its mode. Conversely, model.eval() does not disable gradient recording.

Use both for validation.

Read the learning curves

Imagine training loss falls for 100 epochs. Validation loss falls until epoch 24, then rises while validation accuracy stays nearly flat.

The model is becoming more confident about patterns that do not transfer. Keeping the epoch-24 state is called early stopping.

The runner performs a simple version: after each epoch it copies the state when validation loss reaches a new minimum, then restores the best state before test evaluation.

Early stopping uses validation data and therefore is a model-selection method. It does not make the test set unnecessary. It also does not prove the model will handle distribution shift, such as new writing styles or image devices.

Useful responses to overfitting include collecting more representative data, reducing model size, increasing weight decay, adding modest dropout, using data augmentation that preserves labels, or training for fewer epochs. Choose one change at a time and keep the split fixed. If you change data, architecture, optimizer, and metric together, you cannot tell which change mattered.

Run controlled comparisons

.venv-learning/bin/python examples/pytorch/build_supervised.py \
  --task mlp --optimizer adam --learning-rate 0.01

.venv-learning/bin/python examples/pytorch/build_supervised.py \
  --task mlp --optimizer sgd --learning-rate 0.08

.venv-learning/bin/python -m unittest tests/test_supervised_builds.py

Expected output includes training loss movement, the best validation loss, the majority baseline, test accuracy, macro F1, and a confusion matrix. Do not treat one run as a universal optimizer ranking.

These generated XOR clusters are simple and balanced. Real text can be sparse, noisy, imbalanced, and shifted over time.

A reliable debugging sequence

1. Check one batch. Print dtypes, shapes, target range, and a few examples. Require logits shape (B, C) and targets shape (B,).
2. Check a forward pass. Assert all logits and the loss are finite.
3. Check gradients. After backward, confirm important parameters have nonzero finite gradients.
4. Overfit a tiny batch. A capable model should drive loss very low on perhaps 16 examples. If it cannot, suspect code, optimization, or insufficient capacity.
5. Restore full training. Compare against the baseline and watch training versus validation curves.
6. Inspect errors. A metric can hide one failed class or subgroup.

This sequence moves from local mechanical facts to broader statistical behavior. It is faster than randomly adding layers or changing optimizers.

Failure drill: good training, bad validation

First confirm there is no preprocessing mismatch between splits. Then verify labels have the same meaning and class IDs.

Try to overfit a tiny batch; that should succeed and confirms basic mechanics. Next reduce model width or stop at the best validation epoch.

Compare modest weight decay or dropout one at a time. If nothing helps, inspect whether training and validation come from different populations.

Regularization cannot create information that is absent from input features.

A second failure looks different: both training and validation loss stay high. Adding dropout would likely make it worse. Check the learning rate, gradients, target encoding, and capacity before treating the problem as overfitting.

Practice checks with worked answers

Quick recap

SGD uses a shared learning rate and can use momentum. Adam adapts the scale of each parameter's update.

Initialization controls the starting signal and gradient scale. Weight decay and dropout can reduce fragile fitting, while early stopping keeps the state with the best validation evidence.

None of these tools repairs leakage, wrong labels, or a model that cannot represent the task. Diagnose shapes and gradients first, overfit a tiny batch, then compare complete experiments with fixed splits and useful baselines.

Optional next lesson: CNNs and spatial locality. Integration project: transformer capstone.