NVIDIA研究人员发布了开源语言模型Hymba 1.5B，该模型结合了Transformer和状态空间模型（SSM）架构，旨在解决传统Transformer模型的计算和内存限制，同时提升SSM的召回能力。Hymba通过创新的混合头模块，将注意力头与SSM头并行结合，实现了高效的计算和内存使用，同时保持性能。

主要观点

1. 架构创新：Hymba 1.5B结合了Transformer和SSM的优点，通过混合头模块并行处理高分辨率召回和上下文总结，减少了计算和内存需求。

2. 效率提升：

   • 注意力开销减少：超过50%的注意力计算被SSM处理取代，降低了计算成本。
   • 局部注意力主导：全局注意力被最小化，局部注意力与SSM结合足以总结全局信息。
   • KV缓存优化：引入跨层KV缓存共享，减少了缓存冗余，内存使用量比传统Transformer模型少10倍。
   • 元标记：128个可学习的嵌入作为记忆初始化器，平衡了注意力分布。
3. 性能表现：Hymba 1.5B在多个基准测试中（如MMLU、ARC-C、Hellaswag、SQuAD-C）表现优异，超越了其他参数规模相近的模型（如Llama 3.2 1B、OpenELM 1B、Qwen 2.5 1.5B）。
4. 训练优化：NVIDIA优化了Hymba的训练流程，包括两阶段预训练（早期训练在多样化未过滤数据集上，后续在高质量数据上进行微调）和指令微调（如监督微调和直接偏好优化）。

关键信息

• 开源发布：Hymba 1.5B在Hugging Face和GitHub上开源，供研究者和开发者测试和应用。
• 技术细节：模型通过元标记优化注意力机制，解决了传统注意力机制过分关注句子开头的问题。

重要细节

• 混合头模块：将注意力头和SSM头结合，实现了并行处理，提升了效率。
• 跨层KV缓存共享：显著减少了内存使用，提升了模型的可扩展性。
• 元标记作用：通过可学习的嵌入初始化记忆，平衡了注意力分布，提升了模型性能。

Hymba 1.5B通过结合Transformer和SSM的优势，实现了高效且高性能的语言模型，为小规模语言模型的发展提供了新的方向。