Zamba : Un modèle hybride SSM compact de 7 milliards de paramètres

Zamba: A Compact 7B SSM Hybrid Model

May 26, 2024

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

cs.AI

Résumé

Dans ce rapport technique, nous présentons Zamba, un nouveau modèle hybride SSM-transformer de 7 milliards de paramètres qui atteint des performances compétitives par rapport aux modèles open-weight leaders à une échelle comparable. Zamba est entraîné sur 1 000 milliards de tokens provenant de jeux de données ouvertement disponibles et constitue le meilleur modèle non-transformer à cette échelle. Zamba innove avec une architecture unique combinant une structure principale de type Mamba avec un module d'attention partagé unique, permettant ainsi de bénéficier des avantages de l'attention à un coût en paramètres minimal. Grâce à son architecture, Zamba est significativement plus rapide en inférence que les modèles transformer comparables et nécessite beaucoup moins de mémoire pour la génération de longues séquences. Zamba est pré-entraîné en deux phases : la première phase repose sur des jeux de données web existants, tandis que la seconde consiste à affiner le modèle sur des jeux de données d'instructions de haute qualité et des données synthétiques, caractérisée par une décroissance rapide du taux d'apprentissage. Nous rendons open-source les poids et tous les points de contrôle de Zamba, couvrant à la fois la phase 1 et les phases d'affinage.

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.