RUDOLPH: путь к AGI

Как можно описать последние два года в области машинного обучения и искусственного интеллекта. Пожалуй, подойдёт словосочетание «расцвет генеративных моделей». 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