压缩注意力投影内存，提升大模型训练效率
QKV Projections Require a Fraction of Their Memory
- 提出PAMM技术，对QKV投影激活值进行高达512倍压缩
- 压缩后模型困惑度不变或更好，显著降低内存占用
- 可与FlashAttention等方法无缝组合，适合大规模训练场景
多头注意力机制是大语言模型的核心，现有研究多关注缩放点积计算的效率，却常忽略从输入x生成Q、K、V张量的线性投影所消耗的内存。为此，我们提出点近似矩阵乘法（PAMM），一种新型张量压缩技术，可将注意力层中Q、K、V投影的激活值压缩高达512倍，有效消除其内存开销，同时保持甚至提升最终困惑度。PAMM可与FlashAttention等高效注意力技术完全兼容，是一种实用且互补的内存优化方案，适用于大模型训练。
原文摘要 · Abstract (English)
The Multi-Head Attention mechanism is central to LLM operation, and multiple works target its compute and memory efficiency during training. While most works focus on approximating the scaled dot product, the memory consumption of the linear projections that compute the $Q$, $K$, and $V$ tensors from the input $x$ is often overlooked. To address this, we propose Point-Approximate Matrix Multiplication (PAMM), a novel tensor compression technique that compresses the activations of the $Q,K,V$ projections in attention layers by a factor of up to $\times 512$, effectively erasing their memory footprint, while achieving similar or better final perplexity. PAMM is fully composable with efficient attention techniques such as FlashAttention, making it a practical and complementary method for memory-efficient LLM training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。