Jamba: 하이브리드 트랜스포머-맘바 언어 모델

Jamba: A Hybrid Transformer-Mamba Language Model

March 28, 2024

저자: Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, Yoav Shoham

cs.AI

초록

우리는 새로운 하이브리드 Transformer-Mamba 혼합 전문가(MoE) 아키텍처를 기반으로 한 새로운 대형 언어 모델인 Jamba를 소개합니다. 구체적으로, Jamba는 Transformer와 Mamba 레이어 블록을 교차 배치하여 두 모델 패밀리의 장점을 모두 누릴 수 있습니다. 일부 레이어에는 MoE를 추가하여 모델 용량을 늘리면서도 활성 매개변수 사용을 관리 가능한 수준으로 유지합니다. 이 유연한 아키텍처는 리소스와 목적에 맞는 구성을 가능하게 합니다. 우리가 구현한 특정 구성에서는 단일 80GB GPU에 적합한 강력한 모델을 얻을 수 있습니다. 대규모로 구축된 Jamba는 기존 Transformer 대비 높은 처리량과 작은 메모리 공간을 제공하면서도 표준 언어 모델 벤치마크와 장문맥 평가에서 최첨단 성능을 보여줍니다. 특히, 이 모델은 최대 256K 토큰의 문맥 길이에서도 강력한 결과를 제시합니다. 우리는 Transformer와 Mamba 레이어를 결합하는 방법, 전문가를 혼합하는 방법 등 다양한 아키텍처 결정을 연구하고, 이 중 일부가 대규모 모델링에서 중요한 역할을 한다는 것을 보여줍니다. 또한 Jamba의 훈련과 평가를 통해 밝혀진 이러한 아키텍처의 흥미로운 특성들을 설명하고, 다양한 제거 실험에서 얻은 체크포인트를 공개하여 이 새로운 아키텍처의 추가 탐구를 장려할 계획입니다. 우리는 Jamba 구현의 가중치를 허용적 라이선스 하에 공개합니다.

English

We present Jamba, a new base large language model based on a novel hybrid Transformer-Mamba mixture-of-experts (MoE) architecture. Specifically, Jamba interleaves blocks of Transformer and Mamba layers, enjoying the benefits of both model families. MoE is added in some of these layers to increase model capacity while keeping active parameter usage manageable. This flexible architecture allows resource- and objective-specific configurations. In the particular configuration we have implemented, we end up with a powerful model that fits in a single 80GB GPU. Built at large scale, Jamba provides high throughput and small memory footprint compared to vanilla Transformers, and at the same time state-of-the-art performance on standard language model benchmarks and long-context evaluations. Remarkably, the model presents strong results for up to 256K tokens context length. We study various architectural decisions, such as how to combine Transformer and Mamba layers, and how to mix experts, and show that some of them are crucial in large scale modeling. We also describe several interesting properties of these architectures which the training and evaluation of Jamba have revealed, and plan to release checkpoints from various ablation runs, to encourage further exploration of this novel architecture. We make the weights of our implementation of Jamba publicly available under a permissive license.