Train, Diagnose, and Improve a Small Model

The supervised training loop: mathematics and mechanics

Before training, print the input shape, target shape, data type, value range, and device (CPU or GPU). Pass one batch through the model and inspect the output shape. Confirm that the loss function accepts those outputs and targets. This catches silent mistakes earlier than a long run.

A supervised training loop consists of these steps, applied repeatedly:

The mathematical formulation: For a batch of n samples with inputs X and true labels y:

1. Forward pass: Compute predictions ŷ = model(X; θ), where θ are the current parameters (weights and biases)
2. Loss computation: Calculate scalar loss L = loss_fn(ŷ, y), e.g., L = (1/n) · Σ(y_i - ŷ_i)² for regression (MSE), or L = -(1/n) · Σ[y_i · log(ŷ_i) + (1-y_i) · log(1-ŷ_i)] for binary classification (cross-entropy)
3. Backpropagation: Compute gradients ∂L/∂θ by the chain rule through all layers
4. Parameter update: θ_new = θ_old - learning_rate · ∂L/∂θ
5. Validation: Evaluate on a held-out batch without updating parameters

Key invariant: After each parameter update, loss should generally fall (though with noise). If loss is flat or rising, the learning rate may be too high, the model capacity may be insufficient, or there may be a data/label quality issue.

The minimal code structure ensures this order is never violated:

Example: minimal PyTorch training loop

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt

# Synthetic dataset for demonstration
n_samples = 1000
n_features = 20
X = torch.randn(n_samples, n_features)
y = (X[:, 0] + X[:, 1] > 0).long()  # Simple binary rule

# Split into train and validation
train_size = int(0.8 * n_samples)
X_train, y_train = X[:train_size], y[:train_size]
X_val, y_val = X[train_size:], y[train_size:]

# Data loaders with batch size 32
train_loader = DataLoader(
    TensorDataset(X_train, y_train),
    batch_size=32,
    shuffle=True
)
val_loader = DataLoader(
    TensorDataset(X_val, y_val),
    batch_size=32,
    shuffle=False
)

# Define a small model
class SimpleClassifier(nn.Module):
    def __init__(self, n_features):
        super().__init__()
        self.fc1 = nn.Linear(n_features, 64)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(64, 2)

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleClassifier(n_features)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Define loss and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Storage for metrics
train_losses = []
val_losses = []
val_accuracies = []

# Training loop
n_epochs = 50
for epoch in range(n_epochs):
    # Training phase
    model.train()
    epoch_train_loss = 0.0
    n_batches = 0

    for X_batch, y_batch in train_loader:
        X_batch, y_batch = X_batch.to(device), y_batch.to(device)

        # 1. Forward pass
        logits = model(X_batch)

        # 2. Compute loss
        loss = loss_fn(logits, y_batch)

        # 3. Zero gradients (clear old gradients)
        optimizer.zero_grad()

        # 4. Backpropagation
        loss.backward()

        # 5. Update parameters
        optimizer.step()

        epoch_train_loss += loss.item()
        n_batches += 1

    avg_train_loss = epoch_train_loss / n_batches
    train_losses.append(avg_train_loss)

    # Validation phase (no parameter updates)
    model.eval()
    with torch.no_grad():
        epoch_val_loss = 0.0
        n_correct = 0
        n_val_samples = 0

        for X_batch, y_batch in val_loader:
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)

            logits = model(X_batch)
            loss = loss_fn(logits, y_batch)
            epoch_val_loss += loss.item()

            # Compute accuracy
            preds = logits.argmax(dim=1)
            n_correct += (preds == y_batch).sum().item()
            n_val_samples += y_batch.size(0)

        avg_val_loss = epoch_val_loss / len(val_loader)
        val_accuracy = n_correct / n_val_samples
        val_losses.append(avg_val_loss)
        val_accuracies.append(val_accuracy)

    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1:2d}: train_loss={avg_train_loss:.4f}, "
              f"val_loss={avg_val_loss:.4f}, val_acc={val_accuracy:.3f}")

print("Training complete.")

Read curves as evidence

Training loss falling while validation loss rises is the classic overfitting signal: the model is memorizing the training set instead of learning generalizable patterns. You may need to add regularization, early stopping, or more training data.

Both curves staying flat (not falling at all) suggests a different problem:

• Learning rate too low: parameter updates are tiny; increase it.
• Broken labels: the target variable has errors or is random.
• Saturated activations: ReLU networks with bad initialization can have dead neurons that never fire.
• Model capacity too low: a single neuron cannot fit a nonlinear boundary.

A noisy metric (validation loss bounces wildly) often simply reflects too little validation data. If your validation set is 100 samples and you run it every batch, small changes in which examples are easy/hard cause wild swings. Validate less frequently or use a larger validation set.

Example: loss curve analysis

import matplotlib.pyplot as plt

# Plot training and validation loss curves
plt.figure(figsize=(10, 5))
epochs = range(1, len(train_losses) + 1)

plt.plot(epochs, train_losses, label="Training loss", marker='o', linestyle='-', markersize=3)
plt.plot(epochs, val_losses, label="Validation loss", marker='s', linestyle='--', markersize=3)

plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Training vs Validation Loss Curves")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("loss_curves.png", dpi=150)
plt.show()

# Diagnostic: did validation loss improve in the last 20 epochs?
last_20_val_losses = val_losses[-20:]
if min(last_20_val_losses) == last_20_val_losses[-1]:
    print("✓ Validation loss still improving; training is on track.")
else:
    best_epoch = len(val_losses) - 20 + last_20_val_losses.index(min(last_20_val_losses))
    print(f"⚠ Validation loss plateaued at epoch {best_epoch}; consider early stopping.")

# Overfitting check
gap = val_losses[-1] - train_losses[-1]
if gap > 0.1:
    print(f"⚠ Large gap ({gap:.3f}) between train and validation loss suggests overfitting.")
else:
    print(f"✓ Train/validation gap ({gap:.3f}) is small; model generalizes reasonably.")

Change one factor per experiment. If you increase learning rate AND add regularization at the same time, you will not know which change helped. Record:

• Hypothesis (e.g., "learning rate too low")
• Code and data versions (git hash, dataset date)
• Parameter change (e.g., lr: 1e-3 → 1e-2)
• Result (validation loss after 50 epochs)
• Interpretation (did it confirm your hypothesis?)

Diagnostic table: reading training curves

Use this table to diagnose what went wrong based on the shape of your loss curves:

| Pattern | Train Loss | Val Loss | Diagnosis | Next Step | |---|---|---|---|---| | Both falling, then flatten | ✓ Decreases | ✓ Decreases → plateau | Model is learning; may be limited by capacity or data quality | Try adding features, increasing model size, or collecting more data | | Train falls, val rises | ✓ Decreases | ✗ Increases | Overfitting: model is memorizing training data | Add L2 regularization, dropout, early stopping, or simplify the model | | Both flat, never improve | ✗ High, flat | ✗ High, flat | Underfitting or broken training; learning signal is missing | Check: learning rate (try 10x higher), label quality, feature relevance | | Val noisy, train smooth | ✓ Decreasing | ~ Bounces wildly | Validation set too small; random variation dominates | Use a larger validation set or validate less frequently | | Both decrease, then train rises | ✓ Then ✗ | ✓ Continues falling | Learning rate too high; model is oscillating around the optimum | Reduce learning rate by 2-5x |

Illustrative estimates (typical ranges from benchmark training runs):

• Well-behaved training: validation loss improves for 20-50% of epochs, then plateaus
• Overfitting onset: train/val gap starts widening after ~50% of training
• Convergence: both losses change by <1% per epoch, suggesting diminishing returns

Inspect examples, not just averages

A single validation loss or accuracy number hides important structure. Create slices of your validation set by:

• Class: Is the model weaker on one class than another?
• Data source: Does performance degrade for a particular data source or domain?
• Length: For text or sequences, is the model weaker on short vs long examples?
• Time period: For temporal data, is the model consistent across different time windows?
• User group: For multi-tenant systems, does the model serve some users better than others?

Review confident errors (cases where the model was confident but wrong), borderline cases (where the model was near 0.5 probability), and disagreements with the baseline (cases the baseline got right but your model got wrong). Look for patterns: Do borderline cases have missing values? Do confident errors have unusual feature combinations?

Check for labeling errors before increasing model capacity. A mislabeled example that your model fits well is a signal that you are overfitting, not that the model is learning something useful. If you find systematic labeling errors, fix them and retrain.

Example: error analysis and slicing

import pandas as pd
import numpy as np

# Assume we have validation predictions and true labels
y_val_pred_proba = model.forward(X_val.to(device)).softmax(dim=1)[:, 1].detach().cpu().numpy()
y_val_pred = (y_val_pred_proba >= 0.5).astype(int)

# Create an error analysis DataFrame
val_analysis = pd.DataFrame({
    "true_label": y_val.cpu().numpy(),
    "pred_label": y_val_pred,
    "pred_prob": y_val_pred_proba,
    "error": (y_val_pred != y_val.cpu().numpy()).astype(int),
    "confidence": np.abs(y_val_pred_proba - 0.5),  # Distance from decision boundary
})

print("=" * 60)
print("ERROR ANALYSIS")
print("=" * 60)

# Confident errors: high confidence but wrong
confident_errors = val_analysis[(val_analysis["error"] == 1) & (val_analysis["confidence"] > 0.4)]
print(f"\nConfident errors (n={len(confident_errors)}):")
print(confident_errors.head(10)[["true_label", "pred_prob", "confidence"]])

# Borderline cases: low confidence
borderline = val_analysis[val_analysis["confidence"] < 0.1]
print(f"\nBorderline cases (n={len(borderline)}):")
print(f"  Error rate: {borderline['error'].mean():.3f}")
print(f"  This suggests the model is uncertain about genuinely ambiguous examples.")

# Per-class accuracy
print(f"\nPer-class accuracy:")
for label in [0, 1]:
    acc = val_analysis[val_analysis["true_label"] == label]["error"].apply(lambda x: 1 - x).mean()
    count = (val_analysis["true_label"] == label).sum()
    print(f"  Class {label}: {acc:.3f} ({count} samples)")

# Top 10 misclassified examples by confidence (investigate these first)
top_errors = val_analysis[val_analysis["error"] == 1].nlargest(10, "confidence")
print(f"\nTop 10 misclassified (by model confidence):")
print(top_errors[["true_label", "pred_prob", "confidence"]])

Practice: two-experiment review

1. First run: Train a small model (e.g., logistic regression or a 2-layer neural network) for a fixed number of epochs (e.g., 50). Save the best validation checkpoint.
2. Produce evidence:
   • Loss curves (training and validation)
   • Confusion matrix on validation set
   • Ten concrete examples of misclassified cases
3. Analyze: Based on the curves and misclassified examples, form a single hypothesis (e.g., "model is underfitting because both losses are high and flat" or "model is overfitting; validation loss rose after epoch 20").
4. Second run: Make exactly one change aligned with your hypothesis:
   • If underfitting: increase model capacity, reduce regularization, or increase learning rate.
   • If overfitting: add L2 regularization, drop layers, or use early stopping.
   • If noisy validation: use a larger validation set or validate less frequently.
5. Rerun and write: Did the hypothesis hold? Quote the change in validation loss or error rate as evidence.

Example template

# First experiment: baseline training
print("=" * 60)
print("EXPERIMENT 1: Baseline Model")
print("=" * 60)

model1 = SimpleClassifier(n_features)
model1.to(device)
optimizer1 = torch.optim.Adam(model1.parameters(), lr=1e-3)

# [Train model1 for 50 epochs as above]
# [Evaluate and save best checkpoint]

# Second experiment: based on observed behavior
print("\n" + "=" * 60)
print("EXPERIMENT 2: Increased Learning Rate (hypothesis: underfitting)")
print("=" * 60)

model2 = SimpleClassifier(n_features)
model2.to(device)
optimizer2 = torch.optim.Adam(model2.parameters(), lr=5e-3)  # 5x higher

# [Train model2 with the same loop]
# [Evaluate and compare]

# Comparison
print(f"\nExperiment 1 (baseline, lr=1e-3):")
print(f"  Val loss: {val_losses[-1]:.4f}, Val acc: {val_accuracies[-1]:.3f}")

# Assume you computed model2_val_loss and model2_val_acc similarly
print(f"\nExperiment 2 (lr=5e-3):")
print(f"  Val loss: {model2_val_loss:.4f}, Val acc: {model2_val_acc:.3f}")

print(f"\nConclusion: Higher learning rate {'improved' if model2_val_loss < val_losses[-1] else 'degraded'} performance.")

Stop condition

Stop when additional complexity does not improve the decision-relevant metric enough to justify its latency (slower inference = worse user experience), maintenance burden (more code = more bugs and technical debt), interpretability (deeper models are harder to explain), or data cost (collecting more labeled data is expensive).

If a logistic regression achieves 84% accuracy in 2 ms, and a deep neural network achieves 86% in 500 ms, the 2 percentage point gain might not be worth the 250x latency increase — especially for a real-time system. If a simple model is interpretable enough for your stakeholders and a complex model is not, simplicity wins.

Defer scaling to large models until you have answered:

• Why does the small model fail on specific examples?
• Which examples would larger capacity actually help with?
• Is the next improvement worth the cost in latency, maintainability, or data?

A small, interpretable, well-understood model in production beats a large, opaque model on your laptop.

Learning rate schedules and optimization dynamics

The learning rate controls the step size of parameter updates. Too high, and the model oscillates wildly or diverges. Too low, and training is glacially slow.

A common pattern is learning rate scheduling: start with a moderate learning rate, then decay it as training progresses. Early epochs take larger steps (explore the loss landscape), later epochs take smaller steps (fine-tune near a local minimum).

# Learning rate scheduler example with PyTorch
import torch.optim.lr_scheduler as lr_scheduler

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
# Halve learning rate every 20 epochs
scheduler = lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)

for epoch in range(100):
    epoch_train_loss = 0.0
    for X_batch, y_batch in train_loader:
        X_batch, y_batch = X_batch.to(device), y_batch.to(device)
        logits = model(X_batch)
        loss = loss_fn(logits, y_batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        epoch_train_loss += loss.item()

    # Decay learning rate after each epoch
    scheduler.step()
    current_lr = optimizer.param_groups[0]['lr']
    print(f"Epoch {epoch+1:3d}: loss={epoch_train_loss/len(train_loader):.4f}, lr={current_lr:.6f}")

Illustrative scheduling impact (typical on a benchmark model):

• Epoch 1–20: lr = 1e-3 → loss goes from 2.5 to 0.8
• Epoch 20–40: lr = 5e-4 → loss goes from 0.8 to 0.3
• Epoch 40–60: lr = 2.5e-4 → loss goes from 0.3 to 0.25 (diminishing returns)

This typically accelerates convergence and can improve final validation accuracy by 1–3 points compared to a fixed learning rate.

Common mistakes in training loops

| Mistake | Consequence | How to Catch | Fix | |---------|-------------|-------------|-----| | Forgetting to zero gradients | Gradients accumulate; training is chaotic | Loss is noisy, doesn't decrease systematically | Add optimizer.zero_grad() before loss.backward() | | Evaluating on training data only | Overfitting looks fine; surprises on test data | Validation accuracy is much lower than training | Always keep a separate validation set; validate every N epochs | | Forgetting model.eval() during validation | Dropout and batch norm behave incorrectly | Validation loss is artificially low | Call model.eval() before validation; use with torch.no_grad() | | Using wrong loss function | Model optimizes wrong objective | Training loss is high; accuracy is poor | Match loss to task: Cross-Entropy for classification, MSE for regression | | Not shuffling training data | Model memorizes data order, overfits to batch patterns | Loss plateaus early; sudden jump in validation loss | Set shuffle=True in DataLoader | | Testing on training data | Metrics are optimistically biased | Test accuracy >95% on known-difficult tasks | Always use a held-out test set, never touched during training |

Advanced diagnostic: gradient flow and dead activations

In deep networks, vanishing gradients (gradients become too small) prevent learning in early layers. Dead activations (neurons output zero) also block learning.

Check gradient flow by inspecting layer-by-layer gradient norms:

# Inspect gradient magnitudes per layer (diagnostic for vanishing/exploding gradients)
def print_gradients(model, name="Model"):
    print(f"Gradient norms for {name}:")
    for n, p in model.named_parameters():
        if p.grad is not None:
            print(f"  {n:30s}: grad norm = {p.grad.norm():.6f}")

Healthy gradient norms: Should decrease gradually from later to earlier layers but not drop below 1e-6 (suggests vanishing). Should not exceed 1e-2 (suggests exploding — use gradient clipping).

Dead ReLU units: ReLU(x) = max(0, x) is zero for negative inputs. If a ReLU neuron receives consistently negative pre-activation, it outputs zero forever and cannot learn. Count dead units:

# Count dead ReLU neurons (post-activation outputs are zero)
activations = model(X_batch)  # Some layer output
dead_count = (activations == 0).sum().item()
total_count = activations.numel()
print(f"Dead ReLU units: {dead_count} / {total_count} ({dead_count/total_count*100:.1f}%)")

If >20% of units are dead, the network is undershooting. Try: lower learning rate, batch normalization, or use Leaky ReLU instead of ReLU.