Griffin: Combinando Recorrências Lineares com Portas e Atenção Local para Modelos de Linguagem Eficientes

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

February 29, 2024

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

Resumo

Redes neurais recorrentes (RNNs) possuem inferência rápida e escalam eficientemente em sequências longas, mas são difíceis de treinar e complicadas de escalar. Propomos o Hawk, uma RNN com recorrências lineares com portas, e o Griffin, um modelo híbrido que mistura recorrências lineares com portas e atenção local. O Hawk supera o desempenho relatado do Mamba em tarefas subsequentes, enquanto o Griffin iguala o desempenho do Llama-2 apesar de ter sido treinado com mais de 6 vezes menos tokens. Também mostramos que o Griffin consegue extrapolar em sequências significativamente mais longas do que as vistas durante o treinamento. Nossos modelos igualam a eficiência de hardware dos Transformers durante o treinamento e, durante a inferência, possuem menor latência e uma taxa de transferência significativamente maior. Escalamos o Griffin até 14 bilhões de parâmetros e explicamos como particionar nossos modelos para um treinamento distribuído eficiente.

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.