Hymba: A Hybrid-head Architecture for Small Language Models
Authors: name: Xin Dong; name: Yonggan Fu; name: Shizhe Diao; name: Wonmin Byeon; name: Zijia Chen; name: Ameya Sunil Mahabaleshwarkar; name: Shih-Yang Liu; name: Matthijs Van Keirsbilck; name: Min-Hung Chen; name: Yoshi Suhara; name: Yingyan Lin; name: Jan Kautz; name: Pavlo Molchanov
Published: 2024-11-20 19:51:25

Abstract: 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.