NVIDIA researchers have unveiled Hymba 1.5B, an open-source language model that combines transformer and state-space model (SSM) architectures to achieve unprecedented efficiency and performance. Designed with NVIDIA’s optimized training pipeline, Hymba addresses the computational and memory limitations of traditional transformers while enhancing the recall capabilities of SSMs.
Traditional transformer-based language models excel at long-term recall and parallelization but face substantial challenges with their quadratic computational complexity and large memory demands. On the other hand, SSMs like Mamba and Mamba-2 offer constant complexity and hardware optimization but underperform in memory recall tasks. Hymba resolves these trade-offs by combining the strengths of both architectures.
The hybrid-head module in Hymba fuses attention heads for high-resolution recall with SSM heads for efficient context summarization, enabling both components to work in parallel rather than sequentially. This design reduces computation and memory requirements without sacrificing performance.
The Hymba 1.5B architecture focuses on improving efficiency while maintaining accuracy by introducing several innovative mechanisms:
- Attention Overhead Reduction: over 50% of attention computation is replaced by SSM processing, maintaining task accuracy while reducing costs.
- Local Attention Dominance: global attention is minimized, as local attention paired with SSMs sufficiently summarizes global information.
- KV Cache Optimization: Hymba introduces cross-layer KV cache sharing, which reduces cache redundancy across layers and significantly reduces memory usage—up to 10 times less than comparable transformer models.
- Meta-Tokens: a set of 128 learnable embeddings prepended to prompts acts as memory initializers.
Source: NVIDIA Blog
The role of learnable meta-tokens has been discussed. Daniel Svonava, a machine learning specialist at Superlinked, posed a question:
Can you explain how learnable meta tokens improve the focus of the attention mechanism compared to traditional methods?
Marek Barák, a data scientist, explained:
Attention has this issue where it puts too much focus on the first token in the sentence. This has very little semantic reason since the first token does not really hold a lot of information. With meta tokens, you get a more balanced softmax distribution over tokens.
Hymba 1.5B has proven itself as a top performer in head-to-head comparisons with leading models under 2 billion parameters, including Llama 3.2 1B, OpenELM 1B, and Qwen 2.5 1.5B. Across benchmarks such as MMLU, ARC-C, Hellaswag, and SQuAD-C, Hymba outperformed its competitors.
Source: https://arxiv.org/pdf/2411.13676
NVIDIA optimized Hymba’s training pipeline to balance task performance and efficiency. The pretraining strategy involved a two-stage process: early training on a diverse, unfiltered dataset, followed by fine-tuning on high-quality data. Instruction fine-tuning enhanced the model's capabilities through stages like supervised fine-tuning (SFT) and reinforcement learning via direct preference optimization (DPO).
Hymba 1.5B is available as an open-source release on Hugging Face and GitHub, enabling researchers and developers to test its capabilities in real-world applications.
