Griffin (RG-LRU) Hybrid Language Model

• Griffin (RG-LRU) is a hybrid language model that combines Real-Gated Linear Recurrent Units with sliding-window local multi-query attention for efficient long-sequence processing.
• It alternates RG-LRU and local attention blocks within a residual framework to achieve competitive performance with fewer training tokens and improved throughput.
• Optimized for hardware efficiency, Griffin employs techniques like Megatron-style sharding and bfloat16 operations to deliver lower latency and up to 2× training throughput.

Griffin is a hybrid neural language modeling architecture integrating Real-Gated Linear Recurrent Units (RG-LRU) with sliding-window local multi-query attention. Designed for high hardware efficiency and superior long-sequence processing, Griffin achieves competitive or superior results compared to state-of-the-art Transformer models such as Llama-2, while requiring significantly fewer training tokens and delivering enhanced inference throughput and lower latency. The core innovation lies in the alternation of RG-LRU-based temporal-mixing blocks with local attention within a residual framework, enabling both effective sequence modeling and efficient scaling for large model deployments (De et al., 2024).

1. Architectural Layout and Temporal Mixing

Griffin consists of a stack of $N$ residual blocks configured in a pre-normalization sequence. Each residual block is composed of two sub-layers:

• A temporal-mixing sub-layer for aggregating time-dependent information,
• A feed-forward sub-layer implemented with a gated GeLU expansion and projection.

Both sub-layers utilize a structure of RMSNorm applied to the input, followed by the operator and a skip connection. Griffin alternates between two mechanisms for temporal mixing:

• Recurrent blocks based on RG-LRU,
• Local multi-query attention (MQA) blocks with a sliding window.

The canonical cycle is as follows:

Block Position	Temporal-Mixing Component	Sub-Layer Sequence
(a)	RG-LRU	RG-LRU → MLP
(b)	RG-LRU	RG-LRU → MLP
(c)	Local MQA (window size ≈ 1024)	MQA → MLP

This triadic pattern is repeated $\lfloor N/3 \rfloor$ times. Every temporal-mixing operator (RG-LRU or MQA) is immediately followed by RMSNorm, a gated MLP, and a residual connection.

2. Mathematical Formulation of the RG-LRU

The RG-LRU, or Real-Gated Linear Recurrent Unit, is a recurrent architecture supporting efficient sequence modeling with stable, trainable gating and parallelism over sequence position. For each timestep $t$, given input $x_t \in \mathbb{R}^d$ and hidden state $h_t \in \mathbb{R}^d$:

1. Recurrence gate: $r_t = \sigma(W_a x_t + b_a)$
2. Input gate: $i_t = \sigma(W_x x_t + b_x)$
3. Recurrence weight: $$a = \sigma(\Lambda),\quad \Lambda \in \mathbb{R}^d,\quad 0 \leq a \leq 1$$ (learned)
4. Powered gate: $a_t = a^{c r_t},\quad c=8$
5. Update rule: $$h_t = a_t \odot h_{t-1} + \sqrt{1 - a_t^2} \odot (i_t \odot x_t)$$
6. Output: $\lfloor N/3 \rfloor$0

Here, the gates $\lfloor N/3 \rfloor$1 and $\lfloor N/3 \rfloor$2 depend exclusively on $\lfloor N/3 \rfloor$3, enabling parallel time computation. Recurrence weights $\lfloor N/3 \rfloor$4 are sigmoid-constrained (a=σ(Λ)) for stability. In practice, $\lfloor N/3 \rfloor$5 is computed in log-space for numerical stability using $\lfloor N/3 \rfloor$6. Projection matrices $\lfloor N/3 \rfloor$7 are often block-diagonal, supporting efficient sharding. This design allows interpolation between pure LRU-like recurrences and state preservation.

3. Residual Block Integration: RG-LRU and Local MQA

Within each residual block:

• Input $\lfloor N/3 \rfloor$8 undergoes RMSNorm.
• If RG-LRU:
  • Linear₁ projects $\lfloor N/3 \rfloor$9 to $t$0,
  • Channel-wise separable Conv1D (width=4) produces $t$1,
  • RG-LRU scan computes $t$2,
  • Linear₂ with GeLU and element-wise multiplication: $t$3, $t$4,
  • Linear₃ projects $t$5 to $t$6.
• If local attention:
  • Multi-Query Attention over window $t$7, with RoPE positional encoding.

The output is the sum of the input and $t$8 (residual branch), followed by RMSNorm, a gated-MLP, and an additional residual connection. This structure is used in all RG-LRU and MQA blocks, preserving a uniform interface for distributed and efficient layer application.

4. Scaling, Sharding, and Hardware Efficiency

Griffin is engineered for scalable model and data parallelism:

• Megatron-style tensor slicing with a single all-reduce per layer for MLPs and attention,
• Attention heads sharded across devices,
• RG-LRU weights ($t$9, $x_t \in \mathbb{R}^d$0) block-diagonalized (e.g., 16 blocks per device), with separable Conv1D per channel, eliminating the need for inter-device communication in RG-LRU layers,
• Optimizer states and model parameters are sharded with ZeRO,
• Computations in bfloat16 for both parameters and activations.

On TPU-v3, a custom "Pallas" fused linear-scan kernel for RG-LRU maintains hidden states in fast local memory (VMEM) and batches memory accesses, resulting in a 3× speedup over JAX's scan and a 10–20% overall speedup for the Griffin and Hawk architectures. Associative prefix-sum scans and convolutions were empirically slower and did not decrease memory traffic.

For long sequences ($x_t \in \mathbb{R}^d$1k), Griffin's temporal-mixing complexity is $x_t \in \mathbb{R}^d$2, compared to the Transformer's $x_t \in \mathbb{R}^d$3. For $x_t \in \mathbb{R}^d$4k at 7B scale, Griffin's training throughput per step is up to 2× higher than MQA baselines. In decoder (generation) mode, Griffin requires only streaming of parameters and a small recurrent state ($x_t \in \mathbb{R}^d$5), while Transformers must stream parameters and the full $x_t \in \mathbb{R}^d$6 key-value cache. For $x_t \in \mathbb{R}^d$7k, $x_t \in \mathbb{R}^d$8, Griffin/Hawk achieves $x_t \in \mathbb{R}^d$9 lower latency and 2–3× higher throughput (tokens/sec per TPU-core).

5. Empirical Results and Sequence Extrapolation

Griffin’s empirical benchmarks demonstrate its efficacy across scale and context length:

• Scaling curves: On MassiveText, from 100M to 14B parameters, Griffin outperforms MQA Transformer baselines by approximately 0.1 bits/token at every scale; through higher budgets, its performance approaches that of the pure Hawk (RG-LRU) model.
• Downstream evaluation (on models trained 300B tokens):
  • Mamba-3B (600B tokens): 58.5
  • Llama-2-7B/13B (2T tokens): 65.3 / 69.3
  • Griffin-3B: 60.7
  • Griffin-7B: 65.8 (matches Llama-2-7B with 7× fewer tokens)
  • Griffin-14B: 69.5 (matches Llama-2-13B with ∼7× fewer tokens)
  • Hawk-3B surpasses Mamba-3B despite half the training tokens.
• Long-context extrapolation: When evaluated on held-out Books and ArXiv after training with $h_t \in \mathbb{R}^d$0, Transformer loss plateaus or degrades past 2k tokens, whereas Hawk and Griffin retain improvements up to 32k tokens. Training with $h_t \in \mathbb{R}^d$1 further boosts Griffin's performance for $h_t \in \mathbb{R}^d$2k (at the expense of very short contexts).
• Copy and retrieval tasks: On synthetic "Selective Copy", Griffin matches Transformer learning speed, while Hawk is slower. In "Induction Heads" generalization (train $h_t \in \mathbb{R}^d$3, test $h_t \in \mathbb{R}^d$4), Griffin and Hawk extrapolate perfectly. For phonebook lookup, Griffin retrieves up to the local window (1024), gracefully degrading after, while Transformer and Hawk fail beyond their respective context limits (De et al., 2024).

6. Relationship to Related Architectures and Importance

Griffin’s hybrid approach interpolates between the sequential processing strengths of recurrent models and the global contextualization of attention mechanisms. Unlike pure Transformers, which are quadratic in sequence length for temporal-mixing, Griffin's $h_t \in \mathbb{R}^d$5 complexity enables substantial benefits for long-context tasks and high-throughput inference. Unlike models relying exclusively on attention or recurrence, Griffin’s interleaving of RG-LRU and local attention is central to its capacity for extrapolation, token efficiency, and hardware-aligned deployment at scale.

A plausible implication is that Griffin offers a template for future efficient LLM designs targeting longer context windows, hardware efficiency, and scalable distributed training, while closing the performance gap with well-known large-scale Transformer baselines (De et al., 2024).