Mamba, Jamba, Ху...Hymba!🤣NVIDIA представило новую гибридную модель
Hymba (Химба, а не то,что ты подумал).
Как известно, SSM модели имеют свои архитектурные преимущества по отношению к трансформерам. На самом деле одно преимущество - линейная сложность от длины входной последовательности. При этом SSM модели имеют проблемы с "запоминанием" контекста при большой длине последовательности и не смогли показать иных плюсов кроме сложности.
В новой работе авторы показывают гибридизацию внимания и SSM модели (фото архитектуры будет ниже). Однако, в отличие от предыдущих работ, мы не имеем последовательно блок SSM, далее блок внимания, а их использование происходит параллельно. В таком сетапе, мы к сожалению не уходим полностью от квадратичной сложности по длине последовательности, в отличии от последовательного гибрида.
Однако,авторы уповают на иные "сильные" стороны. Во-первых, используется аналогия с работой мозга. Для запоминания четких образов/мгновенных снимков событий в модели работает SSM блок, в то же время для контекстной связности и длинной памяти используется full attention блок. Также авторы добавили в архитектуру мета-токены, которые используются, по мнению авторов,как навигаторы какие области архитектуры активировать. Интересное замечание,что масштаб векторов вложений после SSM блока на порядки выше, векторов после блоков внимания. Для их выравнивания используется gating norm— привет, нормализация, в классик трансформере после конкатенации multi-head attention. Выходы блоков усредняются.
Кстати, тут же рождается идея у Дяди. И тебе шкет советую положить руки на стол
🔫, а не бежать делать по ней папиру.
Почему бы вдобавок к gating norm, не
использовать gating routing как в MoE. В итоге получаем умную маршрутизацию экспертов SSM блоки или full attention в зависимости от того,что нужно хранить/обработать. Да еще и время инфера апнет (то linear time, то n^2, а не всегда n^2).
🫥Ладно,помечтали и хватит, вернёмся к статье. Там есть еще пару хак. Связка ssm+mha представляет собой hymba block, такой блок чередуется с SWA блоком (аля как в
longformer было). В таком блоке скользящее внимание юзает не все токены, а работает с окном внимания,которое идет "внахлёст". Так, видимо, хотят сократить/устремить сложность от O(N^2) до условно O(N*logN). Далее еще также используется KV-caching,НО с обобщёнными KV-матрицами между блоками.
В статье есть блок с ablation study,советую глянуть. Ну и конечно метрики, где сравниваются small-LM, кажется, до 3b параметров (метрики ниже). Как видим, модель с таким комбо вниманием, проявляет себя лучше на задачах,где важны знания о мире,фактология и тп. Логично, авторы ж работают вокруг концептов "памяти".
Так, вроде, всё. Читаем статью, модельки лежат
тут,
git. Дядя пошел папиру про MoE Hymba писать (но это не точно).
😏 
