FFJORD: непрерывные потоки без ограничений на архитектуру

Проблема, которую решает FFJORD

Нормализующие потоки (normalizing flows) — семейство генеративных моделей, где данные получают через цепочку обратимых преобразований из простого распределения, например из стандартного гаусса. Чтобы обучать такую модель напрямую по правдоподобию, нужно на каждом шаге считать определитель якобиана преобразования. И вот здесь дискретные потоки платят цену: их архитектуру приходится подгонять так, чтобы этот определитель считался дёшево.

Отсюда все ограничения RealNVP, Glow и их родни — coupling-слои, треугольные якобианы, маскирование каналов. Модель обратима и правдоподобие считается быстро, но за это вы отдаёте свободу проектирования сети. FFJORD (Free-form Jacobian of Reversible Dynamics), представленный в 2018 году, ломает этот компромисс: он позволяет использовать произвольную нейросеть в качестве блока преобразования.

От дискретных потоков к непрерывным

Идея выросла из Neural ODE. Вместо конечной цепочки слоёв преобразование задаётся дифференциальным уравнением: состояние z(t) меняется во времени по правилу dz/dt = f(z(t), t), где f — нейросеть. Начальное состояние z(t0) — это точка данных, конечное z(t1) — точка в базовом распределении. Прямой проход и обратный — это просто интегрирование ODE в одну или другую сторону, поэтому преобразование обратимо по построению.

Для такого непрерывного потока (continuous normalizing flow, CNF) плотность меняется не через определитель, а через след якобиана. Логарифм плотности вдоль траектории описывается уравнением:

d log p(z(t)) / dt = -Tr( ∂f / ∂z(t) )

Это уже прогресс: след линеен, определитель — нет. Но лобовой расчёт следа для сети с D входами и D выходами требует D обратных проходов автодиффа — по одному на каждый диагональный элемент якобиана. Для картинок с тысячами измерений это неподъёмно.

Оценка следа по Хатчинсону

Ключевой трюк FFJORD — заменить точный след его несмещённой стохастической оценкой (Hutchinson trace estimator). Для случайного вектора шума ε с нулевым средним и единичной ковариацией выполняется:

Tr(A) = E[ εᵀ A ε ]

Произведение εᵀ (∂f/∂z) считается одним vector-Jacobian product, то есть одним обратным проходом автодифференцирования — независимо от размерности D. Вместо D проходов получается один. Именно это делает свободную архитектуру масштабируемой: динамику f можно брать какой угодно, лишь бы через неё проходил градиент.

Главный сдвиг FFJORD в том, что цена вычислений перестала зависеть от структуры сети. Вы больше не выбираете архитектуру ради дешёвого якобиана — вы выбираете её ради качества модели.

Что это даёт на практике

Сравним три поколения подходов по степени свободы и по тому, как считается изменение плотности.

Обучение и сэмплирование идут через численный ODE-солвер (adaptive solver вроде Dormand–Prince). Отсюда и главные издержки: число вызовов функции (number of function evaluations, NFE) заранее не фиксировано и зависит от жёсткости динамики. Чем сложнее выученное поле f, тем больше шагов делает солвер и тем дольше и прямой, и обратный проход.

Плюсы и минусы, если коротко

• Свобода архитектуры. Можно ставить обычные свёрточные блоки без масок и разбиений на каналы.
• Обратимость без потерь. Не нужно хранить активации всех слоёв — обратный проход восстанавливается интегрированием, память по глубине почти константна.
• Стохастическая оценка следа. Она несмещённая, но добавляет дисперсию в градиент — сходимость может быть шумнее.
• Скорость. Adaptive solver непредсказуем по времени; на инференсе один сэмпл может стоить десятков вызовов сети.

Как это устроено в коде

Концептуально обучение сводится к тому, чтобы интегрировать расширенную систему: саму динамику z и накопление логарифма плотности. Ниже — схематичный набросок на PyTorch-подобном псевдокоде, показывающий идею объединённого состояния.

def odefunc(t, state):
    z, logp = state
    eps = torch.randn_like(z)          # шум для оценки следа
    with torch.enable_grad():
        z.requires_grad_(True)
        dz = f(t, z)                   # свободная сеть-динамика
        # vector-Jacobian product: одним backward
        vjp = torch.autograd.grad(dz, z, eps, create_graph=True)[0]
        trace_est = (vjp * eps).sum(dim=1)
    dlogp = -trace_est
    return dz, dlogp

# интегрируем от t0 к t1 адаптивным солвером
z1, delta_logp = odeint(odefunc, (x, zeros), t=[t0, t1])
log_px = base_logprob(z1) + delta_logp

На практике используют библиотеку интегрирования с поддержкой обратного распространения через ODE (метод сопряжённого состояния, adjoint), чтобы не хранить всю траекторию солвера в памяти. Точные детали и рабочую реализацию лучше смотреть в оригинальном репозитории авторов.

FFJORD не стал самой быстрой генеративной моделью — по времени инференса он проигрывает и дискретным потокам, и более поздним диффузионным моделям. Но как исследовательский шаг он важен: показал, что непрерывная формулировка снимает архитектурные оковы потоков, а стоимость расчёта плотности можно сделать почти независимой от размерности данных. Идеи оценки следа и adjoint-обучения потом переиспользовали далеко за пределами нормализующих потоков.