Автоэнкодирующие GAN: состязательные автоэнкодеры и α-GAN

Заметка о StyleGAN закончилась управлением генерацией, но без пути обратно: чтобы найти код настоящего изображения, StyleGAN приходится решать задачу оптимизации. WGAN и рецепты масштабирования оставили на месте две более старые проблемы. Два семейства моделей, с которыми мы уже познакомились, решают каждое лишь половину задачи. У VAE есть энкодер, так что у каждого изображения есть код, а слагаемое реконструкции заставляет декодер учитывать каждое обучающее изображение; но его сэмплы размыты. У GAN сэмплы чёткие, но пути назад от изображения к его коду нет, и ничто не мешает генератору пропускать моды данных. Эта заметка объединяет оба подхода. Состязательный автоэнкодер (adversarial autoencoder) заменяет слагаемое KL в VAE дискриминатором на кодах. α-GAN заменяет ещё и пиксельное правдоподобие дискриминатором на изображениях, сохраняя слагаемое реконструкции $L_1$, чтобы энкодер по-прежнему защищал от коллапса мод (mode collapse). α-WGAN обучает оба дискриминатора как критиков WGAN-GP.

Обозначения. $x \in \mathbb R^D$ — изображение, $P$ — распределение данных, как в заметке о WGAN. $z \in \mathbb R^d$ — код с априорным распределением $p(z) = \mathcal N(0, I)$. Энкодер $E$ — детерминированная сеть, $\hat z = E(x)$ — код изображения $x$, генератор (декодер) $G$ отображает коды в изображения, $\hat x = G(E(x))$ — реконструкция $x$, а $G(z)$ при $z \sim p(z)$ — сэмпл. Агрегированное апостериорное распределение (aggregate posterior) $q(z)$ — это распределение $E(x)$ при $x \sim P$. $D: \mathbb R^D \to \mathbb R$ — дискриминатор изображений, $C: \mathbb R^d \to \mathbb R$ — дискриминатор кодов. Обе сети выдают вещественное число — своё значение (score). Это логит, когда их обучают как классификаторы, так что $\sigma(D(x))$ — вероятность того, что $x$ — настоящее изображение ($\sigma$ — логистическая сигмоида; в заметке о GAN $D(x)$ обозначало саму эту вероятность), и значение критика, когда их обучают с функцией потерь Вассерштейна. Для $C$ «настоящий» означает «взятый из априорного распределения». $\lVert v \rVert_1 = \sum_i \lvert v_i \rvert$; реализации обычно не суммируют по пикселям, а усредняют, что лишь меняет масштаб веса $\lambda$ слагаемого реконструкции.

Чего не хватает каждому семейству#

VAE обучает энкодер и декодер на ELBO — слагаемом реконструкции минус слагаемое KL (заметка о VAE):

$$\mathcal L(x) = \mathbb E_{q(z \mid x)}\big[\log p(x \mid z)\big] - D_{\mathrm{KL}}\big(q(z \mid x)\,\|\,p(z)\big) .$$

У каждого слагаемого есть цена и польза.

• Слагаемое правдоподобия с гауссовским декодером — это квадратичная ошибка. Если одному коду соответствует несколько правдоподобных изображений, её оптимум — их среднее, а оно размыто (заметка о VAE). Польза этого слагаемого — давление в сторону покрытия: каждое обучающее изображение должно восстанавливаться из своего кода, и это мешает модели игнорировать части данных. Это не гарантия: конечная обученная модель всё равно может плохо покрыть какие-то области данных, а степень размытости зависит от правдоподобия декодера и от того, насколько изображение неопределённо при данном коде.
• Слагаемое KL взимает с каждого изображения плату за информативный код. В среднем по данным оно распадается на взаимную информацию (mutual information) между изображениями и кодами, которая для хорошей реконструкции должна быть большой, и KL-дивергенцию агрегированного апостериорного распределения относительно априорного, которая для сэмплирования должна быть малой (заметка о VAE). Для генерации полезна только вторая часть.

GAN обходится без обоих слагаемых. Его дискриминатор смотрит на изображения целиком, поэтому размытость легко заметить, и сэмплы получаются чёткими. Но у GAN нет энкодера: он не может реконструировать изображение, найти его код или интерполировать между двумя настоящими изображениями. И у GAN нет ограничения на реконструкцию каждого примера: ничто не привязывает генератор к каждому обучающему изображению. Его идеальная целевая функция на уровне распределений, JS при оптимальном дискриминаторе, штрафует за пропущенную массу, но чередующаяся игра этот идеал не обеспечивает, поэтому генератор может свестись к нескольким модам. С ненасыщающейся функцией потерь (non-saturating loss) и оптимальным дискриминатором ожидаемый градиент генератора содержит обратную KL (reverse KL), которая сильно наказывает сэмплы, похожие на подделку, и лишь слабо — пропущенные моды (заметка о GAN). Функция потерь Вассерштейна из предыдущей заметки меняет градиент, но не добавляет энкодер, и генератор по-прежнему учится только через значения критика на собственных сэмплах, так что ничто не привязывает его к каждому обучающему изображению.

Модели ниже сохраняют энкодер VAE и его слагаемое реконструкции, а проблемные части ELBO передают дискриминаторам: сначала слагаемое KL, затем правдоподобие.

Состязательные автоэнкодеры: дискриминатор на кодах#

Состязательный автоэнкодер (AAE) Makhzani et al. (2015) — это обычный автоэнкодер, к коду которого присоединена третья сеть. Распределение кодов, которое он формирует, — агрегированное апостериорное распределение (формула 1),

$$q(z) = \int q(z \mid x)\, P(x)\, dx ,$$

которое для детерминированного энкодера — просто распределение $E(x)$. На каждом мини-батче обучение чередует две фазы (раздел 2).

1. Фаза реконструкции. $E$ и $G$ минимизируют ошибку реконструкции, в экспериментах статьи — квадратичную ошибку: $\mathbb E_{x \sim P}\lVert x - G(E(x))\rVert_2^2$.

2. Фаза регуляризации. Дискриминатор кодов учится помечать сэмплы из априорного распределения как настоящие, а коды — как подделки,

   $$\begin{aligned} \mathcal L_C ={}& -\mathbb E_{z \sim p}\log \sigma\big(C(z)\big)\\ &- \mathbb E_{x \sim P}\log\Big(1 - \sigma\big(C(E(x))\big)\Big), \end{aligned}$$

   а затем энкодер, который играет роль генератора в этом маленьком GAN, обновляется так, чтобы его обмануть, например с ненасыщающейся функцией потерь $-\mathbb E_{x}\log\sigma\big(C(E(x))\big)$.

Что сопоставляет дискриминатор кодов#

$C$ видит по одному коду за раз, вперемешку с кодами других изображений, и никогда не узнаёт, из какого изображения получен код. Поэтому он сравнивает с априорным распределением всё облако кодов: $q(z)$ с $p(z)$, а не код каждого изображения с априорным распределением. Его оптимум следует из рассуждения об оптимальном дискриминаторе из заметки о GAN, применённого к кодам: $\sigma(C^*(z)) = p(z) / \big(p(z) + q(z)\big)$, а в терминах значения это логарифм отношения плотностей:

$$C^*(z) = \log \frac{p(z)}{q(z)} .$$

С минимаксной функцией потерь и оптимальным $C$ энкодер минимизирует $2\,D_{\mathrm{JS}}(p\,\|\,q) - 2\ln 2$ — по той же подстановке, что и в заметке о GAN. По сравнению с VAE, AAE сохраняет ту часть слагаемого KL, которая нужна для сэмплирования, — расхождение между $q(z)$ и $p(z)$, измеряемое здесь через JS, а не через KL, — и отбрасывает часть, которая штрафует информативные коды. Makhzani et al. проводят то же сравнение: они записывают усреднённое слагаемое KL в VAE как кросс-энтропию между $q(z)$ и $p(z)$ минус среднюю энтропию апостериорных распределений $q(z \mid x)$ и заменяют оба слагаемых состязательной игрой (раздел 2.1, формула 2).

Отсюда следуют три вещи.

• Подходит любое априорное распределение, из которого можно сэмплировать. Замкнутой формуле для KL в VAE нужна плотность априорного распределения; дискриминатору достаточно сэмплов из него. Статья накладывает на коды MNIST смесь десяти двумерных гауссиан и «швейцарский рулет» (Swiss roll; разделы 2.1 и 2.3).
• Меньше дыр. С двумерным гауссовским априорным распределением на MNIST коды отложенных (held-out) изображений заполняют априорное распределение без пробелов, тогда как VAE с той же архитектурой оставляет области, куда не кодируется ни одно изображение (рисунок 2). Дыры в $q(z)$ — это места, где сэмплы из априорного распределения плохо декодируются.
• Детерминированного энкодера достаточно — в определённых пределах. Авторы пробовали детерминированный энкодер, гауссовский и энкодер с шумом на входе, получили близкие значения правдоподобия на тестовых данных и поэтому приводят результаты для детерминированного (раздел 2). Его единственный источник случайности — данные, и для конечной обучающей выборки $q(z)$ — облако точек. Поэтому авторы рекомендуют размерность кода, близкую к внутренней размерности (intrinsic dimension) данных (от 5 до 8 для MNIST, от 10 до 20 для лиц и номеров домов), а для гораздо больших кодов — гауссовский энкодер (приложение A.1).

AAE поддерживает также классификацию с частичной разметкой (semi-supervised) и кластеризацию: для этого к категориальной части кода присоединяется второй дискриминатор (разделы 5–6). Здесь мы рассматриваем только его генеративное применение.

Почему дискриминатор, а не штраф на моменты? Более дешёвый способ подтянуть коды к $\mathcal N(0, I)$ — штрафовать средние и дисперсии каждой координаты по батчу. Rosca et al. (2017) попробовали именно это вместо дискриминатора кодов (приложение D). Штраф приблизил средние координат к нулю, но координаты остались сильно коррелированными, поскольку покоординатный штраф не видит корреляций; дискриминатор, который видит векторы кодов целиком, давал декоррелированные коды.

Чего AAE не меняет, так это обучающего сигнала декодера. $G$ по-прежнему учится только на пиксельной ошибке реконструкции, поэтому его выходы страдают тем же усреднением, что и у декодера VAE.

α-GAN: дискриминатор ещё и на изображениях#

Rosca et al. (2017) снова начинают с ELBO и заменяют каждое его слагаемое тем, что может оценить дискриминатор (раздел 3).

• Слагаемое KL превращается в дискриминатор кодов, как в AAE: $-D_{\mathrm{KL}}(q\,\|\,p) = \mathbb E_q \log\big(p(z)/q(z)\big) \approx \mathbb E_q\, C(z)$ — математическое ожидание значения классификатора, оптимум которого равен логарифму отношения (формула 6). Поскольку $C$ сравнивает агрегированное $q(z)$ с $p(z)$, оценивается маргинальная KL (marginal KL) — та часть, которую мы выделили выше.
• Для слагаемого правдоподобия есть два кандидата (формула 7). Первый — явное правдоподобие Лапласа, $p(x \mid z) \propto \exp\big(-\lambda\lVert x - G(z)\rVert_1\big)$, отрицательный логарифм которого — функция потерь реконструкции $L_1$, $\lambda\lVert x - G(z)\rVert_1$, плюс константа. Второй — синтетическое правдоподобие (synthetic likelihood): значение $D(G(z))$ дискриминатора изображений, обученного отличать данные от декодированных изображений; оно велико, когда декодированное изображение выглядит настоящим.

Два кандидата ошибаются противоположным образом. Дискриминатор даёт чёткие изображения, но, как в любом GAN, допускает коллапс мод. Слагаемое $L_1$ отводит массу каждому обучающему изображению, но размывает: для каждого пикселя его оптимум — медиана правдоподобных значений, тогда как у квадратичной ошибки — их среднее. α-GAN использует оба кандидата с весами (формула 8). В записи через значения и с детерминированным энкодером $E$ и $G$ максимизируют

$$\mathbb E_{x \sim P}\Big[-\lambda\big\lVert x - G(E(x))\big\rVert_1 + D\big(G(E(x))\big) + C\big(E(x)\big)\Big] .$$

Метод дополняют два практических изменения (раздел 3, «Improved techniques»).

1. Состязательная функция потерь $E$ и $G$ — значение со знаком минус. Обозначим значение через $s$, а вероятность, которую выдаёт классификатор, через $\sigma(s)$. Поскольку $\sigma(s)/(1 - \sigma(s)) = e^s$,

   $$-\log\sigma(s) + \log\big(1 - \sigma(s)\big) = -s .$$

   Левая часть — функция потерь, которую Rosca et al. используют для обоих генераторов своей игры (алгоритм 1, $R_D$ и $R_C$), а правая показывает, что это просто логит со знаком минус. Её производная по логиту равна $-1$ при любом $s$, поэтому она никогда не насыщается: при $\sigma(s) = 0.05$ минимаксная функция потерь даёт производную $-0.05$ по логиту, а ненасыщающаяся — $-0.95$ (заметка о GAN). При оптимальном классификаторе логит равен $\log\big(P(x)/Q(x)\big)$, где $Q$ — распределение подделок, так что математическое ожидание функции потерь на подделках равно $\mathbb E_Q\big[\log (Q/P)\big] = D_{\mathrm{KL}}(Q\,\|\,P)$ — обратной KL.

2. Сэмплы — тоже подделки. Дискриминатор изображений получает в качестве подделок не только реконструкции, но и сэмплы $G(z)$, $z \sim p(z)$; с обоими видами справляется один дискриминатор. Rosca et al. обосновывают это неравенством Йенсена, $\log p(x) = \log \int p(x \mid z)\,p(z)\,dz \ge \mathbb E_{p(z)}\log p(x \mid z)$, правая часть которого содержит декодированные сэмплы из априорного распределения, и, как и прежде, заменяют это правдоподобие синтетическим.

Четыре функции потерь#

Каждая сеть минимизирует свою функцию потерь при фиксированных остальных (алгоритм 1, формулы 11, 13, 15 и 17, здесь записанные через значения):

$$\begin{aligned} \mathcal L_E &= \mathbb E_{x}\big[\lambda\lVert x - \hat x\rVert_1 - C(\hat z)\big],\\ \mathcal L_G &= \mathbb E_{x}\big[\lambda\lVert x - \hat x\rVert_1 - D(\hat x)\big] - \mathbb E_{z}\big[D(G(z))\big],\\ \mathcal L_D &= -\mathbb E_{x}\Big[\log\sigma\big(D(x)\big) + \log\big(1 - \sigma(D(\hat x))\big)\Big]\\ &\quad - \mathbb E_{z}\Big[\log\big(1 - \sigma(D(G(z)))\big)\Big],\\ \mathcal L_C &= -\mathbb E_{x}\Big[\log\big(1 - \sigma(C(\hat z))\big)\Big] - \mathbb E_{z}\Big[\log\sigma\big(C(z)\big)\Big]. \end{aligned}$$

• Энкодер должен строить коды, из которых $G$ может восстановить изображение и которые дискриминатор кодов принимает за сэмплы из априорного распределения.
• Генератор должен восстанавливать изображение по его коду и делать так, чтобы и реконструкции, и сэмплы казались $D$ настоящими.
• Дискриминатор изображений отделяет данные от подделок обоих видов, а дискриминатор кодов — сэмплы из априорного распределения от кодов; оба используют обычную бинарную кросс-энтропию из заметки о GAN.

Что оценивает оценка дискриминатора изображений. $\mathcal L_D$ учитывает настоящие изображения один раз и каждый из двух видов подделок тоже один раз, так что подделки весят вдвое больше. Поточечный оптимум тогда равен $\sigma(D^*) = P/(P + Q_{\mathrm{rec}} + Q_{\mathrm{gen}})$, где $Q_{\mathrm{rec}}$ и $Q_{\mathrm{gen}}$ — распределения реконструкций и сэмплов. Через их смесь с равными весами $M = \tfrac12(Q_{\mathrm{rec}} + Q_{\mathrm{gen}})$ оптимальный логит равен $D^* = \log(P/M) - \log 2$, а не сбалансированному $\log(P/Q)$ из пункта 1 выше. Сдвиг $-\log 2$ постоянен, поэтому градиенты, которые получают $E$ и $G$, он не меняет, но сдвигает оценки и вероятности: когда подделки обоих видов совпадают с данными, $\sigma(D^*) = 1/3$, а не $1/2$. α-WGAN ниже уравновешивает стороны иначе — множителем 2 при слагаемом на настоящих данных.

В экспериментах $\lambda$ перебирали по значениям 1, 5, 10 и 50 отдельно для энкодера и для генератора, с весом 1 у обоих состязательных слагаемых (приложение H.1). Энкодер и генератор обновлялись дважды на одно обновление двух дискриминаторов (раздел 6). Все модели использовали Adam с $\beta = (0.5, 0.9)$ и батчами по 64, а дискриминатор кодов был трёхслойным MLP с 750 нейронами в слое (приложение H).

Зачем нужны все три части#

Допустим, оба условия согласования выполнены точно: коды распределены как априорное распределение, $q = p$, и каждое изображение восстанавливается, $G(E(x)) = x$ при $x \sim P$. Тогда для $z \sim p(z)$ величина $G(z)$ распределена так же, как $G(E(x))$ при $x \sim P$, поскольку $z$ и $E(x)$ одинаково распределены и к обоим применяется одна и та же функция $G$. Это распределение и есть $P$. Значит, сэмплы распределены как данные. Первое условие обеспечивает дискриминатор кодов, второе — слагаемое реконструкции, а поскольку оба выполняются лишь приближённо, дискриминатор изображений к тому же проверяет сэмплы напрямую.

То же рассуждение показывает, что предотвращает каждое слагаемое.

• Коллапс мод. Генератор, сведшийся к нескольким модам, не смог бы восстанавливать изображения из остальных мод, и слагаемое $L_1$ было бы большим. Rosca et al. приводят это как причину, по которой слагаемое реконструкции препятствует коллапсу мод (раздел 3, после формулы 9).
• Размытость. Реконструкции проверяет $D$, поэтому усреднённую реконструкцию легко отвергнуть, и $G$ подталкивается к чёткому изображению, близкому к $x$ по $L_1$, а не равному попиксельной медиане. Баланс задаёт вес $\lambda$.
• Дыры. Без дискриминатора кодов коды могли бы занимать любую область $\mathbb R^d$, и сэмплы из априорного распределения попадали бы туда, где $G$ ничему не научился, — это провал обычного автоэнкодера.

На ColorMNIST, CelebA и CIFAR-10 α-GAN был сопоставим с DCGAN (Radford, Metz and Chintala, 2016), WGAN-GP и AGE — моделью на основе автоэнкодера (Ulyanov, Vedaldi and Lempitsky, 2017), — но превосходил их не во всём. Независимый критик Вассерштейна, обученный на отложенных данных, поставил WGAN-GP на первое место на всех трёх наборах данных; на CIFAR-10 α-GAN получил лучший Inception Score, вычисленный сетью, обученной на ImageNet, при большом разбросе по гиперпараметрам (раздел 6, рисунки 2 и 5). Его реконструкции были лучше, чем у AGE (приложение C). Самым решающим фактором из всех авторы сочли архитектуры.

α-WGAN: два критика Вассерштейна#

Функции потерь $E$ и $G$ выше уже имеют вид значения со знаком минус, а это ровно вид функции потерь генератора WGAN $-\mathbb E f(G(z))$ (заметка о WGAN). От варианта Вассерштейна α-GAN отличается только обучением дискриминаторов. Kwon, Han and Kim (2019) сохраняют функции потерь энкодера и генератора и обучают и $D$, и $C$ как критиков Вассерштейна со штрафом на градиент (gradient penalty); модель они называют 3D-α-WGAN-GP. Их функции потерь (раздел 2, формулы 1–3; функция потерь энкодера приведена в тексте), где $\lambda_{\mathrm{GP}}$ — вес штрафа, имеют вид

$$\begin{aligned} \mathcal L_D ={}& \mathbb E_{x}\big[D(\hat x)\big] + \mathbb E_{z}\big[D(G(z))\big] - 2\,\mathbb E_{x}\big[D(x)\big]\\ &+ \lambda_{\mathrm{GP}}\,\mathrm{GP}_D,\\ \mathcal L_C ={}& \mathbb E_{x}\big[C(\hat z)\big] - \mathbb E_{z}\big[C(z)\big] + \lambda_{\mathrm{GP}}\,\mathrm{GP}_C,\\ \mathcal L_{EG} ={}& \lambda\,\mathbb E_{x}\lVert x - \hat x\rVert_1 - \mathbb E_{x}\big[D(\hat x)\big]\\ &- \mathbb E_{z}\big[D(G(z))\big] - \mathbb E_{x}\big[C(\hat z)\big]. \end{aligned}$$

Штрафы — те же, что в WGAN-GP, и вычисляются в случайных точках между настоящими входами и подделками:

$$\begin{aligned} \mathrm{GP}_D &= \mathbb E_{\tilde x}\Big[\big(\lVert\nabla_{\tilde x} D(\tilde x)\rVert_2 - 1\big)^2\Big],\\ \mathrm{GP}_C &= \mathbb E_{\tilde z}\Big[\big(\lVert\nabla_{\tilde z} C(\tilde z)\rVert_2 - 1\big)^2\Big]. \end{aligned}$$

Здесь $\tilde x = \alpha x + (1 - \alpha)\,x_{\mathrm f}$ лежит между настоящим изображением $x$ и поддельным $x_{\mathrm f}$, а $\tilde z = \alpha z + (1 - \alpha)\,\hat z$ — между сэмплом из априорного распределения и кодом, причём $\alpha \sim U(0, 1)$ сэмплируется для каждого примера отдельно; тильда заменяет крышку из заметки о WGAN, потому что здесь крышка обозначает реконструкцию. Kwon et al. используют $\lambda = \lambda_{\mathrm{GP}} = 10$.

Рисунок

α-WGAN: энкодер, генератор и два критика

1. Реконструкция. $E$ кодирует $x$, $G$ декодирует код, и $\lambda\lVert x - \hat x\rVert_1$ обучает обе сети.
2. Сэмплирование. Тот же $G$ декодирует сэмплы из априорного распределения в новые изображения.
3. Критик кодов. $\mathcal L_C$ — функция потерь критика WGAN на кодах, где роль настоящего распределения играет априорное, поэтому без штрафа $-\mathcal L_C$ оценивает $W_1(p, q)$ — с оговорками о константе Липшица критика из заметки о WGAN. Слагаемое энкодера $-\mathbb E\, C(\hat z)$ сдвигает коды туда, где критик принимает их за сэмплы из априорного распределения.
4. Критик изображений. $\mathcal L_D$ — это две функции потерь критика WGAN с общим критиком, $\big[\mathbb E D(\hat x) - \mathbb E D(x)\big] + \big[\mathbb E D(G(z)) - \mathbb E D(x)\big]$; слагаемое на настоящих данных входит в оба сравнения, отсюда множитель 2. Обозначим через $Q_{\mathrm{rec}}$ и $Q_{\mathrm{gen}}$ распределения реконструкций и сэмплов. Два слагаемых на подделках вместе равны $2\,\mathbb E D$ по их смеси с равными весами $M = \tfrac12(Q_{\mathrm{rec}} + Q_{\mathrm{gen}})$, так что наилучшее по 1-липшицевым критикам значение $-\mathcal L_D$ без штрафа равно $2\,W_1(P, M)$, а оно в силу выпуклости $W_1$ не больше $W_1(P, Q_{\mathrm{rec}}) + W_1(P, Q_{\mathrm{gen}})$.

$E$ и $G$ оптимизируются как одна сеть по сумме $\mathcal L_{EG}$ (раздел 2, «Training»). Каждое слагаемое доходит до тех сетей, от которых зависит:

Слагаемое $\mathcal L_{EG}$	Обучает	Чего требует
$\lambda\lVert x - \hat x\rVert_1$	$E$ и $G$	каждое изображение восстанавливается по своему коду
$-D(\hat x)$	$G$, а через $\hat x$ и $E$	реконструкции выглядят настоящими
$-D(G(z))$	$G$	сэмплы выглядят настоящими
$-C(\hat z)$	$E$	коды похожи на сэмплы из априорного распределения

Вторая строка отличается от алгоритма 1 α-GAN, где в функции потерь энкодера нет слагаемого на изображениях; при суммарной функции потерь градиент $-D(\hat x)$ попадает и в энкодер.

Зачем критик Вассерштейна и для кодов? В начале обучения коды детерминированного энкодера могут лежать далеко от априорного распределения, сосредоточившись в небольшой области. Тогда классификатор легко их отделяет, и его логит вокруг кодов может стать плоским — это та ситуация, с которой начиналась заметка о WGAN, — а липшицев критик сохраняет наклон с нормой около 1, направляющий коды к априорному распределению. Разобранный пример ниже сравнивает их, когда распределения перекрываются: градиент классификатора уменьшается вместе со сдвигом, а градиент критика — нет. Kwon et al. обосновывают это изменение неустойчивостью своей модели с функциями потерь классификатора.

Их приложение — трёхмерные МРТ-снимки мозга размером $64^3$ вокселей; все четыре сети обучались с Adam с шагом обучения (learning rate) $2 \cdot 10^{-4}$, батчами по 4 и двумя обновлениями генератора на шаг (разделы 2 и 4). Со 100-мерными кодами модель коллапсировала; лучше всего работали 1000-мерные коды. У её сэмплов было наименьшее среди сравниваемых моделей максимальное среднее расхождение (maximum mean discrepancy, MMD) с данными и попарный MS-SSIM 0.829 против 0.846 у настоящих снимков, тогда как базовая WGAN-GP дала 0.996 — признак почти одинаковых сэмплов (таблица 1).

Разобранный пример: критик кодов в одном измерении#

Игрушечный пример показывает, что сообщают энкодеру два вида критика кодов. Возьмём $d = 1$, априорное распределение $p = \mathcal N(0, 1)$ и предположим, что каждый код сдвинут на $\mu = 0.5$, так что $q = \mathcal N(0.5, 1)$.

Критик-классификатор (α-GAN). Оптимальное значение — логарифм отношения двух плотностей,

$$C^*(z) = \log\frac{p(z)}{q(z)} = -\frac{z^2}{2} + \frac{(z - 0.5)^2}{2} = -0.5\,z + 0.125 .$$

Его среднее по кодам $\mathbb E_q C^* = -0.5 \cdot 0.5 + 0.125 = -0.125$ равно $-D_{\mathrm{KL}}(q\,\|\,p) = -\mu^2/2$, как и должно быть для KL между гауссианами с единичной дисперсией, средние которых отличаются на $\mu$. Функция потерь энкодера $-\mathbb E_q C^*(\hat z)$ равна $0.125$, и при фиксированном критике её производная по сдвигу кодов равна $-\mathbb E\,[\partial C^*/\partial z] = 0.5$ — в точности $\partial D_{\mathrm{KL}}/\partial\mu = \mu$. По мере приближения кодов к априорному распределению этот градиент убывает вместе с $\mu$: при $\mu = 0.05$ он равен $0.05$.

Критик Вассерштейна (α-WGAN). Здесь $W_1(p, q) = \lvert\mu\rvert = 0.5$, а $f(z) = -z$ — оптимальный 1-липшицев критик: $\mathbb E_p f - \mathbb E_q f = 0 + 0.5 = 0.5$. Его градиент всюду имеет норму 1, поэтому штраф $\mathrm{GP}_C$ равен нулю, и $\mathcal L_C = \mathbb E_q f - \mathbb E_p f = -0.5 = -W_1$. Функция потерь энкодера $-\mathbb E_q f(\hat z) = \mathbb E_q z = 0.5$ имеет производную 1 по $\mu$ при любой величине сдвига: критик толкает коды с одной и той же силой, пока они не окажутся в центре.

Обучение: четыре сети, три оптимизатора#

Один шаг обучения состоит из двух частей.

1. Энкодер и генератор. Одно обновление $E$ и $G$ по $\mathcal L_{EG}$ при фиксированных $D$ и $C$.
2. Критики. Обновления $D$ по $\mathcal L_D$ и $C$ по $\mathcal L_C$. Реконструкции, сэмплы и коды входят в эти функции потерь как константы, так что шаги критиков не сдвигают $E$ и $G$.

Соотношение между двумя частями — дело выбора. WGAN-GP делает $n_{\text{critic}} = 5$ шагов критика на шаг генератора, чтобы критик оставался близким к оптимальному; α-GAN, наоборот, обновлял энкодер и генератор дважды на шаг дискриминаторов, а Kwon et al. дважды обновляли генератор. Каждый шаг критика должен получать новый мини-батч настоящих данных, как в алгоритме WGAN-GP.

Зачем отдельные оптимизаторы? Каждый игрок минимизирует свою функцию потерь, и эти функции тянут в противоположные стороны. Допустим, $\mathcal L_D$ и $\mathcal L_{EG}$ сложили и стали оптимизировать все параметры вместе. Тогда градиент по параметрам $D$ был бы равен

$$\begin{aligned} &\nabla\big(\mathbb E D(\hat x) + \mathbb E D(G(z)) - 2\,\mathbb E D(x)\big)\\ &\quad + \nabla\big(-\mathbb E D(\hat x) - \mathbb E D(G(z))\big) = -2\,\nabla\,\mathbb E D(x), \end{aligned}$$

если не учитывать штраф: слагаемые на подделках сокращаются, и критик научился бы лишь повышать своё значение на настоящих данных. Поэтому у каждой функции потерь свой оптимизатор по своим параметрам: один для $E$ и $G$ вместе, которые минимизируют одну суммарную функцию потерь, один для $D$ и один для $C$. Отдельные оптимизаторы позволяют также задавать разные шаги обучения; Rosca et al. перебирали шаги обучения отдельно для генератора и энкодера и для двух дискриминаторов (приложение H, таблица 2).

Расписание шага обучения и ранняя остановка#

Распространённое расписание — косинусный отжиг с тёплыми перезапусками (cosine annealing with warm restarts; Loshchilov and Hutter, 2017, формула 5). Внутри $i$-го цикла длиной $T_i$ шаг обучения убывает от $\eta_{\max}$ до $\eta_{\min}$ по половине периода косинуса,

$$\eta_t = \eta_{\min} + \tfrac12\big(\eta_{\max} - \eta_{\min}\big)\Big(1 + \cos\frac{\pi\, T_{\text{cur}}}{T_i}\Big),$$

где $T_{\text{cur}}$ — число эпох, прошедших с последнего перезапуска; затем шаг скачком возвращается к $\eta_{\max}$, а следующий цикл в $T_{\text{mult}}$ раз длиннее. При $T_0 = 10$ эпох, $T_{\text{mult}} = 2$, $\eta_{\max} = 2 \cdot 10^{-4}$, $\eta_{\min} = 0$ и одном обновлении за эпоху циклы длятся 10, 20 и 40 эпох, перезапуски приходятся на начало эпох 11, 31 и 71, а шаг обучения в течение эпохи 10 равен $4.9 \cdot 10^{-6}$. Перезапуски позволяют обучению покинуть плохую область, но временно ухудшают качество, как отмечают Loshchilov и Hutter (раздел 3).

Ранняя остановка завершает обучение, когда валидационная величина не улучшается заданное число эпох — терпение (patience), — и сохраняет лучшую контрольную точку (checkpoint). Для автоэнкодирующего GAN труднее всего выбрать саму величину.

• Ошибка реконструкции на отложенных изображениях — осмысленное и стабильное число, но оно измеряет автоэнкодер, а не сэмплы: в одном запуске реконструкции могут продолжать улучшаться, пока сэмплы коллапсируют, или наоборот.
• Сырые состязательные оценки, например слагаемое генератора $-\mathbb E\, D(G(z))$, не годятся. Оценки критика содержат произвольный аддитивный сдвиг, который дрейфует по мере обучения критика, поэтому их значения нельзя сравнивать между эпохами. Разность оценок на настоящих и поддельных данных этот сдвиг сокращает: при одинаково ограниченном критике, обученном близко к оптимуму, разрыв WGAN можно отслеживать в пределах одного запуска с оговорками из заметки о WGAN, хотя между запусками с разными критиками он по-прежнему несравним.
• С тёплыми перезапусками терпение короче следующего цикла может остановить запуск сразу после перезапуска. В примере выше лучшее значение на эпохе 10 при терпении 10 останавливает обучение на эпохе 20, когда шаг обучения второго цикла всё ещё равен $1.2 \cdot 10^{-4}$.
• Данные, по которым выбирается контрольная точка, должны быть валидационной выборкой, а не тестовой, на которой сообщаются результаты.

Что отслеживать и как оценивать#

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

Автоэнкодирующий GAN проверяют с трёх сторон.

• Реконструкции отложенных изображений, показанные рядом с оригиналами, проверяют $E$ и $G$ как автоэнкодер.
• Сэмплы из априорного распределения проверяют $G$ вместе с согласованием кодов. Их оценивают метриками из заметок об оценке качества, такими как Inception Score и FID. Rosca et al. также обучали независимого критика WGAN-GP отличать отложенные данные от сэмплов (раздел 5).
• Интерполяции между кодами двух настоящих изображений проверяют пространство между кодами — этому посвящён следующий раздел.

Латентная интерполяция: lerp или slerp?#

Энкодер позволяет пройти от одного настоящего изображения к другому: закодировать их, $z_0 = E(x_0)$ и $z_1 = E(x_1)$, двигаться по пути $z_t$ от $z_0$ к $z_1$ и декодировать $G(z_t)$. Плавная последовательность изображений говорит о том, что модель упорядочила пространство кодов, а не запомнила обучающую выборку. Простейший путь — линейная интерполяция (linear interpolation, lerp),

$$z_t = (1 - t)\,z_0 + t\,z_1, \qquad t \in [0, 1].$$

Сферическая линейная интерполяция (spherical linear interpolation, slerp) идёт вместо этого по большому кругу. Если $\Omega$ — угол между $z_0$ и $z_1$, $\cos\Omega = \langle z_0, z_1\rangle / (\lVert z_0\rVert\,\lVert z_1\rVert)$, то

$$z_t = \frac{\sin\big((1 - t)\,\Omega\big)}{\sin\Omega}\, z_0 + \frac{\sin(t\,\Omega)}{\sin\Omega}\, z_1 .$$

Эта формула пришла из анимации вращений (Shoemake, 1985); White (2016) предложил её для латентных пространств генеративных моделей. Когда $z_0$ и $z_1$ имеют одинаковую длину $r$, каждая точка пути имеет длину $r$. Формула перестаёт работать, когда $\sin\Omega$ близок к нулю, а это бывает в двух разных случаях. Если векторы почти параллельны ($\Omega \approx 0$), реализации переходят на lerp: при $\Omega \to 0$ коэффициенты slerp стремятся к $1 - t$ и $t$, так что lerp — приближение для малых углов с точностью до второго порядка по $\Omega$, а не тот же путь. Если они почти противоположны ($\Omega \approx \pi$), большая окружность не единственна, а lerp прошёл бы рядом с началом координат; тогда путь нужно выбрать явно, например через направление, ортогональное $z_0$, либо сообщить о вырожденной паре или пересэмплировать её.

Почему разница важна. В высокой размерности сэмплы $\mathcal N(0, I_d)$ лежат в тонкой оболочке: квадрат их длины имеет среднее $d$ и дисперсию $2d$, так что длина близка к $\sqrt d$ с разбросом около $0.71$ (Многомерное нормальное распределение). При $d = 128$ средняя длина равна $11.29$. Кроме того, два независимых сэмпла почти ортогональны: косинус угла между ними имеет среднее 0 и стандартное отклонение $1/\sqrt d = 0.088$, около $5°$. Их середина по lerp $(z_0 + z_1)/2$ — сумма независимых гауссовских векторов, распределённая как $\mathcal N(0, \tfrac12 I)$, поэтому её длина близка к $\sqrt{d/2} = 8.0$ (среднее $7.98$, разброс $0.50$). При априорном распределении квадрат длины не больше $d/2 = 64$ имеет вероятность $4 \times 10^{-7}$: середина лежит на $4.7$ стандартного отклонения ближе к центру, чем оболочка, в области, которую обучение на $z \sim p(z)$ почти никогда не посещает. Вдоль всего пути средний квадрат длины равен $d\big((1 - t)^2 + t^2\big)$ и минимален в середине. Slerp между двумя почти ортогональными векторами держит длину между $\lVert z_0\rVert$ и $\lVert z_1\rVert$: при $\Omega = 90°$ квадрат длины точки пути равен $\cos^2(\pi t/2)\lVert z_0\rVert^2 + \sin^2(\pi t/2)\lVert z_1\rVert^2$.

Рисунок

Линейная интерполяция уходит внутрь оболочки, сферическая — нет

Разобранный пример. В двух измерениях возьмём $z_0 = (3, 0)$ и $z_1 = (0, 3)$. Середина по lerp — точка $(1.5, 1.5)$ длины $3/\sqrt 2 = 2.12$. Для slerp $\Omega = 90°$, и оба коэффициента при $t = \tfrac12$ равны $\sin 45° / \sin 90° = 0.707$, так что середина — точка $(2.12, 2.12)$ длины 3. В двух измерениях укорочение — лишь иллюстрация, поскольку случайные пары не ортогональны, а оболочка не тонкая. В 128 измерениях почти каждая пара близка к ортогональной, а оболочка тонкая, так что тот же множитель $1/\sqrt 2$ выводит из неё почти каждую середину по lerp.

Для автоэнкодирующего GAN концы пути — коды настоящих изображений, а не сэмплы из априорного распределения. Если критик кодов согласовал $q(z)$ с априорным распределением, то длины кодов типичны для априорного распределения, коды не связанных между собой изображений почти ортогональны, и рассуждение переносится; коды похожих изображений близки по углу, и тогда два пути почти совпадают. Slerp лишь сохраняет типичную длину. Он не гарантирует, что путь обходит области, где в $q(z)$ есть дыры; White называет их мёртвыми зонами (dead zones) и предлагает держаться вблизи кодов отложенных данных (раздел 2.3).

Что изменилось#

	WGAN-GP	AAE	α-GAN	α-WGAN
Сети	$G$, критик	$E$, $G$, $C$	$E$, $G$, $D$, $C$	$E$, $G$, $D$, $C$
Энкодер	нет	детерминированный или гауссовский	детерминированный или неявный	детерминированный
Чем коды согласуются с априорным распределением	—	классификатором на кодах	классификатором на кодах	критиком WGAN-GP на кодах
Пиксельное слагаемое	нет	квадратичная ошибка	$\lambda\lVert x - \hat x\rVert_1$	$\lambda\lVert x - \hat x\rVert_1$
Что видит дискриминатор изображений	сэмплы	—	реконструкции и сэмплы	реконструкции и сэмплы
Состязательная функция потерь $E$, $G$	$-f(G(z))$	обмануть $C$ на кодах	логит со знаком минус	значение критика со знаком минус
Защита от коллапса мод	нет	реконструкция	реконструкция	реконструкция
Восстанавливает и интерполирует настоящие изображения	нет	да	да	да

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

• Много движущихся частей. Четыре сети, три оптимизатора, веса $\lambda$ и $\lambda_{\mathrm{GP}}$, соотношения обновлений и шаги обучения влияют друг на друга. Rosca et al. обнаружили, что все протестированные ими модели чувствительны к архитектурам, причём незначительные изменения приводили к катастрофическому коллапсу мод (раздел 6).
• Нет выигрыша в качестве сэмплов во всех случаях. В их переборах независимый критик ставил WGAN-GP выше α-GAN на всех трёх наборах данных; за энкодер и покрытие приходится платить настройкой.
• Вес реконструкции — компромисс. Большое $\lambda$ возвращает в реконструкции размытость пиксельной функции потерь; малое ослабляет защиту от коллапса мод.
• Распределение кодов согласуется лишь приближённо. При детерминированном энкодере и конечной обучающей выборке $q(z)$ — облако точек, поэтому отношение плотностей и KL, которые оценивает классификатор, существуют только для его сглаженной версии, и дыры остаются возможны. Советы о размере кода расходятся: статья об AAE рекомендует для детерминированных энкодеров размерность кода около внутренней размерности данных, а Kwon et al. понадобилось 1000 измерений, чтобы избежать коллапса.
• Штраф на градиент и батч-нормализация. Штраф вычисляется для каждого примера отдельно, а батч-нормализация (batch normalization) в критике делает каждый выход зависящим от всего батча, поэтому WGAN-GP не рекомендует её использовать (Gulrajani et al., 2017, раздел 4). Kwon et al. тем не менее использовали батч-нормализацию в своих критиках.
• Оценка остаётся косвенной. Правдоподобия по-прежнему нет, ошибка реконструкции измеряет только автоэнкодер, а функции потерь критиков нельзя сравнивать между эпохами или моделями.

Итоги#

• У VAE есть энкодер, и он покрывает данные, но размывает изображения; GAN даёт чёткие сэмплы, но не имеет энкодера и может пропускать моды. Автоэнкодирующие GAN сохраняют энкодер и слагаемое реконструкции, а проблемные части ELBO заменяют дискриминаторами.
• Состязательный автоэнкодер заменяет слагаемое KL дискриминатором на кодах. Он согласует агрегированное апостериорное распределение $q(z)$ с априорным — ту часть KL, которая нужна для сэмплирования, — и требует от априорного распределения только сэмплов.
• α-GAN заменяет ещё и пиксельное правдоподобие дискриминатором изображений, который видит реконструкции и сэмплы, и сохраняет $\lambda\lVert x - \hat x\rVert_1$. Состязательная функция потерь $E$ и $G$ — логит со знаком минус. Согласованные коды вместе с точной реконструкцией сделали бы распределение сэмплов равным распределению данных; слагаемое реконструкции препятствует коллапсу мод.
• α-WGAN обучает оба дискриминатора как критиков WGAN-GP и оптимизирует $E$ и $G$ по одной суммарной функции потерь, с $\lambda = \lambda_{\mathrm{GP}} = 10$.
• Обучению нужны отдельный оптимизатор для каждого игрока, стабильный валидационный сигнал (ошибка реконструкции на отложенных данных, поскольку значения критиков несравнимы во времени) и три проверки: реконструкции, сэмплы и интерполяции.
• Интерполировать следует по slerp: в $d$ измерениях сэмплы из априорного распределения имеют длину около $\sqrt d$, а середина по lerp между двумя из них — длину около $\sqrt{d/2}$, вне оболочки, на которой обучался генератор.

Ноутбук#

Ноутбук об α-WGAN из репозитория HSE-LAMBDA Deep Generative Models обучает α-WGAN на CIFAR-10 (ячейки нумеруются с 1, считая все ячейки подряд, включая текстовые). Сохранённых выходов в нём нет, поэтому он документирует процедуру, а не результаты.

• Ячейки 1–4: подготовка. Ячейка 1 описывает план ноутбука. Ячейка 3 запускает эксперимент Comet.ml с проектом, рабочим пространством и API-ключом, считанными из локального cfg.yaml; глобальная переменная experiment используется затем внутри save_checkpoint и train_alphawgan, так что без учётной записи Comet ноутбук не запустится, если не убрать эти вызовы. Ячейка 4 импортирует библиотеки и выбирает CUDA, Apple MPS или CPU.
• Ячейки 5–6: данные. Обучающая (50 000 изображений) и тестовая (10 000) части CIFAR-10, приведённые от $32 \times 32$ к $64 \times 64$ и отмасштабированные в $[-1, 1]$; размер батча 128, обучающий загрузчик перемешивает данные, так что в эпохе 391 батч, последний — из 80 изображений. Повышение разрешения не добавляет деталей, но учетверяет стоимость каждой свёртки.
• Ячейки 7–8: строительные блоки. Ячейка 7 перечисляет четыре сети и показывает схему архитектуры. Ячейка 8 определяет ChannelsToLinear (вытягивание в вектор, затем линейный слой), LinearToChannels2d (линейный слой, выход которого переформатируется в карты признаков) и ResBlock: две свёртки $3 \times 3$ с батч-нормализацией, нормальная инициализация Ксавье (Xavier-normal) и прибавление входа перед последней активацией.
• Ячейки 9–12: энкодер и генератор. Энкодер — свёртка $5 \times 5$, три усредняющих пулинга ($64 \to 8$), чередующихся с остаточными блоками, и линейное отображение в код размера $d$ без выходной активации. Генератор устроен зеркально: линейное отображение в карты признаков $8 \times 8$, три повышения разрешения $\times 2$ методом ближайшего соседа с остаточными блоками, свёртка $1 \times 1$ в RGB и tanh. При base_channels = 256 из ячейки 24 вытянутые признаки содержат $256 \cdot 8 \cdot 8 = 16\,384$ элемента; комментарии в коде приводят числа для 128 каналов по умолчанию.
• Ячейки 13–14: критики. ImageDiscriminator устроен как энкодер, с LeakyReLU(0.2) и одним выходом; его остаточные блоки содержат батч-нормализацию. CodeDiscriminator — MLP $d \to 700 \to 700 \to 1$ с LeakyReLU(0.2), без нормализации и с инициализацией Ксавье; у Rosca et al. было 750 нейронов.
• Ячейка 15: gradient_penalty. Одно $\alpha \sim U(0, 1)$ на пример, распространяемое (broadcasting) на остальные измерения; интерполяции, отсоединённые от графа и помеченные как требующие градиентов; autograd.grad с create_graph=True; норма $L_2$ для каждого примера и $\lambda_{\mathrm{GP}} \cdot \mathrm{mean}\big((\lVert g\rVert_2 - 1)^2\big)$. Одна и та же функция служит и для изображений, и для кодов.
• Ячейки 16–17: класс AlphaWGAN. sample_prior сэмплирует $z \sim \mathcal N(0, I)$; rec_loss — средняя абсолютная ошибка по всем пикселям, так что $\lambda$ умножает среднее по пикселям (по умолчанию в классе 40, в ячейке 24 задано 10). autoencoder_loss возвращает $\lambda \cdot L_1$, $-\mathrm{mean}\,C(E(x))$ и $-\mathrm{mean}\,D(G(z))$; discriminator_loss — функцию потерь WGAN для $D$ на настоящих изображениях против отсоединённых сэмплов $G(z)$ плюс $\mathrm{GP}_D$; code_discriminator_loss — функцию потерь WGAN для $C$ на сэмплах из априорного распределения против отсоединённых кодов плюс $\mathrm{GP}_C$. forward кодирует, генерирует, сэмплирует (если передано целое число) или реконструирует.
• Ячейки 18–22: вспомогательные функции обучения. save_checkpoint сохраняет веса и логирует их в Comet; evaluate_alphagan усредняет три слагаемых энкодера и генератора по загрузчику в режиме оценки (evaluation mode) без градиентов; EarlyStopping хранит контрольную точку с наименьшим значением отслеживаемой функции потерь и останавливает обучение после patience эпох без улучшения; visualize_samples декодирует сэмплы из априорного распределения. train_alphawgan создаёт три оптимизатора AdamW (энкодер и генератор вместе, $D$, $C$) и три планировщика CosineAnnealingWarmRestarts с $T_0 = 10$ и $T_{\text{mult}} = 2$, которые делают шаг раз в эпоху. На каждом батче выполняется один шаг энкодера и генератора по сумме трёх слагаемых, затем n_critic раундов из одного шага $D$ и одного шага $C$ на том же батче настоящих данных; после каждого обратного прохода в Comet логируется гистограмма градиента каждого параметра только что обновлённых сетей. В конце каждой эпохи функции потерь усредняются, модель оценивается на тестовой части, логируются метрики и рисунок с сэмплами, каждые checkpoint_every эпох сохраняется контрольная точка, и применяется ранняя остановка с терпением 10 по валидационной функции потерь реконструкции.
• Ячейки 23–29: конфигурация и запуск. $d = 128$, base_channels = 256, $\lambda = 10$, $\lambda_{\mathrm{GP}} = 10$; ячейка 26 загружает checkpoints/checkpoint_best.pth, если файл существует (только веса, так что оптимизаторы и планировщики начинают с нуля). Обучение длится до 100 эпох с шагом обучения $2 \cdot 10^{-4}$, betas $(0, 0.9)$, $n_{\text{critic}} = 5$ и контрольной точкой каждые 10 эпох (вместо значений функции по умолчанию $8 \cdot 10^{-4}$ и 3 раундов критика); ячейка 28 логирует гиперпараметры. Расписание даёт циклы в 10, 20 и 40 эпох, вычисленные в основном тексте, с перезапусками на эпохах 11, 31 и 71.
• Ячейки 30–36: оценка. visualize_reconstructions показывает 16 тестовых изображений над их реконструкциями; ячейки 32–33 строят графики валидационной и обучающей функций потерь реконструкции (умноженных на $\lambda$, без подписей осей и легенды); ячейка 34 заново загружает лучшую контрольную точку; ячейки 35–36 показывают 16 сэмплов и реконструкции.
• Ячейки 37–39: интерполяция. slerp реализует формулу выше на ненормированных кодах, с углом, вычисленным по их направлениям, и переходом к lerp при $\sin\Omega < 10^{-6}$; lerp — прямая линия. interpolate_latent кодирует первые num_images изображений батча и декодирует 7 равноотстоящих точек между каждыми двумя последовательными кодами. Ячейка 39 берёт первые четыре тестовых изображения (тестовый загрузчик не перемешивает данные) и даёт три ряда по семь изображений со slerp; method="lerp" переключает на прямую линию.

Проблемы.

• Реконструкции не доходят до критика изображений. discriminator_loss сравнивает настоящие изображения только с сэмплами $G(z)$, а в autoencoder_loss нет слагаемого $-D(\hat x)$. И α-GAN (алгоритм 1), и α-WGAN (Kwon et al., формулы 1–2) считают реконструкции подделками. В нынешнем виде реконструкции обучаются только слагаемым $L_1$, поэтому сохраняют его усреднение, а модель — это состязательный автоэнкодер с дополнительным критиком на сэмплах. Полная версия добавляет D(x_rec.detach()) в функцию потерь критика, удваивая вес слагаемого на настоящих данных, и $-D(\hat x)$ — в функцию потерь энкодера и генератора.
• Текстовая ячейка 16 утверждает, что состязательные функции потерь — бинарные кросс-энтропии; код использует функции потерь Вассерштейна со штрафами на градиент. Строка документации (docstring) train_alphawgan упоминает обрезку градиентов, которой в коде нет.
• Батч-нормализация в критике изображений противоречит штрафу на градиент, вычисляемому для каждого примера отдельно, а настоящие изображения, сэмплы и интерполяции проходят через неё тремя отдельными батчами с разной статистикой. WGAN-GP рекомендует критика без батч-нормализации, например с нормализацией слоя (layer normalization); с критиком кодов, у которого нормализации нет, всё в порядке.
• Шаги критика используют один и тот же батч настоящих данных. Все пять раундов видят один и тот же x; заново сэмплируются только сэмплы из априорного распределения. WGAN-GP берёт новый мини-батч настоящих данных для каждого шага критика.
• Два падения при запуске с нуля. EarlyStopping инициализирует лучшее значение через np.Inf, которое удалено в NumPy 2.0 (актуальное имя — np.inf). Его первый вызов, в конце эпохи 1, сохраняет файл в checkpoints/checkpoint_best.pth раньше, чем что-либо создаст каталог checkpoints (save_checkpoint, который его создаёт, впервые вызывается на эпохе 10), а torch.save выдаёт ошибку, если родительского каталога нет; EarlyStopping нужен собственный os.makedirs.
• Ранняя остановка выбирает модель по тестовой выборке и следит только за функцией потерь реконструкции. Её терпение в 10 эпох короче второго цикла тёплых перезапусков, так что лучшее значение на эпохе 10 останавливает запуск на эпохе 20, когда шаг обучения ещё равен $1.2 \cdot 10^{-4}$. Валидационные состязательные слагаемые в evaluate_alphagan — значения меняющихся критиков, и их нельзя сравнивать между эпохами.
• Устаревшие или хрупкие вызовы. make_grid(..., range=(-1, 1)) в ячейке 38 использует имя аргумента, которое torchvision переименовал в value_range (остальные ячейки уже используют новое имя). У torch.load в ячейках 26 и 34 нет map_location, поэтому контрольная точка, сохранённая на GPU, не загрузится на машине без CUDA.
• Лишняя работа. Обратный проход энкодера и генератора заполняет также градиенты $D$ и $C$; перед шагами критиков они обнуляются, так что на результаты это не влияет. Шаги критиков вычисляют $G(z)$ и $E(x)$ с отслеживанием градиентов, прежде чем отсоединить их, причём в режиме обучения, так что они ещё и обновляют статистики батч-нормализации в $E$ и $G$. Логирование гистограмм градиентов после каждого из 11 обратных проходов батча, с эпохой в качестве шага, копирует градиенты на CPU $391 \times 11$ раз за эпоху.
• Мелочи. AdamW (Loshchilov and Hutter, 2019) применяет своё затухание весов (weight decay) по умолчанию, 0.01, ко всем сетям, включая критиков, тогда как в статьях используется Adam. ResBlock включает смещения свёрток именно тогда, когда за свёрткой следует слой нормализации (bias=(norm_layer is not None)), — обратно обычному выбору; лишние смещения безвредны. Текстовая ячейка 37 объясняет slerp тем, что латентное пространство «часто лежит на гиперсфере после нормировки» (often lies on a hypersphere when normalized); точная причина в том, что гауссовские сэмплы концентрируются вблизи сферы радиуса $\sqrt d$.

Дальше#

До сих пор каждый GAN в курсе генерировал изображения из шума, а энкодер этой заметки лишь придал этому шуму смысл. Rosca et al. указывают на связь со следующей моделью (раздел 4): в пространстве изображений α-GAN согласует реконструкции с изображениями и поточечно, и по распределению, а в пространстве кодов — только по распределению; добавление поточечной реконструкции кода $\lVert z - E(G(z))\rVert$ сделало бы его похожим на CycleGAN. Следующая заметка заменяет пространство кодов вторым доменом изображений: два генератора переводят изображения между доменами без парных примеров, состязательные функции потерь заставляют каждый перевод выглядеть как изображение другого домена, а цикл реконструкций привязывает каждый выход к его входу.

Литература#

• A. Makhzani, J. Shlens, N. Jaitly, I. Goodfellow, B. Frey. Adversarial autoencoders. arXiv:1511.05644, 2015. Агрегированное апостериорное распределение — формула 1; сравнение с VAE — раздел 2.1; детали экспериментов — в приложении A.
• M. Rosca, B. Lakshminarayanan, D. Warde-Farley, S. Mohamed. Variational approaches for auto-encoding generative adversarial networks. arXiv:1706.04987, 2017. α-GAN: раздел 3, формулы 6–9; алгоритм 1 и гиперпараметры — в приложениях B и H.
• G. Kwon, C. Han, D.-S. Kim. Generation of 3D brain MRI using auto-encoding generative adversarial networks. MICCAI, 2019, pp. 118–126; arXiv:1908.02498. Функции потерь α-WGAN — формулы 1–3.
• I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, A. Courville. Improved training of Wasserstein GANs. NeurIPS, 2017.
• T. White. Sampling generative networks. arXiv:1609.04468, 2016. Slerp для латентных пространств — раздел 2.1; мёртвые зоны — раздел 2.3.
• K. Shoemake. Animating rotation with quaternion curves. SIGGRAPH, Computer Graphics 19(3), 245–254, 1985.
• I. Loshchilov, F. Hutter. SGDR: Stochastic gradient descent with warm restarts. ICLR, 2017. Расписание — формула 5, раздел 3.
• I. Loshchilov, F. Hutter. Decoupled weight decay regularization. ICLR, 2019. AdamW.
• A. Radford, L. Metz, S. Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. ICLR, 2016. DCGAN.
• D. Ulyanov, A. Vedaldi, V. Lempitsky. It takes (only) two: adversarial generator-encoder networks. arXiv:1704.02304, 2017. AGE.
• HSE-LAMBDA, репозиторий курса Deep Generative Models: alpha-WGAN.ipynb.