Jamba: Un Modelo de Lenguaje Híbrido Transformer-Mamba

Jamba: A Hybrid Transformer-Mamba Language Model

March 28, 2024

Autores: 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

Resumen

Presentamos Jamba, un nuevo modelo base de lenguaje grande basado en una novedosa arquitectura híbrida Transformer-Mamba de mezcla de expertos (MoE). Específicamente, Jamba intercala bloques de capas Transformer y Mamba, aprovechando los beneficios de ambas familias de modelos. Se añade MoE en algunas de estas capas para aumentar la capacidad del modelo mientras se mantiene un uso manejable de parámetros activos. Esta arquitectura flexible permite configuraciones específicas según los recursos y objetivos. En la configuración particular que hemos implementado, obtenemos un modelo potente que cabe en una única GPU de 80GB. Construido a gran escala, Jamba ofrece un alto rendimiento y una pequeña huella de memoria en comparación con los Transformers convencionales, al mismo tiempo que logra un rendimiento de vanguardia en evaluaciones estándar de modelos de lenguaje y en contextos largos. Notablemente, el modelo presenta resultados sólidos para longitudes de contexto de hasta 256K tokens. Estudiamos varias decisiones arquitectónicas, como cómo combinar capas Transformer y Mamba, y cómo mezclar expertos, y demostramos que algunas de ellas son cruciales en el modelado a gran escala. También describimos varias propiedades interesantes de estas arquitecturas que han sido reveladas durante el entrenamiento y evaluación de Jamba, y planeamos publicar puntos de control de varias ejecuciones de ablación, para fomentar una mayor exploración de esta novedosa arquitectura. Hacemos públicos los pesos de nuestra implementación de Jamba bajo una licencia permisiva.

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.