S4++: Elevating Long Sequence Modeling with State Memory Reply

• The conclusions from Proposition 1 seemed to be based on the assumption that there is asymptotic instability of the states $x(t)$ due to the max eigenvalue $|\lambda_{max}|>1$. However, prior S4-related work (https://arxiv.org/pdf/2202.09729.pdf, https://arxiv.org/pdf/2206.11893.pdf) suggest enforcing the left-half plane condition to ensure stability and most S4-related implementations have this as an option. I would expect this baseline to be cited, discussed and compared to in the main paper.

  • In the appendix of this submission, there is a remark suggesting that enforcing $|\lambda_{max}|<1$ limits the long-term memory capacity of the model. But, I would then expect to see empirical evidence suggesting the proposed method has better long-term memory capacity than vanilla S4. As discussed below, the empirical results do not support this claim.
  • This remark would also seem to overstate the issue, since the role of the learnable timescale parameter in S4 is to extend out the exponential decay to help capture very long-range dependencies.

• The proposed SMR (multiplying a convolution of the input sequence with the input) seems to essentially be the short convolution/shift matrix and multiplicative gating proposed by H3 and Hyena (https://arxiv.org/abs/2302.10866). This connection should be discussed. While this makes the actual proposed method a less novel contribution, discussing this connection in more detail could add more theoretical insight into why this is a good thing to do in these other methods.

• The second contribution of the paper is the proposal to combine SSMs with Attention. However, this is not new and has been proposed before in MEGA (cited in the results but not discussed) and Block-state Transformer (https://arxiv.org/pdf/2306.09539.pdf). The architecture in Figure 5 looks very similar to the proposed architecture in Mega. The similarities and differences between S++ and this work should be discussed thoroughly and ablated.

• The motivation for the need of attention in this paper seems weak. Unless I am misunderstanding Section 3.2, the motivation seems to be the need for bidirectional modeling abilities while it is claimed S4 is simply a unidirectional model. However, almost all prior S4-related papers have proposed and used bidirectional SSMs. This is not to say it is necessarily incorrect that combining attention and SSMs is helpful, but the stated motivation and reasoning for why it is required seems incorrect.

• The empirical results are weak and do not help to support the claim that the proposed modifications lead to performance improvements. The LRA results exclude the most challenging Path-X task which many methods cannot improve over random guessing (this was one of the big contributions of S4). It is suspicious that this task was left out. Showing better performance on this task would also help to support the claim that the proposed method helps to improve the memory capacity over restricting the eigenvalues as mentioned above related to the relevant remark in the appendix.

  • Further, the reported LRA results seem to make no contact with prior reported results on this task. E.g. in this paper (https://arxiv.org/pdf/2206.11893.pdf) S4 has an average score much higher than the average score reported for all the methods in this paper. Similarly, in the Mega paper (a direct comparison to another SSM+attention based method).

• The presentation is poor with typos throughout. Further, while the Figures could be really helpful for the story, the resolutions of Figures 1 and 2 are too small, making it difficult to read and understand. Further, all Figures would benefit from having captions that explain the different components of the figures.

Ultimately I do not believe this paper is close to publication-ready. There may well be a reasonable contribution in this paper as it stands, but I have too many concerns about the pitching and evaluation of the idea to be confident in recommending acceptance. The paper itself is also littered with errors, which to me suggest that the paper was rushed and would greatly benefit from a round of revisions.

I will outline some more of my major concerns below (in no particular order).

Major Concerns

W.1.: Figures quoted: The figures quoted in your results table are incorrect (e.g. S4 should achieve an average of 86.09 as opposed to the 66.48 you quote). The performance gains over bidirectional S4, Liquid S4, S5 etc are marginal at best against the incorrect numbers and are very below the “correct” numbers. Please can the authors clarify?

W.2.: Is NSS a problem?: If it is believed that there is jitter in the sampling timesteps of the system, then shouldn’t that be presented in the training data as well? I am not convinced that the NSS problem is as significant as the authors make out. Does adding some Gaussian noise (for continuous input spaces) to the input remedy the problem? There is also no discussion of S5, which can handle irregularly sampled data (I understand you’re referring to “binned” data, but this is still a necessary point for discussion, even if it is qualitative). Ultimately, I think it is not experimentally shown that NSS is actually a problem that appears outside of a single, very tightly controlled experiment.

W.3.: Unidirectionality of models: You say that S4 is unidirectional, but I don’t believe that is true. S4/S4D/S5 all run an SSM in each direction, and hence gets bi-directional information. You recover this through the attention mechanism in S4++, but I don't think comparing to unidirectional S4 models (either qualitative or quantitative) is fair.

W.4.: Difference to MEGA: My understanding is that MEGA (somewhat reductionist-ly) added an attention layer to an S4 model. Can the authors clarify how their S4++ model is different from MEGA? And why is there no discussion in the main text of MEGA? Similarly, this introduction of attention re-introduces the quadratic complexity, meaning that the original motivation of S4 has been lost. This is not mentioned in the paper.

W.5.: S4+ Motivation and grounding: I just about follow through Theorem 1 that you are effectively looking to control the scale of the inputs to control the stability of the system. (15) is similar to the result, sure, but how do you ensure that the constraints/requirements are met? It would be nice to see somehow empirical validation that the extra complexity is acting as intended (could you compute the value of the message?), and that you’re not just adding extra computational flexibility in the system that is leading to marginally better performance.

Minor weaknesses:

• You describe your method as both State Memory Replay and State Memory Reply.
• ETC theory is not discussed at all.
• Proposition 1 should include a proof sketch in the main paper.
• There is a litany of content/typographical errors, e.g.:
  • Only proper nouns should be capitalized.
  • Figures should be rendered as PDFs.
  • Is it EMC or ETC theory?
  • Fonts in figures should be about the same size as fonts in the body text.
  • Use the backtick (`) when opening quotations (`Pythagorean theorem’ vs ‘Pythagorean theorem’).
  • Use \citet and \citep.
  • Citation style in the bibliography is inconsistent.

1. The presentation has many small errors that make it more difficult to read.
2. The theoretical framework is not explained in a self-contained manner. For example, the ETC theory and Lyapunov functions involve two new parameters P and T, but it is not explained how these are related to the underlying SSM models which do not have these parameters.
3. Overall in the NSS part of the paper (the first main contribution), the connection between the theory (Theorem 1) and the proposed method (equation (15)) seems quite tenuous at worst, and insufficiently explained at best. While the proposed convolutional smoothing technique makes sense intuitively, and has been used by prior work empirically, it is not obviously related to the proposed theoretical motivation. I will raise my score if the connection from Theorem 1 to equation (15) is unpacked in more detail.
4. The second contribution of adding cross-attention is more marginal and has been done before in several related works.

• The paper overlooks various SSM-based approaches that get around keeping $\Delta t$ fixed (for example, S5, which avoids the issue by computing the input-output via a parallel scan).
• The presentation can be significantly improved, some sections (3.2) could be greatly shortened.
• The weakest part of the paper is by far section 3. I elaborate on my concerns in my questions below