Weâll start off where we left off last time. Recall that we have the following two objects:
- The (input) feature space: $\mathscr X\subseteq\mathbb{R}^d$. This is the set of outcomes for the input feature random variable $\mathbf X$.
- The target space: $\mathscr Y\subseteq\mathbb{R}^k$. This is the set of outcomes for the target random variable $\mathbf Y$.
- A dataset
$$ \mathscr D=\{(\mathbf x_i, \mathbf y_i)\in\mathscr X\times\mathscr Y:i\in I\}, $$
which is just a collection of samples from the joint distribution $\mathbf X\times \mathbf Y$. Weâll let $N=\lvert\mathscr D\rvert$.
Preparing the data
Before we do anything with this data, weâll stack them. Currently, we have a bunch of vectors (the input features $\mathbf x_i$ and the targets $y_i$ are both vectors), but storing them as matrices will make computation easier and elegant. In other words, we create the input feature and target matrices $X\in\mathbb{R}^{N\times d}$, $Y\in\mathbb{R}^{N\times k}$ such that for each $i\in I$, the $i$-th row of $X$ is $\mathbf x_i$, and the $i$-th row of $Y$ is $\mathbf y_i$.
Note that in the special case where $k=1$, we can drop the target feature dimension from the notation, so the targets $y_i$ will be scalars and weâll stack them to a target vector $\mathbf y\in\mathbb{R}^N$.
Now that thatâs done, we can finally regress.
Formulating the Model
Recall from last time that we defined regression as the process to find (or âlearnâ) some function $\mathscr X\xrightarrow f\mathscr Y$ such that $f(\mathbf x_i)$ is close to $y_i$ for all $i\in I$.
Our function $f$ will take in weight parameters $\theta$, so youâll often see the notation $f(x; \theta)$ or $f_{\theta}$, which just means that the function acts on $x$, but utilizes the weights $\theta$ in itâs computation.
Weâll start off by making $f_\theta$ an affine transformation.
$$ f(\mathbf x;\theta)=\mathbf x^T\cdot\mathbf w + b $$
where $\mathbf w\in\mathbb{R}^{d}$ is the weight vector and $b\in\mathbb{R}$ is the bias. We have $d':=d+1$ parameters: $\theta'=(\mathbf w, b)$.
We can modify the feature vectors $\mathbf x$ by appending a constant 1 feature: for a feature vector $\mathbf x\in\mathbb{R}^d$, let $\tilde{\mathbf x}=(\mathbf x,1)\in\mathbb{R}^{d'}$. Our model can now be formulated as a linear transformation:
$$ f(\mathbf x; \theta')=\tilde{\mathbf x}^T\cdot\mathbf w' $$
where the weight vector $\mathbf w'=(\mathbf w, b)$.
For simplicity, weâll refactor the notation. From here on, we drop the primes and write $\mathbf w$ for $\mathbf w'$, $d$ for $d'$, and $\theta$ for $\theta'$. The bias $b$ is now absorbed into $\mathbf w$ as its last component, and every $\mathbf x$ is understood to have the appended 1. So weâll just say that our model is simply:
$$ f(\mathbf x; \theta)=\mathbf x^T\cdot\mathbf w $$
Loss Function (MSE)
Now that we have our function $f_{\theta}$, we will now need to learn the parameter vector $\theta \in \mathbb{R}^d$ such that $f(\mathbf x_i)$ is close to $y_i$ for all $i\in I$. But how do we measure how âcloseâ $y_i$ is to $f(\mathbf x_i)$ ?
A conventional approach in ML is to choose a loss function. Given a parameter vector $\theta$ and a dataset entry $(\mathbf x_i,y_i)$, this loss function needs to output an error value $\ell(\mathbf x_i,y_i;\theta)\in\mathbb{R}$. We will then need to somehow tune the parameters $\theta$ such that this is minimized.
For our regression problem, a common choice is Squared Error (SE):
$$ \ell_{SE}(\mathbf x_i,y_i;\theta)=|f_\theta(\mathbf x_i)-y_i|^2. $$
To get overall performance on the dataset, we can take the average of the losses on all the dataset entries:
$$ \ell(\mathscr D;\theta)=\frac{1}{N}\sum_{i\in I}\ell(\mathbf x_i,y_i;\theta) $$
and plugging in SE, we can estimate the total loss through the entire dataset as the Mean Squared Error (MSE)
$$ \ell_{MSE}(\mathscr D;\theta)=\frac{1}{N}\sum_{i\in I}|f_\theta(\mathbf x_i)-y_i|^2. $$
This is the loss function weâll use for regression.
From MSE to Least Squares
Since $f(\mathbf x;\theta) = \mathbf x^T\mathbf w$, we can âvectorizeâ the loss over the whole dataset at once1. Stacking predictions gives $X\mathbf w \in \mathbb R^N$, and stacking targets gives $\mathbf y\in\mathbb R^N$ (using the $k=1$ convention from before), so
$$ \ell_{MSE}(\mathscr D;\mathbf w) = \frac1N\sum_{i\in I}(\mathbf x_i^T\mathbf w - y_i)^2 = \frac1N|X\mathbf w-\mathbf y|^2. $$
Since $\frac1N$ is just a positive constant, it doesnât affect where the minimum occurs, so minimizing $\ell_{MSE}$ over $\mathbf w$ is equivalent to minimizing the unscaled quantity
$$ L(\mathbf w) = |X\mathbf w-\mathbf y|^2. $$
This is actually quite well-studied, and itâs called the least squares problem: we need to find
$$ \mathbf w^\star = \operatorname*{argmin}_{\mathbf w\in\mathbb R^{d}}\;|X\mathbf w - \mathbf y|^2. $$
And for this specific problem, we luckily have a closed form solution:
$$ \mathbf w=(X^TX)^{-1}X^T\mathbf y. $$
Expand to see the full derivation
We want to minimize $L(\mathbf w) = \|X\mathbf w - \mathbf y\|^2$ over $\mathbf w \in \mathbb R^{d}$.
First, expand $L$ using $\|\mathbf v\|^2 = \mathbf v^T\mathbf v$:
$$ \begin{aligned} L(\mathbf w) &= (X\mathbf w-\mathbf y)^T(X\mathbf w-\mathbf y) \\ &= (X\mathbf w)^T(X\mathbf w) - (X\mathbf w)^T\mathbf y - \mathbf y^T(X\mathbf w) + \mathbf y^T\mathbf y \\ &= \mathbf w^TX^TX\mathbf w - 2\mathbf w^TX^T\mathbf y + \mathbf y^T\mathbf y, \end{aligned} $$
where we used $(X\mathbf w)^T\mathbf y = \mathbf y^T(X\mathbf w)$, since both are scalars (so a quantity equals its own transpose), to combine the two middle terms.
Since $L$ is a scalar-valued function of the vector $\mathbf w$, we take its gradient with respect to $\mathbf w$. Using the standard identities $\nabla_{\mathbf w}(\mathbf w^TA\mathbf w) = 2A\mathbf w$ for symmetric $A$, and $\nabla_{\mathbf w}(\mathbf w^T\mathbf b) = \mathbf b$:
$$ \nabla_{\mathbf w}L(\mathbf w) = 2X^TX\mathbf w - 2X^T\mathbf y. $$
Setting this to $\mathbf 0$ to find the critical point:
$$ \begin{aligned} 2X^TX\mathbf w - 2X^T\mathbf y &= \mathbf 0 \\ X^TX\mathbf w &= X^T\mathbf y. \end{aligned} $$
These are the normal equations. Assuming $X^TX$ is invertible (i.e. $X$ has full column rank), we can solve directly:
$$ \mathbf w = (X^TX)^{-1}X^T\mathbf y. $$
Since $L(\mathbf w)$ is convex, its Hessian $\nabla^2_{\mathbf w}L(\mathbf w) = 2X^TX$ is positive semi-definite (and positive definite when $X$ has full column rank), so this critical point is indeed the unique global minimum.
This concludes the proof.
Here weâre quite lucky to have a closed form solution to find the global minimum. For most tasks, we wonât have a closed form solution nor will we arrive at the global minimum. Instead, weâll soon use an optimizer to approximate the minimum.
Now that we familiarized ourselves with the theory, take a moment to build some intuition with the interactive visual below!
Drag either endpoint of the line to change its slope $w$ and intercept $b$. Try to find the line that makes MSE as small as possible by hand before checking "Show optimal fit" to see the closed-form solution.
In Practice
The theory above gives us a nice recipe: compute $\mathbf{w}^\star = (X^TX)^{-1}X^T\mathbf{y}$ and then weâre done. But in practice, there are some things to keep in mind.
Generalization and the Train-Test Split
Recall that we want $f_\theta$ to perform well not just on the data weâve seen, but on new data drawn from the same distribution. The quantity we actually care about is the expected loss
$$ \mathbf{E}\,\ell(X, Y;\theta), $$
but we canât compute this, since we only have a finite dataset. So we estimate it with the average loss on the dataset.
The issue is, if we train on all of $\mathscr{D}$ and then evaluate on the same $\mathscr{D}$, we get a biased estimate. The model has already âseenâ every point, so its train loss can be low even if it has learned noise rather than signal. This is called overfitting.
To get a good estimate of generalization, we need two independently sampled datasets:
- A train set $\mathscr{D}_\text{train}$, used to fit $\theta$.
- A test set $\mathscr{D}_\text{test}$, used only to evaluate $f_\theta$ after training.
Since we rarely have two separate datasets, we simulate this by doing a train-test split on $\mathscr{D}$:
- Choose a split ratio (typically 80â90% train, 10â20% test).
- Randomly partition the index set $I$ into $I_\text{train}$ and $I_\text{test}$ according to that ratio.
- Set $\mathscr{D}_\text{train} = \{(\mathbf{x}_i, y_i) : i \in I_\text{train}\}$ and $\mathscr{D}_\text{test} = \{(\mathbf{x}_i, y_i) : i \in I_\text{test}\}$.
The test set must never influence training in any way! We need to ensure test loss an unbiased estimate of generalization!
Preprocessing
Raw datasets also rarely arrive ready to stack into a matrix. Two common issues:
Categorical features. If a feature takes values in some finite set $C = \{c_0, \ldots, c_{l-1}\}$, it isnât a number, and we canât just directly multiply by a weight. The standard fix is one-hot encoding: map $c_j$ to the $j$-th standard basis vector $e_j \in \mathbf{R}^l$. This expands one categorical column into $l$ binary columns, each indicating membership in one class.
Absorbing the bias. Recall that we folded the bias $b$ into $\mathbf{w}$ by appending a constant 1 to each feature vector. This is the standard convention, and it means $X$ ends up with one extra column of all ones.
Both of these steps happen before the split in terms of schema design, but any statistics-dependent transformations (e.g., mean normalization) must be fit on $\mathscr{D}_\text{train}$ only and then applied to $\mathscr{D}_\text{test}$.
Reproducibility and Random Seeds
The train-test split is random, which creates yet another problem: ML experiments should be reproducible, but they require random sampling. The standard resolution is a pseudorandom generator (PRNG). These are deterministic, stateful objects whose output passes statistical tests for randomness. By fixing a seed (the initial state), you can make every random choice in your experiment deterministic and therefore reproducible. So in practice, always set a seed before any stochastic operation, and record it alongside your results.
Evaluating the Model
After fitting $\mathbf{w}^\star$ on $\mathscr{D}_\text{train}$, we now must compute two losses:
$$ \ell_\text{train} = \frac{1}{|\mathscr{D}_\text{train}|}\,\|X_\text{train}\mathbf{w}^\star - \mathbf{y}_\text{train}\|^2, \qquad \ell_\text{test} = \frac{1}{|\mathscr{D}_\text{test}|}\,\|X_\text{test}\mathbf{w}^\star - \mathbf{y}_\text{test}\|^2. $$
You may encounter one of two scenarios:
- $\ell_\text{train} \approx \ell_\text{test}$: the model generalizes well. If both are high, the model is underfitting, meaning that it lacks the capacity or features to capture the signal.
- $\ell_\text{train} \ll \ell_\text{test}$: the model is overfitting, meaning that it has fit noise in the training data that doesnât transfer.
For linear regression specifically, overfitting is less of a concern than it will be for more expressive models. But the train/test evaluation pattern is universal! Every model we build going forward will be assessed this way.
-
Here $X\in\mathbb R^{N\times d}$ is understood to include the appended column of 1âs (absorbing the bias $b$ into $\mathbf w$), consistent with the refactoring above. ↩