Совсем недавно компания Сбер представила ряд, адаптированных под русский язык, моделей машинного обучения под названием RuCLIP. Суть их работы заключается в возможности сравнения схожести текстового описания и изображения. Рассматривая результаты тестирования, мы заметили, что их модели довольно большие (150+ миллионов параметров) и при этом занимают довольно много места, связи с чем решили сделать свою маленькую, быструю и масштабируемую версию под названием RuCLIP tiny.

Данные

В наших экспериментах мы попробовали датасеты разного качества и остановились на выборе переведённых наборов COCO и CC12M. Они оба содержат пары текст + картинка. На тот момент CC12M ещё не был переведён, связи с чем пришлось обработать его с использованием гугл переводчика. датасет cc12m гугл переводчиком.

Выбор архитектуры

Архитектура CLIP’а состоит из энкодера для картинок и энкодера для текста. Их задача состоит в том, чтобы уменьшить косинусное расстояние между их выходными векторами.

В начале мы сделали архитектуру с cointegrated/rubert-tiny в качестве текстового энкодера и swin_tiny_patch4_window7_224 в качестве энкодера изображений. Но она показала себя не очень относительно следующей архитектуры.

Тогда мы решили обновить архитектуру. Для текстового энкодера мы выбрали недавно вышедший DeepPavlov/distilrubert-tiny-cased-conversational-v1, а для энкодера изображений - convnext_tiny, который быстрее, чем swin transformer.

Всего в новой архитектуре 38317537 параметров.

Тренировка модели

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

Поскольку мы тренировали на гугл колаб про, то у нас была всего одна NVIDIA Tesla P100 на 16 гб. К сожалению, у нас не получалось сделать размер батча выше, чем 256, а для contrastive loss он очень важен. Тогда мы стали искать различные статьи и наткнулись на эту статью. Там авторы рассказывают, как реализовать накопление градиента для contrastive loss. Мы это реализовали и смогли поставить размер батча равный 2048.

Результаты модели

Для оценки модели мы решили взять Google Colab Pro + с NVIDIA Tesla V100.

На датасете CIFAR100 мы получили 0.4662 top1 accuracy и 0.7318 top5 accuracy.

Скорость модели

Подведение итогов

В результате нашей работы мы сделали ru-CLIP-tiny, который способен оценивать схожесть картинок с текстом. Код и веса вы можете найти тут.

Хочется выразить благодарность Sber AI за предоставленные гранты, на которые были проведены исследования, в рамках Artificial Intelligence International Junior Contest (AIIJC).

Авторы

@Арсений Шахматов | @Максим Герасимов