[KongchangAI]
· 2 min read· 1,490 words

The Hidden Cost of Mamba: Escaping Quadratic Memory Only to Face Quadratic Parameter Demands

The Hidden Cost of Mamba: Escaping Quadratic Memory Only to Face Quadratic Parameter Demands

Mamba escapes quadratic memory but may face quadratic parameter demands—complexity shifts, not disappears.

This article examines a provocative argument that Mamba's SSM architecture, while achieving linear inference complexity, may face hidden quadratic parameter requirements during training. Sequential gradient propagation creates compounding optimization burdens, potentially requiring representation space to grow quadratically for equivalent expressiveness. The analysis connects this to information-theoretic limits on compressing long sequences into fixed-size states, while noting this remains an unverified theoretical perspective that hybrid architectures may pragmatically address.

Introduction: The Quadratic Dilemma of Attention Mechanisms

Since the advent of the Transformer, the quadratic complexity of Self-Attention has been an inescapable bottleneck for long-sequence modeling. When the sequence length is N, both the computation and storage costs of the attention matrix grow at O(N²), directly limiting large models' ability to handle ultra-long contexts and spawning a wave of linearization approaches.

To appreciate the severity of this bottleneck, we need to return to the core computational logic of Transformers: the self-attention mechanism computes relevance scores between every token and all other tokens in the sequence, generating an N×N attention matrix where each element represents the association strength between two positions. When N scales from 4,096 to 128K or even 1M, computational costs explode—a 128K-length sequence requires 1,024 times the attention computation of a 4K-length sequence. This bottleneck has given rise to FlashAttention (reducing memory access through IO-aware tiled computation), sparse attention (such as Longformer's sliding window), and various linear attention schemes.

Mamba is among the most prominent of these alternative architectures. Leveraging the selective mechanism of State Space Models (SSM), Mamba achieves long-sequence modeling with linear complexity, positioning it as a formidable challenger to Transformers. However, an analysis published on Substack that sparked heated discussion on Reddit points out: when you try to escape quadratic memory, quadratic parameter demands may find you in another dimension.

reddit source

Mamba's Core Trade-off: From Memory to Parameters

The Mathematical Foundations of State Space Models

Before diving into Mamba's training challenges, it's essential to understand the underlying State Space Model principles. SSMs originate from a set of continuous-time differential equations in control theory: dx/dt = Ax + Bu, y = Cx + Du, where x is the hidden state, u is the input, y is the output, and A/B/C/D are system matrices. After discretization, we get the recurrence form: x_k = Āx_{k-1} + B̄u_k. Mamba's key innovation is the introduction of a "selective mechanism"—making matrices B and C dynamically depend on input content, enabling the model to selectively remember or forget historical information based on current input. This breaks the constraint of traditional Linear Time-Invariant (LTI) SSMs and grants the model content-aware capabilities, but it also means the model can no longer be efficiently trained in parallel using convolution form, instead relying on a selective scan algorithm for sequential computation.

The Optimization Challenge of Sequential Gradients

The core argument of this analysis centers on Mamba's training approach. Unlike Transformers, which can compute attention over the entire sequence in parallel, Mamba's state updates are sequential—each token's representation depends on the accumulated state from all preceding tokens. This is an efficiency advantage during inference, but during training (especially backpropagation), it introduces subtle optimization difficulties.

The author uses a four-token sequence "A, B, C, D" to illustrate this process. In Backpropagation Through Time (BPTT):

  • When the model processes A and B, it optimizes B's representation based solely on A;
  • Then it processes C, which now faces the combined state of (A, B) and must be optimized against this accumulated context;
  • Finally it processes D, which must be optimized based on the state of (A, B, C).

BPTT is the standard method for training such recurrent structures. Its core idea is to unroll the recurrent network across time steps into a deep feedforward network, then apply standard backpropagation. For a sequence of length T, gradients must propagate backward from the last time step to the first, passing through T matrix multiplications. This brings two classic challenges: vanishing gradients (gradients decay exponentially during backpropagation, making early information difficult to learn) and exploding gradients (gradients grow exponentially, causing training instability). While Mamba mitigates these issues to some extent through careful parameterization—particularly the diagonalized A matrix structure—the optimization coupling effect inherent in sequential dependencies remains.

The Cascading Optimization of State Parameters

The key issue is this: when optimizing each new token, the model isn't just optimizing that new token itself—it's simultaneously optimizing the state parameters it depends on. This means tokens in the middle of the sequence bear a heavier optimization burden—they must both express their own information and maintain consistency across the entire accumulated state.

As sequences grow longer, this compounding effect becomes increasingly pronounced. Tokens in the middle of the sequence must, within an increasingly compressed representation space, simultaneously serve "past memory" and "current information." This constitutes Mamba's hidden parameter cost.

The Theoretical Implication of Quadratic Parameter Demands

Why Representation Space Must Grow Quadratically

From this, the author draws a striking conclusion: if you don't want to degrade the model's understanding capability, the representation space must grow quadratically. In other words, while Mamba escapes quadratic growth in memory complexity, it may re-encounter quadratic demands in the parameter dimension.

This is a classic "whack-a-mole" trade-off. Transformers spend their quadratic cost on the memory and computation of the attention matrix; if Mamba wants to maintain equivalent expressiveness on long sequences, it may have to pay a quadratic price in state dimension. Both architectures merely shift the complexity from one place to another.

Does Long-Sequence Modeling Have a Complexity Conservation Law?

This observation touches on a deeper question: does long-sequence modeling obey some kind of "conservation law"? Intuitively, to precisely preserve information over long distances, a model must have sufficient capacity to encode these dependencies. Regardless of architectural design, the information-theoretic lower bound may be fundamentally unavoidable.

The Data Processing Inequality from information theory provides formal support for this intuition: after any deterministic or stochastic transformation, information can only decrease, never increase. When an SSM uses a fixed-dimension d hidden state to compress a history sequence of length N, Rate-Distortion Theory establishes a theoretical upper bound on information retention capacity. If the sequence contains O(N) bits of relevant information but the hidden state has only O(d) bits of capacity, then when N is much greater than d, information loss is inevitable. This also explains why some researchers argue that to losslessly process all dependencies in long sequences, model capacity must somehow scale with sequence length.

SSM-class models compress history into a fixed-size hidden state—this compression is inherently lossy. To reduce information loss, state capacity must expand accordingly. This logic is theoretically compelling, but still requires more empirical support.

A Balanced Perspective: The Value and Limitations of This Argument

An Analytical Perspective Awaiting Verification

It's important to emphasize that this argument comes from a single source (a Substack author's analysis, disseminated via Reddit) and has not undergone peer review or large-scale experimental validation. The proposed "quadratic parameter demand" is more of a theoretical extrapolation and intuitive argument than a rigorous mathematical proof or empirical result.

In fact, academia and industry hold differing views on Mamba's actual performance. Many benchmarks show that Mamba and its variants (such as Mamba-2) can indeed achieve results comparable to Transformers on long-sequence tasks at lower cost. Notably, Mamba-2 reformulates selective SSMs as a Structured State Space Duality (SSD) framework, revealing deep connections between SSMs and attention mechanisms—selective SSMs are essentially equivalent to a special semiseparable matrix form of attention. This discovery not only improves Mamba-2's training speed by 2-8x but also provides new theoretical perspective for understanding the expressiveness boundaries of SSMs. This suggests that even if the "quadratic parameter demand" exists, it may not pose a fatal obstacle in practice—engineering choices like selective mechanisms and hybrid architectures can effectively mitigate this issue.

Three Implications for Model Architecture Design

Regardless of whether this argument ultimately holds, the thinking it provokes has practical value:

  1. There's no free lunch. Any architecture claiming to "completely solve quadratic complexity" should be scrutinized for where it shifts the cost.
  2. Training dynamics and inference efficiency should be evaluated separately. Mamba's inference advantages are clear, but training optimization difficulty and parameter efficiency are a separate dimension requiring independent consideration.
  3. Hybrid architectures may be the pragmatic choice. Many cutting-edge models are currently combining SSM with attention layers for the best of both worlds, rather than pursuing purely linear solutions. Models like Jamba (AI21), Zamba (Zyphra), and Griffin (DeepMind) all interleave SSM layers with attention layers in practice—using SSM layers to efficiently handle local and mid-range dependencies, and a few attention layers to capture critical long-range dependencies, achieving a pragmatic balance between efficiency and expressiveness.

Conclusion: Mamba Is Not a Silver Bullet, But Remains an Important Innovation

"Quadratic parameter demands will find you when you escape quadratic memory"—this somewhat dramatic assertion reminds us to stay clear-eyed about the complexity challenges of long-sequence modeling. Mamba is undoubtedly one of the most important architectural innovations in recent years, but it is not a silver bullet.

This discussion about Mamba's true costs is, at its core, asking a fundamental question: how much does it cost to compress infinite history into a finite hidden state? The answer may have to wait for more rigorous experiments and theoretical work. Until then, maintaining critical thinking—neither deifying linear architectures nor underestimating their value—is the rational stance.

#Mamba#state space model#quadratic complexity#Transformer#SSM#long sequence modeling#hybrid architecture#linear attention
Share:

Related articles