StyleGAN: A Style-Based Generator Architecture for Generative Adversarial Networks

🔗 Ссылки

https://arxiv.org/abs/1812.04948

https://github.com/NVlabs/stylegan — осторожно, внутри TF1 🚩🚩🚩

https://github.com/huangzh13/StyleGAN.pytorch — неофициальная реализация на PyTorch

https://github.com/NVlabs/ffhq-dataset — датасет FFHQ

💎 Контрибьюшн

StyleGAN — одна из первых архитектур, в которой латентный вектор подаётся не напрямую в свёрточный генератор, а через AdaIN'ы. Для того, чтобы это взлетело, авторам понадобилось реализовать несколько интересных идей:

• Нелинейный маппинг (по сути MLP) Z->W — это позволяет unsupervised развязывать признаки в латентном пространстве, что приводит к более гладкой интерполяции
• Аддитивный шум перед AdaIN'ами — чтобы обеспечить стохастичность деталей, в то время как общая суть картинки неизменна
• Style mixing — смешиваем 2 стиля при генерации, "this prevents the network from assuming that adjacent styles are correlated"

Помимо этого в статье предложены 2 автоматические метрики хорошести латентных пространств:

1. Perceptual Path Length — насколько сильно картинка меняется при незначительных изменениях латентного вектора
2. Linear Separability — насколько хорошо работает линейный классификатор (SVM), обученный на латентных векторах из W

В итоге, хотя авторы и не заявляют о SoTA в задаче генерации, они предложили принципиально новый рабочий фреймворк, который послужил фундаментом для дальнейшего развития области и в частности работ StyleGAN2 и StyleGAN3.

🛠 Задача

Главное, что стоит понимать — StyleGAN решает задачу генерации картинок, а не style-transfer'а. То есть на входе нет ничего, а хочется получить реалистичную картинку из определённого домена: лицо (FFHQ), фотографию комнаты (LSUN), машину (LSUN Car), котика (LSUN Cat), etc.

В реальности же вводится латентное пространство Z, из которого сэмплируется вектор и подаётся на вход в генератор. Это нужно для разнообразия синтезируемых картинок — сеть обуславливается на вектор из Z и генерирует разные детали.

🔎 Детали

Статья во многом опирается на Progressive Growing of GANs, предыдущую статью того же автора, поэтому я тезисно повторю главные идеи оттуда:

1. G и D зеркально противоположны
2. На вход в G подаётся латентный вектор из Z
3. В процессе обучения сначала обучаются coarse-слои, затем по расписанию добавляются fine-слои, доходя до требуемого разрешения — 1024x1024
4. Minibatch Discrimination — дискриминатор смотрит на статистики по батчу, чтобы добиваться вариативности генератора
5. Не используется BatchNorm, вместо него:
• PixelNorm — вектор в каждом пикселе имеет норму 1
• Equalized LR — разные слои учатся с разным LR, реализовано вот так

Теперь можно разобрать, что же изменили в StyleGAN'е.

1. Mapping Network

Пространство Z имеет стандартное нормальное распределение по каждой из компонент. Из него легко сэмплировать, но в то же время его структура не отражает структуру признаков в target-домене — оно слишком равномерное.

Поэтому авторы добавили нелинейный маппинг Z->W в надежде, что полученное латентное пространство W будет иметь хорошую структуру. По сути это просто 8-слойный MLP, все FC-слои которого имеют размерность 512.

Мотивационная картинка:

2. AdaIN

AdaIN используется для передачи стиля в генератор. В статье говорится о том, что пространственно-независимые статистики, такие как матрица Грэма или поканальное матожидание / дисперсия, характеризуют глобальное содержание (стиль) изображения. А фичи, зависящие от координаты, кодируют локальные признаки [Demystifying Neural Style Transfer].

Поэтому для каждого свёрточного слоя вектор из W пропускается через свой FC-слой (на схеме: квадратики с буквой A — affine), определяя параметры демодуляции каналов.

3. Noise

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

Как это работает:

• Для каждого свёрточного слоя сэмплируется 1-канальная картинка, каждый пиксель из стандартного нормального распределения
• Обучаемый модуль (B на схеме) домножает эту картинку на своё число для каждого канала фичемапы
• Полученный тензор добавляется к фичемапе

Зачем вообще нужен аддитивный шум?

Многие части картинки могут быть стохастическими, не влияя на наше восприятие. Например, положение отдельных волос, текстура кожи, веснушки, родинки. Без добавочного шума сети пришлось бы изобретать способ, как генерировать случайные числа для каждого пикселя в каждой фичемапе. Это уменьшает capacity сети, и в то же время не удаётся избежать паттернов в генерируемом сигнале.

Почему сеть не опирается на шум при генерации глобального содержания?

Шум добавляется независимо для каждого пикселя, поэтому если сеть обусловится на шум при генерации, скажем, позы, это приведёт к пространственным коллизиям, которые оштрафует дискриминатор.

Локализация эффекта от добавления шума внутри сети

Авторы предполагают, что сеть генерирует новый контент как только это представляется возможным. Наиболее простой способ добавить стохастичности — опереться на шум в текущем слое. Поскольку для каждого слоя добавляется новый шум, генератору нет нужды протаскивать стохастичность из предыдущих слоёв.

4. Style Mixing

Чтобы стабилизировать обучение, применяется смешение стилей:

1. Сэмплируются 2 латентных вектора
2. В генераторе выбирается точка скрещивания
3. До этой точки через AdaIN прокидывается один стиль, а после — другой

Это позволяет сети избавиться от скоррелированности признаков. Например, если в датасете у смуглых людей всегда есть родинки, то сеть может обучиться на то, что смуглость обязательно влечёт наличие родинок. А благодаря смешению стилей можно зафорсить генерацию смуглого человека без родинок, которое будет реалистичным.

Style mixing применяется к определённой доле генерируемых сэмплов, в статье сравнивались разные варианты — 0%, 50% и 90%.

Зависимость стиля от точки скрещивания

Если скрещивание происходит на coarse-слоях, то от первого латентного вектора прокидываются высокоуровневые признаки — пол, возраст, наличие очков. Если же на fine-слоях, то прокидываются текстурные признаки — цвет волос, кожи, фона.

5. Truncation Trick in W

Для визуализаций GAN‘ов часто используется не сэмплированный вектор из Z, а его приближенная к центру масс версия:

Если ψ в интервале (0, 1), можно получить более качественные изображения, т.к. вблизи центра масс латентное пространство обучено лучше. Или можно получить анти-картинку, если положить ψ = -1. Этот подход отлично работает для нормального распределения, т.к. его центр масс в нуле.

Чтобы пользоваться этим приёмом в хорошем пространстве W, нужно знать центр масс распределения. К счастью его легко оценить по Монте-Карло:

После этого можно делать клёвые интерполяции:

6. Обучение

Слои обучаются последовательно, так же как в Progressive Growing of GANs. Используется такая же архитектура дискриминатора и параметры оптимизаторов.

Для FFHQ авторы используют non-saturating loss + R1 regularization вместо WGAN-GP, который использовался во всех других случаях. Они заметили, что с R1 метрика FID падает дольше, поэтому количество обучающих картинок увеличили 12M->25M.

Для обучения сети потребовалась 1 неделя на 8xV100.

🔬 Эксперименты

Авторы посвятили большую часть исследований теме disentanglement — развязанности (не знаю, как это нормально перевести) латентного пространства.

Развязанное латентное пространство состоит из линейных подпространств, каждое из которых контролирует изменение одного определённого признака у генерируемого объекта.

В статье предполагается, что из развязанного пространства проще синтезировать реалистичные изображения. Именно поэтому в архитектуру добавили нелинейный маппинг Z->W — чтобы линеаризировать латентное пространство относительно признаков.

Чтобы понять, действительно ли их сеть выучивает достаточно линеаризованное латентное пространство, авторы предложили 2 метрики, для которых не требуется наличие энкодера.

1. Perceptual Path Length

Представим ситуацию: есть 2 латентных вектора z1 и z2. Им соответствуют синтезированные изображения G(z1) и G(z2). Если пространство Z достаточно линейно, то при интерполяции между векторами полученная картинка будет меняться плавно. И наоборот, если в какой-то промежуточной точке появляются признаки, которых нет на концах отрезка — это говорит об entangled-пространстве.

Предлагается измерить, насколько значительно изменяется синтезируемая картинка при интерполяции между z1 и z2. В качестве метрики d — VGG Loss:

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

2. Linear Separability

Если латентное пространство достаточно линейно, то для бинарных признаков (пол, тёмный цвет кожи, наличие серёжек, наличие очков и т.д.) латентные коды должны хорошо разделяться гиперплоскостью.

Датасет CelebA-HQ содержит 40 бинарных атрибутов, поэтому авторы обучили на нём 40 бинарных классификаторов. Дальше алгоритм следующий:

1. Сгенерировать 200k изображений генератором
2. Отобрать 100k с наибольшим confidence по каждому признаку (сеть уверена в наличии или отсутствии признака на изображении)
3. Зафитить SVM на этой выборке
4. Вычислить условную энтропию H(Y|X), Y — метка класса, полученная от классификатора (аналог GT), X — предикт от SVM
5. Linear Separability score:

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

Результаты

🐈 Картинки с котами (по ржать)