Метод
Почти с момента выхода Attention было предложено много альтернатив с субквадратичной сложностью. Если убрать softmax в Attention - операцию можно посчитать за линейное по длине последовательности число операций. Аналогично, SSM (S4, Mamba-1/2), DeltaNet линейно масштабируются с ростом числа токенов. Тем не менее, несмотря на успехи на отдельных задачах - вытеснить трансформер с пьедестала никому не удалось.
Попытка запихнуть весь контекст в скрытое состояние фиксированного размера, по всей видимости, фундаментально ограничивает модель в возможности знать все в длинном контексте.
Потому предлагается промежуточный вариант - логарифмическая по памяти и времени операция, являющаяся надстройкой над одним из линейных механизмов attention. Токены разбиваются на корзинки с экспоненциально растущим числом токенов. Самые свежие токены обычно важнее для предсказания следующего, потому в одной корзине меньше токенов, и, соответственно, их вес больше, а с отдалением от текущей позиции размер корзинок растет, а вклад индивидуальных токенов убывает. Log-Linear attention сначала вычисляет линейный attention по корзинкам, а затем суммирует с некоторыми обучаемыми коэффициентами результат каждой корзинки (коэффициенты предсказывает отдельная MLP). Число корзинок растет логарифмически с длиной - потому и имеем O(L log L) как итоговую сложность операции. Для эффективной реализации используют деревья Фенвика.
Log-Linear Attention можно представить в виде структурированной матрицы HODLR (Hierarchically Off-Diagonal Low-Rank), где диагональные блоки нижнетреугольные, а внедиагональная часть состоит из блоков ранга-1, где размер блока растет с удалением от диагонали.
Log-Linear Attention можно применить как поверх Linear Attention, так и Mamba-2 и DeltaNet. И для всего написаны соответствующие кернелы.
Эксперименты
Для валидации метода авторы обучают модельки на синтетических и реальных задачах.
На синтетике Log-Linear модификация значительно улучшает качество DeltaNet на MQAR (достать несколько элементов из контекста).
Далее авторы обучают в сопоставимых условиях (700-800M параметров, 50B токенов из Long-Data-Collections с длиной последовательности 16k) Transformer, DeltaNet и Mamba-2 (без и с Log-Linear надстройки). Log-Linear дает небольшой прирост поверх DeltaNet и Mamba-2.
По скорости инференса на длинных контекстах Log-Linear Mamba-2 медленнее Mamba-2 (в ~2 раза на 64k/128k токенах), но быстрее Attention.
На Needle-in-Haystack в бенче, где нужно достать один токен Log-Linear хорош, в multi-key/multi-value задачах Log-Linear лучше линейных бейзлайнов, но хуже Attention.
На LongBench где-то дает прирост, а где-то не дает.
За что уважение авторам - они не утверждают, что предложенная модификация бьет все и всея, а стараются более менее честно все замерить.
Выводы
С точки зрения математики все красиво - вообще вопросов нет, и уважение 🤠 мастерам написания ядер на CUDA. В целом выглядит как неплохой промежуточный вариант между Attention и линейными по длине альтернативами, но как будто требует валидации бюджетах и размерах моделей ближе к production-grade.
>>Click here to continue<<