Hymba: Un'architettura ibrida per modelli linguistici di piccole dimensioni

Hymba: A Hybrid-head Architecture for Small Language Models

November 20, 2024

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

Abstract

Proponiamo Hymba, una famiglia di modelli linguistici di piccole dimensioni che presentano un'architettura parallela ibrida che integra meccanismi di attenzione del trasformatore con modelli dello spazio di stato (SSM) per una maggiore efficienza. Le testate di attenzione forniscono un richiamo ad alta risoluzione, mentre le testate SSM consentono una sintesi efficiente del contesto. Inoltre, introduciamo meta-token apprendibili che vengono anteposti alle istruzioni, memorizzando informazioni critiche e alleviando il peso del "forzato ad attenzione" associato ai meccanismi di attenzione. Questo modello è ulteriormente ottimizzato incorporando la condivisione di chiavi-valore (KV) tra strati e l'attenzione a finestra scorrevole parziale, risultando in una dimensione della cache compatta. Durante lo sviluppo, abbiamo condotto uno studio controllato confrontando varie architetture in impostazioni identiche e osservato significativi vantaggi della nostra architettura proposta. In particolare, Hymba raggiunge risultati all'avanguardia per i piccoli modelli linguistici: il nostro modello Hymba-1.5B-Base supera tutti i modelli pubblici inferiori a 2B in termini di prestazioni e batte persino Llama-3.2-3B con una precisione media superiore del 1,32%, una riduzione della dimensione della cache di 11,67 volte e un throughput di 3,49 volte superiore.

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.