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

Рисунок 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.

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

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

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

LOSS

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

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

$L = D_{\text{CLIP}}(G(w), t_{\text{target}})$

где:

- изображение, cгенерированное генератором по латентному вектору (в пространстве , полученному из );
$t_{\text{target}}$ - текстовое описание нового домена;
$D_{\text{CLIP}}$ - косинусное расстояние между векторами сгенерированного изображения и текста в пространстве CLIP-эмбеддингов.

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

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

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

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

Directional CLIP loss

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

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

Рисунок 2. Инициализируем два предварительно обученных генератора, один из них замораживаем, а веса второго изменяются путем оптимизации, тем самым сдвигая домен в соответствии с заданным пользователем текстовым направлением.

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

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

$\Delta T = E_T(t_{\text{target}}) - E_T(t_{\text{source}})$

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

$\Delta I = E_I(G_{\text{train}}(w)) - E_I(G_{\text{frozen}}(w))$

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

$L_{\text{direction}} = 1 - \frac{\Delta I \cdot \Delta T}{|\Delta I||\Delta T|}$

где:

- текстовый энкодер CLIP;
- визуальный энкодер CLIP;
- латентный вектор из пространства .

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

Рисунок 3. Изображения встраиваются в CLIP-пространство, необходимо, чтобы вектор, соединяющий их, был параллелен вектору, соединяющему тексты.

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

Layer-Freezing

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

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

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

Сэмплируются несколько латентных векторов из пространства и, проходя через mapping network, преобразуются в латентные вектора из пространства .
Латентные вектора преобразуются в пространство , где у каждого слоя свой вектор .
Вектора оптимизируются без изменения весов сети. Здесь как раз используется в качестве функции-потерь $L = D_{\text{CLIP}}(G(w), t_{\text{target}})$ , о которой мы говорили ранее.
Определяются вектора , которые в процессе оптимизации изменились сильнее всего, это говорит о том, какие слои считаются наиболее важными.
Размораживаются k этих наиболее релевантных слоев.
Проводится один шаг их оптимизации с использованием directional CLIP loss.

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

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

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

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

Latent-Mapper

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

Для решения этой проблемы авторы взяли готовую технику из StyleCLIP. Latent-Mapper - это нейросеть, которой на вход подается латентный вектор из пространства , а на выходе получается . После этого генератор работает уже с и дает картинку уже более похожую на целевой домен, например кошку. Latent-Mapper получается «вытягивает» нужный домен из «расширенного» пространства генератора. Он обучается так, чтобы менять латентный вектор в нужном семантическом направлении с помощью текста. Функция потерь здесь тот же CLIP-loss, который описывался ранее, а именно: $L = D_{\text{CLIP}}(G(w), t_{\text{target}})$ .

Few-Shot

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

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

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

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

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

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

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

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

Заключение

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

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

зависимость от CLIP: промпты сильно влияют на результат;
не всегда хорошо передаёт очень сложные или абстрактные стили;
если $\Delta w$ (разница между доменами) слишком большая - могут появляться артефакты или разрушаться исходные черты изображений.

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

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

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

Рисунок 4. Сравнение вне-доменных манипуляций со StyleCLIP. В каждой строке показано: сгенерированное с помощью случайно выбранного кода изображение (слева), результаты редактирования того же кода в направлении текстового запроса вне домена с использованием методов StyleCLIP, а также изображение, полученное для того же кода с генератором, преобразованным методом StyleGAN-NADA (справа).

Ссылки

Оригинальная статья от специалистов из Tel Aviv University и NVIDIA Research: StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators.
Репозиторий от авторов: https://github.com/rinongal/StyleGAN-nada.
Оригинальная статья CLIP: Learning Transferable Visual Models From Natural Language Supervision.
Репозиторий CLIP: https://github.com/openai/CLIP?utm_source=chatgpt.com.
Оригинальная статья StyleGAN: A Style-Based Generator Architecture for Generative Adversarial Networks.
Репозиторий StyleGAN: https://github.com/NVlabs/stylegan.
Оригинальная статья StyleGAN2: Analyzing and Improving the Image Quality of StyleGAN.
Репозиторий StyleGAN2: https://github.com/NVlabs/stylegan2.

Литературный обзор на статью: StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators