State-space Models with Layer-wise Nonlinearity are Universal Approximators with Exponential Decaying Memory

Summary

The paper provides a constructive proof that state-space models (SSMs) are universal approximators of sequence-to-sequence mappings. Moreover, it shows that SSMs (and even empirically S4) suffer from exponentially decaying memory just like standard RNNs.

Strengths

Up to my knowledge this is the first (constructive) proof that SSMs are universal. I find it in particular interesting that SSMs also suffer from exponentially decaying memory (given the provided definition of such).

Weaknesses

Regrettably, the paper exhibits a significant lack of quality in its writing, containing numerous grammatical errors and typos. It is imperative that it undergoes a comprehensive proofreading process to address these issues. From my understanding, the paper proves universality for SSMs, which does not include universality for S4 (i.e., SSMs with specific structure of the matrices). This is a major limitation, as in practice no one uses simple SSMs, but only S4. Can S4 be included in the proof? I presume one has to show that also S4 can approximate element-wise functions up to any precision. The section about the curse of memory is very hard to follow. Can you please rewrite it and properly introduce the concept of memory functions? The paper is not rigorously written. For instance, it would be much more readable, if the authors would define every variable, e.g., $W \in \mathbb{R}^{m\times m}$ and so on. This would also help the authors avoid using variables that haven't been introduced before. The provided full paper in the appendix and the main paper differ: for instance in equation 9: $\mathbf{H}$ depends on $t$ in the full paper, but not in the main paper. Form what I understood, the universality proof of proposition 3.6 approximates functions, where the target value at every index depends on all elements in the full input sequence. However, this is not realistic, as mainly causal operators between sequences have to be learned, i.e., operators that are independent of the future values of the input sequence at every index. Can you please change the proof accordingly?

Questions

See weaknesses

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

1 poor

Contribution

3 good

Limitations

The paper provides an interesting and important proof that SSMs are universal, as well as that the memory functions decay exponentially. However, the proof does not include S4 models. Moreover, the quality of writing is very poor, which makes it hard to follow. Overall, the provided proof is very simple, which makes it appealing.

Summary

This paper attempts to show several properties of SSMs: a) SSMs can approximate element-wise functions and temporal convolutions b) Universality of SSMs c) Exponential memory decay. The authors also show experiments of SSM's exponential memory decay and compare them against variants of RNN architectures.

Strengths

* Thorough background section and clear introduction of the RNN and SMM architecture. * Clear writing. * The usage of the Kolmogorov-Arnold thoerem and Volterra series is interesting.

Weaknesses

* My main concern is most of the main results seems trivial: When taking away the tricks that makes SSMs work such as smarter initialization and diagonal parameterisation, architecture wise, SSMs are simply an RNN without non-linearly between hidden states. Hence, the theory results (i.e. proposition 3.1, 3.2, 3.12 and theorem 3.13) seems to transfer very trivially from RNNs. * The 2 formulations for showing universality of SMMs in section 3.2 seems to deviate greatly from the SSM architecture (i.e. figure 3 and 4 does not follow the definition of SSMs in equation 1 and 2). It is unclear how much the result transfers for standard SSMs used in practice. * In figure 5 and 6 there are no error bars or mentioning of number of seeds.

Questions

Authors may address my comments above.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

4 excellent

Contribution

1 poor

Limitations

N/A

Summary

This paper analyzes properties of state space models (SSMs) and verifies the analytic results with numerical simulations. The primary result is that SSMs have the same basic properties as classic RNNs. They perform temporal convolution on their inputs, with exponentially decaying memory and are universal function approximators.

Strengths

My overall assessment is that this paper is an important contribution to our understanding of SSMs. It raises really important questions about how SSMs are apparently able to learn long-range dependencies and points to the generality of exponentially-decaying memory. The exposition was pretty clear; although the math is challenging for a broad audience, the simulations give confidence that the results are sound.

Weaknesses

It's worth thinking more about presentation of the figures. Figure 4 was pretty helpful but Figure 3 was pretty baffling. Generally speaking my suggestion is that the captions could be more explicit and all of the terms in the expressions should appear in the figures. The paper raises a really important question that should be addressed more explicitly. If SSMs have basically the same properties as generic RNNs, why do they apparently work so much better? I appreciate that it may not be possible to answer this question with certainty, but some thoughtful discussion would really enhance the impact of this paper.

Questions

Why do SSMs work well empirically? Is it just that they more easily find long time constants in their exponential decay? I think it's really important to say something substantive about this.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

Why do SSMs work well?

Summary

The authors set out to analyze state space models, which have been gaining popularity as alternatives to transformer based systems that can better model long range dependencies and are more computationally efficient. Since such models do not utilize a non-linear activation function along the temporal access, it is important to analyze if this poses a restriction to their modeling capabilities. The paper provides a construction-based argument to prove that as long as there are layer-wise nonlinearities in the model, it becomes a universal approximator for any sequence model and thus does not require activation functions in the temporal domain to improve model capacity. Further, the authors also show that like the earlier recurrent networks (RNN, GRU, etc.), these models also suffer from an exponentially decaying memory. While the paper tackles quite a relevant topic on a family of models that are becoming widely popular, the draft itself could use some work. In particular, the presentation and clarity of the work needs to be improved, and it would be very helpful to also get some intuition behind the propositions mentioned and proved. That is, to understand what the propositions imply and why they are important.

Strengths

- The authors show that the state space models (SSMs) which have no non-linearity in their temporal axis can still universally approximate any sequential mapping up to some arbitrary error. This is proved by construction. - The authors also show that like RNNs, such models also suffer from the exponential memory decay problem. - Some preliminary numerical experiments are conducted to test the memory decay problem across different recurrent models.

Weaknesses

- The writing of the draft can be substantially improved. It would be good to include some explanations into the implications of each proposition and theorem, as to understand why exactly it is meaningful and important. The plots are currently taking a lot of extra space which can be freed to add this content. - It would be nice if the authors could provide some kind of rates of memory decay to better understand if it is better or worse in SSMs than other RNN-based methodologies.

Questions

- Line 48: Why is this relationship nonlinear? It looks like a linear function of $x$. - Line 83: Does the definition of $\hat{\rho}_t$ involve $t$ in the exponential or $(t-s)$ which is written in the draft (mistakingly, I believe?) - Could the authors provide some clarification and their reasoning on their construction of the memory function? Why is this reasonable? How is this computed? - For universality results in Section 3.3, could the authors provide some intuition behind Remark 3.4, what it implies and how is it achieved? - Is there a difference in some assumptions between Kolmogorov-theorem based construction with Volterra-Series based construction? Is the latter always superior? What cost does it come with?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The authors have adequately addressed the limitations.

Summary

The paper presents universal approximation results for state-space models (SSM) with layer-wise nonlinearity. In particular, two-layer SSMs with layer-wise nonlinearity can approximate any continuous function over a compact set. Moreover, SSMs can approximate elementwise function and convolution. The paper also characterizes the memory behavior of SSMs, showing that their memory still decays exponentially when the eigenvalues of the A matrix is bounded from 1. These theoretical results are validated with numerical experiments.

Strengths

1. Universal approximation results are helpful to theoretically characterize the expressive power of SSMs. These models are becoming more popular and these results are not yet known to the best of my knowledge. Consequently the paper could help better understand these models. 2. The presentation is quite clear. Even though this paper is theory-heavy, the main proof ideas are summarized in the main paper. This helps readers understand the technical challenges and the idea of the approach.

Weaknesses

1. Lacking connection to linear RNNs. SSMs are a special case of linear RNNs, and universal approximation results of linear RNNs are available. Moreover, one can show that almost any linear RNNs can be written as SSMs (e.g. Sec 2.2 of Gupta et al. 2022). This is simply because almost any matrix A can be diagonalized, and put into RNNs. Could the university approximation results of linear RNNs be used to prove university approximation results of SSM through this connection? 2. Lack of applications, or a direction to help with applications. While the theoretical results could help with the understanding of these models, the paper does not point out potentially applications of these theoretical results. This might limit the significance of these results. [1] Simplifying and Understanding State Space Models with Diagonal Linear RNNs. Ankit Gupta, Harsh Mehta, Jonathan Berant. 2022.

Questions

1. What are the dimensions of input/output of the SSMs? The notation is unclear on this. 2. Eq (9): What's H(x_test). Why is rho_hat defined this way? More generally I don't understand this section. 3. Why does the paper focus on approximation results for elementwise function and convolution? Elementwise function has no recurrent so it's all about the layer-wise non-linearity? 4. I don't understand Prop 3.12. Is it the same constant c0?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Not necessary.

Official Comment

Thanks for the response as well as the clarifications provided. In light of this, I am inclined to raise my score.

Thanks for the detailed response

Thanks for the detailed response. The promised changes by the authors will increase the quality of the paper. I thus increase my rating.

Official Comment

Thank you authors, for your detailed response, I have adjusted the score accordingly