Hymba: Eine Hybrid-Head-Architektur für kleine Sprachmodelle

Hymba: A Hybrid-head Architecture for Small Language Models

November 20, 2024

Autoren: Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabaleshwarkar, Shih-Yang Liu, Matthijs Van Keirsbilck, Min-Hung Chen, Yoshi Suhara, Yingyan Lin, Jan Kautz, Pavlo Molchanov

cs.AI

Zusammenfassung

Wir schlagen Hymba vor, eine Familie von kleinen Sprachmodellen mit einer Hybrid-Head-Parallelarchitektur, die Transformer-Aufmerksamkeitsmechanismen mit Zustandsraummodellen (SSMs) zur Verbesserung der Effizienz integriert. Aufmerksamkeitsköpfe bieten eine hochauflösende Rückrufmöglichkeit, während SSM-Köpfe eine effiziente Kontextzusammenfassung ermöglichen. Darüber hinaus führen wir erlernbare Metatoken ein, die den Anfragen vorangestellt werden, um wichtige Informationen zu speichern und die mit Aufmerksamkeitsmechanismen verbundene "gezwungen-zu-beachten"-Belastung zu lindern. Dieses Modell wird weiter optimiert, indem wir Kreuzschicht-Schlüssel-Wert-(KV)-Teilung und teilweises Schiebefenster-Aufmerksamkeit integrieren, was zu einer kompakten Zwischenspeichergröße führt. Während der Entwicklung führten wir eine kontrollierte Studie durch, in der wir verschiedene Architekturen unter identischen Bedingungen verglichen und signifikante Vorteile unserer vorgeschlagenen Architektur beobachteten. Beachtenswert ist, dass Hymba Spitzenleistungen für kleine Sprachmodelle erzielt: Unser Hymba-1.5B-Base-Modell übertrifft alle öffentlichen Modelle unter 2B in der Leistung und übertrifft sogar Llama-3.2-3B mit 1,32 % höherer durchschnittlicher Genauigkeit, einer 11,67-fachen Reduzierung der Zwischenspeichergröße und einer 3,49-fachen Durchsatzsteigerung.

English

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.