将语音识别模型Whisper的注意力机制改造，大幅降低显存占用。
Whisper-MLA: Reducing GPU Memory Consumption of ASR Models based on MHA2MLA Conversion
- 用潜空间注意力替代原模型的多头注意力，减少键值缓存
- 在长音频场景下显存减少87.5%仍保持高准确率
- 仅需少量微调即可转换预训练模型，适合部署优化
基于Transformer的Whisper模型在自动语音识别（ASR）中表现卓越，但其多头注意力（MHA）机制导致键值（KV）缓存随序列长度线性增长，显著增加GPU显存消耗，限制了长音频应用。为此，本文提出Whisper-MLA，将多头潜空间注意力（MLA）引入Whisper模型，针对绝对位置编码进行适配，并系统研究其在编码器自注意力、解码器自注意力和交叉注意力中的应用。实验表明，仅在解码器自注意力中应用MLA即可实现性能与内存效率的最佳平衡。该方法可对预训练Whisper模型进行最小微调完成转换。在LibriSpeech基准上的大量实验验证了有效性：Whisper-MLA可使KV缓存大小减少高达87.5%，同时保持竞争力的识别准确率。
原文摘要 · Abstract (English)
The Transformer-based Whisper model has achieved state-of-the-art performance in Automatic Speech Recognition (ASR). However, its Multi-Head Attention (MHA) mechanism results in significant GPU memory consumption due to the linearly growing Key-Value (KV) cache usage, which is problematic for many applications especially with long-form audio. To address this, we introduce Whisper-MLA, a novel architecture that incorporates Multi-Head Latent Attention (MLA) into the Whisper model. Specifically, we adapt MLA for Whisper's absolute positional embeddings and systematically investigate its application across encoder self-attention, decoder self-attention, and cross-attention modules. Empirical results indicate that applying MLA exclusively to decoder self-attention yields the desired balance between performance and memory efficiency. Our proposed approach allows conversion of a pretrained Whisper model to Whisper-MLA with minimal fine-tuning. Extensive experiments on the LibriSpeech benchmark validate the effectiveness of this conversion, demonstrating that Whisper-MLA reduces the KV cache size by up to 87.5% while maintaining competitive accuracy.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。