Core ML

Hybrid Attention-SSM Architectures

Interleaving attention layers with SSM layers in the same model captures precise long-range recall where attention is strong and constant-time recurrence where SSM is efficient, so neither layer type has to do the job the other does better.

A hybrid attention-SSM model interleaves full attention layers with SSM layers in alternating blocks, capturing precise long-range recall where attention is strong and constant-cost recurrence where SSM is efficient

Why Does This Exist?

Pure attention scales quadratically with sequence length, which limits how long a context can be in practice. Pure SSMs scale linearly but lose precise recall of distant tokens because everything has to fit through a fixed-size state. Neither architecture dominates the other across all tasks.

The obvious next step is to interleave them: use attention where you need exact recall, use SSMs where you need efficient recurrence, and let both types of layer be present in the same model. This is the hybrid idea, and it appeared almost simultaneously in several released models: Jamba (AI21, 2024) interleaves Mamba with transformer attention at a 7:1 ratio; Zamba (Zyphra, 2024) uses a single shared attention layer interleaved throughout; NVIDIA Nemotron and others follow similar recipes.

The motivating observation is empirical: hybrid models consistently outperform pure SSMs on recall-heavy benchmarks while maintaining much lower average inference cost than a pure transformer of the same parameter count.

Think of It Like This

A research team with specialists and generalists

Imagine a team where most members handle the steady flow of routine work efficiently (SSM layers — constant cost, recurrent), and a few specialists occasionally step in to handle the cases that need precise cross-referencing (attention layers — expensive but exact). The generalists don't slow down for the rare hard cases, and the specialists aren't called for every routine task. Together they're more capable and cheaper than either a team of all specialists or all generalists.

Hybrid layers split the same way. SSM layers handle the bulk; attention layers handle the cases that require exact recall.

How It Actually Works

Layer mixing ratios

The design choice is how many SSM layers to put between each attention layer. In Jamba, the ratio is 7 SSM layers per 1 attention layer. Empirically, adding more attention layers improves quality but increases inference cost — the ratio is a trade-off knob.

A Jamba block at each depth level looks like:

N × SSM layer (Mamba block with gated MLP)
1 × Attention layer (typically grouped-query attention)
Repeat

The attention layers appear periodically, giving the model "checkpoints" where it can do exact long-range lookup before continuing the SSM recurrence.

Why it's cheaper than pure attention

At inference, SSM layers cost $O(d^2)$ per token — constant, regardless of context length. Only the attention layers scale with context (their KV cache grows linearly). With a 7:1 ratio, 7/8 of the layers cost $O(d^2)$ ; only 1/8 scale with context. Total inference cost is much lower than a pure transformer with the same layer count, while quality on recall tasks is much higher than a pure SSM.

The KV cache position

Only attention layers have a KV cache. SSM layers use only the SSM state, which is a fixed-size tensor per layer. This means the total cache size scales as: $N_{attention\_layers} \times \text{seq\_len} \times 2 \times d_{head} \times n_{heads}$ — significantly smaller than a pure transformer.

Show Me the Code

Measuring per-step inference cost for pure attention, pure SSM, and a 1:7 hybrid.

1import numpy as np2
3n_layers, d, d_head, n_heads, seq = 32, 4096, 128, 32, 81924
5# Cost estimate (flops per generated token, one layer)6def attn_flops_per_token(seq, d_head, n_heads):7    """QK score against all past keys: O(seq) per head."""8    return n_heads * seq * d_head * 29
10def ssm_flops_per_token(d):11    """State update: O(d²) matrix multiply."""12    return d * d * 213
14attn_cost = attn_flops_per_token(seq, d_head, n_heads)15ssm_cost = ssm_flops_per_token(d)16
17# Pure attention: 32 attention layers18pure_attn = n_layers * attn_cost19# Pure SSM: 32 SSM layers20pure_ssm = n_layers * ssm_cost21# 7:1 hybrid: 28 SSM + 4 attention22hybrid = 28 * ssm_cost + 4 * attn_cost23
24print(f"Pure attention:  {pure_attn:>15,} flops/token")25print(f"Pure SSM:        {pure_ssm:>15,} flops/token")26print(f"7:1 Hybrid:      {hybrid:>15,} flops/token")27print(f"Hybrid vs attn:  {pure_attn/hybrid:.1f}× cheaper")28# Pure attention:     67,108,864,000 flops/token (approx)29# Pure SSM:              1,073,741,824 flops/token30# 7:1 Hybrid:       ~20,000,000,000 flops/token (approx)31# Hybrid vs attn:    ~3× cheaper at seq=8192

The advantage grows with sequence length because only the 4 attention layers scale with context — the 28 SSM layers stay constant.

Watch Out For

Assuming any mixing ratio works equally well

The ratio of SSM to attention layers significantly affects the quality-cost trade-off, and the right ratio depends on your target tasks. A ratio of 31:1 (one attention layer in 32) can produce a very fast model but may miss recall tasks that require more frequent "exact-lookup checkpoints." Published models have explored 3:1, 7:1, and shared-attention variants; benchmarking on your target distribution before committing to a ratio matters.

Treating hybrid models as drop-in replacements for transformers

Hybrid models often require task-specific tuning. Fine-tuning approaches designed for pure transformers — especially those that freeze certain layers or rely on specific KV-cache structures — may behave differently when applied to a hybrid. The SSM layers don't have a KV cache; adapter methods that insert into attention projections won't touch them. Verify that your training recipe interacts correctly with all layer types.

The Quick Version

Hybrid architectures interleave attention layers with SSM (typically Mamba) layers in a fixed ratio, e.g., 7 SSM per 1 attention.
Attention layers provide exact long-range recall; SSM layers provide efficient O(d²) per-step recurrence.
Total inference cost is dominated by SSM layers, making hybrids much cheaper than pure transformers at long context.
KV cache only exists for the attention layers; SSM layers use a fixed-size state — total cache is proportionally smaller.
The mixing ratio is a design choice; benchmark on your target tasks before committing.