小语言模型的隐式多头注意力
计算与语言
2025-06-17 v2 人工智能
摘要
我们首次对隐式多头注意力（MLA）在小语言模型中的应用进行了全面研究，揭示了有趣的效率-质量权衡。在 100,000 个合成故事上训练 30M 参数的 GPT 模型后，我们对三种架构变体进行了基准测试：标准多头注意力（MHA）、MLA 以及带旋转位置编码的 MLA（MLA+RoPE）。我们的关键发现是，MLA+RoPE 采用半秩隐式维度（r = d/2）可在仅增加 0.3% 验证损失的情况下实现 45% 的 KV 缓存内存减少（本质上与 MHA 质量持平）——这是内存受限部署的帕累托改进。我们进一步表明 RoPE 对小模型中的 MLA 至关重要：没有 RoPE 时 MLA 表现不如标准注意力 3-5%，而有了 RoPE 后则超越标准注意力 2%。在 NVIDIA A100 GPU 上的推理基准测试显示，MLA 在 r=d/2 时相比全秩 MLA 实现 1.4 倍加速，同时保持内存节省。GPT-4 评估证实了困惑度结果，在语法、创造性和一致性指标上我们的模型取得了最高质量分数（7.4/10）。代码和模型将在接收后发布。
引用
@article{arxiv.2506.09342,
title = {Latent Multi-Head Attention for Small Language Models},
author = {Sushant Mehta and Raj Dandekar and Rajat Dandekar and Sreedath Panat},
journal= {arXiv preprint arXiv:2506.09342},
year = {2025}
}
备注
6 pages, 1 figure. 5 tables