23  Chapter 22. Training Vision Models

Empirical Risk, Backpropagation, Optimization, Regularization, and Evaluation

23.1 22.1 Why training deserves its own mathematical chapter

A modern vision model is not only an architecture. It is the result of an interaction among a dataset, a loss function, an optimization algorithm, a regularization strategy, and an evaluation protocol. Two models with the same architecture can behave very differently if they are trained with different sampling schemes, augmentation distributions, learning-rate schedules, or validation criteria.

For MA Applied Mathematics and MS Statistics students, training should be viewed as both an optimization problem and a statistical inference problem. The basic supervised formulation is

\[ \widehat\theta \in \arg\min_{\theta\in\Theta} \widehat R_n(\theta), \qquad \widehat R_n(\theta)=\frac{1}{n}\sum_{i=1}^n \ell\bigl(f_\theta(X_i),Y_i\bigr), \]

where \(X_i\) is an image or visual input, \(Y_i\) is a label or structured target, \(f_\theta\) is a parameterized model, and \(\ell\) is a loss. The corresponding population risk is

\[ R(\theta)=\mathbb E_{(X,Y)\sim P}\,\ell\bigl(f_\theta(X),Y\bigr). \]

Training is therefore the attempt to find parameters with low population risk using only a finite, noisy sample and a computationally constrained optimization method.

NoteCentral idea

Training vision models combines three mathematical questions:

1. Optimization: how do we minimize a large nonconvex empirical objective?
2. Statistics: why should low training error imply useful performance on new images?
3. Geometry: how do architecture, augmentation, and loss functions encode invariances of visual data?

Historically, stochastic gradient methods made large-scale neural training possible (Robbins and Monro 1951; Bottou 2010). Modern deep vision training also depends on backpropagation (Rumelhart et al. 1986), adaptive methods such as Adam (Kingma and Ba 2015), normalization (Ioffe and Szegedy 2015), residual connections (He et al. 2015), and regularization methods including dropout (Srivastava et al. 2014), label smoothing (Szegedy et al. 2016), mixup (Zhang et al. 2018), CutMix (Yun et al. 2019), and augmentation policies (Cubuk et al. 2019, 2020).

23.1.1 Learning goals

After this chapter, you should be able to:

1. formulate supervised vision training as empirical risk minimization;
2. distinguish training risk, validation risk, test risk, and population risk;
3. compute gradients for a simple neural model using the chain rule;
4. explain backpropagation as reverse-mode differentiation on a computational graph;
5. compare gradient descent, stochastic gradient descent, momentum, RMSProp, Adam, and AdamW;
6. explain learning-rate schedules, warmup, early stopping, and checkpoint selection;
7. describe initialization and normalization from a variance-propagation viewpoint;
8. connect weight decay, dropout, data augmentation, label smoothing, mixup, and CutMix to statistical regularization;
9. evaluate classification, detection, and segmentation models using statistically meaningful metrics;
10. design an AI-assisted training workflow that includes debugging, ablation studies, and reproducibility checks.

23.2 22.2 The supervised vision-training problem

Let

\[ \mathcal D_n=\{(X_i,Y_i)\}_{i=1}^n \]

denote a training dataset. In computer vision, the input \(X_i\) may be an image,

\[ X_i\in \mathbb R^{H\times W\times C}, \]

or a video, point cloud, multi-view collection, or image-text pair. The target \(Y_i\) depends on the task:

Task	Target space	Example loss
classification	\(\{1,\ldots,K\}\)	cross-entropy
multilabel classification	\(\{0,1\}^K\)	binary cross-entropy
regression	\(\mathbb R^m\)	squared or Huber loss
segmentation	\(\{1,\ldots,K\}^{H\times W}\)	pixelwise cross-entropy, Dice loss
detection	boxes and class labels	classification + box regression
depth	\(\mathbb R_+^{H\times W}\)	scale-invariant or photometric loss
representation learning	pairs/views/tokens	contrastive or reconstruction loss

A vision model is a map

\[ f_\theta: \mathcal X\to \mathcal Z, \]

where \(\mathcal Z\) may be logits, probabilities, masks, feature maps, boxes, or continuous predictions. The final prediction is often obtained by a decision rule

\[ \widehat Y = a(f_\theta(X)). \]

For classification, \(f_\theta(X)=z\in\mathbb R^K\) may be a vector of logits and

\[ \widehat p_k(X)=\frac{\exp(z_k)}{\sum_{j=1}^K\exp(z_j)}, \qquad \widehat Y=\arg\max_k \widehat p_k(X). \]

23.2.1 22.2.1 Training, validation, and test sets

The dataset is usually split into three parts:

\[ \mathcal D_{\text{train}}, \qquad \mathcal D_{\text{val}}, \qquad \mathcal D_{\text{test}}. \]

The training set is used to update parameters. The validation set is used to select hyperparameters and checkpoints. The test set should be used only for final reporting. If the test set is repeatedly used during model design, it becomes part of the training process and no longer estimates true out-of-sample performance.

WarningStatistical warning

For small or imbalanced vision datasets, the validation estimate

\[ \widehat R_{\text{val}}(\theta) \]

can have high variance. A model that appears better by one percentage point may not be truly better unless the validation set is large enough or the comparison is repeated with confidence intervals.

23.2.2 22.2.2 Empirical risk with regularization

A common training objective is

\[ J(\theta)=\frac{1}{n}\sum_{i=1}^n \ell(f_\theta(X_i),Y_i)+\lambda\Omega(\theta), \]

where \(\Omega(\theta)\) is a regularizer. For example,

\[ \Omega(\theta)=\frac{1}{2}\|\theta\|_2^2 \]

leads to \(L^2\) regularization. In Bayesian language, minimizing

\[ -\sum_{i=1}^n \log p_\theta(Y_i\mid X_i)-\log p(\theta) \]

is maximum a posteriori estimation. A Gaussian prior on weights produces an \(L^2\) penalty.

23.3 22.3 Loss functions for vision

The loss function translates the task into a differentiable objective. It also determines the statistical interpretation of training.

23.3.1 22.3.1 Squared loss and Gaussian noise

For regression targets \(Y\in\mathbb R^m\), squared loss is

\[ \ell(f_\theta(X),Y)=\frac{1}{2}\|f_\theta(X)-Y\|_2^2. \]

If

\[ Y=f_\theta(X)+\varepsilon, \qquad \varepsilon\sim N(0,\sigma^2I), \]

then minimizing squared loss is equivalent to maximizing a Gaussian likelihood.

23.3.2 22.3.2 Cross-entropy and categorical likelihood

For \(K\)-class classification, let \(Y\in\{1,\ldots,K\}\) and let \(p_\theta(k\mid X)\) be the predicted probability. The cross-entropy loss is

\[ \ell(f_\theta(X),Y)=-\log p_\theta(Y\mid X). \]

If \(y\in\{0,1\}^K\) is a one-hot vector, then

\[ \ell(z,y)=-\sum_{k=1}^K y_k\log \operatorname{softmax}(z)_k. \]

For logits \(z\), probabilities \(p=\operatorname{softmax}(z)\), and one-hot target \(y\), the gradient has the simple form

\[ \frac{\partial \ell}{\partial z_k}=p_k-y_k. \]

This identity is one reason softmax cross-entropy is numerically convenient.

23.3.3 22.3.3 Pixelwise losses

For semantic segmentation, the target is a label field

\[ Y\in \{1,\ldots,K\}^{H\times W}. \]

The pixelwise cross-entropy loss is

\[ \ell(\theta)= -\frac{1}{HW}\sum_{u=1}^H\sum_{v=1}^W \log p_\theta(Y[u,v]\mid X)[u,v]. \]

Class imbalance is common in segmentation. A background class may dominate the image, causing the model to obtain low loss while missing small but important objects. Weighted cross-entropy uses

\[ \ell(\theta)= -\frac{1}{HW}\sum_{u,v} w_{Y[u,v]} \log p_\theta(Y[u,v]\mid X)[u,v]. \]

23.3.4 22.3.4 Dice loss and overlap losses

For a binary mask, the soft Dice coefficient is

\[ \operatorname{Dice}(p,y)= \frac{2\sum_j p_jy_j+\epsilon}{\sum_j p_j+\sum_j y_j+\epsilon}. \]

A common loss is

\[ \ell_{\text{Dice}}(p,y)=1-\operatorname{Dice}(p,y). \]

Dice loss is useful when positive pixels are rare. Unlike pixelwise cross-entropy, it directly measures overlap.

23.3.5 22.3.5 Detection losses

Object detection usually combines classification and localization:

\[ \ell=\ell_{\text{cls}}+\alpha\ell_{\text{box}}+\beta\ell_{\text{obj}}. \]

For a predicted box \(b\) and target box \(b^*\), the box loss may use smooth \(L^1\) loss,

\[ \operatorname{smooth}_{L^1}(r)= \begin{cases} \frac{1}{2}r^2, & |r|<1,\\ |r|-\frac{1}{2}, & |r|\ge 1, \end{cases} \]

or an IoU-based loss. Detection training is more complicated than classification because the model must learn both what objects are present and where they are.

23.3.6 22.3.6 Contrastive and self-supervised losses

In self-supervised vision, the target may be generated from the data itself. For example, two augmented views \(X_i'\) and \(X_i''\) of the same image should have nearby embeddings. InfoNCE uses similarities among embeddings:

\[ \ell_i=-\log \frac{\exp(s(z_i,z_i^+)/\tau)} {\exp(s(z_i,z_i^+)/\tau)+\sum_{j\ne i}\exp(s(z_i,z_j^-)/\tau)}. \]

This objective is not merely a heuristic. It encodes a statistical assumption: augmentations should preserve semantically relevant information while removing nuisance variability.

23.4 22.4 Backpropagation as the chain rule on graphs

A neural network is a composition of maps:

\[ f_\theta = f_L\circ f_{L-1}\circ\cdots\circ f_1. \]

Let

\[ h_0=X, \qquad h_\ell=f_\ell(h_{\ell-1};\theta_\ell), \qquad \ell=1,\ldots,L. \]

The loss is

\[ \mathcal L(\theta)=\ell(h_L,Y). \]

The gradient with respect to layer \(\ell\) is obtained by the chain rule:

\[ \frac{\partial \mathcal L}{\partial \theta_\ell} = \frac{\partial \mathcal L}{\partial h_\ell} \frac{\partial h_\ell}{\partial \theta_\ell}. \]

Backpropagation computes these gradients efficiently by propagating adjoints backward:

\[ \delta_L=\frac{\partial \mathcal L}{\partial h_L}, \qquad \delta_{\ell-1} = \delta_\ell \frac{\partial h_\ell}{\partial h_{\ell-1}}. \]

In matrix notation, if

\[ h_\ell=\sigma(W_\ell h_{\ell-1}+b_\ell), \]

then

\[ \delta_{\ell-1}=W_\ell^T\bigl(\delta_\ell\odot \sigma'(a_\ell)\bigr), \qquad \frac{\partial \mathcal L}{\partial W_\ell} = \bigl(\delta_\ell\odot \sigma'(a_\ell)\bigr)h_{\ell-1}^T. \]

For CNNs, the same chain rule holds, but the linear maps are convolutional operators. The weight gradient is a correlation between input activations and output adjoints. The input gradient is a convolution with a flipped kernel, depending on the convention used.

NoteReverse-mode efficiency

If a network has millions of parameters and one scalar loss, reverse-mode automatic differentiation computes all parameter gradients at a cost comparable to a small multiple of the forward pass. This is why backpropagation is essential for deep learning.

23.4.1 22.4.1 Computational graph viewpoint

A computational graph represents the loss as a composition of elementary operations. Each node stores a local derivative. Backpropagation applies the identity

\[ \frac{\partial \mathcal L}{\partial u} = \sum_{v\in \operatorname{children}(u)} \frac{\partial \mathcal L}{\partial v} \frac{\partial v}{\partial u}. \]

This is a dynamic-programming algorithm on the graph.

23.4.2 Interactive demo: backpropagation through a two-layer network

The following demo shows how gradient magnitude changes when activation saturation increases.

23.5 22.5 Gradient descent and stochastic gradient methods

23.5.1 22.5.1 Full-batch gradient descent

Given objective \(J(\theta)\), full-batch gradient descent updates

\[ \theta_{t+1}=\theta_t-\eta_t\nabla J(\theta_t), \]

where \(\eta_t>0\) is the learning rate. For convex smooth objectives, a small enough learning rate guarantees descent. For deep vision models, the objective is nonconvex, but gradient descent remains the basic local search principle.

23.5.2 22.5.2 Mini-batch stochastic gradient descent

Deep vision datasets are large. Instead of computing the full gradient, we sample a mini-batch \(B_t\subset\{1,\ldots,n\}\) and use

\[ g_t=\frac{1}{|B_t|}\sum_{i\in B_t}\nabla_\theta \ell(f_{\theta_t}(X_i),Y_i). \]

The SGD update is

\[ \theta_{t+1}=\theta_t-\eta_t g_t. \]

If the mini-batch is sampled uniformly, then

\[ \mathbb E[g_t\mid \theta_t]=\nabla \widehat R_n(\theta_t). \]

Thus SGD uses an unbiased but noisy gradient estimate.

The covariance of the gradient estimator decreases with batch size:

\[ \operatorname{Cov}(g_t)\approx \frac{1}{B}\Sigma_g. \]

Large batches reduce gradient noise but may require learning-rate tuning and can generalize differently.

23.5.3 22.5.3 Momentum

Momentum introduces a velocity variable:

\[ v_{t+1}=\mu v_t+g_t, \qquad \theta_{t+1}=\theta_t-\eta v_{t+1}. \]

Momentum averages gradients over time. It accelerates motion along consistent directions and damps oscillations in directions with alternating signs.

23.5.4 22.5.4 RMSProp, Adam, and AdamW

Adaptive methods rescale coordinates using estimates of first and second moments. Adam uses

\[ m_t=\beta_1m_{t-1}+(1-\beta_1)g_t, \]

\[ v_t=\beta_2v_{t-1}+(1-\beta_2)g_t\odot g_t, \]

with bias-corrected estimates

\[ \widehat m_t=\frac{m_t}{1-\beta_1^t}, \qquad \widehat v_t=\frac{v_t}{1-\beta_2^t}. \]

The update is

\[ \theta_{t+1}=\theta_t- \eta\frac{\widehat m_t}{\sqrt{\widehat v_t}+\epsilon}. \]

AdamW decouples weight decay from the adaptive gradient update. This distinction matters because adding \(\lambda\theta\) to the gradient is not equivalent to shrinking weights when the learning rate is coordinate-adaptive.

23.5.5 Interactive demo: optimizer paths on a curved loss

23.6 22.6 Learning-rate schedules and training dynamics

The learning rate is one of the most important hyperparameters. If it is too large, training may diverge. If it is too small, optimization may be extremely slow or trapped in poor regions.

Common schedules include:

1. step decay: reduce \(\eta\) by a factor at specified epochs;
2. exponential decay: \(\eta_t=\eta_0\gamma^t\);
3. cosine decay: \[ \eta_t=\eta_{\min}+\frac{1}{2}(\eta_{\max}-\eta_{\min}) \left(1+\cos\frac{\pi t}{T}\right); \]
4. warmup: start with a small learning rate and increase gradually;
5. cyclical schedules: periodically vary the learning rate.

Warmup is especially useful for large-batch training, transformers, and models with normalization layers.

23.6.1 Interactive demo: learning-rate schedules

23.7 22.7 Initialization and variance propagation

Deep networks are compositions. If activations or gradients grow or shrink exponentially with depth, training becomes unstable.

Consider a layer

\[ h_j=\sum_{i=1}^{d} W_{ji}x_i. \]

Suppose \(W_{ji}\) are independent with mean zero and variance \(\sigma_W^2\), and \(x_i\) have variance \(\sigma_x^2\). Then approximately

\[ \operatorname{Var}(h_j)=d\sigma_W^2\sigma_x^2. \]

To preserve variance across layers, choose

\[ \sigma_W^2\approx \frac{1}{d}. \]

For ReLU activations, about half the activations are zero, leading to the He initialization scale

\[ \sigma_W^2\approx \frac{2}{d}. \]

For tanh-like activations, Xavier initialization uses a scale related to both fan-in and fan-out:

\[ \sigma_W^2\approx \frac{2}{d_{\text{in}}+d_{\text{out}}}. \]

23.7.1 22.7.1 Batch normalization

Batch normalization transforms preactivations by

\[ \widehat h=\frac{h-\mu_B}{\sqrt{\sigma_B^2+\epsilon}}, \qquad \operatorname{BN}(h)=\gamma\widehat h+\beta. \]

Here \(\mu_B\) and \(\sigma_B^2\) are mini-batch statistics. Batch normalization can stabilize optimization and allow larger learning rates, although its full theoretical explanation is subtle.

23.7.2 22.7.2 Layer normalization and group normalization

Layer normalization computes statistics across features within an example rather than across examples. Group normalization divides channels into groups. These alternatives are useful when batch sizes are small, as often happens in high-resolution detection and segmentation.

23.8 22.8 Regularization for vision models

Regularization means controlling the model so that it generalizes beyond the training set.

23.8.1 22.8.1 Weight decay

With weight decay, the objective becomes

\[ J(\theta)=\widehat R_n(\theta)+\frac{\lambda}{2}\|\theta\|_2^2. \]

For ordinary gradient descent, this gives

\[ \theta_{t+1}=(1-\eta\lambda)\theta_t-\eta\nabla \widehat R_n(\theta_t). \]

Thus weight decay shrinks parameters toward zero.

23.8.2 22.8.2 Early stopping

Early stopping chooses a checkpoint before the training loss is minimized. If validation risk decreases at first and then increases, the best checkpoint is

\[ \widehat t=\arg\min_t \widehat R_{\text{val}}(\theta_t). \]

Early stopping acts like regularization because it restricts how far optimization moves from the initial parameters.

23.8.3 22.8.3 Dropout

Dropout multiplies hidden activations by random masks:

\[ \widetilde h=m\odot h, \qquad m_j\sim \operatorname{Bernoulli}(q). \]

During training, the network must perform well under random feature removal. This discourages excessive co-adaptation of hidden units.

23.8.4 22.8.4 Label smoothing

Instead of one-hot labels, label smoothing uses

\[ y_k^{(\epsilon)}=(1-\epsilon)y_k+\frac{\epsilon}{K}. \]

This prevents the model from assigning probability one to the observed class. It can improve calibration and reduce overconfidence.

23.8.5 22.8.5 Mixup

Mixup constructs synthetic examples by convex combination:

\[ \widetilde X=\lambda X_i+(1-\lambda)X_j, \qquad \widetilde y=\lambda y_i+(1-\lambda)y_j, \]

where \(\lambda\in[0,1]\). It encourages approximately linear behavior between training examples.

23.8.6 22.8.6 CutMix

CutMix pastes a patch from one image into another and combines labels according to the area fraction. It is especially natural for images because local regions often correspond to visual evidence.

23.8.7 Interactive demo: regularization and validation error

23.9 22.9 Data augmentation as stochastic invariance

In vision, augmentation is often as important as architecture. Let \(G\) be a set or distribution of transformations such as flips, crops, rotations, color jitter, blur, or noise. If a transformation \(g\in G\) preserves the label, then ideally

\[ f_\theta(gX)=f_\theta(X). \]

Training with augmentation minimizes

\[ \widehat R_{\text{aug}}(\theta)= \frac{1}{n}\sum_{i=1}^n \mathbb E_{g\sim \nu} \ell(f_\theta(gX_i),Y_i), \]

where \(\nu\) is the augmentation distribution.

For segmentation, augmentation must act on both image and mask:

\[ (X,Y)\mapsto (gX,gY). \]

For detection, bounding boxes must be transformed consistently. For contrastive learning, different augmentations of the same image define positive pairs.

23.9.1 22.9.1 Good and bad augmentations

A good augmentation preserves the target. A bad augmentation changes the target or destroys task-relevant information.

Examples:

• Horizontal flips may be safe for generic object recognition but unsafe for reading digits or medical laterality.
• Color jitter may help object recognition but harm pathology grading if color is diagnostic.
• Cropping may help robustness but remove small objects in detection.
• Rotation may be safe for satellite imagery but unsafe for scenes where gravity matters.

WarningAugmentation is a modeling assumption

Every augmentation asserts an invariance. The question is not whether augmentation is modern or popular. The question is whether the invariance is valid for the scientific or engineering task.

23.9.2 Interactive demo: augmentation distribution

23.10 22.10 Transfer learning and fine-tuning

Vision models are often initialized from a model pretrained on a large dataset. Let

\[ f_\theta = h_\psi\circ \phi_\omega, \]

where \(\phi_\omega\) is a feature extractor and \(h_\psi\) is a task head. In transfer learning, one may:

1. freeze \(\omega\) and train only \(\psi\);
2. fine-tune both \(\omega\) and \(\psi\) with a small learning rate;
3. fine-tune only selected layers;
4. adapt the model using parameter-efficient methods.

The statistical idea is that pretraining supplies a representation with lower effective sample complexity for the downstream task. The risk is that the pretraining distribution may not match the target distribution.

23.10.1 22.10.1 Linear probing

A linear probe trains

\[ h_\psi(z)=Wz+b, \qquad z=\phi_\omega(X), \]

while keeping \(\omega\) fixed. If the probe performs well, the representation already separates classes linearly. If the probe performs poorly but fine-tuning performs well, the representation contains useful information but must be adapted.

23.10.2 22.10.2 Fine-tuning and catastrophic forgetting

Fine-tuning modifies pretrained features. With small target datasets, unrestricted fine-tuning may overfit or destroy useful pretrained structure. Common strategies include smaller learning rates for earlier layers, weight decay, early stopping, and strong augmentation.

23.11 22.11 Class imbalance and sampling

Vision datasets often have class imbalance. For example, medical images may contain many normal cases and few disease cases; segmentation images may contain mostly background pixels.

Let class prior probabilities be

\[ \pi_k=P(Y=k). \]

If the dataset prior differs from the deployment prior, accuracy can be misleading. Remedies include:

1. class-weighted losses;
2. balanced mini-batch sampling;
3. focal loss;
4. threshold tuning;
5. reporting per-class metrics.

Focal loss for binary classification can be written as

\[ \ell(p,y)= -\alpha y(1-p)^\gamma\log p -(1-\alpha)(1-y)p^\gamma\log(1-p). \]

The factor \((1-p)^\gamma\) downweights easy positive examples. This is useful in dense detection where background examples are abundant.

23.12 22.12 Calibration and uncertainty during training

A model is calibrated if predicted probabilities match empirical frequencies. For example, among images assigned confidence \(0.8\), approximately \(80\%\) should be correct.

For a classifier, the expected calibration error can be approximated by binning predictions:

\[ \operatorname{ECE}=\sum_{m=1}^M \frac{|B_m|}{n} \left| \operatorname{acc}(B_m)-\operatorname{conf}(B_m) \right|. \]

Training can affect calibration. Overparameterized networks trained with cross-entropy may become overconfident. Label smoothing, temperature scaling, mixup, and ensembles can improve calibration.

23.12.1 Interactive demo: confidence and calibration

23.13 22.13 Model selection, ablation studies, and reproducibility

Training a vision model involves many choices:

\[ \text{architecture},\quad \text{optimizer},\quad \text{learning rate},\quad \text{batch size},\quad \text{augmentation},\quad \text{regularization},\quad \text{loss},\quad \text{data split}. \]

An ablation study changes one component at a time to estimate its contribution. A clean ablation table might report:

Change	Validation metric	Interpretation
baseline	82.1	reference
+ stronger augmentation	84.0	improves generalization
+ label smoothing	84.5	improves calibration and accuracy
+ AdamW	84.2	optimization effect smaller here
+ larger model	85.1	capacity helps, but cost rises

For scientific use, one should report random seeds, data splits, preprocessing, training curves, hardware, and confidence intervals when possible.

23.13.1 22.13.1 Validation uncertainty

If a classifier has validation accuracy \(\widehat p\) on \(m\) independent examples, a simple approximate standard error is

\[ \operatorname{SE}(\widehat p)= \sqrt{\frac{\widehat p(1-\widehat p)}{m}}. \]

A difference of \(0.2\%\) accuracy is not meaningful if the standard error is \(0.5\%\).

23.14 22.14 Debugging training failures

Training failures are often mathematical symptoms. The table below lists common issues.

Symptom	Possible cause	Mathematical interpretation
training loss does not decrease	learning rate too small, bad labels, frozen parameters	optimizer not moving in useful descent directions
loss becomes NaN	learning rate too large, unstable normalization	numerical overflow or exploding gradients
training improves but validation fails	overfitting, leakage, distribution shift	low empirical risk, high population risk
model predicts only background	class imbalance	loss dominated by majority class
probabilities are overconfident	poor calibration	predicted posterior not reliable
results vary strongly by seed	high variance, small data	unstable estimator

A useful training debugging order is:

1. overfit a tiny batch;
2. check labels and preprocessing visually;
3. verify loss decreases on a simple model;
4. check gradient norms;
5. check train/validation split leakage;
6. compare to a classical baseline;
7. add complexity only after the simple pipeline works.

23.15 22.15 Python example: supervised image classification workflow

The following example is written for independent study. It uses a small TensorFlow workflow. In a real course, students can replace the synthetic data with CIFAR-10, Fashion-MNIST, medical images, or a domain-specific dataset.

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# Example dataset. Replace with a course dataset if desired.
(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()

# Normalize images to [0, 1].
x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

y_train = y_train.squeeze()
y_test = y_test.squeeze()

# Validation split.
n_val = 5000
x_val, y_val = x_train[-n_val:], y_train[-n_val:]
x_train, y_train = x_train[:-n_val], y_train[:-n_val]

num_classes = 10

augmentation = keras.Sequential([
    layers.RandomFlip("horizontal"),
    layers.RandomTranslation(0.08, 0.08),
    layers.RandomZoom(0.10),
], name="augmentation")

inputs = keras.Input(shape=(32, 32, 3))
x = augmentation(inputs)
x = layers.Conv2D(32, 3, padding="same", activation="relu")(x)
x = layers.BatchNormalization()(x)
x = layers.MaxPooling2D()(x)
x = layers.Conv2D(64, 3, padding="same", activation="relu")(x)
x = layers.BatchNormalization()(x)
x = layers.MaxPooling2D()(x)
x = layers.Conv2D(128, 3, padding="same", activation="relu")(x)
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dropout(0.3)(x)
outputs = layers.Dense(num_classes)(x)

model = keras.Model(inputs, outputs)

loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = keras.optimizers.AdamW(learning_rate=3e-4, weight_decay=1e-4)

model.compile(
    optimizer=optimizer,
    loss=loss,
    metrics=[keras.metrics.SparseCategoricalAccuracy(name="acc")]
)

callbacks = [
    keras.callbacks.ModelCheckpoint(
        "best_ch22_model.keras",
        monitor="val_acc",
        save_best_only=True,
        mode="max"
    ),
    keras.callbacks.EarlyStopping(
        monitor="val_acc",
        patience=8,
        mode="max",
        restore_best_weights=True
    )
]

history = model.fit(
    x_train, y_train,
    validation_data=(x_val, y_val),
    epochs=50,
    batch_size=128,
    callbacks=callbacks
)

print(model.evaluate(x_test, y_test, verbose=0))

The important mathematical objects in this code are:

• the empirical risk, implemented by the loss over mini-batches;
• the optimizer, which approximates a stochastic descent process;
• the augmentation distribution, which encodes invariance assumptions;
• the validation metric, which selects the checkpoint;
• the test metric, which estimates final generalization.

23.16 22.16 Python example: plotting training curves

import matplotlib.pyplot as plt

hist = history.history

epochs = range(1, len(hist["loss"]) + 1)

plt.figure(figsize=(7, 4))
plt.plot(epochs, hist["loss"], label="training loss")
plt.plot(epochs, hist["val_loss"], label="validation loss")
plt.xlabel("epoch")
plt.ylabel("loss")
plt.title("Training and validation loss")
plt.legend()
plt.show()

plt.figure(figsize=(7, 4))
plt.plot(epochs, hist["acc"], label="training accuracy")
plt.plot(epochs, hist["val_acc"], label="validation accuracy")
plt.xlabel("epoch")
plt.ylabel("accuracy")
plt.title("Training and validation accuracy")
plt.legend()
plt.show()

When reading training curves, students should ask:

1. Is the training loss decreasing?
2. Does validation loss track training loss?
3. Is there evidence of overfitting?
4. Is early stopping selecting a sensible checkpoint?
5. Are accuracy and loss telling the same story?

23.17 22.17 Python example: gradient norms

Gradient norms can reveal instability.

import tensorflow as tf

batch_x = x_train[:128]
batch_y = y_train[:128]

with tf.GradientTape() as tape:
    logits = model(batch_x, training=True)
    batch_loss = loss(batch_y, logits)

grads = tape.gradient(batch_loss, model.trainable_variables)

for var, grad in zip(model.trainable_variables, grads):
    if grad is not None:
        print(f"{var.name:45s} gradient norm = {tf.norm(grad).numpy():.4e}")

Very small gradient norms may indicate saturation, vanishing gradients, or frozen layers. Very large norms may indicate exploding gradients, bad scaling, or a learning rate that is too high.

23.18 22.18 Python example: mixup for image classification

import tensorflow as tf

num_classes = 10

# Convert sparse labels to one-hot labels.
y_onehot = tf.one_hot(y_train, depth=num_classes)

@tf.function
def mixup_batch(images, labels, alpha=0.2):
    batch_size = tf.shape(images)[0]
    indices = tf.random.shuffle(tf.range(batch_size))

    # Simple approximation: sample lambda uniformly near the Beta behavior.
    # For production use, sample lambda from Beta(alpha, alpha).
    lam = tf.random.uniform([], 0.0, 1.0)
    lam = tf.maximum(lam, 1.0 - lam)

    mixed_images = lam * images + (1.0 - lam) * tf.gather(images, indices)
    mixed_labels = lam * labels + (1.0 - lam) * tf.gather(labels, indices)
    return mixed_images, mixed_labels

images = tf.convert_to_tensor(x_train[:64])
labels = tf.convert_to_tensor(y_onehot[:64])

mixed_images, mixed_labels = mixup_batch(images, labels)
print(mixed_images.shape, mixed_labels.shape)

With mixup, the loss should use soft labels:

soft_label_loss = keras.losses.CategoricalCrossentropy(from_logits=True)

23.19 22.19 Python example: calibration with temperature scaling

Temperature scaling modifies logits by

\[ z \mapsto z/T. \]

Large \(T\) makes probabilities less sharp; small \(T\) makes them more confident. The temperature is selected on a validation set.

import numpy as np
import tensorflow as tf

val_logits = model.predict(x_val, batch_size=256)
val_labels = y_val

T = tf.Variable(1.0, dtype=tf.float32)
optimizer_T = tf.keras.optimizers.Adam(learning_rate=0.01)
nll = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

for step in range(300):
    with tf.GradientTape() as tape:
        scaled_logits = val_logits / T
        loss_T = nll(val_labels, scaled_logits)
    grad = tape.gradient(loss_T, [T])
    optimizer_T.apply_gradients(zip(grad, [T]))
    T.assign(tf.maximum(T, 0.05))

print("selected temperature:", float(T.numpy()))

23.20 22.20 AI-assisted components for this chapter

AI tools can help students study training, but they should be used as mathematical assistants rather than as substitutes for reasoning.

23.20.1 Prompt 1: Loss-function audit

I am training a vision model for [classification/segmentation/detection/depth]. The target variable is [describe target]. Explain which loss functions are appropriate, their statistical interpretations, and what failure modes I should check.

23.20.2 Prompt 2: Training-curve diagnosis

Here are my training and validation loss curves: [paste values or describe]. Diagnose whether the model is underfitting, overfitting, unstable, or data-limited. Suggest mathematical checks before changing the architecture.

23.20.3 Prompt 3: Augmentation validity

For the following vision task, list augmentations that preserve the label and augmentations that may change the label. Explain each choice in terms of invariance assumptions.

23.20.4 Prompt 4: Optimizer comparison

Compare SGD with momentum, Adam, and AdamW for this model and dataset. Explain the difference using gradient noise, coordinate scaling, and regularization.

23.20.5 Prompt 5: Reproducibility checklist

Create a reproducibility checklist for a computer vision experiment, including data split, random seeds, preprocessing, augmentation, optimizer, learning-rate schedule, checkpoint selection, and metrics.

ImportantAI verification rule

When an AI assistant recommends a training change, translate the recommendation into one of these mathematical categories:

\[ \text{optimization},\quad \text{regularization},\quad \text{data distribution},\quad \text{loss design},\quad \text{evaluation}. \]

If the recommendation cannot be assigned to one of these categories, it is probably too vague to trust.

23.21 22.21 Mini-project: train, diagnose, and report a small vision model

Choose a small image dataset such as Fashion-MNIST, CIFAR-10, a small medical image dataset, or a synthetic segmentation dataset. Your report should include:

1. the mathematical task definition;
2. the model architecture;
3. the loss function and its statistical interpretation;
4. the optimizer and learning-rate schedule;
5. regularization and augmentation choices;
6. training and validation curves;
7. final test metrics;
8. at least one ablation study;
9. at least one error analysis figure;
10. a reproducibility checklist.

A strong report does not merely show high accuracy. It explains why the model was trained in a mathematically defensible way.

23.22 22.22 Chapter summary

Training modern vision models is a mathematical workflow:

\[ \boxed{ \text{data} +\text{model} +\text{loss} +\text{optimizer} +\text{regularization} +\text{evaluation} } \]

The main ideas are:

1. supervised training is empirical risk minimization;
2. backpropagation is reverse-mode differentiation;
3. SGD uses noisy but computationally efficient gradient estimates;
4. optimization choices affect both convergence and generalization;
5. regularization controls overfitting and encodes assumptions;
6. augmentation is stochastic invariance modeling;
7. validation and test metrics must be interpreted statistically;
8. AI tools are useful only when their suggestions are connected to mathematical diagnostics.

The next chapters use these training principles for object detection and segmentation, where the outputs are no longer single labels but structured visual predictions.

23.23 Exercises

23.23.1 Conceptual exercises

1. Explain the difference between empirical risk and population risk in the context of image classification.
2. Why is a validation set not the same as a test set?
3. Give an example of an augmentation that is valid for one vision task but invalid for another.
4. Explain why overfitting can occur even when the training loss is minimized correctly.
5. Why can accuracy be misleading for imbalanced datasets?

23.23.2 Mathematical exercises

6. Let \(\ell(z,y)\) be softmax cross-entropy with one-hot target \(y\). Derive \[ \frac{\partial \ell}{\partial z_k}=p_k-y_k. \]
7. For squared loss regression, show that minimizing empirical squared loss is maximum likelihood under Gaussian noise.
8. Suppose a mini-batch gradient estimator averages \(B\) independent gradients with covariance \(\Sigma\). Show that its covariance is \(\Sigma/B\).
9. Derive the gradient-descent update for an \(L^2\)-regularized objective.
10. For a linear layer \(h=Wx\), compute \(\partial \mathcal L/\partial W\) in terms of \(x\) and \(\partial \mathcal L/\partial h\).
11. Show that label smoothing changes the target distribution from a point mass to a convex combination of the point mass and the uniform distribution.
12. For a validation accuracy \(\widehat p=0.86\) on \(m=1000\) examples, compute the approximate standard error.

23.23.3 Python exercises

13. Train a small CNN on Fashion-MNIST or CIFAR-10 and plot training and validation curves.
14. Compare SGD with momentum and AdamW on the same architecture.
15. Add data augmentation and measure the change in validation accuracy.
16. Implement label smoothing and compare calibration before and after smoothing.
17. Compute a confusion matrix and identify the most common error pairs.
18. Train with and without weight decay and compare the training-validation gap.
19. Implement a simple mixup pipeline and compare validation performance.
20. Estimate uncertainty in test accuracy using a bootstrap or repeated random splits.

23.23.4 AI-assisted exercises

21. Ask an AI assistant to diagnose a training curve. Then classify each recommendation as optimization, regularization, data-distribution, loss-design, or evaluation advice.
22. Ask an AI assistant to propose augmentations for a medical imaging task. Identify which augmentations might be scientifically invalid.
23. Ask an AI assistant to explain AdamW. Verify mathematically how decoupled weight decay differs from adding \(\lambda\theta\) to the gradient.
24. Ask an AI assistant to create an ablation table. Improve the table so that each comparison changes only one variable.
25. Ask an AI assistant to generate a reproducibility checklist for your mini-project. Add missing items related to random seeds, library versions, and data leakage.

23.24 Notes and references

Backpropagation is classically associated with Rumelhart et al. (1986). Stochastic approximation and large-scale stochastic optimization are connected to Robbins and Monro (1951) and Bottou (2010). Adam was introduced in Kingma and Ba (2015), and decoupled weight decay in AdamW is due to Loshchilov and Hutter (2019). Batch normalization was introduced in Ioffe and Szegedy (2015); dropout in Srivastava et al. (2014); label smoothing is discussed in Szegedy et al. (2016); mixup in Zhang et al. (2018); CutMix in Yun et al. (2019); and augmentation-search methods in Cubuk et al. (2019) and Cubuk et al. (2020). General deep learning background can be found in Goodfellow et al. (2016).

Bottou, Léon. 2010. “Large-Scale Machine Learning with Stochastic Gradient Descent.” Proceedings of COMPSTAT 2010, 177–86.

Cubuk, Ekin D., Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V. Le. 2019. “AutoAugment: Learning Augmentation Strategies from Data.” IEEE/CVF Conference on Computer Vision and Pattern Recognition, 113–23. https://arxiv.org/abs/1805.09501.

Cubuk, Ekin D., Barret Zoph, Jon Shlens, and Quoc V. Le. 2020. “RandAugment: Practical Automated Data Augmentation with a Reduced Search Space.” IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, 702–3. https://arxiv.org/abs/1909.13719.

Goodfellow, Ian, Yoshua Bengio, and Aaron Courville. 2016. Deep Learning. MIT Press. https://www.deeplearningbook.org/.

He, Kaiming, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2015. “Deep Residual Learning for Image Recognition.” arXiv Preprint arXiv:1512.03385. https://arxiv.org/abs/1512.03385.

Ioffe, Sergey, and Christian Szegedy. 2015. “Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift.” International Conference on Machine Learning, 448–56. https://arxiv.org/abs/1502.03167.

Kingma, Diederik P., and Jimmy Ba. 2015. “Adam: A Method for Stochastic Optimization.” International Conference on Learning Representations. https://arxiv.org/abs/1412.6980.

Loshchilov, Ilya, and Frank Hutter. 2019. “Decoupled Weight Decay Regularization.” International Conference on Learning Representations. https://arxiv.org/abs/1711.05101.

Robbins, Herbert, and Sutton Monro. 1951. “A Stochastic Approximation Method.” The Annals of Mathematical Statistics 22 (3): 400–407.

Rumelhart, David E., Geoffrey E. Hinton, and Ronald J. Williams. 1986. “Learning Representations by Back-Propagating Errors.” Nature 323 (6088): 533–36.

Srivastava, Nitish, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. 2014. “Dropout: A Simple Way to Prevent Neural Networks from Overfitting.” Journal of Machine Learning Research 15 (56): 1929–58. https://jmlr.org/papers/v15/srivastava14a.html.

Szegedy, Christian, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. 2016. “Rethinking the Inception Architecture for Computer Vision.” IEEE Conference on Computer Vision and Pattern Recognition, 2818–26. https://arxiv.org/abs/1512.00567.

Yun, Sangdoo, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. 2019. “CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features.” IEEE/CVF International Conference on Computer Vision, 6023–32. https://arxiv.org/abs/1905.04899.

Zhang, Hongyi, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. 2018. “Mixup: Beyond Empirical Risk Minimization.” International Conference on Learning Representations. https://arxiv.org/abs/1710.09412.