Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Theoretical Foundations of Deep Selective State-Space Models

Nicola Muca Cirone, Antonio Orvieto, Benjamin Walker, Cristopher Salvi, Terry Lyons

TL;DR

This paper analyzes deep selective state-space models (SSMs) through Rough Path Theory, recasting gated linear recurrences as Linear Controlled Differential Equations driven by input-dependent gates. It proves that dense Linear CDEs achieve universal expressivity via a signature-based expansion, and shows that diagonal, gate-driven recurrences (e.g., Mamba) are expressive but inherently limited unless stacked to recover full power; randomness with a trainable readout can also achieve strong approximation via an RKHS of signatures. The Path-to-Path learning extension demonstrates how an MLP on top of Linear CDEs can approximate time-varying, path-dependent functions, linking architectural design to kernel-method perspectives. Empirically, the framework accounts for observed strengths and weaknesses across S4, Mamba, and linear CDEs on antisymmetric-signature tasks and a state-tracking benchmark, validating the theoretical claims about expressivity, gates, and chaining while guiding future SSM development.

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.

Theoretical Foundations of Deep Selective State-Space Models

TL;DR

This paper analyzes deep selective state-space models (SSMs) through Rough Path Theory, recasting gated linear recurrences as Linear Controlled Differential Equations driven by input-dependent gates. It proves that dense Linear CDEs achieve universal expressivity via a signature-based expansion, and shows that diagonal, gate-driven recurrences (e.g., Mamba) are expressive but inherently limited unless stacked to recover full power; randomness with a trainable readout can also achieve strong approximation via an RKHS of signatures. The Path-to-Path learning extension demonstrates how an MLP on top of Linear CDEs can approximate time-varying, path-dependent functions, linking architectural design to kernel-method perspectives. Empirically, the framework accounts for observed strengths and weaknesses across S4, Mamba, and linear CDEs on antisymmetric-signature tasks and a state-tracking benchmark, validating the theoretical claims about expressivity, gates, and chaining while guiding future SSM development.

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.
Paper Structure (47 sections, 30 theorems, 185 equations, 2 figures)

This paper contains 47 sections, 30 theorems, 185 equations, 2 figures.

Table of Contents

  1. Introduction
  2. Contributions
  3. State-space Models
  4. Review of Modern SSMs
  5. SSM basics and S4.
  6. Mamba.
  7. Known properties of (non-linear) recurrences
  8. SSMs as Linear CDEs
  9. Linear CDEs
  10. Expressivity of Linear CDEs
  11. Characterization of the closure
  12. Intuition on the closure result: role of gates
  13. Warmup -- Linear RNNs.
  14. How does a Linear CDE process inputs?
  15. The Diagonal Case
  16. ...and 32 more sections

Key Result

Theorem 4.1

Let $\mathbb{X}\subset C_{1,0}([0,1],\mathbb{R}^{d})$ be compact and choose continuous gating functions $\omega, \xi$ such that These assumptions are of technical nature and can be relaxed at the cost of talking about tree-like equivalence of paths cf.hambly2010uniqueness. Intuitively $\omega^{\text There exist dense matrices $A_1,...,A_{d_\omega},B$ such that, after a fixed final linear projectio

Figures (2)

  • Figure 1: Comparison of the Linear CDE, Mamba, and S5 on the anti-symmetric signature prediction tasks. For each model, we plotted the mean and range of the validation accuracy over 5 independent runs.
  • Figure 2: For each sequence length, the plot shows the minimum number of blocks required to achieve at least $90\%$ validation accuracy, with each grey band corresponding to a number of blocks. Missing points mean the model did not achieve at least $90\%$ validation accuracy with $4$ blocks or less.

Theorems & Definitions (69)

  • Remark 2.1
  • Definition 3.1: Linear CDE
  • Remark 3.2: Are hidden state components recurrently mixed?
  • Theorem 4.1
  • Theorem 4.2
  • Theorem 4.3: Diagonal Case
  • Corollary 4.4: Mamba Case
  • Proposition 4.5
  • Proposition 4.6
  • Proposition 5.1
  • ...and 59 more