Provable Benefits of Complex Parameterizations for Structured State Space Models

Thank you for your support, and for highlighting the importance of the question we study and the merit of our theory and experiments! We address your comments and questions below.

Addressing selectivity:

We agree with you that addressing selectivity — i.e., theoretically supporting the success of real parameterizations in selective SSMs — is an interesting pursuit. Below we qualitatively discuss the reasoning behind our belief that “our results may serve as an important stepping stone in this pursuit”. A more detailed version of this discussion will be added to the paper. Thank you for raising the matter!

Roughly speaking, the separations we established between real and complex SSMs result from a gap in their ability to express oscillations, i.e., to express frequency components in their impulse response: while the complex SSM can easily express any frequency, the real SSM is required to have exponential dimension or parameter magnitudes. Adding selectivity to the real SSM means that its parameters become input-dependent, leading to what can be viewed as an input-dependent impulse response. It can be shown that this dependence allows “importing” frequency components from the input to the impulse response. Therefore, if the input data is sufficiently rich in terms of its frequency content, selectivity may endow real parameterizations with all the benefits we proved for complex parameterizations.

Recall that, as stated in our introduction, [10] conjectured that complex parameterizations are preferable for continuous data modalities (e.g., audio, video), whereas for discrete data modalities (e.g., text, DNA) real parameterizations suffice. The explanation above aligns with this conjecture: continuous data modalities typically consist of low frequencies only, whereas discrete data modalities have a “whiter spectrum,” i.e. a more uniform mix of frequencies.

As stated, a detailed version of the above discussion will be added to the paper.

Footnote 1:

Per your suggestion, we will remove Footnote 1.

Questions (addressed by order of appearance):

• We will explicitly mention S5 when listing prominent SSM-based neural network architectures.
• See “addressing selectivity” section above.
• Our theory essentially applies as is (more precisely, it applies given very slight modifications in the established bounds) to the case where SSMs include a feedthrough $D$. The reason why incorporation of $D$ does not make a material difference is that it can easily be emulated by: (i) increasing the dimension of the SSM by one; (ii) assigning one to the last diagonal entry of $A$ and the last entry of $B$; and (iii) assigning $D$ to the last entry of $C$. We will mention this in the camera-ready version. Thank you for raising the matter!
• $\Omega (\cdot )$ in line 187 stands for Big-Omega notation, i.e. it signifies a function that is at least linear in its argument. Note that it is used to qualitatively characterize a dependence whose precise form is provided by Theorem 2.
• In Proposition 3 (and Proposition 1), it is indeed possible to improve (relax) the requirement $n_{\mathbb{C}} \geq t$. Namely, by leveraging the fact that the discrete Fourier transform of a real sequence is symmetric, one can update the proof of Proposition 3 to show that $n_{\mathbb{C}} \geq \lceil t / 2 \rceil$ suffices. Parameter counting arguments imply that the latter requirement is tight. We will mention all of this in the camera-ready version. Thank you for asking about it!

Thank you for a very clear and interesting rebuttal.

I think your discussion of selectivity is very interesting and should be included in as much details as possible in the camera ready version.

Your new experiments in the .pdf are great as well! They should also be included in camera ready.

I am raising my score to an 8.

Thank you for your review, and for highlighting the thoroughness and preciseness of our theory. We address your comments and questions below. If you find our responses satisfactory, we would greatly appreciate it if you would consider raising your score.

Significance of real diagonal SSMs:

Real diagonal SSMs (with and without selectivity) underlie some of the most prominent neural network architectures for long sequence processing. This includes the recent Mamba architecture and its variants, as well as other architectures [A, B, C, D, E]. Advantages of real diagonal SSMs include simplicity and speed. The extent to which they compare to complex diagonal SSMs in terms of accuracy has been the subject of much debate (see our introduction). This debate is what motivated our work.

We hope the above addresses your concern regarding the significance of our contribution. If not, please let us know and we will gladly elaborate.

--

[A] “Mamba: Linear-time sequence modeling with selective state spaces”, Gu and Dao, 2023

[B] “Mega: Moving Average Equipped Gated Attention”, Ma et al., 2023

[C] “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality”, Dao and Gu, 2024

[D] “Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length”, Ma et al., 2024

[E] “Jamba: A Hybrid Transformer-mamba language model”, Lieber et al., 2024

Experimentation:

Following your comment, we conducted several experiments, including:

• Demonstration of the gap in practical learnability between real and complex SSMs in the theoretically analyzed setting (i.e., in the setting for which we established such a gap).
• Demonstration that complex parameterizations for (non-selective) SSMs improve performance of prominent neural network architectures (S4) on real-world data (sequential CIFAR-10, included in the standard long range arena benchmark).
• Ablation study showing that among the architectural differences between S6 and S4, the one primarily responsible for closing the gap between real and complex parameterizations for SSMs, is the selectivity of the input and output matrices $B$ and $C$.

The results of these experiments can be found in the PDF attached to our global response.

Impact of discretization:

Our theory essentially applies as is (more precisely, it applies given slight modifications in the established bounds) to cases where the SSMs include conventional discretizations. For example, with the discretization laid out in the original Mamba paper [10]: the state transition matrix $A$ is represented as the exponent of a parameter matrix $A’$ times a scalar $\Delta$; and the input matrix $B$ is represented as a parameter matrix $B’$ times $\Delta$ times some matrix with spectral norm no greater than one. All of our proofs readily adapt to this case (in particular, the proofs of Theorem 2, Corollary 1 and Corollary 2 readily adapt to establishing an exponential lower bound on the magnitudes of $( B’ , C )$ for the real SSM).

We will add a discussion of discretization to the camera-ready version. Thank you very much for raising the matter!

Complex $A$ with real $B$ and $C$:

Your intuition is correct. Namely, all benefits we have proven for the complex SSM apply (given a slight modification in one of the established bounds) when the input matrix $B$ and the output matrix $C$ are constrained to be real. To see this, note that our proofs of Propositions 1, Proposition 2 and Theorem 1 do not make use of complex assignments for $B$ and $C$. Our proof of Proposition 3 does make use of such assignments, but can readily avoid them by replacing the discrete Fourier transform with the discrete cosine transform (in this case a factor of $\sqrt{2}$ is introduced in the second bound). We will mention this in the camera-ready version. Thank you for the question!

Mappings poorly approximated by the real SSM:

While Corollary 1 indeed provides an example (albeit a canonical one) of a mapping poorly approximated by the real SSM, Corollary 2 goes far beyond an example, in the sense that it applies to a random (generic) mapping. Nonetheless, we agree with you that it would be interesting to relate the condition brought forth by Theorem 2 — namely, forward differences not being especially small — to easily interpretable properties of the mappings. We invested considerable efforts trying to establish such a relation (primarily using tools from complex analysis, e.g. Rice’s method), but to no avail. We will mention this pursuit as a direction for future work in the camera-ready version. Thank you for raising the matter!

Thank you for your support!

We conducted several additional experiments, including:

• Demonstration of the gap in practical learnability between real and complex SSMs in the theoretically analyzed setting (i.e., in the setting for which we established such a gap).
• Demonstration that complex parameterizations for (non-selective) SSMs improve performance of prominent neural network architectures (S4) on real-world data (sequential CIFAR-10, included in the standard long range arena benchmark).
• Ablation study showing that among the architectural differences between S6 and S4, the one primarily responsible for closing the gap between real and complex parameterizations for SSMs, is the selectivity of the input and output matrices $B$ and $C$.

The results of these experiments can be found in the PDF attached to our global response.

Thank you for the thoughtful review, for highlighting that “with some tweaks, this can be a very good paper”, and for the willingness to “engage in the discussion so [as] to update [your] borderline score”. Below we address your comments and questions.

Specificity of Theorem 1:

Theorem 1 readily extends to any function realized by the complex SSM whose impulse response oscillates. More precisely, the proof of Theorem 1 can easily be extended to establish a lower bound on $n_{\mathbb{R}}$ that is linear in $t$ whenever the restriction to odd elements of the impulse response of the complex SSM admits a number of sign changes linear in $t$. With this extension, Theorem 1 applies to a positive measure subset of the complex SSM’s function class. Moreover, among the subset of the complex SSM’s function class which is characterized by $A_{\mathbb{C}}$ having entries close to the unit circle, the above-described extension of Theorem 1 can be shown to apply to “most” functions. We will elaborate on this in the camera-ready version. Thank you for bringing it up!

Notwithstanding the above, we note that proving separation in expressiveness through particular counterexamples is very common. Indeed, various important results in deep learning theory fall under this category (see references in “On the Expressive Power of Deep Learning: A Tensor Analysis” by Cohen et al.).

Opening of Section 3.3:

Following your comment, this text was completely refactored. In particular, we broke down Section 3.3 to three subsubsections, titled: “Real Parameterizations Suffer from Exponentiality”, “Exponentiality Impedes Practical Learning”, and “Complex Parameterizations Do Not Suffer from Exponentiality”. We believe Section 3.3 is much clearer now. Thank you for the feedback!

Relation to Orvieto et al.:

Thank you for raising this point! Upon closer inspection, we indeed believe that the result of Orvieto et al. ([25], Section 4.1) can be viewed as a special case of ours. In particular, the task considered in Orvieto et al. — namely, reconstructing past inputs to a real SSM from its current state using a linear mapping — can be viewed as taking an output matrix $C_{\mathbb{R}}$ with which the input-to-output mapping of a real SSM is a canonical copy (delay) mapping. Our Corollary 1 implies that this necessitates exponential parameter magnitudes, therefore essentially recovers the result of Orvieto et al. We stress that our Theorem 2 is far more general than our Corollary 1 (and accordingly, than the result of Orvieto et al.). Indeed, our Theorem 2 applies (i.e., ensures exponential parameter magnitudes for the real SSM) not only with a copy input-to-output mapping, but with any input-to-output mapping whose impulse response has forward differences that are not especially small (this includes, e.g., random input-to-output mappings — see Corollary 2).

We will detail the relation between the result of Orvieto et al. and ours in the camera-ready version. Thanks again!

Empirical demonstration of bounds in Section 3.3:

Following your request, we conducted experiments demonstrating the bounds in Section 3.3. The results of these experiments can be found in the PDF attached to our global response. In a nutshell, the results show that while the complex SSM is able to learn all impulse responses (up to times no greater than its dimension), when the impulse response is such that the real SSM is provably required to have exponential parameter magnitudes, training the real SSM does not converge.

Further experiments:

Following your feedback, we conducted further experiments, including:

• Demonstration that complex parameterizations for (non-selective) SSMs improve performance of prominent neural network architectures (S4) on real-world data (sequential CIFAR-10, included in the standard long range arena benchmark).
• Ablation study showing that among the architectural differences between S6 and S4, the one primarily responsible for closing the gap between real and complex parameterizations for SSMs, is the selectivity of the input and output matrices $B$ and $C$.

The results of these experiments can be found in the PDF attached to our global response.

Success of real parameterizations in selective SSMs:

We agree with you that theoretical support for the success of real parameterizations in selective SSMs is an extremely crucial pursuit. As stated in the paper, we believe our results may serve as an important stepping stone in this pursuit. Below we qualitatively discuss the reasoning behind our belief. A more detailed version of this discussion will be added to the paper. Thank you for raising the matter!

Roughly speaking, the separations we established between real and complex SSMs result from a gap in their ability to express oscillations, i.e., to express frequency components in their impulse response: while the complex SSM can easily express any frequency, the real SSM is required to have exponential dimension or parameter magnitudes. Adding selectivity to the real SSM means that its parameters become input-dependent, leading to what can be viewed as an input-dependent impulse response. It can be shown that this dependence allows “importing” frequency components from the input to the impulse response. Therefore, if the input data is sufficiently rich in terms of its frequency content, selectivity may endow real parameterizations with all the benefits we proved for complex parameterizations.

Recall that, as stated in our introduction, [10] conjectured that complex parameterizations are preferable for continuous data modalities (e.g., audio, video), whereas for discrete data modalities (e.g., text, DNA) real parameterizations suffice. The explanation above aligns with this conjecture: continuous data modalities typically consist of low frequencies only, whereas discrete data modalities have a “whiter spectrum,” i.e. a more uniform mix of frequencies.