/haiku_SSMs

Structured State-Space Models implemented in Haiku

Primary LanguagePython

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