[논문 리뷰] MambaByte: Token-free Selective State Space Model

Junxiong Wang, Tushaar Gangavarapu외 2명|arXiv (Cornell University)|2024. 01. 24.

Cellular Automata and Applications인용 수 7

한 줄 요약

MambaByte는 선택적 상태 공간 모델에 기반한 토큰 없이 바이트 수준의 언어 모델로, 부분 단어 트랜스포머와 경쟁력 있는 성능을 보이며 더 빠르고 확장 가능한 생성 속도를 달성합니다.

ABSTRACT

Token-free language models learn directly from raw bytes and remove the inductive bias of subword tokenization. Operating on bytes, however, results in significantly longer sequences. In this setting, standard autoregressive Transformers scale poorly as the effective memory required grows with sequence length. The recent development of the Mamba state space model (SSM) offers an appealing alternative approach with a fixed-sized memory state and efficient decoding. We propose MambaByte, a token-free adaptation of the Mamba SSM trained autoregressively on byte sequences. In terms of modeling, we show MambaByte to be competitive with, and even to outperform, state-of-the-art subword Transformers on language modeling tasks while maintaining the benefits of token-free language models, such as robustness to noise. In terms of efficiency, we develop an adaptation of speculative decoding with tokenized drafting and byte-level verification. This results in a $2.6 imes$ inference speedup to the standard MambaByte implementation, showing similar decoding efficiency as the subword Mamba. These findings establish the viability of SSMs in enabling token-free language modeling.

연구 동기 및 목표

• 토큰화 편향 및 로버스트성 문제를 완화하기 위한 토큰 없는 언어 모델링의 동기 부여.
• 패치를 거치지 않고도 긴 시퀀스에 효율적으로 확장 가능한 바이트 수준 모델 개발.
• 선택적 상태 공간 아키텍처를 활용해 선형 시간 시퀀스 처리를 가능하게.
• Transformer, MegaByte, 및 SSM 기준선과의 컴퓨트 및 매개변수 매칭 성능 평가.
• 바이트별 자가회귀 설정에서 생성 효율성 이점 입증.

제안 방법

• 바이트 시퀀스에서 직접 작동하도록 Mamba 선택적 상태 공간 모델을 적응.
• 입력에 따라 B, C, Delta를 입력 선택적으로 함수로 하여 연속 시간 SSM 다이내믹스의 이산화.
• 학습 중 선형 재발을 효율적으로 계산하기 위해 병렬 스캔 사용.
• 언어 모델링을 위한 게이트형 네트워크 아키텍처에 SSM 계층 삽입.
• 장문 텍스트 데이터셋(PG19, Stories, Books, ArXiv, Code)에서 컴퓨트 및 매개변수 매칭 실험.
• 바이트 수준 트랜스포머, SSM, MegaByte, 서브워드 기준선과 비교.

Figure 1 : Benchmarking byte-level models with a fixed parameter budget. Language modeling results on PG19 ( $8,192$ consecutive bytes), comparing the standard Transformer (Vaswani et al., 2017 ; Su et al., 2021 ) , MegaByte Transformer (Yu et al., 2023 ) , gated diagonalized S4 (Mehta et al., 2023

실험 결과

연구 질문

• RQ1 토큰 없이 바이트 수준 모델이 긴 형식의 언어 모델링에서 토큰화된 서브워드 트랜스포머와 일치하거나 능가할 수 있는가?
• RQ2 선택적 상태 공간 모델이 바이트 시퀀스에 대해 계산 및 생성 속도 이점을 제공하는가?
• RQ3 컴퓨트 및 매개변수 매칭 비교에서 MambaByte가 기존 바이트 수준 및 서브워드 아키텍처와 어떻게 수행하는가?

주요 결과

바이트 수준 모델	맥락	학습 바이트	테스트 BPB	PG19	Stories	Books	ArXiv	Code
Transformer-320M	1024	80 B	1.057	1.057	1.064	1.097	0.816	0.575
PerceiverAR-248M	8192	80 B	1.104	1.104	1.070	1.104	0.791	0.546
MegaByte-758M+262M (patch: 8)	8192	80 B	1.000	1.000	0.978	1.007	0.678	0.411
MambaByte-353M	8192	30 B*	0.930	0.930	0.908	0.966	0.663	0.396

• 같은 FLOPs-당 바이트 예산으로 여러 장문의 데이터셋에서 MegaByte를 지속적으로 능가.
• 작은 MambaByte(353M)가 BPB에서 평균적으로 바이트 수준 Transformers 및 PerceiverAR보다 우수.
• 972M 파라미터의 MambaByte가 상당한 바이트 수로 학습될 때 PG19에서 최신 서브워드 모델과 경쟁력 있거나 우수한 perplexity를 달성.
• 순환적 상태 공간 기반 설계로 MegaByte보다 빠른 생성 속도.
• 컴퓨트 매칭 설정에서 MambaByte가 더 적은 계산 및 데이터로 Transformer 유사 손실에 도달.
• 토큰 없이 바이트 모델은 긴 문맥 언어 작업에 토크나이저 의존 모델의 실현 가능한 대안이 될 수 있다.

Figure 2 : Illustration of the Mamba SSM. (a) The discrete-time input $x[k]$ , along with input-selective $\Delta[k]$ . (b) The continuous-time signal $x(t)$ . (c) Mathematically, the SSM transforms the continuous-time $x(t)$ through an $n$ -dimensional hidden state (here, $n=4$ ) using parameters $