Griffin: Het combineren van gelineaire recurrenties met lokale aandacht voor efficiënte taalmodelen

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

February 29, 2024

Auteurs: 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

Samenvatting

Recurrente neurale netwerken (RNN's) hebben snelle inferentie en schalen efficiënt op lange sequenties, maar ze zijn moeilijk te trainen en lastig op te schalen. Wij stellen Hawk voor, een RNN met gegate lineaire recurrenties, en Griffin, een hybride model dat gegate lineaire recurrenties combineert met lokale aandacht. Hawk overtreft de gerapporteerde prestaties van Mamba op downstream taken, terwijl Griffin de prestaties van Llama-2 evenaart ondanks dat het getraind is op meer dan 6 keer minder tokens. We laten ook zien dat Griffin kan extrapoleren op sequenties die aanzienlijk langer zijn dan die tijdens de training zijn gezien. Onze modellen evenaren de hardware-efficiëntie van Transformers tijdens de training, en tijdens de inferentie hebben ze een lagere latentie en aanzienlijk hogere doorvoer. We schalen Griffin op tot 14B parameters en leggen uit hoe we onze modellen kunnen sharden voor efficiënte gedistribueerde training.

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.