Energy-based модели: как обучать без явного сэмплера

Что такое energy-based модель и зачем она нужна

Классическая генеративная модель, будь то VAE или GAN, устроена вокруг явного механизма выборки: сеть-декодер или генератор берёт вектор шума и превращает его в картинку за один прямой проход. Energy-based модель (EBM) устроена иначе. Она не строит образец напрямую, а учит скалярную функцию энергии E_theta(x), которая сопоставляет каждому объекту одно число: чем правдоподобнее объект, тем ниже энергия.

Плотность вероятности в такой модели задаётся через распределение Гиббса — Больцмана: p_theta(x) = exp(-E_theta(x)) / Z(theta), где Z(theta) — нормировочная константа, интеграл экспоненты энергии по всему пространству. Вся боль EBM спрятана именно в Z: для картинок 32x32 это интеграл по тысячам измерений, посчитать его в лоб нельзя. Отсюда и термин implicit generation — генерация неявная: чтобы получить образец, вы не вызываете сеть один раз, а запускаете итеративный процесс сэмплирования, который спускается по ландшафту энергии к её минимумам.

EBM не рисует картинку. Она описывает, где на пространстве картинок находятся низины, а достать образец из низины — уже отдельная вычислительная задача.

Как достаётся образец: Langevin dynamics

Стандартный способ неявной генерации в современных EBM — стохастическая динамика Ланжевена (SGLD, stochastic gradient Langevin dynamics^*). Идея простая: начать со случайного шума и шаг за шагом двигаться в сторону меньшей энергии, подмешивая на каждом шаге гауссов шум, чтобы не застрять в одной точке и покрыть распределение.

Один шаг выглядит так:

x_{k+1} = x_k - (alpha / 2) * grad_x E_theta(x_k) + sqrt(alpha) * eps
eps ~ N(0, I)

Здесь alpha — размер шага, grad_x E_theta — градиент энергии по входу, а не по весам. При достаточно малом шаге и большом числе итераций цепочка сходится к выборке из p_theta. На практике никто не ждёт формальной сходимости: берут от нескольких десятков до пары сотен шагов и на этом останавливаются. Это компромисс между качеством образца и временем: каждый шаг требует полного обратного распространения по сети ради градиента по входу.

Contrastive divergence и обучение

Градиент правдоподобия для EBM раскладывается на два слагаемых, и это ключ ко всему обучению:

• Положительная фаза — понижаем энергию на реальных данных из обучающей выборки.
• Отрицательная фаза — повышаем энергию на образцах, которые модель сама сгенерировала через Langevin dynamics.

Модель как бы соревнуется сама с собой: реальные данные тянет вниз, свои фантазии — вверх, пока фантазии не станут неотличимы от данных. Этот приём называется contrastive divergence. Формально шаг по весам выглядит как разность матожиданий градиента энергии по данным и по образцам модели:

grad_theta L = E_data[grad_theta E_theta(x)] - E_model[grad_theta E_theta(x)]

Первое матожидание считается по батчу реальных примеров, второе аппроксимируется образцами из Langevin-цепочки. Отсюда прямая зависимость: качество отрицательной фазы упирается в качество сэмплера. Плохие образцы — неинформативный градиент — расходящееся обучение.

Почему это тяжело обучается

EBM заслужили репутацию капризных, и причины конкретны.

Replay buffer как рабочая лошадка

Ключевой инженерный трюк, сделавший EBM обучаемыми на картинках, — persistent contrastive divergence с буфером. Вместо того чтобы каждый раз запускать Langevin с чистого шума, часть стартовых точек берут из буфера, куда складывают образцы предыдущих итераций. Цепочка продолжается там, где остановилась в прошлый раз. Это резко удешевляет отрицательную фазу: сеть уже успела продвинуть эти образцы, и коротких доработочных шагов достаточно.

Обычная схема: с вероятностью около 95% старт берётся из буфера, с оставшейся — из свежего шума, чтобы буфер не выродился в узкий набор точек.

Обобщение: чем EBM интересны за пределами генерации

Самое любопытное в EBM — то, что одна обученная функция энергии закрывает несколько задач без переобучения. Раз модель знает, где энергия низкая, она умеет:

1. Генерировать — спускаться по энергии из шума.
2. Дорисовывать (inpainting) — запускать Langevin только по замаскированной области, оставляя известную часть фиксированной.
3. Детектировать аномалии — объект с аномально высокой энергией не похож на обучающее распределение.
4. Композировать — сложение энергий нескольких моделей задаёт логическое И условий, то есть образец, удовлетворяющий сразу всем.

Последний пункт — про композиционность — то, что архитектуре с явным генератором даётся тяжело. Сумма энергий E_A(x) + E_B(x) соответствует произведению распределений, и один и тот же сэмплер достаёт образцы из пересечения. Обучили модель на энергии для "улыбка" и отдельно для "очки" — на инференсе складываете энергии и получаете лица с улыбкой и в очках, ничего не дообучая.

Явный генератор выдаёт один образец за проход, но не умеет складываться. EBM платит за генерацию временем, зато функцию энергии можно комбинировать как строительные блоки.

Где это стоит и стоит ли

Честная цена EBM — вычисления на инференсе. Сотня шагов MCMC против одного прохода у GAN означает, что генерация в десятки-сотни раз медленнее. Если вам нужен быстрый семплер картинок в продакшене, EBM почти наверняка не ваш выбор. Но если задача — единая модель под генерацию, дорисовку и поиск аномалий, плюс возможность композиции условий, разница в скорости может окупиться отсутствием зоопарка отдельных сетей.

Отдельно отмечу связь с диффузионными моделями. Score-based подход, на котором стоят современные диффузии, учит не саму энергию, а её градиент по входу — score-функцию. Это близкий родственник EBM: там, где EBM считает grad_x E через дифференцирование скаляра, диффузия предсказывает этот градиент напрямую сетью, обходя нестабильность обучения энергии. Многое из инженерии EBM перетекло именно туда.

^* SGLD (stochastic gradient Langevin dynamics) — метод сэмплирования, добавляющий к градиентному спуску гауссов шум фиксированной дисперсии, связанной с размером шага. Благодаря шуму итерации не сходятся в одну точку минимума, а блуждают по области, давая выборку из целевого распределения, а не единственный оптимум.