近年来，大语言模型（LLM）在各个领域取得了显著成效。然而，现有的 Transformer 架构存在计算复杂度高、内存消耗大等固有瓶颈。相比之下，状态空间模型（SSM）如 Mamba 虽然具有常数时间复杂度和优化的硬件性能，但在长序列记忆回溯任务上表现相对较弱。针对这一痛点，NVIDIA 提出了 Hymba 架构，通过在同一层中结合注意力头和 SSM 头，实现了两种架构优势的互补，旨在平衡推理效率与上下文理解能力。
核心创新
Hymba 的核心创新主要体现在以下三个方面：
-
并行混合头设计： 在同一网络层内并行集成注意力头和 SSM 头。注意力机制负责提供高分辨率的记忆回溯能力，确保关键信息的精准定位；而 SSM 则提供高效的上下文总结能力，降低长序列的处理成本。这种设计相比 Zamba 和 Jamba 等仅在深层或不同层使用单一机制的方法更加灵活，能够动态适应不同层级的需求。
-
可学习的元令牌（Meta Tokens）： 在输入序列前添加一组可学习的元令牌。这些令牌与后续所有序列令牌进行交互，充当知识的压缩表示。它们不仅提高了模型的回溯能力，还增强了通用任务的泛化性能，类似于一种隐式的提示工程。
-
KV 缓存优化： 采用层间共享 KV 缓存的策略。大多数层使用滑动窗口注意力机制（SWA），仅在关键层保留全局视野。这种策略显著减少了显存占用和计算成本，使得模型在长文本处理时更加高效。
架构设计
Hymba 的混合头模块设计遵循统一且对称的原则。其工作流程主要包含以下步骤：
-
输入处理： 在输入序列前添加 Meta Tokens。随后通过投影层将输入转换为查询（Query）、键（Key）、值（Value）以及 SSM 特征。这一过程确保了两种机制接收到的信息源一致。
-
并行处理： 注意力头专注于高精度记忆回溯，捕捉局部依赖关系；SSM 头则进行高效的上下文总结，处理长距离依赖。两种头并行处理相同的输入信息，互不干扰但输出互补。
-
输出融合： 对注意力头和 SSM 头的输出进行归一化处理。通过可学习的向量进行重新缩放，最后取平均得到最终输出。这种融合方式保证了梯度流动的稳定性。
性能优势
相比现有模型，Hymba-1.5B 在多个维度展现出显著优势：
-
与 Llama 3.2 3B 对比：
- 准确率提高 1.32%
- 缓存大小减少 11.67 倍
- 吞吐量提高 3.49 倍
-
与同等规模（2B 以下）模型对比：
- 在常识推理任务上取得最佳性能
- 需要的缓存大小显著减小
- 具有更高的处理速度
-
指令微调后的变体 Hymba-1.5B-Instruct：
- 在 GSM8K 和 GPQA 等基准测试上表现优异
- 经常超越更大规模的基线模型
Hymba 架构实现与实验评估
1. 融合混合头模块设计
Hymba 提出了一个统一且对称的模块设计公式。对于输入序列 X̃（原始输入序列 X 加上元令牌），主要包括：
输入投影： 使用 Win_proj = [WQ, WK, WV, WSSM, WG] 进行投影，生成注意力头的查询、键、值，同时生成 SSM 头的输入特征和门控信号。这一设计简化了参数管理，提高了训练效率。
注意力头输出： 标准自注意力机制计算，重点关注局部上下文和关键信息点的召回。
SSM 头输出： 基于状态空间模型的线性递归计算，擅长处理长序列的全局趋势总结。
输出融合： 其中β1 和β2 是可学习的向量，用于重新缩放各通道的输出，确保不同机制的贡献比例可控。
2. KV 缓存优化策略
全局与局部注意力结合： 仅在关键层（第一层、中间层和最后一层）使用全局注意力，其他层使用滑动窗口注意力（SWA）。该策略在维持性能的同时显著提升效率，避免了全量 KV 缓存带来的显存爆炸。
