Jamba (modelo de lenguaje)
Jamba es una familia de grandes modelos de lenguaje (LLM) desarrollada por la empresa de investigación israelí AI21 Labs. Jamba presenta la primera arquitectura híbrida de su tipo, que combina elementos clave de dos enfoques dominantes en el desarrollo de IA: los transformers y los modelos de espacio de estados (State Space Models, SSM), en particular, la arquitectura Mamba[1].
El objetivo principal de Jamba es resolver el compromiso fundamental de los LLM modernos: la alta calidad y el rendimiento (característicos de los transformers) frente a la eficiencia y la capacidad para procesar contextos extralargos (características de los SSM). Al combinar estos enfoques y añadir escasez mediante Mixture-of-Experts (MoE), Jamba ofrece un modelo que es a la vez potente, eficiente y capaz de trabajar con enormes volúmenes de texto en una sola consulta.
Arquitectura de Jamba en detalle
Jamba no se limita a alternar capas de transformer y Mamba. Utiliza una estructura de bloques cuidadosamente diseñada, donde cada bloque consta de ocho capas.
Estructura de un bloque de Jamba:
- Una capa de Transformer: Esta capa es responsable de la comprensión "profunda" y el razonamiento complejo. En esta capa se integra la arquitectura Mixture-of-Experts (MoE).
- Siete capas de Mamba: Estas capas siguen a la capa de transformer y son responsables del procesamiento eficiente de la secuencia y de "arrastrar" la información a través de un contexto largo[2].
Esta estructura asimétrica permite al modelo gestionar eficientemente los recursos computacionales: las operaciones de transformer, pesadas pero potentes, se ejecutan con menos frecuencia, mientras que las operaciones de Mamba, ligeras y rápidas, se ejecutan más a menudo.
Integración de Mixture-of-Experts (MoE)
En Jamba se utiliza la arquitectura MoE para mejorar aún más la eficiencia.
- El MoE se aplica únicamente a los bloques de avance de red (FFN) dentro de las capas de transformer[3]. Las capas de Mamba permanecen densas.
- El primer modelo Jamba tenía 16 expertos.
- Para cada token, la red de enrutamiento selecciona los 2 mejores expertos (gating Top-2).
Esto significa que, aunque el número total de parámetros del modelo es grande (52 mil millones), en cada paso del procesamiento de un token en la capa de transformer solo están activos 2 de los 16 expertos, lo que hace que los cálculos sean muy rápidos.
Evolución de los modelos Jamba
Jamba-v0.1 (marzo de 2024)
El primer modelo presentado dentro de esta familia tiene las siguientes características:
| Característica | Valor |
|---|---|
| Número total de parámetros | 52 mil millones |
| Parámetros activos | ~12 mil millones |
| Número de expertos (MoE) | 16 (2 activos) |
| Ventana de contexto | 256 000 tokens |
| Licencia | Apache 2.0[4] |
Gracias a su arquitectura híbrida, Jamba-1 es capaz de procesar un contexto de 256 000 tokens, lo que equivale aproximadamente a una novela de 400 páginas, y puede desplegarse en una única GPU de consumo con 80 GB de memoria[5].
Jamba-1.5 (2024)
En 2024, AI21 Labs presentó la familia de modelos actualizada Jamba 1.5, que incluye dos versiones: Jamba 1.5 Mini (12B de parámetros activos de un total de 52B) y Jamba 1.5 Large (94B de parámetros activos de un total de 398B)[6]. Estos modelos demuestran mejoras significativas en el rendimiento:
- Inferencia hasta 2.5 veces más rápida en contextos largos en comparación con la competencia.
- Soporte para nueve idiomas, incluyendo inglés, español, francés y árabe[7].
Ventajas clave y rendimiento
- Enorme ventana de contexto: 256 000 tokens, una de las ventanas más grandes entre todos los modelos disponibles (incluidos los propietarios) en el momento de su lanzamiento. Esto hace que Jamba sea ideal para tareas que requieren el análisis de documentos extensos: contratos legales, artículos científicos, bases de código completas o diálogos largos.
- Alto rendimiento y eficiencia: En las pruebas, Jamba demuestra un rendimiento comparable o superior a los principales modelos abiertos de tamaño similar, como Llama y Mixtral, al tiempo que muestra un rendimiento 3 veces mayor en contextos largos[8].
- Código abierto y accesibilidad: Jamba se distribuye bajo la licencia permisiva Apache 2.0, lo que permite su uso libre con fines comerciales y de investigación. Los pesos del modelo están disponibles en la plataforma Hugging Face.
Resultados en benchmarks
Jamba 1.5 muestra resultados competitivos en diversos benchmarks[9]:
- Jamba 1.5 Mini obtuvo 46.1 puntos en Arena Hard, lo que la convierte en el modelo público líder en su categoría[10].
- Jamba 1.5 Large obtuvo 65.4 puntos en Arena Hard, superando a Llama 3.1 70B y 405B.
Aplicación y disponibilidad
Jamba está optimizado para aplicaciones empresariales y admite capacidades como la llamada a funciones (function calling), la salida estructurada en JSON y el procesamiento de documentos. El modelo está disponible en múltiples plataformas, incluyendo:
- Hugging Face
- Google Cloud Vertex AI
- Microsoft Azure
- NVIDIA API catalog
- Amazon Bedrock[9]
- AI21 Studio
Para soportar una inferencia rentable, AI21 Labs introdujo ExpertsInt8, una nueva técnica de cuantización que permite alojar Jamba 1.5 Large en una máquina con 8 GPU de 80 GB cada una, sin pérdida de calidad al procesar un contexto de 256K tokens[11].
Literatura
- Lieber, O.; et al. (2024). Jamba: A Hybrid Transformer‑Mamba Language Model. arXiv:2403.19887.
- Lieber, O.; et al. (2024). Jamba‑1.5 Models and ExpertsInt8 Quantization. OpenReview JFPaD7lpBD.
- Gu, A.; Dao, T. (2023). Mamba: Linear‑Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.
- Gu, A.; et al. (2021). S4: Efficiently Modeling Long Sequences with Structured State Spaces. arXiv:2111.00396.
- Fedus, W.; et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961.
- Yun, L.; et al. (2024). Toward Inference‑Optimal Mixture‑of‑Expert Large Language Models. arXiv:2404.02852.
- Liu, J.; et al. (2024). A Survey on Mixture of Experts in Large Language Models. arXiv:2407.06204.
- Gupta, V.; et al. (2024). Lynx: Enabling Efficient MoE Inference through Dynamic Batch‑Aware Expert Selection. arXiv:2411.08982.
- Liu, J.; et al. (2024). A Survey on Inference Optimization Techniques for Mixture of Experts Models. arXiv:2412.14219.
- Hsieh, C.‑P.; et al. (2024). RULER: What's the Real Context Size of Your Long‑Context Language Models?. arXiv:2404.06654.
Referencias
- ↑ «Announcing Jamba: AI21's Groundbreaking SSM-Transformer Model». AI21 Labs Blog. [1]
- ↑ Lieber, O., et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. arXiv:2403.19887.
- ↑ «Jamba Documentation». Hugging Face Transformers. [2]
- ↑ «ai21labs/Jamba-v0.1». Hugging Face. [3]
- ↑ «AI21 Labs' Jamba: A New Hybrid LLM Architecture». Gradient Flow. [4]
- ↑ «Announcing the Jamba-1.5 model family». AI21 Labs Blog. [5]
- ↑ «ai21labs/AI21-Jamba-Large-1.5». Hugging Face. [6]
- ↑ «AI21 Labs rompe nuevas barreras con la ayuda de Jamba». ITinAI. [7]
- ↑ 9.0 9.1 «Ha salido Jamba 1.5: el modelo híbrido de AI21 Labs». Dzen. [8]
- ↑ «Jamba-1.5 family of models by AI21 Labs is now available in Amazon Bedrock». AWS What's New. [9]
- ↑ «ExpertsInt8: A new paradigm for efficient inference of MoE-based LLMs». OpenReview. [10]