LayerShuffle: Enhancing Robustness in Vision Transformers by Randomizing Layer Execution Order
Matthias Freiberger, Peter Kun, Anders Sundnes Løvlie, Sebastian Risi
Статья: https://arxiv.org/abs/2407.04513

Недавно вышли две свежие прикольные работы про творческий подход к вычислению трансформерных слоёв. Это перекликается с нежно любимой мной темой про Adaptive Computation Time (ACT) и вообще динамическими вычислениями. Про трансформеры c ACT было тут (https://moocaholic.medium.com/adaptive-computation-time-act-in-neural-networks-3-3-99452b2eff18), но вообще это целая серия постов (в канале она начиналась тут https://www.tgoop.com/gonzo_ML/71).

Сегодня первая интересная работа -- про выкидывание части трансформерных слоёв во время инференса. Если сеть будет устойчива к таким событиям, то это открывает интересные возможности для выполнения вычислений в нестабильной распределённой среде, когда часть вычислений может выпасть или произойти в ином порядке. Способность продолжать работу в таких условиях и не терять сильно в качестве при этом довольно интересна.

Это идейно похоже на Dropout, где выпадает часть нейронов в слое, но делается уровнем выше. В 2019-м одни из моих любимых авторов опубликовали вариант structured dropout под названием LayerDrop (https://arxiv.org/abs/1909.11556), позволявший делать прунинг части слоёв трансформера при инференсе, при этом после прунинга он приводил к моделям с качеством выше, чем если бы модели такого же размера обучались с нуля или даже были получены дистилляцией (!). В этом смысле интересно было бы посмотреть на вариант той же Gemma 2, где малую модель не дистиллировали, а дропали из большой.

LayerShuffle идёт дальше и не просто дропает слои при инференсе, а перемешивает их, меняя порядок вычисления. Делают это на Vision Transformer’ах (ViT, https://www.tgoop.com/gonzo_ML/434), но, кажется, это неважно. Рассматриваются несколько вариантов процедуры.

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

Во втором варианте, LayerShuffle-position, каждый слой также получает эмбеддинг текущей позиции этого слоя (эти дополнительные 32 чиселки конкатенируются с эмбеддингом, прилетающим на вход).

Третий вариант, LayerShuffle-predict, пытается предсказать свою текущую позицию по своему же выходному эмбеддингу. С кроссэнтропийным лоссом при этом, что немного странно, ибо сильный промах и слабый промах в позиции получается штрафуются одинаково.

Всё проверяют на ImageNet2012 с предобученным ViT-B/16 с HuggingFace, который дообучали по описанным процедурам. Стандартный трансформер при переключении в режим случайного порядка слоёв роняет качество с 82% почти в ноль, а варианты LayerShuffle падают примерно с 75% до 63% (интересно, какая была бы разница, если бы модель обучали с нуля по предложенному подходу, а не файнтюнили с классически обученной). LayerShuffle-position лидирует с небольшим отрывом от простого LayerShuffle.

Получается, стандартный трансформер катастрофически неустойчив к смене порядка вычисления слоёв, что было в целом ожидаемо. Интересно, что вариант LayerShuffle-position с передаваемой позицией не очень-то сильно что-то улучшает, все важные данные, видимо, есть и так.

Отдельно проверили эффект прунинга во время инференса. Качество работы примерно соответствует LayerDrop с p=0.2, хотя модель на выкидывание слоёв не обучалась. Если дополнительно к прунингу ещё и перемешивать, то LayerDrop тоже уходит в ноль, а LayerShuffle держится. Получается, если выбирать из этих двух, то последний даёт более устойчивые модели.

Сделали визуализацию выходов слоёв через UMAP, не передавая в алгоритм информацию о позиции слоя. Видно, что с одной стороны явного прерывистого разделения эмбеддингов нет, но таки они группируются по своей реальной позиции, особенно заметно для эмбеддингов близких к началу, но впрочем видно и для конечных.

Интересная идея, раз слои можно вычислять в другом порядке, то что если попробовать смёржить слои из разных моделей, обученных по представленной процедуре? В пределе каждый из 12 слоёв ViT взять из отдельной обученной модели. Намёржили 100 моделей (из 12! возможных комбинаций). Оригинальный ViT-B/16 в такой конфигурации бесполезен, качество крайне низко. LayerShuffle merged в целом ничего, хотя и процента на три ниже оригинальной LayerShuffle модели. Зато если все 12 обученных моделей сансамблировать, то качество процентов на 9 выше.

Когда уже запчасти для GPT-4 будут продаваться на рынке? Можно, наверное, и в микросхемы тогда сразу зашивать и модель рассыпухой набирать 🙂 Ну и распределённую сеть с кусками модели на разных девайсах тоже уже можно строить, ждём здесь своего кубернетеса.

В перерыве между статьями, вот вам свежее интересное видео доклада Жанны Резниковой (https://reznikova.net/) про интеллект животных

https://www.youtube.com/watch?v=aB1u522bNuo

Про новости LLM писать уже поднадоело, но вот краткое саммари анонсов последних дней:

🤖#1 OpenAI выкатил замену старой доброй GPT-3.5 Turbo — GPT-4o mini. То же, но дешевле и качественнее. Контекст 128k, поддержка языков как в большой GPT-4o, цена $0.15/$0.60 за 1M входных/выходных токенов (у 3.5 было $0.50/$1.50).

https://openai.com/index/gpt-4o-mini-advancing-cost-efficient-intelligence/

Теперь у каждого большого игрока есть большая, качественная и медленная модель + маленькая и быстрая (GPT-4o mini, Gemini Flash, Claude Haiku/Sonnet)

🐬 #2 Mistral натренировал с NVIDIA 12B Mistral NeMo. Apache 2.0 лицензия, контекст 128k, вроде как бьют Gemma 2 9B и Llama 3 8B. Ну в целом неудивительно, она и побольше на треть.

Из важного, модель мультиязычная и с function calling! Совпадает с моими ожиданиями в https://www.tgoop.com/gonzo_ML/2821.

https://mistral.ai/news/mistral-nemo/

А ещё Мистраль выкатил MathΣtral для математики (https://mistral.ai/news/mathstral/) и Codestral Mamba на базе Mamba 2 для кода (https://mistral.ai/news/codestral-mamba/).

😼 #3 Apple разродился своей опенсорсной 7B LLM под названием DCLM-7B. Из плюсов, более открытая чем некоторые благодаря открытому датасету. Но в целом непонятно что с ней делать, когда есть Llama3, Gemma2 и Qwen2.

https://x.com/_philschmid/status/1814274909775995087

🐁#4 А, да, HuggingFace выложил открытую SmolLM размерами 135M, 360M, и 1.7B для on-device инференса. Вроде как бьют Phi-1.5 и Qwen2 1.5B.

https://huggingface.co/blog/smollm

А ещё из прикольного, в ЕС не будет доступна новая мультимодальная лама и другие подобные модели. Потому что непонятные регуляции. 😎

https://www.theverge.com/2024/7/18/24201041/meta-multimodal-llama-ai-model-launch-eu-regulations

И ещё из полезного.

Тема про агентов активно развивается. Мой знакомый Victor Dibia, ведёт хороший блог "Designing with Machine Learning" про агентов, мультиагентные фреймворки, в частности AutoGen (он там контрибьютор). Рекомендую. На редкость не булшитно.

Вот примеры интересных постов по теме:
* Multi-Agent LLM Applications | A Review of Current Research, Tools, and Challenges
* Getting Started with AutoGen - A Framework for Building Multi-Agent Generative AI Applications
* Integrating AutoGen Agents into Your Web Application (FastApi + Websockets + Queues)

Также он начал писать книгу "Multi-Agent Systems with AutoGen" в том же издательстве Manning, что и я свою :) Её можно читать по мере написания, что для подобной темы важно — ждать пока выйдет бумажная не имеет смысла, область очень быстро развивается.

Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Sukjun Hwang, Aakash Lahoti, Tri Dao, Albert Gu
https://arxiv.org/abs/2407.09941

Вышла Гидра, двунаправленная (в смысле bidirectional) Mamba! Авторы каждый раз очень основательно подходят к делу — предлагают новый математический фреймворк (здесь это matrix mixer framework), в рамках него раскладывают известные модели и далее находят интересные новые варианты. Очень структурированный подход as usual.

Работа достойна отдельного подробного рассмотрения, но пока хоть так.

Кратко от авторов в твиттере: https://x.com/_albertgu/status/1813252409071968297

Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
[Статья][Код]

Современные нейронные сети, обрабатывающие пространственно-временные данные различной природы будь то текст 📝, изображения 📷, аудио 🎵 и видео 📹 так или иначе обладают механизмом перемешивания каналов (channel mixing), обрабатывающим независимо признаки для каждого элемента последовательности, и механизмом обработки последовательности (sequence mixing), использования взаимосвязей между элементами.

В сегодняшнем рассказе речь пойдет про sequence mixing.

Существуют разнообразные опции sequence mixing. Операция смешивания может не зависеть от входа, как например свертка или обучаемая матрица L x L (L - длина последовательности) в MLP-Mixer, S4 и H3 state-space модели, или зависеть - attention механизм в трансформерах или Mamba (Selective State Spaces).

Кроме того, разные механизмы обладают разной сложностью от длины последовательности. Sequence mixing в Attention или MLP-Mixer требует квадратичного по длине последовательности числа элементарных операций с плавающей точкой (FLOPs), так как используют матричную операцию довольно общего вида. Sequence mixers, обладающие некоторой структурой (низкоранговые, Toeplitz матрицы, DFT, бабочки) позволяют добиваться субквадратичной сложности (обычно с некоторой просадкой в качестве).

И sequence mixing может быть как причинным (causal attention, большинство SSM, в частности, модная нынче Mamba 🐍), где текущий элемент последовательности может смотреть только в прошлое, и двунаправленным (как в masked language modelling, и большинстве задач с ViTами), где элементы последовательности могут изменять свое состояние, как глядя как на прошлые, так и на будущие токены.

И задача, которую, перед собой ставят авторы в данной работе - получение эффективного механизма двунаправленного sequence mixing, такого, чтобы он был с одной стороны субквадратичным (в идеале линейным по длине последовательности) и в то же время выразительным.