Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

читала ли ты на ночь про связь между трансформерами и ссм, Дездемона?

мы уже писали про альтернативные трансформеру методы (здесь и здесь, здесь и здесь, и даже здесь). но не писали про мамбу

Три Дао и Альберт Гу решили продолжить развивать мамбу и выкатили мамбу-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

Restructuring Vector Quantization with the Rotation Trick

Christopher Ré продолжает делать интересное, правда на уже не в рамках HazyResearch, и на сей раз он добрался до VQ-VAE

авторы поднимают вопрос того, что по факту Straight-Through Estimator (STE) во время бекворда третирует производную кодбук вектора по отношению к энкодер аутпуту как identity matrix

но в этом может заключаться и загвоздка, когда бекворд доходит до энкодера, ибо по сути STE не привносит свой импакт по изменению параметров энкодера → положение энкодер аутпутов по отношению к кодбукам не влияет на градиент энкодера

авторы предлагают оч красиво и просто такое исправить - через матрицу поворота, а в итоге метод назвали rotation trick. он не меняет аутпуты при форвард пассе, но когда дело доходит до бекворда, то на градиент уже влияет как раз эта матрица поворота, которая высчитывается в соответствии с энкодер аутпутом и кодбук вектором

и вычисляют к тому же они не наивно, а сразу через хаусхолдера, при том, по заверениям авторов, разницы по времени с классическим вк вае нет

сравнивали на видео и картинках + еще добавили адверсариал сетап вместе с диффузией. много где обыгрывает классику по всем трем фронтам - реконструкции, “утилизации” кодбука и ошибке квантизации

связывают это с тем, что при ротейшн трике, в зависимости от направления градиента, точки в одном регионе (то есть те, которые сматчились на операции квантизации) могут разъезжаться друг от друга, либо же наоборот становится ближе. при STE, обычно, наблюдается только первый сценарий

супер интересно и полезно, а еще оставляет нотку философских размышлений тот факт, что по результатам лучше сохранять исходный угол между кодбук вектором и градиентом, нежели направление вектора

👀LINK

LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

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

в чем основная суть? как оказывается, в основном, финальные токены для генерации формируются не на последнем слое, а где-то внутри модели с дальнейшей возможностью дальше поменяться, но такое происходит не во всех случаях (в среднем токену для полного формирования надо 23.45 слоев из 32-ух слойной лламы).

назвали это early exit - вместо того, чтобы ждать последнего слоя эмбеддингов, мы просто на вход ЛМ голове подаем эмбеддинги из промежуточных слоев (или просто слоев пораньше)

но есть небольшой питфолл в том, что голова научена только декодить эмбеддинги с последнего слоя → возникает необходимость дотренивать голову уметь декодить эмбеддинги с любого слоя

наивно это сделать не получается, чтобы работало + мы хотели бы инферить на как можно мЕньшем количестве слоев ⇒ чтобы токены генерились не на последних слоях. из этого прикрутили layer dropout с гипером, который экспоненциально увеличивается с глубиной сетки

да еще и забустили кроссэнтропию с поправкой на стимуляцию к ealry exit, где циркулярно или постепенно брали эмбеды из ранних слоев и осуществляли early exit

окей, раньше из сетки выходить мы поняли как, но как убедиться в том, что аккураси по сравнению со всей моделью не падает (падает незначительно)? - их навороченный спекулятивный декодинг (уже писали тута и тута), который они назвали self-speculative decoding

в этом декодинге они прогоняют всю модель для предикта следующего токена для каждого драфт токена, и все это в один форвард пасс + на самом деле надо прогонять не всю модель, а только оставшиеся после early exit’a (так еще и кв кэш сохраняется с этих драфт токенов). по экспам наблюдается спидап в 1.86 раз по заверениям авторов

выглядит красиво (и вроде практично) с дополнительными затратами в доп претрейне/файнтюнинге и подборе гиперов под дропаут и новый лосс вместе с лернинг рейтом и тд

код есть, кроме трейна, в принципе как и многого другого

👀LINK

Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators

seems like сделали фаундейшн модель для вычислительной гидродинамики (computational fluid dynamics)

а задача стояла именно так: есть Neural Operator’ы - модельки, которые обучаются маппингу между разными функциональными пространствами. такое особенно четко нужно в физике, в частности для диффуров по времени, как гидродинамика, например

вроде понятно, но чем дальше в лес, тем интереснее: функциональные пространства в случае CFD представить можно двумя разными способами - по Эйлеру и по Лагранжу (в виде гридов или произвольного количества партиклов)

UPT же пытается объединить сразу две эти репрезентации и это благодаря трансформеру в латентном пространстве. есть энкодер/декодер (которые маппят из гридов/партиклов в латенты и обратно) и «аппроксиматор», каждый форвард которого проносит диффур систему на таймстеп вперед

также энкодер и декодер навешаны доп лоссами на реконструкцию по разным временным фазам диффура (до аппроксиматора и после), чтобы каждый из модулей пайплайна изолированно выполнял ту работу, которую пытались внедрить авторы. вроде получилось

оверолл - снижено количество по памяти (ибо все почти происходит в латентном пространстве и не приурочено ни к Эйлеру, ни к Лагранжу) + этот трансформер может выдавать нормальные аутпуты как при выборе мелкого горизонта времени, так и относительно большого

👀LINK