[Mamba-2] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality Tri Dao, Albert Gu Статья: https://arxiv.org/abs/2405.21060 Код: https://github.com/state-spaces/mamba Блог: https://goombalab.github.io/blog/2024/mamba2-part1-model/ Твиттер-тред: https://x.com/_albertgu/status/1797651223035904355 Transformers are RNNs/SSMs Краткий формат. Авторы Mamba (https://t.center/gonzo_ML/2148) выпустили обновлённую версию своей модели, Mamba-2. У неё вектор состояния большего размера (16 -> 256), при этом она в два раза быстрее обучается, а её код проще (30 строк). Старая Мамба была хороша, но авторы не были полностью удовлетворены. Во-первых, механизмы внимания оставались где-то за пределами парадигмы SSM, а было бы интересно их как-то объединить. Во-вторых, хоть Мамба и была уже достаточно быстра (с эффективной имплементацией через selective scan), всё равно она уступала по вычислительной эффективности механизму внимания и перемножению матриц. Авторы подошли к проблеме фундаментально и предложили фреймворк под названием structured state space duality (SSD), который объединяет structured SSM и варианты внимания в новом SSD слое. Показана эквивалентность SSM и семейства структурированных матриц под названием semiseparable matrices. Основная идея работы в том, что различные методы вычисления SSM могут быть выражены как алгоритмы умножения структурированных матриц. Также авторы развили теорию линейного внимания (“Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention”, https://arxiv.org/abs/2006.16236 — название текущей работы, видимо, аллюзия на эту) и обобщили результаты про его рекуррентную форму на structured masked attention (SMA). Наконец они соединили SSM и SMA, показав, что у них есть большое пересечение где они двойственны друг другу и по сути являются моделью, выражаемой одной и той же функцией. Также доказали, что любой kernel attention method с быстрой рекуррентной формой должен быть SSM. В этой статье на 52 страницы много математики, я пока ещё не углублялся, но авторы написали прекрасную серию постов, которую можно читать вместо статьи: Часть 1: В целом про SSD модель (https://goombalab.github.io/blog/2024/mamba2-part1-model/)\ Часть 2: Теория с математическим разбором фреймворка SSD (https://goombalab.github.io/blog/2024/mamba2-part2-theory/) Часть 3: Алгоритмическая часть и код (https://goombalab.github.io/blog/2024/mamba2-part3-algorithm/) Часть 4: Оптимизации на уровне системы для large-scale training, файнтюнинга и инференса (https://goombalab.github.io/blog/2024/mamba2-part4-systems/) Оригинальная Мамба была selective SSM (S6) с диагональной структурой. SSD идёт ещё дальше и ограничивает диагональ матрицы A, теперь все элементы должны иметь одинаковые значения (то есть это скаляр умноженный на единичную матрицу). Старая Mamba применялась к каждому входному каналу отдельно, новая обрабатывает сразу много каналов (например 64) одной общей рекуррентностью. Это повышает вычислительную эффективность и даёт более быстрое обучение. При этом теоретически у Mamba-2 меньшая выразительность, чем у Mamba-1, и первая Мамба также может быть лучше в инференсе. Это всё пока не изучалось и ждёт своих исследователей. Фреймворк позволяет перенести устоявшиеся для внимания техники на архитектуры SSM и реализовать аналоги голов (MHA) в SSM. Также архитектура блока сети (Mamba блок, https://t.center/gonzo_ML/2153) слегка изменена относительно SSM. Появляется grouped-value attention в структуре голов и все data-dependent проекции (параметры A,B,C в SSM) теперь получаются параллельно со входом X, а не последовательно как раньше. Реализованы разные оптимизации, в частности модель сделана Tensor Parallelism-friendly.