Mamba: Selective State Space Models
An educational, from-scratch PyTorch implementation of Mamba, the selective state space model architecture, covering the selective scan mechanism, a Mamba block, and a complete language model built on it.
This repository is a clean, educational PyTorch reimplementation of Mamba, the selective state space model architecture proposed in "Mamba: Linear-Time Sequence Modeling with Selective State Spaces" by Albert Gu and Tri Dao. It implements the core selective SSM, where the state-space parameters B, C, and the timestep are made input-dependent so the model can selectively filter, remember, or forget information as it processes a sequence.
It includes a Mamba block with gated residuals, both parallel and sequential scan implementations, a full language model built from these blocks, tests, training examples, and benchmarks comparing against Transformer-based approaches.
Installation is a standard pip install -r requirements.txt, and the README includes a minimal example instantiating a MambaBlock and running it on a batch of sequences.