Griffin: Combinazione di ricorrenze lineari controllate con attenzione locale per modelli linguistici efficienti

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

February 29, 2024

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

Abstract

Le reti neurali ricorrenti (RNN) offrono inferenza rapida e scalano efficientemente su sequenze lunghe, ma sono difficili da addestrare e complesse da scalare. Proponiamo Hawk, una RNN con ricorrenze lineari controllate da gate, e Griffin, un modello ibrido che combina ricorrenze lineari controllate da gate con attenzione locale. Hawk supera le prestazioni riportate di Mamba su compiti downstream, mentre Griffin eguaglia le prestazioni di Llama-2 nonostante sia stato addestrato su oltre 6 volte meno token. Dimostriamo inoltre che Griffin è in grado di estrapolare su sequenze significativamente più lunghe di quelle osservate durante l'addestramento. I nostri modelli eguagliano l'efficienza hardware dei Transformer durante l'addestramento, mentre durante l'inferenza presentano una latenza inferiore e una produttività significativamente maggiore. Scaliamo Griffin fino a 14 miliardi di parametri e spieghiamo come partizionare i nostri modelli per un addestramento distribuito efficiente.

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.