15일 전

Mamba: 선택적 상태 공간을 활용한 선형 시간 시퀀스 모델링

Albert Gu, Tri Dao
Mamba: 선택적 상태 공간을 활용한 선형 시간 시퀀스 모델링
초록

기초 모델(Foundation models)은 현재 딥러닝 분야에서 가장 흥미로운 응용 프로그램들을 이끄는 핵심 기술로, 거의 전부가 트랜스포머(Transformer) 아키텍처와 그 핵심인 어텐션(attention) 모듈에 기반하고 있다. 긴 시퀀스 처리 시 트랜스포머의 계산 효율성 부족 문제를 해결하기 위해 선형 어텐션(linear attention), 게이트드 컨볼루션(gated convolution), 순환 모델, 구조화된 상태 공간 모델(Structured State Space Models, SSMs) 등 여러 가지 하위 2차 시간 복잡도(subquadratic-time) 아키텍처가 개발되었지만, 언어와 같은 중요한 모달리티에서는 여전히 어텐션 모델의 성능을 따라가지 못하고 있다. 우리는 이러한 모델들의 주요 약점이 콘텐츠 기반 추론(content-based reasoning)을 수행할 수 없다는 점임을 밝혀냈다. 이를 해결하기 위해 몇 가지 개선을 제안한다. 첫째, SSM 파라미터를 입력의 주소(입력 위치)에 대한 함수로 정의함으로써 이산적 모달리티에 대한 약점을 보완하고, 현재 토큰에 따라 시퀀스 길이 방향으로 정보를 선택적으로 전파하거나 잊어버리는 능력을 부여한다. 둘째, 이 변경은 효율적인 컨볼루션 사용을 어렵게 하지만, 순환 모드에서 하드웨어 인식형 병렬 알고리즘을 설계하여 이를 극복한다. 이러한 선택적 SSM을 어텐션 또는 심지어 MLP 블록 없이도 간단화된 엔드투엔드 신경망 아키텍처에 통합함으로써, Mamba라는 새로운 모델을 제안한다. Mamba는 빠른 추론 속도(트랜스포머 대비 5배 높은 처리량), 시퀀스 길이에 대한 선형 스케일링 성능을 가지며, 실제 데이터에서 수백만 길이에 이르는 시퀀스에서도 성능이 향상된다. 일반적인 시퀀스 모델 기반 아키텍처로서 Mamba는 언어, 음성, 유전체학 등 다양한 모달리티에서 최신 기준(SOTA) 성능을 달성한다. 언어 모델링 측면에서, Mamba-3B 모델은 크기가 동일한 트랜스포머보다 우수하며, 크기가 두 배인 트랜스포머와도 동일한 성능을 보이며, 사전 훈련 및 하류 작업 평가 모두에서 뛰어난 성능을 입증하였다.

Mamba: 선택적 상태 공간을 활용한 선형 시간 시퀀스 모델링 | 최신 연구 논문 | HyperAI초신경