Hymba：面向小型语言模型的混合头架构
基本信息
- 标题：Hymba: A Hybrid-head Architecture for Small Language Models
- 作者：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
- 提交日期：2024 年 11 月 20 日
- arXiv 编号：arXiv:2411.13676 [cs.CL]
- 主题分类：Computation and Language (cs.CL)；Artificial Intelligence (cs.AI)；Machine Learning (cs.LG)
- 备注：20 页，模型已在 Hugging Face 上提供
摘要
论文提出 Hymba，一个面向小型语言模型的混合头并行架构（hybrid-head parallel architecture），将 Transformer 注意力机制与状态空间模型（State Space Models, SSMs）相结合，以提升效率。
- 注意力头：提供高分辨率的信息召回能力。
- SSM 头：实现高效的上下文摘要。
此外，模型引入了可学习的元令牌（meta tokens），将其置于提示（prompt）之前，用于存储关键信息，并缓解注意力机制中"被迫关注（forced-to-attend）"的负担。
该模型还通过跨层键值共享（cross-layer KV sharing）与部分滑动窗口注意力（partial sliding window attention）进一步优化，从而实现紧凑的缓存大小。
在开发过程中，作者在相同设置下对多种架构进行了对照研究，观察到所提架构具有显著优势。Hymba 在小模型上取得了当前最优（state-of-the-art）结果：Hymba-1.5B-Base 模型在性能上超越所有 2B 以下的公开模型，甚至优于 Llama-3.2-3B——平均准确率高出 1.32%，缓存大小减少 11.67 倍，吞吐量提升 3.49 倍。
核心要点
- 混合头并行架构：注意力头 + SSM 头并行协作。
- 可学习元令牌：预置于提示前，存储关键信息，减轻注意力负担。
- 紧凑缓存设计：跨层 KV 共享 + 部分滑动窗口注意力。
- 小模型 SOTA：Hymba-1.5B-Base 超越所有 sub-2B 模型及 Llama-3.2-3B。