[논문 리뷰] BlackMamba: Mixture of Experts for State-Space Models

Quentin Anthony, Yury Tokpanov외 2명|arXiv (Cornell University)|2024. 02. 01.

Simulation Techniques and Applications인용 수 10

한 줄 요약

BlackMamba는 Mamba 상태공간 블록과 Mixture-of-Experts 라우팅을 결합하여 선형 시간 생성과 축소된 학습/추론 FLOPs를 달성하며, 300B 토큰으로 학습된 오픈 소스 340M/1.5B 및 630M/2.8B 모델을 제공합니다.

ABSTRACT

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

연구 동기 및 목표

• 상 Dense 트랜스포머보다 효율성과 확장성을 개선하기 위해 상태공간 모델(SSM)과 Mixture-of-Experts(MoE)를 결합하는 것을 동기 부여한다.
• 주의 집중(attention)을 대체하여 BlackMamba를 설계·구현하고, MoE 라우팅을 MLP 구성 요소에 통합한다.
• 성능, 학습 FLOPs 및 추론 효율성 측면에서 BlackMamba를 Mamba 및 트랜스포머 베이스라인과 경험적으로 평가한다.
• 대형 BlackMamba 변형의 확장 가능한 학습을 시연하고 커뮤니티 사용을 위한 모델 가중치와 추론 코드를 공개한다.

제안 방법

• BlackMamba를 주기적으로 교대하는 주의가 없는 Mamba 블록과 라우트된 MoE 계층의 적층 구조로 소개한다.
• 입력 의존 게이팅을 사용하여 Eq. (4)의 A, B, C SSM 행렬이 x(t)에 의존하도록 한다.
• 8개의 MoE 전문가 간의 부하를 균형 있게 분배하기 위해 Sinkhorn 기반 옵티마이저를 사용한 top-1 라우팅을 적용한다.
• 1.8–2.0조 토큰 혼합 데이터셋에서 두 가지 모델 크기(340M/1.5B 및 630M/2.8B 전방/총 매개변수)를 학습한다.
• Megatron-LM 프레임워크를 사용한 bf16 정밀도로 훈련하고; MoE 블록당 8명의 전문가 활성화; 편향 비활성화; 전문가 MLP에서 SwiGLU를 사용한다.
• 추론 코드와 Apache 2.0 하에 오픈 소스 체크포인트를 제공한다.

(c) Transformer-MoE

실험 결과

연구 질문

• RQ1 SSM 블록과 MoE 라우팅을 결합하는 것이 학습 및 추론 FLOPs를 줄이면서 대규모에서 경쟁력 있는 언어 모델링 성능을 낼 수 있는가?
• RQ2 BlackMamba 모델이 생성 지연 및 장문 컨텍스트 처리에서 밀집 트랜스포머 및 순수 Mamba 베이스라인과 어떻게 비교되는가?
• RQ3 깊이에 걸친 BlackMamba에서 MoE 라우팅은 학습 시간에 따라 어떻게 동작하며 Sinkhorn 라우팅 초기화가 수렴에 미치는 영향은 무엇인가?
• RQ4 대규모 다중 데이터셋 사전학습 코퍼스에서 BlackMamba 변형을 학습할 때 데이터 및 매개변수 효율성은 어떤가?

주요 결과

• BlackMamba는 밀집 트랜스포머 및 순수 Mamba 베이스라인에 비해 학습 FLOPs를 크게 줄이고 추론이 더 빠르면서도 경쟁력 있는 평가 성능을 달성한다.
• 특히 긴 시퀀스 길이에서 BlackMamba의 추론 지연은 선형 시간 SSM 생성과 희소 MoE 라우팅의 결합으로 인해 일반적인 트랜스포머 및 비교 가능한 MoE 트랜스포머에 비해 크게 더 빠르다.
• 대부분의 층에서 Sinkhorn 기반의 라우팅으로 전문가 활용이 잘 균형을 이루지만, 최종 층은 특화/비균형 패턴이 나타나 향후 연구가 필요하다.
• 오픈 소스 BlackMamba 구성 두 가지(340M/1.5B 및 630M/2.8B)가 300B 토큰으로 학습되었고 커뮤니티 사용을 위해 가중치와 추론 코드가 공개되었다.

(d) Mamba-MoE