[Paper Review] Theoretical Foundations of Deep Selective State-Space Models

Nicola Muca Cirone, Antonio Orvieto+3 more|arXiv (Cornell University)|Feb 29, 2024

Complex Systems and Decision Making4 citations

TL;DR

This paper establishes a theoretical foundation for deep selective state-space models (SSMs) using Rough Path Theory, proving that input-controlled linear recurrences—such as in Mamba and GLA—provably capture non-linear interactions across temporal scales by projecting inputs onto the signature of the sequence. The key result shows that wide, randomly initialized selective SSMs are fully expressive without needing additional MLPs, while diagonal variants efficiently compute higher-order statistics through depth.

ABSTRACT

Structured state-space models (SSMs) such as S4, stemming from the seminal work of Gu et al., are gaining popularity as effective approaches for modeling sequential data. Deep SSMs demonstrate outstanding performance across a diverse set of domains, at a reduced training and inference cost compared to attention-based transformers. Recent developments show that if the linear recurrence powering SSMs allows for multiplicative interactions between inputs and hidden states (e.g. GateLoop, Mamba, GLA), then the resulting architecture can surpass in both in accuracy and efficiency attention-powered foundation models trained on text, at scales of billion parameters. In this paper, we give theoretical grounding to this recent finding using tools from Rough Path Theory: we show that when random linear recurrences are equipped with simple input-controlled transitions (selectivity mechanism), then the hidden state is provably a low-dimensional projection of a powerful mathematical object called the signature of the input -- capturing non-linear interactions between tokens at distinct timescales. Our theory not only motivates the success of modern selective state-space models such as Mamba but also provides a solid framework to understand the expressive power of future SSM variants.

Motivation & Objective

• To provide a theoretical framework for analyzing modern selective SSMs using tools from Rough Path Theory.
• To prove that input-controlled linear recurrences in SSMs are fully expressive in approximating continuous functions of input sequences.
• To demonstrate that diagonal SSMs (e.g., Mamba) efficiently compute higher-order statistics through depth, matching the expressivity of dense variants.
• To unify the analysis of S4, Mamba, GLA, and related models under a common mathematical formulation based on controlled differential equations (CDEs).
• To offer a generalizable theoretical basis for understanding and designing future SSM variants with enhanced expressivity and efficiency.

Proposed method

• Modeling SSMs as controlled differential equations (CDEs) driven by input sequences, where hidden states evolve via linear, input-dependent transitions.
• Applying Rough Path Theory to formalize the hidden state as a low-dimensional projection of the input's signature, capturing non-linear temporal interactions.
• Defining the state transition matrix $ W_{s,t} $ as a sum over iterated integrals of input-driven vector fields, generalizing the S4 and Mamba recurrences.
• Using the Wronskian matrix and path integral representations to derive explicit solutions for CDEs with additive input terms.
• Proving that random dense input-controlled recurrences achieve full expressivity by capturing sufficient input statistics in the hidden state.
• Demonstrating that chaining diagonal SSM blocks computes higher-order global statistics, approaching the expressivity of dense models in the depth limit.

Figure 1 : Example of the dynamics of the S4 and Mamba CDEs derived in Sec. 3.1 . The continuous-time approximation of these algorithms can be written as a Linear CDE. In S4, $\omega^{\text{X}}_{t}=t$ , while $\xi^{\text{X}}_{t}=\int_{0}^{t}X_{s}ds$ . In Mamba, $\omega^{\text{X}}_{t}=\int_{0}^{t}\si

Experimental results

Research questions

• RQ1 How can the expressive power of modern selective SSMs like Mamba be theoretically grounded?
• RQ2 What mathematical structure underlies the ability of selective SSMs to model non-linear interactions across long sequences?
• RQ3 Can input-controlled linear recurrences achieve full expressivity without additional non-linear components like MLPs?
• RQ4 How do diagonal SSMs compare to dense variants in capturing higher-order input statistics?
• RQ5 Can a unified framework based on controlled differential equations and Rough Path Theory be used to analyze and compare diverse SSM architectures?

Key findings

• Input-controlled linear recurrences in selective SSMs provably represent the hidden state as a low-dimensional projection of the input's signature, capturing non-linear temporal interactions.
• Wide, randomly initialized dense input-controlled SSMs are fully expressive: they can approximate any continuous function from input sequences to outputs without requiring additional MLPs.
• Diagonal SSMs such as Mamba efficiently compute higher-order statistics through depth, achieving expressivity comparable to dense models in the depth limit.
• The hidden state of selective SSMs is mathematically equivalent to solving a controlled differential equation (CDE), with explicit solutions derived via path integral representations and Wronskian matrices.
• The framework unifies S4, Mamba, GLA, and related models under a common theoretical basis, enabling rigorous comparison and analysis of future SSM variants.
• The theory explains why selective SSMs outperform standard S4 and attention mechanisms in long-range modeling: they capture richer input statistics through structured, input-dependent transitions.

Figure 3 : Comparison of Linear CDE, Mamba and S5 on anti-symmetric signature prediction tasks. As suggested by the theory Mamba struggles to immediately capture these high order statistics, but performance improves with chaining . On the other hand S5 does not learn such non-linear features even af