BlackMamba: 状態空間モデルのためのエキスパート混合

要旨

状態空間モデル（SSM）は最近、大規模言語モデリングベンチマークにおいてトランスフォーマーと競争力のある性能を示しつつ、シーケンス長の関数として線形時間およびメモリ複雑性を達成しています。最近リリースされたSSMモデルであるMambaは、言語モデリングおよび長いシーケンス処理タスクにおいて印象的な性能を示しています。同時に、エキスパートの混合（MoE）モデルは、メモリフットプリントの増大を代償に、推論の計算コストとレイテンシを大幅に削減しつつ、顕著な性能を示しています。本論文では、Mamba SSMとMoEを組み合わせた新たなアーキテクチャであるBlackMambaを提案し、両者の利点を享受します。BlackMambaは、Mambaおよびトランスフォーマーのベースラインに対して競争力のある性能を示し、推論およびトレーニングのFLOPsにおいて優れていることを実証します。我々は、カスタムデータセットの300Bトークンに対して、340M/1.5Bおよび630M/2.8BのBlackMambaモデルを完全にトレーニングし、オープンソースとして公開します。BlackMambaは、SSMアーキテクチャの線形複雑性生成とMoEの安価で高速な推論の両方の利点を継承し、組み合わせることを示します。我々は、すべての重み、チェックポイント、および推論コードをオープンソースとして公開します。推論コードは以下にあります：https://github.com/Zyphra/BlackMamba

English

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

BlackMamba: 状態空間モデルのためのエキスパート混合

BlackMamba: Mixture of Experts for State-Space Models

要旨

Support