Hymba: Uma Arquitetura de Cabeça Híbrida para Modelos de Linguagem Pequenos

Hymba: A Hybrid-head Architecture for Small Language Models

November 20, 2024

Autores: 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

Resumo

Propomos o Hymba, uma família de pequenos modelos de linguagem apresentando uma arquitetura paralela híbrida que integra mecanismos de atenção do transformer com modelos de espaço de estado (SSMs) para maior eficiência. As cabeças de atenção fornecem um recall de alta resolução, enquanto as cabeças SSM permitem uma sumarização eficiente do contexto. Além disso, introduzimos meta tokens aprendíveis que são antepostos às sugestões, armazenando informações críticas e aliviando o fardo "forçado a atender" associado aos mecanismos de atenção. Este modelo é ainda otimizado ao incorporar o compartilhamento de chave-valor (KV) entre camadas e atenção de janela deslizante parcial, resultando em um tamanho de cache compacto. Durante o desenvolvimento, realizamos um estudo controlado comparando várias arquiteturas sob configurações idênticas e observamos vantagens significativas de nossa arquitetura proposta. Notavelmente, o Hymba alcança resultados de ponta para pequenos LMs: Nosso modelo Hymba-1.5B-Base supera todos os modelos públicos sub-2B em desempenho e até mesmo supera o Llama-3.2-3B com 1,32% de precisão média mais alta, uma redução de tamanho de cache de 11,67 vezes e uma taxa de transferência 3,49 vezes maior.

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.