On the Expressivity of Selective State-Space Layers: A Multivariate Polynomial Approach

2025-02-04

AI Summary

This work addresses the lack of theoretical characterization of the expressive power of selective state-space models (SSMs)—the core component of Mamba. We establish, for the first time, a rigorous expressivity comparison framework between selective SSMs and linear Transformers. Leveraging multivariate polynomial approximation theory and state-space model analysis, we prove that, under equal parameter budgets, selective SSMs can exactly represent a strictly broader class of long-range dependency functions than linear Transformers; this advantage is particularly pronounced in long-sequence modeling and does not compromise generalization performance. Our theoretical findings are empirically validated across multiple long-sequence benchmarks—including PG19 and ImageNet-64—demonstrating consistent gains in modeling capacity and predictive accuracy. This constitutes the first formal proof of expressive superiority for Mamba-style architectures over linearized Transformer variants.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Natural Language Processing: (Large) Language Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

Recent advances in efficient sequence modeling have introduced selective state-space layers, a key component of the Mamba architecture, which have demonstrated remarkable success in a wide range of NLP and vision tasks. While Mamba's empirical performance has matched or surpassed SoTA transformers on such diverse benchmarks, the theoretical foundations underlying its powerful representational capabilities remain less explored. In this work, we investigate the expressivity of selective state-space layers using multivariate polynomials, and prove that they surpass linear transformers in expressiveness. Consequently, our findings reveal that Mamba offers superior representational power over linear attention-based models for long sequences, while not sacrificing their generalization. Our theoretical insights are validated by a comprehensive set of empirical experiments on various datasets.

Problem

• explore theoretical foundations of selective state-space layers
• compare expressiveness with linear transformers
• validate representational power on diverse datasets

Innovation

• Selective state-space layers
• Multivariate polynomial approach
• Superior representational power