Language models (LMs) based on transformers have become the gold standard in natural language processing, thanks to their exceptional performance, parallel processing capabilities, and ability to retain long-term context via key-value (KV) caches. However, these benefits come at a cost—transformers require quadratic computational resources and large memory footprints, presenting significant efficiency challenges. On the other hand, state space models (SSMs), such as Mamba, boast constant computational complexity and hardware-friendly design, but they struggle with memory recall, which hampers their performance on diverse language tasks.
To address the abovementioned issues, in a new paper Hymba: A Hybrid-head Architecture for Small Language Models, an NVIDIA research team proposes Hymba, a family of small language models that employ a hybrid-head parallel architecture. By blending transformer attention mechanisms with state space models (SSMs), Hymba achieves superior efficiency and performance. Notably, it outperforms the Llama-3.2-3B model with a 1.32% higher average accuracy, while reducing cache size by 11.67× and increasing throughput by 3.49×.
Hymba is a novel LM architecture that integrates attention heads and SSM heads within the same layer, offering parallel and complementary processing of the same inputs. This hybrid-head approach allows each layer to simultaneously harness both the high-resolution recall of attention and the efficient context summarization of SSMs, increasing the model’s flexibility and expressiveness in handling various types of information flows and memory access patterns.
To further enhance the achievable performance of Hymba, the researchers introduce learnable meta tokens that are prepended to the input sequences and interact with all subsequent tokens even in sliding window attention. These meta tokens appear to act as a compressed representation of world knowledge, improving performance across both general and recall-intensive tasks.
Sharing KV cache between attention heads is common practice. Inspired by the idea that consecutive layers have a high correlation in the KV cache, they propose sharing the KV cache between layers as well. Additionally, for most layers, they choose sliding window attention to further minimize cache costs.
Comprehensive evaluations and ablation studies demonstrate that Hymba not only establishes new state-of-the-art (SOTA) benchmark performance across a wide range of representative tasks but also achieves greater efficiency compared to transformers and previous hybrid models. For instance, in commonsense reasoning tasks, Hymba1.5B can outperform Llama-3.2-3B with 1.32% higher average accuracy, while requiring 11.67× smaller cache size and being 3.49× faster.
Overall, this work demonstrates that Hymba sets new SOTA performance across a wide range of tasks, achieving superior results in both accuracy and efficiency. Additionally, it provides valuable insights into the advantages of hybrid-head architectures, offering a promising direction for future research in efficient LMs.
The paper Hymba: A Hybrid-head Architecture for Small Language Models is on .
Author: Hecate He | Editor: Chain Zhang
The post first appeared on .