Modelos y plataformas de IA
Jamba: El nuevo modelo de lenguaje híbrido de AI21 Labs que combina capas Transformer y Mamba
Añade Unite.AI a tus fuentes preferidas en Google
Los modelos de lenguaje han experimentado avances rápidos, con arquitecturas basadas en Transformers liderando el camino en el procesamiento del lenguaje natural. Sin embargo, a medida que los modelos escalan, los desafíos de manejar contextos largos, la eficiencia de la memoria y el rendimiento se han vuelto más pronunciados.
AI21 Labs ha introducido una nueva solución con Jamba, un modelo de lenguaje grande (LLM) de última generación que combina las fortalezas de las arquitecturas Transformer y Mamba en un marco híbrido. Este artículo detalla la arquitectura, el rendimiento y las posibles aplicaciones de Jamba.
Visión general de Jamba
Jamba es un modelo de lenguaje grande híbrido desarrollado por AI21 Labs, que aprovecha una combinación de capas Transformer y Mamba, integradas con un módulo de Mixture-of-Experts (MoE). Esta arquitectura permite a Jamba equilibrar el uso de memoria, el rendimiento y la eficiencia, lo que lo convierte en una herramienta poderosa para una amplia gama de tareas de NLP. El modelo está diseñado para caber en una sola GPU de 80 GB, ofreciendo un alto rendimiento y una pequeña huella de memoria, mientras mantiene un rendimiento de última generación en varios benchmarks.
La arquitectura de Jamba
La arquitectura de Jamba es la piedra angular de sus capacidades. Está construida sobre un diseño híbrido novel que intercala capas Transformer con capas Mamba, incorporando módulos MoE para mejorar la capacidad del modelo sin aumentar significativamente los requisitos computacionales.
1. Capas Transformer
La arquitectura Transformer se ha convertido en el estándar para los modelos de lenguaje modernos debido a su capacidad para manejar el procesamiento paralelo de manera eficiente y capturar dependencias de largo alcance en el texto. Sin embargo, su rendimiento a menudo se ve limitado por los altos requisitos de memoria y cálculo, particularmente al procesar contextos largos. Jamba aborda estas limitaciones integrando capas Mamba, que exploraremos a continuación.
2. Capas Mamba
Mamba es un modelo de espacio de estado (SSM) reciente diseñado para manejar relaciones de largo alcance en secuencias de manera más eficiente que los RNN tradicionales o incluso los Transformers. Las capas Mamba son particularmente efectivas para reducir la huella de memoria asociada con el almacenamiento de cachés de clave-valor (KV) en los Transformers. Al intercalar capas Mamba con capas Transformer, Jamba reduce el uso general de memoria mientras mantiene un alto rendimiento, especialmente en tareas que requieren el manejo de contextos largos.
3. Módulos de Mixture-of-Experts (MoE)
El módulo MoE en Jamba introduce un enfoque flexible para escalar la capacidad del modelo. MoE permite al modelo aumentar la cantidad de parámetros disponibles sin aumentar proporcionalmente los parámetros activos durante la inferencia. En Jamba, MoE se aplica a algunas de las capas MLP, con el mecanismo de enrutamiento que selecciona a los expertos superiores para activar para cada token. Esta activación selectiva permite a Jamba mantener una alta eficiencia mientras maneja tareas complejas.
La imagen a continuación demuestra la funcionalidad de una cabeza de inducción en un modelo híbrido de Atención-Mamba, una característica clave de Jamba. En este ejemplo, la cabeza de atención es responsable de predecir etiquetas como “Positivo” o “Negativo” en respuesta a tareas de análisis de sentimiento. Las palabras resaltadas ilustran cómo la atención del modelo se centra fuertemente en los tokens de etiqueta de los ejemplos de pocos disparos, particularmente en el momento crítico antes de predecir la etiqueta final. Este mecanismo de atención juega un papel crucial en la capacidad del modelo para realizar aprendizaje en contexto, donde el modelo debe inferir la etiqueta adecuada en función del contexto y los ejemplos de pocos disparos.
Las mejoras de rendimiento ofrecidas por la integración de Mixture-of-Experts (MoE) con la arquitectura híbrida de Atención-Mamba se destacan en la Tabla. Al utilizar MoE, Jamba aumenta su capacidad sin aumentar proporcionalmente los costos computacionales. Esto es particularmente evidente en el aumento significativo en el rendimiento en varios benchmarks como HellaSwag, WinoGrande y Preguntas Naturales (NQ). El modelo con MoE no solo logra una mayor precisión (por ejemplo, 66,0% en WinoGrande en comparación con 62,5% sin MoE), sino que también demuestra log-probabilidades mejoradas en diferentes dominios (por ejemplo, -0,534 en C4).
Características arquitectónicas clave
- Composición de capas: La arquitectura de Jamba consiste en bloques que combinan Mamba y capas Transformer en una proporción específica (por ejemplo, 1:7, lo que significa una capa Transformer por cada siete capas Mamba). Esta proporción se ajusta para un rendimiento y eficiencia óptimos.
- Integración de MoE: Las capas MoE se aplican cada pocas capas, con 16 expertos disponibles y los 2 expertos superiores activados por token. Esta configuración permite a Jamba escalar de manera efectiva mientras gestiona los compromisos entre el uso de memoria y la eficiencia computacional.
- Normalización y estabilidad: Para garantizar la estabilidad durante el entrenamiento, Jamba incorpora RMSNorm en las capas Mamba, lo que ayuda a mitigar problemas como los picos de activación grandes que pueden ocurrir a escala.
Rendimiento y evaluación de Jamba
Jamba ha sido probado rigurosamente contra una amplia gama de benchmarks, demostrando un rendimiento competitivo en todos ellos. Las siguientes secciones destacan algunos de los benchmarks clave en los que Jamba ha destacado, mostrando sus fortalezas en tareas de NLP generales y escenarios de contexto largo.
1. Benchmarks de NLP comunes
Jamba ha sido evaluado en varios benchmarks académicos, incluyendo:
- HellaSwag (10-disparos): Una tarea de razonamiento de sentido común en la que Jamba logró una puntuación de 87,1%, superando a muchos modelos competidores.
- WinoGrande (5-disparos): Otra tarea de razonamiento en la que Jamba obtuvo una puntuación de 82,5%, mostrando nuevamente su capacidad para manejar razonamiento lingüístico complejo.
- ARC-Desafío (25-disparos): Jamba demostró un rendimiento sólido con una puntuación de 64,4%, reflejando su capacidad para manejar preguntas de opción múltiple desafiantes.
En benchmarks agregados como MMLU (5-disparos), Jamba logró una puntuación de 67,4%, lo que indica su robustez en tareas diversas.
2. Evaluaciones de contexto largo
Una de las características destacadas de Jamba es su capacidad para manejar contextos extremadamente largos. El modelo admite una longitud de contexto de hasta 256K tokens, la más larga entre los modelos disponibles públicamente. Esta capacidad se probó utilizando el benchmark Needle-in-a-Haystack, donde Jamba mostró una precisión de recuperación excepcional en diferentes longitudes de contexto, incluyendo hasta 256K tokens.
3. Flujo y eficiencia
La arquitectura híbrida de Jamba mejora significativamente el flujo, particularmente con secuencias largas.
En pruebas que comparan el flujo (tokens por segundo) en diferentes modelos, Jamba superó consistentemente a sus pares, especialmente en escenarios que involucran tamaños de lote grandes y contextos largos. Por ejemplo, con un contexto de 128K tokens, Jamba logró 3 veces el flujo de Mixtral, un modelo comparable.
Usando Jamba: Python
Para desarrolladores e investigadores ansiosos por experimentar con Jamba, AI21 Labs ha proporcionado el modelo en plataformas como Hugging Face, haciéndolo accesible para una amplia gama de aplicaciones. El siguiente fragmento de código demuestra cómo cargar y generar texto utilizando Jamba:
<p>from transformers import AutoModelForCausalLM, AutoTokenizer</p> <p>model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1") tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")</p> <p>input_ids = tokenizer("En el reciente Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]</p> <p>outputs = model.generate(input_ids, max_new_tokens=216)</p> print(tokenizer.batch_decode(outputs))
Este script simple carga el modelo Jamba y el tokenizador, genera texto basado en una entrada de prompt dada y imprime la salida generada.
Ajuste fino de Jamba
Jamba está diseñado como un modelo base, lo que significa que se puede ajustar para tareas o aplicaciones específicas. El ajuste fino permite a los usuarios adaptar el modelo a dominios de nicho, mejorando el rendimiento en tareas especializadas. El siguiente ejemplo muestra cómo ajustar Jamba utilizando la biblioteca PEFT:
import torch from datasets import load_dataset from trl import SFTTrainer, SFTConfig from peft import LoraConfig from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments <p>tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1") model = AutoModelForCausalLM.from_pretrained( "ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)</p> <p>lora_config = LoraConfig(r=8, target_modules=[ "embed_tokens","x_proj", "in_proj", "out_proj", # mamba "gate_proj", "up_proj", "down_proj", # mlp "q_proj", "k_proj", "v_proj" # attention], task_type="CAUSAL_LM", bias="none")</p> <p>dataset = load_dataset("Abirate/english_quotes", split="train") training_args = SFTConfig(output_dir="./results", num_train_epochs=2, per_device_train_batch_size=4, logging_dir='./logs', logging_steps=10, learning_rate=1e-5, dataset_text_field="quote") trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=training_args, peft_config=lora_config, train_dataset=dataset, ) trainer.train()
Este fragmento de código ajusta Jamba en un conjunto de datos de citas en inglés, ajustando los parámetros del modelo para que se adapten mejor a la tarea específica de generación de texto en un dominio especializado.
Despliegue e integración
AI21 Labs ha hecho que la familia Jamba esté ampliamente accesible a través de varias plataformas y opciones de despliegue:
- Plataformas en la nube:
- Disponible en los principales proveedores de nube, incluyendo Google Cloud Vertex AI, Microsoft Azure y NVIDIA NIM (NVDA ).
- Próximamente en Amazon Bedrock, Databricks Marketplace y Snowflake Cortex.
- Marco de desarrollo de AI:
- Integración con marcos populares como LangChain y LlamaIndex (próximamente).
- AI21 Studio:
- Acceso directo a través de la propia plataforma de desarrollo de AI21.
- Hugging Face:
- Modelos disponibles para descargar y experimentar.
- Despliegue en el lugar:
- Opciones para despliegue privado en el sitio para organizaciones con necesidades específicas de seguridad o cumplimiento.
- Soluciones personalizadas:
- AI21 ofrece servicios de personalización y ajuste de modelo para clientes empresariales.
Características amigables para desarrolladores
Los modelos Jamba vienen con varias capacidades integradas que los hacen particularmente atractivos para los desarrolladores:
- Llamada a funciones: Integre fácilmente herramientas y API externas en sus flujos de trabajo de AI.
- Salida JSON estructurada: Genere estructuras de datos limpias y parseables directamente desde entradas de lenguaje natural.
- Digestión de objetos de documento: Procese y comprenda de manera eficiente estructuras de documento complejas.
- Optimizaciones RAG: Características integradas para mejorar los flujos de generación aumentada de recuperación.
Estas características, combinadas con la ventana de contexto largo y el procesamiento eficiente del modelo, hacen de Jamba una herramienta versátil para una amplia gama de escenarios de desarrollo.
Consideraciones éticas y AI responsable
Si bien las capacidades de Jamba son impresionantes, es crucial abordar su uso con una mentalidad de AI responsable. AI21 Labs enfatiza varios puntos importantes:
- Naturaleza del modelo base: Los modelos Jamba 1.5 son modelos base preentrenados sin alineación o ajuste de instrucciones específicos.
- Falta de mecanismos de moderación integrados: Los modelos no tienen mecanismos de moderación inherentes.
- Despliegue cuidadoso: Se deben implementar adaptaciones y salvaguardas adicionales antes de utilizar Jamba en entornos de producción o con usuarios finales.
- Privacidad de datos: Al utilizar despliegues en la nube, tenga en cuenta el manejo de datos y los requisitos de cumplimiento.
- Conciencia de sesgo: Al igual que todos los modelos de lenguaje grandes, Jamba puede reflejar sesgos presentes en sus datos de entrenamiento. Los usuarios deben ser conscientes de esto y implementar mitigaciones adecuadas.
Al tener en cuenta estos factores, los desarrolladores y las organizaciones pueden aprovechar las capacidades de Jamba de manera responsable y ética.
¿Un nuevo capítulo en el desarrollo de AI?
La introducción de la familia Jamba por AI21 Labs marca un hito significativo en la evolución de los modelos de lenguaje grandes. Al combinar las fortalezas de los Transformers y los modelos de espacio de estado, integrando técnicas de Mixture-of-Experts y empujando los límites de la longitud de contexto y la velocidad de procesamiento, Jamba abre nuevas posibilidades para aplicaciones de AI en diversas industrias.
A medida que la comunidad de AI continúa explorando y construyendo sobre esta arquitectura innovadora, podemos esperar ver avances adicionales en la eficiencia del modelo, la comprensión de contexto largo y el despliegue práctico de AI. La familia Jamba representa no solo un conjunto de modelos nuevos, sino un posible cambio en la forma en que abordamos el diseño y la implementación de sistemas de AI a gran escala.
