Непрерывные нормализующие потоки

Замена переменных, доведённая до предела

Из блока 3, урок 110: если $y = T(x)$ — обратимое преобразование, то плотности связаны якобианом

$\log p_Y(y) = \log p_X(x) - \log\big|\det \nabla T(x)\big|$

Нормализующий поток строит $T$ как композицию слоёв и следит, чтобы определитель считался дёшево (треугольные матрицы, перестановки, аффинные связки). Ограничение архитектуры — цена за вычислимость.

Непрерывный поток снимает это ограничение. Пусть преобразование задано не композицией слоёв, а решением ОДУ:

\frac{dx}{dt} = \htmlData{k=field}{v_\theta(x, t)}, \qquad \frac{d \log p(x_t)}{dt} = -\htmlData{k=div}{\nabla\!\cdot\! v_\theta(x_t, t)}

Второе уравнение — формула мгновенной замены переменных. Смысл замены огромен: определитель превратился в след. Определитель $d\times d$ матрицы стоит $O(d^3)$ , след — $O(d)$ , если его удаётся посчитать; и никаких ограничений на архитектуру $v_\theta$ больше нет.

Что это даёт и чего стоит

Плюс, которого нет ни у диффузии, ни у flow matching: точное правдоподобие. Проинтегрировав второе уравнение вдоль траектории, получаем

$\log p(x_1) = \log p_0(x_0) - \int_0^1 \nabla\!\cdot\! v_\theta(x_t, t)\,dt$

Не нижнюю границу, как ELBO, а само значение — с точностью решателя. Для задач, где правдоподобие и есть ответ (детекция аномалий, сжатие, сравнение моделей), это существенно.

Минус — там же, где плюс. — это след якобиана, и точное его вычисление требует $d$ обратных проходов, ровно как в уроке 040. Спасает оценка Хатчинсона:

$\nabla\!\cdot\! v = \mathbb{E}_{z}\big[z^\top (\nabla v)\, z\big], \qquad z \sim \mathcal{N}(0, I)$

Одно произведение якобиана на вектор — один обратный проход. Оценка несмещённая, но шумная, и дисперсия растёт с размерностью.

Как обучали CNF и почему это было дорого

Обучение по правдоподобию требует градиента по $\theta$ через решатель ОДУ. Два способа:

подход	память	точность градиента
дифференцировать через шаги решателя	$O(\text{шагов})$	точная
adjoint-метод (решать сопряжённое ОДУ назад)	$O(1)$	приближённая

Adjoint был главным вкладом Neural ODE: память перестала зависеть от числа шагов. Но остались две проблемы, и они оказались решающими:

число шагов не под контролем. Адаптивный решатель сам решает, сколько вычислений ему нужно, и по мере обучения динамика усложняется — шагов становится больше, обучение замедляется само по себе;
симуляция в цикле обучения. Каждый шаг оптимизации требует полного прогона ОДУ вперёд и назад. Сравните с диффузией, где обучающий пример стоит один прогон сети.

Чем это оказалось для flow matching

Вот связь, ради которой урок стоит в этом месте. Flow matching обучает то же самое поле скоростей, что и CNF, но без симуляции:

	CNF по правдоподобию	flow matching
что обучается	$v_\theta$	$v_\theta$
цель	$\log p$ через интеграл дивергенции	MSE по скорости
нужен ли решатель при обучении	да	нет
нужна ли дивергенция	да	нет
стоимость примера	десятки прогонов	один
точное правдоподобие после обучения	да	да

Последняя строка — самое приятное. Обученная модель flow matching это и есть CNF: у неё есть поле скоростей, значит формула правдоподобия применима. Отказ от симуляции произошёл только на обучении.

Такая пара — «объект тот же, цель другая» — уже встречалась: в уроке 060 диффузия обучала score, не считая ни следа, ни нормировочной константы. Приём один и тот же: найти цель, оптимум которой совпадает с нужным, но которая считается локально.

шагов решателя 4

длина пути: 2.702
прямое расстояние: 2.702
кривизна: ×1
ошибка ломаной: 0.0e+0

Прямой путь: решателю достаточно любого числа шагов, потому что ломаная совпадает с траекторией. Для CNF это означало бы дешёвую симуляцию — и именно поэтому спрямление путей интересно не только сэмплированию.

Ошибка ломаной равна нулю: при прямых траекториях число шагов решателя перестаёт быть параметром вовсе.

Итог

CNF заменяет определитель якобиана его следом, снимая ограничения на архитектуру.
Он даёт точное правдоподобие, а не нижнюю границу.
Цена — оценка дивергенции и полная симуляция ОДУ на каждом шаге обучения; число шагов при этом не контролируется.
Flow matching обучает то же поле локальной целью, а правдоподобие остаётся доступным после обучения. Симуляция ушла из обучения, а не из модели.

Источники

Chen и др. — Neural Ordinary Differential Equations — CNF, adjoint-метод, формула изменения плотности
Grathwohl и др. — FFJORD — Оценка Хатчинсона для следа и свободная форма динамики

Что даёт непрерывный поток

Отметьте все верные утверждения о непрерывных нормализующих потоках.
Правдоподобие двумя путями
Возьмём простейшее поле скоростей — диагональное линейное:

$\frac{dx}{dt} = a_{11}x, \qquad \frac{dy}{dt} = a_{22}y$

Его решение известно в замкнутой форме, $x(T) = e^{a_{11}T}x(0)$ , и то же для $y$ . Начальная плотность — стандартная двумерная нормальная.

Реализуйте cnf_likelihood(a11, a22, t_end, x0, y0) — верните [x_end, y_end, log_p_via_divergence, log_p_via_determinant]:
- x_end, y_end — конечная точка;
- log_p_via_divergence = $\log p_0(x_0,y_0) - \displaystyle\int_0^{T}\nabla\!\cdot\! v\,dt$ , где дивергенция поля равна $a_{11} + a_{22}$ и от времени не зависит, поэтому интеграл берётся устно;
- log_p_via_determinant = $\log p_0(x_0,y_0) - \log\big|\det \nabla T\big|$ , где якобиан преобразования диагонален с элементами $e^{a_{11}T}$ и $e^{a_{22}T}$ .
Здесь $\log p_0(x,y) = -\frac{x^2+y^2}{2} - \log(2\pi)$ .

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

Загрузка редактора…

Ctrl/⌘ + Enter