Justify why hybrid SSM models like Jamba still keep a slice of attention layers

MCQ·Medium·4.0 · 0·~1 min·Asked atJasper Linkedin Mphasis

Attempt it

A pure **Mamba** model is fast and memory-light because each layer carries one fixed-size hidden state that summarises everything seen so far. The state updates token by token like an RNN. No KV cache, no `O(n^2)` attention matrix, no memory growth with context length. The price is that the summary is lossy. If you ask the model to copy a 12-digit number that appeared 8000 tokens ago, or to find the exact phrase 'order ID ABC-789' in the conversation history, the fixed-size state has no way to surface that one specific token verbatim. The information was crushed into the summary along with everything else. **Attention** does the opposite. It keeps every past token's K and V available, so a new query can pick out any past token by content. Expensive but exact. **Jamba and similar hybrids** interleave the two. Most layers are Mamba (cheap state updates), but roughly every seventh layer is full attention. That single attention layer gives the network a content-addressable lookup window over past tokens. The hybrid keeps most of Mamba's memory and throughput wins while restoring the exact-retrieval capability that pure SSMs lose.

**The correct answer is option 1.** State-space layers compress history into a fixed-size state, which makes them cheap but lossy; interleaved attention layers restore exact content-addressable access to past tokens where the task demands it. Mamba's selective SSM update is roughly: $$h_t = A_t h_{t-1} + B_t x_t, \quad y_t = C_t h_t$$ The state `h_t` has a fixed dimension regardless of how long the input is. Every past token contributes additively into `h_t`, then gets discarded. There is no way at step 10000 to recover the exact embedding of the token at step 30 unless that token's contribution still dominates the recurrent dynamics. Attention does the opposite: K and V are stored explicitly per past token, so a new Q can attend to any one of them by content. The cache grows with context, but exact retrieval is preserved. **Jamba (AI21 Labs)** interleaves: roughly one attention layer per seven Mamba layers, plus MoE on some FFN sublayers. **Zamba** and **Samba** use similar ratios. The result is a KV cache an order of magnitude smaller than a pure transformer at the same context length, with near-parity on retrieval, copy-from-context, and multi-hop QA benchmarks that pure SSMs visibly underperform on. Options 2, 3, and 4 are wrong: Mamba is fully differentiable, the choice is about capability not backward compatibility, and SSM layers handle the full context not just a slice of it.

TL;DR

Mamba layers compress history into a fixed-size state and lose exact recall; interleaved attention layers restore content-addressable access to past tokens where it matters.

Memory aid

Easy to grasp

Mamba is like a very disciplined note-taker who keeps one small summary card and updates it as new information comes in. Cheap to carry, but the moment you ask for an exact quote from page 3, the summary card cannot deliver. Attention is like keeping all your notes verbatim and flipping through them when needed. It is powerful but the stack of notes grows huge. Jamba and similar hybrids do the obvious thing: most pages use the summary card because it is cheap, but every few pages they keep the full notes too, just in case you need to look up a name or copy a paragraph word for word. The combination gives you mostly cheap summary updates with occasional precise lookups, which is what real tasks actually need.

Key concepts

Transformer Architecture

Concept explanation~2 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

State-space models like Mamba promise a transformer-killer property: linear time and constant memory per token. So why does every production-grade SSM-flavoured LLM in 2026 keep a slice of attention layers? The answer is information theory wearing engineering clothes. A fixed-size recurrent state cannot, by construction, preserve every detail of an arbitrarily long input. Attention can, at the cost of growing memory. Hybrids buy the best of both.

This walkthrough explains the lossy-compression argument, traces it through the associative-recall benchmarks where SSMs visibly struggle, lays out the layer ratios in Jamba, Samba, and Hymba, and ends with where pure SSM remains the right choice.

Why a fixed-size recurrent state is a lossy summary

Mamba's selective SSM update is a recurrence with a state vector of fixed dimension d_state (often 16 to 64). At step t, the state absorbs the new input and discards the previous state's exact contents:

h_t = A_t h_{t-1} + B_t x_t, \quad y_t = C_t h_t

No matter how long the sequence is, h_t has the same size. By the data-processing inequality, no fixed-size summary can losslessly encode an arbitrary-length stream. Information must be discarded somewhere.

Selective SSMs (Mamba's contribution over earlier linear-time models) make A_t, B_t, C_t input-dependent, which lets the state be selective about what to remember. This helps a lot for tasks like language modelling perplexity, which is mostly local. But it does not, and cannot, reach attention's exact-retrieval ability on tasks that depend on a specific past token's content.

The canonical failure mode is associative recall: given a key-value list and a query for the value of a specific key, the model has to retrieve the exact value. A transformer's induction heads do this trivially through attention. An SSM must encode the entire key-value mapping into h_t with enough fidelity to reproduce the right value for an arbitrary query. The state capacity caps how many key-value pairs can be stored cleanly, and benchmarks like Zoology and MQAR show pure SSMs degrade sharply as the number of pairs grows.

What attention does that SSMs cannot

The hybrid recipe in Jamba, Samba, and Hymba

When pure SSM is still the right call

Putting numbers to it: 2026 frontier-model context

Situations where this technique stops working.

2–4 min · Everything important, quickly.

Real products, models, and research that use this idea.

AI21 Jamba-1.5 Large (52B active params, MoE) interleaves Mamba and attention in roughly a 7-to-1 ratio for a 256k context model.
Zamba 2 (Zyphra) uses a hybrid SSM-attention design at the 1B to 7B scale aimed at edge inference.

What an interviewer would ask next. Try answering before peeking at the approach.

QWhat is an associative-recall benchmark and why does it discriminate SSMs from transformers so sharply?

Associative recall asks the model to retrieve a value given a key seen earlier in context (e.g. 'Q: what is the email of Alice? A:' after listing many name-email pairs). Transformers solve it via induction heads in attention; SSMs must encode the key-value mapping into the recurrent state, which has limited capacity. The Zoology paper made this rigorous.

2 more follow-ups an interviewer would ask next. Sign in to reveal them.

Red flags & common mistakes

The phrases that signal junior thinking. Click to expand.

Most common mistake

Believing Mamba layers can match transformer recall by scaling state size, or thinking attention is kept only for legacy reasons. The recurrent state compresses lossily by construction; attention is the only way to do exact content-addressable lookup at decode time.

60 second bullets to scan on the way to the call.

Why a fixed-size recurrent state is a lossy compression of arbitrary-length input
What associative-recall tasks look like and why pure SSMs underperform on them

Primary sources. Browse if you want the original framing.