잠바: 컴팩트한 7B SSM 하이브리드 모델

Zamba: A Compact 7B SSM Hybrid Model

May 26, 2024

저자: Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, Beren Millidge

cs.AI

초록

본 기술 보고서에서는 7B 규모의 SSM-트랜스포머 하이브리드 모델인 Zamba를 소개한다. Zamba는 동일 규모의 주요 오픈 웨이트 모델들과 경쟁력 있는 성능을 달성하며, 공개적으로 이용 가능한 데이터셋에서 1조 개의 토큰으로 학습되었다. 이 규모에서 Zamba는 트랜스포머가 아닌 모델 중 최고의 성능을 보인다. Zamba는 Mamba 백본과 단일 공유 어텐션 모듈을 결합한 독창적인 아키텍처를 도입함으로써, 최소한의 파라미터 비용으로 어텐션의 이점을 얻는다. 이러한 아키텍처 덕분에 Zamba는 유사한 트랜스포머 모델보다 추론 속도가 훨씬 빠르며, 긴 시퀀스 생성 시 메모리 사용량도 상당히 적다. Zamba는 두 단계로 사전 학습되었는데, 첫 번째 단계는 기존 웹 데이터셋을 기반으로 하며, 두 번째 단계는 고품질의 지시 및 합성 데이터셋을 통해 모델을 어닐링(annealing)하고 빠른 학습률 감소를 특징으로 한다. 우리는 Zamba의 가중치와 모든 체크포인트를 첫 번째 단계와 어닐링 단계를 포함하여 오픈소스로 공개한다.

English

In this technical report, we present Zamba, a novel 7B SSM-transformer hybrid model which achieves competitive performance against leading open-weight models at a comparable scale. Zamba is trained on 1T tokens from openly available datasets and is the best non-transformer model at this scale. Zamba pioneers a unique architecture combining a Mamba backbone with a single shared attention module, thus obtaining the benefits of attention at minimal parameter cost. Due to its architecture, Zamba is significantly faster at inference than comparable transformer models and requires substantially less memory for generation of long sequences. Zamba is pretrained in two phases: the first phase is based on existing web datasets, while the second one consists of annealing the model over high-quality instruct and synthetic datasets, and is characterized by a rapid learning rate decay. We open-source the weights and all checkpoints for Zamba, through both phase 1 and annealing phases.