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