通过K-cache技术让注意力机制内存减半,推理速度翻倍且不损失精度。
Slim attention: cut your context memory in half without loss -- K-cache is all you need for MHA
- 用K-cache替代传统注意力存储,实现数学上完全等价的内存压缩。
- 在大上下文场景下推理速度提升2倍,Whisper模型内存减少8倍、生成提速5倍。
- 适合需要长序列推理的场景,如语音转写、超长文本生成等应用。
Slim attention 将使用多头注意力(MHA)的Transformer模型的上下文内存大小减半,可在大上下文窗口下将推理速度提升最高2倍。该方法是标准注意力机制的精确数学等价实现,不牺牲模型准确性,实现了无损内存压缩。对于编码器-解码器结构的模型,压缩效果更显著:以Whisper模型为例,上下文内存可减少8倍,批量大小为64时,令牌生成速度提升5倍;T5-11B模型因其多头投影维度大于嵌入维度,内存可压缩32倍。代码与更多技巧详见 https://github.com/OpenMachine-ai/transformer-tricks,视频讲解见 https://www.youtube.com/watch?v=uVtk3B6YO4Y。
原文摘要 · Abstract (English)
Slim attention shrinks the context memory size by 2x for transformer models with MHA (multi-head attention), which can speed up inference by up to 2x for large context windows. Slim attention is an exact, mathematically identical implementation of the standard attention mechanism and therefore doesn't compromise model accuracy. In other words, slim attention losslessly compresses the context memory by a factor of 2. For encoder-decoder transformers, the context memory size can be reduced even further: For the Whisper models for example, slim attention reduces the context memory by 8x, which can speed up token generation by 5x for batch size 64 for example. And for the T5-11B model for example, the memory can be reduced by 32x because its MHA projection dimension is larger than the embedding dimension. See https://github.com/OpenMachine-ai/transformer-tricks for code and more transformer tricks, and https://www.youtube.com/watch?v=uVtk3B6YO4Y for this paper's YouTube video.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。