MatMamba: A Matryoshka State Space Model Publication date: 9 Oct 2024 Topic: Representation Learning Paper: https://arxiv.org/pdf/2410.06718v1.pdf GitHub: https://github.com/scaledfoundations/matmamba Description: In this work, we present MatMamba: a state space model which combines Matryoshka-style learning with Mamba2, by modifying the block to contain nested dimensions to enable joint training and adaptive inference. MatMamba allows for efficient and adaptive deployment across various model sizes. We train a single large MatMamba model and are able to get a number of smaller nested models for free -- while maintaining or improving upon the performance of a baseline smaller model trained from scratch. We train language and image models at a variety of parameter sizes from 35M to 1.4B. Our results on ImageNet and FineWeb show that MatMamba models scale comparably to Transformers, while having more efficient inference characteristics.