Slim attention：无损将上下文内存减半——MHA 仅需 K-cache
机器学习
2025-06-04 v2
摘要
Slim attention 将具有 MHA（多头注意力）的 Transformer 模型的上下文内存大小缩小了 2 倍，对于大上下文窗口，这可以将推理速度提高多达 2 倍。Slim attention 是标准注意力机制的精确、数学上等价的实现，因此不会损害模型精度。换言之，Slim attention 将上下文内存无损压缩了 2 倍。对于编码器-解码器 Transformer，上下文内存大小可以进一步减小：例如，对于 Whisper 模型，Slim attention 将上下文内存减少了 8 倍，对于批大小为 64 的情况，可以将 Token 生成速度提高 5 倍。再例如，对于 T5-11B 模型，由于其 MHA 投影维度大于嵌入维度，内存可以减少 32 倍。代码和更多 Transformer 技巧请参见 https://github.com/OpenMachine-ai/transformer-tricks，本文的 YouTube 视频请参见 https://www.youtube.com/watch?v=uVtk3B6YO4Y。
引用
@article{arxiv.2503.05840,
title = {Slim attention: cut your context memory in half without loss -- K-cache is all you need for MHA},
author = {Nils Graef and Andrew Wasielewski},
journal= {arXiv preprint arXiv:2503.05840},
year = {2025}
}
备注
18 pages, 7 figures