Архитектуры потоков

Заметка о замене переменных закончилась проблемой вычислительной стоимости. Каждый слой потока вносит в логарифм правдоподобия слагаемое $\pm\log\lvert\det J\rvert$ (со знаком плюс, когда $J$ — матрица Якоби отображения из данных в латентное пространство). Для произвольной матрицы Якоби размера $d\times d$ вычисление этого определителя стоит $O(d^3)$. Для треугольной матрицы Якоби это просто произведение диагональных элементов, которое стоит $O(d)$. Поэтому вопрос в том, как придать $J$ структуру, определитель которой дёшев по построению: треугольную, линейное отображение, хранимое в виде произведения множителей, или единичную матрицу плюс слагаемое малого ранга.

Ответ дают три конструкции: слои связи (coupling layers, часть А) и авторегрессионные слои (часть Б), у которых матрица Якоби треугольная, и обратимые свёртки $1\times1$ (invertible 1×1 convolutions, часть В), у которых она блочно-диагональная с плотными блоками, хранимыми как произведения треугольных множителей. Каждая из них сохраняет обратимость слоя и дешевизну его определителя, а основную работу внутри слоя по-прежнему выполняют произвольные нейронные сети. На них опираются ещё три части. Часть Г заменяет аффинное отображение внутри слоя связи монотонным сплайном, часть Д добавляет ActNorm — слой нормализации из Glow, а часть Е собирает всё это в потоки для изображений. Часть Ж показывает более ранний путь к дешёвому определителю, на котором треугольные матрицы вообще не нужны: планарные и радиальные потоки. При первом чтении достаточно частей А, Б, В и Е; части Г и Ж — расширения.

Обозначения. Один слой потока отображает $z\in\mathbb R^d$ в $z'\in\mathbb R^d$. Его матрица Якоби — $J = \partial z'/\partial z$ с элементами $J_{ij} = \partial z'_i/\partial z_j$: строка $i$ соответствует выходной координате, а столбец $j$ — входной. $\odot$ — поэлементное произведение, а $\exp$ от вектора применяется поэлементно. Здесь неважно, в каком направлении используется слой: из данных в латентное пространство или из латентного пространства в данные. Определитель обратного слоя равен обратной величине определителя прямого, поэтому определитель, дешёвый в одном направлении, дёшев и в другом. Нижний индекс нумерует координаты одного вектора ($z_i$, $z_{<i}$), а не слои потока, как в предыдущей заметке. Тензор-изображение имеет высоту $h$, ширину $w$ и $C$ каналов.

А. Слои связи (NICE, RealNVP)#

Разобьём вход на две части,

$$z = (z_a,\; z_b), \qquad z_a\in\mathbb R^{k},\quad z_b\in\mathbb R^{d-k},$$

и зададим преобразование

$$\begin{aligned} z'_a &= z_a, \\ z'_b &= z_b \odot \exp\bigl(s(z_a)\bigr) + t(z_a), \end{aligned}$$

где $s, t:\mathbb R^{k}\to\mathbb R^{d-k}$ — произвольные нейронные сети, которые называют сетью масштаба (scale) и сетью сдвига (translation). Конструкция работает благодаря двум фактам:

• $z'_a$ вообще не зависит от $z_b$;
• $z'_b$ зависит от $z_b$ лишь аффинно и к тому же поэлементно: в $i$-ю координату $z'_{b,i}$ входит только $z_{b,i}$, умноженная на $\exp(s_i(z_a))$ и сдвинутую на $t_i(z_a)$.

Рисунок

Аффинный слой связи и обратный к нему

Матрица Якоби. Эти два факта задают блочную структуру $J$:

$$J = \frac{\partial z'}{\partial z} = \begin{bmatrix} I_k & 0 \\[4pt] \dfrac{\partial z'_b}{\partial z_a} & \operatorname{diag}\bigl(\exp(s(z_a))\bigr) \end{bmatrix}.$$

• Левый верхний блок: $\partial z'_a/\partial z_a = I$, потому что $z'_a = z_a$.
• Правый верхний блок: $\partial z'_a/\partial z_b = 0$, потому что $z'_a$ не зависит от $z_b$.
• Правый нижний блок: $\partial z'_b/\partial z_b = \operatorname{diag}(\exp(s(z_a)))$, потому что зависимость от $z_b$ поэлементная с угловым коэффициентом $\exp(s_i(z_a))$.
• Левый нижний блок: $\partial z'_b/\partial z_a$ — плотная матрица, в которую входят производные $s$ и $t$. Она сложна, но это неважно.

Значит, $J$ нижнетреугольная, и её определитель равен произведению диагональных элементов:

$$\begin{aligned} \det J &= 1^k\cdot\prod_{i=1}^{d-k} \exp\bigl(s_i(z_a)\bigr) \\ &= \exp\Bigl(\sum_{i=1}^{d-k} s_i(z_a)\Bigr), \\[4pt] \log\lvert\det J\rvert &= \sum_{i=1}^{d-k} s_i(z_a). \end{aligned}$$

Здесь $i$ пробегает $d-k$ координат $z_b$, а $s_i$ — $i$-й выход сети $s$.

Определитель всегда положителен, так что модуль не нужен. Слой обратим при любом выборе сетей $s$ и $t$, и не потому, что $\det J\ne0$ (само по себе это гарантирует лишь локальную обратимость), а потому, что обратное отображение можно выписать явно, как показано ниже. Плотный блок никогда не входит в определитель: $s$ и $t$ могут быть сколь угодно глубокими, и $\det J$ от этого не становится дороже.

Рисунок

Почему матрица Якоби слоя связи нижнетреугольная

Например, при $d = 6$ и разбиении $3+3$ пусть в некоторой точке $s(z_a) = (0.2,\,-0.5,\,0.1)$. Тогда $\log\lvert\det J\rvert = 0.2 - 0.5 + 0.1 = -0.2$ и $\det J = e^{-0.2}\approx 0.819$. Вблизи этой точки слой уменьшает малые объёмы примерно на 18%, что бы ни содержал плотный блок.

Обратное отображение обращает только поэлементный шаг. Поскольку $z'_a = z_a$, сети можно снова вычислить на том же входе:

$$\begin{aligned} z_a &= z'_a, \\ z_b &= \bigl(z'_b - t(z'_a)\bigr)\odot\exp\bigl(-s(z'_a)\bigr), \\[4pt] \log\lvert\det J_{\text{inv}}\rvert &= -\sum_{i=1}^{d-k} s_i(z'_a). \end{aligned}$$

Ни $s$, ни $t$ обращать не нужно. Каждое направление требует одного прохода через обе сети.

Б. Авторегрессионные слои (MAF, IAF)#

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

$$z'_i = f_i(z_1, \dots, z_i), \qquad i = 1,\dots,d.$$

Теперь $J$ автоматически нижнетреугольная:

• зависимости от координат с номерами $> i$ нет, поэтому $J_{ij} = \partial z'_i/\partial z_j = 0$ для всех $j > i$;
• диагональные элементы равны $J_{ii} = \partial f_i/\partial z_i$.

$$J = \begin{bmatrix} \frac{\partial f_1}{\partial z_1} & 0 & \cdots & 0 \\ \ast & \frac{\partial f_2}{\partial z_2} & \cdots & 0 \\ \vdots & & \ddots & \vdots \\ \ast & \ast & \cdots & \frac{\partial f_d}{\partial z_d} \end{bmatrix},$$

так что

$$\det J = \prod_{i=1}^{d} \frac{\partial f_i}{\partial z_i}.$$

Снова определитель — произведение диагональных элементов. Слой связи — частный случай всего с двумя блоками: $z_a$ служит условием для $z_b$.

Конкретный слой. В MAF, маскированном авторегрессионном потоке (masked autoregressive flow; Papamakarios, Pavlakou and Murray, 2017) каждое $f_i$ аффинно по своей координате, а масштаб и сдвиг вычисляются по координатам, стоящим строго перед ней:

$$\begin{aligned} z'_i &= z_i\,\exp\bigl(s_i(z_{<i})\bigr) + t_i(z_{<i}), \\[4pt] J_{ii} &= \exp\bigl(s_i(z_{<i})\bigr), \\[4pt] \log\lvert\det J\rvert &= \sum_{i=1}^{d} s_i(z_{<i}). \end{aligned}$$

Здесь $z_{<i} = (z_1,\dots,z_{i-1})$. Поскольку $s_i$ и $t_i$ не видят $z_i$, диагональный элемент равен просто $\exp(s_i) > 0$, и слой можно обращать по одной координате за раз. Тот же слой часто записывают как $f(x_i) = (x_i - \mu(x_{<i}))/\exp(\alpha(x_{<i}))$; это та же форма, что и выше, с $s = -\alpha$ и $t = -\mu\,e^{-\alpha}$.

Маскирование. Можно было бы дать каждой координате собственную сеть, которая читает $z_{<i}$, ценой $d$ сетей и $d$ последовательных проходов. MAF вместо этого использует одну сеть и маскирует её веса, так что

$$\frac{\partial z'_i}{\partial z_j} = 0 \qquad \text{для всех } j > i .$$

Эта сеть — MADE, маскированный автоэнкодер для оценивания распределений (masked autoencoder for distribution estimation; Germain et al., 2015). Маски строятся по степеням (degrees), которые здесь обозначены $m(\cdot)$, чтобы не путать их с размерностью $d$:

• вход $z_j$ имеет степень $m(j) = j$;
• каждый скрытый нейрон $u$ получает степень $m(u)\in\{1,\dots,d-1\}$ (Germain et al., 2015, §4). Нейрону разрешено зависеть от $z_1,\dots,z_{m(u)}$ и ни от чего более позднего. Степень $d$ была бы бесполезна: такой нейрон мог бы передавать сигнал только выходам, идущим после $z_d$, а их нет;
• выходной нейрон, который вычисляет $(s_i, t_i)$, имеет степень $m(i) = i$.

В скрытом слое скрытый нейрон $u$ может читать нейрон $j$ нижележащего слоя, только если степень отправителя не больше степени получателя:

$$\begin{gathered} M_{uj} = \begin{cases} 1 & \text{если } m(j) \le m(u), \\ 0 & \text{иначе,} \end{cases} \\[6pt] h = Wx + b \;\;\longrightarrow\;\; h = (W\odot M)\,x + b . \end{gathered}$$

Маска имеет ту же форму, что и матрица весов слоя $W$, и $W\odot M$ просто обнуляет запрещённые веса. Выходной слой использует строгое неравенство: его маска $M^{\text{out}}_{iu} = 1$, если $m(u) < i$, и $0$ иначе, так что выход $i$ может читать скрытый нейрон $u$, только если $m(u) < i$. Именно поэтому $(s_i, t_i)$ зависят от $z_{<i}$, а не от самого $z_i$.

Разобранный пример. Возьмём $d = 3$ входа, один скрытый слой из четырёх нейронов со степенями $(1, 2, 1, 2),$ и выходы степеней $(1, 2, 3)$. В обеих масках строки соответствуют принимающим нейронам, а столбцы — отправляющим. Маска скрытого слоя (правило $m(j)\le m(u)$) и маска выходного слоя (правило $m(u) < i$) равны

$$\begin{gathered} M^{(1)} = \begin{bmatrix} 1&0&0\\ 1&1&0\\ 1&0&0\\ 1&1&0 \end{bmatrix}, \\[6pt] M^{(2)} = \begin{bmatrix} 0&0&0&0\\ 1&0&1&0\\ 1&1&1&1 \end{bmatrix}. \end{gathered}$$

Элемент $(i, j)$ их произведения равен числу путей от входа $j$ к выходу $i$:

$$M^{(2)}M^{(1)} = \begin{bmatrix} 0&0&0\\ 2&0&0\\ 4&2&0 \end{bmatrix}.$$

Произведение строго нижнетреугольное. Нелинейность вроде ReLU не может создать путь, который удалили маски, поэтому $s_i, t_i$ зависят только от $z_{<i}$. Затем поэлементный шаг добавляет диагональ $\exp(s_i).$ Выход 1 вообще не получает связей, так что $s_1, t_1$ — обучаемые константы (смещения), как и должно быть: перед $z_1$ ничего нет. Вход $z_3$ никуда не передаётся, потому что ни одной более поздней координате он не нужен.

Рисунок

Маскирование весов делает сеть авторегрессионной

MAF и IAF. Маскированная сеть выдаёт все $s_i, t_i$ за один проход, когда известны переменные, от которых они зависят. В MAF они известны в направлении от данных к латентному пространству, поэтому плотность точки данных вычисляется дёшево. При сэмплировании слой приходится обращать координата за координатой: для $z_i = (z'_i - t_i(z_{<i}))\,e^{-s_i(z_{<i})}$ сначала нужно $z_{<i}$, что занимает $d$ последовательных проходов. IAF, обратный авторегрессионный поток (inverse autoregressive flow; Kingma et al., 2016), меняет два направления местами: его $s_i, t_i$ читают латентные координаты с номерами меньше $i$. Сэмплирование занимает один проход, зато вычисление плотности произвольной заданной точки последовательно. При этом IAF может за один проход вычислить плотность собственных сэмплов, потому что их латентные векторы известны, а именно это и нужно вариационному выводу (variational inference). Слои связи дёшевы в обоих направлениях, потому что $s$ и $t$ читают только $z_a$, который проходит без изменений и известен в обоих направлениях. Плата за это — каждый слой меняет только половину координат.

В. Обратимые свёртки (Glow)#

В слоях Glow используются обратимые свёртки $1\times1$:

$$z' = W z .$$

Для тензора-изображения высоты $h$, ширины $w$ с $C$ каналами свёртка $1\times1$ применяет одну и ту же матрицу $W$ размера $C\times C$ к вектору из $C$ значений каналов в каждом пикселе. Она перемешивает каналы и ничего не делает по пространству. Перемешивание каналов — именно то, что нужно между слоями связи: это обучаемое обобщение фиксированной перестановки каналов (например, обращения их порядка), которую RealNVP использует между слоями связи (Kingma and Dhariwal, 2018, §3.2). Матрица Якоби всего слоя блочно-диагональная, на её диагонали стоят $h\cdot w$ копий $W$, поэтому

$$\log\lvert\det J\rvert = h\cdot w\cdot\log\lvert\det W\rvert .$$

$W$ должна быть обратимой, а вычисление $\det W$ в общем случае стоит $O(C^3)$. Поэтому $W$ обучают через её LU-разложение. Не у всякой обратимой матрицы есть LU-разложение без выбора ведущего элемента (pivoting): у перестановки $\left[\begin{smallmatrix}0&1\\1&0\end{smallmatrix}\right]$ его нет, потому что $(LU)_{11} = U_{11}$ должно было бы равняться $0$, и тогда $(LU)_{21} = L_{21}U_{11} = 0 \ne 1$. Поэтому Glow добавляет фиксированную матрицу перестановки $P$:

$$\begin{gathered} W = PLU, \\[6pt] L = \begin{bmatrix} 1 & & 0 \\ \vdots & \ddots & \\ L_{C1} & \cdots & 1 \end{bmatrix}, \\[6pt] U = \begin{bmatrix} U_{11} & \cdots & U_{1C} \\ & \ddots & \vdots \\ 0 & & U_{CC} \end{bmatrix}. \end{gathered}$$

Элементы $L_{ij}$ и $U_{ij}$ — вещественные числа. $P$ остаётся фиксированной; обучаемые параметры — свободные элементы $L$ (под её единичной диагональю) и $U$ (на её диагонали и выше). Поскольку $\det P = \pm1$, $\det L = 1$ и $\det U = \prod_{i=1}^{C} U_{ii}$,

$$\log\lvert\det W\rvert = \sum_{i=1}^{C} \log\lvert U_{ii}\rvert .$$

Слой обратим ровно тогда, когда все $U_{ii}\neq 0$, а обратная матрица — $W^{-1} = U^{-1}L^{-1}P^{\top}$: перестановка и решение двух треугольных систем. В отличие от $\exp(s_i)$, свободный параметр $U_{ii}$ сам по себе ничем не защищён от обращения в ноль.

Рисунок

Обратимая свёртка 1×1 в Glow

Например, возьмём $C = 3$ и $P = I$ с

$$\begin{gathered} L = \begin{bmatrix} 1&0&0\\ 0.5&1&0\\ -1&2&1 \end{bmatrix},\quad U = \begin{bmatrix} 2&1&-1\\ 0&-0.5&3\\ 0&0&1.5 \end{bmatrix} \\[6pt] \Longrightarrow\quad W = LU = \begin{bmatrix} 2&1&-1\\ 1&0&2.5\\ -2&-2&8.5 \end{bmatrix}. \end{gathered}$$

Непосредственное вычисление $\det W$ даёт $-1.5$, что равно $2\cdot(-0.5)\cdot1.5$. Отсюда $\log\lvert\det W\rvert = \log 1.5 \approx 0.405$. Отрицательный ведущий элемент $U_{22} = -0.5$ показывает, зачем нужен модуль: $\log(-0.5)$ не определён.

Г. Сплайновые слои связи: нейросетевые сплайновые потоки#

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

Внутри слоя связи или авторегрессионного слоя каждая преобразуемая координата проходит через одномерное отображение, параметры которого выдаёт сеть: $z_{b,i}\mapsto z_{b,i}\,e^{s_i}+t_i$ в части А. Аффинное отображение может сдвинуть и растянуть условное распределение, но не может изменить его форму: одномодальное условное распределение после одного слоя остаётся одномодальным. Рассуждению о треугольной матрице Якоби не нужно, чтобы отображение было аффинным. Подойдёт любая строго возрастающая дифференцируемая функция $g_\theta$ с вычислимой обратной. Диагональным элементом $J$ становится $g_\theta'(z_{b,i})$, а $\theta$ выдаёт сеть того же рода, что и раньше: $\theta = \mathrm{NN}(z_a)$ в слое связи, $\theta_i = \mathrm{NN}(z_{<i})$ в авторегрессионном слое.

Рационально-квадратичный сплайн (rational-quadratic spline). Нейросетевые сплайновые потоки (neural spline flows; Durkan et al., 2019, §3.1) используют монотонный сплайн из рационально-квадратичных кусков (отношений двух квадратичных многочленов), построенный по схеме Gregory and Delbourgo (1982). В этой части $x$ — скалярный вход отображения, а не точка данных. Сплайн отображает $[-B, B]$ на себя и тождествен вне этого отрезка. Внутри он задаётся $K$ бинами (bins) с $K + 1$ узлами (knots) $(x_k, y_k)$, возрастающими от $(x_0, y_0) = (-B, -B)$ до $(x_K, y_K) = (B, B)$, и положительными производными $\delta_k$ в узлах. Две крайние производные фиксированы, $\delta_0 = \delta_K = 1$, чтобы наклон совпадал с наклоном тождественных хвостов. В бине $k$ положим

$$s_k = \frac{y_{k+1} - y_k}{x_{k+1} - x_k}, \qquad \xi = \frac{x - x_k}{x_{k+1} - x_k}\in[0, 1]$$

($s_k$ — наклон хорды, стягивающей бин, а не сеть масштаба из части А). Сплайн в бине $k$ имеет вид

$$g(x) = y_k + \frac{(y_{k+1} - y_k)\bigl[s_k\,\xi^2 + \delta_k\,\xi(1-\xi)\bigr]}{s_k + (\delta_{k+1} + \delta_k - 2s_k)\,\xi(1-\xi)} .$$

При $\xi = 0$ он даёт $y_k$, а при $\xi = 1$ дробь равна $s_k/s_k = 1$, так что он даёт $y_{k+1}$: каждый кусок соединяет два соседних узла. Знаменатель можно переписать как $s_k\bigl(\xi^2 + (1-\xi)^2\bigr) + (\delta_k + \delta_{k+1})\,\xi(1-\xi)$, и он положителен. Производная сплайна равна

$$g'(x) = \frac{s_k^2\bigl[\delta_{k+1}\xi^2 + 2s_k\,\xi(1-\xi) + \delta_k(1-\xi)^2\bigr]}{\bigl[s_k + (\delta_{k+1} + \delta_k - 2s_k)\,\xi(1-\xi)\bigr]^2} ,$$

причём все слагаемые числителя положительны, поэтому $g$ строго возрастает. При $\xi = 0$ производная равна $\delta_k$, а при $\xi = 1$ — $\delta_{k+1}$: соседние куски стыкуются с одинаковым наклоном, а с хвостами сплайн стыкуется с наклоном 1. Логарифм определителя слоя равен $\sum_i\log g'(z_{b,i})$ по преобразуемым координатам; координаты в хвостах дают вклад $\log 1 = 0$.

Параметры. Для каждой преобразуемой координаты сеть выдаёт $3K - 1$ чисел без ограничений (Durkan et al., 2019, §3.1): $K$ для ширин бинов и $K$ для их высот, которые проходят через softmax и умножаются на $2B$, чтобы быть положительными и заполнять отрезок, и $K - 1$ для внутренних производных, которые проходят через softplus. Кумулятивные суммы ширин и высот, начиная с $-B$, дают узлы.

Обратное отображение — корень квадратного уравнения. Поскольку $g$ возрастает, бин выхода $y$ находится сравнением $y$ со значениями $y_k$. Приравняв $g(x) = y$ и избавившись от знаменателя, получаем $a\xi^2 + b\xi + c = 0$, где

$$\begin{aligned} a &= (y_{k+1} - y_k)(s_k - \delta_k) + (y - y_k)(\delta_{k+1} + \delta_k - 2s_k), \\ b &= (y_{k+1} - y_k)\,\delta_k - (y - y_k)(\delta_{k+1} + \delta_k - 2s_k), \\ c &= -s_k\,(y - y_k). \end{aligned}$$

Ровно один корень лежит в $[0, 1]$ (Durkan et al., 2019, приложение A.3):

$$\xi = \frac{2c}{-b - \sqrt{b^2 - 4ac}},$$

это обычная формула $(-b + \sqrt{b^2 - 4ac})/(2a)$, переписанная так, чтобы не терять точность при малом $4ac$. Затем $x = x_k + \xi\,(x_{k+1} - x_k)$. И вычисление, и обращение стоят одного поиска бина и нескольких арифметических операций, так что сплайновый слой связи, как и аффинный, требует одного прохода сети в каждую сторону.

Рисунок

Монотонный рационально-квадратичный сплайн и его производная

Разобранный пример (игрушечные числа, сплайн с рисунка). $K = 4$ бина на $[-3, 3]$ с узлами $(-3,-3)$, $(-1.2,-2.4)$, $(0, 1.2)$, $(1.5, 2.2)$, $(3, 3)$, внутренними производными $(\delta_1, \delta_2, \delta_3) = (0.4, 2.5, 0.5)$ и $\delta_0 = \delta_4 = 1$. Вход $x = -0.6$ лежит в бине $k = 1$: $x_1 = -1.2$, ширина $1.2$, $\xi = 0.5$, высота $y_2 - y_1 = 3.6$ и $s_1 = 3$. Тогда $\xi(1-\xi) = 0.25$ и

$$\begin{aligned} g(-0.6) &= -2.4 + \frac{3.6\,(3\cdot0.25 + 0.4\cdot0.25)}{3 + (2.5 + 0.4 - 6)\cdot0.25} \\ &= -2.4 + \frac{3.06}{2.225} \approx -1.025, \\[4pt] g'(-0.6) &= \frac{9\,(2.5\cdot0.25 + 2\cdot3\cdot0.25 + 0.4\cdot0.25)}{2.225^2} \approx 4.04 . \end{aligned}$$

Эта координата добавляет $\log 4.04 \approx 1.40$ к $\log\lvert\det J\rvert$. Обратим $y = -1.0247$: $y - y_1 = 1.3753$, поэтому $a \approx 5.097$, $b \approx 5.703$, $c \approx -4.126$, $b^2 - 4ac \approx 116.6$ и $\xi = 2c/(-b - \sqrt{b^2 - 4ac}) \approx -8.252/-16.503 = 0.500$, что возвращает $x = -1.2 + 0.5\cdot1.2 = -0.6$.

Версии со слоями связи и авторегрессионная. Сплайн заменяет аффинное отображение в любой из двух конструкций. Durkan et al. называют результаты RQ-NSF (C) и RQ-NSF (AR) (§3.2): версия со слоями связи — это Glow, в котором аффинные слои связи заменены сплайновыми, а авторегрессионная — это MAF или IAF, в которых аффинные отображения заменены сплайнами. В их слое связи тождественная половина тоже проходит через поэлементные сплайны, параметры которых обучаются напрямую, а не предсказываются сетью. При достаточном числе бинов монотонный сплайн приближает любую дифференцируемую монотонную функцию на $[-B, B]$; на практике авторы фиксируют $B = 3$ и $K = 8$, не обнаружив большой разницы при $B$ от 1 до 5 (§5).

Д. ActNorm: поканальные масштаб и сдвиг с инициализацией по данным#

Глубокие потоки лучше обучаются со слоем нормализации, а в потоке такой слой должен быть обратимым с известным определителем. RealNVP использует батч-нормализацию (batch normalization; Dinh, Sohl-Dickstein and Bengio, 2017, §3.7): со статистиками батча $\tilde\mu, \tilde\sigma^2$ отображение $x\mapsto(x - \tilde\mu)/\sqrt{\tilde\sigma^2 + \epsilon}$ — покоординатное масштабирование с логарифмом определителя $-\tfrac12\sum_i\log(\tilde\sigma_i^2 + \epsilon)$. Эта формула считает $\tilde\mu$ и $\tilde\sigma^2$ константами, заданными извне примера, как в RealNVP, где они берутся из скользящего среднего по последним мини-батчам (прил. E). Обычная батч-нормализация во время обучения, которую дифференцируют как функцию текущего батча, — другое отображение: каждый выход зависит от всех примеров батча, а стандартизация убирает собственное среднее и масштаб батча, так что по одним выходам входы не восстанавливаются. Батч-нормализация вносит шум через статистики батча, этот шум растёт по мере уменьшения батча на одно устройство, а Glow обучается на больших изображениях по одному изображению на устройство (§3.1). Его слой ActNorm сохраняет аффинную форму и отказывается от статистик батча (Kingma and Dhariwal, 2018, §3.1, таблица 1). С логарифмом масштаба $\gamma_c$ и сдвигом $\tau_c$ для каждого канала $c$

$$\begin{gathered} z'_{ijc} = e^{\gamma_c}\,z_{ijc} + \tau_c, \\[4pt] \log\lvert\det J\rvert = h\cdot w\cdot\sum_{c=1}^{C}\gamma_c . \end{gathered}$$

Матрица Якоби диагональна: элемент $e^{\gamma_c}$ повторяется в каждой из $h\cdot w$ позиций канала $c$. Glow записывает сам масштаб $s_c$ и логарифм определителя $h\cdot w\sum_c\log\lvert s_c\rvert$; логарифм масштаба автоматически делает масштаб положительным. Обратное отображение — $z_{ijc} = (z'_{ijc} - \tau_c)\,e^{-\gamma_c}$. На плоских векторах у того же слоя своя пара $(\gamma_i, \tau_i)$ для каждой координаты и $h = w = 1$; он обобщает слой диагонального масштабирования, который NICE ставит в конце своего потока (Dinh, Krueger and Bengio, 2015, §3.3).

Инициализация по данным (data-dependent initialization). Параметры задаются по первому мини-батчу так, чтобы каждый канал выхода слоя имел на нём нулевое среднее и единичную дисперсию:

$$\gamma_c = -\log\operatorname{std}_c(z), \qquad \tau_c = -\operatorname{mean}_c\bigl(e^{\gamma_c} z\bigr),$$

что даёт $z'_{ijc} = (z_{ijc} - \operatorname{mean}_c)/\operatorname{std}_c$ на этом батче. Дальше $\gamma$ и $\tau$ — обычные обучаемые параметры, которые, в отличие от батч-нормализации, не зависят от данных. Поэтому каждый слой нового потока начинает с хорошо отмасштабированных активаций, какой бы ни была глубина.

Например (игрушечные числа), ActNorm с логарифмами масштаба $\gamma = (0.5, -0.2, 0.1)$ на тензоре $32\times32\times3$ даёт $\log\lvert\det J\rvert = 1024\cdot0.4 = 409.6$ ната (nats).

Один шаг Glow — это ActNorm, затем свёртка $1\times1$, затем аффинный слой связи (Kingma and Dhariwal, 2018, §3). Последняя свёртка каждой сети слоя связи инициализируется нулями, так что каждый слой связи начинает как тождественное отображение (§3.3). Обе инициализации преследуют одну цель: глубокая стопка обратимых слоёв должна начинать с отображения, близкого к хорошо отмасштабированному тождественному.

Е. Потоки на изображениях: squeeze, маски и многомасштабность#

RGB-изображение размером $64\times64$ пикселя имеет $d = 64\cdot64\cdot3 = 12288$ координат. Слои выше работают с любым вектором, но поток для изображений должен использовать в качестве $s$ и $t$ свёрточные сети, а значит, разбиение на $z_a$ и $z_b$ должно учитывать сетку пикселей.

Маски на сетке пикселей#

RealNVP записывает слой связи с бинарной маской $b$ той же формы, что и тензор-изображение (Dinh, Sohl-Dickstein and Bengio, 2017, §3.4):

$$\begin{gathered} z' = b\odot z + (1 - b)\odot\Bigl(z\odot\exp\bigl(s(b\odot z)\bigr) + t(b\odot z)\Bigr), \\[4pt] \log\lvert\det J\rvert = \sum (1 - b)\odot s(b\odot z), \end{gathered}$$

где сумма берётся по всем элементам тензора. Элементы с $b = 1$ проходят без изменений и служат условием для остальных, ровно как $z_a$ в части А, а $s$, $t$ — свёрточные сети, отображающие тензор $h\times w\times C$ в тензор той же формы. Используются две маски:

• шахматная (checkerboard): $b_{ij} = 1$ там, где $i + j$ нечётно, одинаково во всех каналах, так что четыре непосредственных соседа каждого преобразуемого пикселя входят в число условных пикселей;
• поканальная (channel-wise): $b = 1$ на первой половине каналов и $0$ на второй.

Последовательные слои чередуют $b$ и $1 - b$, чтобы преобразовывался каждый элемент.

Squeeze: обмен пространства на каналы#

Операция squeeze разрезает каждый канал на квадраты $2\times2$ и укладывает четыре значения каждого квадрата в четыре канала, превращая тензор $h\times w\times C$ в тензор $\tfrac h2\times\tfrac w2\times 4C$ (§3.6). Если канал $4c + 2p + q$ выхода хранит пиксели со смещением $(p, q)\in\{0,1\}^2$ из канала $c$, то вход $4\times4\times1$ превращается в

$$\begin{bmatrix} 1&2&3&4\\ 5&6&7&8\\ 9&10&11&12\\ 13&14&15&16 \end{bmatrix} \;\longrightarrow\; \begin{bmatrix} 1&3\\ 9&11 \end{bmatrix},\; \begin{bmatrix} 2&4\\ 10&12 \end{bmatrix},\; \begin{bmatrix} 5&7\\ 13&15 \end{bmatrix},\; \begin{bmatrix} 6&8\\ 14&16 \end{bmatrix}.$$

Squeeze лишь переставляет координаты, поэтому её матрица Якоби — матрица перестановки, а логарифм определителя равен $0$. Эта операция объясняет и то, почему две маски дополняют друг друга. С индексами от 0 шахматная маска этого входа $4\times4$ равна

$$b = \begin{bmatrix} 0&1&0&1\\ 1&0&1&0\\ 0&1&0&1\\ 1&0&1&0 \end{bmatrix},$$

и после squeeze её единицы заполняют каналы 1 и 2 (смещения $(0,1)$ и $(1,0)$), а нули — каналы 0 и 3. Поканальная маска делит каналы после squeeze на $\{0, 1\}$ и $\{2, 3\}$ — это другое разбиение: каналы 0 и 3, которые шахматная маска относила к одной стороне, попадают в разные половины, как и каналы 1 и 2, поэтому у слоёв связи после squeeze в условии каждой группы оказываются пиксели, которые шахматные слои связи держали вместе.

Выделение половины переменных на каждом масштабе#

Многомасштабная архитектура (multi-scale architecture) RealNVP (§3.6) повторяет один блок на каждом масштабе: три слоя связи с чередующимися шахматными масками, squeeze и три слоя связи с чередующимися поканальными масками. После каждого блока половина переменных выделяется (factored out): эти переменные покидают поток и становятся частью латентного $z$, а к следующему масштабу переходит только другая половина. В направлении от данных к латентному пространству, где $\Phi^{(i)}$ — блок масштаба $i$ и $h^{(0)} = x$,

$$\begin{aligned} \bigl(z^{(i)},\, h^{(i)}\bigr) &= \Phi^{(i)}\bigl(h^{(i-1)}\bigr), \qquad i = 1,\dots,L-1, \\ z^{(L)} &= \Phi^{(L)}\bigl(h^{(L-1)}\bigr), \\ z &= \bigl(z^{(1)}, \dots, z^{(L)}\bigr). \end{aligned}$$

На последнем масштабе четыре шахматных слоя связи и нет squeeze, а рекурсия останавливается, когда её вход имеет размер $4\times4\times C$ (§4.1). Плотность каждой выделенной части вычисляется по гауссовской базовой плотности сразу, как только часть покидает поток, так что функция потерь распределена по всей сети, а поздние масштабы дешевле, потому что обрабатывают меньше переменных. Для изображения $64\times64\times3$ этот рецепт даёт пять масштабов:

Масштаб	Вход	После squeeze	Выделено	Идёт дальше
1	$64\times64\times3$	$32\times32\times12$	6144	$32\times32\times6$
2	$32\times32\times6$	$16\times16\times24$	3072	$16\times16\times12$
3	$16\times16\times12$	$8\times8\times48$	1536	$8\times8\times24$
4	$8\times8\times24$	$4\times4\times96$	768	$4\times4\times48$
5	$4\times4\times48$	нет	768

Выделенные части в сумме дают $6144 + 3072 + 1536 + 768 + 768 = 12288 = d$: многомасштабный поток по-прежнему сохраняет все измерения. Его 28 слоёв связи обрабатывают $6\cdot(1 + \tfrac12 + \tfrac14 + \tfrac18) + 4\cdot\tfrac1{16} = 11.5$ полных изображений в пересчёте на число переменных, около 41% того, что обработали бы 28 слоёв связи на всех $12288$ переменных. Переменные, выделенные рано, прошли через меньшее число слоёв, и RealNVP описывает их как более мелкие, более локальные признаки (§3.6, приложение D). В RealNVP сеть масштаба заканчивается $\tanh$, умноженным на обучаемый множитель, а сеть сдвига имеет аффинный выход (§4.1).

Шаг Glow на тензорах-изображениях#

Glow сохраняет многомасштабный каркас с $L$ уровнями: каждый уровень — это squeeze, за которым идут $K$ шагов из ActNorm, свёртки $1\times1$ и аффинного слоя связи, и каждый уровень, кроме последнего, заканчивается разделением (split), которое выделяет половину переменных (Kingma and Dhariwal, 2018, §3, рисунок 2). Его слои связи делят тензор только по каналам — это упрощение по сравнению с RealNVP (§3.3), а обучаемые свёртки $1\times1$ заменяют фиксированное обращение порядка каналов между слоями связи. Все три слоя действуют на тензоры $h\times w\times C$ с поканальными параметрами, и их логарифмы определителей содержат множитель $h\cdot w$ (части В и Д). Свёртка $1\times1$ после одного squeeze RGB-изображения имеет $C = 12$ и всего $12^2 = 144$ веса. Сети слоёв связи состоят из трёх свёрток ($3\times3$, $1\times1$, $3\times3$) с 512 скрытыми каналами, а в сравнениях на CIFAR-10 используются $K = 32$ и $L = 3$ (§5). Лучшие результаты Glow — 3.35 бита на размерность (bits per dimension) на CIFAR-10 против 3.49 у RealNVP и 3.81 против 3.98 на ImageNet $64\times64$ (§5, таблица 2). Для лиц CelebA-HQ размером $256\times256$ он использует $K = 32$ и $L = 6$ на 5-битных изображениях (§6).

Альтернатива: развернуть изображение в вектор#

Простейший поток для изображений превращает изображение в вектор из $d = 12288$ чисел и применяет слои частей А–Г с полносвязными сетями. Это работает и не требует ни масок, ни squeeze, но структура сетки теряется. Разбиение по индексам развёрнутого тензора с каналами на первой оси (channel-first) помещает в $z_a$ весь красный канал и верхнюю половину зелёного; сетям приходится с нуля учиться, какие координаты соседние; а первый слой сети слоя связи с 1024 скрытыми нейронами уже имеет $6144\cdot1024\approx6.3$ миллиона весов. Обучаемая плотная матрица перемешивания между слоями связи имела бы $d^2\approx1.5\cdot10^8$ элементов, поэтому на практике выбирают фиксированную перестановку, например случайную. На CIFAR-10 NICE, слои связи которого — полносвязные сети на развёрнутом изображении, достигает 4.48 бита на размерность против 3.49 у свёрточного многомасштабного RealNVP (таблица). Эти две модели различаются также глубиной и видом слоёв связи (аддитивные или аффинные), так что не весь разрыв объясняется сеткой.

Правдоподобие потоков для изображений сообщают в битах на размерность после отмены всех шагов предобработки, то есть к логарифму определителя самой модели добавляется логарифм определителя предобработки; раздел Изображения: деквантование, масштабирование и биты на размерность разбирает это для распространённых способов масштабирования и для логит-преобразования RealNVP.

Интерполяция и кодирование с точным обратным отображением#

У потока есть точный энкодер: у любого изображения $x$, в том числе не из обучающей выборки, ровно один латентный вектор $z = f^{-1}(x)$, и $f(z)$ возвращает $x$ с точностью до ошибок округления. VAE же кодирует изображение в параметры $q_\phi(z\mid x)$, и его восстановление точно настолько, насколько точен декодер: распространённый попиксельный гауссовский декодер или декодер с квадратичной ошибкой возвращает среднее $p_\theta(x\mid z)$, которое обычно размыто, а более выразительное правдоподобие декодера может давать резкие сэмплы (VAE; варианты — в заметке Варианты VAE). Отсюда три применения.

Интерполяция. Закодируем два изображения, пройдём по пути между $z_1$ и $z_2$ и декодируем. Прямая $z_\lambda = (1-\lambda)z_1 + \lambda z_2$ выходит из области, где базовое распределение сосредоточивает свою массу. Для $z \sim \mathcal N(0, I_d)$ длина $\lVert z\rVert$ сосредоточена около $\sqrt d$ с разбросом около $0.7$ (оболочка гауссианы). Если $z_1$ и $z_2$ ведут себя как независимые сэмплы, длина их середины — около $\sqrt{d/2}$. При $d = 12288$ это $78.4$ вместо $110.9$, примерно на 46 разбросов ниже типичной длины. Веса, квадраты которых в сумме дают единицу, сохраняют длину типичной. Пути RealNVP между латентными векторами используют такие веса $\cos$ и $\sin$ (§4.2, формула 19), и середина $\cos\tfrac\pi4\,z_1 + \sin\tfrac\pi4\,z_2$ имеет длину около $\sqrt d$. Glow интерполирует линейно между латентными векторами настоящих лиц и всё же сообщает о реалистичных промежуточных изображениях (§6).

Температура. Glow сэмплирует из модели с пониженной температурой: он умножает стандартное отклонение базового распределения на $T = 0.7$, жертвуя разнообразием ради качества (§6). Для потока, сохраняющего объём, $z\sim\mathcal N(0, T^2 I)$ даёт плотность $\propto\exp\bigl(-\lVert f^{-1}(x)\rVert^2/(2T^2)\bigr) \propto p_\theta(x)^{1/T^2}$ — заострённую версию модели (в статье показатель степени записан как $T^2$).

Кодирование внешнего изображения. Предобработаем его точно так же, как обучающие данные, вычислим $z = f^{-1}(x)$ и сравним с латентными векторами обучающих изображений, например по евклидову расстоянию. Здесь две оговорки. Если латентные векторы ведут себя как независимые сэмплы из базового распределения, все попарные расстояния близки к $\sqrt{2d}$, так что ближайшие соседи лишь немного ближе типичных точек. А точное правдоподобие внешнего изображения не говорит, похоже ли оно на обучающие данные: Glow, обученный на CelebA, даёт номерам домов SVHN — набору данных, которого он никогда не видел, — более высокое правдоподобие, чем лицам CelebA (Nalisnick et al., 2019, §3, рисунок 2).

Ж. Планарные и радиальные потоки: дешёвые определители без треугольных матриц#

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

Первые потоки из работы Rezende and Mohamed (2015), §4.1, приходят к дешёвому определителю другим путём. Их матрица Якоби — единичная матрица плюс матрица ранга один, и лемма об определителе матрицы (matrix determinant lemma) даёт её определитель за $O(d)$ (Papamakarios et al., 2021, §3.3.2): для $a, c\in\mathbb R^d$

$$\det\bigl(I + a\,c^\top\bigr) = 1 + c^\top a .$$

Любой вектор, ортогональный $c$, матрица $I + ac^\top$ оставляет без изменений, поэтому собственное значение $1$ имеет кратность не меньше $d - 1$; тогда след $d + c^\top a$ фиксирует последнее собственное значение равным $1 + c^\top a$, а определитель — произведение собственных значений.

Планарные потоки (planar flows). С $u, w\in\mathbb R^d$ и скалярным смещением $b$ (это не маска из части Е)

$$\begin{gathered} f(z) = z + u\,\tanh(w^\top z + b), \\[4pt] J = I + \tanh'(w^\top z + b)\,u\,w^\top, \\[4pt] \det J = 1 + \tanh'(w^\top z + b)\,w^\top u . \end{gathered}$$

Rezende и Mohamed допускают вместо $\tanh$ любую гладкую поэлементную нелинейность; условие обратимости ниже относится к $\tanh$.

Отображение сдвигает $z$ только вдоль $u$, на величину, которая зависит от знакового расстояния до гиперплоскости $w^\top z + b = 0$, так что оно растягивает или сжимает плотность поперёк этой гиперплоскости.

Когда оно обратимо? Разложим $z$ на часть, ортогональную $w$, и часть $\eta\,w/\lVert w\rVert^2$ вдоль него, так что $w^\top z = \eta$. Умножив $y = f(z)$ слева на $w^\top$, получим скалярное уравнение

$$w^\top y = \eta + w^\top u\,\tanh(\eta + b) =: g(\eta),$$

а когда $\eta$ известно, $z = y - u\tanh(\eta + b)$. Значит, отображение $f$ обратимо ровно тогда, когда обратима функция $g$. Поскольку $0 < \tanh' \le 1$, причём $\tanh' = 1$ только при $\eta = -b$, получаем $g'(\eta) = 1 + w^\top u\,\tanh'(\eta + b)\ge 1 + \min(w^\top u, 0)$. Если $w^\top u\ge -1$, то $g$ строго возрастает и принимает все значения из $\mathbb R$, и $f$ — биекция (Rezende and Mohamed, 2015, приложение A.1). Если $w^\top u < -1$, то $g'(-b) = 1 + w^\top u < 0$, а вдали $g' \to 1$, так что $g$ не монотонна и $f$ складывает пространство само на себя. Диффеоморфизму нужно ещё, чтобы $\det J > 0$ всюду, а это выполняется при $w^\top u > -1$; Rezende и Mohamed обеспечивают это, заменяя $u$ на $u + [m(w^\top u) - w^\top u]\,w/\lVert w\rVert^2$, где $m(x) = -1 + \log(1 + e^x) > -1$. Даже у обратимого планарного потока обратное отображение не выписывается в замкнутом виде: для каждого слоя нужен одномерный поиск корня $\eta$.

Например (игрушечные числа), возьмём $d = 2$, $w = (1, 0)$, $u = (-0.5, 1)$, $b = 0$, так что $w^\top u = -0.5$. В точке $z = (0.5, 0)$: $\tanh'(0.5) \approx 0.786$, $\det J \approx 1 - 0.5\cdot0.786 \approx 0.607$ и $y = f(z) \approx (0.269, 0.462)$. Для обращения нужен корень уравнения $\eta - 0.5\tanh\eta = 0.269$, то есть $\eta = 0.5$, а затем $z = y - u\tanh 0.5 = (0.5, 0)$. Если же взять $u = (-2, 1)$, то $w^\top u = -2$, и $g(\eta) = \eta - 2\tanh\eta$ обращается в ноль при $\eta = 0$ и $\eta\approx\pm1.915$: у каждой точки $y$ с $w^\top y = 0$ три прообраза.

Радиальные потоки (radial flows). С центром $z_0\in\mathbb R^d$, $\alpha > 0$, $\beta\in\mathbb R$ и $r = \lVert z - z_0\rVert$

$$f(z) = z + \frac{\beta}{\alpha + r}\,(z - z_0).$$

Отображение сохраняет направление от $z_0$ и меняет только расстояние, $r\mapsto\rho(r) = r\,\bigl(1 + \tfrac{\beta}{\alpha + r}\bigr)$. Матрица Якоби такого радиального отображения — это $(\rho(r)/r)\,I$ плюс слагаемое ранга один вдоль $z - z_0$: она растягивает $d - 1$ направлений, ортогональных $z - z_0$, с коэффициентом $\rho(r)/r = 1 + \beta/(\alpha + r)$, а радиальное направление — с коэффициентом $\rho'(r) = 1 + \beta/(\alpha + r) - \beta r/(\alpha + r)^2 = 1 + \alpha\beta/(\alpha + r)^2$. Определитель — произведение этих коэффициентов растяжения:

$$\det J = \Bigl(1 + \frac{\beta}{\alpha + r}\Bigr)^{d-1}\Bigl(1 + \frac{\alpha\beta}{(\alpha + r)^2}\Bigr).$$

Оба множителя положительны при всех $r\ge 0$, когда $\beta > -\alpha$, и тогда $\rho$ возрастает от $0$ до $\infty$, так что $f$ — биекция; Rezende и Mohamed формулируют условие $\beta\ge-\alpha$ и обеспечивают его заменой $\beta = -\alpha + \log(1 + e^{\hat\beta})$ (приложение A.2). В отличие от планарного потока, радиальный поток обращается в замкнутом виде: расстояние $R = \lVert y - z_0\rVert$ выхода удовлетворяет уравнению $R(\alpha + r) = r(\alpha + \beta + r)$, квадратному относительно $r$,

$$r^2 + (\alpha + \beta - R)\,r - \alpha R = 0,$$

произведение корней которого равно $-\alpha R\le 0$, так что ровно один корень неотрицателен; затем $z = z_0 + (y - z_0)/\bigl(1 + \tfrac{\beta}{\alpha + r}\bigr)$. Например (игрушечные числа), при $d = 2$, $\alpha = 1$, $\beta = -0.8$ у точки на расстоянии $r = 1$ от центра $\det J = (1 - 0.4)\cdot(1 - 0.2) = 0.48$, и она переходит на расстояние $R = 0.6$; уравнение $r^2 - 0.4r - 0.6 = 0$ имеет неотрицательный корень $r = 1$.

Где они применяются. Каждый слой меняет форму плотности вдоль одного направления или вокруг одной точки, поэтому слоёв нужно много, а обращение планарного слоя итеративно. Ни то ни другое не мешает вариационному выводу, для которого эти потоки и создавались: там берут сэмпл $\varepsilon$ из простого $q_0$, проводят его вперёд через $K$ слоёв до $z_K$ и накапливают $\log q_K(z_K) = \log q_0(\varepsilon) - \sum_k\log\lvert\det J_k\rvert$, используя только прямые матрицы Якоби, как в IAF. Планарные и радиальные потоки в основном так и применяли — как гибкие апостериорные распределения для VAE (Papamakarios et al., 2021, §3.3.2), тогда как слои частей А–Е создавались для точных плотностей данных.

Как слои деформируют пространство и остаются обратимыми#

Матрица Якоби описывает, как поток деформирует пространство. Каждый из слоёв выше растягивает одни области и сжимает другие, а $\lvert\det J\rvert$ в точке — локальный множитель, на который слой умножает малые объёмы. Это становится видно, если пропустить регулярную сетку через обученный поток по одному слою за раз и нарисовать изогнутые линии сетки. Рисунок ниже делает это для одного слоя связи с простыми, выбранными вручную $s$ и $t$.

Рисунок

Матрица Якоби измеряет, как слой деформирует пространство

Обратимость по построению. Точному методу максимального правдоподобия нужны единственная латентная точка $f^{-1}(x)$ для каждой точки данных и матрица Якоби, которая нигде не схлопывает объём; общее рассуждение приведено в разделе о замене переменных предыдущей заметки, а раздел о треугольной матрице Якоби показывает, когда треугольное отображение обратимо. Для слоёв этой заметки:

• Аффинные слои связи и авторегрессионные слои обратимы при любых $s$ и $t$: их диагональные элементы $\exp(s_i)$ положительны, а обратные отображения выписываются явно.
• Сплайновые слои обратимы при любом выходе сети: каждый кусок строго возрастает, а обратное отображение — корень квадратного уравнения.
• ActNorm обратим, потому что его масштабы $e^{\gamma_c}$ положительны; слою, который обучает сам масштаб $s_c$, как в статье о Glow, нужно $s_c \ne 0$.
• Squeeze и фиксированные перестановки каналов лишь переставляют координаты.
• Свёртка $1\times1$ обратима, пока все ведущие элементы $U_{ii}\ne0$, что гарантирует параметризация через $\log\lvert U_{ii}\rvert$ с фиксированными знаками (часть В). У $\mathbf s$ из статьи о Glow такой страховки нет; там $U_{ii}\ne0$ выполняется лишь на практике.
• Планарные и радиальные потоки — диффеоморфизмы только при ограничении $w^\top u > -1$ или $\beta > -\alpha$, которое обеспечивается репараметризацией.

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

• Нет сжатия. Латентный вектор имеет размерность данных: $12288$ чисел для изображения $64\times64\times3$, каким бы простым ни было изображение. Многомасштабное выделение упорядочивает латентный вектор, но не уменьшает его.
• Ограниченные слои. Каждый слой должен быть обратимым с легко вычислимым якобианом, поэтому потоку нужно много простых слоёв, и в каждом шаге Glow два слоя из трёх лишь масштабируют или перемешивают каналы. Glow использует до 320 подпреобразований (Papamakarios et al., 2021, §3.4), и на CIFAR-10 его 3.35 бита на размерность всё ещё уступают 2.92 у авторегрессионной модели PixelCNN++.
• Нет «дыр» в плотности. Диффеоморфизм $\mathbb R^d$, применённый к гауссиане, даёт плотность, положительную всюду. Данные, лежащие в отдельных кластерах или вблизи поверхности меньшей размерности, можно лишь приблизить, с очень большими якобианами в промежутках.
• Правдоподобие — не качество. Точное правдоподобие измеряет сжатие, а не то, как выглядят сэмплы, и может ставить незнакомые изображения выше знакомых (часть Е).

Итоги#

Нормализующие потоки (normalizing flows) — обратимые нейронные сети, которые перемещают массу вероятности контролируемым образом. Каждый слой через свой определитель Якоби сообщает, насколько он растянул или сжал пространство. Авторегрессионные потоки, у которых поэлементное преобразование — достаточно гибкая монотонная функция (нейросетевые или сплайновые преобразователи, neural or spline transformers), при мягких условиях являются универсальными аппроксиматорами плотности (universal density approximators; Huang et al., 2018, §4, теорема 1; Papamakarios et al., 2021, §3.1). Аффинные слои тоже универсальны, хотя и в более слабом смысле. Композиции аффинных слоёв связи и обратимых линейных отображений (достаточно перестановок) приближают любой $C^2$-диффеоморфизм в $L^p$ на компактах, если сети $s$ и $t$ — универсальные аппроксиматоры функций, а образ базовой плотности под действием таких композиций в пределе даёт любое целевое распределение в смысле сходимости по распределению (Teshima et al., 2020, теоремы 2 и 3). Теорема отвечает на вопрос «могут ли аффинные потоки выразить всё?», но практические вопросы оставляет открытыми. Она не оценивает число слоёв; её конструкция преобразует по одной координате в каждом слое связи и относится к плоским векторам, а не прямо к свёрточным сетям Glow; а приближающие потоки могут быть сколь угодно плохо обусловлены. Koehler, Mehta и Risteski (2021) показывают, что неглубокие сети из аффинных слоёв связи универсальны в расстоянии Вассерштейна, если такая плохая обусловленность допускается, а Draxler et al. (2024) доказывают теорему универсальности, в смысле, построенном на KL-дивергенции, для хорошо обусловленных потоков на слоях связи вроде RealNVP.

Слои различаются тем, откуда берётся дешёвый определитель:

Слой	Почему $\det J$ дёшев	$\log\lvert\det J\rvert$	Стоимость двух направлений
Слой связи (NICE, RealNVP)	$z'_a = z_a$; $z'_b$ поэлементно зависит от $z_b$	$\sum_{i=1}^{d-k} s_i(z_a)$ (NICE: $0$)	по одному проходу в каждую сторону
Авторегрессионный (MAF)	$z'_i$ зависит только от $z_1,\dots,z_i$ (маски)	$\sum_{i=1}^{d} s_i(z_{<i})$	плотность: один проход; сэмплирование: $d$ проходов
Авторегрессионный (IAF)	то же, но $s_i, t_i$ читают латентные координаты с номерами меньше $i$	$\sum_{i=1}^{d} s_i$ (от латентных координат с номерами меньше $i$)	сэмплирование: один проход; плотность: $d$ проходов
Свёртка $1\times1$ (Glow)	$W = PLU$, $P$ фиксирована, $L, U$ треугольные	$h\cdot w\cdot\sum_{i=1}^{C}\log\lvert U_{ii}\rvert$	одно матричное произведение в каждую сторону
Сплайновый слой связи (NSF)	как у слоя связи, но с монотонным сплайном $g$ вместо аффинного отображения	$\sum_i \log g'(z_{b,i})$	по одному проходу в каждую сторону; обращение решает квадратное уравнение для каждой координаты
ActNorm	диагональная матрица	$h\cdot w\cdot\sum_{c=1}^{C}\gamma_c$	одно аффинное отображение в каждую сторону
Squeeze, перестановка	переупорядочение координат	$0$	бесплатно
Планарный	$I$ плюс матрица ранга один (лемма об определителе)	$\log\lvert 1 + \tanh'(w^\top z + b)\,w^\top u\rvert$	прямое направление: $O(d)$; обратное: поиск корня скалярного уравнения
Радиальный	$I$ с множителем плюс матрица ранга один	$(d-1)\log\bigl(1 + \tfrac{\beta}{\alpha + r}\bigr) + \log\bigl(1 + \tfrac{\alpha\beta}{(\alpha + r)^2}\bigr)$	прямое направление: $O(d)$; обратное: квадратное уравнение относительно $r$

• На изображениях слои связи используют шахматные и поканальные маски, squeeze обменивает пространство на каналы, а на каждом масштабе выделяется половина переменных; шаг Glow — это ActNorm, свёртка $1\times1$ и аффинный слой связи.
• Поток точно кодирует любое изображение, что позволяет интерполировать и кодировать внешние изображения; прямые пути между латентными векторами выходят из типичной оболочки базового распределения.

Ноутбук#

Ноутбук о нормализующих потоках и его библиотека nflib (HSE-LAMBDA Deep Generative Models, ветка spring-2025-long) реализуют слои частей А–Д на плоских векторах. Ячейки нумеруются с 1, считая все ячейки подряд, включая текстовые. В ноутбуке нет сохранённых выходов ячеек, поэтому он документирует процедуры, а не результаты. В nflib метод forward каждого слоя идёт в направлении данные → латентное пространство и возвращает выход вместе с логарифмом определителя.

• flows.py, слои связи. AffineHalfFlow — это часть А: прямой проход вычисляет x_b * exp(s) + t и log_det = sum(s), обратный — (z_b - t) * exp(-s) и log_det = sum(-s). Его флаг parity меняет, какая половина служит условием для другой, так что слои в стопке чередуются, а scale=False даёт аддитивный слой связи NICE. При нечётном $d$ преобразуется бо́льшая половина.
• flows.py, авторегрессионные слои. SlowMAF даёт каждой координате собственный MLP (первой — обучаемую константу, LeafParam) и в обоих направлениях проходит циклом по координатам. MAF использует один ARMLP, который за один проход возвращает $(s, t)$ для всех координат: MAF.forward работает параллельно, а MAF.inverse проходит циклом по координатам. IAF просто меняет эти два метода местами. У обоих классов MAF есть флаг parity, который обращает порядок координат на выходе, но с противоположным смыслом: MAF переворачивает порядок при parity=True (z.flip), а SlowMAF — при parity=False.
• made.py. MaskedLinear вычисляет F.linear(input, mask * weight, bias), то есть $h = (W\odot M)x + b$, а MADE.update_masks применяет правила степеней из части Б: self.m[l-1][:, None] <= self.m[l][None, :] для скрытых слоёв и self.m[L-1][:, None] < self.m[-1][None, :] для выходного слоя. Они строятся как массивы (вход × выход) со степенями, отсчитываемыми от 0, и транспонируются в MaskedLinear.set_mask.
• nets.py. У MLP три скрытых слоя ширины nh с LeakyReLU(0.2). ARMLP — это MADE с тремя скрытыми слоями ширины nh, одной маской и естественным порядком координат. PosEncMLP подаёт в MLP признаки $\sin$ и $\cos$ на частотах 0.5, 1, 2, 4 и 8; в ноутбуке он нигде не используется.
• flows.py, слои масштабирования. AffineConstantFlow — покоординатный аффинный слой, z = x * exp(s) + t с log_det = sum(s), обобщение слоя масштабирования NICE; его s и t инициализируются сэмплами из стандартного нормального распределения. ActNorm наследует от него и реализует часть Д на плоских векторах: при первом вызове прямого прохода он задаёт s = -log(x.std(0)) и t = -mean(x * exp(s)), а затем обучает их как обычные параметры. Флаг data_dep_init_done, который за этим следит, — обычный атрибут, не входящий в state_dict, поэтому модель, заново собранная и загруженная из чекпойнта, при первом прямом проходе снова выполняет инициализацию и перезаписывает загруженные s и t. После загрузки флаг нужно установить в True.
• flows.py, Invertible1x1Conv. Этот слой следует параметризации из статьи (часть В) и обучает $\mathbf s$ напрямую, без страховки от $s_i\to0$. Он начинает со случайной ортогональной матрицы, раскладывает её с помощью torch.linalg.lu, хранит $P$ как фиксированный буфер, обучает $L$, $\mathbf s$ и $U$, применяет матрицу к векторам-строкам (z = x @ W, что не меняет определитель) и возвращает log_det = sum(log(abs(S))). Обратное отображение использует torch.linalg.inv(W).
• spline_flows.py. unconstrained_RQS и RQS реализуют часть Г: тождественное отображение вне $[-B, B]$ с нулевым логарифмом определителя, крайние производные, равные 1, минимальные ширины и высоты бинов и минимальные производные $10^{-3}$, поиск бина подсчётом узлов ниже входа и обратное отображение 2c / (-b - sqrt(discriminant)) с -logabsdet. NSF_CL — сплайновый слой связи с двумя MLP: f1 читает первую половину и выдаёт $3K - 1$ параметров на каждую координату второй половины, которую он преобразует; затем f2 читает новую вторую половину и преобразует первую. Таким образом, каждый слой обновляет обе половины, тогда как в слое связи Durkan et al. (2019), формулы 9–11, условная половина получает сплайны с напрямую обучаемыми параметрами. Изменения формы тензоров (reshape) в нём предполагают чётное $d$. NSF_AR — авторегрессионная версия с обучаемым вектором параметров для первой координаты и отдельным MLP для каждой следующей; как и SlowMAF, он проходит циклом по координатам в обоих направлениях.
• Две особенности spline_flows.py. NSF_CL и NSF_AR уже пропускают ширины и высоты через softmax и умножают на $2B$, а RQS применяет softmax ещё раз. Поэтому итоговые ширины пропорциональны $\exp(2B\,p_k)$, где $p_k$ — вектор вероятностей, и ширины двух бинов могут различаться не более чем в $e^{2B}\approx403$ раза при $B = 3$. Производные тоже дважды проходят через softplus, поэтому производная во внутреннем узле никогда не бывает меньше $\ln 2 + 10^{-3}\approx 0.694$. Сплайн остаётся монотонным и обратимым, но он менее гибок, чем в статье, где каждое преобразование применяется один раз.
• Ячейка 2 ноутбука перечисляет планарные и радиальные потоки, RealNVP, MAF (в форме $f(x_i) = (x_i - \mu(x_{<i}))/\exp(\alpha(x_{<i}))$ из части Б), свёртку $1\times1$, ActNorm (как $f(x) = Wx + b$ с диагональной $W$) и NSF. nflib не реализует ни планарные, ни радиальные потоки. Ячейка 3 импортирует AffineConstantFlow, ActNorm, AffineHalfFlow, SlowMAF, MAF, IAF, Invertible1x1Conv, NSF_AR и NSF_CL, но ячейка 7 использует только ActNorm, Invertible1x1Conv и NSF_CL ($K = 8$, $B = 3$, ширина скрытого слоя 16): три блока из них на двумерных данных с логистическим базовым распределением. Каждый блок — это шаг Glow из части Д со сплайновым слоем связи, то есть слой RQ-NSF (C); при $C = 2$ и без пространственной сетки «свёртка» — это обучаемая обратимая матрица перемешивания $2\times2$.
• Ячейка 15 пропускает сетку $20\times20$ через обученный поток в обратном направлении (model.inverse) и для каждого слоя рисует перемещения точек и деформацию сетки; именно этот эксперимент рисунок с деформацией сетки имитирует для одного слоя связи. Ячейка 17 готовит ещё 500 шагов обучения с перерисовкой каждые 10 шагов (перемещения точек не более чем для девяти слоёв и деформация сетки одного слоя), но её цикл заканчивается break после первой перерисовки, так что в нынешнем виде она делает один шаг оптимизатора и рисует один раз.

Ноутбук домашнего задания 03_hw_vapnev из того же репозитория просит построить три модели CelebA и сравнить их. Его единственные сохранённые выходы ячеек — устройство (mps), форма изображения [3, 64, 64] и сетка из 25 обучающих изображений; ячейки с моделями оставлены пустыми.

• Ячейки 3–6: данные. CelebADataset из файла репозитория seminars/utils/datasets/celeba.py с Resize(64), CenterCrop(64), ToTensor() и x * 2 - 1, то есть $x' = x/127.5 - 1\in[-1, 1]$ и $d = 12288$. Шум деквантования этот конвейер не добавляет. Поток, обученный на таких значениях, видит лишь 256 различных значений на координату и может неограниченно повышать правдоподобие, сосредоточивая массу на них. Перед обучением нужно добавить равномерный шум ширины $1/127.5$, а к сообщаемому отрицательному логарифму правдоподобия прибавить $\log_2 127.5\approx 6.99$ бита на размерность (откуда берётся этот сдвиг).
• Ячейка 10, задание 1: нормализующий поток на CelebA с оценкой качества по FID и отрицательному логарифму правдоподобия. Допускаются любые слои: свёрточные потоки (часть Е) или потоки на развёрнутом изображении, которые, как предполагает текст задания, проще и слабее.
• Ячейка 12, задание 2: VAE на CelebA с оценкой качества по FID.
• Ячейка 14, задание 3: VAPNEV (Agrawal and Dukkipati, 2016) — VAE, правдоподобие которого $p(x\mid z)$ точно вычисляет real NVP вместо попиксельного гауссовского или бернуллиевского распределения, и его условная версия; см. Варианты VAE.
• Ячейка 16, задание 4: анализ латентного пространства. Какая модель лучше интерполирует, кодирование изображения не из CelebA с поиском ближайших к нему латентных векторов и самый интересный сэмпл. Сторона потоков разобрана в разделе Интерполяция и кодирование с точным обратным отображением.
• Ячейка 18 просит письменно сравнить модели, приведя сэмплы каждой из них.

Дальше#

Потоки получают точное правдоподобие ценой ограничений на архитектуру: каждый слой должен быть обратимым, с дешёвым в вычислении определителем Якоби, а латентный вектор должен сохранять размерность данных. Модель, основанная на энергии (energy-based model), снимает все эти ограничения. Любая сеть $E_\theta(x)$, выдающая число, задаёт плотность $p_\theta(x) = e^{-E_\theta(x)}/Z_\theta$, если $Z_\theta$ конечна, но сама нормирующая константа $Z_\theta = \int e^{-E_\theta(x)}\,dx$ неизвестна, поэтому ни сэмплирование, ни градиент правдоподобия нельзя вычислить напрямую. И для того и для другого тогда нужны марковские методы Монте-Карло (MCMC). Следующая заметка, сэмплирование ненормированных плотностей, строит сэмплеры, которым плотность нужна лишь с точностью до константы: MALA, гамильтонов Монте-Карло и NUTS.

Литература#

• L. Dinh, D. Krueger, Y. Bengio. NICE: Non-linear independent components estimation. ICLR Workshop, 2015. §3.3 (слой масштабирования).
• L. Dinh, J. Sohl-Dickstein, S. Bengio. Density estimation using Real NVP. ICLR, 2017. §3.4 (маски), §3.6 (squeeze и многомасштабная архитектура), §3.7 (батч-нормализация), §4.1 (предобработка и архитектура), приложение D. Эталонный код: tensorflow/models, research/real_nvp.
• M. Germain, K. Gregor, I. Murray, H. Larochelle. MADE: Masked autoencoder for distribution estimation. ICML, 2015.
• 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.
• D. P. Kingma, P. Dhariwal. Glow: Generative flow with invertible 1×1 convolutions. NeurIPS, 2018. §3.1 (ActNorm), §3.2 (свёртка 1×1), §3.3 (слой связи, инициализация нулями), §5 (таблица 2), §6 (температура, интерполяция). Эталонный код: openai/glow.
• C. Durkan, A. Bekasov, I. Murray, G. Papamakarios. Neural spline flows. NeurIPS, 2019. §3.1–3.2, §5, приложение A.
• J. A. Gregory, R. Delbourgo. Piecewise rational quadratic interpolation to monotonic data. IMA Journal of Numerical Analysis, 2(2):123–130, 1982.
• D. J. Rezende, S. Mohamed. Variational inference with normalizing flows. ICML, 2015. §4.1 (планарные и радиальные потоки), приложение A (условия обратимости).
• E. Nalisnick, A. Matsukawa, Y. W. Teh, D. Gorur, B. Lakshminarayanan. Do deep generative models know what they don’t know? ICLR, 2019. §3.
• C.-W. Huang, D. Krueger, A. Lacoste, A. Courville. Neural autoregressive flows. ICML, 2018. §4, теорема 1 (универсальная аппроксимация).
• T. Teshima, I. Ishikawa, K. Tojo, K. Oono, M. Ikeda, M. Sugiyama. Coupling-based invertible neural networks are universal diffeomorphism approximators. NeurIPS, 2020. Теоремы 2 и 3 (универсальность потоков на аффинных слоях связи).
• F. Koehler, V. Mehta, A. Risteski. Representational aspects of depth and conditioning in normalizing flows. ICML, 2021.
• F. Draxler, S. Wahl, C. Schnörr, U. Köthe. On the universality of volume-preserving and coupling-based normalizing flows. ICML, 2024. §5.
• G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, B. Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. JMLR, 22(57), 2021. §3.2, §3.3.2, §3.4.
• HSE-LAMBDA, репозиторий курса Deep Generative Models: nf.ipynb, 03_hw_vapnev.ipynb.