Whisper-MLA: Reducing GPU Memory Consumption of ASR Models based on MHA2MLA Conversion
Abstract
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.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
The paper presents Whisper-MLA, a novel architecture that reduces GPU memory consumption in ASR models while preserving performance. This work is significant as it addresses a critical bottleneck in deploying state-of-the-art ASR systems, particularly for long-form audio applications, thereby enhancing accessibility and usability in various practical contexts.
Methodology
The proposed methodology introduces a novel architecture, Whisper-MLA, which effectively integrates Multi-Head Latent Attention (MLA) into the Whisper model. The authors adapt MLA specifically for absolute positional embeddings, which is a significant innovation given the existing limitations of applying MLA to encoder-decoder architectures. The systematic investigation of MLA's application across different attention modules is commendable, and the decision to focus on decoder self-attention for optimization reflects a well-thought-out approach to balancing memory efficiency and performance.
Results & evaluation
The experiments conducted on the LibriSpeech benchmark are extensive and demonstrate the effectiveness of the Whisper-MLA model in reducing GPU memory consumption significantly while maintaining competitive accuracy. The results clearly show that the proposed model achieves up to 87.5% reduction in KV cache size, which is a critical metric for real-world applications, especially in resource-constrained environments. The comparative analysis with the original Whisper model provides a solid basis for the claims made.
Reproducibility
The paper provides sufficient details regarding the experimental setup, including the model architecture, training parameters, and the dataset used. However, the lack of a publicly available code repository limits the reproducibility of the results. The authors mention that their source code is publicly available, but no specific URL is provided in the text, which is a missed opportunity for enhancing reproducibility.
Limitations
One limitation is that the paper primarily focuses on the decoder self-attention mechanism, potentially overlooking the benefits that could be derived from optimizing other components of the model. Additionally, while the results are promising, the experiments are limited to the LibriSpeech dataset, which may not fully represent the model's performance across diverse ASR tasks and environments.
Broader impact
The Whisper-MLA architecture has significant implications for deploying large-scale ASR models in real-world applications, particularly in scenarios where GPU memory is a limiting factor. By reducing memory consumption while maintaining performance, this work could facilitate the use of advanced ASR technologies in mobile devices, embedded systems, and other resource-constrained environments. The paper presents Whisper-MLA, a novel architecture that reduces GPU memory consumption in ASR models while preserving performance. This work is significant as it addresses a critical bottleneck in deploying state-of-the-art ASR systems, particularly for long-form audio applications, thereby enhancing accessibility and usability in various practical contexts.
Author of this paper? Send a correction.