Hymba: Una Arquitectura de Cabeza Híbrida para Modelos de Lenguaje Pequeños

Hymba: A Hybrid-head Architecture for Small Language Models

November 20, 2024

Autores: Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabaleshwarkar, Shih-Yang Liu, Matthijs Van Keirsbilck, Min-Hung Chen, Yoshi Suhara, Yingyan Lin, Jan Kautz, Pavlo Molchanov

cs.AI

Resumen

Proponemos Hymba, una familia de modelos de lenguaje pequeños que presentan una arquitectura paralela de cabeza híbrida que integra mecanismos de atención de transformadores con modelos de espacio de estado (SSMs) para una mayor eficiencia. Las cabezas de atención proporcionan un recuerdo de alta resolución, mientras que las cabezas de SSM permiten una sumarización eficiente del contexto. Además, introducimos meta tokens aprendibles que se anteponen a las indicaciones, almacenando información crítica y aliviando la carga de "forzado a atender" asociada con los mecanismos de atención. Este modelo se optimiza aún más mediante la incorporación de compartición de clave-valor (KV) entre capas y atención de ventana deslizante parcial, lo que resulta en un tamaño de caché compacto. Durante el desarrollo, realizamos un estudio controlado comparando varias arquitecturas bajo ajustes idénticos y observamos ventajas significativas de nuestra arquitectura propuesta. Notablemente, Hymba logra resultados de vanguardia para LMs pequeños: Nuestro modelo Hymba-1.5B-Base supera a todos los modelos públicos de menos de 2B en rendimiento e incluso supera a Llama-3.2-3B con un 1.32% de precisión promedio más alta, una reducción del tamaño del caché de 11.67 veces y una mayor velocidad de procesamiento de 3.49 veces.

English

We propose Hymba, a family of small language models featuring a hybrid-head parallel architecture that integrates transformer attention mechanisms with state space models (SSMs) for enhanced efficiency. Attention heads provide high-resolution recall, while SSM heads enable efficient context summarization. Additionally, we introduce learnable meta tokens that are prepended to prompts, storing critical information and alleviating the "forced-to-attend" burden associated with attention mechanisms. This model is further optimized by incorporating cross-layer key-value (KV) sharing and partial sliding window attention, resulting in a compact cache size. During development, we conducted a controlled study comparing various architectures under identical settings and observed significant advantages of our proposed architecture. Notably, Hymba achieves state-of-the-art results for small LMs: Our Hymba-1.5B-Base model surpasses all sub-2B public models in performance and even outperforms Llama-3.2-3B with 1.32% higher average accuracy, an 11.67x cache size reduction, and 3.49x throughput.