Zamba: Un 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

Resumen

En este informe técnico, presentamos Zamba, un novedoso modelo híbrido SSM-transformer de 7B que logra un rendimiento competitivo frente a los principales modelos de peso abierto a una escala comparable. Zamba está entrenado con 1T de tokens provenientes de conjuntos de datos disponibles públicamente y es el mejor modelo no basado en transformers a esta escala. Zamba introduce una arquitectura única que combina una estructura principal Mamba con un único módulo de atención compartido, obteniendo así los beneficios de la atención con un costo mínimo de parámetros. Gracias a su arquitectura, Zamba es significativamente más rápido en inferencia que los modelos transformers comparables y requiere sustancialmente menos memoria para la generación de secuencias largas. Zamba se preentrena en dos fases: la primera fase se basa en conjuntos de datos web existentes, mientras que la segunda consiste en un proceso de recocido del modelo sobre conjuntos de datos de instrucciones de alta calidad y datos sintéticos, caracterizado por una rápida disminución de la tasa de aprendizaje. Liberamos los pesos y todos los puntos de control de Zamba, tanto para la fase 1 como para las fases de recocido.

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.