читала ли ты на ночь про связь между трансформерами и ссм, Дездемона?
мы уже писали про альтернативные трансформеру методы (здесь и здесь, здесь и здесь, и даже здесь). но не писали про мамбу
Три Дао и Альберт Гу решили продолжить развивать мамбу и выкатили мамбу-2, где SSM заменился на SSD (state space duality)
а если поконкретнее, то не очень хитрыми махинациями все алгоритмы стейт спейс моделей можно представить в виде матричных перемножений (которые мы упоминали и здесь), что помогает легче понимать структуру модели и ее характеристики + это уже легче подогнать под аналогию с аттеншном и тем самым понимать преимущества ссм над траснформерами (и аналогично с недостатками)
в общем стейт спейс модели сводятся к матрикс миксерам с semiseparable matrices (которые предоставляют сами по себе широкий выбор структуризации, ибо их прикол состоит в том, что они нижнетреугольные и что каждая подматрица в них лоуранговая). здесь они обогащают понятие каузальной маски, которая участвует в финальном аутпуте вместе с линейным аттеншном
при том в стейт спейс моделях, в матричной формулировке, эта каузальная маска L так же включает в себя относительное позиционное кодирование, которое зависит от инпута - тот самый селективный скан в мамбе, завязанный на матрице А. матрица же А отвечает за структурированную проекцию предыдущего хидден стейта на хидден в конкретном таймтсепе, обычно ее делали диагональной, но в ссд теперь она представляется identity * some_scalar (то есть ее можно представить просто в виде скаляра)
и такая потеря в вариативности повышает efficiency. но как быть с потерей качества? поскольку проведены аналогии с аттном, то и мы можем попробовать взять оттуда концепции и применить их на ссд, как те же разные головы аттеншна, вычисления по которым можно распараллелить ⇒ повышаем капасити модели (при том намного по сравнению с обычной мамбой, ибо в ее случае была одна голова с меньшей размерностью, чем это позволяет мамба-2). а так получается можно между ссд и аттеншн терминами провести следующий мэтч:
(A, C, B, X) → (L, Q, K, V), где Х - инпут
так же авторы обобщают концепцию мульти-хеда на multi-query/key/value attention. как оказывается, эффективный тип ssm можно по-другому назвать как multi-value attn
помимо этого архитектуру можно свести на привычный мульти-гпу сетап и иметь только один all-reduce блок, как и аттн. разве что нормировку надо заменять на групп норму и ставить количество групп = tensor parallelism degree (чтобы не замедлять вычисления) + блочное разбиение матриц и все такое
в папире становится жоско быстрее стандартной мамбы, да и даже флеш аттеншн перегоняет при сек лене ≥ 2k
есть неявное замечание, что авторы видимо четко осознали, что раз ссм можно свести только к линейному аттеншну, который может превосходить софтмакс аттеншн только на скейле моделек до ~= 7б, то и нет смысла превзойти ванильный аттн и есть резон только превосходить linear attention-like архитектуры по efficiency и перформансу (что в каком-то смысле очень даже уважаемо)
итого - табличное сравнение по памяти и скорости на последней картинке, мамба-2 старается привнести объединение бенефитов из двух миров аттеншна и ссм. мб займет свою нишу на модельках относительно небольшого скейла
👀LINK