Learn ML

Chapter 20 of 25

The Training Pipeline

Putting every piece together into one training loop

13 min read

Motivation

Every previous chapter handed you one piece of a machine: an architecture, a forward pass, backpropagation, an activation function, a loss function, an optimizer, an initialization scheme, ways to regularize. Each piece works. But a pile of correct engine parts on a workbench isn't a running car. Something has to bolt them together, in the right order, and drive it.

This chapter is that assembly step. It takes every idea covered so far and arranges it into a single, concrete training loop — the actual sequence of operations that runs every time you call .fit() in a real project — and then covers the remaining practical concerns that don't belong to any one earlier chapter: how to size hyperparameters, how much compute and memory a network actually costs, how to adjust the learning rate over time, and how to tell whether the trained result is actually any good.

A Real-Life Scenario

Think of a professional kitchen preparing a complex dish for the first time. Having the best knife, the freshest ingredients, and a perfectly calibrated oven doesn't guarantee a good result if they're used in the wrong order or at the wrong moments — searing before seasoning, or serving before it's rested. The chef needs a recipe: a specific sequence that takes every one of those excellent components and combines them correctly.

Training an MLP is no different. You could have the best initialization scheme, the best optimizer, and the best regularization strategy individually, and still get a broken result if they're wired together incorrectly — say, computing gradients before shuffling the data, or evaluating on a validation set that was contaminated during preprocessing. The training pipeline is the recipe that makes sure every ingredient is used in the right place, at the right time.

Building the Intuition

Before the full algorithm, here's the shape of it:

  • Set up the parameters and the data before touching gradients at all: initialize weights sensibly, and preprocess features so they're all on comparable scales.
  • Then repeat, epoch after epoch: shuffle the data, split it into mini-batches, and for each mini-batch run the whole forward-loss-backward-update cycle.
  • After each epoch, check performance on data the network never trains on — the validation set — and use that check to decide when to stop.
  • Everything else (which optimizer, which learning rate, how much regularization) is a hyperparameter choice layered on top of this same skeleton.

1.Initialize & preprocess

Set weights with He/Xavier initialization; standardize input features using training-set statistics only.

2.Shuffle & batch

At the start of every epoch, shuffle the training data and split it into mini-batches of size B.

3.Forward pass

Propagate each mini-batch through the network to produce predictions.

4.Compute loss

Compare predictions against true labels using the task's loss function.

5.Backward pass

Backpropagate the loss to compute gradients for every weight and bias.

6.Update parameters

Apply the chosen optimizer's update rule using the freshly computed gradients.

7.Validate & check

Evaluate on the held-out validation set after each epoch; apply early stopping if performance has plateaued.

The Full Training Loop

Complete MLP Training Loop (Mini-Batch Gradient Descent)

Input: training data Dtrain\mathcal{D}_{\text{train}}, validation data Dval\mathcal{D}_{\text{val}}, learning rate η\eta, batch size BB, number of epochs EE.

  1. Initialize θ\theta using He or Xavier initialization.
  2. Standardize input features using training-set statistics.
  3. For epoch =1= 1 to EE:
    • Shuffle Dtrain\mathcal{D}_{\text{train}}.
    • For each mini-batch of size BB in Dtrain\mathcal{D}_{\text{train}}:
      • Forward propagate the batch.
      • Compute the batch loss.
      • Backpropagate to compute gradients.
      • Update θ\theta using the chosen optimizer.
    • Evaluate on Dval\mathcal{D}_{\text{val}}; log metrics.
    • If early-stopping criterion is met: break.
  4. Return θ\theta (the best checkpoint by validation performance).

Intuition: every earlier chapter maps onto exactly one line of this loop. Initialization sets the starting point; standardization keeps input scales comparable; the inner loop is forward pass, loss, backprop, and optimizer update in sequence; and the outer loop's validation check is where regularization decisions like early stopping actually get applied.

Data Preprocessing

A four-step checklist, in order

Before training, always: (1) handle missing values, via imputation or removal; (2) encode categorical features, e.g. one-hot encoding; (3) standardize or normalize numeric features so all features sit on comparable scales — especially important for MLPs, since features on very different scales can let some weights dominate the initial gradient and slow convergence; and (4) split data into train/validation/test sets before computing any preprocessing statistics, fitting the scaler only on the training set to avoid data leakage.

Fitting the scaler on the wrong data leaks information

Computing feature means and standard deviations from the entire dataset, then splitting into train/test afterward, lets information from the test set quietly influence how training features are scaled. This is a subtle but real form of data leakage — always fit the scaler on the training split only, then apply that same fitted transform to validation and test data.

Hyperparameters: A Complete Checklist

HyperparameterTypical Range / Guidance
Number of hidden layers1-5 for tabular data; deeper for very complex problems, though MLPs rarely go as deep as CNNs or Transformers
Neurons per layerOften decreasing (e.g., 128 → 64 → 32) or constant; commonly powers of 2
Activation functionReLU (default), or its variants; sigmoid/softmax reserved for the output layer
Loss functionDetermined by the task
OptimizerAdam / AdamW (default first choice)
Learning rate η\etaTypically 10410^{-4} to 10210^{-2}; search on a log scale
Batch size BB32, 64, 128, 256 are common defaults
Number of epochsChosen together with early stopping, not fixed in advance
Dropout rate0.2-0.5 for moderately sized networks
L2 regularization strength λ\lambdaTypically 10510^{-5} to 10210^{-2}

Learning Rate Scheduling

Why a fixed learning rate is often suboptimal

Early in training, a larger learning rate helps make rapid progress across the loss landscape. Later in training, as parameters approach a good minimum, a smaller learning rate helps fine-tune without overshooting it. A single fixed learning rate across all of training is a compromise that is rarely ideal at both stages simultaneously.

ScheduleDescription
Step decayMultiply η\eta by a fixed factor (e.g., 0.1) every fixed number of epochs
Exponential decayηt=η0ekt\eta_t = \eta_0 e^{-kt}, smoothly decreasing every step
Cosine annealingηt\eta_t follows a cosine curve from η0\eta_0 down to near zero over training, often with periodic restarts
Reduce-on-plateauAutomatically reduce η\eta when validation loss stops improving for a set number of epochs
WarmupLinearly increase η\eta from a small value at the very start of training, before switching to another schedule — helps stabilize the first few steps, especially with adaptive optimizers or large batch sizes

Computational and Memory Complexity

Time and Memory Complexity of an MLP

For a network with LL layers of sizes n0,n1,,nLn_0, n_1, \dots, n_L, the number of parameters is:

#params==1L(n1n+n)\#\text{params} = \sum_{\ell=1}^{L} \big(n_{\ell-1}\cdot n_\ell + n_\ell\big)

Time complexity of one forward pass for a single example is O ⁣(=1Ln1n)O\!\left(\sum_{\ell=1}^L n_{\ell-1}n_\ell\right), dominated by the matrix multiplications; backpropagation costs the same order. For a mini-batch of size BB, this becomes O ⁣(Bn1n)O\!\left(B\sum_\ell n_{\ell-1}n_\ell\right).

Memory complexity requires storing all weights and biases (O(n1n)O(\sum_\ell n_{\ell-1}n_\ell)), plus every activation from the forward pass for every example in a batch, needed for backpropagation: O ⁣(Bn)O\!\left(B\sum_\ell n_\ell\right).

Worked example: counting parameters

For a network with architecture 106432110 \to 64 \to 32 \to 1 (fully connected, with biases):

#params=(10×64+64)704+(64×32+32)2080+(32×1+1)33=2817\#\text{params} = \underbrace{(10\times64+64)}_{704} + \underbrace{(64\times32+32)}_{2080} + \underbrace{(32\times1+1)}_{33} = 2817
Always count parameters before training

Use model.summary() in Keras, or sum(p.numel() for p in model.parameters()) in PyTorch, before starting a real training run. This immediately reveals whether a model is unreasonably large for the available data — risking overfitting — or for the available memory.

Evaluation and Error Analysis

TaskCommon Metrics
RegressionMean Squared Error, Mean Absolute Error, R2R^2
Binary classificationAccuracy, Precision, Recall, F1-score, ROC-AUC
Multi-class classificationAccuracy, macro/micro-averaged F1-score, Confusion Matrix
Accuracy alone can be dangerously misleading

On an imbalanced dataset — say, 95% negative and 5% positive — a trivial model that always predicts "negative" achieves 95% accuracy while being completely useless. This is why precision, recall, F1, and the confusion matrix itself are essential complements to accuracy, especially for imbalanced classification problems like fraud or disease detection.

Error analysis is underused but invaluable

Manually examining a sample of the model's actual mistakes is one of the most underused practical skills in machine learning. Look for patterns: are errors concentrated in a particular class, or a particular range of input values? This often reveals data quality issues, missing features, or systematic weaknesses that aggregate metrics alone can hide entirely.

Key Takeaways

Key Takeaways
  • The complete training pipeline integrates initialization, preprocessing, forward/backward propagation, an optimizer, and validation-based stopping into a single repeating loop.
  • Feature standardization before training, and fitting preprocessing statistics only on the training set, are essential steps to avoid slow convergence and data leakage.
  • Key hyperparameters span architecture size, activation choice, optimizer, learning rate, batch size, and regularization strength.
  • Learning rate scheduling — step decay, cosine annealing, reduce-on-plateau, warmup — often outperforms a single fixed learning rate throughout training.
  • Parameter count and computational complexity should be sanity-checked before training begins.
  • Evaluation should go beyond raw accuracy, especially on imbalanced data, and should include manual error analysis alongside aggregate metrics.

Common Mistakes

Skipping preprocessing order

Splitting data into train/validation/test after computing preprocessing statistics on the full dataset silently leaks test-set information into training. Always split first, fit second.

Treating epoch count as a fixed hyperparameter

Choosing a fixed number of epochs in advance ignores the fact that the right stopping point depends on validation performance, which can't be known ahead of time. Pair epoch count with early stopping rather than fixing it upfront.

Relying on accuracy for imbalanced problems

A high accuracy number on an imbalanced dataset can mask a model that never correctly identifies the minority class — the exact class that usually matters most in problems like fraud or disease detection.

Interview Corner

Quick Check
1. Why must the feature scaler be fit only on the training set, not the full dataset?
2. For a fraud-detection dataset with 1% positive examples, why is accuracy alone a poor evaluation metric?

Where This Shows Up in Practice

This exact loop — initialize, preprocess, iterate over mini-batches with forward/backward/update, validate, stop — is what every high-level training API is doing under the hood, whether it's model.fit() in Keras or scikit-learn, or a hand-written loop in PyTorch. The next chapter puts this pipeline into action across all three, so you can see precisely how each framework exposes (or hides) every one of these steps.