Zamba: Um Modelo Híbrido SSM Compacto de 7B

Zamba: A Compact 7B SSM Hybrid Model

May 26, 2024

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

cs.AI

Resumo

Neste relatório técnico, apresentamos o Zamba, um novo modelo híbrido SSM-transformer de 7B que alcança desempenho competitivo em relação aos principais modelos de peso aberto em uma escala comparável. O Zamba foi treinado em 1 trilhão de tokens provenientes de conjuntos de dados abertamente disponíveis e é o melhor modelo não-transformer nessa escala. O Zamba introduz uma arquitetura única que combina uma estrutura principal Mamba com um único módulo de atenção compartilhado, obtendo assim os benefícios da atenção com um custo mínimo de parâmetros. Devido à sua arquitetura, o Zamba é significativamente mais rápido na inferência do que modelos transformer comparáveis e requer substancialmente menos memória para a geração de sequências longas. O Zamba é pré-treinado em duas fases: a primeira fase é baseada em conjuntos de dados web existentes, enquanto a segunda consiste em recozer o modelo sobre conjuntos de dados de instruções e sintéticos de alta qualidade, sendo caracterizada por uma rápida redução da taxa de aprendizado. Disponibilizamos publicamente os pesos e todos os checkpoints do Zamba, tanto da fase 1 quanto das fases de recozimento.

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.