[DeepMind] Flamingo: a Visual Language Model for Few-Shot Learning
Авторы: Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, Iain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katie Millican, Malcolm Reynolds, Roman Ring, Eliza Rutherford, Serkan Cabi, Tengda Han, Zhitao Gong, Sina Samangooei, Marianne Monteiro, Jacob Menick, Sebastian Borgeaud, Andrew Brock, Aida Nematzadeh, Sahand Sharifzadeh, Mikolaj Binkowski, Ricardo Barreira, Oriol Vinyals, Andrew Zisserman, Karen Simonyan
Статья: https://arxiv.org/abs/2204.14198
Пост: https://www.deepmind.com/blog/tackling-multiple-tasks-with-a-single-visual-language-model
Код: нет
Модель: нет
API: нет

Цивилизованный мир идёт вперёд.

🦩Flamingo — это новая зрительно-языковая модель (Visual Language Models, VLM) на 80B параметров от DeepMind. Разница с обычной языковой моделью (LM) в том, что зрительно-языковая модель (VLM) должна уметь принимать мультимодальный prompt, в котором среди текста могут встречаться картинки и видео. На выходе модели текст. В таком режиме с ответом произвольным текстом модель может решать различные задачи: классификацию, captioning, visual dialogue, visual question answering и т.д.

В целом этот подход аналогичен используемому в семействе моделей T5 (https://news.1rj.ru/str/gonzo_ML/761) подходу, когда универсальный выход — это текст, и никакие специализированные головы в модели не нужны.

Идеальная комбинация для чатбота. Голос ещё осталось добавить для тех, кто любит голосовыми сообщениями пулять, но авторы Фламинго, видимо, тоже таковых недолюбливают, потому и не включили такую возможность =) Но явно SVLM или VSLM не за горами.

Модель демонстрирует SoTA на множестве мультимодальных задач в режиме few-shot learning без дополнительного обучения (то есть когда примеры демонстрируются внутри prompt’а, так называемый in-context learning), побив в том числе и специально зафайнтюненные модели. Можете представить это как GPT-3, куда ещё и картинки с видео можно отправлять. Демонстрации очень прикольные, особенно чат, где модель поняла, что ей дают тест на эффект Струпа (https://en.wikipedia.org/wiki/Stroop_effect).

Как оно всё устроено внутри?

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

В отличие от обучения обычных LM, мультимодальность добавляет сложностей.

Во-первых, предобученные модели это хорошо, но у них нет встроенных средств для добавления входов с другими модальностями, и добавление таких входов — это некий дестабилизирующий челлендж для предобученной модели. Для повышения стабильности авторы чередуют слои кросс-внимания с обычным языковым self-attention, который заморожен во время обучения.

Другой челлендж — это поддержка картинок и видео, размерность которых выше и они требуют больше памяти, и которые неочевидно как просто преобразовать в последовательность (иногда для них вводятся специальные inductive biases даже в трансформерах). Авторы используют здесь свой любимый Perceiver (https://news.1rj.ru/str/gonzo_ML/700), хорошо работающий в таких условиях.

Наконец, не так просто собрать хорошие мультимодальные обучающие данные, после обучения на которых потом возникнет способность к few-shot learning (может же и не возникнуть). Авторы собрали для этого коллекцию мультимодальных датасетов: слабо связанный текст/картинка в виде текста веб страниц и layout’а, и сильно связанные датасеты с парами картинка/текст и видео/текст. Правильный микс оказался необходимым.

В качестве зрительной взята модель а-ля CLIP (https://news.1rj.ru/str/gonzo_ML/665). Её роль — выделить семантические признаки из изображения, относительно которых могли бы быть заданы вопросы: цвет, форма, природа объекта, расположение и т.д. Сам картиночный энкодер — это Normalizer-Free ResNet (NFNet, https://arxiv.org/abs/2102.06171), модель F6 (самая тяжёлая). Предобучается модель с контрастной целевой функцией на датасете пар картинки/тексты. Близость считается скалярным произведением от mean pooling выхода NFNet’а и mean pooling выхода BERT’а. Уже обученная модель внутри Flamingo выдаёт двумерную сетку фич, которая через flatten делается одномерной.

В случае видео кадры сэмплятся с частотой 1 FPS, независимо кодируются картиночным энкодером, получая последовательность из T feature maps, которые затем конкатенируются в одну длинную последовательность фичей X_f. К ним добавляются обучаемые временные эмбеддинги, учитывающие номер кадра для видео. Пространственные эмбединги не добавляются, потому что ничего не улучшают (возможно, ценная релевантная информация прилетает из CNN).

В качестве языковой модели взята обученная на большом корпусе авторегрессионная модель Chinchilla (https://arxiv.org/abs/2203.15556) c 70B параметров. Эта недавняя модель знаменита тем, что при сильно меньшем числе параметров побеждает тяжеленные GPT-3 (175B, https://news.1rj.ru/str/gonzo_ML/305), Gopher (280B, https://news.1rj.ru/str/gonzo_ML/742), Jurassic-1 (178B), и Megatron-Turing NLG (530B). И, кстати, Шиншилла сильно подкорректировала предсказания про развитие ИИ в сторону что всё будет быстрее, чем ожидалось.

Дальше челлендж эти модели скрестить. Веса предобученных зрительной и языковой моделей замораживаются и более не меняются. Они соединяются с помощью пары обучаемых компонент.

Первая, Perceiver Resampler, получает пространственно-временные признаки из зрительного энкодера и выдаёт фиксированного размера множество визуальных токенов. Эти токены используются для обусловливания замороженной языковой модели через вторую обучаемую компоненту — свеже инициализированные слои кросс-внимания, которые чередуются со слоями замороженной языковой модели. Такие новые слои дают выразительный способ инкорпорировать зрительную информацию в задачу предсказания следующего токена.

Perceiver Resampler занимается трансформацией последовательности картиночных фичей X_f произвольного размера в фиксированное количество (R=64) визуальных токенов. Как следует из названия, архитектура этого блока аналогична Perciever’у. У него есть заданное фиксированное число (то самое R) обучаемых латентов, обозначаемых вместе как X. На уровне работы механизма внимания эти латенты дальше становятся Q, через кросс-внимание “смотрят” на последовательность картиночных фичей K,V, и генерят на выходе те самые R визуальных токенов (K и V при этом не просто X_f, а конкатенация [X_f, X], что отличается от стандартного персивера — логика чисто эмпирическая, так работает лучше). Про сам Perceiver, Perceiver IO и Hierarchical Perceiver у нас тут есть отдельные посты, если кому интересно.

Визуальные токены и текстовый вход отправляются в блоки декодера трансформера. Каждый блок состоит из двух других под-блоков: кросс-внимание (gated xattn-dense), и замороженная языковая модель (LM layer).

В кросс-внимание в качестве K и V приходят визуальные токены из resampler, в качестве Q приходят токены языка (здесь везде под токенами мы подразумеваем уже их эмбеддинги). После кросс-внимания стоит полносвязный слой (FFW). Ко всем входам Q,K,V и FFW применён LayerNorm.

Для того, чтобы при инициализации вся эта модель работала как обычная замороженная LM используется tahn-гейтинг — выходы свеже добавленных слоёв умножаются на tanh(α) перед тем, как просуммироваться с residual connection. Параметр α обучаемый и инициализирован нулём, так что на старте из новых слоёв просто добавляются нули к эмбеддингам, проходящим через стек трансформерных блоков. Во время обучения этот параметр меняется и модель становится из LM уже VLM, начиная учитывать зрительные входы. Такой гейтинг повышает и стабильность обучения, и итоговое качество.

Эти новые слои добавляются в модели разной глубины с разной частотой. Всего обучили три разных Фламинго, 3B, 9B и основная 80B параметров. Слои соответственно добавлялись в каждый, каждый 4-й или каждый 7-й блок по мере увеличения размера модели.

LM layer — это обычный self-attention + FFW. Языковая часть большой (80B) модели Фламинго инициализирована Шиншиллой (70B), остальные — более мелкими языковыми моделями из той же статьи.

Отдельный вопрос, как кодировать на входе эти мультимодальные данные.

Положение картинки в тексте помечается текстовым тегом <image> (это не специальный новый токен, просто текст). Есть токены <BOS>, <EOS> для начала и конца текста, а также новый токен <EOC> (end of chunk), который ставится перед тегом картинки, а также перед концом документа до <EOS>. Новый токен добавили в словарь языковой модели с рандомным эмбеддингом и обучили его.

Каждый текстовый токен текущего чанка может делать кросс-внимание только на последнее предшествующее изображение. Эмпирически это работало лучше, чем кросс-внимание на все предшествующие картинки. Но важно, что это только про блок кросс-внимания. В блоке обычного языкового трансформера каждый токен по-прежнему смотрит на все предыдущие (а среди них есть и картиночные), так что не напрямую зависимость от всех предшествующих картинок всё равно есть. Да, механизм внимания везде causal, то есть в будущее не заглядывает и смотрит только на токены, которые были до текущего. Из профитов такой схемы — независимость от количества изображений и генерализация на любое их число. Так обучались на входах с максимум пятью картинками, а проверялись потом на последовательностях, где картинок было до 32-х.

Модель обучалась на трёх типах датасетов: 1) датасет с чередующимися картинками и текстом из веба, 2) пары картинка/текст, 3) пары видео/текст.

По вебу собрали датасет MultiModal MassiveWeb (M3W). В нём 185M картинок (≈4.3 на страницу) и 182GB текста (≈4.4 KB на страницу). В обучение идут тексты длиной до 256 токенов и до 5 картинок внутри.

Для пар картинка/текст или видео/текст это смесь разных имеющихся датасетов: ALIGN с 1.8B картинок + alt-текст (но он шумный, https://news.1rj.ru/str/gonzo_ML/679), LTIP (Long Text & Image Pairs) с 312M картинок, и VTP (Video & Text Pairs) с 27M коротких видео (≈22 секунды в среднем). У последних двух датасетов описания покачественнее.

Обучали оптимизируя взвешенную по датасетам сумму negative log-likelihood сгенерированных текстов при заданных входах. Оказалось, что тюнинг весов очень важен.

Модель реализована на JAX и Haiku, обучали на TPUv4. Самую большую модель шардировали на 16 серверов (кроме слоёв NFNet). Использовали смесь float32 и bfloat16. Модель свежая, обучена в марте 2022.

После того как модель обучена, можно проверять её в режиме in-context learning, пробуя разные промпты. По большому счёту вся эта движуха с in-context learning началась с GPT-3, в соответствующем посте (https://news.1rj.ru/str/gonzo_ML/305) больше деталей про суть этой истории.

Для few-shot learning на вход подавался мультимодальный запрос с конкатенированными примерами в случайном порядке. В open-ended задачах генерировали текст до первого <EOC> токена, с beam search размера 3. В close-ended задачах добавляли к запросу все варианты и оценивали log-likelihood получившихся предложений моделью, далее ранжируя результаты. Для zero-shot использовали два примера, но без самих картинок.

Провели много оценок на разных задачах. Фламинго бьёт многие другие zero-shot/few-shot SoTA иногда уже с 4-мя примерами. С 32-мя примерами на 6 задачах побила лучшие методы, которые файнтюнились на тысячах примеров. Надо сказать, что даже 3B модель весьма неплоха.

Не фокус работы, но попробовали пофайнтюнить Flamingo на разные задачи, всё ещё улучшается и получили несколько новых SoTA. Кстати, Florence (https://news.1rj.ru/str/gonzo_ML/734) побили.

Качественные примеры, возможно, более бомбические нежели количественные. Модель реально даёт очень классные ответы. Не знаю, как много здесь черрипикинга.

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

Ни датасеты, ни модель за пределы компании не выйдут :(