Griffin: 効率的な言語モデルのためのゲート付き線形再帰と局所的なアテンションの融合

要旨

リカレントニューラルネットワーク（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.

Griffin: 効率的な言語モデルのためのゲート付き線形再帰と局所的なアテンションの融合

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

要旨

Support