QKV 投影仅需其部分内存
机器学习
2026-03-03 v3
摘要
多头注意力机制是 LLM 运行的核心,多项工作致力于提升其在训练期间的计算和内存效率。尽管大多数工作专注于近似缩放点积,但用于从输入 计算 、 和 张量的线性投影的内存消耗却常被忽视。为了解决这个问题,我们提出了逐点近似矩阵乘法(PAMM),这是一种新颖的张量压缩技术,可将注意力层中 投影的激活值压缩多达 倍,在有效消除其内存占用的同时,实现了相似或更好的最终困惑度。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