Hymba: Een hybride-hoofdarchitectuur voor kleine taalmodellen

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

Samenvatting

Wij stellen Hymba voor, een familie van kleine taalmodellen met een hybride-hoofd parallelle architectuur die transformer aandachtsmechanismen integreert met toestandsruimtemodellen (SSM's) voor verbeterde efficiëntie. Aandachtsmechanismen bieden een hoge-resolutie herinnering, terwijl SSM-hoofden efficiënte contextsamenvatting mogelijk maken. Daarnaast introduceren we leerzame meta-tokens die aan prompts worden toegevoegd, waarin cruciale informatie wordt opgeslagen en de last van "gedwongen aandacht" die gepaard gaat met aandachtsmechanismen verlichten. Dit model is verder geoptimaliseerd door het opnemen van cross-layer key-value (KV) delen en gedeeltelijke schuifraam-aandacht, resulterend in een compacte cache-grootte. Tijdens de ontwikkeling hebben we een gecontroleerde studie uitgevoerd waarin verschillende architecturen onder identieke omstandigheden werden vergeleken en significante voordelen van onze voorgestelde architectuur werden waargenomen. Opmerkelijk genoeg behaalt Hymba state-of-the-art resultaten voor kleine taalmodellen: Ons Hymba-1.5B-Base model overtreft alle sub-2B openbare modellen in prestaties en presteert zelfs beter dan Llama-3.2-3B met 1.32% hogere gemiddelde nauwkeurigheid, een 11.67x cache-groottevermindering en 3.49x doorvoer.

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.