Большому художнику – большой критик: новые версии модели 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 - Сравнение объемов обучающих выборок для различных моделей

Модель	ruCLIP Small	ruCLIP	CLIP	ALIGN	BASIC
Количество пар	3M	240M	400M	1.8B	6.6B

Лидерство по объемам данных принадлежит 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.

Рисунок 1 - Графики изменения loss в процессе обучения (левый столбец) и процесс варьирования learning rate (правый столбец) для 6 реализаций модели ruCLIP

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

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

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

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

Рисунок 3 - Веселый щеночек

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.

Рисунок 4 - Матрица близости изображений и текстовых описаний

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

Рисунок 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 классификации для различных моделей

	ruCLIP Small [rugpt3-small]	ruCLIP Base [vit-base-patch32-224]	ruCLIP Base [vit-base-patch16-224]	ruCLIP Large [vit-large-patch14-224]	ruCLIP Base [vit-base-patch32-384]	ruCLIP Large [vit-large-patch14-336] exclusive	ruCLIP Base [vit-base-patch16-384] exclusive	CLIP [vit-base-patch16-224] original + OPUS-MT	CLIP [vit-base-patch16-224] original
Food101, acc	0.137	0.505	0.552	0.597	0.642	0.712 ?	0.689	0.664	0.883
CIFAR10, acc	0.808	0.818	0.810	0.878	0.862	0.906 ?	0.845	0.859	0.893
CIFAR100, acc	0.440	0.504	0.496	0.511	0.529	0.591	0.569	0.603 ?	0.647
Birdsnap, acc	0.036	0.115	0.117	0.172	0.161	0.213 ?	0.195	0.126	0.396
SUN397, acc	0.036	0.452	0.462	0.484	0.510	0.523 ?	0.521	0.447	0.631
Stanford Cars, acc	0.023	0.433	0.487	0.559	0.572	0.659 ?	0.626	0.567	0.638
DTD, acc	0.169	0.380	0.401	0.370	0.390	0.408	0.421 ?	0.243	0.432
MNIST, acc	0.137	0.447	0.464	0.337	0.404	0.242	0.478	0.559 ?	0.559
STL10, acc	0.910	0.932	0.932	0.934	0.946	0.956	0.964	0.967 ?	0.970
PCam, acc	0.484	0.501	0.505	0.520	0.506	0.554	0.501	0.603 ?	0.573
CLEVR, acc	0.104	0.148	0.128	0.152	0.188	0.142	0.132	0.240 ?	0.240
Rendered SST2, acc	0.483	0.489	0.527	0.529	0.508	0.539 ?	0.525	0.484	0.484
ImageNet, acc	*	0.375	0.401	0.426	0.451	0.488?	0.482	0.392	0.638
FGVC Aircraft, mean-per-class	0.020	0.033	0.043	0.046	0.053	0.075	0.046	0.220 ?	0.244
Oxford Pets, mean-per-class	0.462	0.560	0.595	0.604	0.587	0.546	0.635 ?	0.507	0.874
Caltech101, mean-per-class	0.590	0.786	0.775	0.777	0.834	0.835 ?	0.835 ?	0.792	0.883
Flowers102, mean-per-class	0.063	0.401	0.388	0.455	0.449	0.517 ?	0.452	0.357	0.697
Hateful Memes, roc-auc	0.527	0.564	0.516	0.530	0.537	0.519	0.543	0.579 ?	0.589

Ниже на рисунках 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 оценки моделей.

Рисунок 6 - Усредненный график результатов экспериментов zero-shot (звездочки) и few-shot (линии) классификаций для различных моделей по всем датасетам, участвовавшим в оценке качества

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

Рисунок 7 - График результатов экспериментов zero-shot (звездочки) и few-shot (линии) классификаций для различных моделей для каждого из датасетов, участвовавших в оценке качества

Почему же 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 классификации для различных моделей

	ruCLIP Base [vit-base-patch32-224]	ruCLIP Base [vit-base-patch16-224]	ruCLIP Large [vit-large-patch14-224]	ruCLIP Base [vit-base-patch32-384]	ruCLIP Large [vit-large-patch14-336] exclusive	ruCLIP Base [vit-large-patch16-384] exclusive	ruCLIP Small [rugpt3-small]	CLIP [vit-base-patch16-224] original
Food101	0.765	0.827	0.840	0.851	0.896 ?	0.890	0.874	0.901
CIFAR10	0.917	0.922	0.927	0.934	0.943	0.942	0.948 ?	0.953
CIFAR100	0.716	0.739	0.734	0.745	0.770	0.773	0.794 ?	0.808
Birdsnap	0.347	0.503	0.567	0.434	0.609	0.612 ?	0.584	0.664
SUN397	0.683	0.721	0.731	0.721	0.759 ?	0.758	0.753	0.777
Stanford Cars	0.697	0.776	0.797	0.766	0.831	0.840 ?	0.806	0.866
DTD	0.690	0.734	0.711	0.703	0.731	0.749 ?	0.738	0.770
MNIST	0.963	0.974	0.949	0.965	0.949	0.971	0.985 ?	0.989
STL10	0.957	0.962	0.973	0.968	0.981 ?	0.974	0.977	0.982
PCam	0.827	0.823	0.791	0.835	0.807	0.846 ?	0.833	0.830
CLEVR	0.356	0.360	0.358	0.308	0.318	0.378	0.524 ?	0.604
Rendered SST2	0.603	0.655	0.651	0.651	0.637	0.661 ?	0.568	0.606
FGVC Aircraft	0.254	0.312	0.290	0.283	0.341	0.362	0.500 ?	0.604
Oxford Pets	0.774	0.820	0.819	0.730	0.753	0.856	0.895 ?	0.931
Caltech101	0.904	0.917	0.914	0.922	0.937 ?	0.932	0.937	0.956
HatefulMemes	0.545	0.568	0.563	0.581	0.585	0.578	0.638 ?	0.645

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

Рисунок 8 - Корреляция метрик качества zero-shot и few-shot классификаций

Про скорость

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

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

	ruCLIP Base [vit-base-patch32-224]	ruCLIP Base [vit-base-patch16-224]	ruCLIP Large [vit-large-patch14-224]	ruCLIP Base [vit-base-patch32-384]	ruCLIP Large [vit-large-patch14-336] exclusive	ruCLIP Base [vit-large-patch16-384] exclusive
iter/sec	308.84	155.35	49.95	147.26	22.11	61.79

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

Нам удалось обучить несколько разных версий 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.

Рисунок 9 - Пример работы модели ruCLIP: 'скрипка': 0.68, 'гитара': 0.40, 'пианино': 0.32, 'аккордеон': 0.26, 'труба': 0.28

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