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

GAN порождает сэмплы, но не может сказать, насколько вероятно данное изображение. VAE может, но только через нижнюю оценку. Нормализующие потоки (normalizing flows) идут третьим путём (Rezende and Mohamed, 2015; Papamakarios et al., 2021; Kobyzev, Prince and Brubaker, 2021): берут простую плотность, пропускают её через обратимое отображение и точно вычисляют плотность результата по формуле замены переменных. В этой заметке мы выводим эту формулу, собираем из неё поток, сравниваем потоки с другими семействами моделей и заканчиваем приёмом, который делает потоки вычислимыми: треугольной матрицей Якоби. Как строить слои с этим свойством, разбирается в следующей заметке, Архитектуры потоков.

Обозначения. Отображение $f$ всегда идёт из латентного пространства в пространство данных, $f: \mathcal Z \to \mathcal X$, так что $x = f(z)$ порождает данные, а $z = f^{-1}(x)$ их кодирует. $p_Z$ — базовая (латентная) плотность, $p_\theta(x)$ — плотность данных в модели. Для отображения $g$ через $J_g(u) = \partial g(u)/\partial u$ обозначена его матрица Якоби в точке $u$.

Модель: обратимые отображения гауссианы#

Возьмём непрерывную случайную величину $z$ с простой плотностью — сферической гауссианой (spherical Gaussian):

$$z \sim p_Z(z) = \mathcal N(z;\, 0,\, I).$$

«Сферическая» означает, что ковариационная матрица пропорциональна единичной (здесь в точности равна $I$): разброс одинаков во всех направлениях, а линии уровня — сферы. Данные получаются цепочкой из $K$ отображений,

$$x = f_\theta(z) = f_K \circ \dots \circ f_2 \circ f_1 (z),$$

и каждое $f_i$ обратимо (биективно): каждая точка выходного пространства — образ ровно одной входной точки. Для формулы плотности ниже каждое $f_i$ должно быть ещё и дифференцируемым с дифференцируемым обратным (диффеоморфизмом). Тогда обратима и вся цепочка, причём $f^{-1} = f_1^{-1} \circ f_2^{-1} \circ \dots \circ f_K^{-1}$. Для диффеоморфизма нужно $\dim z = \dim x$. Цепное правило, применённое к $f^{-1}\circ f = \mathrm{id}$ и $f\circ f^{-1} = \mathrm{id}$, даёт $J_{f^{-1}}J_f = I$ и $J_fJ_{f^{-1}} = I$, а ранг произведения матриц размеров $m\times n$ и $n\times m$ не больше $\min(m, n)$, поэтому оба тождества могут выполняться, только когда $m = n$. Значит, в отличие от VAE, латентный вектор имеет ту же размерность, что и данные.

Рисунок

У каждой точки Z ровно одна парная точка в X

Почему нельзя просто взять базовую плотность?#

Нам нужна $p_\theta(x)$. Первая догадка — отобразить $x$ обратно и взять значение базовой плотности, $p_\theta(x) \overset{?}{=} p_Z(f^{-1}(x))$. Это неверно, потому что ответ зависит от того, как ведёт себя $f$: насколько оно растягивает или сжимает пространство вокруг каждой точки.

Контрпример в одну строку: в одном измерении пусть $z \sim \mathcal N(0, 1)$ и $x = f(z) = 2z$. Наивная догадка даёт $p_Z(x/2)$; интеграл этой функции по всей прямой равен $2$, так что это даже не плотность. Истинное распределение — $x \sim \mathcal N(0, 4)$ с плотностью $p_Z(x/2)/2$. В точке $x = 1$ наивное значение равно $0.352$, а правильное — $0.176$: $f$ распределяет ту же вероятность по вдвое большей длине, поэтому плотность должна уменьшиться вдвое.

Замена переменных#

Пусть отображение $f: \mathcal Z \to \mathcal X$ обратимо, дифференцируемо и имеет дифференцируемое обратное, а $p_Z(z)$ — плотность, заданная на $z \in \mathcal Z$.

Одно измерение. Масса вероятности сохраняется. Маленький интервал ширины $|dz|$ вокруг $z$ содержит массу $p_Z(z)\,|dz|$, и $f$ переводит его в интервал ширины $|dx|$ вокруг $x = f(z)$, который должен содержать ту же массу:

$$p_\theta(x)\,|dx| = p_Z(z)\,|dz|,$$

откуда

$$p_\theta(x) = p_Z(z)\left|\frac{dz}{dx}\right| = \frac{p_Z(f^{-1}(x))}{\left|f'(f^{-1}(x))\right|}.$$

Любая размерность. Ширины превращаются в объёмы, и нам нужно знать, во сколько раз $f$ меняет малый объём. Вблизи точки $z$ отображение приблизительно линейно, $f(z + \delta) \approx f(z) + J_f(z)\,\delta$, с матрицей Якоби

$$J_f(z) = \begin{bmatrix} \dfrac{\partial f^{(1)}}{\partial z_1} & \cdots & \dfrac{\partial f^{(1)}}{\partial z_d} \\ \vdots & \ddots & \vdots \\ \dfrac{\partial f^{(d)}}{\partial z_1} & \cdots & \dfrac{\partial f^{(d)}}{\partial z_d} \end{bmatrix},$$

где $f^{(1)},\dots,f^{(d)}$ — выходные координаты $f$: строка $k$ матрицы $J_f$ содержит частные производные выхода $k$, столбец $j$ — производные по входу $z_j$. Поэтому крошечный куб со стороной $\varepsilon$ в точке $z$, натянутый на векторы $\varepsilon e_1,\dots,\varepsilon e_d$, переходит в крошечный параллелепипед, натянутый на $\varepsilon J_f(z) e_1, \dots, \varepsilon J_f(z) e_d$, то есть на столбцы $J_f(z)$, умноженные на $\varepsilon$. Объём параллелепипеда, натянутого на столбцы матрицы $A$, равен $\lvert\det A\rvert$, так что $f$ умножает малые объёмы вблизи $z$ на $\lvert\det J_f(z)\rvert$. В двух измерениях маленький квадрат становится маленьким параллелограммом, площадь которого в $\lvert\det J_f(z)\rvert$ раз больше площади квадрата.

Рисунок

Матрица Якоби — это локальная геометрия

Масса сохраняется так же, как в одном измерении: куб объёма $\varepsilon^d$ содержит массу $p_Z(z)\,\varepsilon^d$, и его образ объёма $\lvert\det J_f(z)\rvert\,\varepsilon^d$ должен содержать ту же массу. Следовательно, при $z = f^{-1}(x)$

$$\begin{aligned} p_\theta(x) &= p_Z\big(f^{-1}(x)\big)\,\left|\det \frac{\partial f^{-1}(x)}{\partial x}\right| \\ &= p_Z(z)\,\left|\det \frac{\partial z}{\partial x}\right| \\ &= p_Z(z)\,\left|\det \frac{\partial f(z)}{\partial z}\right|^{-1}. \end{aligned}$$

Две формы говорят одно и то же. Матрица Якоби обратного отображения вычисляется в точке данных, а матрица Якоби $f$ — в латентной точке $z = f^{-1}(x)$. Их связывают теорема об обратной функции, $J_{f^{-1}}(x) = \big[J_f(z)\big]^{-1}$, и равенство $\det(A^{-1}) = 1/\det A$. Множитель $|\det J_f|$ измеряет расширение или сжатие пространства вокруг точки (Papamakarios et al., 2021, §2.1). Там, где $f$ расширяет пространство, та же вероятность распределяется по большему объёму и плотность падает; там, где $f$ сжимает пространство, плотность растёт:

• $|\det J_f| = 2$: объёмы удваиваются, а плотность уменьшается вдвое;
• $|\det J_f| = 0.5$: объёмы уменьшаются вдвое, а плотность удваивается;
• $|\det J_f| = 1$: объёмы сохраняются (форма при этом может меняться), а с ними и плотность;
• $\det J_f = 0$: объём схлопывается в ноль. Линейное отображение $J_f(z)$ необратимо, и тогда у $f$ нет дифференцируемого обратного в точке $f(z)$, поскольку дифференцируемое обратное давало бы $J_{f^{-1}}J_f = I$. Множитель плотности $1/|\det J_f|$ здесь бесконечен, и формула выше неприменима. Поэтому слои потоков строят так, чтобы всюду было $\det J_f \neq 0$.

Рисунок

Что |det J| делает с маленьким квадратом

Разобранный одномерный пример#

Рисунок

Замена переменных в одном измерении

Возьмём $p_Z = \mathcal N(0, 1)$ и отображение с рисунка,

$$\begin{gathered} f(z) = z + 0.8\tanh(2z), \\ f'(z) = 1 + 1.6\,\operatorname{sech}^2(2z) > 0 . \end{gathered}$$

Производная всюду положительна, поэтому $f$ строго возрастает и, значит, обратимо (обратное не выражается в замкнутой форме, но легко находится численно). Тогда $p_\theta(x) = p_Z(z)/f'(z)$ при $z = f^{-1}(x)$:

$z$	$x = f(z)$	$f'(z)$	$p_Z(z)$	$p_\theta(x) = p_Z(z)/f'(z)$
$0$	$0$	$2.600$	$0.3989$	$0.1534$
$0.5$	$1.109$	$1.672$	$0.3521$	$0.2106$
$1$	$1.771$	$1.113$	$0.2420$	$0.2174$
$2$	$2.799$	$1.002$	$0.0540$	$0.0539$

Вблизи начала координат $f$ растягивает пространство в $2.6$ раза, поэтому плотность там опускается намного ниже пика гауссианы; в хвостах $f' \approx 1$ и плотность почти не меняется. Один горб превращается в два с максимумами в $x \approx \pm 1.50$ (то есть $z \approx \pm 0.77$), $p_\theta \approx 0.234$. Даже одно гладкое обратимое отображение может превратить унимодальное базовое распределение в бимодальную модель. В логарифмической форме при $z = 0.5$: $\log p_\theta(x) = \log p_Z(z) - \log f'(z) = -1.044 - 0.514 = -1.558$.

Полоса $z \in [0.1, 0.5]$ содержит вероятность $\Phi(0.5) - \Phi(0.1) = 0.152$. Её образ — $x \in [0.258, 1.109]$, в $2.13$ раза шире, поэтому средняя плотность на нём падает с $0.379$ до $0.178$, в то же число раз: та же масса вероятности, другая ширина.

Пример: растяжение пространства в 2 раза#

Простейший многомерный случай — линейное отображение $x = Jz$, матрица Якоби которого $\partial x/\partial z = J$ всюду одна и та же. При

$$J = \begin{bmatrix} 2 & 0 \\ 0 & 2 \end{bmatrix}$$

каждое направление растягивается в 2 раза, и $\det J = 2\cdot 2 - 0\cdot 0 = 4$. Единичный квадрат в $\mathcal Z$ становится квадратом $2\times 2$ в $\mathcal X$, так что

$$p_\theta(x) = p_Z(z)\left|\frac{1}{\det\left(\frac{\partial x}{\partial z}\right)}\right| = \frac{p_Z(z)}{4}.$$

Проверим на $z \sim \mathcal N(0, I)$ в двух измерениях: $x = 2z \sim \mathcal N(0, 4I)$, и действительно $p_\theta(0) = \frac{1}{2\pi\cdot 4} = 0.0398 = \frac14\cdot\frac{1}{2\pi}$.

Ориентация не важна. Для

$$J = \begin{bmatrix} 2 & 0 \\ 0 & -2 \end{bmatrix}, \qquad \det J = -4,$$

квадрат ещё и отражается, но его площадь по-прежнему равна $4$. Знак определителя фиксирует ориентацию; изменение объёма, а значит и плотности, зависит только от $|\det J| = 4$. Поэтому в формуле стоит модуль.

Рисунок

Матрица Якоби, растягивающая каждое направление в 2 раза

Композиция отображений: нормализующий поток#

Теперь применим формулу ко всей цепочке $z \sim p_Z(z)$, $x = f_\theta(z) = f_K\circ\dots\circ f_2\circ f_1(z)$. Обозначим промежуточные точки

$$z_0 = z, \quad z_i = f_i(z_{i-1}) \;\; (i = 1,\dots,K),$$

так что $z_K = x$. Обратное отображение проходит цепочку в обратном порядке, $f^{-1} = f_1^{-1}\circ\dots\circ f_K^{-1}$, и по цепному правилу его матрица Якоби — произведение матриц Якоби слоёв, каждая из которых берётся в своей промежуточной точке: $J_{f^{-1}}(x) = J_{f_1^{-1}}(z_1)\,J_{f_2^{-1}}(z_2)\cdots J_{f_K^{-1}}(z_K)$. Поскольку $\det(AB) = \det A\,\det B$,

$$\begin{aligned} p_\theta(x) &= p_Z(z_0)\,\left|\det\frac{\partial f^{-1}(x)}{\partial x}\right| \\ &= p_Z(z_0)\prod_{i=1}^{K}\left|\det\frac{\partial f_i^{-1}(z_i)}{\partial z_i}\right| \\ &= p_Z(z_0)\prod_{i=1}^{K}\left|\det\frac{\partial f_i(z_{i-1})}{\partial z_{i-1}}\right|^{-1}. \end{aligned}$$

Чтобы вычислить её, мы проводим $x$ через слои в обратном направлении, $z_K = x \to z_{K-1} \to \dots \to z_0$, и собираем по одному определителю на слой. Эта конструкция и есть нормализующий поток: плотность перетекает через последовательность обратимых отображений, а цепочка, прочитанная в обратную сторону, нормализует данные, превращая их (приближённо, когда модель обучена) в выборку из стандартного нормального базового распределения. Название популяризовали в машинном обучении Rezende and Mohamed (2015), опираясь на более ранние работы по оцениванию плотности через преобразование данных к гауссиане (Tabak and Vanden-Eijnden, 2010). Потоки Rezende и Mohamed состояли из планарных и радиальных отображений, определители которых дёшевы по другой причине, чем у треугольных слоёв ниже (Архитектуры потоков, часть Ж).

Для сэмплирования нужно только прямое направление: берём $z \sim \mathcal N(0, I)$ и вычисляем $x = f_\theta(z)$; для вычисления плотности нужно обратное. Какое из двух направлений вычисляется дёшево, зависит от архитектуры: маскированные авторегрессионные потоки (masked autoregressive flows, MAF; Papamakarios, Pavlakou and Murray, 2017) быстро вычисляют плотность и медленно сэмплируют, обратные авторегрессионные потоки (inverse autoregressive flows, IAF; Kingma et al., 2016) — наоборот, а слои связи (coupling layers) делают дешёвым и то и другое (см. Архитектуры потоков). Сэмплирование превращает простое $z$ в сэмпл из сложного распределения; формула плотности — то, что позволяет обучить $f_\theta$ так, чтобы это распределение совпало с данными.

Метод максимального правдоподобия#

Возьмём логарифм, и произведение станет суммой, где $J_{f_i^{-1}}(z_i) = \partial f_i^{-1}(z_i)/\partial z_i$:

$$\begin{aligned} \log p_\theta(x) &= \log p_Z(z_0) \\ &\quad + \sum_{i=1}^{K}\log\big|\det J_{f_i^{-1}}(z_i)\big| . \end{aligned}$$

Это точное вычисление логарифма правдоподобия, без оценок снизу и без сэмплирования. Обучение — обычный метод максимального правдоподобия, который минимизирует прямую KL-дивергенцию (forward KL) от распределения данных к модели (почему мы минимизируем именно кросс-энтропию). Как и в заметке о подгонке распределений, плотность данных неизвестна, поэтому мы минимизируем отрицательный логарифм правдоподобия по набору данных $x^{(1)},\dots,x^{(N)}$, теперь уже с нейросетевой моделью:

$$\begin{aligned} \mathcal L(\theta) &= -\frac1N\sum_{n=1}^{N}\Bigg[\log p_Z\big(z_0^{(n)}\big) \\ &\qquad + \sum_{i=1}^{K}\log\left|\det\frac{\partial f_i^{-1}(z_i^{(n)})}{\partial z_i^{(n)}}\right|\Bigg], \end{aligned}$$

где $z_i^{(n)}$ — промежуточные точки для $x^{(n)}$: $z_K^{(n)} = x^{(n)}$ и $z_{i-1}^{(n)} = f_i^{-1}(z_i^{(n)})$, так что $z_0^{(n)} = f_\theta^{-1}(x^{(n)})$.

Два слагаемых тянут в разные стороны. Первое хочет, чтобы каждая закодированная точка лежала вблизи центра гауссианы. Второе штрафует обратное отображение за сжатие пространства, поэтому модель не может улучшить целевую функцию, схлопнув все данные в $z = 0$.

Изображения: деквантование, масштабирование и биты на размерность#

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

Формула даёт плотность, поэтому она точна для непрерывных данных. Изображения дискретны, например 8-битные значения пикселей $x \in \{0,\dots,255\}^d$, и непрерывная плотность могла бы принимать в этих точках сколь угодно большие значения: поток, обученный на исходных значениях, может неограниченно повышать своё правдоподобие, стягивая массу вероятности на эти точки. Поэтому потоки обучают на деквантованных (dequantized) данных $\tilde x = x + u$ с равномерным шумом $u \sim \mathcal U[0, 1)^d$. Вероятность, которую модель даёт дискретному значению $x$, равна $P(x) = \int_{[0,1)^d} p_\theta(x + u)\,du$, и неравенство Йенсена даёт $\mathbb E_u\big[\log p_\theta(x + u)\big] \le \log P(x)$: непрерывный логарифм правдоподобия деквантованных данных — нижняя оценка дискретного логарифма правдоподобия (Theis, van den Oord and Bethge, 2016).

Биты на размерность (bits per dimension). Правдоподобие моделей изображений сообщают как $-\log_2 P(x)/d$ — число бит на координату, которое нужно коду без потерь, построенному по модели (см. биты на размерность в заметке об авторегрессионных моделях). Модель, распределяющая вероятность равномерно по 256 значениям, получает ровно 8. Для потока сообщают число $-\mathbb E_u\big[\log_2 p_\theta(x + u)\big]/d$, которое по неравенству выше является верхней оценкой дискретного значения.

Масштабирование меняет плотность. Сетям редко подают на вход значения пикселей из $[0, 256)$. Распространённый выбор — $x' = \tilde x/127.5 - 1$, что переводит данные примерно в $[-1, 1]$. Тогда поток моделирует плотность $x'$, а замена переменных связывает её с плотностью $\tilde x$: у отображения $\tilde x\mapsto x'$ постоянная матрица Якоби $\tfrac{1}{127.5}I$, поэтому

$$\begin{gathered} p_{\tilde x}(\tilde x) = p_{x'}(x')\cdot 127.5^{-d}, \\[4pt] -\frac{\log_2 p_{\tilde x}(\tilde x)}{d} = -\frac{\log_2 p_{x'}(x')}{d} + \log_2 127.5 . \end{gathered}$$

Число бит на размерность в единицах пикселей равно числу в единицах модели плюс $\log_2 127.5 \approx 6.99$. В общем случае, если после предобработки один уровень квантования имеет ширину $a$, сдвиг равен $-\log_2 a$ бит на размерность (Kingma and Dhariwal, 2018, §2, формула 2): $\log_2 255\approx 7.99$ для $x/255$, как у распространённых загрузчиков изображений, и ровно 8 для $x/256$. Например (игрушечные числа), поток, который сообщает $-3.50$ бита на размерность для $x'$, имеет в единицах пикселей $-3.50 + 6.99 = 3.49$ бита на размерность. Отрицательное значение в единицах модели нормально, потому что плотность на узком диапазоне может превышать 1; с опубликованными результатами, например с таблицей правдоподобий в заметке об авторегрессионных моделях, можно сравнивать только число в единицах пикселей.

Шум должен быть в тех же единицах. Шум деквантования должен покрывать ровно один уровень того, что видит сеть. При $x' = x/127.5 - 1$ это значит прибавлять $u/127.5$ с $u\sim\mathcal U[0,1)^d$, что равносильно прибавлению $u$ до масштабирования. Шум ширины 1 в единицах $x'$ размазал бы каждое значение примерно на 127 уровней, а без шума правдоподобие снова становится неограниченным. NICE, например, деквантует CIFAR-10 равномерным шумом ширины $\tfrac{1}{128}$ на данных, приведённых к $[-1, 1]$, где ширина одного уровня равна $\tfrac{2}{256} = \tfrac{1}{128}$ (Dinh, Krueger and Bengio, 2015, §5.1).

Логит-преобразование. Значения пикселей ограничены, а гауссовская базовая плотность задана на всём $\mathbb R^d$. Чтобы ослабить краевые эффекты, RealNVP моделирует логит масштабированных пикселей (Dinh, Sohl-Dickstein and Bengio, 2017, §4.1):

$$\begin{gathered} y = \operatorname{logit}(p) = \log\frac{p}{1 - p}, \\[4pt] p = \alpha + (1 - 2\alpha)\,\frac{\tilde x}{256}, \qquad \alpha = 0.05 . \end{gathered}$$

Сдвиг $\alpha$ держит $p$ внутри $[0.05, 0.95)$, так что $y$ остаётся в пределах примерно $\pm2.94$, а не уходит в $\pm\infty$ на чёрных и белых пикселях. Это форма из эталонной реализации; в статье напечатано $1 - \alpha$ вместо $1 - 2\alpha$. Производная преобразования по каждой координате равна $dy/d\tilde x = (1 - 2\alpha)/\big(256\,p(1-p)\big)$, поэтому

$$\begin{aligned} -\log_2 p_{\tilde x}(\tilde x) &= -\log_2 p_y(y) \\ &\quad + \sum_{i=1}^{d}\log_2\frac{256\,p_i(1 - p_i)}{1 - 2\alpha} . \end{aligned}$$

В отличие от линейного масштабирования, поправка зависит от пикселя: около $6.15$ бита для среднего серого значения ($\tilde x = 128.5$) и $3.80$ бита для самых тёмных и самых светлых значений ($\tilde x = 0.5$ или $255.5$). RealNVP включает это слагаемое в свои биты на размерность. В статье о Glow сказано, что она следует предобработке RealNVP (§5); её эталонный код вместо этого переводит данные в $[-\tfrac12, \tfrac12)$, добавляет шум ширины $\tfrac{1}{256}$ и прибавляет к логарифму правдоподобия константу $-d\log 256$, то есть реализует линейный случай выше с $a = \tfrac{1}{256}$. Потокам, построенным для изображений, посвящена часть Е заметки «Архитектуры потоков».

Сводка: GAN, VAE и потоки#

Рисунок

Сводка: три способа связать данные x и латентную переменную z

• GAN (Goodfellow et al., 2014; см. заметку о GAN). Логарифм правдоподобия не вычисляется: генератор задаёт своё распределение только через сэмплы. Вывода латентных переменных тоже нет: ничто не отображает $x$ обратно в $z$, и ничто не гарантирует, что генератор покрывает весь носитель (support) распределения данных, так что у настоящего $x$ может вовсе не быть латентного кода.
• VAE (Kingma and Welling, 2014; см. заметку о VAE). Он оптимизирует нижнюю оценку логарифма правдоподобия (evidence lower bound, ELBO) и выводит латентные переменные через приближённое апостериорное распределение $q_\phi(z\mid x)$.
• Генеративные модели на основе потоков, например NICE (Dinh, Krueger and Bengio, 2015) и RealNVP (Dinh, Sohl-Dickstein and Bengio, 2017). Они минимизируют точный отрицательный логарифм правдоподобия, а апостериорный вывод точен: $z = f^{-1}(x)$, детерминированный ответ.
• Диффузионные модели замыкают список; они появятся в курсе позже и, как VAE, выводят свою целевую функцию из нижней оценки правдоподобия.

	Логарифм правдоподобия	Вывод латентных переменных
GAN	нет	нет
VAE	нижняя оценка (ELBO)	приближённый, $q_\phi(z\mid x)$
Поток	точный	точный, $z = f^{-1}(x)$

Ключевой приём: треугольная матрица Якоби#

Всё сказанное выше держится на $|\det J|$: если его можно вычислить, плотность точна; если его нельзя вычислить дёшево, правдоподобие на практике не удаётся ни вычислять, ни использовать для обучения в масштабе изображений. Для произвольной сети это дорого. Вычисление определителя произвольной матрицы $d\times d$ стоит $O(d^3)$ методом Гаусса или через LU-разложение (Golub and Van Loan, 2013); наивное разложение по алгебраическим дополнениям (cofactor expansion) ещё хуже, $O(d!)$. Возьмём $d = 1000$ (у полутонового изображения $32\times 32$ уже $d = 1024$): это порядка $10^9$ операций (около $\tfrac23 d^3 \approx 6.7\times 10^8$ для LU) на каждый пример, каждый слой и каждый шаг обучения. Для произвольной сети даже построение $J$ требует $d$ обратных проходов, по одному на строку.

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

$$\begin{gathered} A = \begin{bmatrix} a_{11} & 0 & \cdots & 0 \\ a_{21} & a_{22} & \ddots & \vdots \\ \vdots & \vdots & \ddots & 0 \\ a_{n1} & a_{n2} & \cdots & a_{nn} \end{bmatrix} \\ \Longrightarrow\quad \det A = a_{11}\, a_{22} \cdots a_{nn}. \end{gathered}$$

Итак, если $J$ треугольная,

$$\begin{gathered} \det J = \prod_{i=1}^{d} J_{ii}, \\ \log|\det J| = \sum_{i=1}^{d} \log|J_{ii}|, \end{gathered}$$

что стоит $O(d)$: около $10^3$ операций вместо примерно $10^9$, и нужны только диагональные элементы.

Треугольная $J$ к тому же обратима тогда и только тогда, когда все её диагональные элементы ненулевые, а это легко гарантировать по построению. Для самого отображения это условие даёт больше. Если $J_f$ всюду нижнетреугольная, выход $i$ зависит только от входов $1,\dots,i$. Если каждый диагональный элемент непрерывен и всюду ненулевой, он сохраняет знак, так что каждый выход строго монотонен по своему входу. Тогда отображение инъективно и обращается по одной координате за раз: $z_1$ из $x_1$, затем $z_2$ из $x_2$ и уже известного $z_1$, и так далее. Оно является биекцией $\mathbb R^d$, когда каждое из этих одномерных отображений сюръективно, то есть принимает все значения из $\mathbb R$, как аффинные отображения $z_i \mapsto e^{s}z_i + t$ в слоях связи.

Рисунок

Почему треугольная матрица Якоби — ключевой приём

Ноутбук#

Ноутбук о нормализующих потоках из репозитория HSE-LAMBDA Deep Generative Models (ссылка закреплена на конкретном коммите; ячейки нумеруются с 1, считая все ячейки подряд, включая текстовые) опирается на библиотеку Андрея Карпаты pytorch-normalizing-flows (Karpathy, 2019), включённую в репозиторий как nflib. В ноутбуке нет сохранённых выходов ячеек, поэтому он документирует процедуры, а не результаты.

• Ячейка 2 записывает замену переменных с порождающей стороны, $\log p(z_K) = \log p(z_0) - \sum_i \log\lvert\det \partial f_i/\partial z_{i-1}\rvert$.
• Ячейки 4–5 задают три двумерных игрушечных набора данных: две луны (two moons; используется именно он), смесь четырёх гауссиан со стандартным отклонением $0.5$ и набор точек SIGGRAPH, которому нужен файл siggraph.pkl, отсутствующий в репозитории.
• Ячейка 7 строит модель: логистическое априорное распределение — равномерное, пропущенное через обратную сигмоиду, с $\log p_Z(z) = \sum_i\big[-\log(1 + e^{z_i}) - \log(1 + e^{-z_i})\big]$ (гауссовское априорное распределение закомментировано). NICE использует то же априорное распределение и предпочитает его, потому что оно даёт градиенты с лучшим поведением (Dinh, Krueger and Bengio, 2015, §3.4). Формула замены переменных верна для любой базовой плотности, так что меняется только $\log p_Z$. Поток состоит из трёх блоков, в каждом из которых идут ActNorm, обратимая «свёртка» $1\times1$ и слой связи на нейросетевых сплайнах (neural spline coupling; 8 интервалов сплайна, граница сплайна (tail bound) 3, ширина скрытого слоя 16); все эти слои описаны в заметке Архитектуры потоков.
• Ячейки 9–11 обучают модель с Adam (шаг обучения $10^{-4}$, затухание весов (weight decay) $10^{-5}$) на функции потерь выше: zs, prior_logprob, log_det = model(x_batch), logprob = prior_logprob + log_det, loss = -torch.sum(logprob), в течение 1000 итераций, каждая на свежем батче из 128 сэмплов. Функция потерь суммируется по батчу, а не усредняется, как в $\mathcal L(\theta)$, что лишь масштабирует градиент. nflib называет направления наоборот по сравнению с этой заметкой: её forward идёт от данных к латентному пространству (наше $f^{-1}$) и накапливает наше $\sum_i \log\lvert\det \partial f_i^{-1}/\partial z_i\rvert$, а model.sample берёт сэмплы из априорного распределения и запускает inverse (наше $f$).
• Ячейки 13–17 отображают данные в латентное пространство, а сэмплы из априорного распределения — в данные и показывают, как сетка деформируется слой за слоем. Ячейка 13 берёт точки из априорного распределения вызовом model.prior.sample([128, 2]); у априорного распределения уже две координаты, поэтому получается тензор $128\times2\times2$, и на диаграмме рассеяния оказывается 256 точек вместо 128. Это безвредно (sample((128,)) даёт задуманную форму). Ячейка 17 выполняет ещё один шаг обучения и рисует деформацию сетки один раз: её цикл на 500 шагов прерывается при $k = 0$.
• Ячейка 19 раскрашивает сетку по prior_logprob.exp(), то есть по $p_Z(f^{-1}(x))$ без логарифма определителя. Это ровно наивная формула из начала заметки; плотность модели — (prior_logprob + log_det).exp().

Дальше#

Потокам нужны слои, которые достаточно выразительны и при этом по построению имеют дешёвый определитель Якоби. Такие слои строит заметка Архитектуры потоков: слои связи (NICE, RealNVP), маскированные авторегрессионные слои (MAF, IAF) и обратимые свёртки $1\times1$ из Glow, затем сплайновые слои связи, ActNorm и многомасштабные потоки для изображений и, наконец, планарные и радиальные потоки Rezende и Mohamed.

Литература#

• D. J. Rezende, S. Mohamed. Variational inference with normalizing flows. ICML, 2015.
• G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, B. Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. JMLR, 22(57), 2021.
• I. Kobyzev, S. J. D. Prince, M. A. Brubaker. Normalizing flows: An introduction and review of current methods. IEEE TPAMI, 43(11), 2021.
• E. G. Tabak, E. Vanden-Eijnden. Density estimation by dual ascent of the log-likelihood. Communications in Mathematical Sciences, 8(1), 2010.
• L. Dinh, D. Krueger, Y. Bengio. NICE: Non-linear independent components estimation. ICLR Workshop, 2015. §3.4 (априорное распределение), §5.1 (деквантование).
• L. Dinh, J. Sohl-Dickstein, S. Bengio. Density estimation using Real NVP. ICLR, 2017. §4.1 (логит-предобработка). Эталонный код: tensorflow/models, research/real_nvp.
• D. P. Kingma, P. Dhariwal. Glow: Generative flow with invertible 1×1 convolutions. NeurIPS, 2018. §2, формула 2 (константа дискретизации), §5. Эталонный код: openai/glow.
• I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, Y. Bengio. Generative adversarial nets. NeurIPS, 2014.
• D. P. Kingma, M. Welling. Auto-encoding variational Bayes. ICLR, 2014.
• G. Papamakarios, T. Pavlakou, I. Murray. Masked autoregressive flow for density estimation. NeurIPS, 2017.
• D. P. Kingma, T. Salimans, R. Jozefowicz, X. Chen, I. Sutskever, M. Welling. Improved variational inference with inverse autoregressive flow. NeurIPS, 2016.
• L. Theis, A. van den Oord, M. Bethge. A note on the evaluation of generative models. ICLR, 2016.
• G. H. Golub, C. F. Van Loan. Matrix Computations, 4th ed. Johns Hopkins University Press, 2013. §3.2: LU-разложение.
• A. Karpathy. pytorch-normalizing-flows. Репозиторий на GitHub, 2019.
• HSE-LAMBDA, репозиторий курса Deep Generative Models: nf.ipynb.