Как можно описать последние два года в области машинного обучения и искусственного интеллекта. Пожалуй, подойдёт словосочетание «расцвет генеративных моделей». Gopher, LaMDA, ChatGPT, LLaMA — большие языковые модели, работающие с текстами. Dalle, Dalle-2, Kandinsky, Kandinsky 2.0, Kandinsky 2.1, Imagen — генеративные модели, способные обрабатывать текст и генерировать изображения. Flamingo, FROMAGe, Kosmos-1, GPT-4 — модели, которые тоже работают с текстами и изображениями, но генерируют только текст.

Тренд на развитие мультимодальности прослеживается очень ярко. Если нашумевшая модель ChatGPT, построенная на основе архитектуры InstructGPT и принципах Reinforcement Learning with Human Feedback, работает только в текстовой модальности и может отвечать на текстовые вопросы различной степени сложности, то её последователь GPT-4 работает уже в двух модальностях, текст и изображение, но выдаёт только текстовый ответ. При этом большинство моделей «однонаправлены», то есть либо переводят текст в соответствующее ему изображение, либо генерируют текстовый ответ по мультимодальному запросу.

Подобное ограничение не позволяет моделям решать разнообразные задачи в двух модальностях, хотя кажется очевидным, что хорошее понимание текстов и изображений, а также способность к их сопоставлению может пригодиться при решении большого спектра задач: описание изображений (Image captioning), Visual QA, Text Recognition in the Wild, Text QA и многих других.

Чтобы адаптировать модель под разнообразные задачи в двух уже обозначенных модальностях, мы задумались над модификацией генеративных моделей, а именно над изменением обработки входной последовательности. Результатом наших экспериментов стала модель RUDOLPH, которая представляет собой авторегрессионную модель, способную решать разные типы задач в текстово‑визуальном домене.

Архитектура RUDOLPH в деталях

Архитектурно модель представляет собой декодер‑блок трансформера, работающий с входной последовательностью токенов, которую можно условно разделить на три основных сегмента: левые текстовые токены, токены изображения и правые текстовые токены (подробное описание на схеме). Именно сегментирование входной последовательности и специально‑организованный механизм внимания позволяют модели работать в мультизадачном формате и осуществлять так называемый fusion модальностей (текста и изображения).

Рассмотрим подробнее роль каждого типа токенов во входной последовательности:

• Левые текстовые токены. Важнейший элемент входной последовательности, отвечающий за понимание текста и, следовательно, любой текстовой инструкции к заданию. Именно левая часть текста использовалась при обучении модели языковому моделированию (text2text).

• Токены изображения. Служат для кодирования входного изображения или для генерации нового изображения по текстовому запросу. Этому типу токенов отведено место между двумя текстовыми последовательностями для решения задач в двух направлениях: текст — изображение (text2image) и изображение — текст (image2text).

• Правые текстовые токены. Служат для генерации ответа в задачах, которые требуют текстовый вывод.

Модель работает в авторегрессионном формате, генерируя сначала левые текстовые токены, затем токены изображения и потом правые текстовые токены, обусловленные левым контекстом. Для токенизации изображения используется Sber‑VQ‑GAN, токенизация текстовой последовательностей осуществляется с помощью YTTM.

Базовые характеристики RUDOLPH

Всего мы обучили три версии модели: 350M, 1.3B, 2.7B. Они различаются количеством параметров и, соответственно, количеством и размером скрытых слоёв:

Можно заметить, что с увеличением размера модели увеличивается и максимальный размер текстовой последовательности. Модель 350М больше подходит для генерации изображений по текстовым запросам (64 токена вполне хватит для подобных запросов), или генерации текста по изображениям. Модель 2.7B, левый текстовый контекст которой вмещает до 384 токенов) может служить уже как языковая модель, продолжать текст или генерировать ответы на текстовые вопросы (например, в абзаце, который вы только что прочитали 112 токенов).

Checkpoint'ы предобученных моделей можно загрузить с Hugging Face ????

• RUDOLPH-350M

• RUDOLPH-1.3B

• RUDOLPH-2.7B

Механизм внимания

В архитектуре RUDOLPH реализован многоуровневый механизм внимания: 1) маскированный механизм внимания (masked self‑attention) для текстовых токенов; 2) механизм разреженного внимания (sparse attention) для визуальных токенов. При генерации текстовых токенов на матрицу скалярных произведений накладывается треугольная маска, чтобы новые токены генерировались только исходя из левого контекста. Когда генерируется новый токен изображения, открытыми остаются не все сгенерированные до этого токены, а лишь его соседи в строке, столбце и некотором контекстом окне.

Подобный механизм разреженного внимания хорошо зарекомендовал себя в генеративных text‑to‑image моделях (например, DALL‑E), поскольку он позволяет учитывать лишь наиболее релевантные входные токены изображения, тем самым значительно снижая требования к памяти при их вычислении.

Маски внимания, которые используются в разных версиях модели RUDOLPH (350В, 1.3М, 2.7М), приведены на рисунке ниже. Слева направо показаны маски внимания для режимов «по строкам», «по столбцам», «свёртка» и «последняя свёртка». Для первого режима можно заметить характерный синий треугольник, который уменьшается с увеличением количества токенов левого контекста. Маска «по столбцам» характеризуется параллельными наклонными линиями в соответствии с позицией токена в столбце для каждого сдвига последовательности. Маски в режиме «свёртки» показывают, как токены изображения используются в окрестности токена в «скользящем окне».

На схеме ниже можно увидеть, что при генерации токенов изображения почти вся обработанная картинка остаëтся закрытой.

Можно отметить, что три версии модели отличаются не только количеством обучаемых параметров, но и спецификой сегментации входной последовательности на текстовые и визуальные токены. Так, например, самая большая модель RUDOLPH 2.7B может кодировать до 384 токенов в левом тексте, что делает еë наиболее подходящей версией модели для решения текстовых задач.

Спецтокены

Разделение токенов во входной последовательности, о которой мы говорили ранее, предполагает, что RUDOLPH может эффективно решать целый спектр задач, связанных с двумя модальностями. Например, модель может описывать изображение и отвечать на вопросы по изображению, генерировать изображения по тексту или работать как языковая модель. Чтобы помочь модели ориентироваться в мультизадачном сетапе, в словарь были добавлены спецтокены, кодирующие тип решаемой задачи.

Соответствующие токены добавляются в начало левой и правой текстовой последовательности. Левый токен обуславливает генерацию всей последовательности, правый токен кодирует информацию о всëм предыдущем контексте и обуславливает генерацию ответа в правых текстовых токенах.

В процессе предобучения использовались следующие спецсимволы:

<LT_T2I> — для задачи генерации изображения по тексту.

<LT_I2T> — для задачи генерации текстового описания изображения.

<LT_T2T> — для задачи языкового моделирования (LM).

<RT_I2T> — для задачи генерации текстового описания по изображению.

<LT_UNK> и<RT_UNK> — спецсимволы неизвестной задачи для левого и правого текстов соответственно. Unknown‑токены заменяли токены определенных задач в 25 % случаев, чтобы модели было легче адаптироваться под решение новых неизвестных задач.

Обучение

Предобучение

RUDOLPH был предобучен на трёх типах задач:

• описание изображения (image captioning i2t);

• генерация изображения по тексту (image generation t2i);

• генерация текста в левых токенах (language modeling t2t).

Для предобучения были доступны следующие данные:

• 119 млн пар текст‑изображение;

• 60 млн текстовых параграфов на русском языке.

В процессе обучения пакет собирался из примеров задачи одного типа (t2i/i2t/t2t). Единственный тип задачи, который модель не видела на предобучении, это генерация текста в правых текстовых токенах по левому тексту. Эту комбинацию мы использовали при тонкой настройке модели.

Предобучение

Самую большую модель RUDOLPH 2.7B мы обучали на 160 GPU в течение 60 дней.

• Размер пакета на 1 GPU: 32.

• Эффективный размер пакета: 5120.

• Пиковый learning rate: 1e-4.

Тонкая настройка с инструкциями (Instruction tuning)

В левых текстовых токенах RUDOLPH обучался как стандартная языковая модель, что означает, что он способен понимать естественный язык. Мы предположили, что модель можно обучать решению новых задач с помощью их описания на естественном языке в виде инструкций (этот подход особенно популярен в последнее время, а выход текстовых моделей FLAN‑T5 и FLAN‑PALM и вовсе показывает, что такой режим обучения модели во многом превосходит prompt‑tuning и few‑shot learning).

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

Первый подход к развитию парадигмы Multimodal Instruction Tuning был реализован нами для модели RUDOLPH при подготовке к соревнованию AI Journey Contest 2022 (AIJC), о котором мы расскажем подробнее в нашей следующей статье. И стоит отметить, что подход оказался довольно удачным, наша базовая модель, предложенная участникам в качестве baseline, заняла почётное третье место.

Примеры инференса модели RUDOLPH 2.7B на pretrain-задачах

Посмотрим, что умеет RUDOLPH. Так как предобучение проходило на трёх типах задач, попросим модель сгенерировать изображение по текстовому запросу, описать изображение, а также продолжить текст по небольшой затравке на русском языке.

• Задача генерации изображения

Пользователь: Морской пейзаж в закате солнца.

Пользователь: Кот в очках.

• Текстовые задачи
  Пользователь: Из окна открывался вид
  RUDOLPH: из окна открывался вид на морское побережье и горы.

Благодаря тому, что в RUDOLPH работает в двух модальностях, её можно попросить продолжить текст и сопроводить его визуальным описанием.

Пользователь: Как приготовить блины?

RUDOLPH: для приготовления нам понадобится: мука пшеничная — 1 ст. л.; разрыхлитель (сода) — 1 ст. л.; растительное масло — 2 ст. л.; масло сливочное — 1 ст. л.; яйцо — 1 шт.; сахар — 1 ст. л.; мука пшеничная — 1 ст. л.; яйца взбить миксером в пышную пену. Для украшения и подачи в блинчики добавить сметану и приступать к подаче. Блины сметанные с измельченным миндалем.

Примеры генераций RUDOLPH 2.7B после тонкой настройки на различных задачах

Выше мы привели примеры работы модели после предобучения в zero‑shot режиме. Посмотрим, чему научилась модель после мультизадачного обучения с инструкциями. Мы ожидали, что после обучения модели в подобном формате она сможет понимать формулировку задачи на естественном языке и, в зависимости от заданной инструкции, генерировать ответ. Проверим.

• Text QA
  Пользователь: Кто такой Чебурашка?
  RUDOLPH: чебурашка — это один из самых известных героев советского телевидения, это маленькая плюшевая игрушка, которую можно увидеть в магазине

• Image Generation
  Пользователь: Сгенерируй изображение по описанию. Поле с одуванчиками в стиле Ван Гога.

• Image Captioning
  Пользователь: Опиши изображение.

• Visual QA
  Пользователь: Проанализируй изображение и дай ответ на вопрос. Какое время суток изображено?

Примеры использования RUDOLPH в zero‑shot режиме, а также тонкой настройки под задачу Text QA можно найти в репозитории на Github.

Заключение и выводы

Предложенное нами решение RUDOLPH как путь к AGI, показывает хорошие результаты на задачах в двух модальностях, а также демонстрирует понимание инструкций и возможность тонкой настройки на нескольких задачах. Последнюю особенность мы исследовали весной 2022 года, когда Instruction tuning ещё не стал мэйнстримом :)

В продолжении развития идеи сильных мультимодальных мультизадачных моделей мы работаем в направлении усиления моделей за счёт использования архитектур encoder‑decoder и устранения чёткого разделения между левым, центральным и правым контекстами. В настоящее время командой Sber AI разработана концепция мультимодальной архитектуры, которая не зависит от типа и количества модальностей, но все описанные подходы к обучению и тонкой настройке RUDOLPH будут применимы в новой концепции, о которой мы также скоро расскажем подробнее.