Structured State Space Models in Haiku

Structured State-Space Models implemented in Haiku. Currently contains S4/S4D and S5 models + dataloaders for sequence prediction and classification on a few image datasets.

References

Repositories

1. https://github.com/HazyResearch/state-spaces - Original S4 implementations in PyTorch

2. https://github.com/srush/annotated-s4 - Flax implementation of S4 models (S4, S4D, DSS)

3. https://github.com/lindermanlab/S5 - Flax implementation of S5 (original paper version)

Papers

1. Efficiently Modeling Long Sequences with Structured State Spaces: https://arxiv.org/abs/2111.00396

2. On the Parameterization and Initialization of Diagonal State Space Models: https://arxiv.org/abs/2206.11893

3. Simplified State Space Layers for Sequence Modeling: https://arxiv.org/pdf/2208.04933.pdf