Training Models

A linear regression model predicts a target $y$ as a weighted sum of features $x$ plus a bias term:

where we append a constant $x_0 = 1$ to absorb the bias into $\boldsymbol{\theta}$. Given a training set of $m$ examples stacked into the design matrix $\mathbf{X} \in \mathbb{R}^{m \times (d+1)}$ and targets $\mathbf{y} \in \mathbb{R}^m$, the mean-squared-error loss is:

Setting $\nabla_{\boldsymbol{\theta}} \text{MSE} = \mathbf{0}$ gives the normal equation:

That is the entire derivation. One line. The catch is that $(\mathbf{X}^T \mathbf{X})^{-1}$ is an $(d+1) \times (d+1)$ matrix inversion. For $d \approx 100$ it is fine. For $d \approx 10^6$ it is impossible.

lin_reg = LinearRegression()
lin_reg.fit(X, y)               # X: (m, d), y: (m, 1)
lin_reg.intercept_, lin_reg.coef_

def fit_linear_regression(X, y):
    m = X.shape[0]
    X_b = np.c_[np.ones((m, 1)), X]                 # prepend bias column
    theta = np.linalg.inv(X_b.T @ X_b) @ X_b.T @ y  # (X^T X)^-1 X^T y
    return theta

The normal equation gives the exact answer. So why does every deep learning framework train by gradient descent instead?

Three reasons, in order of importance.

Reason 1: $\mathbf{X}^T \mathbf{X}$ is too big. With $d$ features the matrix is $(d+1) \times (d+1)$, and inversion is $O(d^3)$. For a model with a million parameters that is $10^{18}$ operations per fit. Gradient descent is $O(md)$ per step and converges in tens to hundreds of steps. The asymptotics flip somewhere around $d \approx 10^4$.

Reason 2: most loss functions are not quadratic. Logistic regression's loss has no closed-form solution. Neither does a general neural network's. The normal equation is a special-case benefit of squared loss; iterative numerical optimization is usually required, and gradient-based methods are the workhorse when derivatives are available.

Reason 3: you can stream the data. Gradient descent (specifically the stochastic and mini-batch variants) processes one example or one mini-batch at a time. You never need to load the full $m \times d$ matrix into memory. The normal equation needs the whole dataset materialized at once.

So gradient descent wins on three axes that all matter at scale. The normal equation is still useful: it is the right answer when $d$ is small, the data fits in RAM, the loss is quadratic, and you want a deterministic result without worrying about learning rates. That covers a lot of textbook regression problems and zero modern ML problems.

The gradient of MSE with respect to $\boldsymbol{\theta}$ is:

Batch gradient descent updates with this gradient, using the full training set every step:

$\eta$ is the learning rate. Pick it too small and you crawl. Pick it too large and you oscillate or diverge. There is a Goldilocks band, and finding it is most of the practical art of training.

Library path (sklearn does not expose pure batch GD for regression; SGDRegressor is the closest, with batch-of-1):

from sklearn.linear_model import SGDRegressor

sgd_reg = SGDRegressor(max_iter=1000, tol=1e-5, penalty=None, eta0=0.01,
                       n_iter_no_change=100, random_state=42)
sgd_reg.fit(X, y.ravel())
print(sgd_reg.intercept_, sgd_reg.coef_)

From-scratch path:

import numpy as np

def batch_gradient_descent(X: np.ndarray, y: np.ndarray, eta: float = 0.1,
                            n_epochs: int = 1000) -> np.ndarray:
    """Full-batch GD. Returns final theta and the loss history."""
    m = X.shape[0]
    X_b = np.c_[np.ones((m, 1)), X]
    np.random.seed(42)
    theta = np.random.randn(X_b.shape[1], 1)
    losses = []
    for epoch in range(n_epochs):
        gradients = 2 / m * X_b.T @ (X_b @ theta - y)
        theta = theta - eta * gradients
        losses.append(float(np.mean((X_b @ theta - y) ** 2)))
    return theta, losses

theta, losses = batch_gradient_descent(X, y, eta=0.1, n_epochs=1000)
print(theta)   # [[4.21509616], [2.77011339]]  — same answer as the normal equation

The from-scratch version returns the same parameters as the normal equation (up to about 1e-8) when run long enough. That is the sanity check. Always verify your gradient descent converges to the closed-form answer on a problem where the closed-form answer exists, before you trust it on problems where it does not.

The learning rate is the most-tuned hyperparameter in machine learning and the one with the cleanest visual story. Géron's Figure 4-8 in Hands-On ML shows three runs of the same problem at $\eta = 0.02$, $0.1$, and $0.5$. The first one creeps toward the answer over hundreds of steps. The middle one lands in twenty. The last one oscillates, overshoots, and never settles.

The math: for a quadratic loss with Hessian eigenvalues bounded between $\mu$ and $L$ (smallest and largest), gradient descent converges if and only if $\eta < 2/L$, and converges fastest at $\eta \approx 2/(L + \mu)$. The Hessian is the matrix of second derivatives of the loss; for a quadratic it is constant, and its eigenvalues measure how sharply the loss curves in each direction. The "condition number" $\kappa = L/\mu$ — the ratio of the steepest to the shallowest direction — determines how fast: when $\kappa$ is large the loss surface is an elongated bowl (a narrow canyon), and gradient descent zig-zags slowly down it.

You do not compute $L$ and $\mu$ in practice. You guess, then read the loss curve. The right behavior looks like:

• Smoothly decreasing, eventually flattening: $\eta$ is right
• Decreasing very slowly: $\eta$ too small
• Decreasing then bouncing around: $\eta$ too high or batch size too small
• Increasing or going to NaN: $\eta$ way too high

The first thing to try when training does not work is to drop $\eta$ by 10x. The second thing is to standardize your features so that $L/\mu$ does not blow up.

Batch gradient descent has a problem: every update touches every training example. For $m = 10^6$ that is a lot of compute per step, and the steps are deterministic, so the trajectory is smooth but slow.

Stochastic gradient descent (SGD) flips the trade-off: one example per step. The gradient is noisy because a single example is a terrible estimate of the full gradient. But the updates are cheap, and the noise itself helps the optimizer escape shallow local minima (and saddle points — flat spots where the gradient is zero but the surface curves up in some directions and down in others, which in high dimensions matter much more than local minima).

Mini-batch gradient descent is the middle path. Use a small subset of examples (16, 32, 64, 128, 256) per step. Reduces gradient variance compared to SGD, processes data faster than full-batch, runs efficiently on GPUs because the batch axis vectorizes cleanly.

Library path (sklearn's SGDRegressor with partial_fit for mini-batches):

from sklearn.linear_model import SGDRegressor
import numpy as np

sgd_reg = SGDRegressor(max_iter=1, warm_start=True, learning_rate="constant",
                       eta0=0.01, penalty=None, random_state=42)

batch_size = 32
for epoch in range(50):
    perm = np.random.permutation(m)
    for i in range(0, m, batch_size):
        idx = perm[i:i + batch_size]
        sgd_reg.partial_fit(X[idx], y[idx].ravel())

From-scratch path (mini-batch with a learning-rate schedule):

def minibatch_gd(X: np.ndarray, y: np.ndarray, batch_size: int = 32,
                 n_epochs: int = 50, t0: float = 200, t1: float = 1000) -> np.ndarray:
    """Mini-batch GD with t0/(t+t1) learning rate schedule."""
    m = X.shape[0]
    X_b = np.c_[np.ones((m, 1)), X]
    np.random.seed(42)
    theta = np.random.randn(X_b.shape[1], 1)
    step = 0
    for epoch in range(n_epochs):
        perm = np.random.permutation(m)
        for i in range(0, m, batch_size):
            idx = perm[i:i + batch_size]
            xi = X_b[idx]
            yi = y[idx]
            gradients = 2 / len(idx) * xi.T @ (xi @ theta - yi)
            eta = t0 / (step + t1)
            theta -= eta * gradients
            step += 1
    return theta

The learning rate schedule $\eta_t = t_0 / (t + t_1)$ is one of many. It satisfies the Robbins-Monro conditions $\sum \eta_t = \infty$ and $\sum \eta_t^2 < \infty$: the rate has to keep summing to infinity so you can still travel any distance, but its squares have to stay finite so the noise eventually dies down. Together these guarantee convergence on convex problems — loss surfaces shaped like a bowl, where any straight line between two points on the curve stays above it, so there is one global optimum and no local minimum to get stuck in. For deep learning, cosine decay with warmup is standard. For convex problems, this old-school schedule still works.

Linear regression with linear features can only fit linear relationships. If your data is quadratic, you cheat: add $x^2$ as a feature and let the linear regression fit a parabola.

from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression

X = 6 * np.random.rand(100, 1) - 3
y = 0.5 * X**2 + X + 2 + np.random.randn(100, 1)

poly = PolynomialFeatures(degree=2, include_bias=False)
X_poly = poly.fit_transform(X)   # adds x^2 as a second column

lin_reg = LinearRegression()
lin_reg.fit(X_poly, y)
print(lin_reg.intercept_, lin_reg.coef_)
# [1.78134581] [[0.93366893 0.56456263]]  — close to (2, 1, 0.5)

For a degree-$k$ polynomial in $d$ features, PolynomialFeatures generates ${d+k \choose k}$ new features. This blows up fast. Degree 2 with 100 features gives 5151 columns; degree 3 gives 176,851. At some point you are no longer doing linear regression; you are doing a non-parametric basis expansion and you need regularization.

The deeper lesson is the bias-variance trade-off. Degree 1: high bias, low variance. The model is too rigid to fit the true curve. Degree 300: low bias, high variance. The model wiggles wildly to fit every training point, including the noise, and generalizes badly. Degree 2 (the truth): minimum total error.

Géron's Figure 4-14 shows all three on the same plot. Look at it. Then internalize that bias-variance is not just a textbook diagram; it is a daily diagnostic. When your test loss is much higher than your train loss, you have variance. When both are high, you have bias.

A learning curve plots train and validation loss as functions of training set size. It is the single most useful diagnostic tool in classical ML.

from sklearn.model_selection import learning_curve
from sklearn.linear_model import LinearRegression
import numpy as np

train_sizes, train_scores, valid_scores = learning_curve(
    LinearRegression(), X, y, train_sizes=np.linspace(0.01, 1.0, 40),
    cv=5, scoring="neg_root_mean_squared_error")

train_errors = -train_scores.mean(axis=1)
valid_errors = -valid_scores.mean(axis=1)

Three shapes to recognize.

Underfit (high bias): Both curves flatten at a high error. Train and valid almost coincide. Adding data does nothing. The model is too simple. Fix: add features, use a deeper model, decrease regularization.

Overfit (high variance): Train error is low, valid error is much higher. There is a gap. Adding data narrows the gap. Fix: more data, more regularization, simpler model.

Good fit: Both curves converge to a low error with a small remaining gap.

The shape of the gap is the diagnostic. The level of the curves is the second diagnostic. Together they tell you whether to add data, add features, or regularize.

To control variance, penalize large weights. Three flavors.

Ridge regression adds an L2 penalty on the weights to the loss:

The bias $\theta_0$ is not penalized (you do not want to shrink the intercept). Ridge shrinks all weights toward zero, smoothly. It has a closed form:

where $\mathbf{A}$ is the identity with $A_{00} = 0$. Notice this also fixes the numerical issue with the plain normal equation: $\mathbf{X}^T \mathbf{X} + \alpha \mathbf{A}$ is invertible even when $\mathbf{X}^T \mathbf{X}$ is singular.

Lasso regression uses L1 instead of L2:

The L1 norm has a corner at zero. The geometry of the optimization (Géron's Figure 4-19 shows this beautifully) means lasso drives some weights exactly to zero, producing a sparse solution. This is the key property: lasso does feature selection, ridge does not.

Elastic net combines both:

with $r \in [0, 1]$ controlling the L1/L2 mix.

ridge = Ridge(alpha=0.7, solver="cholesky")
ridge.fit(X, y)
ridge.intercept_, ridge.coef_

def ridge_regression(X, y, alpha=0.1):
    m, d = X.shape
    X_b = np.c_[np.ones((m, 1)), X]
    A = np.eye(d + 1)
    A[0, 0] = 0                                       # don't regularize the bias
    return np.linalg.inv(X_b.T @ X_b + alpha * A) @ X_b.T @ y

Lasso has no closed form. The standard solvers are coordinate descent (sklearn's default) and proximal gradient descent. Both are short to write but not as short as ridge.

Logistic regression is binary classification with a linear decision boundary. The model outputs a probability:

where $\sigma$ is the sigmoid (also called logistic) function. It squashes any real number into $(0, 1)$. The decision rule is: predict 1 if $\hat{p} \geq 0.5$, else 0.

The loss is binary cross-entropy (also called log-loss), summed over training examples:

This is convex in $\boldsymbol{\theta}$. The gradient is the same surprisingly-clean form as linear regression:

That is the magic. The gradient of the cross-entropy of the sigmoid of a linear function looks like the gradient of MSE of a linear function. The math falls out when you do the chain rule and the $\sigma'$ term cancels with the cross-entropy denominator. Boyd & Vandenberghe (another chapter of Convex Optimization) walks through this.

No closed-form solution for logistic regression. Gradient descent (or Newton's method — a second-order method that uses the Hessian's curvature, not just the gradient, to jump straight toward the minimum; fast on small convex problems, but it has to invert the Hessian, so it does not scale) is the only path.

log_reg = LogisticRegression(penalty=None, max_iter=500)
log_reg.fit(X, y)               # X: (m, d) standardized, y: (m,) 0/1
log_reg.predict_proba(X)[:, 1]  # P(class=1)

def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-np.clip(z, -500, 500)))

def logistic_grad(theta, X_b, y):
    p = sigmoid(X_b @ theta)
    return X_b.T @ (p - y) / X_b.shape[0]

def train(optimizer, X_b, y, n_steps=500):
    theta = np.zeros((X_b.shape[1], 1))
    for _ in range(n_steps):
        g = logistic_grad(theta, X_b, y)
        theta = optimizer.step(theta, g)   # e.g. GD: theta - lr*g
    return theta
# X_b = np.hstack([np.ones((m,1)), X])  -> theta[0] is the bias

For $K$ classes, generalize the sigmoid to the softmax:

One score per class. Softmax converts scores into a probability distribution over classes (positive, sums to 1). The predicted class is $\arg\max_k \hat{p}_k$.

You can sharpen or flatten that distribution with a temperature $\tau$ (the name is borrowed from physics, where a hotter system spreads its energy more evenly), dividing each score before exponentiating: $\hat{p}_k = e^{s_k/\tau} / \sum_j e^{s_j/\tau}$. As $\tau \to 0$ the distribution collapses onto a one-hot vector at the $\arg\max$; as $\tau \to \infty$ it flattens to uniform. Plain softmax is the $\tau = 1$ case. (This is the same scalar as the temperature scaling used for calibration in the safety lens, applied for a different purpose.)

The loss is categorical cross-entropy:

where $y_{ik}$ is a one-hot encoding of the label. Same gradient structure as the binary case: $\nabla_{\boldsymbol{\theta}_k} J = \frac{1}{m} \mathbf{X}^T (\hat{\mathbf{P}}_k - \mathbf{Y}_k)$.

softmax_reg = LogisticRegression(C=30)   # multinomial by default in recent sklearn
softmax_reg.fit(X, y)                     # y: integer class labels 0..K-1
softmax_reg.predict_proba(X)

def softmax(logits):
    logits = logits - np.max(logits, axis=-1, keepdims=True)  # stability
    exp = np.exp(logits)
    return exp / np.sum(exp, axis=-1, keepdims=True)

def softmax_regression(X, y, K, eta=0.1, n_epochs=1000):
    m, d = X.shape
    X_b = np.c_[np.ones((m, 1)), X]
    Theta = np.random.randn(d + 1, K) * 0.01
    Y = np.eye(K)[y]                       # one-hot encode labels
    for _ in range(n_epochs):
        P = softmax(X_b @ Theta)
        gradients = 1 / m * X_b.T @ (P - Y)
        Theta -= eta * gradients
    return Theta

SGD's gradient is noisy. One example tells you very little about the true gradient. Average over batches and you get a better estimate. Average over time and you get momentum.

The momentum update keeps a running velocity vector and updates parameters with that velocity rather than the raw gradient:

$\beta \in [0, 1)$ is the momentum coefficient, typically 0.9 or 0.99. When the gradient is consistent (always pointing the same way), the velocity grows and the optimizer accelerates. When the gradient flips sign (oscillating across a valley), the velocity cancels itself out and the oscillations damp.

The Distill explorable "Why Momentum Really Works" by Goh (2017) is the best visual on this. The takeaway: momentum gives quadratic acceleration on ill-conditioned problems. On a well-conditioned problem (eigenvalues all similar), momentum doesn't help much. On a badly-conditioned problem (one direction much steeper than another), momentum is the difference between converging in 10 steps and converging in 10,000.

Nesterov's accelerated gradient is a smarter variant: evaluate the gradient at the projected position $\boldsymbol{\theta} + \beta \mathbf{v}$ instead of at $\boldsymbol{\theta}$. It is slightly faster in theory and slightly more annoying to code. Most practitioners use plain momentum or Adam.

opt = torch.optim.SGD([theta], lr=0.01, momentum=0.9)             # heavy-ball
# opt = torch.optim.SGD([theta], lr=0.01, momentum=0.9, nesterov=True)
opt.zero_grad(); loss.backward(); opt.step()

class GDMomentum(Optimizer):
    def step(self, theta, grad):
        if self.v is None: self.v = np.zeros_like(theta)
        self.v = self.beta * self.v - self.lr * grad
        return theta + self.v

class Nesterov(Optimizer):
    def step(self, theta, grad):
        if self.v is None: self.v = np.zeros_like(theta)
        v_prev = self.v
        self.v = self.beta * self.v - self.lr * grad
        return theta - self.beta * v_prev + (1 + self.beta) * self.v

Momentum is one trick. Adaptive learning rates are another. Adam combines them.

Adam tracks two exponential moving averages per parameter:

• $\mathbf{m}^{(t)}$ — first moment (mean of gradients), same role as momentum
• $\mathbf{v}^{(t)}$ — second moment (uncentered variance of gradients), used to scale the step size; this per-parameter rescaling on its own, without the momentum term, is the optimizer called RMSProp

Update:

Bias-corrected:

Update:

Default hyperparameters: $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$. They almost always work.

Adam vs AdamW. Plain Adam couples L2 weight decay into the gradient via the loss, which interacts badly with the adaptive scaling. AdamW decouples them: the weight decay is applied directly to the parameters at update time, independent of the gradient. For deep models, AdamW is the standard. For shallow ML, Adam is fine.

opt = torch.optim.Adam([theta], lr=0.05, betas=(0.9, 0.999), eps=1e-8)
for _ in range(n_steps):
    opt.zero_grad(); loss = loss_fn(theta); loss.backward(); opt.step()

def step(self, theta, grad):
    if self.m is None:
        self.m = np.zeros_like(theta); self.v = np.zeros_like(theta)
    self.t += 1
    self.m = self.beta1 * self.m + (1 - self.beta1) * grad
    self.v = self.beta2 * self.v + (1 - self.beta2) * grad ** 2
    m_hat = self.m / (1 - self.beta1 ** self.t)
    v_hat = self.v / (1 - self.beta2 ** self.t)
    return theta - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

opt = torch.optim.AdamW([theta], lr=0.05, weight_decay=0.01, betas=(0.9, 0.999))
for _ in range(n_steps):
    opt.zero_grad(); loss = loss_fn(theta); loss.backward(); opt.step()

class AdamW(Adam):
    def __init__(self, lr=3e-4, weight_decay=0.01, **kwargs):
        super().__init__(lr=lr, **kwargs)
        self.weight_decay = weight_decay

    def step(self, theta, grad):
        theta = theta - self.lr * self.weight_decay * theta  # decoupled WD
        return super().step(theta, grad)                     # then plain Adam

The simplest regularization technique: train until the validation loss stops decreasing, then stop. Even if the training loss is still going down.

The mechanic: keep a held-out validation set. Every epoch (or every $N$ steps), compute the validation loss. Keep a running record of the best validation loss seen so far, and of the model parameters at that point. If you go $K$ evaluations without improving, stop and restore the best-known parameters.

from copy import deepcopy

def train_with_early_stopping(model, optimizer, train_loader, val_loader,
                               max_epochs=100, patience=10):
    best_val_loss = float('inf')
    best_params = None
    epochs_without_improvement = 0
    for epoch in range(max_epochs):
        # ... train one epoch ...
        val_loss = evaluate(model, val_loader)
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_params = deepcopy(model.state_dict())
            epochs_without_improvement = 0
        else:
            epochs_without_improvement += 1
            if epochs_without_improvement >= patience:
                break
    model.load_state_dict(best_params)
    return model

Géron calls this "a beautiful free lunch" and it is. Early stopping reduces overfitting, costs nothing extra to implement, and works for every model class. In sklearn, SGDRegressor(early_stopping=True, n_iter_no_change=10) does this for you.

The expected test error of a model can be decomposed into three components:

Bias is how far the average prediction (over many training sets) is from the truth. A model that is too simple has high bias.

Variance is how much the prediction wiggles across different training sets drawn from the same distribution. A model that is too complex has high variance.

Irreducible noise $\sigma^2$ is what you cannot fix by changing the model. It is the noise in the labels.

The trade-off: simpler models reduce variance, increase bias. More complex models reduce bias, increase variance. The optimal model complexity minimizes the sum. Regularization is a tool for shifting the trade-off.