Hybrid Linear Attention

Hybrid Linear Attention: Balancing Efficiency and Expressiveness in Long-Context Models

Standard transformers use full quadratic attention everywhere, which becomes prohibitively expensive for long sequences. Linear attention mechanisms (like RetNet, GLA, DeltaNet) reduce complexity to linear time through clever kernel tricks, but they make different tradeoffs than standard attention—better efficiency sometimes comes at accuracy cost. A natural idea: mix linear and full attention, using linear layers for efficiency and full attention sparingly for hard problems. The problem: the best standalone linear attention models don't automatically produce the best hybrids. This work systematically evaluates 72 models across three generations of linear attention, showing that hybrid success depends critically on the ratio and that full attention is needed more than expected for tasks requiring recall (like language modeling on long contexts).

When you need to extend transformer context beyond 4K tokens efficiently without quadratic memory growth, hybrid architectures offer a practical path. But design choices matter: putting all full attention at layer boundaries behaves differently than distributed placement; too many linear layers lose critical recall capabilities; too many full layers waste efficiency gains. This analysis reveals the engineering decisions that maximize utility.

Core Concept

Hybrid architectures interleave linear and full attention layers at different mixing ratios (24:1, 12:1, 6:1, 3:1 meaning linear layers to each full layer). Three generations of linear attention are evaluated: (1) gated vector recurrence (HGRN, Hawk)—constant-size hidden states enabling perfect efficiency but limited expressiveness; (2) outer-product states (RetNet, GLA)—matrix-valued hidden states for richer representations; (3) delta-rule mechanisms (DeltaNet)—explicit forgetting control. The key insight: recall tasks (reproducing information seen in context) require more full attention than language modeling. At 3:1 ratios, hybrids perform near full-attention baselines while maintaining O(n) complexity. At 24:1 ratios, recall performance degrades catastrophically even when language modeling remains stable.

Architecture Overview

• Linear Attention Layer (Generation 1): Gated vector recurrence with constant hidden state size
• Linear Attention Layer (Generation 2): Outer-product mechanism maintaining matrix-valued states and decay
• Linear Attention Layer (Generation 3): Delta-rule controlled forgetting for selective information retention
• Full Attention Layer: Standard scaled dot-product attention with quadratic complexity
• Interleaving Strategy: Configurable ratio of linear to full attention layers (24:1 through 3:1)
• Position Encoding: Standard rotary embeddings compatible with both layer types

Implementation

This example demonstrates implementing a generation-2 linear attention layer with outer-product states, which balances efficiency and expressiveness.

# Generation 2 linear attention with matrix-valued states
import torch
import torch.nn as nn
import torch.nn.functional as F

class GenerationTwoLinearAttention(nn.Module):
    def __init__(self, hidden_dim, num_heads=8):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads

        # Projection matrices
        self.q_proj = nn.Linear(hidden_dim, hidden_dim)
        self.k_proj = nn.Linear(hidden_dim, hidden_dim)
        self.v_proj = nn.Linear(hidden_dim, hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, hidden_dim)

        # Gate mechanism for selective information flow
        self.gate = nn.Linear(hidden_dim, hidden_dim)

        # Decay parameter (learnable)
        self.decay = nn.Parameter(torch.ones(num_heads, 1) * 0.99)

    def forward(self, x, cache=None):
        """Process sequence with linear attention.
        Maintains matrix-valued hidden state for richer representations."""

        batch, seq_len, dim = x.shape

        # Project to queries, keys, values
        q = self.q_proj(x)  # [batch, seq_len, dim]
        k = self.k_proj(x)
        v = self.v_proj(x)

        # Reshape for multi-head computation
        q = q.reshape(batch, seq_len, self.num_heads, self.head_dim)
        k = k.reshape(batch, seq_len, self.num_heads, self.head_dim)
        v = v.reshape(batch, seq_len, self.num_heads, self.head_dim)

        # Linear attention kernel (e.g., ELU + 1)
        q = F.elu(q) + 1
        k = F.elu(k) + 1

        # Compute output using outer-product state representation
        outputs = []
        state = None  # Matrix-valued state: [batch, heads, head_dim, head_dim]

        for t in range(seq_len):
            q_t = q[:, t]  # [batch, heads, head_dim]
            k_t = k[:, t]  # [batch, heads, head_dim]
            v_t = v[:, t]  # [batch, heads, head_dim]

            if state is None:
                # Initialize state: outer product of k and v
                state = torch.einsum('bhi,bhj->bhij', k_t, v_t)
                # [batch, heads, head_dim, head_dim]
            else:
                # Apply decay to state (forget old information)
                state = state * self.decay.view(1, self.num_heads, 1, 1)
                # Add new information via outer product
                state = state + torch.einsum('bhi,bhj->bhij', k_t, v_t)

            # Compute attention output: Q @ (K·V) / (Q·K)
            numerator = torch.einsum('bhi,bhij->bhj', q_t, state)
            denominator = torch.einsum('bhi,bhi->bh', q_t, k_t).unsqueeze(-1) + 1e-10
            output_t = numerator / denominator
            outputs.append(output_t)

        # Stack outputs
        output = torch.stack(outputs, dim=1)  # [batch, seq_len, heads, head_dim]
        output = output.reshape(batch, seq_len, dim)

        # Gate output
        gate = torch.sigmoid(self.gate(x))
        output = output * gate

        # Project to output
        output = self.out_proj(output)

        return output, state

This example shows a generation-3 linear attention layer with explicit delta-rule forgetting, enabling selective information retention.

class GenerationThreeLinearAttention(nn.Module):
    """Delta-rule mechanism with explicit content erasure control."""

    def __init__(self, hidden_dim, num_heads=8):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads

        self.q_proj = nn.Linear(hidden_dim, hidden_dim)
        self.k_proj = nn.Linear(hidden_dim, hidden_dim)
        self.v_proj = nn.Linear(hidden_dim, hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, hidden_dim)

        # Delta update weight (controls content erasure)
        self.delta = nn.Parameter(torch.ones(num_heads, self.head_dim))

    def forward(self, x):
        """Process with delta-rule controlled forgetting."""

        batch, seq_len, dim = x.shape

        q = self.q_proj(x).reshape(batch, seq_len, self.num_heads, self.head_dim)
        k = self.k_proj(x).reshape(batch, seq_len, self.num_heads, self.head_dim)
        v = self.v_proj(x).reshape(batch, seq_len, self.num_heads, self.head_dim)

        q = F.elu(q) + 1
        k = F.elu(k) + 1

        outputs = []
        state = None  # Matrix state
        erased_mask = None  # Track what content has been erased

        for t in range(seq_len):
            q_t = q[:, t]
            k_t = k[:, t]
            v_t = v[:, t]

            # Delta-rule update: erase content proportionally to delta weight
            if state is not None:
                # Erase dimension-wise based on delta parameter
                erase_signal = self.delta.view(1, self.num_heads, self.head_dim)
                state = state * (1 - erase_signal).unsqueeze(-1)

                # Add new information
                state = state + torch.einsum('bhi,bhj->bhij', k_t, v_t)
            else:
                state = torch.einsum('bhi,bhj->bhij', k_t, v_t)

            # Output computation
            numerator = torch.einsum('bhi,bhij->bhj', q_t, state)
            denominator = torch.einsum('bhi,bhi->bh', q_t, k_t).unsqueeze(-1) + 1e-10
            output_t = numerator / denominator
            outputs.append(output_t)

        output = torch.stack(outputs, dim=1).reshape(batch, seq_len, dim)
        output = self.out_proj(output)

        return output

This example demonstrates the hybrid interleaving strategy and how to configure mixing ratios.

class HybridAttentionTransformer(nn.Module):
    """Transformer mixing linear and full attention at configurable ratio."""

    def __init__(self, hidden_dim, num_layers=12, num_heads=8, linear_ratio=6):
        """Initialize hybrid architecture.
        linear_ratio: 1 full attention layer per N linear layers (6 = 6:1 ratio)"""
        super().__init__()

        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.linear_ratio = linear_ratio

        self.layers = nn.ModuleList()

        for layer_idx in range(num_layers):
            # Every (linear_ratio + 1)th layer is full attention
            if layer_idx % (linear_ratio + 1) == 0:
                attention = FullAttention(hidden_dim, num_heads)
                layer_type = 'full'
            else:
                # Alternate between different generations of linear attention
                if layer_idx % 3 == 0:
                    attention = GenerationTwoLinearAttention(hidden_dim, num_heads)
                    layer_type = 'gen2_linear'
                else:
                    attention = GenerationThreeLinearAttention(hidden_dim, num_heads)
                    layer_type = 'gen3_linear'

            self.layers.append(nn.ModuleDict({
                'attention': attention,
                'norm1': nn.LayerNorm(hidden_dim),
                'norm2': nn.LayerNorm(hidden_dim),
                'ffn': FeedForwardNetwork(hidden_dim),
                'type': layer_type
            }))

    def forward(self, x):
        """Process sequence through hybrid layers."""

        for layer_module in self.layers:
            # Pre-norm attention
            normed = layer_module['norm1'](x)
            attn_out = layer_module['attention'](normed)
            if isinstance(attn_out, tuple):
                attn_out = attn_out[0]  # Discard state for full attention
            x = x + attn_out

            # Feed-forward
            normed = layer_module['norm2'](x)
            ffn_out = layer_module['ffn'](normed)
            x = x + ffn_out

        return x

    @staticmethod
    def compute_complexity(num_layers, seq_len, hidden_dim, linear_ratio):
        """Compute theoretical complexity of hybrid configuration."""
        num_full = num_layers // (linear_ratio + 1)
        num_linear = num_layers - num_full

        full_complexity = num_full * (seq_len ** 2) * hidden_dim
        linear_complexity = num_linear * seq_len * hidden_dim

        total_complexity = full_complexity + linear_complexity
        return {
            'total_flops': total_complexity,
            'full_attention_layers': num_full,
            'linear_attention_layers': num_linear,
            'complexity_reduction': 1 - (linear_complexity / (num_layers * seq_len ** 2 * hidden_dim))
        }

Practical Guidance

When to use: Apply hybrid linear attention when scaling to very long sequences (>8K tokens) where full attention memory is prohibitive. Use when you need both efficiency and reasonable recall on language modeling tasks. Ideal for retrieval-augmented generation, code completion on large contexts, and long-document understanding where computational budget matters.

When NOT to use: Skip if your sequence length < 4K—overhead of linear attention amortization doesn't justify complexity. Avoid if your task is purely extractive (copying from context without reasoning) where linear attention suffices entirely. Don't use if recall is critical and cannot accept >5% degradation. Skip for small models where engineering complexity isn't worth fractional speedups.

Common pitfalls: Using too many linear layers (>12:1 ratio) degrades recall sharply on language modeling tasks. Not tuning decay/delta parameters for your specific problem leaves performance on the table. Placing all full attention at early layers misses late-stage recall requirements. Forgetting that different linear attention generations have different strengths—gen-3 better on recall, gen-2 more stable. Using incompatible position encodings that don't work with linear kernels. Not benchmarking recall separately from language modeling—hybrid success depends on the task.

Reference

Hybrid Linear Attention Team. (2025). A Systematic Analysis of Hybrid Linear Attention. arXiv preprint arXiv:2507.06457. https://arxiv.org/abs/2507.06457