Quadratic Parameter Requirement Will Find You If You Try to Get Away from Quadratic Memory
Why Mamba and State Space Models trade attention complexity for a narrower gradient path and why the Transformer’s O(N²) cost isn't a bug, but the price of global learning
The Mamba Hype: Why Did It Become This Popular?
From the day that GPT-3 launched, Transformer is undisputed ‘king’ of LLMs. But for almost 9 years, this architecture has a defect: Quadratic Cost (O(N^2)). Since the day we noticed upscaling this models becoming extremely expensive, researchers started to work for achieving new architectures. In this phase, many architectures got popular (Linformer, Performer, RWKV etc.) but all of them failed.
Until Mamba.
Mamba presented token-based selectivity with SSM models, benchmarks showed models keep state well and don’t forget like previous architectures and even get closer to Transformer. So Mamba emerged as the first truly viable alternative after launch of Transformer. Everybody saw future in this models.
A Quick Look at Mamba’s ‘Sequential’ Gradient Flow
To understand why Transformers win, we shouldn’t just look at how models read context, we should also look “how they learn”.
Mamba is a linear costed (O(N)) model unlike Transformer. It provides this linearity by processing data sequential. Yes, Mamba has some “Assosicative Scan” tricks and it makes training parallel, but let’s look at state update formula of Mamba:
The output at “t” depends on the updated state of step “t-1”. So Backpropagation Through Time must flow backwards sequentially through this state matrices. This makes model optimizing to tokens step by step, unlike Transformer. Because Transformer see all sequence at once and optimizes all tokens with O(1) computational distance.
Note: If Mamba were a bidirectional model (like BERT), this wouldn’t strictly ruin gradient quality because the full future context would balance the sequence. But bidirectional training causes loss of causal mask and generative feature of model.
Quadratic Parameter Requirement of Sequence Representation in Mamba
So we said that Mamba optimizes tokens locally one by one. Consider a four-token sequence: “A, B, C, D”. During Backpropagation Through Time:
a. The model processes A and B, optimizing B's representation solely with respect to A.
b. Next, it processes C given the combined state of (A, B), optimizing C for that accumulated context.
c. Finally, it processes D given the state of (A, B, C), attempting to optimize D accordingly.
The problem is, model doesn’t optimizes just the new token, optimizes also the state parameters, so this causes raise of parameter requirement for the tokens middle of sequence. So if you don’t want to ruin capacity of understanding, you should increase representation space quadratic. That’s the tradeoff of Mamba.
Conclusion
Mamba and its linear efficiency is a remarkable engineering breakthrough. But there is no free lunch in deep learning. Mathematically, to represent interaction of all tokens in a sequence, you are bound to quadratic cost. Transformer’s success is representing this interaction pool in best known way in the world.