CLIP (Contrastive Language-Image Pre-Training)

State-of-the-Art модели для CV предобучают на датасетах с фиксированным количеством лейблов. Вспомним Imagenet: в нем больше 21-й тысячи классов с лейблами “water snake”, “triceratops” и многими другими. Создание таких датасетов — очень трудоемкий процесс. А в NLP сейчас для предобучения языковой модели данные не размечают. Берут огромные наборы текстов из интернета и обучают прямо на них. Вот бы претрейн CV моделей можно было бы делать как в NLP…

Так и появилась модель CLIP — Contrastive Language-Image Pre-training. Это модель, которая претренировалась на огромном датасете из пар (Картинка, Текст).

Датасет

Итак, мы не хотим размечать кучу данных для претрейна. Как тогда нам поступить? Можно спарсить картинки из интернета и учиться на них с помощью self-supervised методов (например, поворачивать картинку на n градусов и пытаться предсказать этот n). Очень часто вместе с картинкой идет какой-то текст (описание этой картинки). Такой текст может сообщить полезную информацию об объектах на картинке.

Авторы CLIP создали датасет из 400 миллионов пар <картинка>:<текст>. Они сконструировали около 500 тысяч запросов, с помощью которых постарались максимально покрыть все многообразие визуальных концепций. Каждый запрос аугментировали с использованием синонимов. На каждый из запросов приходилось около 20 тысяч пар <картинка>:<текст>.  Полученный набор данных имеет такое же общее количество слов, как и набор данных WebText, использованный для обучения GPT-2.

Итак, у нас есть датасет с 400 миллионами пар <картинка>:<текст>. Давайте теперь разберемся, как можно обучить претрейн на таком наборе.

Модель

Модель содержить энкодер для текстов, энкодер для картинок и projection head.

Для текстов они использовали GPT2, а для картинок было два варианта:

• модифицированный ResNet;
• модифицированный ViT.

Схематически итоговая модель выглядит так:

Давайте подробнее разберем блоки с моделями.

Модифицированный ResNet

Эту модель изменили больше всего:

• в самом первом блоке до res-блоков теперь 3 сверточных слоя (а не 1), плюс AvgPool вместо MaxPool;
• antialiasing convolutions: добавляем AvgPool перед свертками со страйдом > 1;
• вместо GlobalAvgPool в конце Multihead Attention.

Вот так выглядит вся модель:

Теперь давайте разберем все слои. Изменения в архитектуре выделены зеленым цветом. Начнем со Stem:

А сейчас посмотрим на слои Layer 1, 2, 3 и 4. По сути это Bottleneck-слои из резнета. Они отличаются только количеством фильтров. В этот блок также были внесены изменения:

Далее мы перейдем к самому интересному слою: Multihead Attention вместо GlobalAvgPool. Давайте разберемся, откуда появляются Query, Key и Value на последнем слое:

Итак, мы собрали Query, Key и Value. Затем мы объединяем их с матрицей позиционных эмбеддингов (для Query суммируем только с первой строкой из матрицы позиционных эмбеддингов), а потом умножаем на матрицы проекций (более подробно об этом можно почитать в статье про attention). Рассмотрим теперь, что происходит дальше:

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

Модифицированный ViT

ViT остался прежним, за исключением добавления LayerNorm после блока Linear Projection. Если хотите подробнее ознакомиться с этой архитектурой — можете прочитать нашу статью про ViT 🙂

Очень коротко про GPT2

Это генеративная модель, которая обучалась предсказывать следующий токен в последовательности. Например, если мы на вход GPT2 подадим текст “Берлин — это столица”, то GPT2 продолжит это предложение как “Берлин — это столица Германии”.

Для тех, кто знаком с архитектурой ViT, тут не будет ничего нового. Мы разбиваем предложения на токены, токены подаем в слой эмбеддингов, складываем эмбеддинги со строками матрицы Positional-encoding и подаем на вход трансформеру. Трансформер выглядит так же, как трансформер в ViT, но с одним отличием: в self-attention добавлена маска. Это выглядит примерно таким образом:

Зачем нам нужна маска? Модель генеративная, следовательно, во время инференса мы не можем смотреть “в будущее”, т.е. смотреть на идущие после текущего токены. Например:

Для чего маска конкретно в этой модели? Ведь нам нужны только эмбеддинги, мы не учим модель предсказывать следующий токен. Это сделано для того, чтобы текстовый энкодер можно было инициализировать весами GPT2. Но в итоге авторы CLIP все равно учили текстовый энкодер с нуля)))

А какой именно эмбеддинг мы получаем из модели? Как вы помните, на выходе из трансформера мы имеем эмбеддинги всех токенов, поданных на вход. Перед подачей текста в модель мы в начало и в конец предложения добавляем токены [SOS] — “Start of sequence” и [EOS] — “End of sequence”. Из предложения “Кот прыгнул на шкаф” мы делаем “[SOS]Кот прыгнул на шкаф[EOS]”. Так вот, в качестве эмбеддинга мы берем эмбеддинг токена [EOS]:

Projection Head

Projection Head отвечает за прием эмбеддингов изображения и текста и перевод их в одинаковое пространство. Этот блок — линейный слой, который переводит эмбеддинг из энкодеров в эмбеддинг необходимой для расчета лосса размерности.

Лосс

Картиночная и текстовая модели “встречаются” только во время подсчета лосса. Покажем еще раз картинку из начала нашего поста:

По сути мы даем двум сетям набор картинок и текстов и в каждом батче просим найти именно ту картинку, которая подходит к конкретному тексту. Авторы статьи приложили псевдокод для расчета лосса:

Итак, для расчета лосса мы:

1. нормализуем эмбеддинги;
2. считаем similarity-матрицу скалярных произведений эмбеддингов и умножаем каждый элемент матрицы на температуру;
3. применяем кросс-энтропию по столбцам и строкам similarity-матрицы;
4. складываем эти 2 лосса и делим сумму на 2.

Что такое температура и зачем она нужна? Это обучаемый параметр, который авторы CLIP использовали как способ регулировки “силы” Softmax. Softmax имеет тенденцию усиливать различия между оценками, что может привести к очень уверенным, но не правильным предсказаниям. Для решения этой проблемы авторы CLIP ввели температуру, которую они использовали с целью «смягчения» Softmax. При умножении оценок на температуру вероятности становятся более «разнесенными» и менее чувствительными к небольшим различиям. Это может помочь предотвратить чрезмерную уверенность и повысить точность модели.

Лирическое отступление

Почему бы не учить одну сеть, которая будет на вход получать картинку, а на выходе выдавать текст?

Авторы пробовали так делать. После обучения нескольких вариантов они померили zero-shot качество — оно было не самым лучшим. Затем авторы попытались предсказывать не весь текст, а мешок слов. Мешок слов — это формулирование из текста “Кошка прыгнула на стол” списка слов вида [кошка, прыгнуть, стол]. Качество улучшилось. А потом авторы поменяли задачу с predictive на contrastive, и качество выросло еще больше.

Можно посмотреть на это с такой стороны: описания картинок из интернета иногда очень сложно восстановить при помощи одной картинки. Восстановить только ключевые слова без предлогов — задача более простая. Но еще проще было бы соотнести описания и картинки и найти между ними соответствие.

Такой график авторы CLIP приводят в статье:

Как учили

Учили CLIP с нуля без инициализации картиночного или текстового энкодера какими-либо предобученными весами.

Учили 32 эпохи с помощью оптимизатора Adam с регуляризацией weight decay и косинусным шедулером для learning rate. Отличительная особенность — большие батчи размера 32768. CLIP c самой большой ResNet моделью учился 18 дней на 592 V100 GPU. CLIP c самой большой ViT моделью учился 12 дней на 256 V100 GPU.

Zero-shot

В статье подробно исследуются возможности CLIP успешно работать в zero-shot режиме, то есть на незнакомых датасетах.

Но как из модели, умеющей соотносить текст и картинку, сделать классификатор, который будет различать, например, котов и собак?

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

Ансамбли

Мы можем пойти дальше и заняться prompt-инжинирингом! На наших картинках могут быть большие и маленькие собаки и кошки. Давайте заведем эмбеддинги и этих предложений:

Такой подход авторы статьи называют ансамблем.

Linear probing

Для оценки качества предобученных моделей часто используют подход linear probing. Авторы CLIP используют этот метод для оценки качества эмбеддингов картиночного энкодера.

Суть метода: мы добавляем линейный классификатор к выходам модели, замораживаем все веса модели и учим веса только этого линейного классификатора. Такой метод не требует большого количества данных для обучения.

С ним можно выполнять 1-shot learning, 2-shot learning, 3-shot learning и так далее. Это подходы, когда для обучения модели используются только по 1, 2, 3 и так далее изображений на класс.

Любопытно, что zero-shot CLIP сравним по качеству с 4-shot CLIP. Эта оценка была получена на 20 датасетах.

Почему мы вообще выбрали эту статью для обзора?

1. Это просто очень сильный бекбон, который выдает очень качественные эмбеддинги. Эта модель используется в бекбонах DALLE и Stable-diffusion. С нее в целом начался бум text2image моделей.
2. Если у вас есть большой пул картинок из вашего домена и вы собрались предиктить какой-то новый класс, вам нужно отобрать картинок для разметки. Подход майнить их отсматривая кучу картинок — долгий. Но мы можем проиндексировать нашу базу картинок CV-моделькой и получить эмбеддинги. Потом нагенерировать промптов типа “кошка”, “рыжая кошка”, “кошка с собакой”, сделать запрос к проиндексированным эмбеддингам и получить данные, которые уже можно отправить на разметку. Сэкономим кучу ручного труда.
3. На своих задачах можно дотюнить CV-часть при помощи небольшого числа промптов и картинок и качество будет уже неплохим.

Модель также выложена в open-source. Рекомендуем ознакомиться с кодом, он довольно понятно написан.