{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-1 fw=all -->\n",
    "\n",
    "# Entropy, Cross-Entropy, and KL Divergence\n",
    "\n",
    "Nearly every model in this book is trained by minimizing a cross-entropy. The\n",
    "number your training loop prints is therefore a quantity from *information\n",
    "theory*, the field Claude Shannon created in a single paper published in two\n",
    "parts in 1948 [@Shannon.1948], and information theory is what tells you\n",
    "what that number *means*: it is a code length. Entropy is the irreducible\n",
    "floor, the cost of the data's own randomness; cross-entropy is what your model\n",
    "actually pays; and the gap between them, the Kullback--Leibler divergence, is\n",
    "the waste you can train away. This section builds those three quantities and\n",
    "the one inequality (Gibbs') that relates them, grounds the \"extra bits\"\n",
    "language in an actual coding argument (the Kraft inequality and the Shannon\n",
    "code), explains why language models report *perplexity*, and closes with two\n",
    "modern training tricks, label smoothing and knowledge distillation, that are\n",
    "one-line consequences of the same machinery.\n",
    "\n",
    "The logarithm base is a pure choice of unit: base $2$ gives *bits*, base $e$\n",
    "gives *nats*. **We work in nats throughout**, matching deep-learning practice:\n",
    "the library `log` is natural, so the loss your training loop prints is already\n",
    "in nats. The one place where base $2$ is genuinely natural is coding with\n",
    "binary symbols, and we flag it explicitly when we get there. Since\n",
    "$\\log_2 x = \\ln x / \\ln 2$, converting is a fixed rescaling:\n",
    "$1\\textrm{ bit} = \\ln 2 \\approx 0.693$ nats and\n",
    "$1\\textrm{ nat} = 1/\\ln 2 \\approx 1.443$ bits. No argmin in this section, and\n",
    "no trained model, changes if you switch.\n",
    "\n",
    "We import everything we need once, up front."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-imports"
   },
   "source": [
    "import math\n",
    "from mxnet import autograd, np, npx\n",
    "npx.set_np()"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-2 fw=all -->\n",
    "\n",
    "## Information and Entropy\n",
    "\n",
    "### Surprise and Self-Information\n",
    "\n",
    "What should a numerical measure of \"information\" even be? A friend shuffles a\n",
    "deck of cards, flips some over, and reports what they see. \"I see a card\"\n",
    "tells us nothing (we were already certain of it) and should score zero. \"I\n",
    "see a heart\" narrows four equally likely suits to one: mildly informative.\n",
    "\"It is the three of spades\" pins down one of $52$ equally likely outcomes:\n",
    "more. Reading off the order of the entire shuffled deck selects one outcome\n",
    "out of $52!$: a huge amount of information. Information tracks *surprise*:\n",
    "the less probable the event, the more we learn from observing it. Shannon\n",
    "turned this into a definition. The *self-information* of an event $x$ with\n",
    "probability $p(x)$ is\n",
    "\n",
    "$$\n",
    "I(x) = -\\log p(x),\n",
    "$$\n",
    "\n",
    "measured in nats for the natural logarithm. The four card reports carry\n",
    "$-\\ln 1 = 0$, $\\ln 4 \\approx 1.39$, $\\ln 52 \\approx 3.95$, and\n",
    "$\\ln 52! \\approx 156.4$ nats respectively; in bits (base $2$, the unit of one\n",
    "fair coin flip) they carry the classic numbers $0$, $2$,\n",
    "$\\log_2 52 \\approx 5.70$, and $\\log_2 52! \\approx 225.6$. The function\n",
    "$-\\log p$ is the unique choice (up to the base, i.e., the unit) satisfying\n",
    "the properties we implicitly demanded: certain events carry zero information,\n",
    "rarer events carry more, and independent events add, so that $I$ of a joint\n",
    "outcome with probability $p_1 p_2$ is $I(p_1) + I(p_2)$, which forces a\n",
    "logarithm (@Csiszar.2008 gives the formal axiomatics; Exercise 5\n",
    "walks through the argument).\n",
    "\n",
    "![Self-information $I(x) = -\\log p(x)$ in nats as a function of the probability $p$. Certain events ($p = 1$) carry zero information, the fair coin ($p = \\tfrac{1}{2}$) carries $\\ln 2 \\approx 0.693$ nats, and the curve diverges as $p \\to 0$: rare means surprising.](../img/mdl-it-self-info-curve.svg)\n",
    "\n",
    "the figure plots the curve: decreasing, convex, zero at\n",
    "$p = 1$, infinite at $p = 0$. Computing it is one line."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-self-information"
   },
   "source": [
    "import numpy as onp\n",
    "def self_information(p):\n",
    "    \"\"\"Self-information -log p, in nats.\"\"\"\n",
    "    return -onp.log(p)\n",
    "\n",
    "self_information(1 / 64)"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "4.1588830833596715"
     },
     "metadata": {},
     "execution_count": 2
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-3 fw=all -->\n",
    "\n",
    "An event of probability $1/64$ carries $\\ln 64 \\approx 4.16$ nats (exactly\n",
    "$6$ bits: it takes six fair-coin flips to have probability $1/64$).\n",
    "\n",
    "### Shannon Entropy\n",
    "\n",
    "Self-information scores a single outcome. To score a *random variable*, a\n",
    "whole distribution of outcomes, we average. For $X \\sim P$ with probability\n",
    "mass function (p.m.f.) or probability density function (p.d.f.) $p(x)$, the\n",
    "*entropy* (or *Shannon entropy*) of $X$ is the expected self-information,\n",
    "\n",
    "$$H(X) = - E_{x \\sim P} [\\log p(x)].$$\n",
    "\n",
    "For discrete $X$ this reads $H(X) = -\\sum_i p_i \\log p_i$ with\n",
    "$p_i = P(X = x_i)$. Each term weighs the surprise $-\\log p(x)$ of an outcome\n",
    "by how often it occurs, so entropy is the *average surprise* of observing\n",
    "$X$: a distribution concentrated on one value never surprises us ($H = 0$),\n",
    "while a spread-out distribution surprises us constantly. Why exactly this\n",
    "form? The logarithm is forced by additivity over independent observations,\n",
    "the minus sign makes the measure positive and decreasing in probability, and\n",
    "the expectation is the only consistent way to aggregate outcome-level\n",
    "surprise into a single number for the distribution; this is the content of\n",
    "the axiomatic characterizations mentioned above.\n",
    "\n",
    "For continuous $X$ with density $p$ the sum becomes an integral, and the\n",
    "result is called the *differential entropy*,\n",
    "\n",
    "$$\n",
    "h(X) = -\\int p(x) \\log p(x) \\,dx.\n",
    "$$\n",
    "\n",
    "We write lowercase $h$ for differential entropy throughout, to keep it\n",
    "distinct from the discrete $H$; the two behave differently in ways that\n",
    "matter. Differential entropy is *not* the limit of the discrete entropy, and\n",
    "it is *not coordinate-invariant*: under an invertible change of variables\n",
    "$y = g(x)$ it shifts by $E[\\log|\\det \\partial g/\\partial x|]$, so $h$ can\n",
    "even be *negative* (e.g., a narrow Gaussian with $\\sigma < 1/\\sqrt{2\\pi e}$).\n",
    "What survives a change of variables is anything built from the *ratio* of two\n",
    "densities: both densities pick up the same Jacobian factor, which cancels in\n",
    "$p(x)/q(x)$, so the KL divergence defined below is exactly invariant.\n",
    "Cross-entropy sits in between. It contains a single lone logarithm,\n",
    "$-E_P[\\log q]$, so it inherits the same Jacobian shift as entropy, but that\n",
    "shift depends only on the truth $P$ and the map $g$, never on the model $Q$,\n",
    "so it moves *every* candidate model's loss by the same constant.\n",
    "Reparameterize your data and all the cross-entropy values change, but their\n",
    "differences, their gradients in $Q$, and the argmin do not. That is the\n",
    "precise sense in which objectives that *compare* distributions are safe under\n",
    "reparameterization while raw differential entropy is not.\n",
    "\n",
    "In code, entropy needs one piece of care: the convention $0 \\log 0 = 0$ (an\n",
    "outcome of probability zero contributes nothing). Where $p(x) = 0$ the term\n",
    "$-p \\log p$ has limiting value $0$, but the direct floating-point expression\n",
    "`0 * -inf` evaluates to `nan`, so we sum with `nansum`, which drops those\n",
    "terms: exactly the convention we want."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-definition"
   },
   "source": [
    "import numpy as onp\n",
    "def entropy(p):\n",
    "    \"\"\"Entropy of a probability vector, in nats.\"\"\"\n",
    "    p = onp.asarray(p, dtype=float)\n",
    "    ent = -p * onp.log(p)\n",
    "    # `nansum` encodes the convention 0 log 0 = 0\n",
    "    return float(onp.nansum(ent))\n",
    "\n",
    "entropy(onp.array([0.1, 0.5, 0.1, 0.3]))"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "1.1682824501765625"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-4 fw=all -->\n",
    "\n",
    "The distribution $(0.1, 0.5, 0.1, 0.3)$ has entropy $\\approx 1.168$ nats:\n",
    "less than the $\\ln 4 \\approx 1.386$ nats of a uniform distribution on four\n",
    "outcomes, because the mass is unevenly spread.\n",
    "\n",
    "### Entropy Is Maximized by the Uniform Distribution\n",
    "\n",
    "Two basic properties orient everything that follows. First, for discrete $X$\n",
    "every term $-p_i \\log p_i$ is non-negative, so $H(X) \\geq 0$, with equality\n",
    "exactly for a point mass (a constant \"random\" variable). Second, entropy is\n",
    "largest when the distribution hedges maximally:\n",
    "\n",
    "**Proposition (uniform maximizes entropy).** *If $X$ takes at most $k$ values,\n",
    "then*\n",
    "\n",
    "$$\n",
    "H(X) \\leq \\log k,\n",
    "$$\n",
    "\n",
    "*with equality if and only if $X$ is uniform on $k$ values.*\n",
    "\n",
    "**Proof.** Write the entropy as an expectation of a concave function's\n",
    "argument: $H(X) = E[\\log(1/p(X))]$. Since $\\log$ is concave, Jensen's\n",
    "inequality (that section) gives\n",
    "\n",
    "$$\n",
    "H(X) = E\\!\\left[\\log \\frac{1}{p(X)}\\right]\n",
    "\\leq \\log E\\!\\left[\\frac{1}{p(X)}\\right]\n",
    "= \\log \\sum_{i\\,:\\,p_i > 0} p_i \\cdot \\frac{1}{p_i}\n",
    "= \\log m\n",
    "\\leq \\log k,\n",
    "$$\n",
    "\n",
    "where $m \\leq k$ is the number of outcomes with $p_i > 0$. Equality requires\n",
    "both inequalities to be tight. Because $\\log$ is *strictly* concave, the\n",
    "first is an equality only when $1/p(X)$ is constant almost surely, i.e., $X$\n",
    "is uniform on its $m$ support points; and $\\log m = \\log k$ forces $m = k$.\n",
    "Together: $p_i = 1/k$ for all $i$. $\\blacksquare$\n",
    "\n",
    "For continuous $X$ the analogous statements need constraints to be true at all\n",
    "(on a bounded interval the uniform density again maximizes differential\n",
    "entropy; under a variance constraint the Gaussian does, as we prove in\n",
    "that section).\n",
    "\n",
    "Consider the coin. A Bernoulli variable with heads-probability $p$ has the\n",
    "*binary entropy* $H(p) = -p \\log p - (1-p)\\log(1-p)$ (overloading $H$: the\n",
    "argument here is a parameter, not a random variable), plotted in\n",
    "the figure: zero at the deterministic endpoints\n",
    "$p \\in \\{0, 1\\}$, concave in between, and maximal at the fair coin\n",
    "$p = \\tfrac{1}{2}$, where it equals $\\ln 2 \\approx 0.693$ nats ($= 1$ bit,\n",
    "consistent with the $\\log k$ bound for $k = 2$). A biased coin with $p = 0.1$\n",
    "manages only $H(0.1) \\approx 0.325$ nats: you are rarely surprised by a coin\n",
    "that almost always lands tails.\n",
    "\n",
    "![The binary entropy $H(p) = -p \\log p - (1-p) \\log (1-p)$ in nats, concave and symmetric about $p = \\tfrac{1}{2}$, where it peaks at $\\ln 2 \\approx 0.693$ nats. At the deterministic endpoints $p = 0$ and $p = 1$ the outcome carries no surprise and the entropy vanishes.](../img/mdl-it-bernoulli-entropy.svg)\n",
    "\n",
    "Entropy also extends from one random variable to several: the joint entropy\n",
    "$H(X, Y)$, the conditional entropy $H(Y \\mid X)$, and the *mutual information*\n",
    "$I(X; Y)$ that measures what $X$ and $Y$ share. Those quantities power\n",
    "contrastive and self-supervised learning, and we develop them in their own\n",
    "section, that section.\n",
    "\n",
    "## Cross-Entropy and KL Divergence\n",
    "\n",
    "### The Kullback--Leibler Divergence\n",
    "\n",
    "In that section we measured distances between points with\n",
    "norms. We now want a notion of \"distance\" between *distributions*: how badly\n",
    "does a model $Q$ misrepresent a truth $P$? Information theory's answer is the\n",
    "*Kullback--Leibler (KL) divergence*, also called *relative entropy*\n",
    "[@Kullback.Leibler.1951]. Given a\n",
    "random variable $X \\sim P$ with p.d.f. or p.m.f. $p(x)$, and a second\n",
    "distribution $Q$ with p.d.f. or p.m.f. $q(x)$ that we use to approximate $P$,\n",
    "\n",
    "$$D_{\\textrm{KL}}(P\\|Q) = E_{x \\sim P} \\left[ \\log \\frac{p(x)}{q(x)} \\right].$$\n",
    "\n",
    "The term inside the expectation, $\\log p(x) - \\log q(x)$, is the difference of\n",
    "two surprises: how surprised $Q$ is by the outcome $x$, minus how surprised\n",
    "$P$ is. It is positive where $Q$ underestimates ($q(x) < p(x)$: $Q$ is *more*\n",
    "surprised than it should be) and negative where $Q$ overestimates. KL averages\n",
    "this *relative surprise* over outcomes drawn from the truth $P$. Note the\n",
    "asymmetry baked into the definition: $P$ supplies both the samples and the\n",
    "numerator. In general $D_{\\textrm{KL}}(P\\|Q) \\neq D_{\\textrm{KL}}(Q\\|P)$, and\n",
    "the gap can be dramatic (we exhibit it numerically below); the consequences of\n",
    "*which* direction you optimize are taken up in\n",
    "that section. One more edge case to keep in mind:\n",
    "if some outcome has $p(x) > 0$ but $q(x) = 0$ (the model declares impossible\n",
    "something that actually happens), then $D_{\\textrm{KL}}(P\\|Q) = \\infty$.\n",
    "\n",
    "Here is the discrete case in code, for `p` and `q` given as probability\n",
    "vectors over the same finite outcome set. KL divergence is non-negative on\n",
    "its own (Gibbs' inequality, next), so we do **not** wrap the result in\n",
    "`abs()`; doing so would teach the false idea that KL needs an absolute value\n",
    "to stay non-negative. The `nansum` is deliberate: where $p(x) = 0$ the direct\n",
    "floating-point expression for $0 \\cdot \\log(0/q)$ yields `nan`,\n",
    "which `nansum` drops to encode $0 \\log 0 = 0$; the *other* edge case,\n",
    "$p(x) > 0$ with $q(x) = 0$, instead yields $+\\infty$ (not `nan`), which\n",
    "`nansum` keeps, so the code correctly returns $+\\infty$ for that divergence."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-definition-2"
   },
   "source": [
    "import numpy as onp\n",
    "def kl_divergence(p, q):\n",
    "    \"\"\"KL(P || Q) for two probability vectors, in nats.\"\"\"\n",
    "    p, q = onp.asarray(p, dtype=float), onp.asarray(q, dtype=float)\n",
    "    kl = p * onp.log(p / q)\n",
    "    return float(onp.nansum(kl))"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-5 fw=all -->\n",
    "\n",
    "### Gibbs' Inequality\n",
    "\n",
    "The KL divergence cannot be negative, an observation going back to Gibbs\n",
    "[@Gibbs.1902]. Everything else in this section (cross-entropy as a sound\n",
    "loss, the optimality of code lengths, the label-smoothing optimum) follows\n",
    "from it.\n",
    "\n",
    "**Proposition (Gibbs' inequality).** *For any distributions $P$ and $Q$ on\n",
    "the same space,*\n",
    "\n",
    "$$\n",
    "D_{\\textrm{KL}}(P\\|Q) \\geq 0,\n",
    "$$\n",
    "\n",
    "*with equality if and only if $P = Q$.*\n",
    "\n",
    "**Proof.** One application of Jensen's inequality\n",
    "(that section) to the convex function $-\\log$:\n",
    "\n",
    "$$\n",
    "D_{\\textrm{KL}}(P\\|Q)\n",
    "= E_{x\\sim P}\\!\\left[-\\log\\frac{q(x)}{p(x)}\\right]\n",
    "\\geq -\\log E_{x\\sim P}\\!\\left[\\frac{q(x)}{p(x)}\\right]\n",
    "= -\\log \\!\\!\\sum_{x\\,:\\,p(x)>0}\\!\\! q(x)\n",
    "= -\\log Q(\\textrm{supp}\\,P)\n",
    "\\geq 0.\n",
    "$$\n",
    "\n",
    "The expectation runs only over the support of $P$, where the $p(x)$ factors\n",
    "cancel, and what remains is the $Q$-mass of that support: a number\n",
    "$Q(\\textrm{supp}\\,P) \\leq 1$, whose negative log is therefore $\\geq 0$ (the\n",
    "same computation runs with integrals for densities, with the same convention\n",
    "at $q = 0$). For equality\n",
    "throughout, both inequalities must be tight: because $-\\log$ is *strictly*\n",
    "convex, the first is an equality only when the ratio $q(x)/p(x)$ is constant\n",
    "$P$-almost surely, and the second requires $Q(\\textrm{supp}\\,P) = 1$, i.e.,\n",
    "$Q$ puts no mass outside $P$'s support. Summing $q(x) = c\\,p(x)$ over the\n",
    "support then gives $c = 1$, i.e., $P = Q$.\n",
    "$\\blacksquare$\n",
    "\n",
    "### Maximum Entropy and the Gaussian\n",
    "\n",
    "Gibbs' inequality has an immediate consequence for continuous distributions:\n",
    "among all densities with a given mean and variance, the Gaussian is the\n",
    "hardest to predict. that section granted this fact; here\n",
    "is the proof.\n",
    "\n",
    "**Proposition (Gaussian maximum entropy).** *Among all densities $p$ on\n",
    "$\\mathbb{R}$ with mean $\\mu$ and variance $\\sigma^2$,*\n",
    "\n",
    "$$\n",
    "h(P) \\leq \\tfrac{1}{2}\\log(2\\pi e \\sigma^2) = h(\\mathcal{N}(\\mu, \\sigma^2)),\n",
    "$$\n",
    "\n",
    "*with equality if and only if $P = \\mathcal{N}(\\mu, \\sigma^2)$.*\n",
    "\n",
    "**Proof.** Let $\\varphi(x) = (2\\pi\\sigma^2)^{-1/2}\n",
    "e^{-(x-\\mu)^2/(2\\sigma^2)}$ be the Gaussian density. The key observation is\n",
    "that $-\\log \\varphi(x) = \\tfrac{1}{2}\\log(2\\pi\\sigma^2) +\n",
    "\\tfrac{(x-\\mu)^2}{2\\sigma^2}$ is a quadratic in $x$, so its expectation under\n",
    "*any* density $p$ with mean $\\mu$ and variance $\\sigma^2$ depends on $p$ only\n",
    "through those two moments:\n",
    "\n",
    "$$\n",
    "-E_P[\\log \\varphi(x)]\n",
    "= \\tfrac{1}{2}\\log(2\\pi\\sigma^2) + \\frac{E_P\\big[(x-\\mu)^2\\big]}{2\\sigma^2}\n",
    "= \\tfrac{1}{2}\\log(2\\pi\\sigma^2) + \\tfrac{1}{2}\n",
    "= \\tfrac{1}{2}\\log(2\\pi e \\sigma^2).\n",
    "$$\n",
    "\n",
    "Taking $P = \\mathcal{N}(\\mu, \\sigma^2)$ itself, this computes the Gaussian's\n",
    "own differential entropy en route:\n",
    "$h(\\mathcal{N}(\\mu, \\sigma^2)) = \\tfrac{1}{2}\\log(2\\pi e \\sigma^2)$, the\n",
    "closed form quoted whenever a Gaussian entropy is needed (for instance in\n",
    "that section). Now apply Gibbs' inequality, which\n",
    "holds for densities as noted above, with $Q = \\mathcal{N}(\\mu, \\sigma^2)$.\n",
    "Because $E_P[\\log \\varphi]$ is finite, the divergence splits:\n",
    "\n",
    "$$\n",
    "0 \\leq D_{\\textrm{KL}}\\big(P \\,\\|\\, \\mathcal{N}(\\mu, \\sigma^2)\\big)\n",
    "= E_P[\\log p(x)] - E_P[\\log \\varphi(x)]\n",
    "= -h(P) + \\tfrac{1}{2}\\log(2\\pi e \\sigma^2).\n",
    "$$\n",
    "\n",
    "Rearranging gives the equation. Equality holds\n",
    "exactly when the divergence vanishes, which by Gibbs' equality case means\n",
    "$P = \\mathcal{N}(\\mu, \\sigma^2)$. $\\blacksquare$\n",
    "\n",
    "The bound works because the Gaussian's log-density is a polynomial in exactly\n",
    "the quantities being constrained, the first two moments, so the cross term\n",
    "$-E_P[\\log \\varphi]$ is the same for every competitor $P$.\n",
    "\n",
    "### Gaussians, in Closed Form\n",
    "\n",
    "The `kl_divergence` above works on *discrete* probability vectors. For\n",
    "*continuous* distributions the standard worked example is the KL divergence\n",
    "between two univariate Gaussians, which has a closed form. For\n",
    "$P = \\mathcal{N}(\\mu_1, \\sigma_1^2)$ and $Q = \\mathcal{N}(\\mu_2, \\sigma_2^2)$,\n",
    "\n",
    "$$\n",
    "D_{\\textrm{KL}}\\!\\left(\\mathcal{N}(\\mu_1,\\sigma_1^2) \\,\\big\\|\\, \\mathcal{N}(\\mu_2,\\sigma_2^2)\\right)\n",
    "= \\log\\frac{\\sigma_2}{\\sigma_1} + \\frac{\\sigma_1^2 + (\\mu_1 - \\mu_2)^2}{2\\sigma_2^2} - \\frac{1}{2}.\n",
    "$$\n",
    "\n",
    "These logs are natural logs, so the result is in nats. Read two things off\n",
    "the equation directly. First, it is non-negative\n",
    "and zero exactly when $\\mu_1=\\mu_2$ and $\\sigma_1=\\sigma_2$ (i.e., $P=Q$),\n",
    "just as Gibbs' inequality demands. Second, it is *not symmetric*: swapping the\n",
    "roles of $P$ and $Q$ changes the value, because the variance ratio and the\n",
    "mean gap enter asymmetrically. (Deriving the equation is\n",
    "Exercise 4.)\n",
    "\n",
    "Let's verify the formula against a direct Monte-Carlo estimate of\n",
    "$D_{\\textrm{KL}}(P\\|Q) = E_{x\\sim P}[\\log p(x) - \\log q(x)]$. (Note: no\n",
    "`abs()` anywhere; the divergence comes out non-negative on its own.)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-example-1"
   },
   "source": [
    "import numpy as onp\n",
    "def gaussian_kl(mu1, sigma1, mu2, sigma2):\n",
    "    \"\"\"Closed-form KL(N(mu1, sigma1^2) || N(mu2, sigma2^2)), in nats.\"\"\"\n",
    "    return (math.log(sigma2 / sigma1)\n",
    "            + (sigma1 ** 2 + (mu1 - mu2) ** 2) / (2 * sigma2 ** 2)\n",
    "            - 0.5)\n",
    "\n",
    "def gaussian_log_pdf(x, mu, sigma):\n",
    "    return (-0.5 * math.log(2 * math.pi * sigma ** 2)\n",
    "            - (x - mu) ** 2 / (2 * sigma ** 2))\n",
    "\n",
    "def mc_kl(mu1, sigma1, mu2, sigma2, n=200000):\n",
    "    \"\"\"Monte-Carlo estimate of the same KL by sampling x ~ P.\"\"\"\n",
    "    x = onp.random.normal(loc=mu1, scale=sigma1, size=(n, ))\n",
    "    return (gaussian_log_pdf(x, mu1, sigma1)\n",
    "            - gaussian_log_pdf(x, mu2, sigma2)).mean()"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-6 fw=all -->\n",
    "\n",
    "Take $P = \\mathcal{N}(0, 1)$ and $Q = \\mathcal{N}(1, 1)$ (a unit mean shift,\n",
    "equal variances). The closed form should match the Monte-Carlo estimate to\n",
    "within sampling noise, and both should equal\n",
    "$\\tfrac{(\\mu_1-\\mu_2)^2}{2\\sigma^2} = \\tfrac{1}{2}$ nats here."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-example-2"
   },
   "source": [
    "closed_form = gaussian_kl(0.0, 1.0, 1.0, 1.0)\n",
    "monte_carlo = mc_kl(0.0, 1.0, 1.0, 1.0)\n",
    "\n",
    "closed_form, monte_carlo"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(0.5, 0.5000583438782179)"
     },
     "metadata": {},
     "execution_count": 6
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-7 fw=all -->\n",
    "\n",
    "Now make the asymmetry explicit. Compare $D_{\\textrm{KL}}(P\\|Q)$ with\n",
    "$D_{\\textrm{KL}}(Q\\|P)$ when the two Gaussians have *different variances*, say\n",
    "$P=\\mathcal{N}(0,1)$ and $Q=\\mathcal{N}(0,4)$ (i.e., $\\sigma_2=2$). The two\n",
    "numbers differ ($\\approx 0.318$ vs. $\\approx 0.807$ nats): KL is a\n",
    "divergence, not a distance."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-example-3"
   },
   "source": [
    "kl_pq = gaussian_kl(0.0, 1.0, 0.0, 2.0)\n",
    "kl_qp = gaussian_kl(0.0, 2.0, 0.0, 1.0)\n",
    "\n",
    "kl_pq, kl_qp"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(0.3181471805599453, 0.8068528194400546)"
     },
     "metadata": {},
     "execution_count": 7
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-8 fw=all -->\n",
    "\n",
    "### Cross-Entropy\n",
    "\n",
    "The loss we actually implement in classifiers is not the KL divergence but a\n",
    "close relative. The *cross-entropy* from $P$ to $Q$ scores outcomes drawn from\n",
    "the truth $P$ by the model's surprise,\n",
    "\n",
    "$$\\textrm{CE}(P, Q) = - E_{x \\sim P} [\\log q(x)].$$\n",
    "\n",
    "Adding and subtracting $H(P)$ inside the expectation (legitimate whenever\n",
    "$H(P)$ is finite) splits the model's surprise into the truth's own surprise\n",
    "plus the relative surprise:\n",
    "\n",
    "$$\\textrm{CE} (P, Q) = H(P) + D_{\\textrm{KL}}(P\\|Q).$$\n",
    "\n",
    "This single identity, combined with Gibbs' inequality, gives the chain that\n",
    "explains why cross-entropy is *the* loss to minimize:\n",
    "\n",
    "$$\n",
    "\\underbrace{D_{\\textrm{KL}}(P\\|Q) \\ge 0}_{\\textrm{Gibbs}}\n",
    "\\;\\Longrightarrow\\;\n",
    "\\underbrace{\\textrm{CE}(P,Q) \\ge H(P)}_{\\textrm{cross-entropy} \\ \\ge\\ \\textrm{entropy}},\n",
    "\\quad \\textrm{equality iff } P = Q.\n",
    "$$\n",
    "\n",
    "Entropy $H(P)$ is the irreducible floor; cross-entropy is what you actually\n",
    "pay; and the gap $\\textrm{CE}(P,Q) - H(P) = D_{\\textrm{KL}}(P\\|Q)$ is the\n",
    "waste. Minimizing cross-entropy in $Q$ therefore drives $Q \\to P$ and squeezes\n",
    "the waste to zero, which is the same thing as minimizing KL, since $H(P)$\n",
    "does not depend on $Q$. (In the coding view below, floor, payment, and waste\n",
    "all become exact code lengths.)\n",
    "the figure draws the decomposition to scale\n",
    "for a four-outcome distribution we will code by hand in that coding view.\n",
    "\n",
    "![The decomposition the equation as code lengths, for the truth $P = (\\tfrac{1}{2}, \\tfrac{1}{4}, \\tfrac{1}{8}, \\tfrac{1}{8})$ and a uniform model $Q$. A code matched to the truth pays exactly the entropy floor $H(P) = 1.75$ bits per symbol; a code built for $Q$ pays the cross-entropy $\\textrm{CE}(P, Q) = 2.0$ bits, and the stacked increment is the KL divergence $D_{\\textrm{KL}}(P\\|Q) = 0.25$ bits of pure waste. Minimizing cross-entropy in $Q$ can only shrink the orange block; the blue floor is the data's own randomness.](../img/mdl-it-code-length-bars.svg)\n",
    "\n",
    "Let's verify the decomposition the equation numerically on two\n",
    "small categorical distributions, and check the asymmetry of KL while we are at\n",
    "it."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-kl-categorical"
   },
   "source": [
    "import numpy as onp\n",
    "p = onp.array([0.6, 0.3, 0.1])\n",
    "q = onp.array([0.2, 0.5, 0.3])\n",
    "\n",
    "ce_pq = float(-onp.sum(p * onp.log(q)))\n",
    "print(f'KL(P||Q) = {kl_divergence(p, q):.4f} nats, '\n",
    "      f'KL(Q||P) = {kl_divergence(q, p):.4f} nats')\n",
    "print(f'CE(P, Q) - H(P) = {ce_pq - float(entropy(p)):.4f} nats')"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "KL(P||Q) = 0.3961 nats, KL(Q||P) = 0.3653 nats\nCE(P, Q) - H(P) = 0.3961 nats\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-9 fw=all -->\n",
    "\n",
    "The forward and reverse KL disagree ($\\approx 0.396$ vs. $\\approx 0.365$\n",
    "nats), and $\\textrm{CE} - H$ reproduces the forward KL exactly, as\n",
    "the equation says it must.\n",
    "\n",
    "### The Classification Loss\n",
    "\n",
    "In a $k$-class classification problem, the \"truth\" for one example is the\n",
    "one-hot distribution that puts all its mass on the correct label $y$, and the\n",
    "model supplies a predicted distribution\n",
    "$\\hat{\\mathbf{y}} = (\\hat{y}_1, \\ldots, \\hat{y}_k)$. Cross-entropy\n",
    "the equation between them collapses to a single term,\n",
    "$-\\log \\hat{y}_{y}$: the model's surprise at the correct class. Averaged over\n",
    "$n$ examples, the *cross-entropy loss* is\n",
    "\n",
    "$$\n",
    "\\textrm{CE}(\\mathbf{y}, \\hat{\\mathbf{y}})\n",
    "= -\\frac{1}{n} \\sum_{i=1}^n \\log \\hat{y}_{i, y_i},\n",
    "$$\n",
    "\n",
    "in nats, the mathematical quantity a built-in classification loss computes.\n",
    "\n",
    "This display is the definition. Training code should supply logits to a\n",
    "from-logits loss, which performs the log-softmax without first rounding small\n",
    "probabilities to zero; that section derives that\n",
    "implementation."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-formal-definition-1"
   },
   "source": [
    "import numpy as onp\n",
    "def cross_entropy(y_hat, y):\n",
    "    \"\"\"Mean cross-entropy loss for predicted probabilities, in nats.\"\"\"\n",
    "    ce = -onp.log(y_hat[range(len(y_hat)), y])\n",
    "    return float(ce.mean())"
   ],
   "execution_count": 9,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-10 fw=all -->\n",
    "\n",
    "Two examples, three classes: the first example's true class is $0$, predicted\n",
    "with probability $0.3$; the second's is $2$, predicted with probability $0.5$.\n",
    "The loss is $\\tfrac{1}{2}(-\\ln 0.3 - \\ln 0.5) \\approx 0.949$ nats."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-formal-definition-2"
   },
   "source": [
    "import numpy as onp\n",
    "labels = onp.array([0, 2])\n",
    "preds = onp.array([[0.3, 0.6, 0.1], [0.2, 0.3, 0.5]])\n",
    "\n",
    "cross_entropy(preds, labels)"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "0.9485599924429406"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-11 fw=all -->\n",
    "\n",
    "The built-in loss computes the same number from *logits* $\\mathbf{z}$,\n",
    "applying log-softmax internally. We construct $\\mathbf{z}=\\log\\mathbf{q}$\n",
    "only because the small example above starts from normalized probabilities;\n",
    "in a model, $\\mathbf{z}$ is the raw output of the final layer. Softmax inverts\n",
    "the log for a normalized $\\mathbf{q}$:\n",
    "$\\mathrm{softmax}(\\log \\mathbf{q})_j = e^{\\log q_j} / \\sum_i e^{\\log q_i}\n",
    "= q_j / \\sum_i q_i = q_j$. (More generally, softmax is invariant to adding a\n",
    "constant to all logits, and $\\log$ of a normalized vector is one valid logit\n",
    "vector among many.) Passing probabilities to a loss intended for logits\n",
    "applies softmax twice; applying `log` to a computed softmax can underflow."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-cross-entropy-as-an-objective-function-of-multi-class-classification"
   },
   "source": [
    "preds_mx = np.array(preds.tolist())\n",
    "labels_mx = np.array(labels.tolist())\n",
    "logits = np.log(preds_mx)\n",
    "log_probs = npx.log_softmax(logits, axis=1)\n",
    "-log_probs[np.arange(len(labels_mx)), labels_mx].mean()"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "array(0.9485576)"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-12 fw=all -->\n",
    "\n",
    "The loss also has a decision-theoretic pedigree. A\n",
    "*scoring rule* assigns a penalty $S(\\mathbf{q}, y)$ for having reported the\n",
    "probability vector $\\mathbf{q}$ when class $y$ then occurs. The rule is\n",
    "*proper* if truthful reporting is optimal: when the data-generating distribution is $\\mathbf{p}$, the\n",
    "expected penalty $E_{y \\sim \\mathbf{p}}[S(\\mathbf{q}, y)]$ is minimized by\n",
    "reporting $\\mathbf{q} = \\mathbf{p}$, and *strictly proper* if this report is the\n",
    "*unique* optimum. The cross-entropy loss is the *log score*\n",
    "$S(\\mathbf{q}, y) = -\\log q_y$, and it is strictly proper: its expected\n",
    "penalty is exactly\n",
    "$\\textrm{CE}(\\mathbf{p}, \\mathbf{q}) = H(\\mathbf{p}) +\n",
    "D_{\\textrm{KL}}(\\mathbf{p} \\| \\mathbf{q})$, which Gibbs' inequality minimizes\n",
    "uniquely at $\\mathbf{q} = \\mathbf{p}$. Strict propriety of the log score *is*\n",
    "Gibbs' inequality restated (Exercise 6 asks you to write out the\n",
    "argument).\n",
    "\n",
    "Why insist on *strict* propriety? At population risk, over a model class that\n",
    "can represent the true conditional distribution, the log score is minimized\n",
    "by $p(y \\mid \\mathbf{x})$ itself. That is the basis for calibration\n",
    "(predicted probabilities matching observed frequencies). Finite data, a\n",
    "restricted model class, regularization, imperfect optimization, and\n",
    "distribution shift can still produce miscalibration. Nor is the log score the\n",
    "only strictly proper rule: the *Brier score*\n",
    "$S(\\mathbf{q}, y) = \\sum_j (q_j - \\mathbf{1}[j = y])^2$ penalizes the whole\n",
    "vector by squared error, and its expected penalty exceeds the truthful\n",
    "reporter's by $\\|\\mathbf{q} - \\mathbf{p}\\|^2$: a squared-distance waste term\n",
    "in place of the KL-flavored one. Different strictly proper rules disagree\n",
    "about how severely to punish which mistakes (the log score's penalty diverges\n",
    "on confident errors; the Brier score's stays bounded), but they agree about\n",
    "where the optimum is: the truth [@Gneiting.Raftery.2007].\n",
    "\n",
    "Why is *this* the loss that training a probabilistic classifier by maximum\n",
    "likelihood produces? We proved that already, in\n",
    "that section: the average negative log-likelihood of\n",
    "any i.i.d. dataset *is* the cross-entropy from the empirical distribution\n",
    "$\\hat p_{\\textrm{data}}$ to the model, so maximizing likelihood, minimizing\n",
    "cross-entropy, and minimizing\n",
    "$D_{\\textrm{KL}}(\\hat p_{\\textrm{data}} \\| p_{\\boldsymbol{\\theta}})$ are one\n",
    "and the same optimization, for binary labels, multiclass labels, and\n",
    "densities alike. We will not re-derive it here; this section's contribution is\n",
    "the *interpretation* of that loss as a code length, which we build next.\n",
    "\n",
    "## The Coding View and Perplexity\n",
    "\n",
    "We have repeatedly promised that entropy is a \"floor\" and KL the \"extra\"\n",
    "cost. Those words are statements about\n",
    "compressing data with binary codes, and the proofs are short enough to give in\n",
    "full. Binary codes speak base $2$, so this subsection works in bits; divide by\n",
    "$\\ln 2$ (i.e., reinterpret every $\\log$ below) and every statement holds in\n",
    "nats verbatim.\n",
    "\n",
    "### Prefix Codes and the Kraft Inequality\n",
    "\n",
    "Suppose we must transmit a stream of symbols drawn i.i.d. from a distribution\n",
    "$P$ over $k$ outcomes, encoding each outcome $x_i$ as a binary string\n",
    "(*codeword*) of length $l_i$. So that the receiver can split the stream back\n",
    "into codewords without separators, we require the code to be *prefix-free*: no\n",
    "codeword is a prefix of another (like telephone numbers: once a dialed string\n",
    "matches a number, no longer number starts the same way). Short codewords are\n",
    "precious, and the next proposition [@Kraft.1949] says exactly how\n",
    "precious.\n",
    "\n",
    "**Proposition (Kraft inequality).** *A prefix-free binary code with codeword\n",
    "lengths $l_1, \\ldots, l_k$ exists if and only if*\n",
    "\n",
    "$$\n",
    "\\sum_{i=1}^k 2^{-l_i} \\leq 1.\n",
    "$$\n",
    "\n",
    "**Proof.** Identify each codeword $c$ of length $l$ with the dyadic interval\n",
    "$[0.c,\\, 0.c + 2^{-l}) \\subseteq [0, 1)$ of all real numbers whose binary\n",
    "expansion starts with $c$ (e.g., the codeword $10$ owns\n",
    "$[0.10_2, 0.11_2) = [\\tfrac{1}{2}, \\tfrac{3}{4})$, of length $2^{-2}$). One\n",
    "codeword is a prefix of another exactly when the second's interval sits inside\n",
    "the first's, so a prefix-free code corresponds to *disjoint* intervals, whose\n",
    "total length $\\sum_i 2^{-l_i}$ cannot exceed $1$. Conversely, given lengths\n",
    "with $\\sum_i 2^{-l_i} \\leq 1$, sort them increasingly and lay down intervals\n",
    "of length $2^{-l_i}$ left to right; each starts at a multiple of $2^{-l_i}$,\n",
    "hence is a dyadic interval, and its binary address is a valid codeword.\n",
    "$\\blacksquare$\n",
    "\n",
    "The same inequality constrains every *uniquely decodable* code (one whose\n",
    "concatenations parse unambiguously even without prefix-freeness), so\n",
    "restricting to prefix-free codes costs nothing in code length\n",
    "[@McMillan.1956].\n",
    "\n",
    "![A prefix-free binary code as a tree, for lengths $1, 2, 3, 3$: each codeword is a root-to-leaf path (left edge $= 0$, right edge $= 1$), and prefix-freeness says every codeword ends at a leaf: no codeword continues through another. A leaf at depth $\\ell$ is the tree view of a dyadic interval of length $2^{-\\ell}$ in the proof, and here the Kraft budget is spent exactly: $2^{-1} + 2^{-2} + 2^{-3} + 2^{-3} = 1$.](../img/mdl-it-kraft-tree.svg)\n",
    "\n",
    "the figure redraws the proof's intervals as a binary\n",
    "tree, the picture to keep. The Kraft inequality turns code design into a budget\n",
    "problem: each symbol buys \"address space\" $2^{-l_i}$ out of a total budget of\n",
    "$1$, and shorter codewords cost exponentially more. The optimal spend follows\n",
    "immediately.\n",
    "\n",
    "**Proposition (entropy bounds the code length).** *Every prefix-free binary\n",
    "code for $P$ has expected length $E[l] \\geq H_2(P)$ (entropy in bits), and\n",
    "the* Shannon code *with lengths $l_i = \\lceil \\log_2 (1/p_i) \\rceil$ is\n",
    "prefix-free and achieves $E[l] < H_2(P) + 1$.*\n",
    "\n",
    "**Proof.** For the lower bound, let $Z = \\sum_i 2^{-l_i} \\leq 1$ (Kraft) and\n",
    "define the distribution $q_i = 2^{-l_i}/Z$, the one *implied* by the code\n",
    "lengths. Then $l_i = -\\log_2 q_i - \\log_2 Z$, so\n",
    "\n",
    "$$\n",
    "E[l] = \\sum_i p_i l_i\n",
    "= -\\sum_i p_i \\log_2 q_i - \\log_2 Z\n",
    "= H_2(P) + D_{\\textrm{KL}}^{(2)}(P \\| Q) - \\log_2 Z\n",
    "\\geq H_2(P),\n",
    "$$\n",
    "\n",
    "since the KL term is non-negative by Gibbs and $-\\log_2 Z \\geq 0$ (the\n",
    "superscript $(2)$ marks base-$2$ logs, here and in $\\textrm{CE}_2$ below).\n",
    "For the\n",
    "upper bound, the Shannon lengths satisfy\n",
    "$2^{-l_i} = 2^{-\\lceil \\log_2(1/p_i)\\rceil} \\leq p_i$, so they pass Kraft's\n",
    "test ($\\sum_i 2^{-l_i} \\leq \\sum_i p_i = 1$) and a prefix code with these\n",
    "lengths exists; and $l_i < \\log_2(1/p_i) + 1$ gives\n",
    "$E[l] < H_2(P) + 1$. $\\blacksquare$\n",
    "\n",
    "So entropy is, within one bit, *the* price of communicating draws from\n",
    "$P$: this is (the source-coding half of) Shannon's theorem, and the\n",
    "one-bit slack can be driven to zero by coding long blocks of symbols at once\n",
    "[@Cover.Thomas.1999]. Moreover, look at where the KL divergence appeared\n",
    "in the proof: if you build your code from the *wrong* distribution $Q$\n",
    "(spending $-\\log_2 q_i$ bits on symbol $i$, as the Shannon recipe would if it\n",
    "believed $Q$), your expected cost on data that is really $P$ is\n",
    "\n",
    "$$\n",
    "-\\sum_i p_i \\log_2 q_i = \\textrm{CE}_2(P, Q) = H_2(P) + D_{\\textrm{KL}}^{(2)}(P\\|Q).\n",
    "$$\n",
    "\n",
    "Cross-entropy is the price of coding with the wrong codebook, and the KL\n",
    "divergence is *literally the extra bits per symbol* you waste. Every \"extra\n",
    "nats\" claim earlier in the section is this statement, rescaled by $\\ln 2$.\n",
    "\n",
    "The bound is tight enough to see in a four-line example. Take\n",
    "$P = (\\tfrac{1}{2}, \\tfrac{1}{4}, \\tfrac{1}{8}, \\tfrac{1}{8})$, whose\n",
    "probabilities are exact powers of two, so the Shannon code wastes nothing; a\n",
    "codebook built for the *uniform* distribution instead spends $2$ bits on every\n",
    "symbol."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-coding"
   },
   "source": [
    "p = [1/2, 1/4, 1/8, 1/8]\n",
    "lengths = [math.ceil(math.log2(1 / p_i)) for p_i in p]\n",
    "kraft = sum(2 ** -l for l in lengths)\n",
    "avg_len = sum(p_i * l for p_i, l in zip(p, lengths))\n",
    "h_bits = -sum(p_i * math.log2(p_i) for p_i in p)\n",
    "\n",
    "q = [1/4] * 4   # the wrong codebook, built as if the data were uniform\n",
    "lengths_q = [math.ceil(math.log2(1 / q_i)) for q_i in q]\n",
    "avg_len_q = sum(p_i * l for p_i, l in zip(p, lengths_q))\n",
    "kl_bits = sum(p_i * math.log2(p_i / q_i) for p_i, q_i in zip(p, q))\n",
    "\n",
    "print(f'Shannon lengths {lengths}, Kraft sum = {kraft}')\n",
    "print(f'E[l] = {avg_len} bits = H(P) = {h_bits} bits: no waste')\n",
    "print(f'uniform codebook: {avg_len_q} bits = H + KL = {h_bits + kl_bits} bits')"
   ],
   "execution_count": 12,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Shannon lengths [1, 2, 3, 3], Kraft sum = 1.0\nE[l] = 1.75 bits = H(P) = 1.75 bits: no waste\nuniform codebook: 2.0 bits = H + KL = 2.0 bits\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-13 fw=all -->\n",
    "\n",
    "A possible Shannon code here is $0,\\, 10,\\, 110,\\, 111$, exactly the tree of\n",
    "the figure; check that it is prefix-free and that its\n",
    "lengths $1, 2, 3, 3$ are the ones computed above.\n",
    "\n",
    "### From Symbol Codes to Arithmetic Coding\n",
    "\n",
    "The Shannon code is optimal only up to that \"$+1$\", and the leak is real: the\n",
    "ceiling in $l_i = \\lceil \\log_2(1/p_i) \\rceil$ rounds a fraction of a bit up\n",
    "to a whole bit, once per symbol. For a skewed source the rounding dominates\n",
    "the expected length. A symbol of probability $0.9$ carries only\n",
    "$-\\log_2 0.9 \\approx 0.15$ bits of surprisal, yet no binary codeword is\n",
    "shorter than one bit, so a symbol-by-symbol code overpays more than sixfold\n",
    "on the most common symbol it ever transmits.\n",
    "\n",
    "The fix is to stop coding symbols and code the *whole message*. *Arithmetic\n",
    "coding* [@Rissanen.1976; @Witten.Neal.Cleary.1987] runs the Kraft proof's\n",
    "dyadic-interval picture in reverse. Maintain a\n",
    "current interval, initially $[0, 1)$. To encode the next symbol, partition\n",
    "the current interval into one slice per possible symbol, with lengths\n",
    "proportional to the model's conditional probabilities, and shrink to the\n",
    "slice of the symbol actually observed. After the message $x_1, \\ldots, x_N$\n",
    "the interval has width\n",
    "\n",
    "$$\n",
    "w = \\prod_{i=1}^N q(x_i \\mid x_{<i}),\n",
    "$$\n",
    "\n",
    "and it identifies the message: a decoder replaying the same model recovers\n",
    "each symbol in turn by watching which slice the transmitted number falls\n",
    "into. All that remains is to *name* a point of the final interval in binary,\n",
    "and the dyadic intervals of the Kraft proof do it: any interval of width $w$\n",
    "contains a dyadic interval of length $2^{-l}$ with\n",
    "$l = \\lceil \\log_2(1/w) \\rceil + 1$, because $2^{-l} \\leq w/2$, so some\n",
    "multiple of $2^{-l}$ falls in the interval's left half and the dyadic\n",
    "interval of length $2^{-l}$ starting there fits inside; that dyadic\n",
    "interval's binary address is the codeword. The total length is therefore\n",
    "\n",
    "$$\n",
    "l \\;<\\; \\sum_{i=1}^N \\log_2 \\frac{1}{q(x_i \\mid x_{<i})} + 2,\n",
    "$$\n",
    "\n",
    "the message's total surprisal under the model plus at most two bits: the\n",
    "ceiling is paid once per *message*, not once per symbol. A ten-line\n",
    "demonstration (interval narrowing on a skewed source, then the two bit\n",
    "counts) makes the gap vivid."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mdl-information-theory-from-symbol-codes-to-arithmetic-coding"
   },
   "source": [
    "msg = 'aaaaaaaaba'                    # nine a's, one b\n",
    "q = {'a': 0.9, 'b': 0.1}              # a simple i.i.d. character model\n",
    "\n",
    "lo, hi = 0.0, 1.0\n",
    "for ch in msg:                        # narrow to the observed symbol's slice\n",
    "    width = hi - lo\n",
    "    if ch == 'a':                     # 'a' owns the left 90% of the interval\n",
    "        hi = lo + width * q['a']\n",
    "    else:                             # 'b' owns the right 10%\n",
    "        lo = lo + width * q['a']\n",
    "\n",
    "prod_q = math.prod(q[ch] for ch in msg)\n",
    "surprisal = -math.log2(prod_q)                     # the floor, in bits\n",
    "arith_bits = math.ceil(surprisal) + 1              # one interval, one ceiling\n",
    "shannon_bits = sum(math.ceil(math.log2(1 / q[ch])) for ch in msg)\n",
    "\n",
    "print(f'final width = {hi - lo:.6e}, prod of probs = {prod_q:.6e}')\n",
    "print(f'total surprisal = {surprisal:.2f} bits')\n",
    "print(f'arithmetic code: {arith_bits} bits, symbol-by-symbol Shannon code: '\n",
    "      f'{shannon_bits} bits')"
   ],
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "final width = 3.874205e-02, prod of probs = 3.874205e-02\ntotal surprisal = 4.69 bits\narithmetic code: 6 bits, symbol-by-symbol Shannon code: 13 bits\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-14 fw=all -->\n",
    "\n",
    "The final interval width equals the product of the per-symbol probabilities,\n",
    "as the construction promises, and the bookkeeping lands exactly where the\n",
    "theory says: the message's surprisal is $4.69$ bits, the arithmetic code\n",
    "spends $6$ bits (the floor plus the once-per-message rounding), and the\n",
    "symbol-by-symbol Shannon code spends $13$ bits, because it pays the ceiling's\n",
    "rounded-up whole bit on each of the nine cheap `a`'s. On long messages the\n",
    "two-bit overhead vanishes into the total and the arithmetic code sits *on*\n",
    "the surprisal floor.\n",
    "\n",
    "Notice what the construction never required: the model need not be i.i.d.\n",
    "Any conditional next-symbol distribution $q(\\cdot \\mid x_{<i})$ works,\n",
    "because encoder and decoder both condition on the already-processed prefix.\n",
    "An autoregressive language model therefore supplies the probability model for\n",
    "a lossless arithmetic coder. If encoder and decoder share the tokenizer,\n",
    "model, and finite-precision probability implementation, the ideal code length\n",
    "for a document is its total surprisal plus fewer than two bits. Its token\n",
    "cross-entropy times the token count estimates that model code length; a file\n",
    "format also carries framing and does not include the cost of distributing the\n",
    "model. @Deletang.Ruoss.Duquenne.ea.2023 drive an arithmetic coder with\n",
    "an LLM's next-token probabilities and compress text better than gzip. Thus\n",
    "minimizing cross-entropy improves the code length available to such a coder.\n",
    "\n",
    "How well can any compressor do on a real source such as English text? For a\n",
    "*stationary* source $X_1, X_2, \\ldots$ (one whose statistics do not drift\n",
    "with position), the *entropy rate*\n",
    "\n",
    "$$\n",
    "H_{\\textrm{rate}} = \\lim_{n \\to \\infty} \\frac{1}{n} H(X_1, \\ldots, X_n)\n",
    "= \\lim_{n \\to \\infty} H(X_n \\mid X_{n-1}, \\ldots, X_1)\n",
    "$$\n",
    "\n",
    "is the per-symbol generalization of entropy. That the two limits exist and\n",
    "agree for stationary sources is a theorem we grant [@Cover.Thomas.1999];\n",
    "the second limit is the average surprisal of the next symbol given unlimited\n",
    "context (conditional entropy is developed in\n",
    "that section), and $H_{\\textrm{rate}}$ is the floor\n",
    "on the expected code length per symbol of any lossless code, just as $H$ was\n",
    "for i.i.d. symbols. Dependence lowers the floor (conditioning reduces\n",
    "entropy), so a source with memory compresses better than its letter\n",
    "frequencies suggest: Shannon estimated the entropy rate of printed English at\n",
    "roughly one bit per character [@Shannon.1951], far below the\n",
    "$\\log_2 27 \\approx 4.75$ bits of uniformly random letters and spaces.\n",
    "\n",
    "### Typical Sequences and the Source-Coding Theorem\n",
    "\n",
    "The entropy bound above concerns an expected code length. A stronger statement\n",
    "explains what a **typical long message** looks like. Let\n",
    "$X_1,\\ldots,X_n$ be i.i.d. draws from a finite-alphabet distribution $P$. The\n",
    "probability of the realized block is\n",
    "$P(X_{1:n})=\\prod_i p(X_i)$, so its surprisal per symbol is a sample mean:\n",
    "\n",
    "$$\n",
    "-\\frac{1}{n}\\log P(X_{1:n})\n",
    "=\\frac{1}{n}\\sum_{i=1}^n -\\log p(X_i)\n",
    "\\xrightarrow{P} H(P).\n",
    "$$\n",
    "\n",
    "This is the **asymptotic equipartition property** (AEP): with probability\n",
    "approaching one, a long draw belongs to a **typical set** in which every block\n",
    "has probability approximately $e^{-nH}$. Because the typical set carries almost\n",
    "all the probability, it must contain approximately $e^{nH}$ blocks. The full\n",
    "space may have $k^n$ possible strings, but a source uses only an exponentially\n",
    "smaller effective subset when $H<\\log k$.\n",
    "\n",
    "This counting picture is the operational core of Shannon's source-coding\n",
    "theorem [@Shannon.1948; @Cover.Thomas.1999]. Encode typical blocks by their\n",
    "index using about $nH$ nats (or $nH_2$ bits), and handle the vanishing atypical\n",
    "set separately: rates just above entropy are achievable with failure\n",
    "probability tending to zero. Conversely, fewer than about $e^{nH}$ codewords\n",
    "cannot name almost all typical blocks, so rates below entropy cannot be\n",
    "reliable. For stationary ergodic sources the entropy rate\n",
    "the equation replaces $H$.\n",
    "\n",
    "For a Bernoulli source the counting claim can be checked *exactly*, because\n",
    "strings with the same number of ones share one probability. The cell groups\n",
    "the $2^{200}$ binary strings of length $n=200$ by their count of ones,\n",
    "declares a string typical when its surprisal is within $0.1$ bit per symbol\n",
    "of $H$, and adds up exactly how many strings are typical and how much\n",
    "probability they carry."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mdl-information-theory-typical-sequences-and-the-source-coding-theorem"
   },
   "source": [
    "import numpy as onp\n",
    "n, p = 200, 0.3\n",
    "k = onp.arange(n + 1)\n",
    "log_fact = onp.concatenate([[0.0], onp.cumsum(onp.log(onp.arange(1, n + 1)))])\n",
    "log_count = log_fact[n] - log_fact[k] - log_fact[n - k]     # log #strings, k ones\n",
    "prob_k = onp.exp(log_count + k * onp.log(p) + (n - k) * onp.log(1 - p))\n",
    "h2 = -(p * onp.log2(p) + (1 - p) * onp.log2(1 - p))\n",
    "bits_per_symbol = -(k * onp.log2(p) + (n - k) * onp.log2(1 - p)) / n\n",
    "typical = onp.abs(bits_per_symbol - h2) < 0.1\n",
    "\n",
    "m = log_count[typical].max()\n",
    "log2_typical = (m + onp.log(onp.exp(log_count[typical] - m).sum())) / onp.log(2)\n",
    "print(f'entropy H = {h2:.3f} bits/symbol, so nH = {n * h2:.0f} bits')\n",
    "print(f'typical strings: 2^{log2_typical:.1f} of all 2^{n}')\n",
    "print(f'probability they carry: {prob_k[typical].sum():.4f}')"
   ],
   "execution_count": 14,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "entropy H = 0.881 bits/symbol, so nH = 176 bits\ntypical strings: 2^188.8 of all 2^200\nprobability they carry: 0.9893\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-15 fw=all -->\n",
    "\n",
    "A vanishing fraction of the strings — about $2^{nH}$ of the $2^n$, missing\n",
    "some $24$ orders of magnitude of the count — soaks up almost all of the\n",
    "probability. Indexing just those strings *is* the compressor, and $nH$ bits\n",
    "is the length of the index.\n",
    "\n",
    "### Lossy Compression and Rate--Distortion\n",
    "\n",
    "Lossless recovery is often unnecessary. An image codec may change pixels\n",
    "slightly, and a learned representation may deliberately discard nuisance\n",
    "detail. Choose a distortion measure $d(x,\\hat x)$ and allow average distortion\n",
    "at most $D$. The **rate--distortion function** is\n",
    "\n",
    "$$\n",
    "R(D)=\\inf_{p(\\hat x\\mid x):\\;E[d(X,\\hat X)]\\le D} I(X;\\hat X).\n",
    "$$\n",
    "\n",
    "It is the smallest information rate a representation can retain while meeting\n",
    "the distortion budget. A stricter budget costs more bits; relaxing $D$ lowers\n",
    "the rate. For a Gaussian source $X\\sim\\mathcal N(0,\\sigma^2)$ under squared\n",
    "error,\n",
    "\n",
    "$$\n",
    "R(D)=\\frac12\\log\\frac{\\sigma^2}{D}\n",
    "\\quad (0<D<\\sigma^2),\n",
    "\\qquad R(D)=0\\quad(D\\ge\\sigma^2),\n",
    "$$\n",
    "\n",
    "in nats per sample. The left panel of the figure plots\n",
    "it: bits buy reconstruction quality at a rate that explodes as $D\\to0$ and\n",
    "drops to zero once the allowed distortion exceeds the source variance —\n",
    "at that point transmitting nothing and reconstructing the mean is already\n",
    "good enough. This curve is a useful benchmark for learned compression:\n",
    "reconstruction losses measure distortion, while quantized latents and entropy\n",
    "models determine rate. Objectives of the form\n",
    "$D+\\beta R$ select one supporting point of the trade-off — geometrically, the\n",
    "point where the curve has slope $-1/\\beta$. Autoencoders and\n",
    "variational bottlenecks use approximations to these two terms; they do not evade\n",
    "the trade-off by learning the codec.\n",
    "\n",
    "### Noisy Channels and Capacity\n",
    "\n",
    "Compression asks how few bits describe a source. Communication asks how many\n",
    "bits survive a noisy channel. A memoryless channel is a conditional distribution\n",
    "$p(y\\mid x)$. Its **capacity** is\n",
    "\n",
    "$$\n",
    "C=\\max_{p(x)} I(X;Y),\n",
    "$$\n",
    "\n",
    "where the maximization chooses the input distribution best matched to the\n",
    "channel. Shannon's noisy-channel coding theorem says that long block codes can\n",
    "make decoding error arbitrarily small at every rate below $C$, while rates\n",
    "above $C$ cannot be made reliable [@Shannon.1948; @Cover.Thomas.1999].\n",
    "For a binary symmetric channel that flips each bit independently with\n",
    "probability $\\varepsilon$, the uniform input is optimal and\n",
    "\n",
    "$$\n",
    "C=1-h_2(\\varepsilon)\\quad\\text{bits per channel use},\n",
    "$$\n",
    "\n",
    "where $h_2$ is binary entropy. The right panel of\n",
    "the figure plots this curve, and it rewards a moment's\n",
    "reading. At $\\varepsilon=0$ one bit survives per use. At $\\varepsilon=1/2$\n",
    "the output is independent of the input and capacity is zero — no code, however\n",
    "clever, gets anything through. And at $\\varepsilon=1$ capacity is back to one\n",
    "bit: a channel that *always* lies is as good as one that always tells the\n",
    "truth, because the receiver can simply flip every bit. What destroys\n",
    "information is not corruption but *unpredictability*.\n",
    "\n",
    "![Shannon's two operational curves. Left: the Gaussian rate–distortion function — the number of bits per sample needed to reconstruct within squared error $D$; fidelity gets expensive fast, and beyond the source variance no bits are needed at all. Right: the capacity of a binary symmetric channel with flip probability $\\varepsilon$; a coin-flip channel carries nothing, while a channel that always flips is again perfect.](../img/mdl-it-capacity-rd.svg)\n",
    "\n",
    "The same mutual-information lens applies when a representation passes through\n",
    "noise, quantization, dropout, or a bandwidth bottleneck. Data processing limits\n",
    "what later layers can recover, while coding or redundancy determines how close\n",
    "a designed system comes to the limit.\n",
    "\n",
    "### Perplexity\n",
    "\n",
    "Language models are evaluated by exactly the quantity this section has been\n",
    "studying: the per-token cross-entropy of held-out text,\n",
    "$\\textrm{CE} = -\\tfrac{1}{N}\\sum_{i=1}^N \\log q(x_i \\mid x_{<i})$ nats, where\n",
    "$q$ is the model's predicted next-token distribution. The community\n",
    "exponentiates it and calls the result *perplexity*\n",
    "[@Jelinek.Mercer.Bahl.ea.1977]:\n",
    "\n",
    "$$\n",
    "\\textrm{PPL} = \\exp(\\textrm{CE})\n",
    "= \\exp\\!\\Big(-\\frac{1}{N}\\sum_{i=1}^N \\log q(x_i \\mid x_{<i})\\Big)\n",
    "= \\Big(\\prod_{i=1}^N q(x_i \\mid x_{<i})\\Big)^{-1/N},\n",
    "$$\n",
    "\n",
    "the inverse *geometric mean* of the probabilities the model assigned to what\n",
    "actually came next. Exponentiating buys an interpretation: a model that is\n",
    "uniformly undecided among $V$ tokens at every step has\n",
    "$\\textrm{CE} = \\ln V$ and hence $\\textrm{PPL} = V$, so a perplexity of,\n",
    "say, $20$ means the model is, on average, as uncertain as if it were choosing\n",
    "uniformly among $20$ tokens: an *effective branching factor*. A perfect model\n",
    "of a deterministic sequence has perplexity $1$, and the floor is\n",
    "$\\textrm{PPL} \\geq e^{H}$, where $H$ is the entropy rate\n",
    "the equation of the language itself, the per-token cost of\n",
    "its irreducible randomness (Gibbs applied to each next-token prediction,\n",
    "averaged, plus the entropy-rate limit granted above). Because $\\textrm{PPL} = b^{\\textrm{CE}_b}$ for *any* log base $b$,\n",
    "perplexity is also the rare information-theoretic quantity with no unit\n",
    "ambiguity: nats and bits give the same number (Exercise 7).\n",
    "\n",
    "A five-token toy stream makes the bookkeeping concrete."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-perplexity"
   },
   "source": [
    "q_tok = [0.2, 0.1, 0.4, 0.25, 0.05]   # model prob. of each observed token\n",
    "nll = [-math.log(q_i) for q_i in q_tok]\n",
    "mean_nll = sum(nll) / len(nll)\n",
    "\n",
    "print(f'per-token NLL (nats): {[round(v, 3) for v in nll]}')\n",
    "print(f'mean NLL = {mean_nll:.4f} nats, perplexity = {math.exp(mean_nll):.2f}')"
   ],
   "execution_count": 15,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "per-token NLL (nats): [1.609, 2.303, 0.916, 1.386, 2.996]\nmean NLL = 1.8421 nats, perplexity = 6.31\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-16 fw=all -->\n",
    "\n",
    "A mean of $\\approx 1.84$ nats per token corresponds to a perplexity of\n",
    "$\\approx 6.3$: this model is as confused as a uniform choice among six or so\n",
    "tokens. Note how the geometric mean punishes confident mistakes: the single\n",
    "$0.05$ contributes $3.0$ nats, as much as several good predictions\n",
    "combined. When you encounter perplexity as the headline metric in\n",
    "that section and the Transformer chapters, it is this number.\n",
    "\n",
    "Perplexity is also the quantity behind the defining empirical fact of the\n",
    "LLM era: held-out cross-entropy falls as a *power law* in model parameters,\n",
    "dataset size, and training compute, smoothly across many orders of magnitude\n",
    "[@kaplan2020scaling], with exponents stable enough that one can budget\n",
    "compute against a target loss in advance [@hoffmann2022training]. Read\n",
    "through this section's lens, the scaling laws say that the approach to the\n",
    "language's entropy rate (the floor the equation that no\n",
    "model can beat) is empirically\n",
    "*predictable*: each constant multiple of compute removes a roughly constant\n",
    "fraction of the remaining excess over the floor. One caution when comparing\n",
    "reported numbers: per-token quantities depend on what a token is, so\n",
    "perplexities are comparable only between models that share a tokenizer.\n",
    "To compare across tokenizers, convert to bits per character, the\n",
    "compression rate of the previous subsection, which is tokenizer-free.\n",
    "\n",
    "## Modern Uses\n",
    "\n",
    "Two modern training techniques follow directly from these identities; both\n",
    "are corollaries of Gibbs' inequality.\n",
    "\n",
    "### Learning by Compression: Minimum Description Length\n",
    "\n",
    "Likelihood measures the bits needed to encode data **after** a model is shared.\n",
    "Model selection must also account for sharing the model. The minimum description\n",
    "length (MDL) principle chooses the hypothesis that gives the shortest complete\n",
    "code,\n",
    "\n",
    "$$\n",
    "\\hat h_{\\textrm{MDL}}\n",
    "=\\mathop{\\mathrm{argmin}}_h\n",
    "\\bigl[L(h)+L(\\mathcal D\\mid h)\\bigr].\n",
    "$$\n",
    "\n",
    "The second term is usually a negative log-likelihood; the first pays for model\n",
    "complexity. A highly flexible model may compress the training residuals while\n",
    "losing overall because its parameters or structure are expensive to describe.\n",
    "This is Occam's razor in operational units rather than a parameter-count slogan.\n",
    "\n",
    "There is a close Bayesian connection. Choosing\n",
    "$L(h)=-\\log P(h)$ and $L(\\mathcal D\\mid h)=-\\log P(\\mathcal D\\mid h)$ gives the\n",
    "MAP objective. A one-part Bayesian code instead uses the marginal likelihood\n",
    "$-\\log\\int P(\\mathcal D\\mid\\theta)P(\\theta)\\,d\\theta$, averaging rather than\n",
    "plugging in one parameter. For real-valued neural weights, a literal two-part\n",
    "code also needs a quantization precision and a concrete coder; counting\n",
    "parameters alone is not a description length. MDL is therefore a disciplined\n",
    "way to compare complete predictive codes, not an automatic theorem that every\n",
    "compressed network generalizes.\n",
    "\n",
    "### Label Smoothing\n",
    "\n",
    "A classifier trained on one-hot targets is asked to assign probability $1$ to\n",
    "the true class, which a softmax can approach only by driving the true class's\n",
    "logit infinitely far above the rest. The result is overconfident models and\n",
    "ever-growing weights. *Label smoothing* [@Szegedy.Vanhoucke.Ioffe.ea.2016]\n",
    "replaces the one-hot target $\\mathbf{e}_y$ with a mixture that reserves a\n",
    "small probability $\\epsilon$ for the other classes:\n",
    "\n",
    "$$\n",
    "\\mathbf{p}^\\epsilon = (1 - \\epsilon)\\, \\mathbf{e}_y + \\frac{\\epsilon}{k}\\, \\mathbf{1},\n",
    "$$\n",
    "\n",
    "i.e., target probability $1 - \\epsilon + \\epsilon/k$ for the true class and\n",
    "$\\epsilon/k$ for each of the other $k - 1$. The training loss is the ordinary\n",
    "cross-entropy $\\textrm{CE}(\\mathbf{p}^\\epsilon, \\mathbf{q})$ to this softened\n",
    "target. What prediction does the smoothed loss actually ask for?\n",
    "\n",
    "**Proposition (the optimal smoothed prediction).** *Over all probability\n",
    "vectors $\\mathbf{q}$, the cross-entropy\n",
    "$\\textrm{CE}(\\mathbf{p}^\\epsilon, \\mathbf{q})$ is minimized uniquely at\n",
    "$\\mathbf{q}^* = \\mathbf{p}^\\epsilon$, where it equals $H(\\mathbf{p}^\\epsilon)$.*\n",
    "\n",
    "**Proof.** By the equation,\n",
    "$\\textrm{CE}(\\mathbf{p}^\\epsilon, \\mathbf{q}) = H(\\mathbf{p}^\\epsilon) +\n",
    "D_{\\textrm{KL}}(\\mathbf{p}^\\epsilon \\| \\mathbf{q})$. The first term does not\n",
    "depend on $\\mathbf{q}$, and by Gibbs' inequality the second is non-negative\n",
    "and zero iff $\\mathbf{q} = \\mathbf{p}^\\epsilon$. $\\blacksquare$\n",
    "\n",
    "So the optimum is no longer at infinity: the model is asked to predict\n",
    "$1 - \\epsilon + \\epsilon/k$ on the true class, which a softmax reaches with the\n",
    "*finite* logit gap\n",
    "\n",
    "$$\n",
    "z_y - z_j = \\log\\frac{1 - \\epsilon + \\epsilon/k}{\\epsilon/k}\n",
    "= \\log\\!\\Big(1 + \\frac{k(1-\\epsilon)}{\\epsilon}\\Big)\n",
    "$$\n",
    "\n",
    "between the true class and each other class. For $k = 10$ classes and\n",
    "$\\epsilon = 0.1$ the gap is $\\ln 91 \\approx 4.51$, a concrete, attainable\n",
    "target instead of a runaway one. (A logit difference is a log-odds ratio, so\n",
    "we leave it unitless rather than calling it nats.) The loss at the optimum\n",
    "is the (nonzero) entropy $H(\\mathbf{p}^\\epsilon)$, which is why a\n",
    "label-smoothed loss curve plateaus above zero even when the model is doing\n",
    "everything right. Both numbers are easy to check:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-label-smoothing"
   },
   "source": [
    "k, eps = 10, 0.1\n",
    "p_eps = [1 - eps + eps / k] + [eps / k] * (k - 1)\n",
    "\n",
    "gap = math.log(p_eps[0] / p_eps[1])              # optimal logit gap\n",
    "h_p_eps = -sum(t * math.log(t) for t in p_eps)   # CE at the optimum q* = p_eps\n",
    "\n",
    "z = [10.0] + [0.0] * (k - 1)                     # an overconfident prediction\n",
    "z_norm = sum(math.exp(v) for v in z)\n",
    "q_conf = [math.exp(v) / z_norm for v in z]\n",
    "ce_conf = -sum(t * math.log(s) for t, s in zip(p_eps, q_conf))\n",
    "\n",
    "print(f'optimal logit gap = ln(91) = {gap:.4f}')\n",
    "print(f'CE at q* = H(p_eps) = {h_p_eps:.4f} nats')\n",
    "print(f'CE at the overconfident q = {ce_conf:.4f} nats')"
   ],
   "execution_count": 16,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "optimal logit gap = ln(91) = 4.5109\nCE at q* = H(p_eps) = 0.5003 nats\nCE at the overconfident q = 0.9004 nats\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-17 fw=all -->\n",
    "\n",
    "The overconfident prediction (logit gap $10$) is *worse* under the smoothed\n",
    "target than the calibrated one: label smoothing penalizes the very\n",
    "behavior that one-hot targets demand.\n",
    "\n",
    "### Knowledge Distillation\n",
    "\n",
    "*Knowledge distillation* [@Hinton.Vinyals.Dean.2015] trains a small\n",
    "*student* network to imitate a large *teacher*, transferring the teacher's\n",
    "full distribution over wrong answers rather than its argmax alone (a \"2\" that\n",
    "looks somewhat like a \"7\" is valuable information). Because a trained\n",
    "teacher's softmax is nearly one-hot, both distributions are first *softened*\n",
    "with a temperature $T$:\n",
    "\n",
    "$$\n",
    "p_j^{(T)} = \\mathrm{softmax}(\\mathbf{z}^{\\textrm{tea}}/T)_j, \\qquad\n",
    "q_j^{(T)} = \\mathrm{softmax}(\\mathbf{z}^{\\textrm{stu}}/T)_j,\n",
    "$$\n",
    "\n",
    "and the distillation loss is the KL divergence between the softened\n",
    "distributions, scaled by $T^2$:\n",
    "\n",
    "$$\n",
    "\\mathcal{L}_{\\textrm{distill}}\n",
    "= T^2\\, D_{\\textrm{KL}}\\big(\\mathbf{p}^{(T)} \\,\\|\\, \\mathbf{q}^{(T)}\\big).\n",
    "$$\n",
    "\n",
    "Where does the $T^2$ come from? Differentiate the KL term with respect to a\n",
    "student logit. Since only $-\\sum_j p_j^{(T)} \\log q_j^{(T)}$ depends on the\n",
    "student, and the derivative of a log-softmax at temperature $T$ is the usual\n",
    "softmax-minus-target expression scaled by $1/T$,\n",
    "\n",
    "$$\n",
    "\\frac{\\partial}{\\partial z^{\\textrm{stu}}_j}\n",
    "D_{\\textrm{KL}}\\big(\\mathbf{p}^{(T)} \\| \\mathbf{q}^{(T)}\\big)\n",
    "= \\frac{1}{T}\\big(q_j^{(T)} - p_j^{(T)}\\big).\n",
    "$$\n",
    "\n",
    "For large $T$ the softened distributions flatten toward uniform,\n",
    "$\\mathrm{softmax}(\\mathbf{z}/T)_j \\approx \\tfrac{1}{k}\\big(1 + (z_j - \\bar z)/T\\big)$\n",
    "with $\\bar z$ the mean logit, so the difference $q_j^{(T)} - p_j^{(T)}$ itself\n",
    "shrinks like $1/T$ and the gradient like $1/T^2$:\n",
    "\n",
    "$$\n",
    "\\frac{1}{T}\\big(q_j^{(T)} - p_j^{(T)}\\big)\n",
    "\\;\\approx\\; \\frac{1}{k\\,T^2}\\Big(\\big(z^{\\textrm{stu}}_j - \\bar z^{\\textrm{stu}}\\big)\n",
    "- \\big(z^{\\textrm{tea}}_j - \\bar z^{\\textrm{tea}}\\big)\\Big).\n",
    "$$\n",
    "\n",
    "Multiplying the loss by $T^2$ therefore keeps the gradient magnitude roughly\n",
    "constant as the temperature is tuned, so $T$ can be chosen for the *quality*\n",
    "of the transferred soft targets without silently rescaling the learning rate\n",
    "against any hard-label loss it is mixed with. Let's verify both claims\n",
    "numerically: the closed-form gradient matches autograd, and $T^2$ times the\n",
    "gradient approaches the limit above instead of vanishing."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "information-theory-distillation"
   },
   "source": [
    "z_tea = np.array([5.0, 2.0, -1.0])\n",
    "z_stu = np.array([3.0, 3.0, 0.0])\n",
    "z_stu.attach_grad()\n",
    "\n",
    "for T in (1.0, 2.0, 5.0, 10.0):\n",
    "    with autograd.record():\n",
    "        p, q = npx.softmax(z_tea / T), npx.softmax(z_stu / T)\n",
    "        kl = np.sum(p * (np.log(p) - np.log(q)))\n",
    "    kl.backward()\n",
    "    closed = (npx.softmax(z_stu / T) - npx.softmax(z_tea / T)) / T\n",
    "    err = float(np.abs(z_stu.grad - closed).max())\n",
    "    print(f'T={T:4.1f}  |autograd - closed| = {err:.1e}  '\n",
    "          f'T^2 grad = {[round(g, 3) for g in (T**2 * z_stu.grad).tolist()]}')"
   ],
   "execution_count": 17,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "T= 1.0  |autograd - closed| = 0.0e+00  T^2 grad = [-0.462, 0.441, 0.022]\nT= 2.0  |autograd - closed| = 3.0e-08  T^2 grad = [-0.672, 0.549, 0.123]\nT= 5.0  |autograd - closed| = 2.8e-09  T^2 grad = [-0.741, 0.478, 0.263]\nT=10.0  |autograd - closed| = 9.3e-10  T^2 grad = [-0.719, 0.413, 0.306]\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-18 fw=all -->\n",
    "\n",
    "The closed form matches autograd to floating-point precision at every\n",
    "temperature. The raw gradient decays like $1/T^2$, while the $T^2$-scaled\n",
    "gradient stays order-one: its $T \\to \\infty$ limit is\n",
    "$\\tfrac{1}{k}\\big((\\mathbf{z}^{\\textrm{stu}} - \\bar z^{\\textrm{stu}}) -\n",
    "(\\mathbf{z}^{\\textrm{tea}} - \\bar z^{\\textrm{tea}})\\big) = (-0.667, 0.333,\n",
    "0.333)$ here, but the convergence is $O(1/T)$ and still visibly incomplete in\n",
    "the table: the $T = 10$ row reads $(-0.719, 0.413, 0.306)$, and the first\n",
    "component even overshoots its limit on the way. The scale-matching, which is\n",
    "what the $T^2$ factor was designed for, holds regardless. As\n",
    "$T \\to 1$ the loss reduces to the ordinary KL (and, with a one-hot teacher,\n",
    "to the ordinary cross-entropy loss), recovering standard training as a special\n",
    "case.\n",
    "\n",
    "### One Principle, Many Losses\n",
    "\n",
    "Stepping back: maximum likelihood *is* cross-entropy minimization *is*\n",
    "KL-projection of the empirical distribution onto the model family\n",
    "(that section), and this section has now equipped\n",
    "that one principle with its operational meaning (code length and waste), its\n",
    "evaluation metric (perplexity), and two of its modern refinements (smoothed\n",
    "targets and distilled teachers). The story continues in two directions: KL is\n",
    "just one member of a whole family of divergences, each inducing a different\n",
    "generative-modeling objective (that section), and\n",
    "applying KL to a joint distribution versus the product of its marginals yields\n",
    "mutual information, the quantity behind contrastive representation learning\n",
    "(that section).\n",
    "\n",
    "## Summary\n",
    "\n",
    "* Self-information $I(x) = -\\log p(x)$ measures the surprise of an outcome;\n",
    "  entropy $H(P) = -E_{x\\sim P}[\\log p(x)]$ is the average surprise of a\n",
    "  distribution. We measure both in nats (natural log); bits differ by a\n",
    "  factor of $\\ln 2$.\n",
    "* On $k$ outcomes, $0 \\leq H(P) \\leq \\log k$, with the maximum exactly at the\n",
    "  uniform distribution (Jensen). Differential entropy $h$ of continuous\n",
    "  variables is *not* coordinate-invariant and can be negative, one reason deep\n",
    "  learning prefers relative quantities; under fixed mean and variance it is\n",
    "  maximized by the Gaussian (that section).\n",
    "* The KL divergence $D_{\\textrm{KL}}(P\\|Q) = E_{x\\sim P}[\\log p(x)/q(x)]$ is\n",
    "  asymmetric and, by Gibbs' inequality, non-negative with equality iff\n",
    "  $P = Q$.\n",
    "* Cross-entropy decomposes as\n",
    "  $\\textrm{CE}(P, Q) = H(P) + D_{\\textrm{KL}}(P\\|Q)$. Minimizing it in $Q$ is\n",
    "  minimizing KL, and on empirical data it is maximum likelihood\n",
    "  (that section).\n",
    "* As a scoring rule, the cross-entropy loss (the log score) is *strictly\n",
    "  proper*: at population risk, a model class containing the truth is uniquely\n",
    "  minimized by the true conditional probabilities. Finite data, restricted\n",
    "  models, and distribution shift can still produce miscalibration.\n",
    "* In the coding view, entropy is the optimal expected code length (Kraft\n",
    "  inequality + Shannon code) and KL is the extra bits per symbol from coding\n",
    "  with the wrong distribution.\n",
    "* Arithmetic coding removes the symbol code's per-symbol rounding: with a\n",
    "  shared model and coder, a whole message costs its total surprisal plus at\n",
    "  most two bits. An autoregressive model supplies the probabilities for that\n",
    "  lossless code; its cross-entropy times the token count is the model code\n",
    "  length before file-format overhead. The entropy rate is the per-symbol floor\n",
    "  for any such code. The AEP explains this floor by showing that almost all\n",
    "  length-$n$ messages occupy a typical set of size about $e^{nH}$.\n",
    "* **Rate--distortion** gives the smallest retained information compatible with\n",
    "  a distortion budget, while **channel capacity** $C=\\max_{p(x)}I(X;Y)$ gives\n",
    "  the largest reliable communication rate through noise.\n",
    "* **Minimum description length** selects a complete code for model plus data;\n",
    "  negative log-likelihood is only the data-given-model part.\n",
    "* Perplexity $\\textrm{PPL} = \\exp(\\textrm{CE})$ is the exponentiated\n",
    "  per-token cross-entropy of a language model: an effective branching factor,\n",
    "  independent of the log base. Empirically it falls as a power law in\n",
    "  parameters, data, and compute, and it is comparable only under a shared\n",
    "  tokenizer.\n",
    "* Label smoothing and knowledge distillation are corollaries of the same\n",
    "  identities: the smoothed target makes the cross-entropy optimum finite and\n",
    "  equal to the deliberately softened target; it does not by itself guarantee\n",
    "  calibration to the original data distribution, and the $T^2$ factor on the distillation KL keeps\n",
    "  gradients scale-matched across temperatures.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Verify the card-deck numbers from the first subsection: the four reports\n",
    "   carry $0$, $\\ln 4$, $\\ln 52$, and $\\ln 52!$ nats. Use Stirling's\n",
    "   approximation to estimate $\\ln 52!$ and check it against the exact value.\n",
    "1. Give a second proof that the uniform distribution maximizes entropy on $k$\n",
    "   outcomes, this time from Gibbs' inequality: show that\n",
    "   $D_{\\textrm{KL}}(P \\,\\|\\, U) = \\log k - H(P)$ for the uniform $U$, and\n",
    "   conclude.\n",
    "1. Let's compute the entropy of a few data sources:\n",
    "    * Suppose that you are watching the output generated by a monkey at a\n",
    "      typewriter that hits any of the $44$ keys uniformly at random. How many\n",
    "      nats (and bits) of randomness per character do you observe?\n",
    "    * Unhappy with the monkey, you replace it with a drunk typesetter that\n",
    "      picks a random word out of a vocabulary of $2{,}000$ words, with an\n",
    "      average word length of $4.5$ letters. How many nats per character now?\n",
    "    * Still unhappy, you replace the typesetter with a language model that\n",
    "      achieves a per-word perplexity of $15$. How many nats per character\n",
    "      does that correspond to?\n",
    "1. Derive the closed form the equation for the KL divergence\n",
    "   between two univariate Gaussians by writing\n",
    "   $E_{x\\sim P}[\\log p(x) - \\log q(x)]$ and using\n",
    "   $E_{x \\sim P}[x] = \\mu_1$, $E_{x\\sim P}[(x-\\mu_1)^2] = \\sigma_1^2$.\n",
    "1. Show that self-information is essentially unique: if a continuous,\n",
    "   decreasing function $I(p)$ on $(0, 1]$ satisfies $I(1) = 0$ and\n",
    "   $I(p_1 p_2) = I(p_1) + I(p_2)$ for all $p_1, p_2$, then\n",
    "   $I(p) = -c \\log p$ for some constant $c > 0$.\n",
    "1. A scoring rule $S(\\mathbf{q}, y)$ is *strictly proper* if the expected\n",
    "   penalty $E_{y \\sim \\mathbf{p}}[S(\\mathbf{q}, y)]$ is uniquely minimized at\n",
    "   $\\mathbf{q} = \\mathbf{p}$. Prove from Gibbs' inequality that the log score\n",
    "   $S(\\mathbf{q}, y) = -\\log q_y$ is strictly proper. Then prove that the\n",
    "   Brier score $S(\\mathbf{q}, y) = \\sum_j (q_j - \\mathbf{1}[j = y])^2$ is\n",
    "   strictly proper too, by showing that its expected penalty exceeds the\n",
    "   truthful reporter's by exactly $\\|\\mathbf{q} - \\mathbf{p}\\|^2$\n",
    "   [@Gneiting.Raftery.2007].\n",
    "1. Show that perplexity does not depend on the log base: if\n",
    "   $\\textrm{CE}_b$ denotes the per-token cross-entropy in base-$b$ units,\n",
    "   then $b^{\\textrm{CE}_b}$ is the same for every $b$. Then compute the\n",
    "   perplexity of a model whose per-token cross-entropy is $\\ln 5$ nats.\n",
    "1. For label smoothing with $k$ classes and smoothing parameter $\\epsilon$,\n",
    "   derive the optimal prediction $\\mathbf{q}^*$ and the optimal logit gap,\n",
    "   and evaluate the gap for $k = 1000$, $\\epsilon = 0.1$. What happens to the\n",
    "   gap as $\\epsilon \\to 0$, and why does that recover the one-hot pathology?\n",
    "1. For the distillation loss, derive the gradient formula\n",
    "   $\\partial D_{\\textrm{KL}}(\\mathbf{p}^{(T)} \\| \\mathbf{q}^{(T)}) / \\partial\n",
    "   z^{\\textrm{stu}}_j = (q_j^{(T)} - p_j^{(T)})/T$, and show that as\n",
    "   $T \\to \\infty$ the $T^2$-scaled gradient tends to\n",
    "   $\\tfrac{1}{k}\\big((z^{\\textrm{stu}}_j - \\bar z^{\\textrm{stu}}) -\n",
    "   (z^{\\textrm{tea}}_j - \\bar z^{\\textrm{tea}})\\big)$. What does the loss\n",
    "   reduce to at $T = 1$ with a one-hot teacher?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-information-theory-md-19 fw=mxnet -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/420)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-mxnet",
   "display_name": "d2l (MXNet)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
