Новый подход к CLIP guidance для задачи backlit image enhancement.

Тут я расскажу о своей статье, которая опубликована на ECCV-2024. Она называется RAVE: Residual Vector Embedding for CLIP-Guided Backlit Image Enhancement. Это статья о новом способе использовать CLIP guidance для обучения модели улучшения засвеченных картинок (backlit image enhancement). На картинке выше — пример засвеченных и восстановленных картинок. Тут надо, правда, сказать, что идею из статьи можно использовать не только для задачи backlit image enhancement, но и для других задач, связанных с улучшением картинок. Но в статье эксперименты ставились именно на этой задаче, поэтому называется статья именно так.

Сразу ссылки:

Project page

Paper

GitHub (буду рада звездочкам🌝)

И небольшой тизер: по сравнению с SOTA методом CLIP-LIT, который RAVE улучшает, RAVE выдает картинки лучшего визуального качества без артефактов. Ниже пара примеров. А еще вектор, который строится в RAVE, интерпретируем! Об этом всем подробно — ниже в этой статье.

Описание метода

Идея RAVE основана на другой статье (не моей), которая была представлена на ICCV-2023. Это статья Iterative Prompt Learning for Unsupervised Backlit Image Enhancement. Она тоже про то, как использовать CLIP guidance для задачи backlit image enhancement. Метод из этой статьи называется CLIP-LIT. Я прочитала эту статью и у меня возник вопрос: а зачем так сложно? Можно же проще. И, как оказалось, еще и лучше по качеству. Так и родилась моя статья.

Сейчас кратко разберу, в чем идея подхода CLIP-LIT. И затем перейдем к двум идеям его улучшения, которые и вошли в мою статью.

CLIP-LIT

Итак, как работает CLIP-LIT.

Давайте мы будем учить UNet (enhancement network на картинке), который принимает на вход засвеченную (backlit) картинку, а выдает на выходе восстановленную (enhanced). Если у нас в обучающих данных есть пары (backlit img, well-lit img), которые соответствуют дргу другу, то мы могли бы учить этот UNet втупую, подавая ему backlit картинку на вход и сравнивая выход с well-lit картиной. Но давайте представим, что у нас в тренировочных данных нет таких пар. Есть кучка каких-то backlit картинок и кучка well-lit картинок, которые могут иметь совершенно другую семантику. Назовем этот сетап unpaired data.

В этом случае давайте учить UNet с помощью двух лоссов:

• identity loss, который не позволяет картинке сильно далеко отходить от изначальной (т.е. помогает сохранить семантику нетронутой)
• CLIP-Guided Loss, который заставляет UNet делать из backlit картинки well-lit картинку.

Вот в  CLIP-Guided Loss и вся соль. Авторы CLIP-LIT строят его так:

• Давайте заведем два псевдо-промпта: positive и negative. На основе тренировочных backlit и well-lit картинок будем учить эти промпты так, чтобы CLIP-эмбеддинг positive промпта был близок по косинусному расстоянию к CLIP-эмбеддингам well-lit картинок, а CLIP-эмбеддинг negative промпта был близок к CLIP-эмбеддингам backlit картинок.
• Теперь используем выученные эмбеддинги для обучения UNet. Делаем это так: подаем в UNet backlit картинку, получаем enhanced картинку на выходе. Получаем CLIP-эмбеддинг enhanced картинки. Сравниваем его по косинусному расстоянию с эмбеддингами выученных positive и negative промптов. Требуем, чтобы расстояние до positive промпта было ближе, чем расстояние до negative промпта.

Но это еще не все: после этих двух шагов результаты получаются такие себе. Поэтому авторы предлагают сделать несколько итераций этих двух пунктов. Т.е. после первичного обучения промптов и UNet вернуться к дообучению псевдо-промптов, используя картинки из датасета и картинки, которые UNet выдает на данном этапе обучения. Там аж 4 лосс-функции, которые используются для более оптимального дообучения этих промптов. И чтобы получить UNet, который будет хорошо восстанавливать засвеченные картинки, нужно где-то 10-12 итераций этих двух пунктов. Вот иллюстрация к этим итерациям дообучения:

Сразу можно сказать, что обучение CLIP-LIT довольно нестабильное: от итерации к итерации улучшенные картинки получаются то пересвеченные, то обратно слишком темные. Скорее всего, это потому, что от итерации к итерации меняются псевдо-промпты, а их обучение довольно нестабильно и неточно. Неточно в том плане, что при таком виде обучения псевдо-промпты никогда не будут идеально отражать смыслы backlit и well-lit картинок. Из-за этой нестабильности и требуется так много итераций обучения, прежде чем получится нормальный результат.

Вот такая идея у CLIP-LIT. Теперь перейдем к тому, что в этой идее странного и как ее улучшить. Идеи будут две — CLIP-LIT-Latent, и, собственно, RAVE. Идея CLIP-LIT-Latent очень простая и тупая, а RAVE уже поинтереснее.

Начнем с первой.

CLIP-LIT-Latent

Давайте посмотрим еще раз на то, что делает CLIP-LIT, и подумаем: а зачем нам псевдо-промпты? Они же нужны, чтобы сравнивать их CLIP-эмбеддинги с эмбеддингами картинок. Почему бы тогда не учить что-то прямо внутри латентного пространства CLIP?

То есть, первая идея упрощения подхода приходит такая: давайте вместо псевдо-промптов учить сразу псевдо-векторы внутри CLIP embedding space. Т.е. типа сразу выучивать эмбеддинги этих промптов. Тогда мы сможем вообще выкинуть CLIP text encoder и не делать лишние движения типа "прогоним псевдо-промпты через CLIP text encoder, чтобы получить их эмбеддинги". Назовем этот подход CLIP-LIT-Latent.

Общая идея, получается, такая: зачем учить что-то, что дает хорошие фичи при прогоне через какое-то сложное преобразование, когда можно учить эти фичи сразу напрямую? Под сложным преобразованием здесь понимается CLIP Text Encoder.

Вот картинка со сравнением CLIP-LIT и CLIP-LIT-Latent:

CLIP-LIT-Latent дает такие же по качеству результаты, что и оригинальный CLIP-LIT (как по метрикам, так и по визуальной оценке получаемых картинок). Обучение только занимает чуть-чуть меньше времени, потому что нам больше не надо кучу раз проталкивать псевдо-промпты через CLIP Text Encoder. Заметим, что для получения хороших результатов в CLIP-LIT-Latent тоже нужно делать итерации дообучения латентных векторов и UNet, и обучение тоже не очень стабильное.

Получается, CLIP-LIT-Latent просто чуть упрощает модель, сохраняя все остальные ее свойства.

Перейдем теперь ко второй идее улучшения CLIP Guidance, которую мы назвали RAVE.

RAVE

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

RAVE использует это свойство латентного пространства CLIP. В RAVE мы не будем обучать никакие псевдоним-векторы или псевдо-промпты. Вместо это сделаем вот что:

• Возьмем backlit картинки из обучающего датасета. Получим их CLIP-эмбеддинги. Усредним, получим один усредненный CLIP-эмбеддинг backlit картинок v_backlit;
• Возьмем well-lit картинки из обучающего датасета. Получим их CLIP-эмбеддинги. Усредним, получим один усредненный CLIP-эмбеддинг well-lit картинок v_well-lit;
• Отнимем от вектора v_well-lit вектор v_backlit. Получим вектор v_residual. Этот вектор в латентном пространстве CLIP будет иметь смысл перехода от backlit к well-lit изображениям. Действительно, этот вектор содержит в себе те смыслы, что есть в well-lit картинках, но нет в backlit картинках. А это и есть смысл "backlit -> well-lit".

В итоге формула получения v_residual такая:

f_norm здесь — L2-нормализация вектора. Ф_image — CLIP Image Encoder. N_p — количество well-lit картинок в обучающих данных, N_n — количество backlit картинок в обучающих данных.

Проиллюстрировать это можно так:

Этот вектор v_residual мы и будем использовать для обучения UNet. Guidance с помощью v_residual устроен следующим образом: получаем картинку, которую выдает UNet, получаем ее CLIP-эмбеддинг v_enhanced. Лосс-функция заставляет косинусое расстояние между v_enhanced и v_residualбыть таким же, как и косинусое расстояние между v_well-lit и v_residual.

Иными словами, мы хотим, чтобы картинка, которую выдает UNet, имела CLIP-эмбеддинг, близкий к CLIP-эмббедингам well-lit картинок. Мы "двигаем" поданную на вход backlit картинку по пространству эмбеддингов CLIP вдоль вектора v_residual от части пространства, где находятся эмбеддинги backlit картинок, к части пространства, где находятся эмбеддинги well-lit картинок.

Общая иллюстрация метода тогда выглядит так:

Преимущества RAVE:

• Не надо обучать промпты/векторы, только обучаем UNet;
• Не нужно делать много итераций обучения, обучение сходится очень быстро;
• Обучение более стабильно (и, как следствие, быстрее) за счет того, что CLIP guidance устроен на основе вектора, который имеет достаточно четкий смысл, а не на основе псевдо-векторов или псевдо-промптов, которые неоптимально обучаются;
• Визуальная оценка результатов и метрики показывают, что RAVE работает лучше, чем CLIP-LIT и CLIP-LIT-Latent.

Результаты

Таблица со сравнением подходов по классическим метрикам для image enhancement (PSNR, LPIPS, SSIM, FID). Результаты для двух вариантов обучающих данных: paired (где для каждой backlit картинки есть соответствующая ей well-lit картинка) и unpaired (где backlit и well-lit картинки могут быть совершенно разные). В случае paired данных в таблице больше моделей для сравнения, чем в случае unpaired: это потому, что некоторые подходы просто нельзя обучить на unpaired данных.

Видно, что в случае paired данных RAVE показывает лучшие результаты. В случае же unpaired данных у обычного RAVE результаты смешанные, но вот вариант RAVE shifted впереди по большинству метрик. Давайте поясню, почему так и что такое RAVE shifted.

Интерпретация v_residual

Давайте попробуем проинтерпретировать смысл вектора v_residual, который мы строим. Благо у CLIP совместное латентное пространство для картинок и текстов. Что мы можем сделать: давайте переберем все токены из словаря CLIP Text Encoder, и найдем те, эмбеддинги которых по косинусом расстоянию ближе всего и дальше всего от v_residual. Получим вот что:

Видно, что если считать v_residual на основе paired данных, то этот вектор очень хорошо отражает смысл перехода от backlit к well-lit картинкам: токены с самыми далекими от v_residual эмбеддингами имеют смысл вокруг "dark" — solhouette, darkness и т.п., а токены с самыми близкими к v_residual эмбеддингами не имеют какой-то общей семантики. Поэтому этот вектор так хорош при обучении UNet, и в случае paired данных результаты получаются хорошими.

А вот в случае unpaired данных у v_residual появляются какие-то странные смыслы. Токены с самыми далекими от v_residual эмбеддингами имеют смысл вокруг Азии, а токены с самыми близкими к v_residual эмбеддингами имеют смысл вокруг природы. Это произошло потому, что в unpaired обучающих данных backlit картинки были в основном картинками азиатских городов и людей, а на well-lit картинках была природа. Поэтому при вычислении v_residual семантика картинок не до конца схлопнулась, оставив ненужные смыслы. Поэтому в этом случае v_residual не отражает идеального направления в латентном пространстве CLIP, куда нужно двигать backlit картинки, и результаты RAVE в случае unpaired данных получаются чуть хуже (хотя все равно SOTA).

Но тут есть простой хак, как это исправить. Давайте возьмем токены, эмбеддинги которых наиболее близки и наиболее далеки от v_residual, и удалим их смысл из v_residual. Сделать это можно, например, так: возьмем 15 токенов, эмбеддинги которых наиболее близки к v_residual, и 15 токенов, эмбеддинги которых наиболее далеки от v_residual, и построим на их основе еще один вектор v_add_residual:

И дальше отнимем этот вектор от изначального v_residual:

Теперь если мы похожим образом проинтерпретируем новый v_residual, то увидим, что ненужная семантика действительно исчезла, и вектор стал более похож по смыслу на тот, что получался в случае paired дата:

Ну и результаты (как метрики, так и визуальное сравнение) показывают, что RAVE-shifted действительно работает лучше.

Вот так вот. Ну и ниже немного картинок со сравнениями подходов.

Визуальное сравнение

Больше примеров можно найти на странице проекта и в статье.