OmniFusion: выходим за границы текста

Кто-то ещё сомневается, что в мире машинного обучения происходит революция? Уверен, мы являемся свидетелями преобразования привычного взаимодействия с данными, поиска информации, да и вообще работы как таковой. Ведь умные ассистенты (ChatGPT, GigaChat, Bard) готовы взять на себя даже самые сложные задачи.

Но не всегда возможно сформулировать проблему в виде текстового запроса, иногда требуется информация из других «модальностей» — картинка, звук, 3D и тд. Ниже я разберу какие именно есть способы соединения больших языковых моделей (LLM) с дополнительными форматами данных, а также опишу как устроена наша новая модель OmniFusion.

Dalle 3

Dalle 3

Мультимодальные модели

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

  1. Мультимодальная архитектура, которая сразу обучается на разнообразных форматах данных (Kosmos-1, RUDOLPH, ONE-PEACE)

  2. Сильная языковая модель, предобученная исключительно на тексте, затем обучается понимать дополнительные модальности.

Наиболее перспективен, да и в целом более простой в реализации, именно второй вариант. Поэтому про такие подходы, базирующиеся на сильных языковых моделях, я расскажу подробнее.

Как языковую модель сделать мультимодальной

Идейно все решения тут очень похожи — есть некоторая большая языковая модель, которая берёт на себя роль главного мыслительного центра, и есть энкодеры дополнительных модальностей, например: для картинок — ViT, для звука Wav2Vec и тп. Информация из новой модальности при помощи таких (обычно предобученных и замороженных) энкодеров попадает в текстовую модель через обучаемые адаптеры (FROMAGe, LLaVA) или обучаемый cross-attention (Flamingo). При этом в некоторых решениях языковая модель может быть заморожена (FROMAGe), либо дообучаться вместе с адаптером (miniGPT-4, Shikra, LLaVA). 

Забавно, но даже в случае, когда языковая модель НЕ дообучается, в ней всё равно просыпаются «мультимодальные нейроны», которые связывают текстовые и визуальные (не только) концепты. То есть большие LM очень хорошо подходят для такого сетапа, так как они выучивают внутри себя структурированную картину мира, а не просто запоминают тонны текста в режиме стохастического попугая. Но это только косвенное свидетельство, а однозначно ответить, почему это работает и как пока никто не может. Но от этого даже интереснее!

OmniFusion

Для создания своей мультимодальной модели мы решили взять самые лучшие и наиболее стабильные практики — сильная LLM, мощный адаптер, спецтокены новых модальностей и многоступенчатое инструктивное обучение с постепенной разморозкой LLM. 

Рисунок 1 — Архитектура модели OmniFusion

Рисунок 1 — Архитектура модели OmniFusion

Устройство модели

Мозговое ядро OmniFusion — это последняя версия GigaChat-7B. Для начала мы решили сосредоточиться только на картинках, поэтому в качестве визуального энкодера взяли CLIP-ViT-L, т.к. по нашим тестам он показал наиболее полную передачу информации из картинок в языковую модель (возможно из-за бóльшего количества эмбеддингов). 

Из визуального энкодера мы берём все эмбеддинги до пуллинга, кроме CLS-токена (576 штук) и отображаем их в 576 текстовых эмбеддингов подающихся на вход языковой модели вместе с текстовым запросом. Для такого отображения мы используем однослойный, четырехголовый трансформер — это даёт результаты лучше, чем всего лишь один линейный слой как в Shikra и Fromage, либо MLP как в LLaVa. Для более качественного выделения визуальной «инъекции» в последовательности текстовых эмбеддингов LLM и в качестве задела на добавление большего количества модальностей мы используем обучаемые спецтокены начала и конца картинки.

Процесс обучения

Размораживать языковую модель считается крайне рискованной операцией, так как это неизбежно приводит к деградации в текстовом домене (да и в целом к ослабеванию модели в смысле качества генерации, если делать это слишком резко). Учитывая сказанное выше, обучение OmniFusion делится на 2 этапа:

  1. Предобучение адаптера на задаче Image Captioning. Перед тем как приступить к дообучению LLM, нужно максимально качественно подготовить адаптер, связывающий модальности, поэтому тут обучается только он (и эмбеддинги спецтокенов). На этом этапе мы использовали большое количество НЕ инструктивных данных из фильтрованных датасетов Conceptual Captions-4M и LAION. OmniFusion на данном этапе учится предсказывать описания картинок, а лосс функция —  это обычная кросс-энтропия. Гиперпараметры: 1 эпоха, ≈600k примеров (bs = 128), ≈45 часов на 2x Tesla A100 (80Gb)

  2. После того, как адаптер научился отображать визуальные эмбеддинги ViT в пространство текстовых эмбеддингов языковой модели пришла пора разморозить GigaChat, чтобы он лучше понимал диалоговый формат, сложные запросы и даже OCR. На данном этапе помимо адаптера целиком размораживается и сама языковая модель (но с маленьким LR). При этом лосс мы учитывали только для ответных реплик, иначе модель будет учиться задавать вопросы, а не отвечать на них. Состав датасетов второго этапа обучения представлен на диаграмме. Гиперпараметры: 1 эпоха, ≈620k примеров (bs = 128), ≈52 часов на 2 × Tesla A100 (80Gb)

Рисунок 2 — Графики лосса на обеих фазах обучения

Рисунок 2 — Графики лосса на обеих фазах обучения

Рисунок 3 — Состав текстово-картиночного датасета для второй стадии обучения

Рисунок 3 — Состав текстово-картиночного датасета для второй стадии обучения

Результаты

Мы сравнили OmniFusion с последними мультимодальными SOTA моделям. Наша модель получилась довольно сильной, обходит QWEN-VL-7b, LLaVa-1.5–7b и LLaVa-13b не только на генеративных метриках (см. круговую диаграмму), но и на классификационном бенчмарке VisualDialog (см. таблицу). Но ещё есть куда расти: например, LVIS победить пока не получилось. Также немного проседают OCR возможности нашей модели, но мы над этим уже работаем.

Рисунок 4 — Круговая диаграмма оценки качества модели OmniFusion в ряду ближайших конкурентов на 8 популярных бенчмарках

Рисунок 4 — Круговая диаграмма оценки качества модели OmniFusion в ряду ближайших конкурентов на 8 популярных бенчмарках

Следует отметить, что датасеты, на которых проводилась оценка качества модели представляют собой довольно широкий список доменов:

  1. MMBench — бенчмарк на различные темы с вопросами по изображениям, GPT-4 используется в качестве метрики для сопоставления правильных ответов и сгенерированных

  2. Text-VQA — бенчмарк для задачи OCR с вопросами по изображениям, содержащим текст

  3. Viz-wiz — датасет с вопросами слабовидящих и слепых людей, которые они задавали, чтобы понять, что нарисовано на картинке 

  4. MM-Vet — датасет комплексных вопросов по изображениям, GPT-4 также используется в качестве метрики

  5. POPE — бенчмарк для оценки галлюцинаций

  6. LLaVA-bench — небольшой датасет из статьи про LLaVA (синтетические диалоги на основе описаний изображений с использованием  GPT4)

  7. VQA v2.0 — ещё один датасет с вопросами по изображениям, для ответов на вопросы требуются знания не только из самого изображения, но и commonsense

  8. ScienceQA — бенчмарк с вопросами на различные научные темы, основанный на различных лекциях, вопросы с несколькими вариантами ответов

Рисунок 5 — примеры бенчмарков

Рисунок 5 — примеры бенчмарков

В этой таблице представлены результаты сравнения моделей по качеству ответов на бенчмарке Visual Dialog

Model

NDCG

MRR

Recall@1

Recall@5

Recall@10

OmniFusion

25.91

10.78

4.74

13.80

20.53

LLaVA-13B

24.74

8.91

2.98

10.80

18.02

Ниже представлены скриншоты использования OmniFusion в различных юзкейсах — от анализа медицинских снимков до определения съедобности предметов на фотографиях. Чуть позже будет открыт доступ к телеграм-боту с нашей моделью.

Рисунок 6 — Примеры работы OmniFusion на различных доменах

Рисунок 6 — Примеры работы OmniFusion на различных доменах

Планы

В ближайшее время мы выложим публичную версию OmniFusion, базирующуюся на открытой языковой модели. Также мы начали готовить версию, которая понимает русский язык, работает на энкодерах ImageBind и принимает большее количество модальностей (звук, 3D, видео). Следите за обновлениями на GitHub!

Авторы

Обучением модели занималась научная группа FusionBrain Института AIRI при участии учёных из Sber AI.

Следите за новостями в каналах: AbstractDL, CompleteAI, Dendi Math&AI, Канал Сергея Маркова.

© Habrahabr.ru