Hymba : Une architecture à tête hybride pour de petits modèles de langage

Hymba: A Hybrid-head Architecture for Small Language Models

November 20, 2024

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

Résumé

Nous proposons Hymba, une famille de petits modèles de langage présentant une architecture parallèle à tête hybride qui intègre des mécanismes d'attention de transformateur avec des modèles d'espace d'état (SSM) pour une efficacité accrue. Les têtes d'attention fournissent un rappel haute résolution, tandis que les têtes SSM permettent une synthèse efficace du contexte. De plus, nous introduisons des jetons méta apprenables qui sont préfixés aux invites, stockant des informations critiques et soulageant le fardeau de "devoir assister" associé aux mécanismes d'attention. Ce modèle est en outre optimisé en incorporant le partage de clés-valeurs (KV) entre les couches et une attention à fenêtre coulissante partielle, ce qui se traduit par une taille de cache compacte. Au cours du développement, nous avons mené une étude contrôlée comparant diverses architectures dans des paramètres identiques et avons observé des avantages significatifs de notre architecture proposée. Notamment, Hymba atteint des résultats de pointe pour les petits modèles de langage : Notre modèle Hymba-1.5B-Base surpasse tous les modèles publics de moins de 2B en performances et surpasse même Llama-3.2-3B avec une précision moyenne supérieure de 1,32 %, une réduction de la taille du cache de 11,67 fois et un débit supérieur de 3,49 fois.

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.