MatMamba将嵌套模型结构融入Mamba2，实现大小模型一体训练与按需部署。
MatMamba: A Matryoshka State Space Model
- 在Mamba2基础上引入嵌套维度，支持多尺度子模型联合训练
- 35M到1.4B参数模型均能从单一大模型中提取，性能不降反升
- 适合资源受限场景的弹性推理，兼顾效率与精度
状态空间模型（SSMs）如Mamba2是Transformer的有力替代方案，尤其在长序列处理上理论训练和推理速度更快。受马特里什卡表征学习启发，本文提出MatMamba：一种结合马特里什卡思想与Mamba2的新型状态空间模型。通过修改模块结构引入嵌套维度，实现多尺度子模型的联合训练与自适应推理。训练一个大型MatMamba模型后，可免费获得多个更小的嵌套模型，其性能优于同等规模从头训练的基准模型。我们在语言与图像任务上分别验证了35M至1.4B参数范围内的模型表现，结果表明，MatMamba在ImageNet与FineWeb上的扩展能力可媲美Transformer，同时具备更优的推理效率。这使其成为基于可用算力弹性部署大规模模型的实用选择。代码与模型已在GitHub开源。
原文摘要 · Abstract (English)
State Space Models (SSMs) like Mamba2 are a promising alternative to Transformers, with faster theoretical training and inference times -- especially for long context lengths. Recent work on Matryoshka Representation Learning -- and its application to Transformer backbones in works like MatFormer -- showed how to introduce nested granularities of smaller submodels in one universal elastic model. 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. This makes MatMamba a practically viable option for deploying large-scale models in an elastic way based on the available inference compute. Code and models are open sourced at \url{https://github.com/ScaledFoundations/MatMamba}
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。