Exploring Linear Attention Variants

Beyond $O(N^2)$: Working Through Linear Attention and State Space Models

I’ve been spending a lot of time lately thinking about the attention scaling problem, and I wanted to write up where my understanding has landed. This is partly for my own reference, partly because I think the landscape has shifted enough in 2025 that it’s worth revisiting assumptions.

The Problem We’re Actually Solving

Standard self-attention computes pairwise relationships between every token:

$\text{Attention}(Q, K, V) = \sigma\left(\frac{QK^T}{\sqrt{d}}\right)V$

That $QK^T$ creates an $N \times N$ matrix. At 1K tokens, that’s a million entries. At 100K tokens, it’s 10 billion. The quadratic scaling isn’t just expensive, it fundamentally limits what you can do with long sequences without throwing hardware at the problem.

What I find interesting is that we actually had linear-scaling sequence models before Transformers. RNNs process tokens sequentially with fixed memory. So why did we move away from them?

Three things killed RNNs:

Vanishing gradients. Gradients shrink exponentially as they backpropagate through time. LSTMs helped, but realistically topped out around 100-200 tokens of effective memory.

Sequential training. You can’t parallelize across the time dimension. On modern GPUs optimized for parallel ops, this is a massive bottleneck.

The compression problem. Squeezing an entire document into a fixed-size hidden state loses information. There’s no way around it.

Transformers fixed all three. No vanishing gradients (every token directly attends to every other). Fully parallelizable. No compression (just keep everything in the KV cache). The tradeoff was quadratic scaling, which seemed acceptable when contexts were 512 tokens.

Now we want 100K+ token contexts, and the tradeoff doesn’t look so good anymore.

First Attempts: Linear Attention

The first obvious fix was algebraic. Matrix multiplication is associative:

Standard order: $(QK^T)V$ — build the $N \times N$ matrix first Linear order: $Q(K^TV)$ — compute the $d \times d$ term first

This drops complexity to $O(N)$. Models like Performer and Linformer used kernel tricks to approximate the softmax and make this work.

The results were… okay. They worked on some tasks but consistently underperformed on language modeling. The issue seems to be that softmax provides a “sharp” attention distribution. When you approximate it, you get something fuzzier. Tasks requiring precise retrieval (“what was the third word?”) suffer.

I spent some time with Performer variants and my takeaway was: useful for specific applications, but not a general replacement. The approximation quality matters a lot and degrades in ways that are hard to predict.

State Space Models: A Different Framing

The more interesting direction came from thinking about sequences as continuous signals rather than discrete token lists. This is the state space model (SSM) approach.

The core equations:

$h'(t) = Ah(t) + Bx(t)$          $y(t) = Ch(t) + Dx(t)$

This is essentially an RNN, but the S4 work (Gu et al., 2021) introduced two ideas that made it actually work.

The HiPPO Matrix

In standard RNNs, the state transition matrix $A$ is initialized randomly. S4 uses a specifically structured matrix called HiPPO (High-order Polynomial Projection Operators).

The intuition: HiPPO is designed to optimally project the input history onto a polynomial basis. Instead of the hidden state being an arbitrary compression, it’s structured to preserve historical information in a mathematically principled way. The coefficients of these polynomials can reconstruct the input signal.

This is what lets S4 handle long sequences without gradient issues. The structure of $A$ prevents the exponential decay that kills traditional RNNs.

The Convolution Connection

The other insight: when $A$, $B$, $C$ are time-invariant (LTI), the recurrence can be expressed as a convolution. Convolutions parallelize beautifully on GPUs.

So S4 gets the best of both:

• Training: Convolutional mode, parallelizes like a Transformer

• Inference: Recurrent mode, constant memory like an RNN

S4 crushed benchmarks like Long Range Arena. But language modeling was still a weak point. Something was missing.

Mamba: Making SSMs Content-Aware

The Mamba paper (Gu & Dao, December 2023) identified the problem: S4 is linear time-invariant. It processes every token identically. The matrices don’t change based on what you’re looking at.

For language, this is wrong. Some tokens matter (“the defendant is guilty”), some are filler (“um, like, basically”). An LTI system can’t distinguish them.

Selectivity

Mamba makes the SSM parameters input-dependent. Specifically, it generates $\Delta$, $B$, and $C$ as functions of the current input:

• $\Delta \approx 0$: ignore this token, preserve current state

• $\Delta \approx 1$: this token matters, update the state significantly

This is the key innovation. The model can learn to selectively propagate or forget information based on content.

The Parallel Scan

Making parameters input-dependent breaks the convolution trick (the system is no longer LTI). But Mamba removes nonlinearities from the recurrence, which means the state updates are associative. You can use a parallel associative scan to compute cumulative effects in $O(N \log N)$ parallel time.

The result: training speed comparable to Transformers, inference with constant memory per step. No KV cache.

What Mamba Gets You

From the paper:

• 3B Mamba outperforms same-size Transformers on language modeling

• 5x higher inference throughput

• Linear scaling in sequence length

• Works up to million-token sequences

I’ve been running some Mamba models and the inference characteristics are genuinely different. The memory profile stays flat as context grows. You can batch much more aggressively.

Mamba-2 and the Theoretical Unification

The Mamba-2 paper (May 2024) is worth reading carefully because it reveals something deeper.

SSMs and Attention Are Related

The core insight: structured SSMs and linear attention are different decompositions of the same mathematical object (semiseparable matrices). When you restrict Mamba’s state matrix to scalar-diagonal form, the SSM becomes equivalent to a specific form of linear attention.

This isn’t just theoretical elegance. It means:

• You can use matrix multiplication as the primary primitive (16x faster on tensor cores vs general arithmetic)

• Transformer optimization techniques (tensor parallelism, sequence parallelism) transfer directly

• The “state size vs speed” tradeoff becomes more favorable

Practical Improvements

Mamba-2 gets 2-8x faster training than Mamba-1, supports larger state dimensions (N=64 or 128 vs N=16), and has a much simpler implementation. The core SSD layer is about 25 lines of code.

One interesting finding from their ablations: the extra expressivity of larger states matters less when you’re hybridizing with attention anyway. The attention layers can handle the cases where large state would help.

Other Approaches Worth Knowing

RWKV

RWKV (Receptance Weighted Key Value) has been developing in parallel. It’s now at version 7 (“Goose”), trained up to 14B parameters, and is actually deployed in production (Windows Copilot uses it for on-device inference).

The architecture uses linear attention that can be computed either as parallel attention or sequential RNN. The key components are:

• Receptance: gates how much past information to accept

• Weighted: time-decay weighting

• Key/Value: similar to attention’s K/V

RWKV-7 claims to be doing something like in-context gradient descent, essentially meta-learning its state on the fly. I haven’t dug into this deeply enough to have a strong opinion, but the claim is interesting.

The tradeoffs seem similar to Mamba: great efficiency, some struggles with associative recall over very long contexts. The community is active and iteration is fast.

xLSTM

Sepp Hochreiter (co-inventor of LSTM) released xLSTM in May 2024. It’s essentially “what if we redesigned LSTM with everything we’ve learned since 1997?”

Key innovations:

• Exponential gating: sharper on/off control than sigmoid gates, helps revise stored information

• Matrix memory (mLSTM): instead of scalar cell state, use a matrix with covariance update rules

• Parallelizable variant: removes memory mixing to enable parallel training

They’ve trained up to 7B parameters on 2.3T tokens. The scaling curves look competitive with Transformers and SSMs.

I find this interesting because it suggests the LSTM concept wasn’t fundamentally limited, just underexplored. The exponential gating in particular seems to help with the “can’t revise old decisions” problem that plagued original LSTMs.

RetNet, Griffin, etc.

Microsoft’s RetNet, Google’s Griffin, and various other architectures are exploring similar territory. The common pattern: parallel training with recurrent inference, some form of linear or gated attention. I haven’t spent enough time with these to have strong takes.

The Hybrid Approach

Here’s where I’ve landed after working through this: pure SSMs and pure linear attention still have a fundamental limitation.

The Recall Problem

Attention can directly look up any token from any other token. It’s basically a database query. SSMs compress history into fixed-size state. That compression is lossy.

For tasks like “find the exact phrase mentioned 5000 tokens ago” or needle-in-haystack retrieval, attention wins. This isn’t a bug in SSMs that can be engineered away. It’s inherent to the architecture.

Mixing Attention and SSMs

The practical solution is to use both. Jamba (AI21, March 2024) interleaves Mamba and Transformer layers at a 7:1 ratio:

• 256K token context (longest among open-weight models at release)

• Fits on a single 80GB GPU

• Competitive benchmarks

• 3x faster than Mixtral at long contexts

Zamba takes a different approach: use a shared attention block that all Mamba layers can access. Minimizes KV cache overhead while still providing lookup capability.

The pattern emerging: use SSM/linear attention for most of the processing (cheap, scales well), add just enough attention to handle precise retrieval. The ratio seems to be somewhere between 6:1 and 8:1 Mamba-to-attention.

IBM’s Bamba work (collaboration with Gu and Dao) explored optimal configurations and found that Mamba-2’s larger state dimensions matter less in hybrids because attention handles those cases anyway.

---

What I’m Still Uncertain About

Why does softmax sharpness matter so much? The theoretical connection between SSMs and linear attention is beautiful, but it doesn’t fully explain why the softmax’s sharp focus is so important for language. We’re engineering around something we don’t fully understand.

How do these scale? Most comparisons are at 3-7B parameters. Scaling behavior matters enormously. Mamba-2 claims to follow Chinchilla scaling laws, but we don’t have the same density of evidence as for Transformers.

Real-world task performance. Perplexity correlates with quality but isn’t everything. How do these architectures perform on complex reasoning, multi-step tasks, tool use? The benchmarks are still sparse.

Training stability. SSMs are sensitive to initialization and numerical precision (the Mamba repo explicitly warns about needing FP32 for parameters). Is this going to bite people at scale?

Practical Takeaways

If I were starting a new project today:

Short contexts (<8K), maximum capability needed: Standard Transformer. Well-understood, extensively optimized, proven.

Long contexts, inference-heavy, constrained hardware: Mamba or hybrid. The efficiency gains are real and significant.

General purpose, need to balance capability and efficiency: Hybrid architecture (Jamba-style). This seems to be where things are converging.

Edge deployment, streaming, real-time: Pure SSM or RWKV. Constant memory, no KV cache growth, predictable latency.

The main thing I’d avoid is assuming Transformers are always the right answer. The scaling wall is real, and the alternatives have matured enough to be practical.

Where This Is Going

My guess is that 2025 models will mostly be hybrids. Pure Transformers will be reserved for applications where cost doesn’t matter and contexts are short. The default architecture will be “mostly SSM with some attention.”

The interesting open question is whether someone will find a way to get attention-quality recall without attention-cost scaling. The theoretical connection between SSMs and attention suggests there might be something there. But for now, mixing them seems to be the pragmatic answer.

---

References I found useful:

• Mamba paper: arxiv.org/abs/2312.00752

• Mamba-2 / SSD: arxiv.org/abs/2405.21060

• Jamba: arxiv.org/abs/2403.19887

• xLSTM: arxiv.org/abs/2405.04517

• RWKV: arxiv.org/abs/2305.13048

• HiPPO: arxiv.org/abs/2206.12037

• Maarten Grootendorst’s visual guide: newsletter.maartengrootendorst.com/p/a-visual-guide-to-mamba-and-state