Provable Benefits of Complex Parameterizations for Structured State Space Models

Yuval Ran-Milo*, Eden Lumbroso, Edo Cohen-Karlik, Raja Giryes, Amir Globerson, Nadav Cohen

*Corresponding author for this work

Research output: Contribution to journalConference articlepeer-review

Abstract

Structured state space models (SSMs), the core engine behind prominent neural networks such as S4 and Mamba, are linear dynamical systems adhering to a specified structure, most notably diagonal. In contrast to typical neural network modules, whose parameterizations are real, SSMs often use complex parameterizations. Theoretically explaining the benefits of complex parameterizations for SSMs is an open problem. The current paper takes a step towards its resolution, by establishing formal gaps between real and complex diagonal SSMs. Firstly, we prove that while a moderate dimension suffices in order for a complex SSM to express all mappings of a real SSM, a much higher dimension is needed for a real SSM to express mappings of a complex SSM. Secondly, we prove that even if the dimension of a real SSM is high enough to express a given mapping, typically, doing so requires the parameters of the real SSM to hold exponentially large values, which cannot be learned in practice. In contrast, a complex SSM can express any given mapping with moderate parameter values. Experiments corroborate our theory, and suggest a potential extension of the theory that accounts for selectivity, a new architectural feature yielding state of the art performance.

Original languageEnglish
JournalAdvances in Neural Information Processing Systems
Volume37
StatePublished - 2024
Event38th Conference on Neural Information Processing Systems, NeurIPS 2024 - Vancouver, Canada
Duration: 9 Dec 202415 Dec 2024

Funding

FundersFunder number
Blavatnik Family Foundation
Google
Yandex Initiative in Machine Learning
Adelis Research Fund for Artificial Intelligence
Tel Aviv University
Amnon and Anat Shashua
Google Research Gift
European Research CouncilHOLI 819080, NN4C 101164614
Israel Science Foundation1780/21

    Fingerprint

    Dive into the research topics of 'Provable Benefits of Complex Parameterizations for Structured State Space Models'. Together they form a unique fingerprint.

    Cite this