BlackMamba: 상태-공간 모델을 위한 전문가 혼합(MoE)
BlackMamba: Mixture of Experts for State-Space Models
February 1, 2024
저자: Quentin Anthony, Yury Tokpanov, Paolo Glorioso, Beren Millidge
cs.AI
초록
상태 공간 모델(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