Griffin: 効率的な言語モデルのためのゲート付き線形再帰と局所的なアテンションの融合
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)は、推論が高速で長いシーケンスに対して効率的にスケールしますが、学習が難しく、スケーリングも困難です。本論文では、ゲート付き線形再帰を持つRNNであるHawkと、ゲート付き線形再帰と局所的なアテンションを組み合わせたハイブリッドモデルであるGriffinを提案します。Hawkは、下流タスクにおいてMambaの報告性能を上回り、Griffinは、Llama-2の性能に匹敵しながら、6倍以上少ないトークンで学習されています。また、Griffinは、学習中に見たシーケンスよりも大幅に長いシーケンスに対して外挿できることを示します。提案モデルは、学習中にTransformerと同等のハードウェア効率を達成し、推論時には低レイテンシと大幅に高いスループットを実現します。Griffinを140億パラメータまでスケールアップし、効率的な分散学習のためのモデルのシャーディング方法についても説明します。
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.