Provable Benefits of Complex Parameterizations for Structured State Space Models (2410.14067v2)

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.

• The paper demonstrates that complex SSMs can represent real SSM mappings within moderate dimensions.
• The paper finds that real parameterizations require exponentially large values, making learning computationally impractical.
• The paper validates its theory with experiments showing marked performance improvements on both synthetic tasks and real-world data.

Insights into Complex Parameterizations for Structured State Space Models

This paper explores the intricacies of Structured State Space Models (SSMs), particularly focusing on the role of complex parameterizations. SSMs with diagonal structures are pivotal in modern neural architectures such as S4 and Mamba. Unlike typical neural modules, these models employ complex rather than real parameters, which has been the subject of ongoing theoretical exploration.

The authors present a significant contribution by establishing formal distinctions between real and complex diagonal SSMs. Their theoretical results are twofold. Firstly, the paper demonstrates that complex SSMs can express any mapping of a real SSM within moderate dimensions, while real SSMs require a substantially higher dimensionality to express mappings of complex SSMs. Secondly, even if a real SSM can theoretically express a mapping, doing so in practice necessitates exponentially large parameter values, rendering learning impractical. Complex SSMs avoid this issue, maintaining reasonable parameter magnitudes.

Key Findings

1. Expressiveness and Dimensionality:
   • Complex SSMs are shown to possess superior expressiveness due to their ability to succinctly represent oscillatory mappings, a feature less efficiently captured by real SSMs.
2. Practical Learnability:
   • The paper highlights a critical scalability issue with real SSMs, where either extremely high dimensionality or exponential parameter magnitudes are necessary to achieve similar mappings as their complex counterparts. Learning such parameters is computationally infeasible.
3. Experimental Corroboration:
   • The authors support their theoretical assertions with controlled experiments. Complex parameterizations displayed significantly improved performance on both synthetic tasks designed to mimic theoretical assumptions and real-world data settings.

Implications

The implications are both theoretical and practical. Theoretically, the results offer a deeper understanding of when and why complex parameterizations outperform real ones, particularly in capturing the frequency components essential in linear dynamical systems. Practically, the findings suggest a potential shift towards incorporating complex parameters in architectural designs, especially in models handling continuous modalities such as audio and video.

Future Directions

The paper suggests promising avenues for future research, especially concerning how selectivity in SSMs—where parameter values are input-dependent—might bridge the performance gap, enabling real SSMs to approximate complex behavior more effectively. Exploring this could provide comprehensive insights into harnessing the full potential of SSM architectures.

Overall, the paper presents a thorough examination of complex parameterizations in SSMs, offering foundational insights and practical recommendations for both designing and applying these models across various AI domains. The detailed theoretical analysis combined with robust empirical validation strengthens its contribution to ongoing research in neural network architecture.