PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2402.19427v1cs.LGcs.CL

TL;DR

Long-sequence Transformers face quadratic attention and growing KV-cache costs, while RNNs offer fixed-state efficiency but are harder to train and scale. The paper introduces Hawk and Griffin, using gated linear recurrences alone or with local attention. Griffin matches Llama-2 with roughly 7 times fewer tokens, while both models retain comparable training efficiency and improve inference throughput and latency.

  • Problem

    Transformers are difficult to scale efficiently to long sequences because attention is quadratic and the KV cache grows linearly, while recurrent models must establish comparable performance and hardware efficiency.

  • Method

    The paper proposes the RG-LRU gated linear recurrent layer and builds Hawk as a recurrent model and Griffin as a hybrid combining recurrent blocks with local attention.

  • Results

    Griffin-7B and Griffin-14B match the performance of Llama-2 despite being trained on roughly 7 times fewer tokens, while Hawk and Griffin achieve comparable training efficiency to Transformers and higher inference throughput with lower latency.

  • Takeaways & Limitations

    Griffin provides a recurrent-attention alternative with strong language-modeling performance and efficient long-sequence inference.

  • Takeaways & Limitations

    Pre-trained Hawk and Griffin perform less well than Transformers on copying and exact-retrieval tasks without fine-tuning, and the training-efficiency conclusion does not necessarily apply to other accelerators.

Abstract

from arXiv · show

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.

1. Introduction

The paper develops recurrent and hybrid alternatives to Transformers for efficient long-sequence modeling, targeting both scalable performance and hardware efficiency. Griffin combines gated recurrences with local attention and achieves strong loss, downstream, training-efficiency, and inference results.

  • Motivation: Transformers are difficult to scale for long sequences because global attention has quadratic complexity and the KV cache grows linearly with sequence length.Recurrent models instead compress the sequence into a fixed-sized hidden state, but must also demonstrate comparable performance and hardware efficiency.
  • Approach: The paper proposes the RG-LRU gated linear recurrent layer and builds Hawk from recurrent blocks and Griffin from recurrent blocks mixed with local attention.The recurrent block is designed to replace MQA.
  • Scaling and performance: Griffin achieves slightly lower held-out loss than strong Transformer baselines at all model scales.Hawk and Griffin also exhibit power law scaling between held-out loss and training FLOPs up to and beyond 7B parameters.
  • Scaling and performance: Hawk-3B exceeds the reported performance of Mamba-3B on downstream tasks despite using half as many tokens, while Griffin-7B and Griffin-14B match Llama-2 with roughly 7 times fewer tokens.The models are overtrained on 300B tokens across multiple scales.
  • Efficiency: Both Hawk and Griffin achieve comparable training efficiency to Transformers on TPU-v3.A Pallas kernel minimizes RG-LRU memory transfers because diagonal RNN layers are memory bound.
  • Efficiency: During inference, Hawk and Griffin achieve significantly higher throughput than MQA Transformers and lower latency when sampling long sequences.The throughput advantage is especially pronounced as sample length increases.
  • Long-context behavior: Griffin performs better than Transformers on sequences longer than those seen during training, but pretrained Hawk and Griffin perform less well than Transformers on copying and exact-retrieval tasks without fine-tuning.The paper separately investigates learned copying and retrieval capabilities and inference-time behavior of pretrained models.

2. Model Architecture

The models use residual backbones with gated MLPs and one of three temporal-mixing blocks: global MQA, local MQA, or the proposed RG-LRU recurrent block. Griffin combines recurrent blocks with local attention to balance sequence mixing and efficiency.

  • Overall architecture: All models contain a residual block, an MLP block, and a temporal-mixing block, with global MQA, local MQA, or the proposed recurrent block used for temporal mixing.The RG-LRU is the recurrent layer used within the proposed recurrent block.
  • Residual block: The residual backbone stacks N pre-norm-inspired blocks and applies RMSNorm, a final linear layer, and softmax to produce token probabilities.The final linear layer shares weights with the input embedding layer.
  • Residual block: Each residual block applies normalized temporal mixing and normalized MLP processing, merging both outputs with skip connections.The temporal-mixing and MLP components are applied in sequence within the residual structure.
  • Gated MLP block: The gated MLP creates two expanded branches, applies GeLU to one, multiplies them element-wise, and projects the result back to dimension D.The expansion factor is M=3 throughout the work.
  • Temporal-mixing blocks: Local sliding-window attention restricts each position to a fixed number of past tokens, reducing computational FLOPs and bounding the KV cache by the window size.This addresses global attention’s quadratic sequence-length complexity.
  • RG-LRU: RG-LRU gates both the input and recurrent weight using element-wise operations, constrains the diagonal recurrent weight to 0≤a≤1, and avoids dependence of gates on the previous state.These design choices ensure stable recurrence and efficient on-device execution.
  • RG-LRU: Unlike the original LRU, RG-LRU uses real rather than complex recurrences because complex recurrences were not beneficial for language modeling in practice.The paper notes that complex-valued variants may be more expressive, but reports no practical language-modeling benefit here.

3. Recurrent Models Scale as Efficiently as Transformers

The paper evaluates MQA Transformers, Hawk, and Griffin across model scales, finding that Griffin combines recurrent blocks with local attention while maintaining strong scaling and downstream performance. Griffin achieves lower validation loss than the Transformer baseline and matches Llama-2 performance with substantially fewer training tokens.

  • Scaling studies: The study compares an MQA-Transformer baseline, pure recurrent Hawk, and hybrid Griffin across scales from 100M to 7B parameters, plus Griffin-14B.Training tokens are increased roughly in proportion to parameter count, following Chinchilla scaling laws.
  • Model designs: Griffin mixes recurrent blocks with local MQA attention, using a fixed 1024-token local-attention window.Its layered structure alternates recurrent and local-attention residual blocks.
  • Scaling curves: All three model families show a linear relationship between validation loss and training FLOPs on log-scaled axes.This scaling behavior extends the pattern previously observed for Transformers.
  • Scaling curves: Griffin achieves lower validation loss than the Transformer baseline across all FLOPs budgets, while Hawk’s higher loss gap appears to close at larger budgets.The comparison is made despite Griffin using no global attention layers.
  • Downstream evaluation: Griffin-7B and Griffin-14B match Llama-2 performance despite training on roughly 7 times fewer tokens.Hawk-3B exceeds reported Mamba-3B performance despite using half as many tokens, while Griffin also outperforms the in-house MQA baseline.

4. Training Recurrent Models Efficiently on Device

The paper addresses distributed sharding and TPU execution challenges for recurrent models. A custom Pallas linear-scan kernel reduces memory transfers, while Griffin’s training advantage grows on sufficiently long sequences.

  • Distributed training: Large models are sharded across devices with model parallelism because they do not fit on one device, making communication-efficient partitioning critical for fast training.Megatron-style sharding, ZeRO parallelism, and bfloat16 reduce communication and memory pressure.
  • Distributed training: The recurrent block matches the MLP block’s communication requirements through sharded linear layers, channel-wise Conv1D partitioning, and block-diagonal RG-LRU gates.The diagonal recurrence permits parameter sharding and computation without cross-device communication.
  • Efficient linear recurrences on device: 0.75 FLOPs-to-byte ratio makes the RG-LRU update memory bound on TPU-v3, below the device’s 4.2 capacity for elementwise operations.Each element requires 6 FLOPs, 6 bytes loaded, and 2 bytes written under bfloat16 assumptions.
  • Efficient linear recurrences on device: A custom Pallas linear-scan kernel delivers almost 3x speedup over native JAX and reduces full-Hawk training time per step by 10-20%.The kernel keeps the hidden state in VMEM and transfers data in larger chunks.
  • Efficient linear recurrences on device: The RG-LRU gating mechanism prevents using the convolutional view of linear recurrences, motivating the custom linear-scan implementation.Associative scan remains possible in principle, but the gating changes the computational structure.
  • Training speed on longer sequences: At longer sequences, Griffin’s runtime remains stable while the Transformer baseline becomes slower under a fixed token count per batch.The gain is largest when sequence length is sufficiently large relative to model width, because global attention scales as O(T^2D) versus O(TD) for RG-LRU.

5. Inference Speed

During autoregressive decoding, recurrent models avoid the sequence-growing KV-cache costs of Transformers, yielding lower latency and higher throughput, especially for long sequences.

  • Decode-stage behavior: Prefill is compute-bound and parallelizable, whereas decode samples tokens autoregressively and becomes especially favorable to recurrent models at longer sequence lengths.The decoding advantage is linked to the large attention KV cache at long contexts.
  • Metrics: Latency measures generation time for a specified token count and batch size, while throughput measures the maximum tokens per second on one device.Throughput can increase through lower latency or larger feasible batch sizes.
  • Decode-stage behavior: At batch sizes below roughly 128, decoding is memory-bound, with parameter and cache loading dominating token-generation time.The analysis assumes B≲128 for the remainder of the section.
  • Cache sizes: Recurrent models use a fixed-sized recurrent state, whereas Transformer KV caches grow linearly with sequence length.Cache size is compared at batch size 1 for each architecture.
  • Results: Hawk and Griffin achieve faster sampling latency than MQA Transformers for long sequences, especially with longer sequence lengths and 4096-token prefills.The comparison uses batch size 16 with empty and 4096-token prefills.
  • Results: Both Griffin and Hawk achieve significantly higher throughput than the MQA Transformer baseline because their smaller caches permit larger batch sizes on one device.Hawk exceeds Griffin at large batch sizes because Griffin’s local-attention cache eventually becomes comparable to parameter size.

6. Long Context Modeling

Hawk and Griffin improve next-token prediction with longer contexts and extrapolate beyond training lengths, while Griffin combines local attention with recurrent modeling for copying and retrieval.

  • Improving next token prediction with longer contexts: The long-context evaluation measures held-out books loss across sequence lengths, testing prediction on contexts longer than those used during training.Extrapolation is defined as accurately predicting the next token from longer-than-training contexts.
  • Improving next token prediction with longer contexts: Hawk and Griffin extrapolate to sequences at least 4x longer than their training sequences and improve next-token prediction up to a maximal length.Griffin extrapolates well even when its local-attention layers use RoPE.
  • Improving next token prediction with longer contexts: Training 1B models on 8192-token sequences lowers evaluation loss at sequence lengths of 8192 or larger relative to 2048-token training.The models trained on 2048 tokens perform slightly better at short sequence lengths.
  • Copy and retrieval capabilities: Selective Copying and Induction Heads test whether models learn to copy or retrieve context information and extrapolate on these synthetic tasks.The experiments use 5-block networks with roughly 250K parameters, with Griffin placing one local-attention layer in the middle block.
  • Copy and retrieval capabilities: On Selective Copying, all models solve the task perfectly, while Griffin matches Transformer learning speed and Hawk learns significantly more slowly.Griffin uses only a single local-attention layer in this comparison.
  • Copy and retrieval capabilities: On Induction Heads, Hawk perfectly extrapolates to evaluation sequences several orders of magnitude longer than training sequences, unlike the Transformer baseline.All three models solve the task perfectly through the training sequence length.
  • Copy and retrieval capabilities: In pre-trained phonebook lookup, Griffin solves contexts up to its 1024-token local-attention window, then degrades when that window no longer covers the phonebook.Griffin extrapolates better to longer sequences than the Transformer baseline, while Hawk fails as phonebooks grow.

7. Related Works

Related work contrasts Transformers’ parallel training and global attention with recurrent and state-space alternatives that reduce long-sequence computation and inference costs.

  • Transformers and recurrent alternatives: Transformers interleave MLPs with multi-head attention and became dominant as RNNs became difficult to scale for modern deep learning and NLP.Classical RNNs process sequences sequentially, slowing forward and backward training.
  • State-space models: State-space models use linear recurrences computed through parallel scans or convolutions, enabling training speeds comparable to Transformer models.SSMs combine state-space modeling with recurrent computation for long inputs.
  • State-space models: The LRU line of work found that simplified linear RNNs with improved parameterization and initialization perform comparably to other SSM variants on long-range tasks.RetNet instead uses a gated SSM design that parallelizes through a multi-head-attention variant.
  • State-space models: Mamba introduced input-dependent selection and reported Transformer-comparable performance with efficient inference, while later work proposed extensions for different applications.Gateloop also uses input-dependent gating.
  • Efficient attention: Linear attention reduces self-attention cost by linearizing attention into a recurrent linear RNN, but this efficiency can trade off against model performance.Flash Attention improves GPU attention training speed through efficient memory-hierarchy use.

8. Conclusion

Hawk and Griffin are presented as efficient alternatives to Transformers, combining strong language-modeling performance with favorable inference characteristics and long-sequence capabilities.

  • Hawk incorporates the novel gated linear recurrent RG-LRU layer, while Griffin mixes RG-LRU recurrence with local attention.
  • Hawk exceeds reported Mamba downstream performance using half as many training tokens, while Griffin slightly exceeds Llama-2 using over 6 times fewer tokens.
  • Hawk and Griffin show power-law scaling of held-out loss as compute resources increase.
  • The models reduce inference latency and substantially increase throughput compared with Transformer baselines.
  • Hawk and Griffin extrapolate beyond training sequence lengths and efficiently learn to copy and retrieve information over long horizons.

A. RG-LRU Recurrence Gate

The RG-LRU recurrence gate differs from standard gates by favoring retention of the previous hidden state and using a distinct nonlinear interpolation of recurrent updates and inputs.

  • Implementation: The gate’s implementation computes log a_t and then exponentiates it for numerical stability, rather than applying a sigmoid before taking a power.
  • Gate behaviour: RG-LRU gating is biased toward retaining information and cannot fully discard the previous hidden state, depending on Λ.
  • Gate behaviour: The relative contribution of the new input x_t and previous hidden state h_t−1 is used to analyze different recurrence gates.
  • Gate behaviour: Mamba gating is almost identical to GRU gating when A is near 1, with minor deviations at smaller values.
  • Gate behaviour: RG-LRU gating performs a substantially different nonlinear interpolation between discarding x_t and applying the LRU update.

B. Complex-Gated Linear Recurrent Unit (CG-LRU)

The CG-LRU extends the recurrent layer to complex-valued diagonal recurrences by splitting real inputs into complex components and reconstructing a real output.

  • The complex recurrence is parameterized as ã=σ(Λ)e^iθ, with θ learned separately.
  • The input x_t is split across channels, with its first half interpreted as the real part and its second half as the imaginary part.
  • The LRU equations are rewritten for the complex-valued recurrence, with complex variables marked by a tilde.
  • The dimensions of r_t, i_t, ã_t, and h̃_t are half those of the real input x_t.
  • The output y_t is formed by stacking the real and imaginary parts of h_t into a single vector.
  • The experiments use shared hyperparameters across the model families and report TPU-v3 hardware specifications and scan implementations for efficiency analysis.

D.2. Scan runtimes

The efficiency analysis covers scan runtimes, local-attention windows, and decode-time memory behavior, showing optimized recurrent computation and bounded recurrent state sizes.

  • Scan runtimes: A Pallas scan kernel on TPU-v3 achieves nearly 3x speedup over the naive Jax implementation.
  • Scan runtimes: The associative scan is significantly slower than the tested alternatives and can make runtime up to 50% slower.
  • Scan runtimes: The Pallas scan gains improve full Hawk training time per step, including at the 7B scale.
  • Local attention windows: With a fixed local-attention window of 1024, Griffin outperforms the global-attention MQA Transformer across all tested sequence lengths.
  • Local attention windows: The performance gap between Griffin with a 1024-token window and global-attention MQA decreases as sequence length grows.
  • Decode-time memory: Transformer KV-cache size and attention computation grow linearly with sequence length, while recurrent state size does not grow with sequence length.
  • Decode-time memory: Griffin’s local-attention cache remains small relative to parameter count when its window satisfies T_WS≲D^2/(B d_head), supporting decode speed similar to Hawk.

H. Additional Details of the Copy and Retrieval Tasks

The section illustrates two retrieval tasks: Selective Copying requires copying designated data while ignoring noise, whereas Induction Heads requires recalling the token following a special token.

  • Selective Copying: Selective Copying requires copying data tokens while ignoring noise tokens.Colored tokens are copied, whereas white tokens are treated as noise.
  • Crossed-out output tokens indicate positions masked out in the loss for both tasks.
  • Figure 11 illustrates the Selective Copying and Induction Heads tasks.
  • Induction Heads: Induction Heads requires recalling the token immediately following a special token.The special token is represented as a black token in Figure 11.
Loading 2402.19427v1…