arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2402.01771cs.CLcs.AIcs.DCcs.LG
收藏

BlackMamba:状态空间模型的专家混合

BlackMamba: Mixture of Experts for State-Space Models

Zyphra(齐弗拉公司)

Quentin Anthony, Yury Tokpanov, Paolo Glorioso, Beren Millidge

AI总结:

BlackMamba结合Mamba状态空间模型与专家混合,实现线性复杂度生成和高效推理,在语言建模中性能媲美基线,并开源多种规模模型。

摘要:

状态空间模型(SSM)近期在大规模语言建模基准测试中展现出与Transformer竞争的性能,同时实现了随序列长度线性增长的时间和内存复杂度。最近发布的SSM模型Mamba在语言建模和长序列处理任务中均表现出色。与此同时,专家混合(MoE)模型在显著降低推理计算和延迟成本的同时,展现出卓越性能,但代价是更大的内存占用。在本文中,我们提出了BlackMamba,一种新颖的架构,将Mamba SSM与MoE相结合,以兼得两者优势。我们证明BlackMamba在与Mamba和Transformer基线的竞争中表现优异,并在推理和训练FLOPs方面超越它们。我们在自定义数据集的300B tokens上完整训练并开源了340M/1.5B和630M/2.8B的BlackMamba模型。我们展示了BlackMamba继承了并融合了SSM和MoE架构的优势,结合了SSM的线性复杂度生成和MoE的廉价快速推理。我们开源了所有权重、检查点和推理代码。推理代码见:https://github.com/Zyphra/BlackMamba

英文摘要:

State-space models (SSMs) have recently demonstrated competitive performance to transformers at large-scale language modeling benchmarks while achieving linear time and memory complexity as a function of sequence length. Mamba, a recently released SSM model, shows impressive performance in both language modeling and long sequence processing tasks. Simultaneously, mixture-of-expert (MoE) models have shown remarkable performance while significantly reducing the compute and latency costs of inference at the expense of a larger memory footprint. In this paper, we present BlackMamba, a novel architecture that combines the Mamba SSM with MoE to obtain the benefits of both. We demonstrate that BlackMamba performs competitively against both Mamba and transformer baselines, and outperforms in inference and training FLOPs. We fully train and open-source 340M/1.5B and 630M/2.8B BlackMamba models on 300B tokens of a custom dataset. We show that BlackMamba inherits and combines both of the benefits of SSM and MoE architectures, combining linear-complexity generation from SSM with cheap and fast inference from MoE. We release all weights, checkpoints, and inference code open-source. Inference code at: https://github.com/Zyphra/BlackMamba

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑