Jamba: A Hybrid Transformer-Mamba Language Model

Abstract

Transformer + Mamba + MoE

 

[AI21 labs Blog]

[arXiv](2024/03/28 version v1)

 

이거 완전 혁명인 거 같다;;

Positional encoding도 필요 없고 속도나 메모리면에서 훨씬 효율적이면서도 transformer랑 성능 비슷함;;

 

Introduction

Transformer의 단점:

1. 높은 메모리, 컴퓨팅 비용
2. 단일 요약 상태가 없다. (항상 시퀀스로 존재한다.)

 

Mamba는 Transformer보다 성능은 뒤처지지만 더 효율적이고 장거리 모델링에 적합하다.

 

Jamba에는 MoE도 포함된다.

 

 

 

Model Architecture

Jamba block은

• number of layers (default l = 8)
• ratio of attention-mamba (a:m = 1:7)
• how often to use MoE (e = 2)
• number of experts (e = 16)
• top-k experts at each token (K = 2)

등 다양하게 조절 가능.

 

매우 적은 KV 캐시의 수

 

 

 

Reaping the Benefits

Context 길이

 

토큰 처리량

 

 

 

Evaluation

표준 benchmark

 

Naturalistic long-context evaluation

 

 

 

Ablations and Insights

비율에 따른 성능 차이는 거의 없다.

 

Mamba의 성능이 많이 뒤처지는 특정 작업에서도 attention과 비슷한 성능을 보여준다.

 

MoE

 

Mamba에 RMSNorm을 추가하면 안정적으로 대규모 훈련이 가능하다.

원본이 화질구지임

 

Mamba는 원래 positional encoding을 사용하지 않는다. 

Jamba에서는 attention layer에도 PE 필요 없음!