GAN Вассерштейна и ограничение Липшица

Заметка о GAN закончилась проблемой. При оптимальном дискриминаторе генератор минимизирует $2\,D_{\mathrm{JS}}(P\,\|\,Q) - 2\ln 2$, а когда носители (supports) настоящего распределения $P$ и распределения генератора $Q$ не пересекаются, $D_{\mathrm{JS}} = \ln 2$, где бы ни находилось $Q$. Тогда (почти) оптимальный дискриминатор постоянен вокруг сгенерированных сэмплов, поэтому минимаксная функция потерь генератора не получает градиента, а ненасыщающаяся (non-saturating) даёт неустойчивые обновления (Arjovsky and Bottou, 2017, теоремы 2.1–2.2, 2.4 и 2.6; см. заметку о GAN). GAN Вассерштейна (Wasserstein GAN; Arjovsky, Chintala and Bottou, 2017) заменяет дивергенцию Йенсена–Шеннона (Jensen–Shannon divergence) расстоянием Вассерштейна-1 (Wasserstein-1 distance), которое продолжает уменьшаться по мере того, как $Q$ движется к $P$, даже когда они не перекрываются. Цена — ограничение на дискриминатор: он должен быть 1-липшицевым. Дальше в заметке разобраны функции потерь и два стандартных способа выполнить это ограничение.

Обозначения. $P$ — распределение настоящих данных, $Q$ — распределение генератора ($x = G(z)$ с шумом $z \sim p_z$, как в заметке о GAN), $f$ — критик (critic): так в WGAN называется дискриминатор, сеть с вещественным выходом и без сигмоиды. Его функция потерь сохраняет обозначение функции потерь дискриминатора, $\mathcal L_D$. Настоящие сэмплы — $x_r \sim P$, подделки — $x_f = G(z)$. $\|\cdot\|_2$ — евклидова норма.

Расстояние Вассерштейна-1#

Представим $P$ и $Q$ как две кучи земли одинаковой массы. Транспортный план (transport plan) $\gamma(x, y)$ показывает, сколько массы перемещается из $x$ в $y$; его маргинальные распределения должны совпадать с $P$ и $Q$, а $\Pi(P, Q)$ — множество всех таких планов (Peyré and Cuturi, 2019, глава 2). Перемещение массы $m$ на расстояние $\|x - y\|$ стоит $m\,\|x - y\|$, и $W_1$, которое также называют расстоянием перемещения земли (earth mover’s distance), — наименьшая суммарная стоимость:

$$W_1(P, Q) = \inf_{\gamma \in \Pi(P, Q)} \mathbb{E}_{(x, y)\sim\gamma}\,\|x - y\| .$$

Сеть не может оптимизировать напрямую по всем транспортным планам. Для распределений на $\mathbb R^d$ с конечными первыми моментами, $\mathbb E\|x\| < \infty$ для обоих, $W_1$ конечно, и двойственность Канторовича–Рубинштейна (Kantorovich–Rubinstein duality; Villani, 2009, замечание 6.5) переписывает инфимум как супремум по функциям:

$$\begin{aligned} W_1(P, Q) = \sup_{\|f\|_L \le 1}\; \Big[\, &\mathbb{E}_{x\sim P}\,f(x)\\ &- \mathbb{E}_{x\sim Q}\,f(x) \Big], \end{aligned}$$

где $\|f\|_L \le 1$ ($f$ является 1-липшицевой) означает $|f(x) - f(y)| \le \|x - y\|$ для всех $x, y$. Отсюда и берутся критик и его ограничение: критик $f$ играет роль функции под знаком супремума. Он ставит настоящим сэмплам высокие оценки, а подделкам — низкие, но не может меняться быстрее расстояния. Без ограничения и при $P \ne Q$ супремум был бы равен $+\infty$, поскольку любую разделяющую $f$ можно было бы умножить на сколь угодно большое число.

Почему это помогает. Возьмём две точечные массы: $P = \delta_0$ (данные) и $Q_\theta = \delta_\theta$ (модель). При любом $\theta \ne 0$ носители не пересекаются. Смесь $M = \tfrac12(P + Q_\theta)$ помещает массу $\tfrac12$ в каждый атом, поэтому каждое слагаемое KL в $D_{\mathrm{JS}}$ сравнивает массу $1$ с массой $\tfrac12$ в одном атоме:

$$\begin{aligned} D_{\mathrm{JS}}(P\,\|\,Q_\theta) &= \tfrac12 D_{\mathrm{KL}}(P\,\|\,M)\\ &\quad + \tfrac12 D_{\mathrm{KL}}(Q_\theta\,\|\,M)\\ &= \tfrac12\ln\frac{1}{1/2} + \tfrac12\ln\frac{1}{1/2}\\ &= \ln 2 \approx 0.693,\\ D_{\mathrm{KL}}(P\,\|\,Q_\theta) &= \infty,\\ W_1(P, Q_\theta) &= |\theta| . \end{aligned}$$

Единственный план перемещает всю массу на расстояние $|\theta|$. $D_{\mathrm{JS}}$ не зависит от $\theta$, поэтому её производная равна $0$; у $W_1$ наклон $\pm 1$, и это расстояние всегда указывает модели направление к данным. Это вариант примера с параллельными прямыми из статьи о WGAN, записанный для точечных масс (Arjovsky, Chintala and Bottou, 2017, пример 1). В общем случае обозначим через $Q_\theta$ распределение $G_\theta(z)$. Если данные лежат в компактном множестве, как везде предполагает статья, и $G_\theta$ непрерывна по $\theta$, то непрерывно и $W_1(P, Q_\theta)$. На неограниченном пространстве одной непрерывности $G_\theta$ мало, нужно ещё условие интегрируемости. Если $G_\theta$ — сеть прямого распространения (feedforward network), а априорное распределение удовлетворяет условию $\mathbb E\|z\| < \infty$, которое и даёт такое условие, то $W_1(P, Q_\theta)$ непрерывно и ещё и дифференцируемо по $\theta$ почти всюду. Ни одно из этих утверждений не верно для $D_{\mathrm{JS}}$ или KL, которые могут меняться скачком (Arjovsky, Chintala and Bottou, 2017, теорема 1 и следствие 1).

Рисунок

D_JS постоянна при непересекающихся носителях; у W₁ остаётся наклон

Функции потерь WGAN#

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

$$\begin{aligned} \mathcal L_D &= -\mathbb{E}_{x\sim P}\big[f(x)\big] + \mathbb{E}_{z\sim p_z}\big[f(G(z))\big],\\ \mathcal L_G &= -\mathbb{E}_{z\sim p_z}\big[f(G(z))\big]. \end{aligned}$$

Что измеряет $-\mathcal L_D$? Семейство сетей содержит лишь часть $K$-липшицевых функций, поэтому даже его лучший представитель даёт $-\mathcal L_D \le K\,W_1(P, Q)$ — нижнюю оценку, которая становится точной, когда семейство содержит оптимального 1-липшицева критика, умноженного на $K$ ($K = 1$ для 1-липшицева семейства). При обучении $-\mathcal L_D$ — оценка этой величины по мини-батчу. В отличие от дискриминатора GAN, критика следует обучать почти до оптимума: чем лучше критик, тем надёжнее градиент $W_1$, тогда как более сильный дискриминатор GAN насыщается. Поэтому критика обновляют несколько раз на каждый шаг генератора, $n_{\text{critic}} = 5$ в исходном алгоритме (Arjovsky, Chintala and Bottou, 2017, раздел 3 и алгоритм 1).

Если на критика наложено ограничение Липшица и он обучен почти до оптимума, $-\mathcal L_D$ в ходе одного запуска коррелирует с качеством сэмплов, чего нельзя сказать о функции потерь дискриминатора GAN (раздел 4.2, рисунки 3–4). Константа $K$ зависит от архитектуры критика, поэтому это число трудно сравнивать между моделями с разными критиками, и авторы не предлагают его в качестве метрики качества.

Сами по себе функции потерь не делают $f$ липшицевой. Исходный WGAN после каждого шага обрезал каждый вес критика до отрезка $[-c, c]$; такая обрезка весов (weight clipping) делает критика $K$-липшицевым для некоторого $K$, которое зависит от $c$ и архитектуры. Авторы назвали это «откровенно ужасным способом» (a clearly terrible way) обеспечить ограничение: при большом $c$ критик медленно обучается до оптимума, а при малом $c$ градиенты могут затухать (Arjovsky, Chintala and Bottou, 2017, раздел 3). Статья о WGAN-GP показала, кроме того, что обрезка смещает критика в сторону слишком простых функций и что без тщательного подбора $c$ градиенты взрываются или затухают (Gulrajani et al., 2017, разделы 3.1–3.2). Два метода ниже заменяют обрезку.

WGAN-GP: штраф на градиент#

WGAN-GP (Gulrajani et al., 2017) сохраняет функции потерь WGAN и добавляет к функции потерь критика штраф на градиент (gradient penalty):

$$\begin{aligned} \mathcal L_D = {}& \mathbb{E}_{z\sim p_z}\big[f(G(z))\big] - \mathbb{E}_{x\sim P}\big[f(x)\big]\\ &+ \lambda\,\mathbb{E}_{\hat x}\Big[\big(\|\nabla_{\hat x} f(\hat x)\|_2 - 1\big)^2\Big]. \end{aligned}$$

Первые два слагаемых — функция потерь критика WGAN, записанная в обратном порядке, а функция потерь генератора $\mathcal L_G$ не меняется. Символ $\mathcal L_D$ теперь включает штраф, поэтому $-\mathcal L_D$ уже не является оценкой нижней границы $K\,W_1$ из предыдущего раздела: это та оценка минус $\lambda$, умноженное на штраф. Для наблюдения за обучением стоит записывать разность оценок на настоящих и поддельных данных $\mathbb E_P f - \mathbb E_Q f$ и штраф по отдельности. Разность читается с прежними оговорками и ещё одной: мягкое ограничение делает критика лишь приближённо 1-липшицевым, так что разность не обязана быть нижней оценкой $W_1$.

Почему норма градиента. Для дифференцируемой $f$ на $\mathbb R^d$ быть 1-липшицевой — то же самое, что иметь $\|\nabla f(x)\|_2 \le 1$ всюду (одна импликация следует из определения производной, другая — из теоремы о среднем на отрезке). Поэтому, подталкивая $\|\nabla_{\hat x} f(\hat x)\|_2$ к $1$, мы сохраняем $f$ приблизительно 1-липшицевой. Буквальному ограничению соответствовал бы односторонний штраф $\max(0, \|\nabla f\|_2 - 1)^2$. WGAN-GP использует двусторонний штраф $(\|\nabla f\|_2 - 1)^2$, который не исключает оптимального критика. На компактном пространстве, если оптимальный критик $f^*$ дифференцируем, а оптимальный транспортный план (optimal coupling) перемещает каждую точку ($\pi(x = y) = 0$), то норма градиента $f^*$ равна $1$ на отрезках между точками, которые этот план связывает друг с другом, а также почти всюду по мере $P$ и $Q$ (Gulrajani et al., 2017, предложение 1 и следствие 1). В экспериментах авторов двусторонний штраф к тому же работал немного лучше (раздел 4 и приложение C). Ограничение мягкое: оно проверяется только в сэмплированных точках $\hat x$.

Откуда берётся $\hat x$. Сопоставим каждому настоящему сэмплу подделку, сэмплируем $\alpha \sim U(0, 1)$ и возьмём случайную точку на отрезке между ними:

$$\hat x = \alpha\,x_r + (1 - \alpha)\,x_f .$$

Предложение 1 говорит о парах, взятых из оптимального транспортного плана, а он неизвестен. Вместо этого WGAN-GP составляет пары из настоящих сэмплов и подделок случайным образом. Это эвристика, мотивированная предложением: обеспечить единичную норму всюду вычислительно неразрешимо (intractable), а требовать её только вдоль этих отрезков на практике достаточно (Gulrajani et al., 2017, раздел 4). В коде $\alpha$ сэмплируется один раз на изображение и распространяется на все пиксели (broadcasting). Градиент выхода критика по $\hat x$ вычисляется с помощью torch.autograd.grad(..., create_graph=True), чтобы через сам штраф тоже можно было выполнить обратное распространение ошибки. Затем для каждого сэмпла градиент вытягивается в вектор $g$, и его норма входит в штраф ${(\|g\|_2 - 1)^2}$. Во всех экспериментах статьи $\lambda = 10$ (раздел 4).

Рисунок

Штраф на градиент вычисляется между настоящими сэмплами и подделками

Без батч-нормализации в критике. Штраф рассматривает критика как функцию одного изображения: он ограничивает $\|\nabla_{\hat x} f(\hat x)\|_2$ для каждого $\hat x$ по отдельности. С батч-нормализацией (batch normalization) критик отображает целый мини-батч входов в мини-батч выходов, и штраф на градиент по каждому входу уже не описывает ту функцию, которую ограничивают. Поэтому WGAN-GP исключает батч-нормализацию из критика и допускает только нормализации, которые не смешивают примеры; авторы рекомендуют нормализацию слоя (layer normalization; Gulrajani et al., 2017, раздел 4). Подходит и нормализация по экземпляру (instance normalization), которая нормализует каждый канал каждого изображения по его пространственным позициям.

SN-GAN: спектральная нормализация#

Второй подход ограничивает не градиенты, а веса (Miyato et al., 2018). Авторы обучали свои SN-GAN со стандартной функцией потерь GAN и с кусочно-линейной функцией потерь (hinge loss), а не с функцией потерь Вассерштейна; WGAN-GP служил им базой для сравнения (раздел 4, формулы 15–17). Сама нормализация — это ограничение Липшица, и она так же применима к критику WGAN. Матрица весов каждого слоя делится на свою спектральную норму (spectral norm) — наибольшее сингулярное число:

$$\begin{gathered} \hat W = \frac{W}{\sigma_{\max}(W)},\\ \sigma_{\max}(W) = \max_{h \ne 0}\frac{\|W h\|_2}{\|h\|_2} . \end{gathered}$$

Что такое сингулярное число. Линейное отображение растягивает пространство. Оно переводит единичную окружность в эллипс, а сингулярные числа говорят, насколько сильно отображение растягивает пространство в каждом направлении: правые сингулярные векторы $v_i$ переходят в $\sigma_i u_i$ — полуоси эллипса. Наибольшее из них, $\sigma_{\max}$, — наибольший коэффициент, с которым может растянуться какой-либо вектор, поэтому оно равно константе Липшица отображения $h \mapsto W h$. После деления на него ни одно направление не растягивается больше чем в $1$ раз.

Рисунок

Матрица растягивает единичную окружность в эллипс; её наибольшее сингулярное число — самая длинная полуось

Для сети, в которой линейные слои чередуются с 1-липшицевыми активациями (ReLU, LeakyReLU), константы Липшица при композиции перемножаются, поэтому (Miyato et al., 2018, формула 7)

$$\|f\|_L \le \prod_{l} \sigma_{\max}(W_l),$$

а после нормализации это произведение не превосходит $1$.

В отличие от штрафа на градиент, ограничение накладывается на веса, а не проверяется в сэмплированных точках, и настраивать $\lambda$ не нужно. В PyTorch это обёртка torch.nn.utils.spectral_norm(nn.Conv2d(...)) (в новом API — torch.nn.utils.parametrizations.spectral_norm), которая вместо полного SVD оценивает $\sigma_{\max}$ одним шагом степенного метода (power iteration) на каждый прямой проход в режиме обучения. На практике эта граница выполняется лишь приближённо, и тому есть три причины. Степенной метод даёт лишь приближённое значение $\sigma_{\max}$. Для свёртки обёртка нормализует ядро, развёрнутое в матрицу, а это лишь суррогат операторной нормы свёртки, не равный ей; истинная норма зависит ещё и от шага (stride) и дополнения (padding) (Gouk et al., 2021, раздел 3.2; Miyato et al., 2018, сноска 3). Наконец, слои нормализации, стоящие между свёртками, в общем случае не 1-липшицевы. Поэтому спектральная нормализация держит константу Липшица критика под контролем, но не гарантирует, что она не превосходит $1$. Кроме того, граница через произведение не обязана быть точной, так что нормализованный критик $K$-липшицев для некоторого $K$, которое может быть заметно меньше $1$; тогда критик WGAN со спектральной нормализацией оценивает нижнюю границу для $K\,W_1$, а не само $W_1$.

Два двойственных взгляда#

И f-GAN из заметки о GAN, и WGAN обучают дискриминатор как функцию под знаком супремума. f-GAN использует оценку f-дивергенции (f-divergence) через сопряжённую по Фенхелю функцию (convex conjugate); f-дивергенция сравнивает отношения плотностей и поэтому наследует проблему носителей, присущую $D_{\mathrm{JS}}$ и KL. WGAN использует двойственную форму Канторовича–Рубинштейна для транспортной стоимости, которая сравнивает положения и остаётся информативной, когда носители не пересекаются. Ограничение Липшица — цена второго взгляда: именно оно связывает критика с геометрией пространства данных.

Ноутбук#

Ноутбук о WGAN из репозитория HSE-LAMBDA Deep Generative Models (закреплённый коммит; ячейки нумеруются с 1, считая все ячейки подряд, включая текстовые) реализует все три варианта на CelebA с разрешением $64 \times 64$. Все три функции обучения используют Adam с шагом обучения (learning rate) $2\cdot10^{-4}$ и betas $(0.5, 0.9)$ и $n_{\text{critic}} = 5$ шагов критика на шаг генератора; в статье о WGAN-GP шаг обучения равен $10^{-4}$, а betas — $(0, 0.9)$ (Gulrajani et al., 2017, алгоритм 1). Базовый запуск и запуск WGAN-GP рассчитаны на 2 эпохи, запуск со спектральной нормализацией — на 20.

• Ячейки 1–5: введение, расстояние перемещения земли и его двойственная форма Канторовича–Рубинштейна, а также целевая функция WGAN с критиком, выдающим вещественное число.
• Ячейки 6–9: данные. CelebA: 202 599 выровненных и обрезанных изображений лиц размером 178 × 218 пикселей (ширина × высота), принадлежащих 10 177 людям; каждое размечено 40 бинарными атрибутами, например «улыбается» или «в очках». Класс набора данных из курса возвращает атрибуты, а этот ноутбук их не использует. Изображения приводятся к размеру $64 \times 64$ и масштабируются в $[-1, 1]$, размер батча 1024 с drop_last=True (197 батчей за эпоху), а ячейка 9 показывает сетку настоящих изображений. Ячейка 7 выбирает устройство (CUDA, Apple MPS или CPU). Resize((64, 64)) с парой размеров масштабирует две стороны независимо, поэтому соотношение сторон меняется с 178 : 218 на 1 : 1, и лица становятся примерно на 22% шире относительно своей высоты; Resize(64) с последующим CenterCrop(64) сохраняет пропорции.
• Ячейки 10–16: базовый WGAN. Ячейка 12 определяет сети. Генератор следует рекомендациям DCGAN: $z \in \mathbb R^{100}$ проходит через линейный слой без смещения с $512 \cdot 4 \cdot 4 = 8192$ выходами, BatchNorm1d и ReLU, переформируется в $512 \times 4 \times 4$, и четыре транспонированные свёртки (transposed convolutions) с ядром 5, шагом 2, дополнением 2 и выходным дополнением (output padding), равным 1, проводят его через 256, 128 и 64 канала к $3 \times 64 \times 64$, с батч-нормализацией и ReLU после первых трёх и Tanh после последней. Ядро 5 не делится на шаг 2, поэтому эти слои могут давать шахматные артефакты (checkerboard artifacts). Критик применяет четыре свёртки с шагом 2 и ядром 5, $3 \to 64 \to 128 \to 256 \to 512$ каналов, после каждой — LeakyReLU(0.2), с InstanceNorm2d(affine=True) после последних трёх и без нормализации на входе, а затем свёртку $4 \times 4$, дающую одно значение без сигмоиды. train_wgan (ячейка 14) использует d_loss = -mean(real_score) + mean(fake_score) и g_loss = -mean(fake_score). Базовый вариант не накладывает никакого ограничения Липшица — ни обрезки, ни штрафа (ячейка 11). Поэтому целевая функция его критика не ограничена (супремум выше равен $+\infty$), его выходы и D_loss могут неограниченно дрейфовать, а функция потерь не оценивает $W_1$. Ячейка 13 отмечает, что обрезку весов можно добавить, обрезая веса критика до отрезка (clamp) после каждого обновления. Каждая функция обучения один раз сэмплирует 16 латентных векторов (fixed_noise, ячейки 14, 19 и 24) и после каждой эпохи декодирует их генератором в режиме инференса (evaluation mode), так что сетки прослеживают одни и те же коды на протяжении обучения.
• Ячейки 17–21: WGAN-GP. gradient_penalty (ячейка 18) сэмплирует $\alpha \sim U(0,1)$ отдельно для каждого элемента батча, интерполирует, дифференцирует с create_graph=True и штрафует норму $L_2$ каждого элемента; lambda_gp = 10. Критик — та же сеть, а её InstanceNorm2d нормализует каждый сэмпл отдельно, поэтому штраф по отдельным сэмплам остаётся корректным (см. «Без батч-нормализации в критике»). В train_wgan_gp (ячейка 19) штраф получает fake_imgs без detach(), поэтому каждый шаг критика выполняет обратное распространение и через генератор. Эти градиенты обнуляются перед шагом генератора, так что на результаты это не влияет, но вычисления тратятся впустую.
• Ячейки 22–27: спектральная нормализация. DiscriminatorSN оборачивает каждую свёртку в spectral_norm; во внутренних блоках за нормализованной свёрткой следуют InstanceNorm2d(affine=True) и LeakyReLU, поэтому неравенство выше выполняется лишь приближённо. Обучение использует обычные функции потерь WGAN.
• Ячейка 28: итоги.

Сохранённые выходы есть только у ячеек 7, 9, 26 и 27: устройство (mps), сетка настоящих изображений, индикаторы прогресса запуска со спектральной нормализацией, прерванного на 11-й эпохе из 20, и сэмплы из этого частично обученного генератора. Значения функций потерь не сохранены. У ячеек обучения и сэмплирования базового варианта и WGAN-GP нет сохранённых выходов, поэтому для этих вариантов ноутбук документирует только процедуры.

Дальше#

WGAN дал генератору полезный градиент, а спектральная нормализация удержала критика под контролем. Но встреченные до сих пор GAN всё ещё порождают маленькие изображения одного вида: лица 64×64, цифры, спальни. Следующая заметка прослеживает, что понадобилось, чтобы дойти до больших, разнообразных изображений высокого разрешения: постепенное наращивание сетей, self-attention для структуры, охватывающей всё изображение, и BigGAN, который масштабирует тот же рецепт на весь ImageNet.

Литература#

• M. Arjovsky, S. Chintala, L. Bottou. Wasserstein GAN. arXiv:1701.07875, 2017; опубликована на ICML 2017 под названием «Wasserstein generative adversarial networks». Номера разделов относятся к версии arXiv.
• M. Arjovsky, L. Bottou. Towards principled methods for training generative adversarial networks. ICLR, 2017.
• I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, A. Courville. Improved training of Wasserstein GANs. NeurIPS, 2017.
• T. Miyato, T. Kataoka, M. Koyama, Y. Yoshida. Spectral normalization for generative adversarial networks. ICLR, 2018.
• C. Villani. Optimal Transport: Old and New. Springer, 2009. Главы 5–6: двойственность Канторовича и расстояния Вассерштейна.
• G. Peyré, M. Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning, 11(5–6), 2019.
• H. Gouk, E. Frank, B. Pfahringer, M. J. Cree. Regularisation of neural networks by enforcing Lipschitz continuity. Machine Learning, 110, 2021.
• HSE-LAMBDA, репозиторий курса Deep Generative Models: WGAN.ipynb.