Overview

There are a lot of settings we can adjust when training a neural network. Do we need to try all of them? Thankfully, no. Think of the ideas in this tutorial as options to draw on when you encounter a problem, not a checklist you need to complete for every model.

We'll start with a few checks that help us see whether our model is learning, then look at changes that might help when it isn't. These tips can help us make better use of our data, but they aren't a replacement for good examples, reliable labels, or a sensible evaluation setup. A small, working baseline is a good place to begin.

Resources

Yoav Goldberg's Neural Network Methods for Natural Language Processing [2017] provides practical advice for training neural networks.
Sebastian Ruder's Deep Learning for NLP Best Practices [2017] provides broader task-specific background. We'll discuss RNNs and transformers in later units.

Outcomes

After completing this lesson, you'll be able to …

read training and validation losses and decide what to check next
select a checkpoint using validation performance and an explicit early stopping rule
distinguish sampling, loss weighting, and regularization
read PyTorch code for evaluation, accumulation, and gradient clipping
investigate vanishing or exploding gradients and choose a reasonable next step

Before changing the model

Before reaching for a bigger network or another training trick, let's make sure the basics are working.

Check input shapes, labels, padding, and the loss function. Are the examples and labels aligned?
Try to overfit a tiny training batch, without dropout or other regularization. If our model can't learn even a few examples, let's check the code and data before making it larger.
Compare training and validation performance. If learning becomes unstable, check gradients and outputs for values such as NaN or infinity.
Change one thing at a time. Keep the data split and evaluation procedure the same, and jot down what changed.

If we can fit that tiny batch, that's a useful first sign: our model and training code can learn something. We still need to see how it does on examples it hasn't trained on.

Overfitting and early stopping

Our training loss is falling. Great! What about validation? If that loss keeps rising, our model may be getting better at fitting the training examples without getting better at the task. That's a sign of possible overfitting, and a reason to pause before simply training longer.

Reading learning curves

Let's look at the learning-curve visual. We've made up these losses to illustrate the idea; they aren't results from the code below. Validation loss rises a little at epoch 3, improves at epoch 4, then stops improving. Meanwhile, training loss keeps falling.

Defining the stopping rule

What should count as progress? First, choose a metric and decide whether lower or higher is better. For validation loss, we want lower values. Patience tells us how many checks in a row we'll wait without enough improvement before stopping. When the metric improves enough, we reset the counter.

We'll use patience 2 and strict improvement: any lower validation loss counts^†. With losses 0.60, 0.50, 0.51, 0.49, 0.50, 0.52, we stop after epoch 6 and restore the checkpoint from epoch 4. Patience counts checks, not necessarily epochs. For instance, we could validate every 10 batches.

Should a tiny change count as progress? We can use min_delta to set how much improvement we're looking for. Decide this and how often we'll validate before starting the run.

A small PyTorch training example

Let's try a small example with two classes. We'll use integer labels and apply cross entropy to raw logits, then train, validate, and return to our chosen checkpoint. It runs on the CPU without a wrapper library. The toy data lets us follow the steps; it doesn't tell us how well this model would do on an NLP task.

import copy

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(42)

# Two input features and an integer class label for each example.
x = torch.randn(80, 2)
y = (x[:, 0] + x[:, 1] > 0).long()

# Keep validation examples out of the training loader.
train_loader = DataLoader(
    TensorDataset(x[:64], y[:64]), batch_size=8, shuffle=True
)

validation_loader = DataLoader(
    TensorDataset(x[64:], y[64:]), batch_size=8
)

# Two output scores (logits), one for each class.
model = nn.Sequential(
    nn.Linear(2, 8),
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(8, 2),
)

# Cross entropy expects logits, not softmax probabilities.
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=0.01)

# Stop after two validation checks without a lower loss.
patience = 2
min_delta = 0.0
# we'll start with the lowest of bars...
# Any finite validation loss will improve on infinity.
best_loss = float("inf")
best_state = None
checks_without_improvement = 0

# Thirty epochs is a maximum, not a promise to train that long.
for epoch in range(1, 31):
    # Training mode enables dropout.
    model.train()
    training_total = 0.0

    for inputs, targets in train_loader:
        # Don't carry gradients over from the previous update.
        optimizer.zero_grad()

        logits = model(inputs)
        loss = criterion(logits, targets)

        # Compute gradients, then use them to update parameters.
        loss.backward()
        optimizer.step()

        # Weight the batch mean by its number of examples.
        training_total += loss.item() * len(targets)

    # Evaluate without dropout or gradient recording.
    model.eval()
    validation_total = 0.0
    with torch.no_grad():
        for inputs, targets in validation_loader:
            logits = model(inputs)
            loss = criterion(logits, targets)
            validation_total += loss.item() * len(targets)

    training_loss = training_total / len(train_loader.dataset)
    validation_loss = validation_total / len(validation_loader.dataset)

    # A nonfinite loss needs investigation, not more patience.
    if not torch.isfinite(torch.tensor(validation_loss)):
        raise RuntimeError("Nonfinite validation loss: inspect the run.")

    print(epoch, training_loss, validation_loss)

    if validation_loss < best_loss - min_delta:
        best_loss = validation_loss
        # Copy the weights so later updates can't alter this checkpoint.
        best_state = copy.deepcopy(model.state_dict())
        checks_without_improvement = 0
    else:
        # Only consecutive checks without improvement count.
        checks_without_improvement += 1
        if checks_without_improvement >= patience:
            break

# Return to the selected checkpoint, not necessarily the last epoch.
assert best_state is not None
model.load_state_dict(best_state)
model.eval()

Why multiply each batch's mean loss by its size? Imagine the final batch has only two examples while the others have eight. We don't want those two examples to count as much as a whole batch of eight. Here, we're using ordinary, unweighted cross entropy with no ignored labels.

We copy the model's state so later updates won't change our saved checkpoint. For a longer experiment, we'd save it to disk. If we want to resume training, we'll also need the optimizer state and other run information. For evaluating our chosen model here, restoring its weights is enough.

Imbalanced data

Suppose 90% of our examples have the same label. A model that always predicts that label gets 90% accuracy without learning to recognize the other class at all. If that rare class is important to our task, we'll want to look more closely.

Some labels are much more common than others. That's normal in many NLP tasks, so we don't necessarily want to make every dataset perfectly balanced. First, check how a majority-class baseline does, and look at performance for each class. A model can have high accuracy while doing poorly on the rare class we care about.

Sampling versus loss weighting

We have a few options. Undersampling removes examples from the majority class, though we might lose useful variation. Weighted sampling lets us draw some examples more often. Loss weighting lets some examples contribute more to the loss instead. These change different things, and we still need to check how well each class is doing.

The WeightedRandomSampler takes one weight per example. Class weights must first be mapped onto the examples:

import torch
from torch.utils.data import DataLoader, TensorDataset, WeightedRandomSampler

features = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
labels = torch.tensor([0, 0, 0, 1])
counts = torch.bincount(labels)
class_weights = 1.0 / counts.float()
sample_weights = class_weights[labels]
sampler = WeightedRandomSampler(
    sample_weights, num_samples=len(labels), replacement=True
)
loader = DataLoader(
    TensorDataset(features, labels), batch_size=2, sampler=sampler
)

Each majority-class example has weight 1/3; the minority example has weight 1. That gives both classes the same total sampling weight. Does every batch now contain equal numbers of each class? No. The probabilities are balanced in expectation, but individual draws can differ. With replacement, we can see an example more than once and miss another entirely.

Another option is nn.CrossEntropyLoss(weight=class_weights). Don't automatically combine it with weighted sampling: we'd be changing how often we see examples and how much they count. Keep validation and test data representative of the task rather than balancing them just because we changed the training data.

Regularization

We want our model to learn something useful, rather than just memorize the training examples. Regularization gives us a few ways to encourage that. Try a change, then check whether it actually helps on validation data.

Fewer parameters

Try a shallower network or fewer hidden units if your model fits training examples but generalizes poorly. Compare using the same validation procedure.

Weight decay and L2 regularization

Our model fits the training examples well, but validation isn't keeping up. One thing we might try is discouraging very large weights, rather than letting the model rely on them to fit those examples.

An L2 penalty adds a term proportional to the sum of squared weights to the loss, controlled by a hyperparameter $\lambda$ . With ordinary SGD, the resulting update can be expressed as weight decay. For adaptive optimizers, this equivalence does not generally hold: AdamW applies decoupled weight decay rather than adding it to Adam's gradient statistics.

# Uses the model from the training example.
optimizer = torch.optim.AdamW(
    model.parameters(), lr=0.01, weight_decay=0.01
)

How much decay should we use? That's something to check on validation data. Also, watch out for adding a separate L2 penalty when the optimizer already applies weight decay.

Dropout

Suppose our model does well on the training examples but struggles with new ones. This could indicate that it's overfitting: learning the quirks of those examples rather than patterns that carry over to new ones. That's not uncommon when a model has so many parameters available, especially when training data is limited.

What if we asked it to make predictions without always having every activation available? That gives it a reason to learn useful combinations rather than rely too heavily on one activation.

During training, dropout sets activations to zero with some probability $p$ . This introduces variation in the internal representations the network uses to make predictions. The aim is to make the network more robust: it should learn to make useful predictions without relying too heavily on any single activation. This can reduce overfitting, though it does not guarantee better performance.

After training, we disable dropout with model.eval() so predictions use the full network without randomly removing activations. We also disable it during validation, then re-enable it with model.train() when we return to training. torch.no_grad() alone does not disable dropout.

PyTorch rescales the retained activations during training; in evaluation mode, the dropout layer passes its input through unchanged. We don't need to rescale it manually. The probability is a hyperparameter; 0.5 is an example, not a universal default.

from torch import nn

dropout = nn.Dropout(p=0.5)

Unstable learning

Is the loss bouncing around? Batches can contain quite different examples, so some variation is normal. Before changing the learning rate, look at the overall trend, validation results, and whether the data and gradients look as expected.

Small batches can make updates more variable, and dropout introduces randomness during training. Inputs with very different scales or an unfortunate initialization can also make learning harder. Occasionally, the problem is in the code or data rather than a setting: mislabeled examples, unexpected input values, or an incorrect loss calculation can send us looking for the wrong fix.

We don't need a perfectly smooth loss curve. If the overall trend and validation results look good, a little batch-to-batch variation may be nothing to worry about. If learning really is unstable, start with those basic checks before changing several settings at once.

Learning rate scheduling

The learning rate that helped us make progress early in training may give us steps that are too large later on. We might want to reduce it as training progresses, rather than use one fixed rate throughout.

A scheduler changes the learning rate according to a rule. It is not the same as early stopping, and it need not always decrease the rate. This example halves the rate every five training epochs:

# Uses the optimizer from the training example.
scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer, step_size=5, gamma=0.5
)
# In the epoch loop, after that epoch's optimizer updates:
# scheduler.step()

One noisy batch isn't a reason to reach for a scheduler. If you do use one, check when it expects to be called. See PyTorch's scheduling guidance for the call order and other options.

Gradient accumulation

What is the training data, taken together, telling us about how to adjust the model? Different examples can send competing signals: one might favor increasing a parameter, while another favors decreasing it. That doesn't necessarily mean either example is wrong. We want an update that takes their contributions into account, rather than reacting too strongly to one small group's quirks.

As you may recall from previous courses, a larger batch lets us average gradients over a larger sample of the training data before making an update. That can give us a less noisy estimate of the average gradient. It won't make the examples "agree", but it will combine their guidance for a single update.

What if that larger batch doesn't fit in memory? We can process smaller batches and accumulate their gradients before updating. Four batches of eight examples give one update based on 32 examples, without needing to keep all 32 examples' forward-pass activations in memory at once.

This replacement training loop uses the model, criterion, optimizer, and loader from the earlier example. It assumes a mean loss with equal contribution from each example. The final group can be smaller:

from itertools import islice

model.train()
batches = iter(train_loader)
accumulate_for = 4

while True:
    group = list(islice(batches, accumulate_for))
    if not group:
        break
    examples_in_group = sum(len(targets) for _, targets in group)
    optimizer.zero_grad()
    for inputs, targets in group:
        loss = criterion(model(inputs), targets)
        scaled_loss = loss * len(targets) / examples_in_group
        scaled_loss.backward()
    optimizer.step()

We scale each batch's mean loss by its share of the group's examples. A final group of two equal batches divides each loss by two, not four. Keeping detached numbers such as loss.item() for logging is separate from accumulating gradients.

Layer normalization

Imagine we're passing values through several hidden layers, and their scale varies enough to make learning difficult. One option is to normalize the values passed on to the next part of the network.

LayerNorm normalizes activations over the specified dimensions, not the layer's weights. It has learned scale and bias parameters by default.

from torch import nn

normalized_layer = nn.Sequential(
    nn.Linear(4, 8),
    nn.LayerNorm(8),
    nn.ReLU(),
    nn.Linear(8, 2),  # two class logits
)

Vanishing and exploding gradients

Gradients tell us how parameter changes affect the loss. Very small gradients in early layers can make those layers learn slowly^†, and very large gradients can make updates unstable. Inspect gradients after backward() and before clearing them.

Why might gradients shrink in earlier layers?↩

Think back to our computation graph. A change in an early parameter reaches the loss through later operations. Working backward along a path, we multiply their local rates of change.

Imagine each of four operations contributes a factor of 0.5. All of that halving gives us a product of only 0.0625! A longer path can make that product smaller still. A sigmoid near either end of its output range can contribute a particularly small factor.

The deeper the network, the more opportunities there can be for small factors to multiply along a path, increasing this risk. Several tricks have been developed to help, including careful initialization, different activation functions, and residual connections that give gradients shorter paths to follow.

The factors aren't always small, and contributions from different paths add together. It's one reason to compare gradients in early and late layers rather than assume they're equally large.

If learning is slow, have the gradients vanished? Not necessarily. Look at their magnitudes across layers and steps; the weights themselves being zero doesn't answer that question. Saturated sigmoid activations can make gradients small. Depending on what we find, ReLU, careful initialization, a shallower network, or residual connections may help. ReLU has its own limitations. Namely, it has zero derivative for negative inputs, so a unit that stays inactive can stop receiving useful gradients.

One activation function to rule them all?

If only that were true! Fortunately, we don't need to use the same activation everywhere (and for good reason). ReLU is a useful option for hidden layers: positive inputs aren't squeezed into a bounded range, though negative inputs get zeroed out. Other choices have their own tradeoffs.

At the output, think about the task. A sigmoid can turn one score into a probability for a binary classification task. When we're choosing one of several classes, softmax can turn their scores into probabilities that sum to one.

Should we always add those to the model? No. PyTorch's BCEWithLogitsLoss includes sigmoid, and CrossEntropyLoss works directly with logits. Adding the activation first would give these losses the wrong kind of input. Our earlier two-class example therefore ends with a Linear layer, not softmax.

NaN means not a number. This is different from a value of infinity. If either shows up, or the loss starts making large jumps, let's investigate. Exploding gradients are one possibility; invalid operations or problems in the data can also cause trouble.

Gradient clipping

We've checked the data and noticed occasional very large gradients. Before those gradients produce another large parameter update, can we limit their magnitude?

Norm clipping rescales the combined parameter gradients when their norm exceeds a threshold. It limits their magnitude while preserving direction. It does not directly clip weights or repair every cause of instability.

$g' = \begin{cases} \gamma g / \lVert g \rVert_2 & \text{if } \lVert g \rVert_2 > \gamma \\ g & \text{otherwise} \end{cases}$

import torch

# Inside a training step, using the earlier model and optimizer:
optimizer.zero_grad()
loss = criterion(model(inputs), targets)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

Not enough data

Data augmentation, transfer learning, and sharing information across related tasks can help when labeled data is limited, but they come with assumptions. Augmented examples should retain the intended labels, and knowledge from another task may not always transfer well. We'll return to these approaches later.

Pretrained embeddings

We only have a small labeled dataset, but our model needs a representation for many different words. Fortunately, we don't have to learn every embedding from the few examples that may constitute our primary task dataset. We can start with vectors learned from a larger collection of text.

This connects to word2vec and the embedding lookup we used earlier. A pretrained matrix gives us starting representations rather than requiring us to learn them all from scratch:

import torch
from torch import nn

# A tiny illustrative matrix, not actual pretrained word vectors.
vectors = torch.tensor([[0.0, 0.0], [0.2, 0.8], [0.7, 0.1]])
frozen = nn.Embedding.from_pretrained(vectors.clone(), freeze=True)
trainable = nn.Embedding.from_pretrained(vectors.clone(), freeze=False)
ids = torch.tensor([1, 2], dtype=torch.long)
print(frozen(ids))
print(trainable.weight.requires_grad)

Both versions perform the same lookup initially. Freezing prevents gradient-based updates to this embedding matrix; it does not prevent other layers from learning. With freeze=False, the matrix can adapt to the task. The row-to-token mapping must match your vocabulary.

Tokenization and vocabulary size

Smaller vocabularies reduce the number of embedding parameters when embedding dimension is fixed. Character or subword tokenization can share representations across words, but often creates longer sequences. This trades one cost for another. We'll explore tokenization in the LLM unit.

Sometimes, though, we want to expand the vocabulary. Suppose we want to mark the start of a new dialogue turn with a control token^† such as <TURN>, or give a programming language operator such as := its own representation. We may want the tokenizer to keep that symbol as one unit, rather than split it into smaller pieces. This is similar to giving the boundary symbols ^ and $ their own IDs in our text CNN example.

A dedicated token needs an entry in the vocabulary and a corresponding embedding row. Adding it doesn't automatically teach the model what it means: we still need to initialize that representation and give it useful training examples. If we're extending pretrained embeddings, keep the existing token IDs aligned with their rows. The goal isn't always the smallest vocabulary; it's a vocabulary that works for the task.

Other useful habits

Shuffle ordinary training batches each epoch. When using a sampler, let that sampler determine ordering.
Record seeds, package versions, splits, settings, and metrics. A fixed seed alone does not guarantee identical results across platforms or releases. See PyTorch's reproducibility notes.
Residual connections give signals another path through a network. Addition, however, requires compatible shapes. There's no real silver bullet to DL (other than maybe data improvements).
Start with a simple baseline. An optimizer, a larger model, or another trick should earn its place through a useful comparison.

Don't sweat it

That's a lot of options. You don't need to adjust every setting, and a model that's working well doesn't need another trick just because we've listed one here.

If you're wondering where to start, check your data and labels, get a small baseline working, and use validation to see where it struggles. Then try one change that matches the problem you've actually observed. Keep a note of what helped. If a change doesn't help, there's no obligation to keep it (in fact, it's probably best to KonMari and declutter).

These tips can improve training and sometimes performance, but they aren't a replacement for better data, a well-defined task, or a sensible evaluation setup. More reliable labels or more useful examples may help more than another round of tuning. We want to understand what our model needs, not collect every possible setting.

References

Goldberg, Yoav (2017). Neural Network Methods for Natural Language Processing. Morgan & Claypool Publishers
Ruder, Sebastian (2017). Deep Learning for NLP Best Practices. https://www.ruder.io/deep-learning-nlp-best-practices/