gonzo-обзоры ML статей

[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://www.group-telegram.com/ms/gonzo_ML.com/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://www.group-telegram.com/ms/gonzo_ML.com/2153) слегка изменена относительно SSM. Появляется grouped-value attention в структуре голов и все data-dependent проекции (параметры A,B,C в SSM) теперь получаются параллельно со входом X, а не последовательно как раньше. Реализованы разные оптимизации, в частности модель сделана Tensor Parallelism-friendly.

www.group-telegram.com/ms/gonzo_ML.com/2718

4.5K viewsJun 3, 2024 at 21:51

[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://www.group-telegram.com/ms/gonzo_ML.com/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://www.group-telegram.com/ms/gonzo_ML.com/2153) слегка изменена относительно SSM. Появляется grouped-value attention в структуре голов и все data-dependent проекции (параметры A,B,C в SSM) теперь получаются параллельно со входом X, а не последовательно как раньше. Реализованы разные оптимизации, в частности модель сделана Tensor Parallelism-friendly.

BY gonzo-обзоры ML статей