Репост из: КПД
Метод
Общий фреймворк выглядит следующим образом:
▶️ Функция преобразования данных f_X (X - входные данные)
▶️ Функция конструкции смешивающей матрицы f_M L x L, которая может быть постоянной или зависеть от входов
▶️ Результат sequence mixing имеет вид f_M(f_X (X))
Далее авторы вводят термин Sequence Aligned Matrices (SAM, еще один… 🥱) означающий, что матрица смешивания зависит от входных данных. Такие sequece миксеры хороши с одной стороны тем, что более адаптивно подстраиваются под входы, и, кроме того, работают с последовательностями разной длины.
Авторы рассматривают разные механизмы из литературы:
⭐️ MLP-Mixer, S4, H3, Monarch, Vandermonde и Cauchy миксеры - не SAM
⭐️ Attention, Linear Attention, S6, SSD - SAM
Потому хороший двунаправленный sequence mixer должен быть SAM, и представлять собой некоторую структурированную матрицу. В частности, предлагаются способы сделать Vandermonde и Cauchy миксеры зависящими от входов, но основной упор делается на прокачку SSD (не твердотельного жесткого диска, а механизма в Mamba-2!) под двунаправленность.
Напомним, что SSD во второй мамбе является полуразложимыми (semi-separable) матрицей - каждый блок является низкоранговой матрицей. Для двунаправленности можно было бы чередовать слои SSD (один бегущий слева направо, другой справа налево), но здесь предлагают использование одной матрицы смешивания, такой что любой ее блок в верхне-треугольной и нижне-треугольной части является низкоранговой матрицей. Иначе говоря, получается нечто типа суммы исходной SSD из Mamba-2 (нижнетреугольной матрицы) и транспонированной по длине последовательности (верхнетреугольной матрицы) и диагональной части. Такие матрицы называют квазиразложимыми (QS). Данная модификация требует всего пары дополнительных строчек в реализации по сравнению с исходным SSD слоем (shift - сдвиг на один элемент, flip - разворот последовательности задом наперед, DX - диагональная добавка).
QS(X) = shift(SS(X)) + flip(shift(SS(flip(X)))) + DX
Называют гидрой, потому что много голов, как в SSD, и звучит красиво 😹.
Эксперименты
Метод валидируют на задаче Masked Language Modelling, где в качестве бейзлайнов берется BERT, обученный по рецепту от MosaicML, и иные варианты sequence mixerов из литературы. Для оценки качества моделей смотрят на валидационную кросс-энтропию на C4 (на train set которого обучают) и точность на бенчах из GLUE. Все модели имеют размер порядка 70M параметров (несколько меньше, чем BERT-Base), так что хрен вам SOTA на LMSYS. Hydra модели глубже трансформеров примерно в 2️⃣ раза при примерно том же числе параметров.
SAM модели стабильно опережают свои не SAM версии (Toeplitz, Cauchy, Vandermonde с параметрами, зависящими от входа, заметно точнее версии с обучаемыми, не зависящими). Hydra (естесна), лучше всех, и на втором месте любимый нами трансформер. Однако, памятуя о недавнем результате MobileLLM, где более глубокие и тонкие трансформеры, оказываются лучше по качеству более коротких и жирных при том же числе параметров, задаешься вопросом - можно ли устранить разрыв в качестве за счет изменения конфигурации трансформера 🤔?
Исходная Mamba-2 не очень сильна в MLM, так как умеет обрабатывать информацию только в одном направлении, и предложенный способ (Hydra) лучше вариантов с суммой, конкатенацией, и перемножением результатов двух мамб-2.
Далее метод проверяют на задаче классификации ImageNet-1k, где обучают модели размера порядка ViT-Base (87-91M параметров) и Hydra опережает ViT-B, Hyena и S4. Однако, ViT бейзлайн вызывает вопросы, ибо согласно их результатам ViT-B имеет top-1 точность 78.8%, а его EMA 80.6%, в то время, как c рецептом обучения из Swin, на который они ссылаются (унаследованный в свою очередь из DeiT) выдает 81.8% (их лучший результат 81.6%)
Общий фреймворк выглядит следующим образом:
▶️ Функция преобразования данных f_X (X - входные данные)
▶️ Функция конструкции смешивающей матрицы f_M L x L, которая может быть постоянной или зависеть от входов
▶️ Результат sequence mixing имеет вид f_M(f_X (X))
Далее авторы вводят термин Sequence Aligned Matrices (SAM, еще один… 🥱) означающий, что матрица смешивания зависит от входных данных. Такие sequece миксеры хороши с одной стороны тем, что более адаптивно подстраиваются под входы, и, кроме того, работают с последовательностями разной длины.
Авторы рассматривают разные механизмы из литературы:
⭐️ MLP-Mixer, S4, H3, Monarch, Vandermonde и Cauchy миксеры - не SAM
⭐️ Attention, Linear Attention, S6, SSD - SAM
Потому хороший двунаправленный sequence mixer должен быть SAM, и представлять собой некоторую структурированную матрицу. В частности, предлагаются способы сделать Vandermonde и Cauchy миксеры зависящими от входов, но основной упор делается на прокачку SSD (не твердотельного жесткого диска, а механизма в Mamba-2!) под двунаправленность.
Напомним, что SSD во второй мамбе является полуразложимыми (semi-separable) матрицей - каждый блок является низкоранговой матрицей. Для двунаправленности можно было бы чередовать слои SSD (один бегущий слева направо, другой справа налево), но здесь предлагают использование одной матрицы смешивания, такой что любой ее блок в верхне-треугольной и нижне-треугольной части является низкоранговой матрицей. Иначе говоря, получается нечто типа суммы исходной SSD из Mamba-2 (нижнетреугольной матрицы) и транспонированной по длине последовательности (верхнетреугольной матрицы) и диагональной части. Такие матрицы называют квазиразложимыми (QS). Данная модификация требует всего пары дополнительных строчек в реализации по сравнению с исходным SSD слоем (shift - сдвиг на один элемент, flip - разворот последовательности задом наперед, DX - диагональная добавка).
QS(X) = shift(SS(X)) + flip(shift(SS(flip(X)))) + DX
Называют гидрой, потому что много голов, как в SSD, и звучит красиво 😹.
Эксперименты
Метод валидируют на задаче Masked Language Modelling, где в качестве бейзлайнов берется BERT, обученный по рецепту от MosaicML, и иные варианты sequence mixerов из литературы. Для оценки качества моделей смотрят на валидационную кросс-энтропию на C4 (на train set которого обучают) и точность на бенчах из GLUE. Все модели имеют размер порядка 70M параметров (несколько меньше, чем BERT-Base), так что хрен вам SOTA на LMSYS. Hydra модели глубже трансформеров примерно в 2️⃣ раза при примерно том же числе параметров.
SAM модели стабильно опережают свои не SAM версии (Toeplitz, Cauchy, Vandermonde с параметрами, зависящими от входа, заметно точнее версии с обучаемыми, не зависящими). Hydra (естесна), лучше всех, и на втором месте любимый нами трансформер. Однако, памятуя о недавнем результате MobileLLM, где более глубокие и тонкие трансформеры, оказываются лучше по качеству более коротких и жирных при том же числе параметров, задаешься вопросом - можно ли устранить разрыв в качестве за счет изменения конфигурации трансформера 🤔?
Исходная Mamba-2 не очень сильна в MLM, так как умеет обрабатывать информацию только в одном направлении, и предложенный способ (Hydra) лучше вариантов с суммой, конкатенацией, и перемножением результатов двух мамб-2.
Далее метод проверяют на задаче классификации ImageNet-1k, где обучают модели размера порядка ViT-Base (87-91M параметров) и Hydra опережает ViT-B, Hyena и S4. Однако, ViT бейзлайн вызывает вопросы, ибо согласно их результатам ViT-B имеет top-1 точность 78.8%, а его EMA 80.6%, в то время, как c рецептом обучения из Swin, на который они ссылаются (унаследованный в свою очередь из DeiT) выдает 81.8% (их лучший результат 81.6%)