中文

QKV 投影仅需其部分内存

机器学习 2026-03-03 v3

摘要

多头注意力机制是 LLM 运行的核心,多项工作致力于提升其在训练期间的计算和内存效率。尽管大多数工作专注于近似缩放点积,但用于从输入 xx 计算 QQKKVV 张量的线性投影的内存消耗却常被忽视。为了解决这个问题,我们提出了逐点近似矩阵乘法(PAMM),这是一种新颖的张量压缩技术,可将注意力层中 Q,K,VQ,K,V 投影的激活值压缩多达 ×512\times 512 倍,在有效消除其内存占用的同时,实现了相似或更好的最终困惑度。PAMM 与 FlashAttention 等高效注意力技术完全可组合,使其成为内存高效 LLM 训练的实用且互补的方法。

关键词

引用

@article{arxiv.2506.02939,
  title  = {QKV Projections Require a Fraction of Their Memory},
  author = {Malik Khalaf and Yara Shamshoum and Nitzan Hodos and Yuval Sieradzki and Assaf Schuster},
  journal= {arXiv preprint arXiv:2506.02939},
  year   = {2026}
}

备注

Accepted to ICLR 2026