Модуль 3. Кросс-энтропия и KL-дивергенция

Пусть

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

1. Несовпадение поддержек

Кросс-энтропия в битах равна

Соответственно,

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

потому что $P$ допускает второй исход, а $Q$ объявляет его невозможным.

2. $\varepsilon$-сглаживание

Для

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

получаем

Энтропия $P$ равна

Поэтому

3. Асимптотика при $\varepsilon\to0^+$

Поскольку

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

главный расходящийся член равен

Для KL коэффициент тот же:

Даже редкий исход с вероятностью $0.1$ создаёт неограниченный штраф, если модель всё сильнее приближает его вероятность к нулю.

4. Что делает ограничение вероятностей (clipping)

Замена

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

создаёт некоторое новое сглаженное распределение $\widetilde Q_\varepsilon$ (для показанного кода это не буквально $(1-\varepsilon,\varepsilon)$). Это может быть разумной защитой численной реализации, но функция потерь после такого ограничения — уже логарифмическая потеря сглаженной модели, а не исходного $Q$. Если ограничение применяется только внутри логарифма без повторной нормировки, объект вообще может перестать быть корректным вероятностным распределением.

На практике лучше вычислять логарифмы вероятностей из логитов через устойчивый log_softmax: конечные логиты автоматически дают строго положительные вероятности softmax и не требуют ручного ограничения.

Пусть

Используем

Тогда

Дифференцируем по $z_m$:

Следовательно,

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

Для целевого one-hot-вектора $r=e_y$:

$\nabla_z\ell=q-e_y.$

Матрица Гессе (Hessian)

Из стандартной производной softmax

получаем

Для любого $v\in\mathbb R^K$

Значит, матрица Гессе положительно полуопределена, и кросс-энтропия выпукла по логитам.

Нулевое направление — добавление одной и той же константы ко всем логитам: softmax не меняется при $z\mapsto z+c\mathbf1$. Но логиты нелинейно зависят от весов глубокой сети, поэтому композиция функции потерь с сетью в общем случае невыпукла по параметрам.

1. Разложение условной кросс-энтропии

Начнём с

Добавим и вычтем $\log p(Y\mid C)$:

Первое слагаемое — условная энтропия:

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

Во втором сначала зафиксируем $C=c$:

После усреднения по $C$:

2. Цепное правило для авторегрессионной KL-дивергенции

Факторизуем оба распределения:

Тогда

Усредняем по $P$ и меняем местами конечную сумму и ожидание:

Последнее равенство — закон полного ожидания: при фиксированном префиксе внутреннее ожидание берётся по истинному распределению следующего токена.

3. Почему loss с teacher forcing не описывает всё поведение при генерации

При teacher forcing префиксы распределены как $P(X_{<t})$: это реальные префиксы из набора данных. При свободной генерации распределение префиксов определяется уже $Q$. Раннее отклонение меняет последующие контексты, и модель может попасть в область, редко представленную при валидации с teacher forcing. Поэтому хороший NLL на отложенных данных необходим для качественной вероятностной модели на распределении данных, но не является полной оценкой свободной авторегрессионной генерации.

Пусть

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

а

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

1. Аналитический оптимум прямой KL

При фиксированном $P$ минимизация $D_{\mathrm{KL}}(P\|Q)$ эквивалентна минимизации кросс-энтропии

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

Для гауссиана

Следовательно, с точностью до константы нужно минимизировать

По $\mu$ минимум достигается при

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

После этого

Для симметричной смеси

Итак,

2. Обратная KL через квадратуру Гаусса—Эрмита

Для обратной KL удобно выполнить репараметризацию

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

и считать

Ниже используется roots_hermitenorm, чьи узлы и веса непосредственно подходят для интегралов с весом $e^{-z^2/2}$.

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}"
    )

Типичный вывод:

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

Последние цифры центрального решения слегка меняются с порядком квадратуры, потому что это численный интеграл. Симметричные минимумы около $\pm2.9843$ устойчивы и имеют меньшее значение целевой функции.

3. Визуализация

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. Что именно показывает пример

Прямая KL в выбранном гауссовском семействе совпадает с целевым распределением по среднему и дисперсии и потому «накрывает» обе моды. Обратная KL имеет два более выгодных решения возле отдельных мод и менее выгодное широкое центральное решение. Это сильная иллюстрация асимметрии, но не универсальный закон о каждой модели и каждом оптимизаторе.

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))

Для токенной перплексии нужно экспоненцировать средний обычный NLL по целевым токенам, измеренный в натах:

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

В среднее должны входить ровно нужные немаскированные целевые токены, каждый с обычным весом. Функция потерь со сглаживанием меток (label smoothing), весами классов, произвольным перевзвешиванием примеров или другой нормировкой может быть полезна для обучения, но её экспонента не является стандартной перплексией.