В этой статье я хочу остановиться на разборе предложенной мной архитектуры декодера и тех вариантов, с которыми я сравниваю её в исследовании, но сделать это проще и интуитивнее, чем в самой работе. На мой взгляд, существующие объяснения архитектур декодеров часто подаются разрозненно. Каждый подход описывают отдельно, без общей опоры. А ведь всё можно свести к одному фундаменту, и тогда становятся гораздо заметнее как сильные стороны каждого решения, так и их ограничения.

Для начала приведу все необходимые ссылки.

Само исследование:https://arxiv.org/abs/2604.18580Код:https://github.com/LibratioAI/sessa

Структура статьи:

1. Transformer
2. S4D и Mamba
3. Sessa: Selective State Space Attention

Transformer

Начнём с классики — трансформера, но посмотрим на него под другим углом: не на то, как он устроен и какие у него компоненты, а на то, как идейно к нему прийти. Для этого мы воспользуемся свёрткой. В нашем случае для простоты будем оперировать дискретными функциями со скалярными элементами.

К примеру, рассмотрим на входе последовательность из пяти элементов (слова, эмбеддинги и т. п.) и ядро из трех элементов. В классическом определении ядро свёртки отражается по нулю (переворачивается), но будем считать, что это уже сделано.

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

Пусть вход

Тогда выходы миксера на позициях могут быть, например,

В decoder-only модели такой выход в позицииtзатем используется для предсказания следующего элементаНо дальше нас будет интересовать именно сам миксер: какую информацию из префикса он извлекает в позицииt. Поэтомудалее обозначает выход миксера, а не финальное предсказание следующего токена.

И, как видим, у нас есть две проблемы: ядро фиксировано по коэффициентам и имеет фиксированную длину окна. Такая свёртка в теории цифровой обработки сигналов называется FIR (Finite Impulse Response, конечный импульсный отклик) фильтром: его реакция на единичный импульс длится конечное число шагов, потому что ядро имеет конечную поддержку.

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

и выход миксера в позицииt:

Например, если

При этом есть одно "но": даже если сделать коэффициенты зависящими от входа, это само по себе не снимает ограничения по длине. Пока мы работаем в фиксированном окне радиуса, оператор остаётся локальным и видит только последниеэлементов. А если расширить суммирование на большее число позиций (вплоть до всего контекста), то при увеличении длины меняется режим вычисления: в сумме появляется больше слагаемых, и без нормировки выход начинает дрейфовать по масштабу и по распределению вкладов. Поэтому веса нужно сделать устойчивыми к росту, например, так, чтобы их сумма была фиксированной. И здесь естественно появляется softmax: он нормирует сырые коэффициенты в веса, суммирующиеся к единице, хотя, конечно, это ещё не attention как механизм как таковой.

В нашем примере:

Например, при

Несмотря на то что мы исправили две ключевые проблемы, такая система всё ещё не справляется со многими задачами. Возьмём, например, простую подзадачу для миксера: хотим, чтобы выход в позицииtизвлекал предыдущий вход,

Это означает, что на каждом шаге вес на позициидолжен быть больше, чем на позициии на позицииВ частности, достаточно следующих требований:

Но softmax монотонен по своим аргументам: еслито обязательноЗначит из требований получаем:

Получается противоречие: одновременно требуетсяи

Если интуитивно, то когда оценки зависят только от важности токена самого по себе, то есть фактически только от источника, модель не может менять предпочтение, кого копировать, в зависимости от шагаИменно здесь и появляется необходимость в том, чтобы оценки зависели от того, кто читает и на что смотрит. И вот тут уже естественно появляется классический attention (с позиционным кодированием): вместо важности токена самой по себе, мы нормируем совместимость парыt

Гдеотвечает за позиционную информацию

И на нашем примере с конечным окноммы смотрим только наПри

В общем виде, когдаявляется вектором, вместо скаляраиспользуются линейные отображения для построенияис произвольнымВот и всё: теперь к этому attention добавляем MLP и получаем слой трансформера.

А теперь поговорим немного о свойствах этой архитектуры.

Чтобы честно сравнить память трёх архитектур, нужно выбрать согласованный режим анализа. Потому что можно подобрать такие параметры, при которых каждая из них будет лучше другой, или наоборот. В контексте исследования я рассматривал контролируемый общий режим, в котором моделям сложно сфокусироваться на определённом элементе входной последовательности, то есть ломаем sharp retrieval. В случае трансформера это диффузный режим. А для самой оценки нужно понять, как выход в моментtчувствителен к токену в прошломиспользуя якобианы.

Для создания диффузного режима в трансформере достаточно, чтобы разброс логитов, то есть разница между максимальным и минимальным значением перед softmax, был ограничен константой, не растущей с длиной контекста. Тогда softmax не может устойчиво сконцентрировать почти всю массу на одном токене, и веса остаются распределёнными по большому числу позиций.

При этом я доказываю оценки для режимов:

Однослойный:Если рассматривать режим без заморозки коэффициентов attention, верхняя оценка равна:

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

Многослойный:рассматриваем стек изслоёв attention и смотрим полный градиент от выхода на позицииtк входу на позиции

В этом режиме верхняя оценка равна:

Глубина добавляет дополнительные маршруты через промежуточные позиции и даёт логарифмическое усиление по сравнению с одним слоем:

Но при фиксированной глубиневлияние старых токенов всё равно затухает:

А нижняя оценка, как и в однослойном режиме, остаётся тривиальной: она больше или равна нулю.

В исследовании я также ввожу терминыone path / many paths(пути) иone-hop / multi-hop(переходы). Они интуитивно объясняют, почему в диффузных режимах память ведёт себя по-разному у трёх архитектур. В контексте трансформера и механизма внимания удобно представить маршрутизацию влияния как ориентированный ациклический граф (DAG) по временным индексам: для трансформера это граф с прямыми рёбрами(при), где вес ребра задаётся вниманиемТогда влияние извtвнутри одного слоя реализуется одним переходом (one-hop) и по одному пути (one path). Так как интуитивно, обратиться к определённому элементу attention может только одним способом в одном слое — дать ему максимальный вес, то есть один путь и один шаг.

S4D и Mamba

Целью данных SSM является исправить недостаток attention, а именно квадратичную сложность вычислений при росте длины последовательности, при этом оставаясь примерно на том же уровне качества на длинных контекстах. И в целом цель хорошая, если бы не несколько фундаментальных "но", которые мы в этом разделе обсудим. Но, чтобы прийти интуитивно к Mamba, мы сделаем это через промежуточный этап в виде S4D, возьмём нашу предыдущую LTI-систему и посмотрим на неё с другой стороны.

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

В цифровой обработке сигналов такая система известна как IIR-фильтр, а эта форма записи — как разностное уравнение. И теперь у нас есть кроме forward-ветки, которую мы рассматривали в предыдущем пункте ещё и feedback. И отталкиваясь от предыдущего пункта, мы учитываем как в нашем примере три последнихи взвешиваем их коэффициентами

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

И тут всплывает довольно важный момент про устойчивость системы. Есть два близких, но разных способа на это смотреть.

BIBO-стабильность:если вход ограничен, то выход тоже должен оставаться ограниченным. Обычно это рассматривают при нулевом начальном состоянии.

Внутренняя устойчивость:если вход равен нулю, но начальное состояние ненулевое, то собственная динамика системы не должна разгоняться.

Рассмотрим простой случай:

Коэффициент обратной связи равен 1.

Если для наглядности взять постоянный входто при

То есть выход растёт линейно. Поэтому приглобальную BIBO-стабильность в общем случае гарантировать нельзя.

Коэффициент по модулю больше 1.

Еслито система неустойчива. Это можно увидеть двумя способами.

Во-первых, как нарушение внутренней устойчивости. Пусть вход равен нулю:

но начальное состояние ненулевое. Тогда

Приэто растёт экспоненциально.

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

Тогда выход будет содержать множитель

Прион становится неограниченным. Значит ограниченный вход может породить неограниченный выход.

Коэффициент по модулю меньше 1.

Еслито собственная динамика затухает:

Кроме того, импульсный отклик такой системы имеет вид геометрической прогрессии. Поэтому вклад прошлого суммируется с экспоненциально убывающими весами, и при ограниченном входе выход остаётся ограниченным. В этом случае система является BIBO-стабильной.

Но при этом нам никто не мешает использовать произвольное число элементов для взвешивания в обратной связи. В общем виде разностное уравнение будет выглядеть как:

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

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

Здесь​ являются параметрами перехода состояния, и их не следует путать с коэффициентамив разностной форме

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

Рассмотрим два состояния без входа:Для простоты возьмём

Подставляем динамику в

Теперь у нас система:

Умножим первое уравнение на​и вычтем из второго:

Сдвинем на один шаг назад:

Но по динамике второго состоянияДомножим предыдущее равенство на

Значит правые части совпадают, и получаем:

И теперь если мы рассмотрим разностное уравнение для IIR фильтра:

При этомявляются полюсами, а значит достаточно, чтобы они были по модулю меньше 1. Что и описывается в исследовании, посвящённом S4D.

Теперь, чтобы увидеть, что такое Mamba, достаточно просто сделать коэффициенты внутри состояний зависимыми от входа, то есть:

При этом создатели архитектуры Mamba заложили подобную селективность:

И тогда, если представить миксер Mamba в виде разностного уравнения, то мы получим:

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

Рассмотрим одно состояние:

В Mamba параметрызависят только от текущего токена

Умножая наполучаем

Локальная невырожденность:

Тогда существует единственный набор коэффициентов

такой, что

Следовательно,

По построению

Правая часть является линейной комбинацией входов

Поэтому существуют коэффициентытакие что

Следовательно,

Так какдиагональна,

Поэтомузависит только от параметров на индексах

а матрицазависит только от параметров на индексах

Следовательно

зависит только от параметров на индексах

Но параметры индексовзависят только отпоэтому

Здесь все множителизависят только от параметров на индексах

а коэффициентызависят от

Следовательно

Поэтому, приполучаем

И теперь можно перейти к анализу свойств данной архитектуры.

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

У Mamba есть режим, который в одном исследовании назвалиfreezing time, суть в том, что еслито, асостояние сохраняется, а новый вход не попадает. А значит оно может в этом случае хранить его сколь угодно долго, напоминая режим attention с one-hot. Для контролируемого режима достаточно предположить, что этот механизм плохо работает и не может корректно определить, где заморозить время, а где нет (failedfreezing time), то есть ломаем sharp retrieval. И для этого режима в исследовании я показываю оценки также для однослойного и многослойного вариантов.

Однослойный.В этом режиме верхняя оценка равна:

Нижняя оценка в этом режиме тривиальна и будет больше либо равна нулю.

Многослойный.Рассматриваем стек изслоёв Mamba и полный градиент. В этом режиме верхняя оценка равна:

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

Но при фиксированной глубиневлияние старых токенов всё равно затухает:

Нижняя оценка, как и в однослойном режиме, остаётся тривиальной: она больше либо равна нулю.

Но допустим мы рассмотрим отдельно реальную ситуацию с freezing time, и тут есть некоторые подводные камни, о которых я рассказываю в исследовании. Механизм заморозки опирается на распознавание по входным представлениям, а посколькувычисляется как функция от входа, то когда разделимость сигнала падает и появляется шум в последовательности, распознавание может стать ненадёжным: модель не поддерживает режимна всём релевантном промежутке, и возникают шагисущественно больше нуля. Поскольку реальные последовательности часто содержат шум, подобный режим может встречаться довольно часто.

Ну и напоследок перед тем, как перейти к Sessa, представим маршрутизацию влияния в Mamba в виде ориентированного ациклического графа. Но он устроен иначе, чем в трансформере: это цепочка с рёбрами только между соседними шагамигде вес ребра задаётся переходом состояния. Тогда влияние извtвнутри одного слоя проходит по одному пути(one path), но требуетпоследовательных переходов (multi-hop). Интуитивно, при фиксированной модели и произвольной последовательности, чтобы обратиться к далёкому элементумодели нужно пройти через цепочку обратной связи. Как мы разбирали выше, динамика на каждом шаге учитывает только конечное число элементов в forward- и feedback-ветках. Поэтому получается один путь, но много переходов. Это и объясняет, почему Mamba часто имеет экспоненциальное затухание.

Sessa: Selective State Space Attention

Теперь перейдём к архитектуре, которую я предлагаю как альтернативу Mamba и Transformer.

Идея заключается в том, чтобы объединить два направления: добавить feedback и одновременно сделать forward- и feedback-ветки адаптивными к текущему токену и длине последовательности. Иными словами, добавить feedback в Transformer используя attention.

Сам слой Sessa имеет вид:

после чего изстроятся две ветки миксера: forward-ветка

и feedback-ветка

или, в параллельной матричной форме,

Финальный выход слоя равен

Если записать миксер в разностной форме, как в предыдущем пункте, получаем:

Подставляяво вторую формулу, получаем

Теперь обозначим

В случае трансформера это эквивалентно

В Mamba коэффициенты зависят от конечного окна входов, а более далёкая история передаётся через состояние. В Sessa же feedback-коэффициенты напрямую строятся по всему префиксу.

Теперь перейдём к свойствам.

Для сравнения памяти в случае Sessa применяется диффузный режим, аналогичный трансформеру. Для этого режима однослойный и многослойный вариант:

Однослойный:Если рассматривать полный градиент одного блока без заморозки весов attention:

У Sessa в этом режиме хвостсто есть медленнее хвоста Transformer

Нижняя оценка здесь тоже тривиальна и будет больше либо равна нулю.

Многослойный:рассматриваем стек изслоёв Sessa и полный градиент от выхода на позицииtк входу на позицииВ этом режиме верхняя оценка равна:

Эквивалентно, если смотреть на доминирующий по порядку член при фиксированной глубине, то можно писать:

То есть глубина здесь даёт не просто логарифмическое усиление, как у трансформера, а более сильное накопление вклада по мере композиции слоёв. Поэтому затухание может быть заметно медленнее.

При этом, если

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

Нижняя оценка, как и в однослойном режиме, остаётся тривиальной: она больше либо равна нулю.

Теперь перейдём к рассмотрению маршрутизации влияния в Sessa как ориентированный ациклический граф.

Допустим, мы применяем к начальному состояниюрекурентно миксер Sessa получая еще дополнительно три состоянияИсходя из формул миксера, это будет выглядеть следующим образом:

подставимв

подставимив

подставимв

Если посмотреть на коэффициент прито есть влияние из момента 0 в момент 3:

Можно выделить many paths и multi-hop:

Путьс одним переходом:,

Путьc двумя переходами:и путьтакже с двумя переходами:,

Путьс тремя переходами:

За счёт того, что модель может добираться до элементов различными путями, это и даёт более медленное затухание.

Если более интуитивно, то за счёт attention в feedback-части модель может обратиться к любому предыдущему feedback-шагу. Этот шаг, в свою очередь, уже учитывает предыдущие шаги, а также все входы через attention в forward-части. Поэтому и возникает множество путей с возможностью реализации множества переходов.

Гибкое селективное извлечение (flexible selective retrieval):

Это один из наиболее интересных и практически важных результатов. Суть в том, что в каждой из трёх моделей можно настроить параметры так, чтобы она без проблем извлекала необходимый вход. Но при анализе селективного извлечения лучше не смотреть только на якобианы, потому что память, транспорт нужного сигнала и селективность — это немного разные вещи. То есть память, которую мы рассматривали выше, отвечает за то, дошло ли влияние вообще. Транспорт нужного сигнала отвечает за то, дошёл ли именно нужный компонент. А селективность — за то, победил ли нужный источник всех конкурентов. И все эти три вещи нужны для данного анализа.

При этом у селективности есть три профиля:

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

Замороженный профиль. Качество retrieval не деградирует с расстоянием, то есть по мере роста лага разница между текущим токеном и остальной массой остаётся одного и того же порядка.

Возрастающий профиль.Разница между текущим токеном и остальной массой может не просто не падать, а расти, то есть идет накопление преимущества.

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

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

Каждая из архитектур с одним слоем может реализовать только затухающий профиль. Но когда слоёв становится два или больше, только Sessa может реализовать замороженный и возрастающий профиль. То есть диффузный режим для неё не проблема. Интуитивно это связано с тем, что в многослойной модели возникает столько путей ко входу, что, несмотря на диффузность внимания, сама структура памяти может сфокусировать всю эту массу на нужном входе.

Позиционное кодирование и универсальная аппроксимация.

Sessa, как и Mamba, может кодировать APE внутри себя без явных таблиц или экстраполяции, в отличие от трансформера. Следовательно, её можно обучать без встроенного позиционного кодирования, но стоит учитывать, что к нему или к его альтернативному варианту модель должна ещё прийти в процессе обучения.

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

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

Что касается Sessa, в исследовании я был больше сфокусирован на теории. Но в ближайших планах обучить модель на несколько миллиардов параметров и посмотреть, что там уже с длинным контекстом на больших масштабах.

Если вы хотите поддержать это исследование, я был бы очень благодарен за вашголосна Hugging Face. Спасибо!

В этой статье я хочу остановиться на разборе предложенной мной архитектуры декодера и тех вариантов, с которыми я сравниваю её в исследовании, но сделать это проще и интуитивнее, чем в самой работе. На мой взгляд, существующие объяснения архитектур декодеров часто подаются разрозненно. Каждый подход описывают отдельно, без общей опоры. А ведь всё можно свести к одному фундаменту, и тогда становятся гораздо заметнее как сильные стороны каждого решения, так и их ограничения.

Для начала приведу все необходимые ссылки.

Само исследование:https://arxiv.org/abs/2604.18580Код:https://github.com/LibratioAI/sessa

Структура статьи:

1. Transformer
2. S4D и Mamba
3. Sessa: Selective State Space Attention

Transformer

Начнём с классики — трансформера, но посмотрим на него под другим углом: не на то, как он устроен и какие у него компоненты, а на то, как идейно к нему прийти. Для этого мы воспользуемся свёрткой. В нашем случае для простоты будем оперировать дискретными функциями со скалярными элементами.

К примеру, рассмотрим на входе последовательность из пяти элементов (слова, эмбеддинги и т. п.) и ядро из трех элементов. В классическом определении ядро свёртки отражается по нулю (переворачивается), но будем считать, что это уже сделано.

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

Пусть вход

Тогда выходы миксера на позициях могут быть, например,

В decoder-only модели такой выход в позицииtзатем используется для предсказания следующего элементаНо дальше нас будет интересовать именно сам миксер: какую информацию из префикса он извлекает в позицииt. Поэтомудалее обозначает выход миксера, а не финальное предсказание следующего токена.

И, как видим, у нас есть две проблемы: ядро фиксировано по коэффициентам и имеет фиксированную длину окна. Такая свёртка в теории цифровой обработки сигналов называется FIR (Finite Impulse Response, конечный импульсный отклик) фильтром: его реакция на единичный импульс длится конечное число шагов, потому что ядро имеет конечную поддержку.

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

и выход миксера в позицииt:

Например, если

При этом есть одно "но": даже если сделать коэффициенты зависящими от входа, это само по себе не снимает ограничения по длине. Пока мы работаем в фиксированном окне радиуса, оператор остаётся локальным и видит только последниеэлементов. А если расширить суммирование на большее число позиций (вплоть до всего контекста), то при увеличении длины меняется режим вычисления: в сумме появляется больше слагаемых, и без нормировки выход начинает дрейфовать по масштабу и по распределению вкладов. Поэтому веса нужно сделать устойчивыми к росту, например, так, чтобы их сумма была фиксированной. И здесь естественно появляется softmax: он нормирует сырые коэффициенты в веса, суммирующиеся к единице, хотя, конечно, это ещё не attention как механизм как таковой.

В нашем примере:

Например, при

Несмотря на то что мы исправили две ключевые проблемы, такая система всё ещё не справляется со многими задачами. Возьмём, например, простую подзадачу для миксера: хотим, чтобы выход в позицииtизвлекал предыдущий вход,

Это означает, что на каждом шаге вес на позициидолжен быть больше, чем на позициии на позицииВ частности, достаточно следующих требований:

Но softmax монотонен по своим аргументам: еслито обязательноЗначит из требований получаем:

Получается противоречие: одновременно требуетсяи

Если интуитивно, то когда оценки зависят только от важности токена самого по себе, то есть фактически только от источника, модель не может менять предпочтение, кого копировать, в зависимости от шагаИменно здесь и появляется необходимость в том, чтобы оценки зависели от того, кто читает и на что смотрит. И вот тут уже естественно появляется классический attention (с позиционным кодированием): вместо важности токена самой по себе, мы нормируем совместимость парыt

Гдеотвечает за позиционную информацию

И на нашем примере с конечным окноммы смотрим только наПри

В общем виде, когдаявляется вектором, вместо скаляраиспользуются линейные отображения для построенияис произвольнымВот и всё: теперь к этому attention добавляем MLP и получаем слой трансформера.

А теперь поговорим немного о свойствах этой архитектуры.

Чтобы честно сравнить память трёх архитектур, нужно выбрать согласованный режим анализа. Потому что можно подобрать такие параметры, при которых каждая из них будет лучше другой, или наоборот. В контексте исследования я рассматривал контролируемый общий режим, в котором моделям сложно сфокусироваться на определённом элементе входной последовательности, то есть ломаем sharp retrieval. В случае трансформера это диффузный режим. А для самой оценки нужно понять, как выход в моментtчувствителен к токену в прошломиспользуя якобианы.

Для создания диффузного режима в трансформере достаточно, чтобы разброс логитов, то есть разница между максимальным и минимальным значением перед softmax, был ограничен константой, не растущей с длиной контекста. Тогда softmax не может устойчиво сконцентрировать почти всю массу на одном токене, и веса остаются распределёнными по большому числу позиций.

При этом я доказываю оценки для режимов:

Однослойный:Если рассматривать режим без заморозки коэффициентов attention, верхняя оценка равна:

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

Многослойный:рассматриваем стек изслоёв attention и смотрим полный градиент от выхода на позицииtк входу на позиции

В этом режиме верхняя оценка равна:

Глубина добавляет дополнительные маршруты через промежуточные позиции и даёт логарифмическое усиление по сравнению с одним слоем:

Но при фиксированной глубиневлияние старых токенов всё равно затухает:

А нижняя оценка, как и в однослойном режиме, остаётся тривиальной: она больше или равна нулю.

В исследовании я также ввожу терминыone path / many paths(пути) иone-hop / multi-hop(переходы). Они интуитивно объясняют, почему в диффузных режимах память ведёт себя по-разному у трёх архитектур. В контексте трансформера и механизма внимания удобно представить маршрутизацию влияния как ориентированный ациклический граф (DAG) по временным индексам: для трансформера это граф с прямыми рёбрами(при), где вес ребра задаётся вниманиемТогда влияние извtвнутри одного слоя реализуется одним переходом (one-hop) и по одному пути (one path). Так как интуитивно, обратиться к определённому элементу attention может только одним способом в одном слое — дать ему максимальный вес, то есть один путь и один шаг.

S4D и Mamba

Целью данных SSM является исправить недостаток attention, а именно квадратичную сложность вычислений при росте длины последовательности, при этом оставаясь примерно на том же уровне качества на длинных контекстах. И в целом цель хорошая, если бы не несколько фундаментальных "но", которые мы в этом разделе обсудим. Но, чтобы прийти интуитивно к Mamba, мы сделаем это через промежуточный этап в виде S4D, возьмём нашу предыдущую LTI-систему и посмотрим на неё с другой стороны.

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

В цифровой обработке сигналов такая система известна как IIR-фильтр, а эта форма записи — как разностное уравнение. И теперь у нас есть кроме forward-ветки, которую мы рассматривали в предыдущем пункте ещё и feedback. И отталкиваясь от предыдущего пункта, мы учитываем как в нашем примере три последнихи взвешиваем их коэффициентами

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

И тут всплывает довольно важный момент про устойчивость системы. Есть два близких, но разных способа на это смотреть.

BIBO-стабильность:если вход ограничен, то выход тоже должен оставаться ограниченным. Обычно это рассматривают при нулевом начальном состоянии.

Внутренняя устойчивость:если вход равен нулю, но начальное состояние ненулевое, то собственная динамика системы не должна разгоняться.

Рассмотрим простой случай:

Коэффициент обратной связи равен 1.

Если для наглядности взять постоянный входто при

То есть выход растёт линейно. Поэтому приглобальную BIBO-стабильность в общем случае гарантировать нельзя.

Коэффициент по модулю больше 1.

Еслито система неустойчива. Это можно увидеть двумя способами.

Во-первых, как нарушение внутренней устойчивости. Пусть вход равен нулю:

но начальное состояние ненулевое. Тогда

Приэто растёт экспоненциально.

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

Тогда выход будет содержать множитель

Прион становится неограниченным. Значит ограниченный вход может породить неограниченный выход.

Коэффициент по модулю меньше 1.

Еслито собственная динамика затухает:

Кроме того, импульсный отклик такой системы имеет вид геометрической прогрессии. Поэтому вклад прошлого суммируется с экспоненциально убывающими весами, и при ограниченном входе выход остаётся ограниченным. В этом случае система является BIBO-стабильной.

Но при этом нам никто не мешает использовать произвольное число элементов для взвешивания в обратной связи. В общем виде разностное уравнение будет выглядеть как:

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

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

Здесь​ являются параметрами перехода состояния, и их не следует путать с коэффициентамив разностной форме

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

Рассмотрим два состояния без входа:Для простоты возьмём

Подставляем динамику в

Теперь у нас система:

Умножим первое уравнение на​и вычтем из второго:

Сдвинем на один шаг назад:

Но по динамике второго состоянияДомножим предыдущее равенство на

Значит правые части совпадают, и получаем:

И теперь если мы рассмотрим разностное уравнение для IIR фильтра:

При этомявляются полюсами, а значит достаточно, чтобы они были по модулю меньше 1. Что и описывается в исследовании, посвящённом S4D.

Теперь, чтобы увидеть, что такое Mamba, достаточно просто сделать коэффициенты внутри состояний зависимыми от входа, то есть:

При этом создатели архитектуры Mamba заложили подобную селективность:

И тогда, если представить миксер Mamba в виде разностного уравнения, то мы получим:

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

Рассмотрим одно состояние:

В Mamba параметрызависят только от текущего токена

Умножая наполучаем

Локальная невырожденность:

Тогда существует единственный набор коэффициентов

такой, что

Следовательно,

По построению

Правая часть является линейной комбинацией входов

Поэтому существуют коэффициентытакие что

Следовательно,

Так какдиагональна,

Поэтомузависит только от параметров на индексах

а матрицазависит только от параметров на индексах

Следовательно

зависит только от параметров на индексах

Но параметры индексовзависят только отпоэтому

Здесь все множителизависят только от параметров на индексах

а коэффициентызависят от

Следовательно

Поэтому, приполучаем

И теперь можно перейти к анализу свойств данной архитектуры.

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

У Mamba есть режим, который в одном исследовании назвалиfreezing time, суть в том, что еслито, асостояние сохраняется, а новый вход не попадает. А значит оно может в этом случае хранить его сколь угодно долго, напоминая режим attention с one-hot. Для контролируемого режима достаточно предположить, что этот механизм плохо работает и не может корректно определить, где заморозить время, а где нет (failedfreezing time), то есть ломаем sharp retrieval. И для этого режима в исследовании я показываю оценки также для однослойного и многослойного вариантов.

Однослойный.В этом режиме верхняя оценка равна:

Нижняя оценка в этом режиме тривиальна и будет больше либо равна нулю.

Многослойный.Рассматриваем стек изслоёв Mamba и полный градиент. В этом режиме верхняя оценка равна:

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

Но при фиксированной глубиневлияние старых токенов всё равно затухает:

Нижняя оценка, как и в однослойном режиме, остаётся тривиальной: она больше либо равна нулю.

Но допустим мы рассмотрим отдельно реальную ситуацию с freezing time, и тут есть некоторые подводные камни, о которых я рассказываю в исследовании. Механизм заморозки опирается на распознавание по входным представлениям, а посколькувычисляется как функция от входа, то когда разделимость сигнала падает и появляется шум в последовательности, распознавание может стать ненадёжным: модель не поддерживает режимна всём релевантном промежутке, и возникают шагисущественно больше нуля. Поскольку реальные последовательности часто содержат шум, подобный режим может встречаться довольно часто.

Ну и напоследок перед тем, как перейти к Sessa, представим маршрутизацию влияния в Mamba в виде ориентированного ациклического графа. Но он устроен иначе, чем в трансформере: это цепочка с рёбрами только между соседними шагамигде вес ребра задаётся переходом состояния. Тогда влияние извtвнутри одного слоя проходит по одному пути(one path), но требуетпоследовательных переходов (multi-hop). Интуитивно, при фиксированной модели и произвольной последовательности, чтобы обратиться к далёкому элементумодели нужно пройти через цепочку обратной связи. Как мы разбирали выше, динамика на каждом шаге учитывает только конечное число элементов в forward- и feedback-ветках. Поэтому получается один путь, но много переходов. Это и объясняет, почему Mamba часто имеет экспоненциальное затухание.

Sessa: Selective State Space Attention

Теперь перейдём к архитектуре, которую я предлагаю как альтернативу Mamba и Transformer.

Идея заключается в том, чтобы объединить два направления: добавить feedback и одновременно сделать forward- и feedback-ветки адаптивными к текущему токену и длине последовательности. Иными словами, добавить feedback в Transformer используя attention.

Сам слой Sessa имеет вид:

после чего изстроятся две ветки миксера: forward-ветка

и feedback-ветка

или, в параллельной матричной форме,

Финальный выход слоя равен

Если записать миксер в разностной форме, как в предыдущем пункте, получаем:

Подставляяво вторую формулу, получаем

Теперь обозначим

В случае трансформера это эквивалентно

В Mamba коэффициенты зависят от конечного окна входов, а более далёкая история передаётся через состояние. В Sessa же feedback-коэффициенты напрямую строятся по всему префиксу.

Теперь перейдём к свойствам.

Для сравнения памяти в случае Sessa применяется диффузный режим, аналогичный трансформеру. Для этого режима однослойный и многослойный вариант:

Однослойный:Если рассматривать полный градиент одного блока без заморозки весов attention:

У Sessa в этом режиме хвостсто есть медленнее хвоста Transformer

Нижняя оценка здесь тоже тривиальна и будет больше либо равна нулю.

Многослойный:рассматриваем стек изслоёв Sessa и полный градиент от выхода на позицииtк входу на позицииВ этом режиме верхняя оценка равна:

Эквивалентно, если смотреть на доминирующий по порядку член при фиксированной глубине, то можно писать:

То есть глубина здесь даёт не просто логарифмическое усиление, как у трансформера, а более сильное накопление вклада по мере композиции слоёв. Поэтому затухание может быть заметно медленнее.

При этом, если

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

Нижняя оценка, как и в однослойном режиме, остаётся тривиальной: она больше либо равна нулю.

Теперь перейдём к рассмотрению маршрутизации влияния в Sessa как ориентированный ациклический граф.

Допустим, мы применяем к начальному состояниюрекурентно миксер Sessa получая еще дополнительно три состоянияИсходя из формул миксера, это будет выглядеть следующим образом:

подставимв

подставимив

подставимв

Если посмотреть на коэффициент прито есть влияние из момента 0 в момент 3:

Можно выделить many paths и multi-hop:

Путьс одним переходом:,

Путьc двумя переходами:и путьтакже с двумя переходами:,

Путьс тремя переходами:

За счёт того, что модель может добираться до элементов различными путями, это и даёт более медленное затухание.

Если более интуитивно, то за счёт attention в feedback-части модель может обратиться к любому предыдущему feedback-шагу. Этот шаг, в свою очередь, уже учитывает предыдущие шаги, а также все входы через attention в forward-части. Поэтому и возникает множество путей с возможностью реализации множества переходов.

Гибкое селективное извлечение (flexible selective retrieval):

Это один из наиболее интересных и практически важных результатов. Суть в том, что в каждой из трёх моделей можно настроить параметры так, чтобы она без проблем извлекала необходимый вход. Но при анализе селективного извлечения лучше не смотреть только на якобианы, потому что память, транспорт нужного сигнала и селективность — это немного разные вещи. То есть память, которую мы рассматривали выше, отвечает за то, дошло ли влияние вообще. Транспорт нужного сигнала отвечает за то, дошёл ли именно нужный компонент. А селективность — за то, победил ли нужный источник всех конкурентов. И все эти три вещи нужны для данного анализа.

При этом у селективности есть три профиля:

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

Замороженный профиль. Качество retrieval не деградирует с расстоянием, то есть по мере роста лага разница между текущим токеном и остальной массой остаётся одного и того же порядка.

Возрастающий профиль.Разница между текущим токеном и остальной массой может не просто не падать, а расти, то есть идет накопление преимущества.

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

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

Каждая из архитектур с одним слоем может реализовать только затухающий профиль. Но когда слоёв становится два или больше, только Sessa может реализовать замороженный и возрастающий профиль. То есть диффузный режим для неё не проблема. Интуитивно это связано с тем, что в многослойной модели возникает столько путей ко входу, что, несмотря на диффузность внимания, сама структура памяти может сфокусировать всю эту массу на нужном входе.

Позиционное кодирование и универсальная аппроксимация.

Sessa, как и Mamba, может кодировать APE внутри себя без явных таблиц или экстраполяции, в отличие от трансформера. Следовательно, её можно обучать без встроенного позиционного кодирования, но стоит учитывать, что к нему или к его альтернативному варианту модель должна ещё прийти в процессе обучения.

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

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

Что касается Sessa, в исследовании я был больше сфокусирован на теории. Но в ближайших планах обучить модель на несколько миллиардов параметров и посмотреть, что там уже с длинным контекстом на больших масштабах.

Если вы хотите поддержать это исследование, я был бы очень благодарен за вашголосна Hugging Face. Спасибо!