Summary
This work analyses the modeling capability of different SSMs (S4-S6 and others) using Rough Path Theory by viewing SSMs as (input-) controlled differential equations (CDE). To this end, the authors show that SSMs with with dense transition matrices (A) are able to approximate arbitrarily close any continuous functions of the inputs. In contrast, SSMs with diagonal transition matrices lack this property. However, when stacking/chaining multiple diagonal CDEs (as is the case for the deep architectures), full expressivity is gained again as the depth approaches infinity. The work is mostly theoretical, and, hence, there is only a single toy experiment showing RMSE for Mamba and S5 models with depth 1 and 2, showing that the more limited models are not able to solve the toy task, as they cannot approximate the corresponding function.
Strengths
- The paper is generally well written and the authors managed to make the very involved theory and proofs more accessible by walking through the interpretation and takeaways on a more high level.
- The theoretical results are important contributions for categorizing the expressivity of existing SSM models and may potentially inform about future architecture design of SSMs and related models.
Weaknesses
Despite the well-described theoretical contributions and insights, it is still not very clear to me what to take away practically.
In particular, it is not clear how useful the result regarding chaining diagonal CDEs is in practice: while it is important to have a proof that an infinite number of chains of diagonal CDE with mixing can recover the full expressiveness of the dense CDEs, the infinite case is not practically important. Instead, it would be important for practical considerations (computational efficiency vs model expressiveness) to know whether one should use dense CDEs or diagonal CDEs with mixing to achieve a certain level of expressivity. However, as far as I understand, there is no result that quantifies e.g. how many diagonal plus mixing layers have a similar expressiveness as a certain number of dense CDEs. The authors speculate in the summary that architectures with non-diagonality might improve performance, but I do not see how this actually follows from this work, since there is no comparison between diagonal+mixing and dense in the regime of finite number of parameters.
The experimental section is very weak. While this work is mostly theoretical, this work would nevertheless be improved a lot by a greater experimental section that validates the claims. Showing error curves with a single random seed is not that interesting. I would be better to have a table of final performance that includes error bars. Furthermore, it would be great to explore a few different tasks and several ablations such as number of layers. This would make it at least empirically possible to compare dense and diagonal + mixing and random dense + learned C (or MLP) in the finite parameter setting, where its not clear what to take away from the theorems.
Questions
- what is NCDE? This was not introduced.
- does random linear CDE perform best for Area computation and is also a strong baseline for Volume? I understand that with high probability we can just take a random initialized A matrix and train only the C matrix (Theorem 4.2) to achieve full expressivity. However, that model should still not be better than a dense CDE where everything is trained such as in case of S5. Or did you choose S5 with diagonal transition?
Limitations
Just a single limitation is mentioned in the conclusion. I am sure, the authors can think of many more limitations, such as my main concern raised above regarding the comparison of different mixed diagonal + mixing vs dense architectures. Or the limited experimental evidence.