"Hymba: A Hybrid-head Architecture for Small Language Models"

Super cool model release from @NVIDIA 👏

Small but mighty, Hymba combines the best of attention and SSMs to make small LLMs smarter and faster.

Hymba's dual-processing brain outsmarts larger models.

Hymba introduces a hybrid architecture combining Transformer attention with State Space Models (SSMs) in parallel within each layer. This novel approach solves the efficiency-performance trade-off in small LLMs by leveraging attention for high-resolution recall and SSMs for efficient context processing, while using meta tokens to enhance performance.

-----

https://arxiv.org/abs/2411.13676

🤔 Original Problem:

Traditional Transformers have quadratic computational costs and high memory demands, while SSMs struggle with memory rec tasks. Existing hybrid models that stack these layers sequentially create bottlenecks when one layer type underperforms.

-----

🔧 Solution in this Paper:

→ Hymba fuses attention and SSM heads in parallel within each layer, allowing simultaneous processing of inputs through both mechanisms

→ The architecture introduces learnable meta tokens prepended to input sequences, acting as compressed world knowledge and improving attention distribution

→ Cross-layer key-value cache sharing and partial sliding window attention optimize memory usage and computational efficiency

→ Meta tokens serve dual purposes: mitigating attention drain through backstop tokens and encapsulating domain-specific knowledge

-----

💡 Key Insights:

→ Parallel fusion of attention and SSM heads outperforms sequential stacking

→ Using global attention in just three layers (first, middllast) maintains recall accuracy while reducing cache size

→ Meta tokens help redistribute attention more effectly across different types of tokens

→ SSM heads focus on current tokens while attention heads handle cross-token relatioips

-----

📊 Results:

→ Hymba-1.5B outperforms all sub-2B public models in performance

→ Achieves 1.32% higher average accuracy than Llama-3.2-3B

→ Delivers 11.67x smaller cache size and 3.49x faster throughput compared to Llama-3.2-3B