Аттеншн - та же добыча радия.
В грамм добыча - в год труды.
Изводишь единого токена ради
Тысячи тонн словесной руды.
Маяковский (Поэзия - та же добыча радия)
мы уже упоминали как картезиа дистиллирует трансформеры в ссм, правда прошлая статья вызывала некоторые вопросы с точки зрения реализации
на этот раз получилось у них еще круче - в лламе 3.1 заменили аттеншн блоки на mamba блоки и получилась лламба (1б 3б и 8б параметров)
обучали в стиле MOHAWK по разным стадиям, где в knowledge distillation манере постепенно фиттили блоки все бОльшего размера: от выхода аттеншна (ссм) до конечных логитов модели
теперь только на замену аттеншну пришел дискретный блок мамбы 2, а дискретный он потому, что получает матрицу А напрямую из инпута, без махинаций с параметров дельта (который теоретически и дискретизирует стейт спейс модели для непрерывного потока времени), потому что задача здесь заключается в дистилляции аттеншна, который умеет оперировать только над дискретными объектами (токенами). короче подогнали
ну и здесь еще есть отличие от классической мамбы 2 в том, что оригинал является multi-value, а здесь же опять из-за подгона сделали multi-head + в оверолл архитектуре получается меньше мамба 2 блоков → больше батч сайз и меньше оверхеда на лаунч кернелов
по результатам даже у 8б модели потребовалось 12б токенов чтобы выдавать сравнимый результат на дефолтных бенчмарках (намного меньше (0.1%) чем у других гибридных или около того методов, из интересного есть еще наверное lolcats от Christopher Re, которые сворачивают софтмакс аттеншн в линейный через свд разложения)
еще через MLX погоняли модели на айфоне и макбуке что прикольно
судя по всему этот стартап в ближайшее время будет показывать много интересного своими on-device подходами
👀 link, code