Содержание курса 11 / 33

11 / 33 · GAN

Масштабирование GAN: прогрессивный рост, self-attention и BigGAN

Как GAN дошли до мегапиксельных лиц и условной по классу генерации ImageNet: ProGAN выращивает обе сети от 4 × 4 с плавным включением новых слоёв, SAGAN добавляет self-attention, спектральную нормализацию в обеих сетях, два шага обучения и кусочно-линейную функцию потерь, а BigGAN увеличивает батч и ширину, передаёт класс через общий эмбеддинг и проекционный дискриминатор и обменивает разнообразие на качество с помощью трюка усечения (truncation trick).

На этой странице 12

Предыдущая заметка дала обучению GAN функцию потерь, которая сохраняет градиент, когда два распределения не перекрываются, и два способа сохранить критика липшицевым: штраф на градиент и спектральную нормализацию. Вместе с рекомендациями DCGAN для устройства сетей этого достаточно, чтобы обучаться на лицах или спальнях с разрешением 64×6464 \times 64. Для мегапиксельных изображений или набора данных из тысячи классов этого мало. С 2017 по 2019 год три статьи закрыли бо́льшую часть этого разрыва, каждая своей идеей. ProGAN выращивает обе сети в ходе обучения, от 4×44 \times 4 до 1024×10241024 \times 1024 (Karras et al., 2018). SAGAN добавляет self-attention (внимание карты признаков к самой себе), чтобы согласовывать далёкие друг от друга части изображения, и закрепляет набор стабилизирующих решений (Zhang et al., 2019). BigGAN сохраняет рецепт SAGAN и масштабирует его: батч в восемь раз больше, более широкие слои и лучший способ передавать метку класса (Brock, Donahue and Simonyan, 2019). Заметка проходит эти три шага и то, что измерил каждый из них.

Обозначения. GG — генератор, DD — дискриминатор, z∼pzz \sim p_z — латентный код, xx — изображение. Как в заметке о GAN, PP — распределение данных с плотностью pdatap_{\text{data}}, а pgp_g — плотность генератора. c∈{1,…,K}c \in \{1, \dots, K\} — метка класса (K=1000K = 1000 для ImageNet), e(c)e(c) — её обучаемый вектор эмбеддинга. У карты признаков CC каналов и H×WH \times W позиций, N=HWN = HW. Inception Score (IS, чем больше, тем лучше) и Fréchet Inception Distance (FID, чем меньше, тем лучше) — два числа, которые приводит каждая из этих статей. Они определены в заметках об оценке качества дальше в курсе (IS, FID); пока достаточно знать, что IS поощряет сэмплы, которые предобученный классификатор уверенно распознаёт и которые распределены по многим классам, а FID измеряет, насколько признаки сгенерированных изображений далеки от признаков настоящих.

Почему большие изображения и много классов — это трудно#

Как только растёт разрешение или число классов, возникают три проблемы.

  • Дискриминатор побеждает слишком легко. При большем разрешении сгенерированные изображения легче отличить от настоящих, и это усиливает проблему с градиентом из заметки о GAN. Кроме того, чтобы большие изображения уместились в памяти, приходится брать маленькие мини-батчи, а это делает обучение менее устойчивым (Karras et al., 2018, раздел 1).
  • Свёртки локальны. На ImageNet свёрточные GAN хорошо рисовали классы, которые определяются текстурой, например океан или небо, и не справлялись с классами с устойчивой геометрией: у собак получался реалистичный мех, но не было чётко разделённых лап (Zhang et al., 2019, раздел 1).
  • Один генератор должен покрыть много мод. Тысяча классов — это тысяча групп мод, и коллапс на несколько из них случается легко.

ProGAN: выращивание сетей#

Ответ ProGAN на первую проблему — учебный план (curriculum). Обучение начинается с крошечных изображений, в которых меньше информации и меньше мод, а разрешение добавляется по ходу обучения, так что сети сначала выучивают крупномасштабную структуру, а затем всё более мелкие детали (Karras et al., 2018, раздел 2). Обе сети начинают с 4×44 \times 4. Каждый раз, когда разрешение удваивается, в конец GG добавляется блок слоёв, а его зеркальное отражение — в начало DD. Все уже существующие слои продолжают обучаться. При 1024×10241024 \times 1024 у генератора позади восемь таких удвоений.

Плавное включение новых слоёв#

Новый блок начинается со случайных весов. Резкое включение стало бы ударом для уже обученных слоёв. ProGAN вводит его плавно (fade-in; рисунок 2 статьи). Пусть ara_r — признаки генератора при разрешении rr, up\mathrm{up} — повышение разрешения в 2 раза методом ближайшего соседа, toRGB\mathrm{toRGB} — свёртка 1×11 \times 1 из признаков в цвета, а block2r\mathrm{block}_{2r} — новые свёртки. Во время перехода к разрешению 2r2r выходное изображение равно

x2r=(1−α) up(toRGBr(ar))+α toRGB2r(block2r(up(ar))),\begin{aligned} x_{2r} ={}& (1 - \alpha)\,\mathrm{up}\big(\mathrm{toRGB}_r(a_r)\big) \\ &+ \alpha\,\mathrm{toRGB}_{2r}\big(\mathrm{block}_{2r}(\mathrm{up}(a_r))\big), \end{aligned}

где α\alpha линейно растёт от 00 до 11. При α=0\alpha = 0 сеть совпадает со старой, а её выход просто увеличен; при α=0.25\alpha = 0.25 три четверти каждого пикселя всё ещё приходят по старому пути. Дискриминатор зеркально повторяет это с усредняющим пулингом (average pooling) в 2 раза и слоями 1×11 \times 1 fromRGB\mathrm{fromRGB}. Настоящие изображения уменьшаются до текущего разрешения и во время перехода интерполируются между двумя разрешениями так же, как выход генератора.

Расписание отсчитывается в настоящих изображениях, показанных DD: 800 тыс. изображений при 4×44 \times 4, затем для каждого нового разрешения 800 тыс. изображений плавного включения и 800 тыс. стабилизации. Мини-батч равен 16 вплоть до 128×128128 \times 128 и уменьшается до 14, 6 и 3 при 2562256^2, 5122512^2 и 102421024^2, чтобы уместиться в памяти (приложение A.1).

Стандартное отклонение по мини-батчу#

Для борьбы с коллапсом мод минибатч-дискриминация (minibatch discrimination) позволяет дискриминатору видеть статистики всего батча через обучаемый тензор. ProGAN заменяет её слоем без параметров — стандартным отклонением по мини-батчу (minibatch standard deviation; раздел 3). Ближе к концу DD, при 4×44 \times 4, этот слой вычисляет стандартное отклонение каждого признака в каждой позиции по мини-батчу, усредняет их в одно число и добавляет это число как дополнительную постоянную карту признаков, так что 512 каналов превращаются в 513 (таблица 2). У сколлапсировавшего батча разброс мал, и DD может научиться его отвергать.

Игрушечный пример. Один признак в одной позиции, батч из четырёх. Настоящие изображения дают (1,−1,2,0)(1, -1, 2, 0) со стандартным отклонением 1.1181.118. Сколлапсировавший генератор даёт (1.0,1.1,0.9,1.0)(1.0, 1.1, 0.9, 1.0) со стандартным отклонением 0.0710.071. Дополнительный канал несёт 1.1181.118 для настоящего батча и 0.0710.071 для поддельного, и одного веса на нём достаточно, чтобы их разделить.

В абляционном эксперименте этот слой улучшил расстояние Вассерштейна по срезам (метрику на статистиках фрагментов изображений; см. ниже), а минибатч-дискриминация — нет (таблица 1, строки e и e*).

Выравненный шаг обучения#

ProGAN инициализирует каждый вес из N(0,1)\mathcal N(0, 1) и при каждом прямом проходе умножает его на константу слоя из инициализации Хе (He initialization; раздел 4.1):

w^=2/nin  w,w∼N(0,1),\hat w = \sqrt{2 / n_{\text{in}}}\;w, \qquad w \sim \mathcal N(0, 1),

где ninn_{\text{in}} — число входов (fan-in) слоя. В момент инициализации это и есть инициализация Хе. Разница проявляется во время обучения. Adam делит каждое обновление на скользящую оценку масштаба градиента, поэтому за шаг он сдвигает каждый хранимый параметр примерно на величину шага обучения, каким бы ни был масштаб самого параметра. Если хранимые параметры имеют разные масштабы, один и тот же шаг для одних оказывается большим относительным изменением, а для других — малым.

Разобранный пример. ProGAN использует шаг обучения 10−310^{-3}. У свёртки 3×33 \times 3 с 512 входными каналами nin=4608n_{\text{in}} = 4608 и масштаб Хе 2/4608=0.0208\sqrt{2/4608} = 0.0208; у слоя 1×11 \times 1 fromRGB\mathrm{fromRGB} с 3 входными каналами масштаб 2/3=0.816\sqrt{2/3} = 0.816. Если хранимые веса инициализированы по Хе, один шаг меняет первый слой примерно на 10−3/0.0208=4.8%10^{-3}/0.0208 = 4.8\% его масштаба, а второй — на 0.12%0.12\%, в 39 раз меньше. При масштабировании во время прохода каждый хранимый вес имеет масштаб 1, и каждый слой меняется примерно на 0.1%0.1\% за шаг: скорость обучения во всех слоях одинакова, отсюда и название (equalized learning rate).

Попиксельная нормализация признаков#

В GAN величины сигналов склонны разрастаться по мере соперничества двух сетей (раздел 4). Вместо батч-нормализации генератор ProGAN после каждой свёртки 3×33 \times 3 нормирует вектор признаков в каждом пикселе (pixelwise feature normalization; раздел 4.2):

bx,y=ax,y1C∑j=1C(ax,yj)2+ε,b_{x,y} = \frac{a_{x,y}}{\sqrt{\frac1C\sum_{j=1}^{C}\big(a^{j}_{x,y}\big)^2 + \varepsilon}},

где ax,y∈RCa_{x,y} \in \mathbb R^C — вектор из CC каналов в пикселе (x,y)(x, y), а ε=10−8\varepsilon = 10^{-8}. Для a=(3,4)a = (3, 4) при C=2C = 2 средний квадрат равен 12.512.5 и b=(0.849,1.131)b = (0.849, 1.131): среднеквадратичное значение каналов становится равным 1, поэтому длина вектора равна C=1.414\sqrt C = 1.414. У слоя нет параметров, и он не смешивает примеры. Авторы обнаружили, что он редко меняет результаты, но предотвращает разрастание, когда оно возникает.

Ни одна из сетей не использует батч-нормализацию, нормализацию слоя или нормализацию весов (приложение A.1). Функция потерь — WGAN-GP с одним шагом критика на шаг генератора (ncritic=1n_{\text{critic}} = 1) и дополнительным членом 0.001 Ex∼P[D(x)2]0.001\,\mathbb E_{x\sim P}[D(x)^2], который не даёт выходу критика уходить далеко от нуля; у Adam β1=0\beta_1 = 0, β2=0.99\beta_2 = 0.99. Латентный вектор 512-мерный и нормирован так, чтобы лежать на гиперсфере. Изображения сэмплируются из экспоненциального скользящего среднего весов генератора с коэффициентом затухания 0.999. Та же сеть обучалась и с функцией потерь наименьших квадратов (least-squares loss), так что вклад этих приёмов в основном не зависит от функции потерь (раздел 6.3).

Что дал прогрессивный рост#

Таблица 1 статьи добавляет нововведения по одному на CelebA при 128×128128 \times 128, в намеренно маленькой сети, обученной на 10 млн настоящих изображений. Метрика — расстояние Вассерштейна по срезам (sliced Wasserstein distance, SWD) между распределениями фрагментов 7×77 \times 7 настоящих и сгенерированных изображений на нескольких уровнях пирамиды Лапласа, усреднённое по уровням; чем меньше, тем лучше.

Конфигурация (изменения накапливаются)Среднее SWD ×103\times 10^3
(a) базовая модель WGAN-GP, мини-батч 649.28
(b) + прогрессивный рост (progressive growing)4.28
(c) + мини-батч 1646.23
(d) + пересмотренные гиперпараметры, без батч-нормализации и нормализации слоя8.07
(e) + стандартное отклонение по мини-батчу7.04
(f) + выравненный шаг обучения4.39
(g) + попиксельная нормализация3.56
(h) полная сеть, обученная до сходимости2.96

Один только прогрессивный рост уменьшает расстояние больше чем вдвое, с 9.28 до 4.28. Строка (c) показывает, зачем понадобились остальные изменения: уменьшение мини-батча до размера, который допускает высокое разрешение, ломает базовую модель. Прогрессивный рост к тому же быстрее, поскольку большинство итераций проходит при низком разрешении: при 1024×10241024 \times 1024 он дошёл до 6.4 млн изображений за 96 часов на одном GPU, тогда как сети фиксированного размера понадобилось бы около 520 часов, — ускорение в 5.4×5.4\times (раздел 6.2).

Для главного результата авторы собрали CelebA-HQ — 30 000 изображений лиц размера 1024×10241024 \times 1024, отобранных из CelebA и обработанных (приложение C), — и обучали модель на 8 GPU Tesla V100 в течение 4 дней (раздел 6.3). Модель также генерировала изображения 256×256256 \times 256 в 30 категориях LSUN и установила рекорд Inception Score для обучения без учителя — 8.80 на CIFAR-10 (раздел 6.5).

SAGAN: self-attention для дальней структуры#

ProGAN работал на лицах — выровненных и однотипных. На ImageNet трудность другая: много классов, объекты в любом месте кадра и структура, охватывающая всё изображение.

Почему свёртки упускают дальнюю структуру#

Свёртка 3×33 \times 3 с шагом 1 смешивает каждую позицию с её непосредственными соседями. Стопка из LL таких слоёв видит окно (2L+1)×(2L+1)(2L + 1) \times (2L + 1), так что информация продвигается не более чем на одну позицию за слой: выход в одной позиции начинает зависеть от входа, отстоящего на 16 позиций, только через 16 слоёв. Сеть может связать далёкие области, но лишь через множество слоёв, которые должны действовать согласованно, а маленькая модель может оказаться неспособной это представить, и оптимизатор может такое решение не найти. Помогли бы ядра побольше, но ценой эффективности локальных свёрток (Zhang et al., 2019, раздел 1). Self-attention связывает каждую пару позиций за один слой.

Блок self-attention#

SAGAN переносит в обе сети GAN нелокальный блок (non-local block), который раньше использовался для моделирования зависимостей в видео (Zhang et al., 2019, раздел 3). В этом разделе xx — не изображение, а карта признаков на входе блока, рассматриваемая как NN векторов xi∈RCx_i \in \mathbb R^C, по одному на позицию. Три свёртки 1×11 \times 1, которые действуют на каждую позицию отдельно, дают

f(xi)=Wfxi,g(xi)=Wgxi,h(xi)=Whxi,f(x_i) = W_f x_i, \quad g(x_i) = W_g x_i, \quad h(x_i) = W_h x_i,

где Wf,Wg,Wh∈RCˉ×CW_f, W_g, W_h \in \mathbb R^{\bar C \times C}. Для каждой пары позиций вычисляются оценка (score) и вес внимания:

sij=f(xi)⊤g(xj),βj,i=exp⁡(sij)∑k=1Nexp⁡(skj).\begin{gathered} s_{ij} = f(x_i)^\top g(x_j),\\ \beta_{j,i} = \frac{\exp(s_{ij})}{\sum_{k=1}^{N}\exp(s_{kj})} . \end{gathered}

βj,i\beta_{j,i} показывает, насколько блок обращает внимание на позицию ii, когда вычисляет позицию jj; для каждого jj веса в сумме по ii дают 1. Выход в позиции jj — взвешенная сумма значений h(xi)h(x_i) со всех позиций, возвращённая к CC каналам четвёртой свёрткой 1×11 \times 1 Wv∈RC×CˉW_v \in \mathbb R^{C \times \bar C} и прибавленная ко входу с обучаемым скаляром γ\gamma:

oj=Wv∑i=1Nβj,i h(xi),yj=γ oj+xj.\begin{gathered} o_j = W_v \sum_{i=1}^{N} \beta_{j,i}\, h(x_i),\\ y_j = \gamma\, o_j + x_j . \end{gathered}
Рисунок

Блок self-attention в SAGAN

Поток данных через блок self-attentionСлева карта признаков x. Три свёртки 1×1 W_f, W_g, W_h дают f(x), g(x) и h(x). f(x) и g(x) сходятся в softmax, который даёт карту внимания β размера N на N. β умножается на h(x), результат через W_v даёт o. o умножается на γ и складывается с x, пришедшим по обходной связи сверху; получается выход y.xC × NWfсвёртка 1×1Wgсвёртка 1×1Whсвёртка 1×1f(x)C̄ × Ng(x)C̄ × Nh(x)C̄ × N1softmaxs = f(x)ᵀ g(x)βN × N2×WvoC × N3·γx (обход)+yγo + xв начале γ = 04Поток данных через блок self-attentionСверху карта признаков x. Под ней три свёртки 1×1 и карты f(x), g(x), h(x). f(x) и g(x) сходятся в softmax, который даёт карту внимания β. β умножается на h(x), через W_v получается o; o умножается на γ и складывается с x, пришедшим по обходной связи слева; внизу выход y.xC × NWfсвёртка 1×1Wgсвёртка 1×1Whсвёртка 1×1f(x)C̄ × Ng(x)C̄ × Nh(x)C̄ × N1softmaxβN × N2×WvoC × N3·γx (обход)+yγo + xв начале γ = 04
Зелёное — входная карта признаков x с C каналами и N = H·W позициями; жёлтое — обучаемые свёртки 1×1 и карта внимания β; синее — промежуточные карты; красное — выход блока. Строка j карты β говорит, с какими весами позиция j собирает значения h(x_i) со всех позиций i, сколь угодно далёких. Выход o масштабируется обучаемым скаляром γ и прибавляется к входу; γ начинается с 0, так что в начале обучения блок пропускает x без изменений. В SAGAN C̄ = C/8.
  1. Три свёртки 1×11 \times 1 превращают карту признаков в f(x)f(x), g(x)g(x) и h(x)h(x), каждая с Cˉ\bar C каналами.
  2. Оценки sijs_{ij} образуют карту N×NN \times N; softmax по ii превращает каждую строку в веса внимания.
  3. Каждая позиция собирает значения всех позиций со своими весами; WvW_v восстанавливает CC каналов.
  4. Результат умножается на γ\gamma и прибавляется ко входу.

Игрушечный пример. Три позиции и Cˉ=1\bar C = 1. Пусть f(x)=(2,0,−1)f(x) = (2, 0, -1) по позициям i=1,2,3i = 1, 2, 3; возьмём позицию j=1j = 1 с g(x1)=1g(x_1) = 1. Оценки равны si1=(2,0,−1)s_{i1} = (2, 0, -1), и softmax даёт β1,⋅=(0.844,0.114,0.042)\beta_{1,\cdot} = (0.844, 0.114, 0.042), в сумме 1. При значениях h(x)=(1,3,5)h(x) = (1, 3, 5) позиция 1 собирает 0.844⋅1+0.114⋅3+0.042⋅5=1.3960.844 \cdot 1 + 0.114 \cdot 3 + 0.042 \cdot 5 = 1.396, и WvW_v переводит это в o1o_1. Пока γ=0\gamma = 0, y1=x1y_1 = x_1, каким бы ни было o1o_1.

Почему γ\gamma начинается с 0. При инициализации блок тождественный, поэтому сеть сначала учится на локальных свёрточных признаках, что проще, и подмешивает нелокальные по мере роста γ\gamma (раздел 3). Три отображения узкие: SAGAN берёт Cˉ=C/8\bar C = C/8 для экономии памяти, не увидев значимых потерь для Cˉ=C/k\bar C = C/k при k=1,2,4,8k = 1, 2, 4, 8.

На языке заметки о self-attention g(xj)g(x_j) — запрос (query) позиции jj, f(xi)f(x_i) — ключ (key), а h(xi)h(x_i) — значение (value) позиции ii. Два отличия от версии в Transformer: оценки не делятся на Cˉ\sqrt{\bar C}, а блок — одна голова внутри свёрточной сети, а не стопка слоёв внимания.

Куда его ставить. Карта внимания содержит N2N^2 элементов на изображение: 1 048 5761\,048\,576 на карте признаков 32×3232 \times 32 и 16 777 21616\,777\,216 на 64×6464 \times 64. В экспериментах SAGAN блок стоит на одном уровне сетей, и выбор уровня важен. На больших картах он помогал сильнее: с блоком на 8×88 \times 8 FID был 22.98, на 32×3232 \times 32 — 18.28, на 64×6464 \times 64 — 18.65, против 22.96 без внимания. Остаточные блоки с тем же числом параметров работали хуже, например 27.33 вместо 18.28 на 32×3232 \times 32, так что выигрыш — не просто дополнительная ёмкость (раздел 5.2, таблица 1). На маленьких картах внимание играет роль, похожую на роль локальной свёртки. Выученные карты внимания следуют форме объектов и похожим цветам и текстурам, а не фиксированным окрестностям (рисунок 5).

Спектральная нормализация в обеих сетях#

SN-GAN применял спектральную нормализацию к дискриминатору, чтобы сохранить его липшицевым. SAGAN применяет её и к каждому слою генератора, опираясь на наблюдение, что хорошо обусловленный генератор обучается лучше: она предотвращает разрастание величин параметров и необычные градиенты (раздел 4.1). Для генератора это не требование липшицевости со стороны функции потерь, а стабилизатор. Со спектральной нормализацией в обеих сетях обучение было устойчивым при одном шаге дискриминатора на шаг генератора, тогда как базовая модель со спектральной нормализацией только в DD рано коллапсировала при обновлениях 1:1; статья о SN-GAN избегала этого, делая 5 шагов дискриминатора на шаг генератора (раздел 5.1).

Два шага обучения#

Регуляризованный дискриминатор учится медленно, поэтому WGAN-GP и SN-GAN делают несколько шагов дискриминатора на шаг генератора. SAGAN вместо этого даёт двум сетям разные шаги обучения — это правило обновления с двумя масштабами времени (two time-scale update rule, TTUR) из работы Heusel et al. (2017): 4⋅10−44 \cdot 10^{-4} для DD и 10−410^{-4} для GG, с Adam β1=0\beta_1 = 0, β2=0.9\beta_2 = 0.9 и по одному шагу каждой сети (раздел 5). Один более быстрый шаг дискриминатора дешевле пяти медленных. Эффекты двух изменений складываются (рисунок 4): со спектральной нормализацией в обеих сетях, но без TTUR FID достиг 33.39 через 160 тыс. итераций, а затем ухудшился до 72.41 к 260 тыс.; с TTUR он продолжал улучшаться, до 22.96 через миллион итераций.

Кусочно-линейная функция потерь#

SAGAN обучается с кусочно-линейной версией состязательной функции потерь (hinge loss), которую использовали также SN-GAN и проекционный дискриминатор (projection discriminator; формула 4):

LD=Ex∼P[max⁡(0, 1−D(x))]+Ez∼pz[max⁡(0, 1+D(G(z)))],LG=−Ez∼pz[D(G(z))].\begin{aligned} \mathcal L_D ={}& \mathbb E_{x\sim P}\big[\max(0,\, 1 - D(x))\big]\\ &+ \mathbb E_{z\sim p_z}\big[\max(0,\, 1 + D(G(z)))\big],\\ \mathcal L_G ={}& -\mathbb E_{z\sim p_z}\big[D(G(z))\big]. \end{aligned}

DD выдаёт вещественное число без сигмоиды. От него требуется давать настоящим изображениям значение не меньше 11, а подделкам — не больше −1-1; за этими границами (margins) пример ничего не вносит. Функция потерь генератора такая же, как в WGAN. В условной по классу версии DD также получает метку: D(x,c)D(x, c) на настоящих парах и D(G(z,c),c)D(G(z, c), c) на сгенерированных.

Игрушечный пример. Значения на настоящих изображениях (1.5,0.4)(1.5, 0.4) и на подделках (−2,0.3)(-2, 0.3) дают LD=12(0+0.6)+12(0+1.3)=0.95\mathcal L_D = \tfrac12(0 + 0.6) + \tfrac12(0 + 1.3) = 0.95. Вклад вносят только второе настоящее изображение и вторая подделка, и только они дают градиенты для DD; два других примера уже за своими границами. Функция потерь генератора равна −12(−2+0.3)=0.85-\tfrac12(-2 + 0.3) = 0.85.

В отличие от функции потерь критика WGAN, кусочно-линейная функция перестаёт раздвигать значения, как только они перешли границу, так что дискриминатор ничего не выигрывает, делая свои значения всё больше. Граница 1 имеет смысл только потому, что спектральная нормализация фиксирует масштаб DD; сеть без ограничений могла бы вместо этого перемасштабировать свой выход (Brock, Donahue and Simonyan, 2019, приложение G.2, сноска 3).

Результаты SAGAN на ImageNet#

SAGAN генерирует изображения ImageNet размера 128×128128 \times 128 по заданному классу, с условной батч-нормализацией (conditional batch norm) в генераторе и проекционным дискриминатором; оба приёма описаны в следующем разделе. Метрики считаются по 50 тыс. сэмплов; intra-FID — это FID между сгенерированными и настоящими изображениями одного класса, агрегированный по классам (раздел 5, таблица 2):

Модель, ImageNet 128×128128 \times 128ISintra-FIDFID
AC-GAN28.5260.0—
SN-GAN с проекцией36.892.427.62
SAGAN52.5283.718.65

Выигрыш был наибольшим на классах с геометрической структурой, таких как золотая рыбка и сенбернар, и наименьшим на текстурных классах, таких как долина или каменная стена (рисунок 6), — как и предсказывает рассуждение о свёртках. Каждая модель обучалась около двух недель на 4 GPU.

BigGAN: тот же рецепт в большом масштабе#

BigGAN начинает с SAGAN: та же архитектура ResNet с блоком self-attention, кусочно-линейная функция потерь, спектральная нормализация в обеих сетях, условная батч-нормализация и проекционный дискриминатор (Brock, Donahue and Simonyan, 2019, раздел 3). Он вдвое уменьшает шаги обучения SAGAN, до 2⋅10−42 \cdot 10^{-4} для DD и 5⋅10−55 \cdot 10^{-5} для GG, делает два шага DD на шаг GG, использует Adam с β1=0\beta_1 = 0, β2=0.999\beta_2 = 0.999, ортогональную инициализацию и сэмплирует из скользящего среднего весов GG с коэффициентом затухания 0.9999 (приложение C). Обучение идёт на 128–512 ядрах TPU v3 Pod, а статистики батч-нормализации вычисляются по всем устройствам. Прогрессивный рост оказался не нужен, даже при 512×512512 \times 512.

Метка класса: условная батч-нормализация и проекция#

В генераторе. Слой батч-нормализации стандартизует каждый канал, а затем применяет обучаемые масштаб и сдвиг. Условная батч-нормализация делает масштаб и сдвиг функциями вектора условия uu (de Vries et al., 2017):

CBN(hk;u)=(1+ak⊤u) hk−μkσk2+ε+bk⊤u,\begin{aligned} \mathrm{CBN}(h_k; u) ={}& \big(1 + a_k^\top u\big)\,\frac{h_k - \mu_k}{\sqrt{\sigma_k^2 + \varepsilon}}\\ &+ b_k^\top u, \end{aligned}

для канала kk со статистиками батча μk\mu_k, σk2\sigma_k^2 и обучаемыми векторами aka_k, bkb_k. Коэффициент масштаба центрирован в 1, а сдвиг — в 0 (приложение B). При u=e(c)u = e(c) (BigGAN дописывает к нему часть zz, см. ниже) метка меняет масштаб и сдвиг каждого канала каждого слоя, так что она действует на весь генератор, а не только на его вход. Более ранние модели учили отдельный эмбеддинг класса для каждого слоя; BigGAN использует один общий эмбеддинг (shared embedding) e(c)e(c) и линейно проецирует его в коэффициенты масштаба и сдвиги каждого слоя. Это уменьшило затраты памяти и вычислений и, по словам авторов, повысило «скорость обучения (в числе итераций, нужных для достижения заданного качества) на 37%» (раздел 3).

В дискриминаторе. Условный GAN из заметки о GAN подаёт метку в DD конкатенацией. Проекционный дискриминатор (Miyato and Koyama, 2018) выводит другую форму из оптимального дискриминатора. На парах (x,c)(x, c) оптимальный логит для исходной функции потерь GAN равен логарифму отношения плотностей, который раскладывается по формуле умножения p(x,c)=p(c∣x) p(x)p(x, c) = p(c \mid x)\,p(x):

ln⁡pdata(x,c)pg(x,c)=ln⁡pdata(c∣x)pg(c∣x)+ln⁡pdata(x)pg(x).\begin{aligned} \ln\frac{p_{\text{data}}(x, c)}{p_g(x, c)} ={}& \ln\frac{p_{\text{data}}(c \mid x)}{p_g(c \mid x)}\\ &+ \ln\frac{p_{\text{data}}(x)}{p_g(x)} . \end{aligned}

Предположим, что оба апостериорных распределения классов — softmax-классификаторы на общем векторе признаков ϕ(x)\phi(x), ln⁡p(c∣x)=vc⊤ϕ(x)−ln⁡Z(x)\ln p(c \mid x) = v_c^\top\phi(x) - \ln Z(x), с векторами классов vcdatav_c^{\text{data}} и vcgv_c^{g} и нормировочными множителями ZdataZ_{\text{data}}, ZgZ_g. Тогда первое слагаемое равно (vcdata−vcg)⊤ϕ(x)−ln⁡(Zdata(x)/Zg(x))(v_c^{\text{data}} - v_c^{g})^\top\phi(x) - \ln\big(Z_{\text{data}}(x)/Z_g(x)\big). Нормировочные множители и второе слагаемое зависят только от xx и поглощаются одной скалярной функцией ψ(ϕ(x))\psi(\phi(x)); разность векторов классов становится обучаемым эмбеддингом e(c)e(c) (раздел 3, формулы 2–7):

D(x,c)=ψ(ϕ(x))+e(c)⊤ϕ(x).D(x, c) = \psi\big(\phi(x)\big) + e(c)^\top\phi(x) .

Метка входит через скалярное произведение с признаками изображения: безусловное значение реалистичности плюс слагаемое, показывающее, насколько признаки соответствуют классу. В BigGAN ϕ(x)\phi(x) — последние признаки дискриминатора после ReLU и глобального суммирующего пулинга, а ψ\psi — линейный слой (таблица 4).

Размер батча и ширина#

Первый эксперимент просто увеличивал батч (раздел 3, таблица 1). В таблице также указана итерация, на которой каждая конфигурация коллапсировала, или отмечено, что она оставалась устойчивой 10610^6 итераций; строки 5–8 — средние по 8 запускам с разными начальными значениями (seeds).

Батч, изменениеШирина chПараметрыУстойчиво до / коллапс наISFID
2566481.5 млнустойчиво52.5218.65
5126481.5 млнустойчиво58.7715.30
10246481.5 млнустойчиво63.0314.88
20486481.5 млн732 тыс.76.8512.39
204896173.5 млн295 тыс.92.989.54
+ общий эмбеддинг96160.6 млн185 тыс.94.949.18
+ иерархический латентный вектор96158.3 млн152 тыс.98.768.73
+ ортогональная регуляризация96158.3 млн165 тыс.99.318.51

Батч в восемь раз больше поднял IS на 46%. Авторы предполагают, что больший батч покрывает больше мод и даёт обеим сетям лучшие градиенты. Увеличение ширины на 50%, с 64 до 96 базовых каналов, примерно удвоило число параметров и подняло IS ещё на 21%. Удвоение глубины поначалу не помогло; более поздняя BigGAN-deep сделала глубину выгодной благодаря другому остаточному блоку — с узким местом (bottleneck). Цена видна в столбце коллапса: большие конфигурации достигают лучших метрик за меньшее число итераций, а затем коллапсируют, и их метрики берутся из контрольных точек непосредственно перед коллапсом.

Иерархический латентный вектор#

В DCGAN латентный вектор входит только в первый слой. BigGAN разбивает zz на части, по одной на разрешение, и конкатенирует каждую часть с общим эмбеддингом класса, получая вектор условия uu одного остаточного блока (приложение B). При 128×128128 \times 128 вектор z∈R120z \in \mathbb R^{120} разбивается на шесть 20-мерных частей: одну для входного слоя и по одной для каждого из пяти остаточных блоков генератора. Тогда латентный вектор может напрямую влиять на признаки любого масштаба. Эта связь skip-z дала примерно на 4% лучшее качество и на 18% более быстрое обучение (раздел 3).

Трюк усечения#

GAN никогда не распространяют градиент через априорное распределение, поэтому распределение латентного вектора при сэмплировании не обязано совпадать с распределением при обучении. BigGAN обучается с z∼N(0,I)z \sim \mathcal N(0, I), а сэмплирует из усечённого нормального распределения (truncated normal): каждая координата zz, модуль которой превышает порог τ\tau, сэмплируется заново (resampling), пока не попадёт в [−τ,τ][-\tau, \tau] (раздел 3.1). Тогда каждая координата имеет распределение N(0,1)\mathcal N(0, 1) при условии ∣zk∣≤τ|z_k| \le \tau. Меньшие пороги стягивают zz к моде априорного распределения, а сэмплы — к самым типичным выходам GG: отдельные изображения выглядят лучше, но набор становится менее разнообразным. Заметка о точности и полноте даёт этому компромиссу имя: точность против полноты (precision against recall).

Рисунок

Усечение одной координаты: доля сохранённой массы и разброс

Усечение одной координаты: доля сохранённой массы и разбросГрафик по порогу τ от 0 до 3. Синяя кривая — доля сохранённой массы — растёт от 0 до почти 1. Красная кривая — стандартное отклонение — растёт от 0 до почти 1, оставаясь ниже синей. Жёлтая пунктирная кривая — вероятность сохранить все 120 координат — почти равна нулю до τ около 2 и поднимается до 0.72 при τ = 3. Отмечены точки при τ = 0.5, 1 и 2.00.250.50.75100.511.522.53порог τ0.38290.2840.68270.5400.95450.8800.00370.72P(|z| ≤ τ) сохранённая массаст. отклонение после усечениявсе 120 координат не пересэмплированыУсечение одной координаты: доля сохранённой массы и разбросГрафик по порогу τ от 0 до 3. Синяя кривая — доля сохранённой массы — растёт от 0 до почти 1. Красная кривая — стандартное отклонение — растёт от 0 до почти 1, оставаясь ниже синей. Жёлтая пунктирная кривая — вероятность сохранить все 120 координат — почти равна нулю до τ около 2 и поднимается до 0.72 при τ = 3. Отмечены точки при τ = 0.5, 1 и 2.00.250.50.75100.511.522.53порог τ0.38290.2840.68270.5400.95450.8800.00370.72P(|z| ≤ τ) · сохранённая массаст. отклонение · после усечениявсе 120 координат · не пересэмплированы
Координата z ~ N(0, 1) усекается порогом τ: значения с |z| > τ пересэмплируются, пока не попадут в [−τ, τ]. Синяя кривая — вероятность, что значение сохранится, 2Φ(τ) − 1; красная — стандартное отклонение усечённой координаты (у исходной оно равно 1); жёлтая пунктирная — вероятность, что ни одну из 120 координат латентного вектора BigGAN для 128 × 128 не придётся пересэмплировать. Точки отмечают пороги 0.5, 1 и 2 из иллюстрации усечения в статье BigGAN. При τ = 2 сохраняется 0.9545 массы каждой координаты, но нетронутой остаётся лишь доля 0.0037 векторов; при τ = 0.5 разброс падает до 0.284, а среднеквадратичная норма вектора — с √120 ≈ 10.95 до 3.11.

Разобранный пример. Одна координата сохраняет значение с вероятностью Pr⁡(∣zk∣≤τ)=2Φ(τ)−1\Pr(|z_k| \le \tau) = 2\Phi(\tau) - 1, а после усечения её дисперсия равна 1−2τφ(τ)/(2Φ(τ)−1)1 - 2\tau\varphi(\tau)/(2\Phi(\tau) - 1), где Φ\Phi и φ\varphi — функция распределения и плотность стандартного нормального распределения.

τ\tauсохранённая масса на координатустандартное отклонениеE∥z∥2\sqrt{\mathbb E\lVert z\rVert^2}, d=120d = 120Pr⁡(ни одна координата не сэмплирована заново)\Pr(\text{ни одна координата не сэмплирована заново})
∞\infty1110.951
20.95450.8809.640.0037
10.68270.5405.911.3⋅10−201.3 \cdot 10^{-20}
0.50.38290.2843.119.4⋅10−519.4 \cdot 10^{-51}

Даже мягкий порог τ=2\tau = 2 меняет почти каждый вектор, потому что у 120-мерного zz почти всегда есть координата за пределами 2. При τ=1\tau = 1 векторы настолько короткие, что генератор практически никогда таких не видел: при обучающем априорном распределении ∥z∥2\lVert z\rVert^2 имеет распределение χ2\chi^2 со 120 степенями свободы, у 90% обучающих кодов норма лежит между 9.78 и 12.11, а вероятность нормы меньше 5.91 — около 10−1510^{-15}. Поэтому усечение — это сдвиг входного распределения, и выдержать его может только генератор, который ведёт себя гладко вдали от своих обучающих входов.

Что происходит с метриками. IS и FID можно вычислить для ряда порогов, и тогда получается кривая, а не точка (приложение D, рисунки 17–18). В условных по классу моделях IS не штрафует за нехватку разнообразия внутри класса, поэтому он стабильно растёт по мере уменьшения τ\tau. FID штрафует за потерю разнообразия и одновременно поощряет качество, поэтому сначала немного улучшается, а затем сильно ухудшается, когда τ\tau приближается к 0. Некоторые модели вовсе не выдерживали усечения и давали артефакты насыщения (рисунок 2b) — для этого и нужен следующий регуляризатор.

Ортогональная регуляризация#

Чтобы GG стал настолько гладким, что всё латентное пространство, включая короткие векторы усечения, отображается в хорошие изображения, BigGAN регуляризует веса генератора в сторону ортогональности (orthogonal regularization; раздел 3.1). Для матрицы весов WW, столбцы которой — фильтры, полный штраф за неортогональность равен λ∥W⊤W−I∥F2\lambda\lVert W^\top W - I\rVert_F^2, где ∥⋅∥F\lVert\cdot\rVert_F — норма Фробениуса. Известно, что он слишком ограничителен, поэтому BigGAN отбрасывает диагональ:

Rλ(W)=λ ∥W⊤W⊙(1−I)∥F2,R_\lambda(W) = \lambda\,\big\lVert W^\top W \odot (\mathbf 1 - I)\big\rVert_F^2,

где 1\mathbf 1 — матрица из единиц, а ⊙\odot — поэлементное произведение. На диагонали W⊤WW^\top W стоят квадраты норм фильтров, вне диагонали — их скалярные произведения, так что изменённый штраф подталкивает разные фильтры к ортогональности, не ограничивая их норм. В статье сила штрафа обозначена β\beta и равна 10−410^{-4}.

Игрушечный пример. Два фильтра, w1=(2,0)w_1 = (2, 0) и w2=(0.6,0.8)w_2 = (0.6, 0.8). Тогда у W⊤WW^\top W диагональ (4,1)(4, 1) и внедиагональные элементы 1.21.2. Полный штраф равен (4−1)2+2⋅1.22+0=11.88(4 - 1)^2 + 2 \cdot 1.2^2 + 0 = 11.88, и бо́льшая его часть — требование ∥w1∥=1\lVert w_1\rVert = 1. Изменённый штраф равен 2⋅1.22=2.882 \cdot 1.2^2 = 2.88: ему важно только перекрытие w1⊤w2=1.2w_1^\top w_2 = 1.2.

По всем запускам таблицы 1 усечение выдерживали 16% моделей без ортогональной регуляризации и 60% — с ней (раздел 3.1).

Коллапс обучения в большом масштабе#

Большие конфигурации из таблицы выше хорошо обучаются, а затем внезапно ломаются: за несколько сотен итераций качество сэмплов падает от пикового до наихудшего (приложение G.1). Авторы BigGAN искали причину в сингулярных числах весов (раздел 4).

  • Генератор. У большинства слоёв спектральные нормы ведут себя хорошо, но у нескольких, обычно у первого, полносвязного слоя, наибольшее сингулярное число растёт во время обучения и взрывается в момент коллапса. Его ограничение или регуляризация убирали взрыв, но не предотвращали коллапс (раздел 4.1).
  • Дискриминатор. Его спектры шумные, с выбросами, и скачком меняются при коллапсе. Штраф на градиент с центром в нуле на настоящих данных с рекомендованной силой 10 сделал обучение устойчивым, но снизил IS на 45%; сила 1, наименьшая без внезапного коллапса, всё равно стоила 20%. Дропаут (dropout) и более сильные ортогональные или L2-штрафы на DD вели себя так же: устойчивость покупалась ценой качества (раздел 4.2).
  • Запоминание. Функция потерь DD во время обучения стремится к нулю. При проверке на ImageNet не сколлапсировавшие дискриминаторы правильно относили к настоящим более 98% обучающих изображений, но лишь 50–55% валидационных, что не лучше случайного угадывания. DD запоминает (memorization) обучающую выборку. Ближайшие соседи сгенерированных изображений внутри их классов указывают на то, что GG этого не делает (разделы 4.2 и 5.1).
  • Кто должен побеждать. Заморозка GG перед коллапсом оставляла DD устойчивым; заморозка DD сразу приводила к коллапсу GG, а функция потерь DD подскакивала выше 300 при обычном диапазоне от 0 до 3. Увеличение любого из шагов обучения от контрольной точки вызывало немедленный коллапс (приложение G.1).

Практический итог: устойчивость возникает из взаимодействия двух сетей, а не из какой-то одной из них, и существующие методы могут обеспечить её только ценой большой потери качества. Поэтому BigGAN позволяет коллапсу случиться поздно и использует раннюю остановку, сохраняя контрольную точку до коллапса (раздел 4.3).

Результаты BigGAN на ImageNet#

Поскольку усечение обменивает разнообразие на качество, BigGAN приводит метрики при нескольких порогах (раздел 5.1, таблица 2): без усечения, при пороге с лучшим FID, при пороге, где IS равен IS настоящих валидационных изображений, и при максимальном IS. Настоящие изображения ImageNet размера 128×128128 \times 128 имеют IS 166 на валидационной выборке и 233 на обучающей, на которой обучалась сеть Inception (приложение C.3).

Модель, ImageNet 128×128128 \times 128СэмплированиеISFID
SN-GAN с проекциейбез усечения36.827.62
SAGANбез усечения52.5218.65
BigGANбез усечения98.88.7
BigGANпорог с лучшим FID126.57.7
BigGANIS валидационных данных166.39.6
BigGANмаксимальный IS20625
BigGAN-deepIS валидационных данных166.57.4

Строки BigGAN — средние не менее чем по трём запускам. Последняя строка даёт главный результат статьи: IS 166.5 и FID 7.4 против 52.52 и 18.65 у SAGAN. BigGAN-deep в четыре раза глубже BigGAN, использует остаточные блоки с узким местом и имеет меньше параметров: 50.4 млн в GG и 34.6 млн в DD против 70.4 млн и 88.0 млн (приложение B). При 256×256256 \times 256 BigGAN-deep достиг IS 232.5 и FID 8.1, а при 512×512512 \times 512 — IS 241.5 и FID 11.5. На JFT-300M, наборе данных на два порядка больше ImageNet, те же архитектурные решения помогли, и модели не коллапсировали в течение сотен тысяч итераций (раздел 5.2).

Покрытие мод и оценка качества#

Все три статьи заботились о разнообразии и сочли стандартные меры недостаточными.

  • Разнообразие сэмплов — не близость к данным. MS-SSIM в том виде, как его использует статья о ProGAN, измеряет, насколько сгенерированные изображения отличаются друг от друга. В таблице 1 он оставался около 0.28 от базовой модели (a) до сошедшейся сети (h), хотя изображения заметно улучшились; SWD, который сравнивает фрагменты сгенерированных изображений с фрагментами настоящих, отслеживал улучшение (Karras et al., 2018, раздел 6.1).
  • IS игнорирует разнообразие внутри класса. Условная по классу модель может поднять IS, выдавая одно идеальное изображение на класс (Brock, Donahue and Simonyan, 2019, раздел 3.1). Усечение толкает именно к этому: при максимальном IS у BigGAN FID равен 25, примерно втрое больше, чем без усечения.
  • Один FID смешивает два вида неудач. FID улучшается с ростом качества и ухудшается с потерей разнообразия, поэтому одно его значение — лишь точка на кривой. Привести всю кривую усечения, как делает BigGAN, — честный вариант; точность и полнота позже явно разделяют эти две оси.
  • Метрики по классам. Intra-FID, который использует SAGAN, оценивает каждый класс отдельно и выявляет классы, с которыми модель справляется плохо, — единый FID по всем классам их скрывает. BigGAN наблюдал новые сбои такого рода: утечку классов (class leakage), когда изображение одного класса несёт признаки другого, и в некоторых моделях коллапс к одному-двум сэмплам для нескольких классов, тогда как остальные оставались разнообразными (Brock, Donahue and Simonyan, 2019, раздел 5.1 и приложение H).
  • Проверки на запоминание. Дискриминатор, запоминающий обучающую выборку, ставит вопрос о том, не копирует ли её генератор; ближайшие соседи сэмплов среди обучающих изображений их класса, в пространстве пикселей и в признаках классификатора, — это проверка, которой пользуется BigGAN (Brock, Donahue and Simonyan, 2019, приложение A).

Inception Score, FID и точность и полнота выводятся в разделе курса об оценке качества.

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

DCGANProGANSAGANBigGAN
Данные и выходотдельные домены, 64×6464 \times 64один домен, до 1024×10241024 \times 1024ImageNet, 128×128128 \times 128, условная по классуImageNet до 512×512512 \times 512, условная по классу
Функция потерьненасыщающаяся функция потерь GANWGAN-GP, ncritic=1n_{\text{critic}} = 1, член против дрейфакусочно-линейнаякусочно-линейная
Нормализация в GGбатч-нормализацияпопиксельная нормализация признаков, выравненный шаг обученияспектральная нормализация, условная батч-нормализацияспектральная нормализация, условная батч-нормализация с общим эмбеддингом
Нормализация в DDбатч-нормализациявыравненный шаг обучения; стандартное отклонение по мини-батчуспектральная нормализацияспектральная нормализация
Метка в DD——проекцияпроекция
Дальняя структуратолько свёрткитолько свёрткиодин блок self-attentionодин блок self-attention
ОптимизацияAdam 2⋅10−42 \cdot 10^{-4}, β1=0.5\beta_1 = 0.5Adam 10−310^{-3}, β1=0\beta_1 = 0; разрешение растёт с плавным включениемTTUR: 4⋅10−44 \cdot 10^{-4} для DD, 10−410^{-4} для GG, шаги 1:1батч 2048, ширина ×1.5, два шага DD на шаг GG
Латентный вектор100-мерный равномерный, во входной слой512-мерный на гиперсфереN(0,I)\mathcal N(0, I)N(0,I)\mathcal N(0, I), разбит по блокам; при сэмплировании усекается

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

  • Вычисления. BigGAN обучался на 128–512 ядрах TPU по 24–48 часов на модель, SAGAN — две недели на 4 GPU, модель ProGAN для CelebA-HQ — 4 дня на 8 GPU. Абляционные эксперименты, стоящие за этими рецептами, обошлись гораздо дороже одного запуска.
  • Коллапс отложен, а не устранён. Большие конфигурации BigGAN коллапсируют; средства, которые их стабилизируют, стоят 20–45% IS, поэтому на практике полагаются на раннюю остановку.
  • Усечение — это сдвиг распределения. Его сэмплы приходят не из того распределения, на котором обучался генератор. Ему нужна ортогональная регуляризация, и даже с ней усечение выдерживали только 60% моделей в абляционном эксперименте.
  • Метки и отбор данных. Качество BigGAN зависит от меток классов, на которые опирается генерация; мегапиксельные лица ProGAN зависят от отобранного, выровненного набора данных.
  • Self-attention квадратичен. Карта внимания размера N2N^2 ограничивает блок одним умеренным разрешением в каждой сети.
  • Нет управления. Латентный код — один запутанный (entangled) вектор. Ничто в обучении не связывает направление в zz с одним атрибутом изображения, поэтому сдвиг zz обычно меняет несколько атрибутов сразу.

Итоги#

  • ProGAN выращивает GG и DD от 4×44 \times 4 до 1024×10241024 \times 1024, плавно включая каждый новый блок с весом α\alpha от 0 до 1. Стандартное отклонение по мини-батчу выявляет сколлапсировавшие батчи, выравненный шаг обучения умножает веса на константу Хе во время прохода, так что каждый слой учится с одинаковой относительной скоростью, а попиксельная нормализация держит под контролем величины сигналов в генераторе.
  • SAGAN добавляет блок self-attention, y=γ o+xy = \gamma\,o + x с γ\gamma, инициализированным нулём, так что любые две позиции взаимодействуют за один слой. Он применяет спектральную нормализацию в обеих сетях, использует TTUR (4⋅10−44 \cdot 10^{-4} для DD, 10−410^{-4} для GG) вместо дополнительных шагов дискриминатора и обучается с кусочно-линейной функцией потерь: IS 52.52 и FID 18.65 на ImageNet при 128×128128 \times 128.
  • BigGAN сохраняет рецепт SAGAN и масштабирует его: батч 2048, слои на 50% шире, общий эмбеддинг класса для условной батч-нормализации, латентный вектор, разбитый по блокам генератора, проекционный дискриминатор и ортогональная регуляризация.
  • Трюк усечения сэмплирует заново латентные координаты за порогом и обменивает разнообразие на качество; при пороге, где его IS совпадает с IS настоящих валидационных данных, BigGAN-deep достигает IS 166.5 и FID 7.4.
  • В большом масштабе обучение коллапсирует; дискриминатор запоминает обучающую выборку, а устойчивость можно купить только ценой качества, поэтому модели останавливают рано.

Код#

Блок self-attention на PyTorch и числа из примеров с усечением и выравненным шагом обучения:

import math
import torch, torch.nn as nn
from scipy import stats

# блок self-attention из SAGAN: f, g, h, v — свёртки 1x1, gamma начинается с 0
class SelfAttention(nn.Module):
    def __init__(self, C, k=8):
        super().__init__()
        self.f, self.g, self.h = (nn.Conv2d(C, C // k, 1) for _ in range(3))
        self.v = nn.Conv2d(C // k, C, 1)
        self.gamma = nn.Parameter(torch.zeros(1))

    def forward(self, x):                                   # x: (B, C, H, W)
        B, C, H, W = x.shape
        f, g, h = (m(x).flatten(2) for m in (self.f, self.g, self.h))   # (B, C/k, N)
        s = f.transpose(1, 2) @ g                           # s[b, i, j] = f(x_i)^T g(x_j), (B, N, N)
        beta = s.softmax(dim=1)                             # нормировка по i для каждого j
        o = self.v((h @ beta).view(B, -1, H, W))            # o_j = W_v sum_i beta[i, j] h(x_i)
        return self.gamma * o + x, beta

torch.manual_seed(0)
block = SelfAttention(C=64)
x = torch.randn(2, 64, 16, 16)
y, beta = block(x)
print(tuple(beta.shape), torch.allclose(beta.sum(dim=1), torch.ones(2, 256)), torch.equal(y, x))
# (2, 256, 256) True True

# трюк усечения, одна координата: сохранённая масса, стандартное отклонение; 120-мерный вектор: P(не тронут),
# RMS-норма и P(||z|| <= этой нормы) при обучающем априорном N(0, I), где ||z||^2 — хи-квадрат со 120 степенями свободы
for tau in (0.5, 1.0, 2.0):
    kept = 2 * stats.norm.cdf(tau) - 1
    var = 1 - 2 * tau * stats.norm.pdf(tau) / kept
    print(tau, round(kept, 4), round(math.sqrt(var), 3), f"{kept ** 120:.1e}",
          round(math.sqrt(120 * var), 2), f"{stats.chi2(120).cdf(120 * var):.0e}")
# 0.5 0.3829 0.284 9.4e-51 3.11 1e-43 | 1.0 0.6827 0.54 1.3e-20 5.91 1e-15 | 2.0 0.9545 0.88 3.7e-03 9.64 3e-02
print(round(math.sqrt(stats.chi2(120).ppf(0.05)), 2), round(math.sqrt(stats.chi2(120).ppf(0.95)), 2))   # 9.78 12.11

# выравненный шаг обучения: один шаг Adam 1e-3 относительно масштаба Хе sqrt(2 / n_in) у двух слоёв
for n_in in (3 * 3 * 512, 3):
    print(n_in, round(1e-3 / math.sqrt(2 / n_in), 4))     # 4608 0.048 | 3 0.0012

В блоке beta[b, i, j] — это βj,i\beta_{j,i} из текста, а torch.equal(y, x) подтверждает, что при γ=0\gamma = 0 блок в точности тождественный.

Дальше#

BigGAN получает качество за счёт масштаба и метки класса, но почти не даёт управлять тем, что он рисует. Метка выбирает категорию; всё остальное решает один запутанный вектор zz. Его части попадают в разные блоки, но ничто не связывает направление в zz с одним атрибутом, так что сдвиг zz в любом направлении меняет сразу много атрибутов изображения. Следующая заметка перестраивает генератор вокруг этой проблемы: StyleGAN отображает zz в промежуточное латентное пространство и внедряет его как «стиль» на каждом разрешении, что разделяет грубые атрибуты, такие как поза, и тонкие, такие как текстура.

Где посмотреть изображения#

Эти модели лучше всего судить по их сэмплам. В каждой строке сказано, где их найти; номера рисунков относятся к указанной версии на arXiv.

  • ProGAN (Karras et al., 2018; arXiv v3): лица CelebA-HQ размера 1024×10241024 \times 1024 на рисунке 5 и ещё на рисунке 11; категории LSUN при 256×256256 \times 256 на рисунке 7; ближайшие соседи из обучающей выборки на рисунке 10; репозиторий с видео результатов и сгенерированными изображениями.
  • SAGAN (Zhang et al., 2019; arXiv v2): сэмплы ImageNet при 128×128128 \times 128, по одному классу в строке, на рисунке 6; карты внимания на рисунках 1 и 5; сравнение базовой модели со спектральной нормализацией и TTUR на рисунке 4.
  • BigGAN (Brock, Donahue and Simonyan, 2019; arXiv v2): условные по классу сэмплы на рисунке 1; усиливающееся усечение на рисунке 2(a); сэмплы при 256×256256 \times 256 и 512×512512 \times 512 на рисунках 5 и 6; лёгкие и трудные классы на рисунке 7; интерполяции на рисунках 8 и 9.

Литература#