La arquitectura híbrida Hymba mejora el rendimiento de los modelos de lenguaje pequeños

Los investigadores de NVIDIA propusieron recientemente Hymba , una familia de modelos de lenguaje pequeños (SLM) que presenta una arquitectura paralela de cabezales híbridos que integra mecanismos de atención de transformadores con SSM para lograr una mayor eficiencia y un mejor rendimiento.

Los transformadores, con su arquitectura basada en la atención, se han convertido en la opción dominante para los modelos de lenguaje (LM) debido a su alto rendimiento, capacidades de paralelización y recuperación a largo plazo a través de cachés de clave-valor (KV). Sin embargo, su costo computacional cuadrático y sus altas demandas de memoria plantean desafíos de eficiencia. Por el contrario, los modelos de espacio de estado (SSM) como Mamba y Mamba-2 ofrecen una complejidad constante y una optimización eficiente del hardware, pero tienen dificultades con las tareas de recuperación de memoria, lo que afecta su rendimiento en los puntos de referencia generales.

Los investigadores de NVIDIA propusieron recientemente Hymba , una familia de modelos de lenguaje pequeños (SLM) que presenta una arquitectura paralela de cabezales híbridos que integra mecanismos de atención de transformadores con SSM para lograr una mayor eficiencia y un mejor rendimiento. En Hymba, los cabezales de atención brindan una recuperación de alta resolución, mientras que los cabezales SSM permiten un resumen de contexto eficiente.

La novedosa arquitectura de Hymba revela varias ideas:

1. Sobrecarga en atención: más del 50% del cálculo de atención se puede reemplazar con un cálculo SSM más económico.
2. Dominio de la atención local: la mayor parte de la atención global puede ser reemplazada por atención local sin sacrificar el rendimiento en tareas generales y que requieren un uso intensivo de la memoria, gracias a la información global resumida por los jefes del SSM.
3. Redundancia de caché KV: la caché clave-valor está altamente correlacionada entre cabezas y capas, por lo que se puede compartir entre cabezas (atención de consulta de grupo) y capas (uso compartido de caché KV entre capas).
4. Limitación de atención de Softmax: los mecanismos de atención están restringidos a sumar uno, lo que limita la escasez y la flexibilidad. Introducimos metatokens aprendibles que se anteponen a las indicaciones, almacenan información crítica y alivian la carga de “atención forzada” asociada con los mecanismos de atención.

Esta publicación muestra que Hymba 1.5B tiene un rendimiento favorable frente a los modelos de código abierto de última generación de tamaño similar, incluidos Llama 3.2 1B, OpenELM 1B, Phi 1.5, SmolLM2 1.7B, Danube2 1.8B y Qwen2.5 1.5B. En comparación con los modelos Transformer de tamaño similar, Hymba también logra un mayor rendimiento y requiere 10 veces menos memoria para almacenar la caché.

Hymba 1.5B se lanza a la colección Hugging Face y GitHub .

Rendimiento del Hymba 1.5B

La Figura 1 compara Hymba 1.5B con modelos sub-2B (Llama 3.2 1B, OpenELM 1B, Phi 1.5, SmolLM2 1.7B, Danube2 1.8B, Qwen2.5 1.5B) en términos de precisión promedio de la tarea, tamaño de caché (MB) en relación con la longitud de la secuencia y rendimiento (tok/seg).

En este conjunto de experimentos, las tareas incluyen MMLU, ARC-C, ARC-E, PIQA, Hellaswag, Winogrande y SQuAD-C. El rendimiento se mide en una GPU NVIDIA A100 con una longitud de secuencia de 8K y un tamaño de lote de 128 utilizando PyTorch. Para los modelos que encontraron problemas de falta de memoria (OOM) durante la medición del rendimiento, el tamaño del lote se redujo a la mitad hasta que se resolvió el OOM para medir el rendimiento máximo alcanzable sin OOM.

Diseño del modelo Hymba 

Los SSM como Mamba se introdujeron para abordar la complejidad cuadrática y los grandes problemas de caché KV en el tiempo de inferencia de los transformadores. Sin embargo, debido a su memoria de baja resolución, los SSM tienen dificultades con la recuperación de la memoria y el rendimiento. Para superar estas limitaciones, proponemos una hoja de ruta para desarrollar LM pequeños eficientes y de alto rendimiento en la Tabla 1.

Configuración	Razonamiento de sentido común (%) ↑	Recordatorio (%) ↑	Rendimiento (token/seg) ↑	Tamaño de caché (MB) ↓	Razón de diseño
Ablaciones en modelos de tamaño 300M y tokens de entrenamiento de 100B
Transformador (Llama)	44.08	39,98	721.1	414.7	Recuerdo preciso aunque ineficiente
Modelos de espacio de estados (Mamba)	42,98	19.23	4720.8	1.9	Recuperación eficiente aunque imprecisa
A. + Atención cabezas (secuenciales)	44.07	45,16	776.3	156.3	Mejorar las capacidades de recuperación
B. + Cabezales multicabezal (paralelos)	45,19	49,90	876.7	148.2	Mejor equilibrio de dos módulos
C. + Atención local/global	44,56	48,79	2399.7	41.2	Aumente la eficiencia de procesamiento y almacenamiento en caché
D. + Uso compartido de caché KV	45,16	48.04	2756.5	39.4	Eficiencia de la caché
E. + Metatokens	45,59	51,79	2695.8	40.0	Inicialización de memoria aprendida
Escalado a un tamaño de modelo de 1,5 mil millones y tokens de entrenamiento de 1,5 T
F. + Tamaño/datos	60,56	64,15	664.1	78.6	Mejorar aún más el rendimiento de las tareas
G. + Longitud de contexto extendida (2K→8K)	60,64	68,79	664.1	78.6	Mejorar las tareas de recuperación y disparo múltiple

Tabla 1. Hoja de ruta de diseño del modelo Hymba

Módulos híbridos fusionados 

Según el estudio de ablación, la fusión de las cabezas de atención y SSM en paralelo dentro de un módulo de cabeza híbrida supera el apilamiento secuencial. Hymba fusiona las cabezas de atención y SSM en paralelo dentro de un módulo de cabeza híbrida, lo que permite que ambas cabezas procesen la misma información simultáneamente. Esta arquitectura mejora el razonamiento y la precisión de la memoria.

Eficiencia y optimización de caché KV

Si bien los cabezales de atención mejoran el rendimiento de las tareas, aumentan los requisitos de caché KV y reducen el rendimiento. Para mitigar esto, Hymba optimiza el módulo de cabezal híbrido combinando la atención local y global y empleando el uso compartido de caché KV entre capas. Esto mejora el rendimiento en 3 veces y reduce el caché en casi 4 veces sin sacrificar el rendimiento. 

Metatokens

Un conjunto de 128 incrustaciones entrenadas previamente que se anteponen a las entradas y que funcionan como inicialización de caché aprendida para mejorar el enfoque en la información relevante. Estos tokens cumplen una doble función: 

• Mitigar la pérdida de atención actuando como tokens de respaldo, redistribuyendo la atención de manera efectiva
• Encapsulando el conocimiento mundial comprimido

Análisis de modelos

En esta sección se presenta una comparación de igual a igual entre diferentes arquitecturas bajo la misma configuración de entrenamiento. Luego, visualizamos los mapas de atención de SSM y Attention en diferentes modelos entrenados previamente. Por último, realizamos un análisis de importancia de la cabeza para Hymba mediante poda. Todos los análisis de esta sección ayudan a ilustrar cómo y por qué las opciones de diseño para Hymba son efectivas. 

Comparación de manzanas con manzanas 

Realizamos una comparación de igual a igual de las arquitecturas Hymba, Mamba2 pura, Mamba2 con FFN, estilo Llama3 y estilo Samba (Mamba-FFN-Attn-FFN). Todos los modelos tienen mil millones de parámetros y se entrenan desde cero para 100 mil millones de tokens de SmolLM-Corpus con exactamente la misma receta de entrenamiento. Todos los resultados se obtienen a través de lm-evaluation-harness utilizando una configuración de disparo cero en modelos Hugging Face. Hymba tiene el mejor desempeño en razonamiento de sentido común, así como en tareas de respuesta a preguntas y de memoria intensiva. 

La Tabla 2 compara varias arquitecturas de modelos en tareas de modelado de lenguaje y de razonamiento de sentido común y de memoria intensiva, y Hymba logra un sólido desempeño en todas las métricas. Hymba demuestra la menor perplejidad en tareas de lenguaje (18,62 para Wiki y 10,38 para LMB) y sólidos resultados en tareas de memoria intensiva, particularmente en SWDE (54,29) y SQuAD-C (44,71), lo que lleva al puntaje promedio más alto en esta categoría (49,50). 

Modelo	Idioma (PPL) ↓	Recordatorio intensivo (%) ↑	Razonamiento de sentido común (%) ↑
Mamba2	15,88	43.34	52,52
Mamba2 con FFN	17.43	28,92	51.14
Llama3	16.19	47.33	52,82
Samba	16.28	36.17	52,83
Himba	14.5	49,5	54,57

Tabla 2. Comparación de arquitecturas entrenadas con 100 mil millones de tokens bajo la misma configuración

En razonamiento de sentido común y respuesta a preguntas, Hymba supera a otros modelos en la mayoría de las tareas, como SIQA (31,76) y TruthfulQA (31,64), con una puntuación media de 54,57, ligeramente por encima de Llama3 y Mamba2. En general, Hymba se destaca como un modelo equilibrado, que destaca tanto en eficiencia como en rendimiento de tareas en diversas categorías.

Visualización del mapa de atención

Además, clasificamos los elementos del mapa de atención en cuatro tipos: 

1. Meta: puntuaciones de atención de todos los tokens reales a los metatokens. Esta categoría refleja la preferencia del modelo por prestar atención a los metatokens. En los mapas de atención, suelen estar ubicados en las primeras columnas (por ejemplo, 128 para Hymba) si un modelo tiene metatokens. 
2. BOS: puntuaciones de atención de todos los tokens reales hasta el token de inicio de secuencia. En el mapa de atención, suelen estar ubicados en la primera columna justo después de los metatokens. 
3. Yo: puntuaciones de atención de todos los tokens reales hacia ellos mismos. En el mapa de atención, suelen estar ubicados en la línea diagonal. 
4. Cruz: puntuaciones de atención de todos los tokens reales a otros tokens reales. En el mapa de atención, normalmente se ubican en el área fuera de la diagonal. 

El patrón de atención de Hymba es significativamente diferente al de los Transformers tradicionales. En los Transformers tradicionales, las puntuaciones de atención están más concentradas en BOS, lo que es coherente con los hallazgos en Attention Sink. Además, los Transformers tradicionales también tienen una mayor proporción de puntuaciones de atención en Self. En Hymba, los metatokens, las cabezas de atención y las cabezas SSM funcionan de manera complementaria entre sí, lo que conduce a una distribución más equilibrada de las puntuaciones de atención entre los diferentes tipos de tokens. 

En concreto, los metatokens descargan las puntuaciones de atención de BOS, lo que permite que el modelo se centre más en los tokens reales. Los cabezales SSM resumen el contexto global, que se centra más en los tokens actuales (puntuaciones de atención propia). Los cabezales de atención, por otro lado, prestan menos atención a los tokens propios y de BOS, y más atención a otros tokens (es decir, puntuaciones de atención cruzada). Esto sugiere que el diseño de cabezal híbrido de Hymba puede equilibrar eficazmente la distribución de la atención entre diferentes tipos de tokens, lo que potencialmente conduce a un mejor rendimiento.

Análisis de la importancia de las cabezas 

Analizamos la importancia relativa de la atención y de los cabezales SSM en cada capa, eliminándolos y registrando la precisión final. Nuestro análisis revela lo siguiente: 

• La importancia relativa de las cabezas de atención/SSM en la misma capa es adaptativa a la entrada y varía según las tareas, lo que sugiere que pueden cumplir diferentes funciones al manejar distintas entradas.
• El cabezal SSM en la primera capa es fundamental para el modelado del lenguaje, y eliminarlo provoca una caída sustancial en la precisión a niveles de conjeturas aleatorias.
• En general, quitar un cabezal de atención/SSM da como resultado una caída de precisión promedio de 0,24 %/1,1 % en Hellaswag, respectivamente.

Arquitectura de modelos y mejores prácticas de capacitación

Esta sección describe las decisiones arquitectónicas clave y las metodologías de capacitación para Hymba 1.5B Base y Hymba 1.5B Instruct.

Arquitectura del modelo

• Arquitectura híbrida: Mamba es excelente para realizar resúmenes y, por lo general, se centra más en el token actual, mientras que la atención es más precisa y actúa como memoria instantánea. Al combinarlos en paralelo, se fusionan estos beneficios, pero la fusión secuencial estándar no lo hace. Elegimos una relación de parámetros de 5:1 entre los cabezales SSM y de atención.
• Atención de ventana deslizante: los cabezales de atención completos se conservan en tres capas (primera, última y media), y los cabezales de atención de ventana deslizante se utilizan en el 90 % restante de las capas.
• Uso compartido de caché KV entre capas: se implementa entre cada dos capas de atención consecutivas. Se realiza además del uso compartido de caché KV de GQA entre los nodos.
• Metatokens: estos 128 tokens se pueden aprender sin supervisión, lo que ayuda a evitar problemas de colapso de entropía en modelos de lenguaje grandes (LLM) y a mitigar el fenómeno de pérdida de atención. Además, el modelo almacena conocimiento general en estos tokens. 

Mejores prácticas de capacitación 

• Preentrenamiento: optamos por un entrenamiento del modelo base en dos etapas. La etapa 1 mantuvo una tasa de aprendizaje constante y utilizó datos de corpus grandes menos filtrados. Luego, se realizó una reducción continua de la tasa de aprendizaje a 1e-5 utilizando datos de alta calidad. Este enfoque permite el entrenamiento continuo y la reanudación de la etapa 1.
• Ajuste fino de las instrucciones: el ajuste del modelo de instrucciones se realiza en tres etapas. En primer lugar, SFT-1 proporciona al modelo sólidas capacidades de razonamiento mediante el entrenamiento en código, matemáticas, llamadas a funciones, juegos de roles y otros datos específicos de la tarea. En segundo lugar, SFT-2 enseña al modelo a seguir instrucciones humanas. Por último, se aprovecha la DPO para alinear el modelo con las preferencias humanas y mejorar su seguridad.

Evaluación del desempeño y eficiencia 

Con solo 1,5 T de tokens de preentrenamiento, el modelo Hymba 1.5B tiene el mejor desempeño entre todos los LM pequeños y logra un mejor rendimiento y eficiencia de caché que todos los LM basados ​​en transformadores. 

Por ejemplo, al realizar una evaluación comparativa con la línea base más sólida, Qwen2.5, que está entrenada previamente con 13 veces más tokens, Hymba 1.5B logra una mejora de precisión promedio del 1,55 %, un rendimiento de 1,41 veces y una eficiencia de caché de 2,90 veces. En comparación con el LM pequeño más sólido entrenado con menos de 2T tokens, concretamente h2o-danube2, nuestro método logra una mejora de precisión promedio del 5,41 %, un rendimiento de 2,45 veces y una eficiencia de caché de 6,23 veces.

Modelo	# Parámetros	Fichas de tren	Token por segundo	Caché (MB)	MMLU 5 disparos	ARC-E 0 disparos	ARC-C 0 disparos	PIQA 0 disparos	borracho. 0 disparos	Hella. 0 disparos	Escuadrón -C de un disparo	Promedio
Abrir ELM-1	1.1B	1,5 toneladas	246	346	27.06	62.37	19.54	74,76	61.8	48.37	45.38	48,57
René v0.1	1.3B	1,5 toneladas	800	113	32,94	67.05	31.06	76,49	62,75	51,16	48.36	52,83
Pi 1,5	1.3B	0,15 toneladas	241	1573	42,56	76,18	44,71	76,56	72,85	48	30.09	55,85
Pequeño LM	1.7B	1 tonelada	238	1573	27.06	76,47	43.43	75,79	60,93	49,58	45,81	54,15
Cosmo	1.8B	.2T	244	1573	26.1	62,42	32,94	71,76	55.8	42.9	38,51	47.2
h20 dan-ube2	1.8B	2T	271	492	40.05	70,66	33.19	76.01	66,93	53.7	49.03	55,65
Llama 3.2 1B	1.2B	9T	535	262	32.12	65,53	31.39	74,43	60,69	47,72	40,18	50,29
Qwen 2.5	1.5B	18T	469	229	60,92	75,51	41.21	75,79	63.38	50.2	49,53	59,51
AMD OLMo	1.2B	1,3 toneladas	387	1049	26,93	65,91	31,57	74,92	61,64	47.3	33,71	48,85
Pequeño LM2	1.7B	11T	238	1573	50,29	77,78	44,71	77.09	66.38	53,55	50,5	60.04
Llama 3.2 3B	3.0B	9T	191	918	56.03	74,54	42.32	76,66	69,85	55,29	43,46	59,74
Himba	1.5B	1,5 toneladas	664	79	51,19	76,94	45.9	77.31	66,61	53,55	55,93	61.06

Tabla 2. Resultados del modelo base Hymba 1.5B

Modelos instruidos 

El modelo Hymba 1.5B Instruct logra el mayor rendimiento en promedio de todas las tareas, superando al modelo de última generación anterior, Qwen 2.5 Instruct, en alrededor de un 2 %. En concreto, Hymba 1.5B supera a todos los demás modelos en GSM8K/GPQA/BFCLv2 con una puntuación de 58,76/31,03/46,40, respectivamente. Estos resultados indican la superioridad de Hymba 1.5B, en particular en áreas que requieren capacidades de razonamiento complejas.

Modelo	# Parámetros	↑ MMLU-L	IFEval ↑	GSM8K ↑	Acreditación de calidad de la salud ↑	BFCLv2 ↑	Promedio ↑
Pequeño LM	1.7B	27,80	25.16	1.36	25,67	-*	20.00
OpenELM	1.1B	25,65	6.25	56.03	21.62	-*	27.39
Llama 3.2	1.2B	44.41	58,92	42,99	24.11	20.27	38.14
Qwen2.5	1.5B	59,73	46,78	56.03	30.13	43,85	47.30
SmolLM2	1.7B	49.11	55.06	47,68	29.24	22,83	40,78
Hymba 1.5B	1.5B	52,79	57,14	58,76	31.03	46,40	49,22

Tabla 3. Resultados del modelo Instruct de Hymba 1.5B

Conclusión

La nueva familia Hymba de LM pequeños presenta una arquitectura de cabezal híbrido que combina las capacidades de recuperación de alta resolución de los cabezales de atención con el resumen de contexto eficiente de los cabezales SSM. Para optimizar aún más el rendimiento de Hymba, se introducen metatokens aprendibles para que actúen como una memoria caché aprendida tanto para los cabezales de atención como para los SSM, lo que mejora el enfoque del modelo en la información destacada. A través de la hoja de ruta de Hymba, evaluaciones integrales y estudios de ablación, Hymba establece un nuevo rendimiento de vanguardia en una amplia gama de tareas, logrando resultados superiores tanto en precisión como en eficiencia. Además, este trabajo proporciona información valiosa sobre las ventajas de las arquitecturas de cabezal híbrido, lo que ofrece una dirección prometedora para futuras investigaciones en LM eficientes.

Obtenga más información sobre Hybma 1.5B Base y Hymba 1.5B Instruct .

Expresiones de gratitud

Este trabajo no habría sido posible sin las contribuciones de muchas personas de NVIDIA, incluidos Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabaleshwarkar, Shih-Yang Liu, Matthijs Van Keirsbilck, Min-Hung Chen, Yoshi Suhara, Nikolaus Binder, Hanah Zhang, Maksim Khadkevich, Yingyan Celine Lin, Jan Kautz, Pavlo Molchanov y Nathan Horrocks.

Recursos relacionados

• Sesión de GTC: Optimización de modelos de lenguaje de gran tamaño: un enfoque experimental para la poda y el ajuste fino de LLama2 7B
• Sesión de GTC: Aceleración de sistemas de modelos de lenguaje de extremo a extremo utilizando una arquitectura de inferencia unificada y FP8
• Sesión de GTC: Entrenamiento y ajuste de modelos de autoinstrucciones con GPU NVIDIA y marcos de software
• Contenedores NGC: Llama-3.1-Nemotron-70B-Instrucciones
• Contenedores NGC: Llama-3-Swallow-70B-Instruct-v0.1
• Kit de desarrollo de software: NeMo Megatron

NVIDIA Blog. X. D., Y. F., S. D. y P. M. Traducido al español