StyleGAN-NADA (No Annotation Domain Adaptation) - метод, разработанный специалистами из Tel Aviv University и NVIDIA Research, позволяющий адаптировать генеративную модель (StyleGAN2) к новому домену без единого изображения, использующий только семантическую силу больших моделей, предварительно обученных контрастивным методом на тексте и изображениях (СLIP). На рисунке 1 представлены возможности метода StyleGAN-NADA.

Что было до?

Для решения задачи генерации изображений по текстовому описанию ранее применяли следующие подходы:

• оптимизация латентного представления изображения, поступающего на вход генератора;

• обучение генератора с нуля;

• fine-tuning - дообучение предобученного генератора.

Первый подход имеет следующие недостатки:

• при оптимизации под текст могут меняться незапланированные детали;

• изменения возможны в рамках домена, на котором обучен генератор;

• правильный выбор исходного латентного кода влияет на итоговый результат;

• появление артефактов при генерации.

Второй подход требует большого количества изображений, их предобработки и ресурсов. 

Третий метод рискует привести к переобучению и mode collapse.

Авторы StyleGAN-NADA предложили свой уникальный метод, который не требует ни одного изображения из целевого домена, а также предлагает адаптивный выбор обучаемых слоев - с учетом состояния сети и целевого текста.

Архитектура

В основе метода лежат два компонента - CLIP и StyleGAN2. Далее рассмотрим каждый из них по отдельности. 

CLIP

CLIP - это двухмодальная нейросеть, разработанная компанией OpenAI в 2021 году. Состоит из двух энкодеров:

• image encoder - преобразует изображение в вектор;

• text encoder - преобразует текст в вектор.

Оба вектора нормализуются по L2 норме до единичной длины, чтобы использовать косинусное расстояние. 

Модель контрастивным способом обучалась на 400 миллионах пар «изображения - текст», как двусторонний сопоставитель:

• картинка → текст;

• текст → картинка. 

Цель - научить сопоставлять тексты картинке так, чтобы «правильный» текст был ближе к изображению, а все остальные - дальше.

Модель не обучалась классифицировать пары, вместо этого она училась понимать семантику, что например: «кот лежит на диване» ближе к «кот лежит на кресле», чем «машина на трассе».

По сути, CLIP создает единое смысловое пространство для текстов и изображений. 

Подробнее о CLIP вы можете почитать в оригинальной статье на arxiv.org: Learning Transferable Visual Models From Natural Language Supervision и посмотреть их репозиторий на GitHub: CLIP. 

StyleGAN2

Говоря про StyleGAN2, для начала стоит вспомнить первый StyleGAN. Это генеративная модель, разработанная специалистами из NVIDIA в 2018 году, она сильно изменила подход к генерации изображений, введя модель стилизации через промежуточное латентное пространство и механизм нормализации AdaIN. StyleGAN2 - это улучшение первой модели, разработанное теми же специалистами в 2020 году. Он устраняет артефакты, делает генерацию более стабильной и реалистичной.

StyleGAN2 состоит из следующих блоков:

• mapping network - 8-слойный MLP, преобразует латентный вектор из пространства в латентный вектор из пространства ;

• генератор - включает в себя операции свертки, модуляцию и демодуляцию их весов (вместо AdaIN), также инъекцию шума. 

В StyleGAN2 генератор всегда начинается с фиксированной обучаемой константы размером 4×4×512. Разнообразие изображений обеспечивается не самим входом, а модуляцией сверточных фильтров латентным кодом из пространства (через weight demodulation) и добавлением стохастического шума.

StyleGAN2 обучался на чистых, хорошо выровненных, однотипных изображениях. Основные датасеты - FFHQ, LSUN (cars, cats, bedrooms, churches), MetFaces (лица, сцены, животные). Принцип обучения, как и в классическом GAN: Дискриминатор обучается отличать реальные изображения от подделок, а Генератор пытается обмануть Дискриминатор, только архитектура и loss в StyleGAN2 более сложные. 

Также стоит отметить про инверсию - это процесс поиска такого латентного вектора, который при генерации максимально приближает изображение к исходному. Существует несколько способов инверсии:

• оптимизация: градиентный спуск по -вектору (или -вектору) для минимизации ошибки между исходным и сгенерированным изображениями;

• encoder-based: обучается нейросеть, предсказывающая латентный вектор из изображения;

• гибридный: энкодер + оптимизация для доработки результата.

Стоит отметить, что все перечисленные методы инверсии также применимы для получения латентного вектора целевого изображения при работе с StyleGAN-NADA. 

Подробнее о StyleGAN, StyleGAN2 вы можете почитать в оригинальных статьях на arxiv.org: A Style-Based Generator Architecture for Generative Adversarial Networks, Analyzing and Improving the Image Quality of StyleGAN и посмотреть их репозиторий на GitHub: stylegan, stylegan2. 

Метод StyleGAN-NADA

Авторы предложили следующую идею: сдвинуть предобученный генератор из исходного домена в целевой, который описан только с помощью текстового промпта, без использования изображений. В качестве источника «обучающего сигнала» они используют только предобученную модель CLIP. 

Возникает два вопроса, которые предстоит решать:

1. Как извлечь семантическую информацию, содержащуюся в CLIP?

2. Как наилучшим образом оптимизировать генератор, чтобы он не просто учился обманывать CLIP, а создавал качественные, разнообразные изображения?

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

LOSS

Наивный подход

Первое, что приходит на ум для решения этой задачи - использовать в качестве функции-потерь косинусное расстояние, а именно:

где:

• - изображение, cгенерированное генератором по латентному вектору (в пространстве , полученному из );

• - текстовое описание нового домена;

• - косинусное расстояние между векторами сгенерированного изображения и текста в пространстве CLIP-эмбеддингов.

Задача состоит в минимизации этого расстояния. 

Как оказалось, такой loss не дает качественного результата, а именно, он приводит к «адверсариальным» результатам:

• оптимизация «читерит», добавляя мелкие пиксельные искажения на исходное изображение, чтобы приблизить его к текстовому описанию;

• mode collapse - одно изображение может идеально соответствовать текстовому описанию и генератору, нет смысла создавать разнообразные картинки. 

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

Directional CLIP loss

Чтобы справиться с выше описанными проблемами, авторы используют идею «глобального направления». Ее суть в том, чтобы определить направление в CLIP-пространстве, которое соответствует переходу от исходного домена к целевому. Например: «dog» → «cat» - это вектор направления в CLIP-пространстве текстов. Далее следует обучить генератор так, чтобы изображения отличались от исходных только в этом направлении.

Этот подход реализуется следующим образом: берется генератор , обученный на каком-то домене и клонируются, итого получается G_frozen - не обучается и выдает исходное изображение для любого латентного вектора, G_train - обучается выдавать измененное изображение, отличающееся от исходного только в нужном направлении. Важно отметить, оба генератора используют один и тот же mapping network, который не обучается, это обеспечивает совместимость латентных пространств. Схема обучения показана на рисунке 2.

Теперь посмотрим на сам directional loss, он состоит из следующих частей:

• СLIP-вектор текстового направления:

• СLIP-вектор визуального направления (между изображением после и до fine-tuning):

Итоговая функция-потерь - это косинусное расстояние между этими векторами:

где:

• - текстовый энкодер CLIP;

• - визуальный энкодер CLIP;

• - латентный вектор из пространства . 

Иллюстрация направлений показана на рисунке 3.

Выше описанный loss обеспечивает более разнообразные выходные изображения, то есть борется с mode collapse, потому что если генератор будет выдавать одно и то же изображение, то CLIP-вектора будут разными для разных исходов и не смогут согласовываться с текстовым направлением. Также решается проблема обмана CLIP путем добавления мелких пиксельных искажений, так как теперь необходимо, чтобы изменения между двумя изображениями совпадали с направлением в CLIP-пространстве между двумя текстами, а это уже сложнее, особенно для множества пар векторов из пространства и созданных обучаемым генератором.

Layer-Freezing

Еще одна особенность, которую заметили авторы, заключается в адаптивном изменении весов обучаемого генератора. Интуиция заключается в следующем:

• некоторые слои генератора уже подходят для нового домена, и их нет смысла изменять;

• изменяя меньше количество весом, снижается риск переобучения. 

Разберем пошаговый алгоритм, как это работает:

1. Сэмплируются несколько латентных векторов из пространства и, проходя через mapping network, преобразуются в латентные вектора из пространства .

2. Латентные вектора преобразуются в пространство , где у каждого слоя свой вектор . 

3. Вектора оптимизируются без изменения весов сети. Здесь как раз используется в качестве функции-потерь , о которой мы говорили ранее. 

4. Определяются вектора , которые в процессе оптимизации изменились сильнее всего, это говорит о том, какие слои считаются наиболее важными.

5. Размораживаются k этих наиболее релевантных слоев.

6. Проводится один шаг их оптимизации с использованием directional CLIP loss.

Плюсы такого подхода:

• результаты генерации становятся более реалистичными;

• уменьшается риск переобучения;

• адаптация применяется на каждом шаге обучения, что оптимизирует слои по мере изменения модели. 

Дополнительные методы для улучшения качества

В этом разделе я расскажу еще про два метода, которые предложили авторы для улучшения качества результата. 

Latent-Mapper

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

Для решения этой проблемы авторы взяли готовую технику из StyleCLIP. Latent-Mapper - это нейросеть, которой на вход подается латентный вектор из пространства , а на выходе получается . После этого генератор работает уже с и дает картинку уже более похожую на целевой домен, например кошку. Latent-Mapper получается «вытягивает» нужный домен из «расширенного» пространства генератора. Он обучается так, чтобы менять латентный вектор в нужном семантическом направлении с помощью текста. Функция потерь здесь тот же CLIP-loss, который описывался ранее, а именно: . 

Few-Shot

Хоть главная «фишка» StyleGAN-NADA, это переход в целевой домен без использования изображений из него, стоит отметить, что идея направленной loss-функции применима и к картинкам.

Не углубляясь в подробности, идея та же, что и с текстом. Функция потерь строится на использовании направления в СLIP-пространстве между изображениями, созданными исходным генератором, и изображениями, представляющими целевой домен. 

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

Латентное пространство

Подходя к концу обзора, стоит сказать пару слов о существующих техниках инверсии и редактировании изображений.

У генераторов типа StyleGAN латентное пространство устроено таким образом, что позволяет изменениями в нем осмысленно изменять изображения (возраст, выражение лица, прическу и т.д.). Это делает возможным:

• редактировать изображения;

• переводить изображения в изображение.

Касаемо метода StyleGAN-NADA, после его применения для адаптации генератора к новому домену, структура латентного пространства сохраняется, что делает возможным применение тех же осмысленных изменений изображений через латентное пространство, а также использование существующих методов инверсии для работы с конкретным изображением. 

Заключение

В заключение скажу, что метод StyleGAN-NADA в свое время стал мощным прорывом в генерации изображений, позволяющим всего за несколько минут обучения, используя только лишь текст, перевести изображение в нужный домен. 

У данного подхода также имеются и ограничения:

• зависимость от CLIP: промпты сильно влияют на результат;

• не всегда хорошо передаёт очень сложные или абстрактные стили;

• если (разница между доменами) слишком большая - могут появляться артефакты или разрушаться исходные черты изображений.

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

Благодарю за внимание к обзору! В дальнейших публикациях я планирую подробно разобрать реализацию StyleGAN-NADA, а также провести собственные исследования по его улучшению.

Исходный код и дополнительные материалы будут доступны в моём GitHub-репозитории: https://github.com/KBaron34.

Ссылки

1. Оригинальная статья от специалистов из Tel Aviv University и NVIDIA Research: StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators.

2. Репозиторий от авторов: https://github.com/rinongal/StyleGAN-nada.

3. Оригинальная статья CLIP: Learning Transferable Visual Models From Natural Language Supervision. 

4. Репозиторий CLIP: https://github.com/openai/CLIP?utm_source=chatgpt.com.

5. Оригинальная статья StyleGAN: A Style-Based Generator Architecture for Generative Adversarial Networks. 

6. Репозиторий StyleGAN: https://github.com/NVlabs/stylegan.

7. Оригинальная статья StyleGAN2: Analyzing and Improving the Image Quality of StyleGAN. 

8. Репозиторий StyleGAN2: https://github.com/NVlabs/stylegan2.