Модели, основанные на энергии

Предыдущая заметка сэмплировала из плотности, известной лишь с точностью до константы, с помощью MALA, гамильтонова Монте-Карло и NUTS. Там целевое распределение было фиксированным. Модель, основанная на энергии (energy-based model, EBM), делает его обучаемым: нейронная сеть $E_\theta(x)$ с одним скалярным выходом задаёт плотность $p_\theta(x) = e^{-E_\theta(x)}/Z_\theta$, а обучение сдвигает $\theta$, пока эта плотность не приблизит данные (LeCun et al., 2006). Архитектура ничем не ограничена: не нужны ни обратимость, как в потоках, ни порядок пикселей, как в авторегрессионных моделях, ни энкодер, как в VAE. Цена — нормировочная константа $Z_\theta$. Правдоподобие нельзя вычислить, для его градиента нужны сэмплы из модели, и поэтому каждый шаг обучения запускает марковскую цепь на той самой модели, которая обучается. В этой заметке мы выводим этот градиент, строим вокруг него стандартный рецепт (короткие цепочки Ланжевена, буфер воспроизведения, штраф на энергии), объясняем, почему обучение неустойчиво, приводим две оценки, которым не нужны сэмплы во время обучения, и заканчиваем тем, для чего энергии используются помимо генерации.

Обозначения. $x \in \mathbb R^d$ — точка данных, например изображение MNIST с $d = 784$ значениями пикселей в $[-1, 1]$. $P$ — распределение данных с плотностью $p$. Энергия (energy) $E_\theta: \mathbb R^d \to \mathbb R$ — сеть с параметрами $\theta$, а плотность модели — $p_\theta(x) = e^{-E_\theta(x)}/Z_\theta$ со статистической суммой (partition function) $Z_\theta = \int e^{-E_\theta(x)}\,dx$. Сэмплы данных $x^+ \sim P$ — позитивные примеры (positive examples); сэмплы $x^-$, которые должны происходить из $p_\theta$, — негативные примеры (negative examples). $\nabla_\theta$ — градиент по параметрам, $\nabla_x$ — по входу. Один шаг Ланжевена размера $\varepsilon$ — это $x \leftarrow x - \tfrac{\varepsilon}{2}\nabla_x E_\theta(x) + \sqrt{\varepsilon}\,z$ с $z \sim \mathcal N(0, I)$, как в курсе теории вероятностей, а $K$ — число шагов цепочки. Второй аргумент $\mathcal N$ — дисперсия.

Энергия задаёт плотность#

Низкая энергия означает высокую вероятность. Экспонента превращает любую вещественнозначную $E_\theta$ в положительную функцию, а деление на $Z_\theta$ делает её интеграл равным единице. Это распределение Больцмана, или Гиббса, из статистической физики при температуре 1. Единственное условие — конечность $Z_\theta$: энергия должна достаточно быстро расти на бесконечности, или область должна быть ограниченной, как у изображений с пикселями в $[-1, 1]$. Квадратичная энергия $E(x) = x^2/2$ даёт $Z = \sqrt{2\pi}$ и $p = \mathcal N(0, 1)$; сеть может нарисовать любой ландшафт, с таким числом долин, сколько мод у данных.

Значение имеют только разности энергий:

$$\frac{p_\theta(x)}{p_\theta(y)} = e^{E_\theta(y) - E_\theta(x)} .$$

Прибавление константы $c$ к энергии умножает $Z_\theta$ на $e^{-c}$ и не меняет $p_\theta$. Эта свобода вернётся ниже как проблема сдвига.

Почему бы не вычислить $Z_\theta$? Это интеграл по $\mathbb R^d$. Сетка всего из 10 точек на пиксель изображения MNIST содержит $10^{784}$ узлов. Поэтому всё дальнейшее устроено так, чтобы $Z_\theta$ сокращалась или вообще не требовалась.

Для решений нормировка не нужна#

Многие применения энергии сравнивают её значения и никогда их не нормируют. Для энергии $E_\theta(x, y)$ входа $x$ и ответа $y$ инференс — это $\hat y = \arg\min_y E_\theta(x, y)$: классификация, когда $y$ — метка, и восстановление, когда $x$ — зашумлённое изображение, а $y$ — чистое (LeCun et al., 2006, раздел 1.1 и рисунок 2). Важен только порядок энергий, поэтому $Z$ не появляется. EBM для изображений делают то же самое динамикой Ланжевена: цепочка, начатая с испорченного тестового изображения, убирает импульсный шум («соль и перец», salt-and-pepper) или заполняет закрытую маской область (Du and Mordatch, 2019, раздел 4.2). Плотность, а с ней и $Z$, нужна, когда модель должна порождать сэмплы, сообщать правдоподобие или сочетаться с другими моделями в общей шкале.

Градиент правдоподобия: позитивная и негативная фазы#

Метод максимального правдоподобия минимизирует отрицательный логарифм правдоподобия, который с точностью до константы равен KL-дивергенции распределения данных относительно модели (заметка об энтропии). Для EBM $\log p_\theta(x) = -E_\theta(x) - \log Z_\theta$, поэтому

$$\begin{aligned} \mathcal L(\theta) &= -\mathbb E_{x\sim P}\log p_\theta(x)\\ &= \mathbb E_{x\sim P}\,E_\theta(x) + \log Z_\theta . \end{aligned}$$

Первое слагаемое — среднее по данным; его градиент равен $\mathbb E_{x\sim P}\nabla_\theta E_\theta(x)$. Вся трудность — во втором слагаемом. Три шага:

1. Производная логарифма сложной функции: $\nabla_\theta \log Z_\theta = \nabla_\theta Z_\theta / Z_\theta$.
2. Дифференцируем под знаком интеграла; это законно, когда $E_\theta$ дифференцируема по $\theta$, а градиент подынтегральной функции ограничен интегрируемой функцией: $\nabla_\theta Z_\theta = \int \nabla_\theta e^{-E_\theta(x)}\,dx = \int e^{-E_\theta(x)}\big(-\nabla_\theta E_\theta(x)\big)\,dx$.
3. Делим на $Z_\theta$ и узнаём плотность модели $e^{-E_\theta(x)}/Z_\theta = p_\theta(x)$:

$$\begin{aligned} \nabla_\theta \log Z_\theta &= \int p_\theta(x)\big(-\nabla_\theta E_\theta(x)\big)\,dx\\ &= -\mathbb E_{x\sim p_\theta}\nabla_\theta E_\theta(x) . \end{aligned}$$

Вместе:

$$\begin{aligned} \nabla_\theta \mathcal L(\theta) = {}& \underbrace{\mathbb E_{x^+\sim P}\,\nabla_\theta E_\theta(x^+)}_{\text{позитивная фаза}}\\ &- \underbrace{\mathbb E_{x^-\sim p_\theta}\,\nabla_\theta E_\theta(x^-)}_{\text{негативная фаза}} , \end{aligned}$$

а градиент логарифма правдоподобия — то же выражение с противоположным знаком, $-\mathbb E_P \nabla_\theta E_\theta(x^+) + \mathbb E_{p_\theta}\nabla_\theta E_\theta(x^-)$. Шаг градиентного спуска $\theta \leftarrow \theta - \eta\,\nabla_\theta\mathcal L$ понижает энергию в точках данных и повышает её в собственных сэмплах модели: негативная фаза поднимает энергию каждой точки пропорционально её вероятности в модели (LeCun et al., 2006, раздел 2.2.4). Обновление останавливается, когда два математических ожидания совпадают, $\mathbb E_P\nabla_\theta E_\theta = \mathbb E_{p_\theta}\nabla_\theta E_\theta$: метод максимального правдоподобия приравнивает средние статистик $\nabla_\theta E_\theta(x)$ по данным и по модели. Вывод верен для любой архитектуры (Song and Kingma, 2021, раздел 3).

Разобранный пример: три состояния. Пусть $x \in \{a, b, c\}$, а энергия задана таблицей, $\theta = (E_a, E_b, E_c) = (0, 1, 2)$. Тогда $Z = 1 + e^{-1} + e^{-2} = 1.503$ и $p_\theta = (0.665, 0.245, 0.090)$. Пусть данные — $P = (0.5, 0.3, 0.2)$. Для таблицы $\partial E_\theta(x)/\partial E_j$ равно 1 при $x = j$ и 0 иначе, поэтому две фазы равны $P(j)$ и $p_\theta(j)$:

$$\frac{\partial \mathcal L}{\partial E_j} = P(j) - p_\theta(j) = (-0.165,\ 0.055,\ 0.110) .$$

Шаг против этого градиента повышает $E_a$, где модель помещает слишком много массы, и понижает $E_b$ и $E_c$. С шагом 1 энергии становятся $(0.165, 0.945, 1.890)$, а $p_\theta = (0.611, 0.280, 0.109)$ — ближе к $P$. Неподвижная точка — $p_\theta = P$, то есть $E_j = -\log P(j) + \text{const}$. Три компоненты градиента в сумме дают ноль, поэтому средняя энергия остаётся равной 1: градиент никогда не сдвигает общий уровень энергий.

Зачем негативной фазе сэмплы из модели#

Позитивная фаза — среднее по обучающей выборке. Негативная фаза — среднее по $p_\theta$: по распределению, которое обучается, меняется после каждого обновления и нормировочная константа которого неизвестна. Выборке по значимости (importance sampling) понадобилось бы предлагающее распределение, близкое к $p_\theta$ в $d$ измерениях, а найти его так же трудно, как решить исходную задачу. Остаются марковские методы Монте-Карло (Markov chain Monte Carlo, MCMC), которым нужен только

$$\nabla_x \log p_\theta(x) = -\nabla_x E_\theta(x) - \underbrace{\nabla_x \log Z_\theta}_{=\,0} ,$$

поскольку $Z_\theta$ зависит от $\theta$, но не от $x$. Динамика Ланжевена (Langevin dynamics),

$$\begin{gathered} x_{k+1} = x_k - \frac{\varepsilon}{2}\,\nabla_x E_\theta(x_k) + \sqrt{\varepsilon}\,z_k,\\ z_k \sim \mathcal N(0, I), \end{gathered}$$

— обычный выбор для изображений: каждый шаг стоит одного градиента сети энергии по входу, то есть одного прямого и одного обратного прохода (Du and Mordatch, 2019, раздел 3). Распределение $x_K$ приближается к $p_\theta$ при $\varepsilon \to 0$ и $K \to \infty$; MALA и HMC из предыдущей заметки устраняют смещение, вызванное размером шага, ценой отношения плотностей, которое для EBM равно $e^{E_\theta(x) - E_\theta(x')}$ и тоже не содержит $Z_\theta$. В литературе по EBM сэмплер часто называют SGLD — по стохастической градиентной динамике Ланжевена (stochastic gradient Langevin dynamics) из работы Welling and Teh (2011), где шаги Ланжевена с градиентами по мини-батчам сэмплировали параметры модели из апостериорного распределения. Здесь градиент точный, а цепочка движется в пространстве данных.

Препятствие — стоимость. На мультимодальной энергии, например такой, которая описывает изображения, цепочки Ланжевена перемешиваются медленно: цепочки из разных стартов застревают в разных модах (Nijkamp et al., 2019, раздел 1.1). После каждого обновления параметров понадобилась бы длинная цепочка. Методы ниже укорачивают цепочки и платят за это смещением.

Контрастивная дивергенция: ранняя остановка цепочек#

Hinton (2002) предложил останавливать цепочки рано. Контрастивная дивергенция (contrastive divergence, CD-$k$) запускает каждую цепочку из точки данных, делает лишь $k$ шагов (в исходной работе о произведениях экспертов — один полный шаг Гиббса) и использует конечные точки как негативные примеры. Она не следует градиенту правдоподобия. Обозначим через $P_0 = P$ данные, через $P_1$ — распределение после одного шага цепи, а через $P_\infty = p_\theta$ — её равновесное распределение. CD-1 следует, если отбросить слагаемое, обычно малое, градиенту

$$D_{\mathrm{KL}}(P_0\,\|\,P_\infty) - D_{\mathrm{KL}}(P_1\,\|\,P_\infty) .$$

Разность никогда не бывает отрицательной, потому что шаг марковской цепи не может увеличить KL-дивергенцию относительно её стационарного распределения. Для цепи, все переходы которой имеют положительную вероятность, она равна нулю, только когда $P_1 = P_0$, а это означает, что данные уже имеют равновесное распределение: модель идеальна (Hinton, 2002, раздел 3). Иными словами, CD требует, чтобы цепь не уходила от данных.

Современные EBM для изображений сохраняют короткие цепочки, но запускают их из шума: равномерного шума в диапазоне пикселей, за которым следуют $K$ шагов Ланжевена на текущей энергии, где $K$ — примерно от 10 до 100 (Du and Mordatch, 2019; Nijkamp et al., 2019). Оценка градиента та же, только $x^-$ — конечные точки цепочек. Поскольку конечные точки не являются сэмплами из $p_\theta$, оценка смещена (Song and Kingma, 2021, раздел 3).

Чему учатся короткие цепочки#

У смещения есть определённое направление, и пример с одним параметром показывает его точно. Пусть данные — $\mathcal N(0, 1)$, а энергия — $E_v(x) = x^2/(2v)$, так что $p_v = \mathcal N(0, v)$ и правильный ответ — $v = 1$. Поскольку $\partial E_v/\partial v = -x^2/(2v^2)$, градиент CD равен

$$\frac{\partial \mathcal L_{\text{CD}}}{\partial v} = -\frac{1}{2v^2}\Big(\mathbb E\,(x^+)^2 - \mathbb E\,(x^-)^2\Big) .$$

Обучение останавливается там, где второй момент негативных сэмплов равен второму моменту данных, $\mathbb E(x^-)^2 = 1$, какое бы $v$ для этого ни потребовалось. Пусть негативные сэмплы получаются $K$ шагами Ланжевена из $U(-3, 3)$, дисперсия которого равна 3. Каждый шаг умножает $x$ на $a = 1 - \varepsilon/(2v)$, где $0 < a < 1$ при $v > \varepsilon/2$, и добавляет независимый шум дисперсии $\varepsilon$, поэтому после $K$ шагов

$$\operatorname{Var} x^- = 3\,a^{2K} + \varepsilon\,\frac{1 - a^{2K}}{1 - a^2} .$$

При $\varepsilon = 0.01$ и $K = 10$ цепочка едва успевает сжать шум, и довести дисперсию её выхода до 1 можно только крутой энергией: CD останавливается на $v = 0.088$, то есть на модели со стандартным отклонением 0.30 вместо 1. При $K = 100$ она останавливается на $v = 0.58$. При $K = 1000$ — на $0.997$, практически на предельном значении $0.9975$ для нескорректированной динамики Ланжевена с этим размером шага; оно немного меньше 1 из-за смещения от дискретизации. Сэмплы короткой цепочки совпадают с данными по дисперсии — единственной статистике, которую видит эта энергия; плотность модели с данными не совпадает. Nijkamp et al. (2019) описывают тот же эффект для EBM изображений, обученных короткими цепочками из шума: получается плохая плотность, но хороший сэмплер — фиксированное число шагов Ланжевена, превращающее шум в реалистичные изображения.

Рисунок

Короткие цепочки учат неправильную плотность

Персистентные цепочки и буфер воспроизведения#

Цепочки, перезапускаемые на каждом шаге, выбрасывают работу предыдущего шага. Персистентная контрастивная дивергенция (persistent contrastive divergence; Tieleman, 2008) их сохраняет: каждый шаг обучения продолжает цепочки с того места, где остановился предыдущий. Параметры между шагами меняются мало, поэтому цепочка, близкая к равновесию для старой модели, остаётся близкой к нему и для новой, а эффективная длина цепочек растёт в ходе обучения. Буфер воспроизведения (replay buffer) обобщает это до пула прошлых сэмплов (Du and Mordatch, 2019, раздел 3.2). Цепочка начинается с сохранённого сэмпла с вероятностью 0.95 и с равномерного шума в остальных случаях, делает $K$ шагов Ланжевена, и её конечная точка возвращается в буфер. Сохранённые сэмплы позволяют продолжать уточнение от одного шага обучения к другому, так что хватает коротких цепочек. Свежий шум продолжает исследовать пространство и не даёт буферу схлопнуться на нескольких изображениях. Du и Mordatch сочли буфер необходимым: на непрерывном MNIST обучение с обычными персистентными цепочками вместо буфера дало заметно худшее правдоподобие, а к размеру буфера свыше примерно 10 000 сэмплов обучение было нечувствительно (приложения A.3 и A.4).

В гауссовском примере буфер делает возраст сэмпла геометрическим: сэмпл прошёл $n$ раундов по $K$ шагов с вероятностью $0.05 \cdot 0.95^{\,n-1}$, в среднем 20 раундов. Тогда CD останавливается на $v = 0.52$ при $K = 10$ вместо $0.088$ и на $0.94$ при $K = 100$ вместо $0.58$ — синяя кривая на рисунке выше. Свежие 5% сохраняют смещение: они попадают в негативную фазу всего после $K$ шагов.

Один шаг обучения#

Конфигурация, которая обучает небольшую свёрточную энергию на MNIST (четыре свёртки с шагом 2 и активациями Swish и два линейных слоя, 76 993 параметра; пиксели в $[-1, 1]$), собирает эти части вместе. Рисунок показывает один шаг; номера пунктов ниже совпадают с номерами на рисунке.

Рисунок

Шаг обучения EBM: буфер, цепочки Ланжевена, две фазы

1. Старты. 128 цепочек. Каждая начинается с сохранённого сэмпла с вероятностью 0.95 или с равномерного шума на $[-1, 1]^{784}$ с вероятностью 0.05 — примерно 6.4 свежих старта на батч. Буфер хранит до 8192 изображений.
2. Шаги Ланжевена. 60 шагов при фиксированных параметрах сети. Каждый шаг добавляет гауссов шум со стандартным отклонением 0.005, ограничивает значения пикселей отрезком $[-1, 1]$, сдвигает точку против градиента энергии с размером шага 10, предварительно обрезав каждую компоненту градиента до $\pm 0.03$, и снова ограничивает значения пикселей. Градиентный шаг сдвигает пиксель не более чем на 0.3.
3. Две фазы. Сеть вычисляет энергии батча из 128 изображений данных $x^+$ и 128 сэмплов $x^-$.
4. Функция потерь и обновление. $\mathcal L = \text{mean}\,E_\theta(x^+) - \text{mean}\,E_\theta(x^-) + \alpha\,\text{mean}\big(E_\theta(x^+)^2 + E_\theta(x^-)^2\big)$ с $\alpha = 0.1$ (этот штраф объяснён в следующем разделе); её минимизирует Adam с $\beta_1 = 0$ и шагом обучения $10^{-4}$.
5. Запись обратно. Сэмплы помещаются в начало буфера; сверх 8192 изображений самые старые выбрасываются.

Настройки Du и Mordatch для CIFAR-10 близки: 60 шагов размера 10, компоненты градиента обрезаются до 0.01, шум со стандартным отклонением 0.005, буфер на 10 000 изображений, вес штрафа 1, Adam с $\beta_1 = 0$, $\beta_2 = 0.999$ и шагом обучения $10^{-4}$, спектральная нормализация во всех слоях; изображения отмасштабированы в $[0, 1]$ (Du and Mordatch, 2019, приложение A.11).

Генерация после обучения. Новые изображения получаются цепочками, которые начинаются с равномерного шума и работают дольше, чем при обучении, например 256 шагов вместо 60: генерация выполняется один раз, а не на каждом обновлении, и буфера, из которого можно было бы стартовать, у неё нет. Du и Mordatch, напротив, на этапе тестирования сэмплируют из буфера воспроизведения с несколько большим шумом, чем при обучении, потому что при слабом шуме цепочки оседают на немногих изображениях высокой вероятности (приложение A.2).

Проблема сдвига: почему энергии дрейфуют#

Функция потерь CD зависит только от разностей энергий. Если сдвинуть все энергии на одну и ту же константу, $\mathcal L_{\text{CD}}$ не изменится, как и $p_\theta$. В примере с тремя состояниями компоненты градиента в сумме дают ноль; для сети, последний слой которой имеет смещение $b$, каждая энергия зависит от $b$ с производной 1, поэтому $\partial\mathcal L_{\text{CD}}/\partial b = 1 - 1 = 0$ точно. Ничто в функции потерь не удерживает общий уровень энергий, и остальные параметры свободно двигают его как побочный эффект своих обновлений. Масштаб тоже не удерживается: всякий раз, когда сэмплер отстаёт от модели, сеть может увеличить разрыв $E_\theta(x^-) - E_\theta(x^+)$, увеличив все энергии по модулю. На практике разность между позитивными и негативными энергиями остаётся осмысленной, а сами значения дрейфуют к численно неустойчивым величинам (Du and Mordatch, 2019, раздел 3.3).

Штраф L2. Добавим малое кратное $\alpha \ll 1$ среднего квадрата энергии обоих батчей:

$$\begin{aligned} \mathcal L(\theta) = {}& \frac1N\sum_{i=1}^N \Big[E_\theta(x_i^+) - E_\theta(x_i^-)\Big]\\ &+ \frac{\alpha}{N}\sum_{i=1}^N\Big[E_\theta(x_i^+)^2 + E_\theta(x_i^-)^2\Big] . \end{aligned}$$

Он делает две вещи. Когда сэмплы похожи на данные, у обоих батчей одинаковое распределение энергий, слагаемое CD в среднем обращается в ноль, и штраф притягивает общий уровень к нулю, так что энергии данных оказываются около 0. Когда сэмплы не похожи на данные, штраф ограничивает разрыв.

Разобранный пример. Предположим, что сеть может задать одну энергию $E^+$ для всех данных и одну энергию $E^-$ для всех сэмплов. Она минимизировала бы $E^+ - E^- + \alpha\big((E^+)^2 + (E^-)^2\big)$. Приравнивание обеих производных к нулю даёт $1 + 2\alpha E^+ = 0$ и $-1 + 2\alpha E^- = 0$, так что $E^\pm = \mp 1/(2\alpha)$. При $\alpha = 0.1$ энергии останавливаются на $\mp 5$, разрыв равен 10; при $\alpha = 1$ — на $\mp 0.5$. Без штрафа минимума нет: функция потерь неограниченно убывает при $E^+ \to -\infty$ и $E^- \to +\infty$.

Штраф не входит в правдоподобие и немного сдвигает оптимум, поэтому $\alpha$ держат малым: 0.1 подходит для конфигурации MNIST выше, а Du и Mordatch используют 1 на CIFAR-10. Они сочетают штраф со спектральной нормализацией всех слоёв, которая ограничивает константу Липшица энергии и не даёт её градиентам резко меняться (заметка о WGAN). Du и Mordatch также приписывают обоим приёмам сохранение интегрируемости $e^{-E_\theta}$ (раздел 3.3).

Практические решения#

• Гладкие активации. Сэмплер следует $\nabla_x E_\theta$. Для сети с ReLU энергия кусочно-линейна по $x$, и её градиент по входу скачком меняется на границах линейных областей. Гладкие активации, такие как Swish, $x\,\sigma(x)$, сохраняют поле градиентов гладким; Du и Mordatch сообщают, что Swish заметно улучшает устойчивость (приложение A.12).
• Никаких слоёв нормализации. Du и Mordatch обнаружили, что батч-нормализация (batch normalization), групповая нормализация, нормализация слоя, попиксельная и другие нормализации, по-видимому, ухудшают сэмплирование, и предполагают, что они делают шаг Ланжевена зависимым от других точек данных (приложение A.12). Этот механизм подходит для батч-нормализации, статистики которой смешивают примеры батча. Групповая, послойная и попиксельная нормализации считают статистики внутри одного примера, такой зависимости не создают, и почему они вредят, не установлено.
• Без импульса. Негативные сэмплы на каждом шаге порождаются заново из текущей модели, поэтому поверхность функции потерь движется вместе с параметрами, а импульс (momentum) продолжает толкать в направлениях, вычисленных для прежней поверхности. Обычная настройка — Adam с $\beta_1 = 0$ (приложение A.11).
• Немного шума на данных. Сэмплы всегда несут шум сэмплера, а многие пиксели данных лежат точно на границе диапазона (чёрный фон MNIST). Добавление к данным сравнимого шума — со стандартным отклонением 0.005 в конфигурации MNIST и 0.03 в JEM (Grathwohl et al., 2019, приложение A) — не даёт энергии различать их по такому тривиальному признаку.
• Смотреть на сэмплы, а не на функцию потерь. Слагаемое CD и штраф рано устанавливаются около нуля и остаются там, потому что негативные сэмплы следуют за моделью. Прогресс виден в другом: в изображениях, которые каждые несколько эпох генерируются из шума, в содержимом буфера — это то, что модель сейчас считает негативными примерами, — и в средней энергии равномерного шума, которая должна продолжать расти.

Неустойчивость#

Обучение EBM легко расходится (Grathwohl et al., 2019, раздел 6 и приложение H.3). Типичный сбой начинается, когда на поверхности энергии появляется много острых локальных минимумов. Цепочки из шума падают в ближайший из них и остаются похожими на шум. Негативная фаза продолжает поднимать энергию там, где они находятся, а позитивная — понижать её на данных, так что разрыв растёт, а сэмплы перестают улучшаться. В итоге модель приписывает высокую вероятность обучающим данным, сэмплер выдаёт только шум, а поле градиентов становится бесполезным для сэмплирования. Тревожный признак — негативные сэмплы, энергии которых намного больше, чем у обучающих данных (Du and Mordatch, 2019, приложение A.12).

Средства, примерно от самых дешёвых:

• штраф L2 и спектральная нормализация, которые ограничивают энергии и их градиенты;
• меньший шаг обучения или больше шагов Ланжевена на обновление. По наблюдениям Grathwohl et al., только эти два изменения надёжно стабилизировали расходящуюся модель. В их масштабе цепочкам, перезапускаемым из шума, требовалось 80 шагов на обновление для устойчивого обучения, и даже такие модели расходились на поздних этапах обучения; персистентные цепочки с 20 шагами были в четыре раза быстрее, но менее устойчивы (приложения H.2 и H.3);
• перезапуск с более ранней контрольной точки, которая ещё не разошлась, возможно, с другим начальным значением генератора (seed) (приложение H.3).

Самые чувствительные гиперпараметры — параметры сэмплера: размер шага, число шагов и уровень шума, за ними следуют шаг обучения и ширина сети. Du и Mordatch нашли самым чувствительным размер шага, с пригодным диапазоном в два-три порядка величины (приложение A.12).

Оценка качества модели без нормировочной константы#

Без $Z_\theta$ правдоподобие сообщить нельзя. Выборка по значимости с отжигом (annealed importance sampling; Neal, 2001) оценивает нормировочную константу, переходя от простого распределения $p_0$ к $p_\theta$ через промежуточные плотности, пропорциональные $p_0(x)^{1-\beta}\,e^{-\beta E_\theta(x)}$, для $\beta$ от 0 до 1, с шагами MCMC на каждом этапе и весами значимости (importance weights), которые перемножаются вдоль пути; средний вес — несмещённая оценка $Z_\theta/Z_0$. Метод работает на небольших задачах. Du and Mordatch (2019) получили так нижнюю и верхнюю оценки правдоподобия непрерывного MNIST, но на CIFAR-10 более 300 000 цепочек, работавших два дня, всё ещё оставляли большой зазор между оценками (приложение A.3). Поэтому EBM изображений оценивают по их сэмплам, например с помощью FID, и по задачам из последних разделов.

Шумоконтрастное оценивание: Z через классификацию#

Двум оценкам из этого и следующего разделов не нужны сэмплы из модели во время обучения. Шумоконтрастное оценивание (noise-contrastive estimation, NCE; Gutmann and Hyvärinen, 2010) считает логарифм нормировочной константы ещё одним параметром. Модель —

$$\log p_m(x) = -E_\theta(x) + c ,$$

где $c$ оценивает $-\log Z_\theta$. Выберем шумовое распределение (noise distribution) $p_n$, из которого легко сэмплировать и плотность которого можно вычислить, например гауссово. Возьмём столько же сэмплов шума, сколько сэмплов данных, и обучим логистическую регрессию их различать, зафиксировав логит равным логарифму отношения плотностей модели и шума:

$$\begin{gathered} h(x) = \sigma\big(\log p_m(x) - \log p_n(x)\big) = \frac{p_m(x)}{p_m(x) + p_n(x)},\\ J(\theta, c) = \tfrac12\,\mathbb E_{x\sim P}\log h(x)\\ + \tfrac12\,\mathbb E_{y\sim p_n}\log\big(1 - h(y)\big) . \end{gathered}$$

$J$ — логарифм правдоподобия меток «данные» и «шум» (Gutmann and Hyvärinen, 2010, раздел 2.2). Почему классификатор выучивает плотность? При равном числе сэмплов данных и шума истинная вероятность того, что точка $x$ взята из данных, равна $p(x)/(p(x) + p_n(x)) = \sigma\big(\log p(x) - \log p_n(x)\big)$. У лучшего классификатора $\log p_m = \log p$ в точности, включая нормировку. Gutmann и Hyvärinen показывают, что $J$ как функция от $\log p_m$ имеет максимум в $\log p$ и других экстремумов не имеет, если $p_n > 0$ всюду, где $p > 0$ (теорема 1), и что оценка $(\theta, c)$ состоятельна при условиях, аналогичных условиям для метода максимального правдоподобия (теорема 2). Поэтому в популяционном оптимуме правильно специфицированной модели, то есть такой, семейство которой содержит плотность данных, подогнанная функция $e^{-E_\theta + c}$ интегрируется в единицу без всяких ограничений, а $c$ равно $-\log Z_\theta$. На конечной выборке $c$ — лишь оценка, и нормировка выполняется приближённо, а для неверно специфицированной энергии лучшая подгонка может и вовсе не интегрироваться в единицу: тогда $c$ — обучаемый параметр, а не численная гарантия.

Метод максимального правдоподобия так не может: при свободном $c$ правдоподобие неограниченно растёт при $c \to \infty$, поскольку ничто не заставляет $e^{-E_\theta + c}$ интегрироваться в единицу (раздел 1). Целевая функция классификации заставляет. Слишком большое $c$ заставляет классификатор называть «данными» и шум, и за это он платит на сэмплах шума.

Разобранный пример. Данные — $\mathcal N(0, 1)$, шум — $p_n = \mathcal N(0, 4)$, энергия имеет правильную форму, $E(x) = x^2/2$, так что неизвестно только $c$. Правильное значение — $c^* = -\log\sqrt{2\pi} = -0.919$. В точке $x = 0$ истинная вероятность «данных» равна $0.399/(0.399 + 0.199) = 0.667$, и с $c^*$ классификатор даёт $h(0) = \sigma(-0.919 + 1.612) = 0.667$. Ненормированная модель, $c = 0$, даёт $h(0) = \sigma(1.612) = 0.834$: она слишком уверена, поэтому логистическая функция потерь уменьшает $c$. Численная максимизация по $c$ точной целевой функции $J$ (по всей генеральной совокупности, а не по выборке) даёт $-0.919$.

Шумовое распределение определяет, сколько модель выучит. Из него должно быть легко сэмплировать, его плотность должна вычисляться, и оно должно быть близко к данным. Если шум совпадает с распределением данных, асимптотическая ковариация оценки вдвое больше границы Крамера–Рао (раздел 2.4). Далёкий шум делает задачу тривиальной: равномерный шум отличается от цифр по нескольким грубым признакам, таким как гладкость или насыщенные пиксели фона, и энергия не выучивает ничего тоньше. Лучший шум даёт предварительная модель данных, например поток (Song and Kingma, 2021, раздел 5). С точки зрения заметки о GAN NCE — это дискриминатор, обученный против фиксированного генератора $p_n$, логит которого ограничен видом $\log p_m - \log p_n$; это ограничение и превращает дискриминатор в плотность.

Score matching: без сэмплов и без Z#

Score модели (градиент логарифма плотности по входу) у EBM не содержит $Z_\theta$: $\nabla_x \log p_\theta(x) = -\nabla_x E_\theta(x)$, как показано выше. Поэтому подгонка score модели к score данных минимизацией дивергенции Фишера (Fisher divergence) $\tfrac12\,\mathbb E_P\big\|\nabla_x \log p(x) + \nabla_x E_\theta(x)\big\|^2$ обучает энергию без нормировочной константы и без сэмплов (Hyvärinen, 2005); этот подход называют score matching (сопоставление градиентов логарифма плотности). Score данных неизвестен. Заметка о score matching убирает его интегрированием по частям и выводит более дешёвые срезанную и денойзинговую версии. Для энергии целевая функция Хюваринена — это $\mathbb E_P\big[\tfrac12\|\nabla_x E_\theta(x)\|^2 - \Delta_x E_\theta(x)\big]$ с точностью до константы, где $\Delta_x$ — лапласиан, след гессиана; денойзинговый score matching регрессирует $-\nabla_x E_\theta(x + \sigma\varepsilon)$ на $-\varepsilon/\sigma$. В обоих случаях функция потерь содержит градиент энергии по входу, так что каждый шаг обучения дифференцирует дважды, а для генерации всё равно нужна динамика Ланжевена. Две оценки без сэмплирования связаны: NCE, шум которого — данные, сдвинутые на малый вектор, со слегка изменённым классификатором переходит в срезанный score matching, когда сдвиг стремится к нулю (Song and Kingma, 2021, раздел 5.1).

Обнаружение данных вне распределения#

Энергия, обученная с негативной фазой, низка на данных и поднята везде, куда приходили её собственные сэмплы. Поэтому $E_\theta(x)$ — естественный детектор необычных входов: будем считать $x$ лежащим вне распределения (out of distribution), если его энергия превышает порог, выбранный на отложенных данных. Уверенность классификатора — слабая замена, поскольку современные классификаторы склонны к излишней уверенности (Grathwohl et al., 2019, раздел 5.2). Есть четыре оговорки.

• Высокая плотность — не типичность. Потоки, VAE и PixelCNN, обученные на CIFAR-10, приписывают изображениям номеров домов из SVHN более высокое правдоподобие, чем тестовым изображениям CIFAR-10 (Nalisnick et al., 2019). Модель правдоподобия не становится автоматически детектором выбросов.
• EBM справляются лучше, но не хорошо. По AUROC отделения тестовых изображений CIFAR-10 от других данных (0.5 — уровень случайного угадывания) EBM получает 0.63 против SVHN, где у PixelCNN++ 0.32, а у Glow 0.24; в среднем по пяти видам выбросов — 0.62 против 0.47 и 0.42. Та же EBM получает 0.30 на однотонных изображениях и 0.48 на текстурах — хуже случайного угадывания (Du and Mordatch, 2019, раздел 4.4 и рисунок 10).
• Поднято только то, что посетил сэмплер. Равномерный шум, с которого стартуют цепочки, получает высокую энергию. Изменения, которых сэмплер никогда не порождал, например повёрнутая или уменьшенная цифра, могут повысить энергию лишь немного.
• У энергий нет общей шкалы. Порог нужно калибровать для каждой модели; штраф L2 держит энергии данных около нуля, и это помогает его выбрать.

Grathwohl et al. (2019) предлагают для EBM другой показатель — норму градиента по входу $\|\nabla_x \log p_\theta(x)\|$, исходя из интуиции, что типичный вход лежит там, где плотность меняется медленно, так что большой градиент выдаёт точку с высокой плотностью, но малой массой (раздел 5.3.3). Эта интуиция — не общее свойство типичных множеств. Для $\mathcal N(0, I_d)$ score равен $-x$: на типичной оболочке его норма около $\sqrt d$, 28 при $d = 784$, а в центре — точке наибольшей плотности, причём нетипичной, — она равна 0. Поэтому норма градиента — эмпирический детектор, поведение которого зависит от модели и данных, и его порог и направление нужно калибровать на отложенных примерах из распределения и вне его, как и для самой энергии.

JEM: классификатор — скрытая EBM#

Классификатор с $C$ классами выдаёт логиты $f_\theta(x) \in \mathbb R^C$ и вероятности классов $p_\theta(y \mid x) = \operatorname{softmax}(f_\theta(x))_y$. Grathwohl et al. (2019) читают те же логиты как отрицательную энергию пары $(x, y)$:

$$p_\theta(x, y) = \frac{e^{f_\theta(x)[y]}}{Z_\theta}, \qquad E_\theta(x, y) = -f_\theta(x)[y] .$$

Суммирование по $y$ даёт плотность одного только входа,

$$\begin{aligned} p_\theta(x) &= \frac{1}{Z_\theta}\sum_y e^{f_\theta(x)[y]},\\ E_\theta(x) &= -\log\sum_y e^{f_\theta(x)[y]}\\ &= -\operatorname{logsumexp}_y f_\theta(x)[y] , \end{aligned}$$

а деление совместного распределения на маргинальное сокращает $Z_\theta$ и возвращает softmax: $p_\theta(y \mid x)$ — тот же классификатор, что и раньше. Прибавление одной и той же константы ко всем логитам входа не меняет классификатор, но меняет $p_\theta(x)$. В этой степени свободы, которую обычное обучение не использует, и живёт плотность; такую модель называют совместной моделью, основанной на энергии (joint energy-based model, JEM).

Разобранный пример. Логиты $(2, 0, -1)$ дают вероятности классов $(0.844, 0.114, 0.042)$, $\operatorname{logsumexp} = 2.170$ и $E(x) = -2.170$. Логиты $(3, 1, 0)$ дают те же вероятности и $E(x) = -3.170$. Классификатор не различает эти два входа; по $p_\theta(x)$ второй в $e \approx 2.72$ раза вероятнее.

JEM обучается на $\log p_\theta(x, y) = \log p_\theta(x) + \log p_\theta(y \mid x)$: кросс-энтропия для слагаемого классификатора, которое вычисляется точно, плюс градиент EBM для $\log p_\theta(x)$ с энергией $-\operatorname{logsumexp}$; негативные сэмплы получаются шагами Ланжевена из буфера воспроизведения на 10 000 изображений, 20 шагов на обновление, 5% цепочек перезапускаются из равномерного шума, шум со стандартным отклонением 0.01 (раздел 4 и приложение A). WideResNet-28-10 без батч-нормализации, обученная так на CIFAR-10, достигает доли правильных ответов (accuracy) 92.9% при Inception Score 8.76 и FID 38.4 (таблица 1). Совместное обучение также улучшает калибровку, обнаружение данных вне распределения и устойчивость к состязательным примерам (adversarial robustness) по сравнению с тем же классификатором, обученным отдельно (раздел 5). Цена — та же, что у любой EBM: каждый шаг Ланжевена стоит примерно столько же, сколько шаг обучения классификатора, а модели регулярно расходились, и их приходилось перезапускать с меньшим шагом обучения или более сильной регуляризацией (разделы 6 и 7).

Композиция энергий#

Сложение энергий перемножает плотности:

$$\begin{aligned} E_{A+B}(x) &= E_A(x) + E_B(x)\\ \Longrightarrow\quad p_{A+B}(x) &\propto e^{-E_A(x)}\,e^{-E_B(x)} \propto p_A(x)\,p_B(x) , \end{aligned}$$

— это произведение экспертов (product of experts; Hinton, 2002). Произведение велико только там, где велики оба множителя, поэтому сумма энергий действует как логическое И для ограничений, которые кодируют энергии. Score суммы равен сумме score, поэтому динамике Ланжевена на составной модели нужны только градиенты частей и никакого нового обучения.

Разобранный пример. $E_A(x) = x^2/2$ — это $\mathcal N(0, 1)$, а $E_B(x) = (x - 2)^2/2$ — это $\mathcal N(2, 1)$. Их сумма равна $x^2 - 2x + 2 = (x - 1)^2 + 1$, то есть энергии $\mathcal N(1, \tfrac12)$: произведение лежит между ними и уже каждого из них.

Du and Mordatch (2019) обучают отдельные условные энергии для масштаба, положения и поворота фигуры и генерируют изображения со всеми тремя атрибутами динамикой Ланжевена на сумме; суммы энергий, дообученные на обучающих выборках, порождают и сочетания размера и положения, которые при обучении вместе не встречались (раздел 7). Части должны иметь сравнимые масштабы. Отдельно обученные энергии измеряются в произвольных единицах (LeCun et al., 2006, раздел 1.3), и вес одного слагаемого — или, что то же самое, температура — меняет результат.

Что изменилось#

	Критик WGAN	EBM, обучение с MCMC	NCE	Score matching
Выход сети	значение критика $f(x)$	энергия $E_\theta(x)$	энергия и константа $c$	энергия; её градиент по входу — score
Негативные примеры	от сети-генератора	цепочки Ланжевена на $E_\theta$, часто из буфера	сэмплы фиксированной шумовой плотности	нет
Градиент через сэмплер	да, генератор обучается	нет	нет	сэмплера нет
Нормировочная константа $Z_\theta$	не определена	никогда не вычисляется	выучивается как $-c$	сокращается
Что держит обучение под контролем	ограничение Липшица	L2 на энергиях, спектральная нормализация	шум, близкий к данным	вторые производные или уровень шума
Сэмплирование после обучения	один проход генератора	динамика Ланжевена	динамика Ланжевена	динамика Ланжевена
Типичный сбой	критик не липшицев	расходимость; плотность, сформированная сэмплером	тривиальная классификация против далёкого шума	стоимость следа гессиана; score зашумлённых данных

Ограничения#

• Сэмплирование внутри обучения. Каждое обновление выполняет для сэмплера $K$ прямых и обратных проходов сети, 60 в конфигурации MNIST, так что обучение в десятки раз медленнее обучения классификатора того же размера (Grathwohl et al., 2019, раздел 7).
• Смещённые градиенты. Короткие цепочки, как и ослабленный сэмплер с температурой около $10^{-6}$, не сэмплируют $p_\theta$. Выученная плотность сформирована сэмплером, как показывает гауссовский пример, и не является калиброванным правдоподобием.
• Неустойчивость. Обучение расходится без штрафа, спектральной нормализации и тщательной настройки сэмплера, а иногда и несмотря на них.
• Правдоподобия бесплатно нет. Чтобы сообщить $\log p_\theta(x)$, нужна $Z_\theta$, а выборка по значимости с отжигом слишком медленна для изображений размера CIFAR-10.
• Медленная генерация. Сотни шагов Ланжевена на сэмпл, а длинные цепочки медленно исследуют моды; по словам Du и Mordatch, исследование мод динамикой Ланжевена для безусловного CIFAR-10 «заняло очень много времени» (took a very long time; раздел 4.1).
• У оценок без сэмплирования свои издержки. NCE зависит от шумового распределения, близкого к данным, а для изображений такое трудно найти; score matching требует вторых производных или уровня шума, а для сэмплирования всё равно нужна динамика Ланжевена.

Итоги#

• Сеть энергии задаёт $p_\theta(x) = e^{-E_\theta(x)}/Z_\theta$. Значение имеют только разности энергий; $Z_\theta$ — вычислительно неразрешимый (intractable) интеграл.
• Градиент правдоподобия равен $\mathbb E_P\nabla_\theta E_\theta(x^+) - \mathbb E_{p_\theta}\nabla_\theta E_\theta(x^-)$: понизить энергию на данных, повысить на сэмплах модели. Негативной фазе нужны сэмплы из модели, и их даёт динамика Ланжевена по $-\nabla_x E_\theta$.
• Контрастивная дивергенция останавливает цепочки через $K$ шагов. Короткие цепочки смещают модель: она выучивает энергию, $K$-шаговый сэмплер которой воспроизводит данные, а не плотность данных. Персистентные цепочки и буфер воспроизведения (95% старых сэмплов, 5% свежего шума) удлиняют цепочки без дополнительных затрат.
• Функция потерь видит только разности энергий, поэтому уровень и масштаб энергий дрейфуют; штраф L2 $\alpha\,\text{mean}(E^2)$ прижимает уровень к нулю и ограничивает разрыв величиной порядка $1/\alpha$.
• Обучение неустойчиво; помогают гладкие активации, отказ от слоёв нормализации, Adam без импульса, малые шаги обучения и больше шагов сэмплирования.
• NCE выучивает $E_\theta$ и $\log Z_\theta$, классифицируя данные против шума; score matching подгоняет $-\nabla_x E_\theta$ к score данных. Ни один из этих методов не сэмплирует во время обучения.
• Энергии обнаруживают выбросы (с оговорками), любой классификатор — это EBM с $E(x) = -\operatorname{logsumexp}_y f_\theta(x)[y]$ (JEM), а сложение энергий перемножает плотности.

Ноутбук#

Ноутбук об EBM из репозитория HSE-LAMBDA Deep Generative Models обучает конфигурацию MNIST из этой заметки (ячейки нумеруются с 1, считая все ячейки подряд, включая текстовые). Он следует Tutorial 8, Deep Energy-Based Generative Models, из курса глубокого обучения Амстердамского университета, автор которого — Phillip Lippe, и использует PyTorch Lightning. В ноутбуке плотность модели обозначена $q_\theta$, а в этой заметке — $p_\theta$.

• Ячейки 1–3: теория. Энергия и статистическая сумма; score и score matching, с упоминанием денойзингового и срезанного score matching; NCE с отношением $r_\theta(x) = e^{-E_\theta(x)}/q_0(x)$ для шумовой плотности $q_0$; вывод градиента правдоподобия, как в этой заметке; сэмплирование Ланжевена с $K$ от 10 до 60. Ячейка 2 также упоминает SGLD и выборку по значимости с отжигом, которые ноутбук не реализует.
• Ячейки 4–6: подготовка. Импорты, начальное значение генератора 42, устройство (CUDA, Apple MPS или CPU) и загрузка предобученной контрольной точки MNIST из исходного руководства вместе с журналом TensorBoard. Если контрольная точка есть, ячейка 27 загружает её и пропускает обучение.
• Ячейка 7: применения помимо генерации. Классификация и шумоподавление как минимизация $E(X, Y)$ по ответу $Y$, со встроенным рисунком.
• Ячейки 8–9: данные. MNIST, 60 000 обучающих и 10 000 тестовых изображений $28 \times 28$, отмасштабированных в $[-1, 1]$, чтобы сэмплер мог ограничивать значения фиксированным диапазоном. Загрузчик обучающих данных использует размер батча 128, перемешивание и drop_last=True (468 батчей за эпоху); тестовый загрузчик с размером батча 256 служит валидационной выборкой.
• Ячейки 10–12: модель. CNNModel: свёртки $1 \to 16$ (ядро 5, шаг 2, дополнение 4: $28 \to 16$), $16 \to 32$, $32 \to 64$ и $64 \to 64$ (ядро 3, шаг 2: $8$, $4$, $2$), Swish после каждой, затем $256 \to 64 \to 1$ со Swish между ними; 76 993 параметра, слоёв нормализации нет. Ячейка 12 указывает, что сеть выдаёт $f_\theta = -E_\theta$, так что «score» в ноутбуке — это отрицательная энергия.
• Ячейки 13–14: сэмплер и буфер. Sampler начинает с буфера из 128 равномерных изображений. Каждый вызов берёт $n_{\text{new}} \sim \text{Binomial}(128, 0.05)$ свежих равномерных изображений, а остальные выбирает из буфера с возвращением (random.choices), затем запускает generate_samples: параметры заморожены, модель в режиме оценки, 60 шагов, каждый из которых — шум (стандартное отклонение 0.005), ограничение диапазона, градиентный шаг размера 10 по $E_\theta = -f_\theta$ с обрезкой каждой компоненты до $\pm 0.03$ и снова ограничение диапазона. Новые сэмплы добавляются в начало, а буфер обрезается до 8192 изображений; этого размера он достигает после 63 шагов обучения.
• Ячейки 15–19: обучение. Ячейка 16 объясняет проблему сдвига, штраф L2 с весом $\alpha$ и даёт практические советы (баланс размера шага и шума, $\alpha$, равное 0.1 или 0.01, короткие и длинные цепочки, большой буфер). DeepEnergyModel (ячейка 18) использует Adam с шагом обучения $10^{-4}$ и betas $(0, 0.999)$ и ступенчатое уменьшение шага обучения: умножение на 0.97 после каждой эпохи. Каждый шаг обучения добавляет к данным шум со стандартным отклонением 0.005, берёт 128 сэмплов и вычисляет cdiv_loss = mean f(x⁻) − mean f(x⁺), то есть $\text{mean}\,E(x^+) - \text{mean}\,E(x^-)$, и reg_loss $= \alpha\,\text{mean}(f(x^+)^2 + f(x^-)^2)$ со значением по умолчанию $\alpha = 0.1$. Валидация сравнивает тестовые изображения с равномерным шумом и записывает в журнал $\text{mean}\,f(\text{noise}) - \text{mean}\,f(\text{test})$. Тестового шага нет (ячейка 19).
• Ячейки 20–25: колбэки. GenerateCallback каждые 5 эпох делает 256 шагов размера 10 из равномерного шума для 8 изображений и записывает 8 промежуточных стадий; SamplerCallback каждые 5 эпох записывает 32 случайных изображения из буфера; OutlierCallback каждую эпоху записывает среднее $f_\theta$ по 1024 равномерным изображениям.
• Ячейки 26–28: запуск. Trainer из Lightning на 60 эпох с обрезкой нормы градиента на уровне 0.1, контрольной точкой, выбираемой по наименьшему значению на валидации (это эпоха, которая лучше всего отделяет тестовые цифры от равномерного шума, а не эпоха с лучшими сэмплами), тремя колбэками и монитором шага обучения. Одна эпоха — это 468 шагов обучения по 60 проходов сэмплера, то есть 28 080 проходов сети на батчах из 128 изображений.
• Ячейки 29–34: TensorBoard. Загружается скачанный журнал. Ячейка 33 встраивает снимок экрана; ячейка 34 отмечает, что слагаемое CD и штраф быстро сходятся примерно к нулю, а прогресс виден в сэмплах и в падающем «score» случайных изображений.
• Ячейки 35–39: генерация. С начальным значением генератора 43 четыре цепочки из равномерного шума делают по 256 шагов размера 10; график показывает старт и каждый 32-й шаг. Ячейка 39 описывает графики исходного руководства (восьмёрку и двойку), а не графики, сохранённые здесь.
• Ячейки 40–54: обнаружение данных вне распределения. Ячейка 40 обосновывает задачу и ссылается на JEM. Ячейки 41 и 43 печатают среднее $f_\theta$ по 128 равномерным изображениям и по обучающему батчу; ячейка 44 говорит, что второе близко к 0 из-за штрафа. Ячейки 45–53 сравнивают первое тестовое изображение с тремя изменёнными версиями: с гауссовым шумом со стандартным отклонением 0.3 и ограничением диапазона (ячейка 49); с flip(dims=(1, 2)) (ячейка 51); с цифрой, прореженной в 2 раза и помещённой в правый нижний квадрант чёрного изображения (ячейка 53). Текст сообщает о большом падении «score» для шума и о небольших — для двух других версий.
• Ячейка 55: неустойчивость. Расхождение, как описано в этой заметке; перезагрузка устойчивой контрольной точки как крайняя мера; чувствительные гиперпараметры (размер шага, число шагов, шум, шаг обучения, ширина).
• Ячейки 56–57: заключение; ячейка 57 пустая.

Ни у одной ячейки с кодом нет сохранённых выходов, так что ноутбук документирует только процедуру. Отрисованная страница исходного руководства показывает запуск предобученной контрольной точки: среднее значение «score» −17.88 для равномерного шума и −0.00 для обучающих изображений, а для тестового изображения 0.03 против −0.07 (шум), −0.00 (отражение) и −0.02 (уменьшение) (Lippe, 2025).

Проблемы ноутбука.

• Сэмплер — не динамика Ланжевена. Текстовые ячейки 3 и 13 записывают обновление как $x - \eta\nabla_x E + \omega$, $\omega \sim \mathcal N(0, \sigma)$, не связывая $\sigma$ с $\eta$, и утверждают, что малое $\eta$ и большое $K$ дают сэмплы из модели. Это верно, только когда дисперсия шума равна $2\eta$, то есть для $x - \frac{\varepsilon}{2}\nabla_x E + \sqrt{\varepsilon}\,z$. С шагом 10 и шумом 0.005 это обновление формально связано с $p_\theta^{1/T}$, $T = 1.25 \cdot 10^{-6}$, но с обрезкой, ограничением диапазона и 60 шагами его сэмплы не следуют никакой такой плотности, как объяснено во врезке выше.
• Контрастивная дивергенция или нет. Ячейка 3 называет метод контрастивной дивергенцией. CD у Hinton запускает цепочки из данных; ноутбук запускает их из буфера и шума, а это персистентная CD с буфером воспроизведения (Du and Mordatch, 2019).
• Шум для NCE. Ячейка 3 предлагает буферы воспроизведения как источник шума для NCE. NCE нужен логарифм плотности шума, а у буфера сэмплов MCMC его нет.
• Денойзинговый score matching (ячейка 3) описан как предсказание чистого изображения по зашумлённому; целевая функция регрессирует score модели на $-(\tilde x - x)/\sigma^2$, что эквивалентно предсказанию шума или чистого изображения с точностью до известного аффинного отображения.
• Формула L2 в ячейке 16 ставит сумму по $i$ только перед первым квадратом; она должна охватывать оба, $\frac1N\sum_i\big[E_\theta(x_i^+)^2 + E_\theta(x_i^-)^2\big]$, что и вычисляет код.
• Колбэки не запускаются в текущей версии Lightning. Три колбэка переопределяют on_epoch_end — хук, которого в PyTorch Lightning 2.x больше нет, поэтому во время обучения ничего не генерируется и не записывается; нужный хук — on_train_epoch_end.
• make_grid(..., range=(-1, 1)) в ячейках 21, 23, 38 и 45 использует аргумент, который в текущей версии torchvision заменён на value_range; вызов выбрасывает TypeError, так что графики ячейки 38 и сравнения ячеек 49, 51 и 53 не работают, пока аргумент не переименован.
• «Отражение» — это поворот. flip(dims=(1, 2)) на изображении $1 \times 28 \times 28$ переворачивает и высоту, и ширину, то есть поворачивает изображение на 180°; зеркальное отражение — это flip(dims=(2,)).
• Комментарии, не соответствующие коду. Ячейка 9 обещает разбить обучающую выборку на обучающую и валидационную части и использовать для обучения батчи меньшего размера; разбиения нет, валидационной выборкой служит тестовая, а обучение использует тот же загрузчик. Ячейка 13 упоминает равномерные или гауссовы старты; код использует только равномерный шум. Ячейка 7 повторяет номер раздела 2.5.

Дальше#

Обучение модели, основанной на энергии, запускает смещённые короткие цепочки MCMC внутри каждого обновления параметров, и получаемая модель сформирована сэмплером не меньше, чем данными. Диффузионные модели обходят этот цикл. Они фиксируют процесс зашумления, превращающий данные в гауссов шум, и учатся обращать его маленькими шагами, один за другим, а это регрессия без сэмплирования во время обучения. Следующая заметка строит этот прямой процесс. Заметка о score matching позже показывает, что сеть, которую обучает диффузионная модель, с точностью до известного множителя и есть score зашумлённых данных — отрицательный градиент энергии.

Литература#

• Y. LeCun, S. Chopra, R. Hadsell, M. Ranzato, F. J. Huang. A Tutorial on Energy-Based Learning. In G. Bakır, T. Hofmann, B. Schölkopf, A. Smola, B. Taskar (eds.), Predicting Structured Data, MIT Press, 2006. Раздел 1.1: инференс минимизацией энергии; раздел 1.3: распределение Гиббса и некалиброванные энергии; раздел 2.2.4: отрицательный логарифм правдоподобия как функция потерь.
• G. E. Hinton. Training Products of Experts by Minimizing Contrastive Divergence. Neural Computation, 14(8), 1771–1800, 2002. Раздел 3: контрастивная дивергенция.
• T. Tieleman. Training Restricted Boltzmann Machines using Approximations to the Likelihood Gradient. ICML, 2008. Персистентная контрастивная дивергенция.
• M. Gutmann, A. Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. AISTATS, PMLR 9, 297–304, 2010. Теоремы 1 и 2; раздел 2.4: выбор шума.
• A. Hyvärinen. Estimation of Non-Normalized Statistical Models by Score Matching. Journal of Machine Learning Research, 6, 2005.
• Y. Du, I. Mordatch. Implicit Generation and Modeling with Energy-Based Models. NeurIPS, 2019. Раздел 3: сэмплирование Ланжевена, буфер воспроизведения и штраф L2 (алгоритм 1); раздел 4.4: обнаружение данных вне распределения; раздел 7: композиция; приложения A.3, A.4, A.11 и A.12: оценки правдоподобия, чувствительность, гиперпараметры и советы по обучению.
• E. Nijkamp, M. Hill, S.-C. Zhu, Y. N. Wu. Learning Non-Convergent Non-Persistent Short-Run MCMC Toward Energy-Based Model. NeurIPS, 2019.
• Y. Song, D. P. Kingma. How to Train Your Energy-Based Models. arXiv:2101.03288, 2021. Раздел 3: метод максимального правдоподобия с MCMC; раздел 5: шумоконтрастное оценивание и его связь со score matching.
• W. Grathwohl, K.-C. Wang, J.-H. Jacobsen, D. Duvenaud, M. Norouzi, K. Swersky. Your Classifier is Secretly an Energy Based Model and You Should Treat it Like One. ICLR, 2020. Таблица 1, раздел 5.3 и приложения A и H.
• M. Welling, Y. W. Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamics. ICML, 2011.
• R. M. Neal. Annealed Importance Sampling. Statistics and Computing, 11(2), 125–139, 2001.
• E. Nalisnick, A. Matsukawa, Y. W. Teh, D. Gorur, B. Lakshminarayanan. Do Deep Generative Models Know What They Don’t Know? ICLR, 2019.
• P. Lippe. Tutorial 8: Deep Energy-Based Generative Models. In UvA Deep Learning Tutorials, University of Amsterdam, 2025 (год из рекомендуемой ссылки на серию).
• HSE-LAMBDA, репозиторий курса Deep Generative Models: EBM.ipynb.