arXiv:2505.13544cs.LGcs.AI2025-05NeurIPS被引 2

通过动态压缩时序键值缓存,显著提升自注意力推理速度与内存效率。

Multi-head Temporal Latent Attention

论文配图:Multi-head Temporal Latent Attention
图 1 · 摘自论文原文
  • 用超网络动态合并相邻时间步的键值缓存向量,降低内存占用。
  • 在英德语音翻译任务中实现5.3倍加速,显存降低8.3倍。
  • 适合高并发、长序列推理场景,如语音翻译与文本摘要。

尽管Transformer自注意力具备强并行性,但键值(KV)缓存随序列长度线性增长,成为推理效率瓶颈。多头潜在注意力近期被提出,将KV缓存压缩至低秩潜在空间。本文提出多头时序潜在注意力(MTLA),进一步沿时序维度压缩KV缓存,大幅降低自注意力推理的内存开销。MTLA采用超网络动态合并时间上相邻的KV缓存向量,并设计一种步长感知因果掩码,以解决压缩后缓存与处理序列长度不匹配的问题,确保训练高效且与推理行为一致。在语音翻译、语音识别、语音理解及文本摘要等任务上的实验表明,MTLA相比标准多头注意力(MHA)性能相当,但显著提升推理速度与GPU内存利用率。例如,在英语-德语语音翻译任务中,MTLA实现5.3倍加速,显存使用减少8.3倍,同时保持翻译质量。

原文摘要 · Abstract (English)

While Transformer self-attention offers strong parallelism, the Key-Value (KV) cache grows linearly with sequence length and becomes a bottleneck for inference efficiency. Multi-head latent attention was recently developed to compress the KV cache into a low-rank latent space. This paper proposes Multi-head Temporal Latent Attention (MTLA), which further reduces the KV cache size along the temporal dimension, greatly lowering the memory footprint of self-attention inference. MTLA employs a hyper-network to dynamically merge temporally adjacent KV cache vectors. To address the mismatch between the compressed KV cache and processed sequence lengths, a stride-aware causal mask is proposed to ensure efficient parallel training and consistency with inference behaviour. Experiments across tasks, including speech translation, speech recognition, speech understanding and text summarisation, demonstrate that MTLA achieves competitive performance compared to standard Multi-Head Attention (MHA), while greatly improving inference speed and GPU memory usage. For example, on a English-German speech translation task, MTLA achieves a 5.3x speedup and a reduction in GPU memory usage by a factor of 8.3 compared to MHA, while maintaining translation quality.

自注意力推理优化语音翻译内存压缩

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。