BlackMamba：状态空间模型的专家混合
计算与语言
2024-02-06 v1 人工智能
分布式、并行与集群计算
机器学习
摘要
状态空间模型（SSMs）最近在大规模语言建模基准上展示了与Transformer相竞争的性能，同时实现了与序列长度成线性关系的时间和内存复杂度。Mamba是一个最近发布的SSM模型，在语言建模和长序列处理任务中均表现出令人印象深刻的性能。同时，专家混合（MoE）模型在显著降低推理的计算和延迟成本（以更大的内存占用为代价）的同时，展示了卓越的性能。在本文中，我们提出了BlackMamba，一种结合Mamba SSM和MoE的新颖架构，以获得两者的优势。我们证明BlackMamba在与Mamba和Transformer基线竞争时表现优异，并在推理和训练FLOPs方面胜出。我们在自定义数据集的300B token上完整训练并开源了340M/1.5B和630M/2.8B的BlackMamba模型。我们展示了BlackMamba继承并结合了SSM和MoE架构的优势，将SSM的线性复杂度生成与MoE的廉价快速推理相结合。我们开源了所有权重、检查点和推理代码。推理代码见：https://github.com/Zyphra/BlackMamba
引用
@article{arxiv.2402.01771,
title = {BlackMamba: Mixture of Experts for State-Space Models},
author = {Quentin Anthony and Yury Tokpanov and Paolo Glorioso and Beren Millidge},
journal= {arXiv preprint arXiv:2402.01771},
year = {2024}
}