多矩阵分解注意力
机器学习
2025-01-15 v2 计算与语言
摘要
我们提出了新的注意力架构——多矩阵分解注意力（Multi-matrix Factorization Attention, MFA）和MFA-Key-Reuse（MFA-KR）。现有的多头注意力（Multi-Head Attention, MHA）变体，包括像MLA这样的SOTA方法，在面对严格的键值缓存（KV cache）约束时，未能保持同样的性能。MFA通过在查询-键（Query-Key, QK）电路中采用低秩矩阵分解，高效地扩展了注意力头的数量和维度，从而提升了模型容量。扩展MFA，MFA-KR通过价值投影重参数化，将键缓存复用为价值，从而进一步降低内存需求。MFA的设计使其在受限KV缓存预算下仍能保持强大的模型容量，而MFA-KR则适用于更严苛的KV缓存限制，虽有轻微的性能权衡。在我们的大规模实验中，所提出的架构超过了MLA，并在性能上与MHA相当，同时将KV缓存使用量分别降低了最高可达56%和93.7%。
引用
@article{arxiv.2412.19255,
title = {Multi-matrix Factorization Attention},
author = {Jingcheng Hu and Houyi Li and Yinmin Zhang and Zili Wang and Shuigeng Zhou and Xiangyu Zhang and Heung-Yeung Shum and Daxin Jiang},
journal= {arXiv preprint arXiv:2412.19255},
year = {2025}
}