Гриффин: Смешивание линейных рекуррент с воротной функцией и локального внимания для эффективных языковых моделей

Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models

February 29, 2024

Авторы: 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

cs.AI

Аннотация

Рекуррентные нейронные сети (RNN) обладают быстрым выводом и эффективно масштабируются на длинных последовательностях, но их сложно обучать и трудно масштабировать. Мы предлагаем Hawk, RNN с затворными линейными повторениями, и Griffin, гибридную модель, которая сочетает в себе затворные линейные повторения с локальным вниманием. Hawk превосходит отчетную производительность Mamba на последующих задачах, в то время как Griffin соответствует производительности Llama-2, несмотря на то, что обучалась на более чем в 6 раз меньшем количестве токенов. Мы также показываем, что Griffin способна экстраполировать последовательности значительно длиннее тех, что были использованы во время обучения. Наши модели соответствуют аппаратной эффективности Трансформеров во время обучения, а во время вывода они имеют более низкую задержку и значительно большую пропускную способность. Мы масштабируем Griffin до 14 миллиардов параметров и объясняем, как разделить наши модели для эффективного распределенного обучения.

English

Recurrent neural networks (RNNs) have fast inference and scale efficiently on long sequences, but they are difficult to train and hard to scale. We propose Hawk, an RNN with gated linear recurrences, and Griffin, a hybrid model that mixes gated linear recurrences with local attention. Hawk exceeds the reported performance of Mamba on downstream tasks, while Griffin matches the performance of Llama-2 despite being trained on over 6 times fewer tokens. We also show that Griffin can extrapolate on sequences significantly longer than those seen during training. Our models match the hardware efficiency of Transformers during training, and during inference they have lower latency and significantly higher throughput. We scale Griffin up to 14B parameters, and explain how to shard our models for efficient distributed training.