Module 3. Cross-Entropy and KL Divergence

Let

$P=(0.9,0.1),\qquad Q=(1,0).$

1. Support mismatch

In bits,

Consequently,

$D_{\mathrm{KL},2}(P\|Q)=+\infty,$

because $P$ allows the second outcome while $Q$ declares it impossible.

2. $\varepsilon$-smoothing

For

$Q_\varepsilon=(1-\varepsilon,\varepsilon),$

we obtain

The entropy of $P$ is

Therefore,

3. Asymptotics as $\varepsilon\to0^+$

Since

$-\log_2(1-\varepsilon)=O(\varepsilon),$

the leading diverging term is

KL has the same coefficient:

Even an outcome with probability only $0.1$ creates an unbounded penalty as the model drives its probability toward zero.

4. What clipping does

Replacing probabilities by

q = np.clip(q, eps, 1.0)
q = q / q.sum()

creates a new smoothed distribution $\widetilde Q_\varepsilon$ (for the code shown, it is not literally $(1-\varepsilon,\varepsilon)$). This may be a sensible implementation safeguard, but the resulting loss evaluates the smoothed model, not the original $Q$. If clipping is applied only inside the logarithm and the values are not renormalized, the object may cease to be a probability distribution at all.

In practice, computing log-probabilities from logits through a stable log_softmax is preferable: finite logits already produce strictly positive softmax probabilities, with no manual clipping required.

Let

Using

we get

Differentiate with respect to $z_m$:

Hence

$\boxed{\nabla_z\ell=q-r}.$

For a one-hot target $r=e_y$,

$\nabla_z\ell=q-e_y.$

Hessian

The standard softmax derivative is

so

For any $v\in\mathbb R^K$,

Thus the Hessian is positive semidefinite and cross-entropy is convex in the logits.

The null direction corresponds to adding the same constant to every logit: softmax is invariant under $z\mapsto z+c\mathbf1$. But logits depend nonlinearly on a deep network's weights, so composing this loss with the network is not generally convex in the parameters.

1. Conditional cross-entropy decomposition

Start with

Add and subtract $\log p(Y\mid C)$:

The first term is conditional entropy:

$\mathbb E_P[-\log p(Y\mid C)]=H_P(Y\mid C).$

For a fixed $C=c$, the inner expectation in the second term is

Averaging over $C$ gives

2. Autoregressive KL chain rule

Factor both distributions:

Then

Taking expectation under $P$ and interchanging a finite sum with expectation,

The last step is the law of total expectation: for each fixed prefix, the inner expectation is taken under the true next-token distribution.

3. Why teacher-forced loss is not the whole deployment story

Teacher forcing evaluates prefixes distributed as $P(X_{<t})$, i.e. real prefixes from the dataset. During free generation, the prefix distribution is induced by $Q$. An early deviation changes later contexts and may move the model into regions poorly covered by teacher-forced validation. Good held-out NLL is therefore necessary for a good probabilistic model on the data distribution, but it is not a complete closed-loop generation evaluation.

Let

$p(x)=\tfrac12\mathcal N(x;-3,1)+\tfrac12\mathcal N(x;3,1),$

and

$q_{\mu,\sigma}(x)=\mathcal N(x;\mu,\sigma^2).$

1. Analytic forward-KL optimum

For fixed $P$, minimizing $D_{\mathrm{KL}}(P\|Q)$ is equivalent to minimizing the cross-entropy

$-\mathbb E_P[\log q_{\mu,\sigma}(X)].$

For a Gaussian,

Up to an additive constant, we minimize

The optimum in $\mu$ is

$\mu^*=\mathbb E_P[X]=0.$

Then

For the symmetric mixture,

Therefore,

2. Reverse KL with Gauss–Hermite quadrature

For reverse KL, reparameterize

$X=\mu+\sigma Z,\qquad Z\sim\mathcal N(0,1),$

and evaluate

The code below uses roots_hermitenorm; its nodes and weights directly match integrals with the $e^{-z^2/2}$ weight.

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize
from scipy.special import logsumexp, roots_hermitenorm
from scipy.stats import norm

LOG_2PI = np.log(2.0 * np.pi)

# E[f(Z)] for Z ~ N(0, 1)
z_nodes, z_weights = roots_hermitenorm(800)
z_weights = z_weights / np.sqrt(2.0 * np.pi)


def log_p(x: np.ndarray) -> np.ndarray:
    """Log-density of 0.5 N(-3,1) + 0.5 N(3,1)."""
    components = np.stack([
        np.log(0.5) - 0.5 * (x + 3.0) ** 2 - 0.5 * LOG_2PI,
        np.log(0.5) - 0.5 * (x - 3.0) ** 2 - 0.5 * LOG_2PI,
    ])
    return logsumexp(components, axis=0)


def reverse_kl(params: np.ndarray) -> float:
    mu, log_sigma = params
    sigma = np.exp(log_sigma)

    # Reparameterization X = mu + sigma Z, Z ~ N(0,1)
    x = mu + sigma * z_nodes
    log_q = -0.5 * z_nodes**2 - log_sigma - 0.5 * LOG_2PI

    return float(np.dot(z_weights, log_q - log_p(x)))


starts = [
    (-4.0, 0.0),
    (-2.0, 0.0),
    (0.0, 0.0),
    (2.0, 0.0),
    (4.0, 0.0),
    (0.0, 1.0),
]

solutions = []
for start in starts:
    result = minimize(
        reverse_kl,
        x0=np.array(start),
        method="L-BFGS-B",
        bounds=[(-8.0, 8.0), (-4.0, 3.0)],
        options={"ftol": 1e-13, "gtol": 1e-9, "maxiter": 2000},
    )
    if not result.success:
        raise RuntimeError(result.message)

    mu, log_sigma = result.x
    solutions.append((start, mu, np.exp(log_sigma), result.fun))

for start, mu, sigma, value in solutions:
    print(
        f"start={start!s:>12} -> "
        f"mu={mu:+.6f}, sigma={sigma:.6f}, reverse_KL={value:.6f}"
    )

Typical output is

start= (-4.0, 0.0) -> mu=-2.984306, sigma=1.023369, reverse_KL=0.688769
start= (-2.0, 0.0) -> mu=-2.984306, sigma=1.023369, reverse_KL=0.688769
start=  (0.0, 0.0) -> mu=+0.000000, sigma=2.7438,   reverse_KL=0.8406
start=  (2.0, 0.0) -> mu=+2.984306, sigma=1.023369, reverse_KL=0.688769
start=  (4.0, 0.0) -> mu=+2.984306, sigma=1.023369, reverse_KL=0.688769
start=  (0.0, 1.0) -> mu=+0.000000, sigma=2.7438,   reverse_KL=0.8406

The last digits of the central solution vary slightly with quadrature order because this is numerical integration. The symmetric minima near $\pm2.9843$ are stable and have the lower objective.

3. Visualization

x = np.linspace(-7.0, 7.0, 2000)
p = 0.5 * norm.pdf(x, -3.0, 1.0) + 0.5 * norm.pdf(x, 3.0, 1.0)

q_forward = norm.pdf(x, 0.0, np.sqrt(10.0))
q_reverse_right = norm.pdf(x, 2.984306, 1.023369)
q_reverse_center = norm.pdf(x, 0.0, 2.7438)

plt.figure(figsize=(10, 5))
plt.plot(x, p, linewidth=2.5, label="target mixture p")
plt.plot(x, q_forward, linewidth=2, label="forward KL optimum")
plt.plot(x, q_reverse_right, linewidth=2, label="reverse KL global optimum")
plt.plot(
    x,
    q_reverse_center,
    linewidth=2,
    linestyle="--",
    label="reverse KL central local optimum",
)
plt.xlabel("x")
plt.ylabel("density")
plt.legend()
plt.grid(alpha=0.25)
plt.show()

4. What the example establishes

Within the chosen Gaussian family, forward KL exactly matches mean and variance and therefore covers both modes. Reverse KL has two better solutions near the individual modes and one worse, broad central solution. This is a powerful illustration of asymmetry, not a universal law about every model family and optimizer.

import torch
import torch.nn.functional as F


torch.manual_seed(7)
logits = torch.randn(4, 5, dtype=torch.float64, requires_grad=True)
target = torch.tensor([0, 3, 1, 4])

# 1. Fused CE == log_softmax + NLL for class-index targets
loss_fused = F.cross_entropy(logits, target, reduction="mean")
loss_manual = F.nll_loss(
    F.log_softmax(logits, dim=-1),
    target,
    reduction="mean",
)
torch.testing.assert_close(loss_fused, loss_manual)

# 2. Per-example gradient is (q - one_hot) / batch_size for mean reduction
loss_fused.backward()
q = logits.detach().softmax(dim=-1)
one_hot = F.one_hot(target, num_classes=5).to(torch.float64)
expected_grad = (q - one_hot) / logits.shape[0]
torch.testing.assert_close(logits.grad, expected_grad)

# 3. Soft-label cross-entropy
soft_target = 0.9 * one_hot + 0.1 / logits.shape[-1]
loss_soft_fused = F.cross_entropy(
    logits.detach(),
    soft_target,
    reduction="mean",
)
loss_soft_manual = -(
    soft_target * F.log_softmax(logits.detach(), dim=-1)
).sum(dim=-1).mean()
torch.testing.assert_close(loss_soft_fused, loss_soft_manual)

# 4. Label smoothing changes the target distribution and the loss
loss_smooth = F.cross_entropy(
    logits.detach(),
    target,
    label_smoothing=0.1,
)
print("ordinary CE:", float(loss_fused.detach()))
print("smoothed CE:", float(loss_smooth))

For token perplexity, exponentiate the mean ordinary NLL of the observed target tokens in nats:

$\operatorname{PPL}=e^{\text{mean token NLL}}.$

The mean must include exactly the intended unmasked targets, each with its ordinary weight. A loss with label smoothing, class weights, arbitrary sample reweighting, or a different normalization may be useful for training, but its exponential is not standard perplexity.