Подгонка распределений с помощью KL

В заметке об энтропии мы вывели KL-дивергенцию (KL divergence) из неожиданности и показали, что минимизация $D_{\mathrm{KL}}(P\,\|\,Q)$ по модели — то же самое, что максимизация правдоподобия. В этой заметке мы применяем её к гауссовской модели: градиентный спуск притягивает $Q$ с обучаемыми параметрами к целевому распределению $P$. Заметка построена вокруг трёх вопросов:

1. Как обучаются параметры? Что на самом деле происходит за одну итерацию цикла подгонки?
2. Что предполагает KL? Что должно выполняться, чтобы формула имела смысл, и где каждое допущение проявляется в коде?
3. Что, если модель не может представить целевое распределение? Одна гауссиана, подогнанная к двум модам, показывает, чем различаются KL, обратная KL (reverse KL), дивергенция Йенсена–Шеннона (Jensen–Shannon divergence, JS) и функция потерь наименьших квадратов (least-squares loss).

Обозначения. $P$ — истинное (целевое) распределение с плотностью $p$, $Q = Q_\theta$ — модель с плотностью $q$ и параметрами $\theta$, $\eta$ — шаг обучения (learning rate). В одномерном случае $\theta = (\mu, \log\sigma)$: $\sigma$ хранится через свой логарифм, чтобы $\sigma = e^{\log\sigma}$ оставалось положительным. В двумерном случае среднее равно $\boldsymbol\mu$, а ковариация — $\boldsymbol\Sigma = \mathbf L\mathbf L^\top$, где $\mathbf L$ — нижнетреугольный множитель Холецкого (Cholesky factor) из заметки о гауссианах.

Как обучаются параметры#

Одна итерация цикла, шаг за шагом:

• Оптимизатор получает список обучаемых тензоров. Каждый из них создан с requires_grad=True, чтобы backward() заполнял его .grad. В одномерном случае это среднее и $\log\sigma$; в двумерном — среднее и три исходных элемента множителя Холецкого из заметки о гауссианах. У параметров целевого распределения нет градиента, и они остаются фиксированными.
• Q заново строится из текущих параметров на каждой итерации: $Q = \mathcal N(\mu, \sigma^2)$ с $\sigma = e^{\log\sigma}$.
• qx содержит плотность модели в точках выборки, $q(x_i)$; px содержит $p(x_i)$ — плотность истинного распределения в тех же точках.
• loss — выборочная оценка KL-дивергенции, построенная по px и qx.
• loss.backward() вычисляет градиенты $\partial\,\text{loss}/\partial\mu$ и $\partial\,\text{loss}/\partial\log\sigma = \sigma\,\partial\,\text{loss}/\partial\sigma$ методом обратного распространения ошибки и сохраняет их в .grad.
• optimizer.step() сдвигает параметры против градиента. Для обычного SGD это

$$\begin{gathered} \mu \leftarrow \mu - \eta\,\frac{\partial\,\text{loss}}{\partial \mu}, \\ \log\sigma \leftarrow \log\sigma - \eta\,\frac{\partial\,\text{loss}}{\partial \log\sigma}. \end{gathered}$$

Adam (Kingma and Ba, 2015) движется в похожем направлении, но масштабирует шаг по каждой координате скользящими оценками среднего и величины градиента. Подгонки по KL ниже выполняются с помощью Adam.

Рисунок

Одна итерация цикла подгонки

Одномерная подгонка в коде:

# вспомогательные функции: kl_loss (оценка ниже), get_parameters_1d (mu = 0, log_std = 0),
# create_distr_1d и target_1d — одномерная смесь с двумя модами, определённая ниже
train_mu, train_log_std = get_parameters_1d()   # requires_grad=True
optimizer = torch.optim.Adam([train_mu, train_log_std], lr=0.01)

x  = target_1d.sample((1000,))      # фиксированные сэмплы x_i ~ P
px = target_1d.log_prob(x).exp()    # p(x_i), известна для игрушечной цели

for i in range(2000):
    optimizer.zero_grad()           # очищаем .grad от предыдущего шага
    Q  = create_distr_1d(train_mu, train_log_std)   # Normal(mu, exp(log_std))
    qx = Q.log_prob(x).exp()        # q(x_i); может обратиться в 0, см. «В коде»
    loss = kl_loss(px, qx)          # mean(log2 p(x_i) - log2 q(x_i))
    loss.backward()                 # d loss / d mu, d loss / d log_std
    optimizer.step()                # шаг против градиента

Функция потерь — оценка KL методом Монте-Карло#

KL — математическое ожидание по $P$, поэтому её можно оценить усреднением по сэмплам $x_1, \dots, x_N \sim P$:

$$\begin{aligned} &D_{\mathrm{KL}}(P\,\|\,Q) = \int p(x)\log\frac{p(x)}{q(x)}\,dx \\ &\quad= \mathbb E_{x\sim P}\big[\log p(x) - \log q(x)\big] \\ &\quad\approx \frac1N\sum_{i=1}^{N}\big[\log p(x_i) - \log q(x_i)\big]. \end{aligned}$$

По закону больших чисел с ростом $N$ среднее сходится к математическому ожиданию. Код выше усредняет $\log_2 p(x_i) - \log_2 q(x_i)$, так что его функция потерь измеряется в битах; смена основания лишь умножает функцию потерь и её градиенты на константу $1/\ln 2 \approx 1.443$.

Рисунок

Выборочная оценка KL(P‖Q), используемая в ноутбуке

Сэмплы из другого распределения. Если сэмплы взяты не из $P$, а из распределения $R$ с плотностью $r$, каждое слагаемое перевзвешивается множителем $p/r$ (выборка по значимости (importance sampling)):

$$\begin{aligned} D_{\mathrm{KL}}(P\,\|\,Q) &\approx \frac1N\sum_{i=1}^{N} w_i\big[\log p(x_i) - \log q(x_i)\big], \\ w_i &= \frac{p(x_i)}{r(x_i)}, \qquad x_i \sim R , \end{aligned}$$

что требует $r > 0$ всюду, где $p > 0$. При $R = P$ все веса равны $1$, и мы возвращаемся к обычной оценке.

Оценка может быть отрицательной. KL никогда не бывает отрицательной, но её оценка на конечной выборке может. При подгонке двумерной гауссианы $Q$ с полной ковариацией (пять параметров) к цели $\mathcal N\big((-5,-5), \mathbf I\big)$ с помощью Adam на 1000 фиксированных сэмплах функция потерь после 4800 шагов равна $-0.0020$ бита. Причина в том, что $Q$ подгоняется на тех же сэмплах, на которых вычисляется функция потерь. Сама цель входит в число гауссиан-кандидатов, поэтому в точке максимального правдоподобия среднее $\log q$ по сэмплам не меньше среднего $\log p$, и оценка $\le 0$. Насколько она опускается ниже нуля, следует из теоремы Уилкса (Wilks’ theorem): удвоенный прирост суммарного логарифма правдоподобия подогнанной модели по сравнению с истинными параметрами приближённо имеет распределение $\chi^2$ с $k$ степенями свободы, где $k$ — число подгоняемых параметров (Wilks, 1938), так что в среднем оценка лежит на $k/(2N)$ ната ниже нуля. Здесь это $5/2000 = 0.0025$ ната, или $0.0036$ бита. В симуляции 2000 таких подгонок оценка в среднем равна $-0.0037$ бита со стандартным отклонением $0.0024$ бита, так что $-0.0020$ — типичное значение. Честная оценка дивергенции подогнанной модели требует новых сэмплов — так же, как любой модели нужны отложенные данные (held-out data).

Чему равны градиенты#

Для одномерной гауссианы $\log q(x) = -\log\sigma - \tfrac12\log 2\pi - \frac{(x-\mu)^2}{2\sigma^2}$. Слагаемое $\log p(x_i)$ не зависит от параметров, поэтому вклад даёт только $-\log q$ (натуральные логарифмы):

$$\begin{gathered} \frac{\partial\,\text{loss}}{\partial \mu} = -\frac1N\sum_{i=1}^{N}\frac{x_i-\mu}{\sigma^2}, \\ \frac{\partial\,\text{loss}}{\partial \sigma} = \frac1N\sum_{i=1}^{N}\left[\frac1\sigma - \frac{(x_i-\mu)^2}{\sigma^3}\right]. \end{gathered}$$

Оба градиента обращаются в ноль, когда $\mu$ — выборочное среднее, а $\sigma^2 = \frac1N\sum_i (x_i-\mu)^2$ — (смещённая) выборочная дисперсия. Это оценки максимального правдоподобия, как и предсказывала заметка об энтропии. Градиенту плотность $p$ вообще не нужна: значения $p(x_i)$ лишь делают выдаваемое число равным KL. На реальных данных, где $p$ неизвестна, мы отбрасываем эту константу и вместо KL минимизируем отрицательный логарифм правдоподобия $-\frac1N\sum_i \log q(x_i)$.

Направление шагов проще всего увидеть на точной KL между двумя одномерными гауссианами:

$$\begin{aligned} &D_{\mathrm{KL}}\big(\mathcal N(\mu_P, \sigma_P^2)\,\|\,\mathcal N(\mu, \sigma^2)\big) \\ &\quad= \log\frac{\sigma}{\sigma_P} + \frac{\sigma_P^2 + (\mu_P - \mu)^2}{2\sigma^2} - \frac12 . \end{aligned}$$

Она следует из $D_{\mathrm{KL}} = \mathbb E_P[\log p(x) - \log q(x)]$, если подставить выражение для $\log q$ выше и такое же выражение для $\log p$ (с $\mu_P, \sigma_P$) и учесть, что $\mathbb E_P[(x-\mu_P)^2] = \sigma_P^2$ и $\mathbb E_P[(x-\mu)^2] = \sigma_P^2 + (\mu_P-\mu)^2$. Константы $\tfrac12\log 2\pi$ сокращаются, а слагаемое $-\tfrac12$ — это $-\sigma_P^2/(2\sigma_P^2)$ из $\log p$.

Рисунок

Подгонка Q = 𝒩(μ, σ²) к P = 𝒩(3, 1) градиентным спуском по KL(P‖Q)

При обычном градиентном спуске по $(\mu, \log\sigma)$, как на рисунке, траектория показывает приоритеты прямой KL (forward KL). Пока $Q$ узкое и находится далеко, большая часть массы $P$ лежит там, где $q \approx 0$, и $\log p - \log q$ там огромно. Дешевле всего расширить $Q$: на первых шагах $\sigma$ резко растёт, а $\mu$ едва сдвигается. Когда $Q$ накрывает $P$, $\mu$ смещается к цели, а $\sigma$ сжимается обратно до ширины $P$. Порядок определяется величиной градиентов: из формулы в замкнутом виде $\partial/\partial\log\sigma = 1 - \frac{\sigma_P^2 + (\mu_P-\mu)^2}{\sigma^2}$ растёт как $(\mu_P-\mu)^2/\sigma^2$, тогда как $\partial/\partial\mu = \frac{\mu-\mu_P}{\sigma^2}$ растёт лишь как $\lvert\mu_P-\mu\rvert/\sigma^2$. Adam, который нормирует шаг по каждой координате, сдвигает оба параметра примерно с одинаковой скоростью и приходит к тому же оптимуму другим путём.

Что предполагает KL#

Начнём с формулы для непрерывных распределений (Kullback and Leibler, 1951):

$$D_{\mathrm{KL}}(P\,\|\,Q) = \int p(x)\,\log\frac{p(x)}{q(x)}\,dx .$$

Чтобы она была осмысленным конечным числом, должны выполняться условия (a)–(c); (d) и (e) — это выбор, который нужно делать согласованно (Cover and Thomas, 2006, §8.5; Murphy, 2023, §5.1).

(a) Общий носитель#

Носитель (support) распределения — множество его возможных значений, $\operatorname{supp} p = \lbrace x : p(x) > 0\rbrace$. Если $p(x) > 0$, но $q(x) = 0$ на области, которой $P$ даёт положительную вероятность, то там $\log\frac{p(x)}{q(x)} = \infty$ и дивергенция бесконечна.

Формально $D_{\mathrm{KL}}(P\,\|\,Q)$ конечна, только если $P$ абсолютно непрерывно относительно $Q$, что записывается как $P \ll Q$: каждое множество, которому $Q$ даёт нулевую вероятность, имеет нулевую и $P$-вероятность. Где у $P$ есть масса, она должна быть и у $Q$. Это одна из причин, по которым модели, обучаемые методом максимального правдоподобия, используют всюду положительные плотности, например гауссианы: плотность без «дыр» избавляет от бесконечностей из-за несовпадения носителей: потеря на сэмпле $-\log q(x_i)$ конечна в каждой наблюдённой точке. Конечности ожидаемой потери или KL положительность не гарантирует: в условии (b) ниже обе плотности всюду положительны, и всё же $D_{\mathrm{KL}} = \infty$. Сама по себе положительность не делает модель размытой. Но в сочетании с прямой KL, которая наказывает за любую пропущенную массу, модель, неспособная представить данные, размазывает вероятность между модами и вокруг них (см. раздел «Одна гауссиана, две моды»). Это одно из предлагаемых объяснений размытых сэмплов у моделей изображений, обученных по правдоподобию (Theis, van den Oord and Bethge, 2016).

Рисунок

Общий носитель: у Q не должно быть дыр там, где живёт P

Условие одностороннее. Точки, где $q(x) > 0$, но $p(x) = 0$, вносят в $D_{\mathrm{KL}}(P\,\|\,Q)$ вклад $0\cdot\log\frac{0}{q(x)} = 0$; они важны только для обратного направления.

(b) Интегрируемость#

Интеграл должен сходиться. Совпадения носителей для этого недостаточно. Пусть, например, $P$ — стандартное распределение Коши, а $Q = \mathcal N(0, \sigma^2)$. Обе плотности всюду положительны, но $\log\frac{p}{q}$ растёт как $\frac{x^2}{2\sigma^2}$, а $\mathbb E_P[x^2] = \infty$, поэтому $D_{\mathrm{KL}}(P\,\|\,Q) = \infty$. Модель с лёгкими хвостами не может подогнать по KL данные с тяжёлыми хвостами.

(c) Нормированные плотности#

$p$ и $q$ должны быть плотностями вероятности (или функциями вероятности): неотрицательными и с интегралом (или суммой), равным $1$. KL не определена для ненормированных значений, например логитов классификатора; их сначала нужно пропустить через softmax. Объект torch.distributions возвращает из log_prob нормированные логарифмы плотности.

(d) Направление#

KL не симметрична, $D_{\mathrm{KL}}(P\,\|\,Q) \ne D_{\mathrm{KL}}(Q\,\|\,P)$, поэтому нужно решить, какое распределение играет роль истины, а какое — приближения:

• $D_{\mathrm{KL}}(P\,\|\,Q)$, прямая KL, усредняет по $P$ и наказывает за пропущенные моды: $Q$ должно покрывать всё, где живёт $P$. Это направление максимального правдоподобия.
• $D_{\mathrm{KL}}(Q\,\|\,P)$, обратная KL, усредняет по $Q$ и наказывает за лишнюю массу: $Q$ не должно растекаться туда, где $P \approx 0$. Это направление используется в вариационном выводе (variational inference), в том числе в VAE.

Рисунок

Какое направление что наказывает: пропущенные моды и лишняя масса

(e) Непрерывные и дискретные распределения#

Для дискретных распределений интеграл становится суммой, $D_{\mathrm{KL}}(P\,\|\,Q) = \sum_{x} p(x)\log\frac{p(x)}{q(x)}$. Смешение двух случаев даёт бесконечную дивергенцию. Если у $P$ есть атом в точке $x$, то $Q(\lbrace x\rbrace) = 0$ для любого распределения $Q$ с плотностью, поэтому $P$ не абсолютно непрерывно относительно $Q$ и $D_{\mathrm{KL}}(P\,\|\,Q) = \infty$ по условию (a). Вот почему для непрерывных данных максимальное правдоподобие понимают как выборочную оценку $D_{\mathrm{KL}}(P\,\|\,Q)$ для лежащей в основе плотности $p$, а не как KL эмпирического распределения сэмплов (ср. заметку об энтропии, где для дискретных исходов такое прочтение точно).

Допущения моделирования#

Когда KL используется как функция потерь при обучении, на веру принимается ещё несколько вещей:

• Представимость. Семейство $Q_\theta$ способно выразить нужную форму. Одна гауссиана может точно совпасть с целевым распределением из одной гауссианы, но не со смесью двух компонент; что происходит тогда, показывает последний раздел.
• Дифференцируемость. $\log q_\theta(x)$ должна быть дифференцируема по $\theta$, чтобы можно было применить градиентный спуск. Репараметризация $\sigma = e^{\log\sigma}$ к тому же снимает ограничения с параметров, так что шаг градиента никогда не даст отрицательное $\sigma$.
• Вычислимость. Прямой KL нужны сэмплы из $P$ и значения $\log q$ в них. Обратной KL нужны сэмплы из $Q$ (репараметризованные, чтобы получить градиенты) и значения $\log q$ и $\log p$ в них. Для реальных данных доступно только первое, поэтому максимальное правдоподобие использует прямое направление.

В коде. У каждого допущения есть конкретное место в коде. $Q$ — это Normal или MultivariateNormal, всюду положительное, поэтому (a) математически выполняется при любом целевом распределении. Численно оно всё же может нарушиться. Код вычисляет qx = Q.log_prob(x).exp(), а затем qx.log2(), и в 32-битной арифметике с плавающей точкой экспонента обращается в $0$ (underflow), как только $\log q(x)$ опускается ниже примерно $-103$: при $\sigma \approx 1$ это около $14$ стандартных отклонений от $\mu$, а порог сдвигается вместе с $\log\sigma$. Из-за этого функция потерь становится бесконечной, когда $Q$ стартует далеко от $P$. Работа напрямую с log_prob позволяет этого избежать. Та же осторожность нужна и для двухкомпонентных целевых распределений ниже. Их логарифм плотности $\log\big(\tfrac12 p_1(x) + \tfrac12 p_2(x)\big)$ вычисляется из логарифмов плотности компонент $a = \log p_1(x) + \log\tfrac12$ и $b = \log p_2(x) + \log\tfrac12$ как $\operatorname{logaddexp}(a, b) = \max(a, b) + \log\big(1 + e^{-|a - b|}\big)$; экспонента большого отрицательного числа в этой формуле никогда не стоит отдельно, поэтому результат остаётся конечным вдали от обеих компонент. Ковариация $\mathbf L\mathbf L^\top$ с экспоненцированной диагональю сохраняет $Q$ корректной нормированной плотностью, как требует (c). Целевое распределение — известная гауссиана или смесь, поэтому px можно вычислить. Наконец, функция потерь усредняет по x ~ P, что фиксирует направление $D_{\mathrm{KL}}(P\,\|\,Q)$.

Одна гауссиана, две моды#

Если целевое распределение — одна гауссиана, подгонка по KL просто восстанавливает его. Интересен случай, когда семейство моделей не может представить целевое распределение. В двумерном случае целевое распределение — смесь с равными весами $\mathcal N\big((-5,-5), \mathbf I\big)$ и $\mathcal N\big((4,4), \mathbf I\big)$. В одномерном оно равно

$$P = \tfrac12\,\mathcal N(-3,\, 0.5^2) + \tfrac12\,\mathcal N(-8,\, 0.5^2),$$

а модель стартует с $Q = \mathcal N(0, 1)$. Сравниваются четыре функции потерь: прямая KL, обратная KL, JS и функция потерь наименьших квадратов на значениях плотности. Ни одна из них не удовлетворительна, и у каждой на то своя причина. Те же компромиссы снова возникают в разных генеративных моделях, поэтому эту игрушечную задачу стоит разобрать подробно (Huszár, 2015; Theis, van den Oord and Bethge, 2016).

Вдали от цели доминирует пропущенная мода#

Начальное $Q = \mathcal N(0, 1)$ едва достаёт до дальней моды в $-8$. Гауссиана никогда не равна нулю в точности, но в восьми стандартных отклонениях от среднего она практически равна нулю. Там $\log\frac{p(x)}{q(x)}$ огромен, и дальняя мода доминирует во всей дивергенции. Если бы $Q$ было там в точности нулём, дивергенция была бы бесконечной, как говорит условие (a). Для гауссовского $Q$ прямая KL тем не менее остаётся конечной, и её градиент по $\mu$, $(\mu - \mathbb E_P[x])/\sigma^2$, здесь велик. Перекрытие важно, во-первых, численно, потому что в 32-битной арифметике с плавающей точкой $q$ обращается в $0$ (см. В коде выше), а во-вторых, для ограниченных функций потерь — JS и наименьших квадратов, — градиенты которых без перекрытия исчезают (см. таблицу ниже).

Рисунок

Когда Q не покрывает часть носителя P

Прямая KL покрывает обе моды, обратная выбирает одну#

Минимизация $D_{\mathrm{KL}}(P\,\|\,Q)$ заставляет одну гауссиану покрыть обе моды с большим разбросом. Функция потерь усредняет $\log\frac{p}{q}$ по сэмплам $P$, поэтому любая область, где у $P$ есть масса, а $Q$ мало, наказывается очень сильно. Для гауссовского $Q$ оптимум — точное сопоставление моментов (moment matching). Если записать прямую KL как

$$D_{\mathrm{KL}}(P\,\|\,Q) = -H(P) - \mathbb E_P[\log q(x)],$$

то от $Q$ зависит только второе слагаемое, а для гауссианы $\log q$ зависит от $x$ только через $x$ и $xx^{\top}$. Приравнивая нулю производные по $\boldsymbol\mu$ и $\boldsymbol\Sigma$, получаем

$$\boldsymbol\mu^\ast = \mathbb E_P[x],\qquad \boldsymbol\Sigma^\ast = \operatorname{Cov}_P[x]$$

(Bishop, 2006, §10.7). Для одномерной цели это даёт $\mu^\ast = -5.5$ и $\sigma^{\ast2} = 0.25 + 2.5^2 = 6.5$, так что $\sigma^\ast \approx 2.55$: каждая компонента вносит собственную дисперсию $0.25$, а разброс двух средних вокруг $-5.5$ добавляет $2.5^2$. Для двумерной цели получаются среднее $(-0.5, -0.5)$ и ковариация $\begin{pmatrix} 21.25 & 20.25 \\ 20.25 & 21.25\end{pmatrix}$: одна вытянутая гауссиана вдоль диагонали, проходящей через обе моды, с самой плотной частью в пустом промежутке между ними. В этом оптимуме дивергенция равна $1.170$ ната ($1.69$ бита). Подгонка с Adam (шаг обучения $0.1$, 2000 фиксированных сэмплов, старт с $\mathcal N(\mathbf 0, \mathbf I)$) приходит именно к этой картине с оценкой $1.669$ бита после 4800 шагов; разрыв в $0.019$ бита меньше стандартного отклонения оценки по 2000 сэмплам — около $0.023$ бита.

Обратная дивергенция $D_{\mathrm{KL}}(Q\,\|\,P)$ усредняет по сэмплам $Q$. Она наказывает $Q$ за массу там, где $P$ мало, и не обращает внимания на те части $P$, которые $Q$ игнорирует, поэтому $Q$ захватывает одну моду (Bishop, 2006, §10.1.2, рисунок 10.3; Minka, 2005). Для двух равных, хорошо разделённых мод положение точно на одной из них стоит всего $D_{\mathrm{KL}}(Q\,\|\,P) \approx \ln 2$: на этой моде $p \approx q/2$, поэтому $\log\frac{q}{p} \approx \ln 2$ всюду, где у $Q$ есть масса.

Рисунок

Подгонка одной гауссианы к двум модам: прямая и обратная KL

JS тоже не покрывает моды#

JS никогда не обращается в бесконечность, но проблему покрытия мод не решает. JS-подгонка двумерной смеси с Adam (шаг обучения $0.1$, 5000 итераций), начатая с $\boldsymbol\mu = (-2,-2)$, заканчивается на $\boldsymbol\mu \approx (-4.99, -5.06)$ с почти единичной ковариацией при значении функции потерь $0.273$ бита. Это гауссиана на моде в $(-5,-5)$; мода в $(4,4)$ отброшена. Функция потерь в этом запуске — взвешенный вариант, $0.7\,D_{\mathrm{KL}}(P\,\|\,M) + 0.3\,D_{\mathrm{KL}}(Q\,\|\,M)$ в битах. Для него точный захват одной моды даёт $0.270$ бита, что близко к достигнутому значению; для симметричной JS получилось бы $\tfrac34\log_2\tfrac43 \approx 0.311$ бита. С одномерной целью всё заканчивается так же. Обычный градиентный спуск (шаг обучения $0.1$, 5000 новых сэмплов каждого распределения на каждом шаге) из $\mathcal N(0, 1)$ после 1800 шагов находится в $\mu = -2.995$, $\sigma = 0.495$, на ближней моде в $-3$, с тем же значением взвешенной функции потерь $0.270$ бита. Пиковая плотность $Q$ там, $0.81$, вдвое больше, чем у цели, $0.40$: вся масса $Q$ сосредоточена на моде, которая содержит половину массы $P$. Симметричная JS ведёт себя так же: для одномерной цели её лучшая одиночная гауссиана тоже располагается на одной моде (см. таблицу ниже).

Оценка JS по сэмплам. Каждая половина JS — математическое ожидание по своему распределению,

$$\begin{aligned} D_{\mathrm{JS}}(P\,\|\,Q) &= \tfrac12\,\mathbb E_{x\sim P}\Big[\log\frac{p(x)}{m(x)}\Big] + \tfrac12\,\mathbb E_{y\sim Q}\Big[\log\frac{q(y)}{m(y)}\Big], \\ m &= \tfrac12(p + q), \end{aligned}$$

поэтому для оценки нужны две выборки: $x_1, \dots, x_N \sim P$ для первого слагаемого и $y_1, \dots, y_N \sim Q$ для второго, и каждое слагаемое усредняется так же, как в оценке KL выше. Плотность смеси вычисляется в логарифмической шкале, $\log m = \operatorname{logaddexp}\big(\log p + \log\tfrac12,\ \log q + \log\tfrac12\big)$, как и для двухкомпонентной цели. Второе слагаемое зависит от параметров $Q$ ещё и через сэмплы $y_j$. Чтобы его градиент был верным, сэмплы нужно брать с помощью трюка репараметризации (reparameterization trick), $y = \mu + \sigma\varepsilon$ с $\varepsilon \sim \mathcal N(0, 1)$: тогда обратное распространение ошибки учитывает, как они смещаются при изменении $\mu$ и $\sigma$. В приведённых выше запусках JS сэмплы $y_j$ берутся без него, поэтому их градиенты эту зависимость не учитывают.

У ограниченности есть и вторая цена. Когда $Q$ и $P$ почти не перекрываются, JS уже близка к своему максимуму $\ln 2$ и почти плоская, поэтому её градиент слаб. Для одномерной цели при начальном $\mathcal N(0,1)$ она равна $0.645$ ната — уже $93\%$ от $\ln 2$. То же насыщение — причина, по которой исходная минимаксная целевая функция GAN, равная $2\,D_{\mathrm{JS}} - \ln 4$ при оптимальном дискриминаторе, даёт затухающие градиенты генератора (Arjovsky and Bottou, 2017, теорема 2.4; см. заметку о GAN и GAN Вассерштейна).

Функция потерь наименьших квадратов на плотностях чувствительна к масштабу#

Наивная альтернатива — напрямую сравнивать значения плотности:

$$\mathcal L_{\mathrm{LS}} = \frac1N \sum_{i=1}^{N} \bigl(q(x_i) - p(x_i)\bigr)^2 , \qquad x_i \sim P .$$

Значения плотности зависят от масштаба данных. Если растянуть каждую координату в $c$ раз, все плотности в $d$ измерениях делятся на $c^{d}$. Тогда квадраты разностей, а с ними и функция потерь, уменьшаются в $c^{2d}$ раз, а градиенты — в сопоставимое число раз. При обычном SGD шаг обучения, подходящий для одного масштаба, не годится для другого. Адаптивные оптимизаторы вроде Adam компенсируют постоянный множитель, но не решают вторую проблему: когда $Q$ не перекрывается с $P$, в точках данных $q(x_i) \approx 0$, и функция потерь остаётся около константы $\mathbb E_P[p(x)^2]$, что почти не даёт градиента. У KL нет ни одной из этих проблем. Отношение $p/q$, а значит, и KL не меняются при изменении масштаба. А $-\log q(x_i)$ вдали от $\mu$ растёт как $(x_i-\mu)^2/2\sigma^2$, так что при удалении $Q$ от $P$ градиент KL растёт, а не исчезает ($\lvert\partial/\partial\mu\rvert = 5.5$ при начальном $Q$ в таблице ниже).

У функции потерь наименьших квадратов к тому же несколько минимумов. Для одномерной цели её глобальный минимум — широкая гауссиана из таблицы ниже ($\sigma \approx 2.53$), но гауссиана на одной моде, $\mathcal N(-3, 1.21^2)$, — локальный минимум (значение $0.0495$ против $0.0470$). В двумерном случае гауссиана на одной моде, $\mathcal N\big((4,4), 1.58^2\,\mathbf I\big)$, — локальный минимум со значением $0.00113$; наименьшее значение, которое находит численный поиск, равно $0.00102$, у гауссианы с центром в $(-0.5,-0.5)$, вытянутой вдоль диагонали (стандартное отклонение $6.5$ вдоль неё и $0.43$ поперёк). К какому минимуму придёт градиентный спуск, зависит от начальной точки. Подгонка с Adam (шаг обучения $0.001$, 5000 итераций, 1000 сэмплов), начатая с $\mathcal N\big((1,1), \mathbf I\big)$, которое перекрывается только с модой в $(4,4)$, заканчивается гауссианой $Q$ с центром на этой моде, более широкой, чем сама мода; мода в $(-5,-5)$ игнорируется. Её сэмплы почти ничего не вносят в градиент, потому что там $q \approx 0$. $Q$ получается шире компоненты, потому что гауссиана, вся масса которой сосредоточена на моде с весом $\tfrac12$, была бы там вдвое плотнее $P$, а расширение понижает её пик. Больший шаг обучения или вчетверо большее число итераций приводят к той же моде.

Сравнение четырёх функций потерь#

Одномерная цель выше, натуральные логарифмы. Функция потерь наименьших квадратов $\mathbb E_P\big[(q-p)^2\big]$ измеряется не в натах: её единица — квадрат плотности, поэтому её значение зависит от масштаба данных.

Одна гауссиана $Q$	$D_{\mathrm{KL}}(P\,\Vert\,Q)$	$D_{\mathrm{KL}}(Q\,\Vert\,P)$	$D_{\mathrm{JS}}$	LS
начальная $\mathcal N(0, 1)$	$17.88$	$19.50$	$0.645$	$0.0901$
с сопоставленными моментами $\mathcal N(-5.5,\, 6.5)$, $\sigma \approx 2.55$	$\mathbf{0.936}$	$3.709$	$0.255$	$\mathbf{0.0470}$
одна мода $\mathcal N(-3,\, 0.5^2)$	$24.31$	$\mathbf{0.693}$	$\mathbf{0.216}$	$0.0919$
$\lvert\partial/\partial\mu\rvert$ при начальном $Q$	$5.5$	$12.0$	$0.070$	$0.0045$

Жирным для каждой функции потерь выделен лучший из трёх кандидатов. Для KL, обратной KL и JS выделенное значение — также глобальный оптимум функции потерь по всем $(\mu, \sigma)$, найденный перебором по сетке; оптимум наименьших квадратов находится при $\sigma \approx 2.53$, то есть практически совпадает с гауссианой с сопоставленными моментами. Значения вычислены численным интегрированием. Прямая KL и наименьшие квадраты предпочитают широкую гауссиану, обратная KL и JS — одну моду. Последняя строка показывает, насколько сильно каждая функция потерь тянет начальное $Q$ к данным: KL даёт сильный сигнал, а JS и наименьшие квадраты почти плоские.

Подгонка по прямой KL с Adam (шаг обучения $0.01$, 1000 фиксированных сэмплов) приходит к гауссиане с сопоставленными моментами: $\mu = -5.494$, $\sigma = 2.566$, с оценкой $1.300$ бита. Точное значение в оптимуме — $0.936$ ната, или $1.350$ бита; разница в $0.05$ бита меньше двух стандартных отклонений оценки по 1000 сэмплам (каждое около $0.03$ бита).

Как ведут себя градиенты во время обучения#

Дополнительный материал: как меняются модули градиентов в двух запусках обучения, один из которых использует неполную оценку JS. Остальная часть заметки на него не опирается.

Последняя строка таблицы — снимок в начальной точке. На протяжении всей подгонки градиенты KL и JS ведут себя по-разному (одномерная цель, модули в битах на единицу $\mu$ и на единицу $\log\sigma$):

• Прямая KL на фиксированных сэмплах: большие и гладкие, в конце нулевые. При начальном $\mathcal N(0, 1)$ градиенты равны $7.9$ по $\mu$ и $51.6$ по $\log\sigma$ (по формулам для градиентов выше это $5.5$ и $35.75$ ната, делённые на $\ln 2$). Градиент по $\log\sigma$, намного больший, обваливается первым — как только $Q$ становится достаточно широким, чтобы накрыть обе моды. Градиент по $\mu$ убывает медленнее и держится на уровне нескольких десятых бита, пока $Q$ перемещается к цели. Поскольку сэмплы фиксированы, функция потерь — детерминированная функция от $(\mu, \log\sigma)$, и по мере приближения к её минимуму оба градиента стремятся к нулю.
• JS с новыми сэмплами: маленькие и шумные. Этот запуск — диагностика, а не измерение настоящего градиента JS: он сэмплирует $y_j$ без приёма репараметризации, поэтому его градиенты не учитывают зависимость второго слагаемого от сэмплов. В описанном выше запуске обычного градиентного спуска (шаг обучения $0.1$, 5000 новых сэмплов $P$ и $Q$ на каждом шаге) ни один градиент не превышает примерно $0.11$ бита. Модули дважды достигают пика: около шага 40, пока $Q$ расширяется, и около шага 720, когда оно сжимается к моде в $-3$. После этого они не затухают, а колеблются между $0$ и примерно $0.04$ бита вплоть до последнего из 2000 шагов. На каждом шаге функция потерь оценивается по новым сэмплам, поэтому её градиент несёт шум оценки Монте-Карло, который не исчезает в оптимуме, и при постоянном шаге обучения параметры продолжают на него реагировать.

Здесь смешаны три эффекта, и их стоит разделять. Геометрия целевой функции: JS почти плоская вдали от цели, что показывает последняя строка таблицы по точным интегралам, а KL даёт там сильный сигнал и чистую остановку в оптимуме. Оценка: градиенты запуска с JS неполны, поэтому точная форма его траектории мало говорит о настоящем градиенте JS. Сэмплирование: при новых сэмплах на каждом шаге (resampling) любой оценённый градиент сохраняет шум Монте-Карло и в оптимуме, так что подгонка с постоянным шагом обучения никогда точно не останавливается.

Ноутбук#

Эксперименты следуют ноутбуку Distances between distributions (HSE-LAMBDA Deep Generative Models, закреплённый коммит; ячейки нумеруются с 1, считая все ячейки, включая текстовые). Результаты запусков, приведённые выше, взяты из его сохранённых выходов ячеек; точные значения дивергенций и градиентов, минимумы функции потерь наименьших квадратов и симуляция отрицательных оценок вычислены отдельно, как и запуски с функцией потерь наименьших квадратов с пересчитанным px и возобновлённая кривая градиентов KL. Ячейки 1–11 (вспомогательные функции, целевое распределение из одной гауссианы и начальное $Q$) описаны в заметке о гауссианах.

• Ячейки 12–15: выборочная оценка $D_{\mathrm{KL}}(P\,\|\,Q)$ (kl_loss, в битах; её аргумент weights принимает веса значимости (importance weights) и по умолчанию состоит из единиц) и подгонка двумерной модели с Adam (шаг обучения 0.1, 5000 итераций) к одиночной цели $\mathcal N((-5,-5), \mathbf I)$ на 1000 сэмплах x из ячейки 8. Последний график, на шаге 4800, озаглавлен KL = -0.0020.
• Ячейки 16–19: двухкомпонентная двумерная смесь со средними $(-5,-5)$ и $(4,4)$ и та же подгонка по KL, применённая к ней, на 2000 фиксированных сэмплах. log_prob_mixture использует torch.logaddexp. sample_from_mixture берёт ровно половину сэмплов из каждой компоненты, то есть строит стратифицированную выборку (stratified sample); при независимых одинаково распределённых (i.i.d.) сэмплах из смеси их число в каждой компоненте распределялось бы по биномиальному закону. Последний график, на шаге 4800, показывает KL mixture = 1.6688.
• Ячейки 20–23: JS-дивергенция и её оценка методом Монте-Карло (js_div: x = P.sample(...) для первого слагаемого, y = sample(Q, ...) для второго, torch.logaddexp для $\log m$), а также JS-подгонка смеси, начатая с $\boldsymbol\mu = (-2,-2)$, с 5000 сэмплами на каждом шаге. В последней распечатке, на шаге 4900, стоят JS loss = 0.2734 и mu = tensor([-4.9879, -5.0583]). При повторном запуске важны две детали реализации. js_div возвращает взвешенную комбинацию $0.7/0.3$, а не симметричную JS с весами ½–½. Кроме того, сэмплы $Q$ она берёт через sample(), а не rsample(), поэтому градиент не учитывает, как эти сэмплы смещаются при изменении $Q$.
• Ячейки 24–26: LSE_loss — среднеквадратичная ошибка между $q(x_i)$ и $p(x_i)$ — и подгонка с Adam (шаг обучения 0.001, 5000 итераций), начатая с $\mathcal N((1,1), \mathbf I)$. Для контекста ноутбук ссылается на статью о GAN с наименьшими квадратами (least-squares GAN; Mao et al., 2017), хотя в ней наименьшие квадраты применяются к выходам дискриминатора, а не к плотностям. Ячейка обучения берёт новые сэмплы x, но повторно использует px, оставшееся от ячейки с JS, — плотности цели в 1000 других точках, поэтому px нужно пересчитать, чтобы функция потерь сравнивала две плотности в одних и тех же точках. Последний график, на шаге 4800, показывает LSE = 0.0013, и $Q$ находится на моде в $(4,4)$. С пересчитанным px та же подгонка всё равно заканчивается на этой моде, со значением функции потерь около $0.0011$ в повторной реализации.
• Ячейки 27–38: вспомогательные функции кода выше (get_parameters_1d и create_distr_1d в ячейке 29, target_1d в ячейках 30–31) и одномерная смесь $\tfrac12\mathcal N(-3, 0.5^2) + \tfrac12\mathcal N(-8, 0.5^2)$, к которой гауссиана подгоняется по KL (Adam, шаг обучения 0.01, 2000 итераций, 1000 фиксированных сэмплов) и, начиная с $\mathcal N(0, 1)$, по JS (SGD, шаг обучения 0.1, 2000 итераций); модули градиентов по $\mu$ и $\log\sigma$ построены на графиках по ходу обучения (ячейки 35 и 38). JS-подгонка использует ту же взвешенную js_div с 5000 новыми сэмплами на каждом шаге. Сохранённые графики показывают подгонку по KL на шаге 1500 (Loss (KL) = 1.3001, mu = -5.494, std = 2.566) и JS-подгонку на шаге 1800 (JS = 0.2698, mu = -2.995, std = 0.495). В легенде ячейки 35 написано |grad sigma|, но кривая — это градиент по $\log\sigma$, как и подписано в ячейке 38.
• Сохранённая кривая градиентов KL начинается не с $\mathcal N(0, 1)$. Ячейка 33 создаёт параметры, соответствующие $\mathcal N(0, 1)$; ячейка 34 не создаёт их заново. Кривая ячейки 35 начинается с $\lvert\partial/\partial\mu\rvert \approx 0.85$ и $\lvert\partial/\partial\log\sigma\rvert \approx 3.6$ бита, тогда как старт с $\mathcal N(0, 1)$ даёт $7.9$ и $51.6$. Таких значений новый запуск достигает примерно через 200 шагов, при $\mu \approx -1.1$, $\sigma \approx 2.7$, а в номерах выполнения (execution count) между ячейками 33 и 34 пропущено одно число: сохранённый запуск продолжился с того места, где параметры оставил более ранний, прерванный запуск ячейки 34. Повторная реализация, которая возобновляет обучение с этой точки с новым оптимизатором, воспроизводит сохранённую кривую: $\lvert\partial/\partial\log\sigma\rvert$ падает почти до $0$ примерно за 100 шагов, $\lvert\partial/\partial\mu\rvert$ держится около $0.25$ примерно до шага 400 и достигает $0$ к шагу 750 или около того, а к шагу 1500 подгонка находится в $\mu \approx -5.50$, $\sigma \approx 2.56$. Новый запуск медленнее: к шагу 1500 он всё ещё в $\mu \approx -4.1$. Чтобы получить задуманный старт, ячейку 33 нужно выполнить заново перед ячейкой 34.

Дальше#

До сих пор у модели было пять параметров и плотность в замкнутом виде. Реальные генеративные модели заменяют эти пять параметров нейронной сетью, а многие отказываются и от плотности в замкнутом виде. Обучение таких сетей требует собственного инструментария. Следующая заметка собирает элементы, которые используются на протяжении всего курса: батч-нормализацию (batch normalization), то, как градиенты проходят через глубокие стеки слоёв, и размерности свёрток.

Литература#

• S. Kullback, R. A. Leibler. On Information and Sufficiency. Annals of Mathematical Statistics, 22(1), 1951.
• T. M. Cover, J. A. Thomas. Elements of Information Theory, 2nd ed. Wiley, 2006. §8.5: относительная энтропия для непрерывных распределений.
• C. M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006. §10.1.2 и рисунок 10.3: прямая и обратная KL на бимодальном целевом распределении; §10.7: сопоставление моментов.
• K. P. Murphy. Probabilistic Machine Learning: Advanced Topics. MIT Press, 2023. §5.1: KL-дивергенция, её свойства и различие прямого и обратного направлений.
• T. Minka. Divergence measures and message passing. Microsoft Research Technical Report MSR-TR-2005-173, 2005.
• F. Huszár. How (not) to train your generative model: scheduled sampling, likelihood, adversary? arXiv:1511.05101, 2015.
• L. Theis, A. van den Oord, M. Bethge. A note on the evaluation of generative models. ICLR, 2016.
• M. Arjovsky, L. Bottou. Towards principled methods for training generative adversarial networks. ICLR, 2017. §2, теорема 2.4: затухающие градиенты минимаксной функции потерь генератора.
• D. P. Kingma, J. Ba. Adam: A method for stochastic optimization. ICLR, 2015.
• S. S. Wilks. The large-sample distribution of the likelihood ratio for testing composite hypotheses. Annals of Mathematical Statistics, 9(1), 60–62, 1938.
• X. Mao, Q. Li, H. Xie, R. Y. K. Lau, Z. Wang, S. P. Smolley. Least squares generative adversarial networks. ICCV, 2017.
• HSE-LAMBDA, репозиторий курса Deep Generative Models: distances.ipynb.