Hymba: A Hybrid-head Architecture for Small Language Models

TL;DR

Hymba introduces a hybrid-head architecture combining transformers and state space models for small language models, achieving superior performance and efficiency through innovative design and optimization techniques.

Contribution

The paper presents a novel hybrid-head architecture with learnable meta tokens and optimized mechanisms, significantly improving small LMs' performance and efficiency.

Findings

Hymba-1.5B-Base outperforms sub-2B models.

Achieves 1.32% higher accuracy than Llama-3.2-3B.

Reduces cache size by 11.67x and increases throughput by 3.49x.

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.