This repository contains simple implementations of the various SSM-based models discussed in https://tinkerd.net/blog/machine-learning/state-space-models/

These implementations are optimized for readability rather than performance and are thus not meant for production use.

Minimal implementations of the following model architectures are included:

Linear State-Space Layers (LSSL)

Combining recurrent, convolutional, and continuous-time models with linear state space layers

Gu, A., Johnson, I., Goel, K., Saab, K., Dao, T., Rudra, A. and Ré, C., 2021.

Paper: https://arxiv.org/abs/2110.13985

S4D

On the parameterization and initialization of diagonal state space models

Gu, A., Goel, K., Gupta, A. and Ré, C., 2022.

Paper: https://arxiv.org/abs/2206.11893

Hungry Hungry Hippos (H3)

Hungry hungry hippos: Towards language modeling with state space models.

Fu, D.Y., Dao, T., Saab, K.K., Thomas, A.W., Rudra, A. and Ré, C., 2022.

Paper: https://arxiv.org/abs/2212.14052

Mamba

Mamba: Linear-time sequence modeling with selective state spaces.

Gu, A. and Dao, T., 2023.

Paper: https://arxiv.org/abs/2312.00752