Supervised learning as KL minimization

The model is a distribution, not a point estimate

$$ (x,y)\sim p_{\text{data}}(x,y)=p_{\text{data}}(x)\,p_{\text{data}}(y\mid x), \qquad \text{model: } p_\theta(y\mid x) \text{ (not a single number).} $$

A regression network outputs the mean of a Gaussian over $y$; a classifier outputs the full parameter vector of a Categorical over $y$ (Common distributions).

Intuition: "the network predicts $y$" really means "predicts a distribution over $y$," reporting its mean/mode; keeping that distribution in view is what makes the loss fall out of probability, not hand-design.

The objective: forward KL, averaged over $x$

At each $x$, training should pull $p_\theta(y\mid x)$ toward $p_{\text{data}}(y\mid x)$. Since $p_{\text{data}}$ is the fixed target being approximated, this is a forward KL (Lebesgue measure and KL divergence), averaged over $x\sim p_{\text{data}}(x)$:

$$ \min_\theta\ \mathbb{E}_{x\sim p_{\text{data}}}\Big[KL\big(p_{\text{data}}(y\mid x)\,\|\,p_\theta(y\mid x)\big)\Big]. $$

Intuition: forward KL is mode-covering and zero-avoiding: the right choice here, since missing mass that $p_{\text{data}}$ actually puts on some $y$ (for a given $x$) is exactly what a bad prediction looks like.

This is maximum likelihood

Expand the inner KL and drop the term with no $\theta$ in it:

$$ KL\big(p_{\text{data}}(y\mid x)\,\|\,p_\theta(y\mid x)\big) = \underbrace{\mathbb{E}_{y\sim p_{\text{data}}(y\mid x)}\big[\log p_{\text{data}}(y\mid x)\big]}_{-H(p_{\text{data}}(y\mid x)),\ \text{no }\theta} \ -\ \mathbb{E}_{y\sim p_{\text{data}}(y\mid x)}\big[\log p_\theta(y\mid x)\big]. $$

Averaging over $x\sim p_{\text{data}}$ and dropping that constant term, minimizing the KL objective is exactly equivalent to:

$$ \max_\theta\ \mathbb{E}_{(x,y)\sim p_{\text{data}}}\big[\log p_\theta(y\mid x)\big], $$

i.e. maximum likelihood estimation (MLE). $p_{\text{data}}$ is unavailable directly, so the expectation becomes its empirical average over the training set (law of large numbers):

$$ \max_\theta\ \frac1N\sum_{i=1}^N \log p_\theta(y_i\mid x_i) \qquad\Longleftrightarrow\qquad \min_\theta\ \underbrace{-\frac1N\sum_{i=1}^N \log p_\theta(y_i\mid x_i)}_{\text{negative log-likelihood (NLL)}}. $$

Intuition: forward-KL minimization, MLE, and "minimize the NLL loss" are three names for the same training objective, each one step more concrete than the last: distributional closeness, then an expectation, then a sum you can actually compute.

Regression: Gaussian output gives MSE

Let $p_\theta(y\mid x) = \mathcal{N}(y;\,f_\theta(x),\,\sigma^2)$ (Common distributions), with $f_\theta(x)$ the network's output and $\sigma^2$ fixed (not a function of $\theta$):

$$ -\log p_\theta(y\mid x) = \frac{(y-f_\theta(x))^2}{2\sigma^2} + \underbrace{\tfrac12\log(2\pi\sigma^2)}_{\text{const. in }\theta}. $$

Minimizing NLL over $\theta$ therefore only depends on the first term, and since $\sigma^2>0$ is a positive constant multiplier, it is equivalent to minimizing the sum of squared errors:

$$ \min_\theta\ \frac1N\sum_{i=1}^N\big(y_i-f_\theta(x_i)\big)^2, \qquad \text{the MSE loss.} $$

Intuition: MSE is not an arbitrary engineering choice: it is exactly what maximum likelihood gives you the moment you assume $y$ given $x$ is Gaussian-distributed around the network's prediction.

Classification: Categorical output gives cross-entropy

Let $p_\theta(y\mid x) = \text{Categorical}\big(y;\,\text{softmax}(f_\theta(x))\big)$, with $f_\theta(x)=z\in\mathbb{R}^K$ the network's logits (Common distributions). For true label $y=k$:

$$ -\log p_\theta(y\mid x) = -\log\big(\text{softmax}(f_\theta(x))_k\big), $$

which is precisely the cross-entropy loss between the one-hot label and the predicted class probabilities. This is not a coincidence of naming: for any two distributions $p,q$ on the same space, the cross-entropy $H(p,q):=-\mathbb{E}_{y\sim p}[\log q(y)]$ decomposes as

$$ H(p,q) = -\mathbb{E}_{y\sim p}[\log p(y)] - \mathbb{E}_{y\sim p}\left[\log\frac{q(y)}{p(y)}\right] = \underbrace{H(p)}_{\text{entropy of }p} + \underbrace{KL(p\|q)}_{\text{forward KL}}. $$

With $p=p_{\text{data}}(y\mid x)$ (one-hot at the true label) and $q=p_\theta(y\mid x)$: $H(p_{\text{data}})$ has no $\theta$, so $H(p_{\text{data}},p_\theta)$ and $KL(p_{\text{data}}\|p_\theta)$ differ by a constant: same minimizer.

Intuition: cross-entropy loss is not a separate idea from KL divergence; it is forward KL plus a constant that doesn't move under training, so "minimize cross-entropy" and "minimize forward KL from the true label distribution" are the same instruction.

Why forward KL: both losses seek a mean

Minimize each loss pointwise at a fixed $x$, over the free parameter of the output family, and check what the minimizer actually is.

Regression

$$ m^\star = \operatorname*{arg\,min}_{m\in\mathbb{R}}\ \mathbb{E}_{y\sim p_{\text{data}}(y\mid x)}\big[(y-m)^2\big] \ \ \Longrightarrow\ \ \frac{d}{dm}\mathbb{E}\big[(y-m)^2\big] = -2\,\mathbb{E}[y-m] = 0 \ \ \Longrightarrow\ \ m^\star = \mathbb{E}_{p_{\text{data}}}[y\mid x]. $$

Minimizing MSE sets $f_\theta(x)\to\mathbb{E}_{p_{\text{data}}}[y\mid x]$: the conditional mean of $y$, not any single observed value.

Classification

Minimize cross-entropy over $q\in\mathbb{R}^K$ on the simplex, $\sum_k q_k=1$, via a Lagrangian (The Lagrangian and KKT conditions):

$$ \min_q\ -\sum_{k=1}^K p_{\text{data}}(y{=}k\mid x)\log q_k \ \ \text{s.t.}\ \ \sum_k q_k = 1 \ \ \Longrightarrow\ \ q_k^\star = p_{\text{data}}(y{=}k\mid x). $$

Since $p_{\text{data}}(y{=}k\mid x) = \mathbb{E}_{p_{\text{data}}}\big[\mathbb{1}[y{=}k]\mid x\big]$, the optimal softmax output is exactly the mean of the one-hot label vector:

$$ \text{softmax}(f_\theta(x)) \ \to\ \mathbb{E}_{p_{\text{data}}}\big[e_y \mid x\big], \qquad e_y := \text{one-hot}(y). $$

Intuition: $T(y)=y$ for regression, $T(y)=e_y$ for classification; in both cases forward KL/MLE sets $f_\theta(x)\to\mathbb{E}_{p_{\text{data}}}[T(y)\mid x]$, the mean of whatever the output family is built to predict. That is why forward KL, not some other divergence: it is the one choice that reduces to "match the mean" for these families.

One objective, two losses

Same forward-KL/MLE objective, specialized only by output family: Gaussian $\to$ MSE, Categorical $\to$ cross-entropy. Fix the family and forward KL, and the loss is derived, not designed.

Intuition: Generative modeling runs the same machinery one level up, modeling $x$ itself instead of $y\mid x$.