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

TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

yet another attempt на упрощение скейлинга трансформеров

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

было и стало (f - нелинейность)

O = f(X W_1) W_2
→
O = f(X K^T) V

плюс перед нелинейностью GeLU добавили л2 нормализацию по строчкам

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

а скейлиться лучше в увеличении параметров вот так: было n x embedding_dim. хотим увеличить → делаем (n+m) x embedding_dim, где инитим новые параметры нулями или как-то еще

нлп и вижн трансформеры (в частности особенно pythia) обыгрывает при примерно том же наборе параметров, так еще и флопсов меньше требуется

идея довольно интересная но при этом немножко топорная. смущает разве что отсутствие аблаций по типу “давайте попробуем обычный гелу без какой-либо нормализации”.

👀LINK