Reading club I. Linear regression

Notes from an ML reading club session on machine-learning interatomic potentials.

The electronic Hamiltonian depends parametrically on the nuclear positions, so its eigenvalues, the potential energy surface, are determined once the nuclei are fixed (the Born-Oppenheimer approximation underpins this separation of electronic and nuclear degrees of freedom). The goal of a machine-learning interatomic potential is to fit a parameterized function that reproduces this potential energy surface across all nuclear configurations. Before any of the modern machinery, it is worth getting linear regression completely right, because most of the field’s kernel methods are linear regression in disguise.

Setup

We propose that the target function is linear. In one dimension, we seek the best $wx + b$ for a given dataset. In $D$ dimensions, the model becomes

$$f(\hat{\mathbf{x}}) = \sum_{j=1}^{D} w_j x_j$$

where $\mathbf{w}$ is the weight vector and $D$ is the dimensionality of the feature space (the bias $b$ is absorbed by appending a constant feature $x_0 = 1$ to each input). The objective is to learn, that is, to fit the weights $\mathbf{w}$ given a set of $N$ labeled training pairs $\{(\mathbf{x}^{(n)}, y^{(n)})\}_{n=1}^N$.

Primal form

We minimize the least-squares loss, which measures the squared deviation between predictions and labels,

$$L = \frac{1}{2} \sum_{n=1}^{N} \left(\sum_{j=1}^{D} w_j x_j^{(n)} - y^{(n)}\right)^2$$

Setting $\partial L / \partial w_i = 0$ yields the normal equations in matrix form, where $\mathbf{X}$ is the $N \times D$ design matrix with rows $(\mathbf{x}^{(n)})^T$,

$$\mathbf{X}^T\mathbf{X}\,\mathbf{w} = \mathbf{X}^T\mathbf{y}$$

Solving directly gives the primal form

$$\mathbf{w} = \left(\mathbf{X}^T\mathbf{X}\right)^{-1}\mathbf{X}^T\mathbf{y}, \qquad f(\hat{\mathbf{x}}) = \mathbf{w}^T\hat{\mathbf{x}}$$

Regularization

Adding an $\ell_2$ penalty $\lambda \|\mathbf{w}\|^2$ to the loss trades a slightly higher training error for better generalization. This is ridge regression. The modified normal equations become $(\mathbf{X}^T\mathbf{X} + \lambda \mathbf{I})\,\mathbf{w} = \mathbf{X}^T\mathbf{y}$, which is always invertible for $\lambda > 0$. The regularizer biases toward smaller weights, preventing overfitting to noise.

Dual form

We can rewrite the prediction as a linear combination of the training labels. The dual form is properly derived via the representer theorem, which states that the optimal weight vector lies in the span of the training data, $\mathbf{w} = \mathbf{X}^T\boldsymbol{\alpha}$. Substituting yields

$$f(\hat{\mathbf{x}}) = \hat{\mathbf{x}}^T \mathbf{X}^T \left(\mathbf{X}\mathbf{X}^T\right)^{-1} \mathbf{y}$$

The quantity $\mathbf{x}^{(n)} \cdot \hat{\mathbf{x}}$ appearing inside can be viewed as projecting the test point onto basis vectors defined by the training data. Intuitively, the solution already lives in the span of the training data, and we are finding the right linear combination. Notice that the predictions depend on the inputs only through dot products $\mathbf{x}^{(n)} \cdot \hat{\mathbf{x}}$, which is a prelude to kernelization. Replace the dot products with a kernel $k(\mathbf{x}^{(n)}, \hat{\mathbf{x}})$ and we implicitly work in a higher-dimensional feature space. That move gets its own session.

Maximum likelihood

Why least squares and not some other loss? Suppose the training labels are noisy observations of the true function, drawn from a Gaussian distribution centered at the model prediction with variance $\sigma^2$,

$$y^{(n)} \sim \mathcal{N}\!\left(\mathbf{w}^T\mathbf{x}^{(n)},\; \sigma^2\right)$$

which is equivalent to assuming additive i.i.d. Gaussian noise. For a given $\mathbf{w}$, the likelihood of observing the full training set is

$$P(\mathbf{y}\mid\mathbf{X}, \mathbf{w}) = \prod_{n=1}^{N} \frac{1}{\sqrt{2\pi\sigma^2}}\exp\!\left(-\frac{(\mathbf{w}^T\mathbf{x}^{(n)} - y^{(n)})^2}{2\sigma^2}\right)$$

The normalization factor does not depend on $\mathbf{w}$, so maximizing the likelihood is the same as maximizing the log-likelihood, which up to constants reduces to minimizing

$$-\log P = \frac{1}{2\sigma^2}\sum_{n=1}^{N}\left(\mathbf{w}^T\mathbf{x}^{(n)} - y^{(n)}\right)^2$$

and this is the least-squares loss $L$ again. The maximum likelihood estimator coincides with the ordinary least-squares solution, and under Gaussian noise least squares is thus the statistically optimal estimator.