Griffin: Mixing Gated Linear Recurrences with Local Attention...
Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models
Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando De Freitas, Caglar Gulcehre
Статья: https://arxiv.org/abs/2402.19427
На днях вышла открытая RecurrentGemma (https://arxiv.org/abs/2404.07839), построенная на архитектуре Griffin. Сам Griffin был опубликован DeepMind’ом в конце февраля 2024. Разберём же его.
Работа вертится вокруг нового рекуррентного блока, RG-LRU, на котором построены архитектуры Hawk (чередование RG-LRU и MLP) и Griffin (чередование MLP со смесью RG-LRU и локального внимания). Hawk при этом бьёт Mamba (https://t.me/gonzo_ML/2148) аналогичного размера, а Griffin обходит Llama-2, обучаясь на вшестеро меньших данных.
Архитектура строится на повторяющихся residual blocks, похожих на используемые в pre-norm трансформерах: (RMSNorm + Temporal mixing block) и (RMSNorm + MLP block), оба с residual connection поверх.
В качестве MLP block используется gated блок аналогичный GeGLU имени Ноама Шазира (https://arxiv.org/abs/2002.05202, в текущей работе его назвали GeGeLU): с двумя ветвями размерности M*D каждая (в работе выбрано M=3, то есть эмбеддинги расширяются), в одной ветви сидит нелинейность GeLU, а в другой считаются коэффициенты для поэлементного умножения, после которого слитые ветви обрабатываются ещё одним линейным слоем.
Самая интересная и вариабельная часть -- Temporal mixing block. Их три варианта: 1) global Multi-Query Attention (MQA), 2) local (sliding-window) MQA 3) и новый рекуррентный блок.
Вариант 1 (MQA, https://arxiv.org/abs/1911.02150 тоже имени Ноама Шазира) это замена классического Multi-Head Attention (MHA), где K и V общие для всех голов. Используются позиционные эмбеддинги RoPE.
Вариант 2 с локальным вниманием (оно же sliding window attention) аналогичен локальному вниманию в Longformer (https://t.me/gonzo_ML/294). Окно локального внимания установлено в 1024 токена.
Наконец вариант 3 напоминает блок из Мамбы (https://t.me/gonzo_ML/2153), где тоже две ветви, в одном из которых всё тот же GeLU как в MLP, а в другом одномерная свёртка + RG-LRU слой.
Сам RG-LRU (Real-Gated Linear Recurrent Unit) -- это развитие LRU (https://t.me/gonzo_ML/1734) с двумя добавленными гейтами, не зависящими от предыдущего рекуррентного состояния, только от входа. Input gate 𝑖_t аналогичен таковому из LSTM, он фильтрует или масштабирует вход. Второй гейт, recurrence gate 𝑟_t, нов и может приближённо интерполировать между стандартным апдейтом LRU из оригинальной работы и предыдущим скрытым состоянием, тем самым отбрасывая входные данные и сохраняя информацию из прошлого. В приложении A подробнее разбирается поведение рекуррентного гейта.
Для инициализации RG-LRU _не_ используются полиномы по типу HiPPO или дискретизация по типу SSM. Также не используются комплексные числа в рекуррентности, как это было в LRU. В приложении B также рассмотрен комплекснозначный вариант под названием CG-LRU (Complex-Gated Linear Recurrent Unit). Он более выразителен, но языковому моделированию на практике не помогает.
Результаты интересные. Рассматриваются три варианта:
1) MQA-Transformer в качестве бейзлайна
2) Hawk с тем же residual и MLP как у трансформерного бейзлайна, но с рекуррентным блоком с RG-LRU в качестве temporal mixing block.
3) Griffin с тем же residual и MLP как у трансформерного бейзлайна, но с миксом рекуррентных и локальных MQA блоков (через каждые два residual блока c RG-LRU, один блок с локальным вниманием).
Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando De Freitas, Caglar Gulcehre
Статья: https://arxiv.org/abs/2402.19427
На днях вышла открытая RecurrentGemma (https://arxiv.org/abs/2404.07839), построенная на архитектуре Griffin. Сам Griffin был опубликован DeepMind’ом в конце февраля 2024. Разберём же его.
Работа вертится вокруг нового рекуррентного блока, RG-LRU, на котором построены архитектуры Hawk (чередование RG-LRU и MLP) и Griffin (чередование MLP со смесью RG-LRU и локального внимания). Hawk при этом бьёт Mamba (https://t.me/gonzo_ML/2148) аналогичного размера, а Griffin обходит Llama-2, обучаясь на вшестеро меньших данных.
Архитектура строится на повторяющихся residual blocks, похожих на используемые в pre-norm трансформерах: (RMSNorm + Temporal mixing block) и (RMSNorm + MLP block), оба с residual connection поверх.
В качестве MLP block используется gated блок аналогичный GeGLU имени Ноама Шазира (https://arxiv.org/abs/2002.05202, в текущей работе его назвали GeGeLU): с двумя ветвями размерности M*D каждая (в работе выбрано M=3, то есть эмбеддинги расширяются), в одной ветви сидит нелинейность GeLU, а в другой считаются коэффициенты для поэлементного умножения, после которого слитые ветви обрабатываются ещё одним линейным слоем.
Самая интересная и вариабельная часть -- Temporal mixing block. Их три варианта: 1) global Multi-Query Attention (MQA), 2) local (sliding-window) MQA 3) и новый рекуррентный блок.
Вариант 1 (MQA, https://arxiv.org/abs/1911.02150 тоже имени Ноама Шазира) это замена классического Multi-Head Attention (MHA), где K и V общие для всех голов. Используются позиционные эмбеддинги RoPE.
Вариант 2 с локальным вниманием (оно же sliding window attention) аналогичен локальному вниманию в Longformer (https://t.me/gonzo_ML/294). Окно локального внимания установлено в 1024 токена.
Наконец вариант 3 напоминает блок из Мамбы (https://t.me/gonzo_ML/2153), где тоже две ветви, в одном из которых всё тот же GeLU как в MLP, а в другом одномерная свёртка + RG-LRU слой.
Сам RG-LRU (Real-Gated Linear Recurrent Unit) -- это развитие LRU (https://t.me/gonzo_ML/1734) с двумя добавленными гейтами, не зависящими от предыдущего рекуррентного состояния, только от входа. Input gate 𝑖_t аналогичен таковому из LSTM, он фильтрует или масштабирует вход. Второй гейт, recurrence gate 𝑟_t, нов и может приближённо интерполировать между стандартным апдейтом LRU из оригинальной работы и предыдущим скрытым состоянием, тем самым отбрасывая входные данные и сохраняя информацию из прошлого. В приложении A подробнее разбирается поведение рекуррентного гейта.
Для инициализации RG-LRU _не_ используются полиномы по типу HiPPO или дискретизация по типу SSM. Также не используются комплексные числа в рекуррентности, как это было в LRU. В приложении B также рассмотрен комплекснозначный вариант под названием CG-LRU (Complex-Gated Linear Recurrent Unit). Он более выразителен, но языковому моделированию на практике не помогает.
Результаты интересные. Рассматриваются три варианта:
1) MQA-Transformer в качестве бейзлайна
2) Hawk с тем же residual и MLP как у трансформерного бейзлайна, но с рекуррентным блоком с RG-LRU в качестве temporal mixing block.
3) Griffin с тем же residual и MLP как у трансформерного бейзлайна, но с миксом рекуррентных и локальных MQA блоков (через каждые два residual блока c RG-LRU, один блок с локальным вниманием).
Источник: gonzo-обзоры ML статей
2024-04-15 08:36:34