Базированная модель, часть 1
или почему RWKV/Mamba/RetNet не работали, но заработают.

Есть такая группа в Стэнфорде, HazyResearch. Это они сделали первые SSM (state space models) моделии их современные версии (H3, Hyena). Ну и всякие мелочи вроде FlashAttention.

На этот раз ребята начали с того, что обучили трансформеры / H3 / Hyena / RWKV не очень больших одинаковых размеров на 10B токенах из The Pile. Трансформеры выиграли! 🤔

Возникает два вопроса: "кто виноват?" и "что делать?". На первый вопрос отвечает Zoology, на второй вопрос отвечает Based.

Кто виноват?
Zoology: Measuring and Improving Recall in Efficient Language Models, статья, пост 1, пост 2.

А виноваты оказались... повторяющиеся N-граммы 🤨
То есть словосочетания, которые уже встречались в контексте. Ну, не все, только достаточно редкие, потому что частотные N-граммы любая модель хорошо запоминает. Пример такого редкого словосочетания: фамилия и имя какого-то человека, которые несколько раз встречаются в одном документе, и никогда больше не встречаются в обучающем корпусе.

Предсказание последнего токена такой повторяющейся N-граммы и вызывает трудности у моделей без внимания. В статье такие токены называют associatve recall hits, AR hits, и по ним отдельно считают перплексию. Для Гиены и RWKV разница перплексии на этих токенах полностью покрывает разницу в общей перплексии.

Самое забавное, что мы уже такое видели несколькими постами ранее, тут. Хронологически "Repeat After Me" вышла двумя месяцами позже, но эту статью они явно не читали, иначе бы их статья скорее всего не вышла бы 😂

Теперь можно сформулировать, какая задача мешает не-трансформерам захватить мир. Задача называется multi-query associative recall (MQAR), и заключается она в поиске нескольких "иголок". Упрощенно она выглядит так:

Вход:
A 4 B 3 C 6 F 1 E 2

Запрос:
A ? C ? F ? E ? B ?

Ожидаемый выход:
4 6 1 2 3

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

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

Есть и намёки, как это исправить! Нужно всего лишь добавить капельку внимания, то есть делать гибриды 😂
Однако ж нам не нужна полная маска внимания, и мы можем либо точечно влиять на AR hits ("programmatic selection" метод в статье), либо на основе входов предсказывать, для каких k токенов нужно включить внимание. Втыкают 3 слоя внимания на 6% параметров, и этого достаточно, чтобы добить бОльшую часть перплексии.

И наконец, 30 страниц доказательств! 😱
Что вообще доказывают:
- Обзывают все используемые в статье не-трансформерные модели вентильными свёртками (gated convolutions).
- Вводят архитектурный блок BaseConv: y = linear(x) ⊙ conv(h, x), где x - вход, y - выход, ⊙ - покомпонентное произведение, h - обучаемый фильтр.
- Доказывают эквивалентность между Гиеной и BaseConv в смысле симуляции за константное количество слоёв.
- Доказывают, что BaseConv может симулировать арифметические схемы, то есть вычислять многочлены.
- Доказывают эквивалентность между RetNet и BaseConv, в которой BaseConv нужно в log(d) раз больше слоёв.
- Выводят теоретическую оценку сложности для BaseConv на нашей задаче, MQAR, через построение алгоритма с параллельным бинпоиском и оценку его сложности.

К сожалению, я не слишком хорош в математике, чтобы всё это нормально осознать.

Вывод
RWKV не работает, совсем без внимания никак, гибриды победят, синтетические бенчмарки рулят.

Второй пост будет, и будет про непосредственно Based.

Базированная модель, часть 2

Первая часть: ссылка. В ней мы определили основную проблему не-трансформерных моделей: запоминание повторяющихся N-грамм. Остался второй вопрос...

Что делать?
Simple linear attention language models balance the recall-throughput tradeoff, статья, пост 1, пост 2, твиттер, репо.

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

Для начала, напомню: основная цель большинства не-трансформеров сделать быстрее при том же качестве. В этой и прошлой статье у трансформеров лучше перплексия. Единственное место, где это не так — первый пост про Based. Кажется, они замерили там некорректные числа, но к выпуску основной статьи исправились 💀
Ещё раз: за исключением первого старого поста со старыми числами, здесь нигде не заявляется превосходство по перплексии над трансформерами. Зато заявляется превосходство по пропускной способности и скорости генерации, даже против FlashAttention-2 😳

У Based есть 3 типа слоёв, смешивающих эмеддинги токенов:
- BaseConv из предыдущего поста
- Линейное внимание
- Полное внимание на маленьком скользящем окне
В одном слое только один конкретный смеситель, например в первом слое линейное внимание, во втором скользящее окно, в третьем BaseConv, и так далее.
Кроме смесителей по токенам есть и стандартные MLP для смешивания каналов, одинаковые для каждого типа слоёв.

Что за линейное внимание? Сама концепция не нова: убираем софтмакс и получаем рекуррентную форму. Софтмакс убираем через аппроксимацию exp(qk) через ряд Тейлора до второго порядка малости: 1 + qk + (qk)^2 / 2. Подробнее об этом написано в другой статье, но тут всё равно показывают, что другие варианты аппроксимаций хуже.

Псевдокод вычисления:

qk0 = [1]

# Считаем слагаемые от q
q_first = [q1, ..., qd] # q
q_second = [q1 * q1, ..., q1 * qd, ..., qd * q1, ... qd * qd] # qq^T
q_new = cat(qk0, q_first,  q_second / sqrt(2))

# Считаем слагаемые от k
k_first = [k1, ..., kd] # k
k_second =  [k1 * k1, ..., k1 * kd, ..., kd * k1, ... kd * kd]  # kk^T
k_new = cat(qk0, k_first,  k_second / sqrt(2))

# Разложение экспоненты до второго порядка малости: 1 + qk + (qk)^2 / 2
y = (q_new * k_new).sum()

Мотивация добавления чего-либо к BaseConv всё та же: улучшить вспоминание информации, добавив входозависимые смесители. При этом, в отличие от трансформеров, у которых KV-кэш растёт линейно от длины последовательности, мы можем варьировать объём доступной нам памяти через размер скрытого состояния в линейном внимании. И таким образом можем разменивать точность вспоминания на скорость.

В статье есть фундаментальная теорема: решение MQAR (а значит и языкового моделирования на реальных текстах) требует размер скрытого состояния, линейный от длины последовательности. Для внимания таким скрытым состоянием является KV-кэш, который имеет размер O(Nd). Для Based это KV-состояние для числителя и K-сотояние для знаменателя в симуляции софтмакса, с общей размерностью O(d^3), но в статье эту размерность несколько уменьшают проекциями. Это всё означает, что чем длиннее последовательность мы хотим корректно обрабатывать, тем жирнее нам нужно делать размерность модели 😭

Для того, чтобы конкурировать с эффективными реализациями внимания, ребята написали кастомные ядра, которые стараются выполнять операции в SRAM, быстрой памяти GPU, по аналогии с FlashAttention. Получилось и правда быстро 😘

В итоге имеем модель, которая по качеству на уровне трансформеров, но работает в 25 раз быстрее.

Что мне кажется сомнительным во всей этой истории:
- А что произошло с перплексией в первом посте? Что авторы изначально сделали не так?
- Нафига было разрабатывать BaseConv, чтобы потом вернуться к линейному вниманию? Зачем нам теперь вообще BaseConv нужен? Нет, в статье конечно есть секция C, в которой добавление BaseConv обосновывается выигрышем в перплексии, но это ничего не объясняет.
- Почему только 1b параметров? Спонсоров вроде много, люди богатые.

P.S. Немного покринжевал со слова "смеситель" в этом контексте, но решил оставить. Тоже смешивает же.