Большому художнику – большой критик: новые версии модели ruCLIP для ранжирования результатов ruDALL-E (и не только)

В ​2021 году особенно активно развивались языковые и визуальные трансформеры, что привело к появлению новых разработок в сфере мультимодального анализа данных. Ярким примером стала модель синтеза изображений по текстовым описаниям DALL-E от OpenAI. Команды Сбера - Sber AI и SberDevices - успешно ворвались в эту гонку трансформеров и обучили модели: ruDALL-E XL, которая доступна на GitHub, эксклюзивную ruDALL-E XXL (доступна в DataHub SberCloud ML Space), ruDALL-E Emojich для синтеза эмоджи, первую версию мультимодальной мультитаск архитектуры Fusion Brain, а с введением понятия гипермодальности создали универсальную модель ruDOLPH (про которую обязательно расскажем подробнее в одной из следующих статей).

Мы также не могли обойти стороной другое полезное приложение языковых и визуальных трансформеров, – модель CLIP от OpenAI, задача которой – определять “смысловую” близость текстов на естественном (английском) языке и изображений. Успешное обучение модели командой OpenAI уже позволило с уверенностью говорить о возможности решения таких задач компьютерного зрения, как, например, zero-shot классификация или zero-shot детекция объектов. В итоге мы поставили перед собой цель качественно “русифицировать” CLIP, чтобы модель можно было использовать напрямую на русском языке и обходить применение переводчиков в связке с англоязычной моделью CLIP. Первые шаги в этом направлении были сделаны нами ещё в начале 2021 года - тогда мы представляли модель ruCLIP Small (Russian Contrastive Language – Image Pre-training). Сейчас  же мы хотим рассказать об ещё 6 версиях модели ruCLIP, которые мы обучили за это время (для более детального погружения в терминологию и суть проблемы советуем перечитать нашу прошлую статью по ruCLIP Small).

В релиз вошли следующие версии модели ruCLIP, которые отличаются количеством слоев энкодера ViT, размером использованного патча (14×14, 16×16, 32×32) и размерами входных изображений — 224×224, 336×336 и 384×384 пикселя, соответственно. Семантика названия моделей выглядит следующим образом:

• ruCLIP Base [vit-base-patch16-224]

• ruCLIP Base [vit-base-patch32-224]

• ruCLIP Base [vit-base-patch32-384] 

• ruCLIP Large [vit-large-patch14-224]

• ruCLIP Base [vit-base-patch16-384] exclusive ☁️DataHub SberCloud only ️ 

• ruCLIP Large [vit-large-patch14-336] exclusive ☁️DataHub SberCloud only

Первые 4 модели доступны для открытого пользования на GitHub, две из них - в хабе предобученных моделей, датасетов и контейнеров – DataHub SberCloud ML Space, а последние две, – самые большие можно скачать только через DataHub.

Все новые версии ruCLIP были обучены на открытых данных, собранных из Рунета, а также на данных экосистемы Сбера. Всего удалось собрать около 240 млн уникальных пар “изображение-описание на естественном языке”. Модель мы обучали в течение 12 дней на платформе SberCloud ML Space и суперкомпьютере Christofari Neo, используя распределенное обучение на 256 GPU Tesla A100, что существенно превосходит ресурсы, затраченные в ходе первого обучения ruCLIP Small. Важно подчеркнуть, что в новых версиях при обучении использовался батч размера 32768, что соответствует тому, как обучалась оригинальная модель CLIP (в версии ruCLIP Small размер батча был равен 16).

В основе новых версий ruCLIP по-прежнему лежат две составляющие:

1. Image Encoder — энкодер, формирующий векторное представление изображений. В основе нашей модели остался хорошо известный ViT.

2. Text Encoder — энкодер, формирующий векторное представление текстовых описаний. В отличие от модели ruCLIP Small, мы не стали использовать RuGPT3Small, а взяли текстовый энкодер - трансформер со следующими параметрами:

• для Base версий
  

  • Context Length: 77

  • Transformer Layers: 12

  • Transformer Width: 512

  • Transformer Heads: 8

• для Large версий
  

  • Context Length: 77

  • Transformer Layers: 12

  • Transformer Width: 768

  • Transformer Heads: 12

Данные

Как мы уже упоминали ранее, для обучения ruCLIP использовался датасет, состоящий из 240М пар “текстовое описание-изображение” и представляющий из себя расширенную версию обучающего датасета ruDALL-E. Следует отметить, что наш датасет является одним из самых больших русскоязычных датасетов, содержащих пары “текстовое описание-изображение”, но на фоне существующих англоязычных конкурентов наш сет далек по объёму от лидеров списка, что можно видеть в таблице 1.

Таблица 1 - Сравнение объемов обучающих выборок для различных моделей

Лидерство по объемам данных принадлежит Google с их моделями ALIGN и BASIC. В первой из них данные собраны на основе Conceptual Captions с упрощённой фильтрацией данных. Во второй модели набор данных для ALIGN дополняется приватным датасетом JFT-3B.

Особенности обучения

Процесс обучения был довольно интересный, и о его этапах хочется рассказать подробнее. Начнем с модели, основанной на ViT-Large. На первом этапе проводилось обучение ruCLIP Large [vit-large-patch14-224] с нуля на основе собранного датасета. Обучение этой модели выполнялось в течение 380К итераций. Далее, изменив размер входных данных и проведя обучение в течение еще 15К итераций была получена модель ruCLIP Large [vit-large-patch14-336] exclusive. Что касается ViT-Base версий модели, то они обучались путем тюнинга английской модели в течение 140К итераций. В результате были получены модели ruCLIP Base [vit-base-patch16-224] и ruCLIP Base [vit-base-patch32-224]. Изменив размер входных изображений до 384 и выполнив тюнинг в течение еще 20К итераций были получены модели ruCLIP Base [vit-base-patch32-384] и ruCLIP Base [vit-base-patch16-384] exclusive.

Наглядно процесс обучения представлен ниже на графиках изменения значений loss (рисунок 1). Также для каждого процесса обучения показана методика изменения значений learning rate.

На рисунке 2 показан график изменения среднего значения accuracy для всех датасетов в процессе обучения модели ruCLIP Large [vit-large-patch14-224].

Пример работы

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

import ruclip
import torch
import requests
from PIL import Image
from io import BytesIO
 
url_dog = "https://ichef.bbci.co.uk/news/640/cpsprodpb/475B/production/_98776281_gettyimages-521697453.jpg"
response = requests.get(url_dog)
image = Image.open(BytesIO(response.content))
 
device = 'cuda'
clip, processor = ruclip.load('ruclip-vit-base-patch32-384', device=device)
classes = ['кошка', 'собачка', 'мышка', 'машина', 'стол', 'дом', 'жидкость']
predictor = ruclip.Predictor(clip, processor, device, bs=8)
 
with torch.no_grad():
    text_latents = predictor.get_text_latents(classes)
    pred_labels = predictor.run([image], text_latents)
 
print(f'На изображении {classes[pred_labels[0]]}')

--------------------------------------------------
На изображении собачка

В результате получили, что с близкой к наибольшей вероятностью изображение было отнесено к классу “собачка”, что не может не радовать. Дальше мы покажем результаты работы нашей самой большой доступной модели ruCLIP Large [vit-large-patch14-224] на данных из разных доменов и продемонстрируем ее жизнеспособность в задачах zero-shot и few-shot классификации.

В случае, если захочется воспользоваться эксклюзивными версиями модели ruCLIP из DataHub SberCloud ML Space, то с инструкциями можно ознакомиться по ссылкам ruCLIP Base [vit-base-patch16-384] exclusive и ruCLIP Large [vit-large-patch14-336] exclusive.

Оценка качества. Сравнение моделей 

В ходе экспериментов мы преследовали несколько целей. Первая – посмотреть насколько хороши разные модели ruCLIP на данных различной природы, а вторая – сравнить качество нашего решения с использованием комбинации переводчика OPUS-MT ru-en и оригинальной модели CLIP. Цифры ниже все скажут за себя.

Про качество

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

Для демонстрации снова возьмем те же самые изображения, которые использовались в ноутбуке оригинальной модели CLIP и в статье про ruCLIP Small. Среди наших моделей возьмем самую большую открытую модель ruCLIP Large [vit-large-patch14-224] и проверим ее качество. Косинусная близость между текстами и изображениями для модели ruCLIP Large [vit-large-patch14-224] также показана ниже на рисунке 4.

Предсказания топ 5 классов для изображений с помощью ruCLIP Large  [vit-large-patch14-224] показаны на рисунке 5.

В таблице 2 приведено сравнение реализаций модели ruCLIP в рамках задачи zero-shot классификации для датасетов разных доменов. Символом ? выделено лучшее значение метрики для каждого из датасетов. В таблице представлены результаты сравнения первой обученной модели ruCLIP Small, всех 6 реализаций модели ruCLIP, модели CLIP original с переводчиком OPUS-MT ru-en и модели CLIP original. В 8 из 18 датасетов модель ruCLIP Large [vit-large-patch14-336] exclusive показывает лучший результат среди всех сравниваемых, а в тех, где проигрывает, разница не очень значительная (пожалуй, кроме MNIST и FGVC Aircraft - здесь еще есть над чем потрудиться). Стоит отметить, что на датасете MNIST даже самым большим моделям ALIGN и BASIC не удается обойти CLIP по качеству. Для большинства датасетов (11 из 18) обученные нами модели обходят по качеству модель CLIP [vit-base-patch16-224] original + OPUS-MT. Преимущество модели CLIP original (правый столбец) можно объяснить существенно большим объемом обучающей выборки (400М vs 240М).

Таблица 2 - Значения точности zero-shot классификации для различных моделей

Ниже на рисунках 6 и 7 показаны результаты экспериментальных исследований zero-shot и few-shot классификации 10 моделями на 16 датасетах. Так же, как и в оригинальной статье, на признаках, которые извлекает CLIP для изображений, были обучены линейные классификаторы с использованием 1, 2, 4, 8 и 16 изображений для каждого класса. Это достаточно простой и понятный подход transfer learning, называемый linear probe, при котором обученная модель используется для извлечения признаков, и далее доучивается “голова” для необходимого количества классов при ограниченном наборе изображений для каждого класса. Поскольку признаки, которые извлекаются у моделей CLIP original и CLIP original + OPUS-MT одинаковые — для CLIP original + OPUS-MT нет отдельного графика few-shot классификации. Также мы посчитали усредненный few-shot график для лучшей модели ruCLIP Large exclusive без учета трех датасетов - PCam, Oxford Pets и FGVC Aircraft, на которых модель проигрывает достаточно сильно, и можно видеть (пунктирная линия), что среднее качество даже немного превосходит ruCLIP Small. Звездочками показаны значения средней zero-shot оценки моделей.

Отдельно для каждого датасета картина выглядит следующим образом (рисунок 7). 

Почему же ruCLIP Small имеет такие хорошие показатели качества в результате few-shot классификации? Такой результат обусловлен тем, что при реализации ruCLIP Small использовался Image Encoder, состоящий из замороженных весов модели CLIP original и двух дополнительных линейных слоев, которые и подвергались обучению. Ввиду того, что для linear probe эксперимента используется только визуальная составляющая модели, качество ruCLIP Small должно быть очень близко к CLIP original. Посмотреть детальное сравнение всех 6 обученных моделей можно в нашем репозитории на GitHub.

В таблице 3 приведено сравнение значений метрики few-shot классификации для восьми моделей на разных датасетах. В составе сравнительной базы также выступают модели ruCLIP Small и CLIP original. И снова преимущество модели CLIP original (правый столбец) можно объяснить существенно большим объемом обучающей выборки (400М vs 240М).

Таблица 3 - Усредненные значения качества few-shot классификации для различных моделей

В дополнение мы еще посчитали корреляцию значений качества zero-shot и few-shot классификаций для разных моделей. Значение коэффициента корреляции нашей лучшей модели ruCLIP Large [vit-large-patch14-336] exclusive ближе других к значению оригинальной модели CLIP (рисунок 8), что говорит о хорошей обобщающей способности при переходе от zero-shot к few-shot задаче.

Про скорость

Мы не ограничились лишь оценкой качественных характеристик реализаций модели ruCLIP, но и оценили их производительность на датасете CIFAR100 с использованием Nvidia-V100. Результаты оценки инференса реализаций показаны в таблице 4.

Таблица 4 - Оценка скорости работы (в iter/sec) реализаций модели ruCLIP на датасете CIFAR100

Что получили в итоге?

Нам удалось обучить несколько разных версий ruCLIP, которые на ряде датасетов успешно обошли оригинальную англоязычную модель CLIP с русско-английским переводчиком. Все обучение строилось на самостоятельно собранном датасете из 240 млн пар и заняло 12 дней на 256 Tesla GPU A100 суперкомпьютера Christofari Neo. Обширные исследования на датасетах из различных доменов показали применимость ruCLIP в задачах zero-shot и few-shot классификации. Это безусловно можно назвать успешным переносом на русский язык модели CLIP (еще один пример работы модели ruCLIP на рисунке 9). 

В задачах zero-shot и few-shot классификации версии модели ruCLIP проигрывают оригинальной модели CLIP в связи с тем, что для обучения использовался датасет меньшего объема. Следует также отметить тенденцию к росту качества моделей с ростом объема обучающих выборок, что подтверждается преимуществом ALIGN и BASIC над CLIP ввиду существенно большего объема использованных данных.

Успешное обучение ruCLIP и доступность моделей в open source позволит эффективно решать многие задачи компьютерного зрения в различных продуктах и сервисах в режиме «zero-shot», то есть без необходимости дорогостоящего дообучения. Нам бы очень хотелось, чтобы все заинтересованные читатели загрузили себе доступные модели, потестировали их в режиме zero-shot и предложили свои идеи по дальнейшему развитию/применению проекта. В дополнение на платформе SberCloud ML Space для первых 100 клиентов-юридических лиц доступен грант в 180 000 руб. на тестовый доступ и дообучение моделей, а в модуле DataHub размещены open source и эксклюзивные реализации модели ruCLIP.

Коллектив авторов: Алексей Шоненков, Андрей Кузнецов, Денис Димитров, Татьяна Шаврина, Даниил Чесаков, Анастасия Мальцева, Алёна Феногенова, Игорь Павлов, Антон Емельянов, Сергей Марков, Дарья Бакшандаева, Вера Шибаева, Андрей Черток