通过低秩分解提升注意力模型容量，显著降低显存占用。
Multi-matrix Factorization Attention
- 用低秩矩阵分解扩展查询-键电路中的注意力头数量与维度
- 在严格缓存限制下，显存占用减少最高达93.7%
- 适合大模型部署场景，尤其对内存受限系统友好
我们提出新型注意力架构——多矩阵分解注意力（MFA）及其改进版本MFA-Key-Reuse（MFA-KR）。现有标准多头注意力（MHA）变体，包括最先进的MLA方法，在严格的键值缓存（KV cache）约束下表现下降明显。MFA通过在查询-键（QK）电路中采用低秩矩阵分解，高效提升注意力头的数量与维度，从而增强模型容量。在此基础上，MFA-KR进一步通过值投影重参数化，将键缓存复用为值缓存，显著降低内存开销。大规模实验表明，所提架构在紧致的KV缓存预算下仍保持强性能，优于MLA，且接近MHA表现；相比后者，其KV缓存使用量最高可降低93.7%。
原文摘要 · Abstract (English)
We propose novel attention architectures, Multi-matrix Factorization Attention (MFA) and MFA-Key-Reuse (MFA-KR). Existing variants for standard Multi-Head Attention (MHA), including SOTA methods like MLA, fail to maintain as strong performance under stringent Key-Value cache (KV cache) constraints. MFA enhances model capacity by efficiently scaling up both the number and dimension of attention heads through low-rank matrix factorization in the Query-Key (QK) circuit. Extending MFA, MFA-KR further reduces memory requirements by repurposing the key cache as value through value projection re-parameterization. MFA's design enables strong model capacity when working under tight KV cache budget, while MFA-KR is suitable for even harsher KV cache limits with minor performance trade-off. Notably, in our extensive and large-scale experiments, the proposed architecture outperforms MLA and performs comparably to MHA, while reducing KV cache usage by up to 56% and 93.7%, respectively.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。