Whisper MLA 메모리 절감으로 장시간 음성 인식 가능하게 하는 새로운 어텐션 구조

초록

본 논문은 Whisper 모델의 멀티‑헤드 어텐션(MHA)에서 발생하는 KV 캐시 메모리 증가 문제를 해결하기 위해, 절감된 KV 캐시를 제공하는 멀티‑헤드 라텐트 어텐션(MLA)을 도입한 Whisper‑MLA를 제안한다. 절대 위치 임베딩을 유지하면서 키와 값을 저‑랭크 잠재 공간으로 압축하고, 특히 디코더 자체 어텐션에만 MLA를 적용함으로써 최대 87.5%의 메모리 절감과 최소한의 인식 정확도 손실을 달성한다.

상세 분석

Whisper‑MLA는 기존 Whisper 모델의 구조적 특성을 보존하면서 메모리 효율성을 극대화한다. 핵심 아이디어는 KV 캐시가 자동 회귀 디코더 단계에서 선형적으로 증가한다는 점을 이용해, 키와 값 행렬을 저‑랭크 잠재 공간으로 압축하는 MLA를 적용하는 것이다. 기존 MHA2MLA 프레임워크는 RoPE와 같은 상대 위치 인코딩을 전제로 했지만, Whisper는 절대 위치 인코딩(사인‑코사인 주파수 서브스페이스)을 사용한다. 따라서 저자들은 두 가지 차원 보존 전략을 설계하였다. 첫 번째는 균등 샘플링으로, 전체 주파수 서브스페이스를 일정 간격으로 선택해 r개의 서브스페이스를 유지한다. 두 번째는 헤드‑와이즈 2‑노름 기여도 기반 선택으로, 쿼리와 키 벡터의 2‑노름 곱을 평균해 가장 큰 기여도를 보이는 서브스페이스를 보존한다. 실험 결과, 전체 차원 압축보다 6.25% 정도의 키 차원을 보존하는 것이 성능 저하를 크게 완화한다는 것이 확인되었다.

모델 변환 과정에서는 기존 Whisper의 쿼리 가중치 Wq는 그대로 두고, 키 가중치 Wk를 보존 파트 Wkp와 압축 파트 Wkc로 분리한다. 이후 Wkc와 전체 값 가중치 Wv를 하나의 행렬로 결합해 공동 SVD를 수행하고, 얻어진 저‑랭크 행렬을 통해 잠재 투영 Uk와 Uv를 정의한다. 이 방식은 사전 학습된 파라미터를 최대한 재활용하면서도 새로운 저‑랭크 표현을 효율적으로 학습한다.

두 가지 아키텍처 변형이 제안되었다. Whisper‑MLA(Full)은 인코더, 디코더 자체 어텐션, 교차 어텐션 모두에 MLA를 적용해 전반적인 메모리 절감을 목표로 한다. 반면 Whisper‑MLA(DSO, Decoder Self‑attention Only)는 디코더 자체 어텐션에만 MLA를 적용한다. DSO는 KV 캐시가 디코더 단계에서만 동적으로 증가한다는 사실에 기반해, 인코더와 교차 어텐션을 그대로 유지함으로써 사전 학습된 음향 특성 추출 능력을 보존한다. 실험 결과 DSO가 Full 변형보다 WER 상승이 현저히 적으며, 메모리 절감 효과는 동일하게 유지되는 것으로 나타났다.

성능 평가에서는 Whisper‑small(244M 파라미터)을 기준 모델로 사용하고, LibriSpeech 960시간 데이터셋으로 3 epoch 미세조정을 수행하였다. KV 캐시 감소율은 DSO‑Uniform 설정에서 81.25%(실제 메모리 사용량은 87.5% 감소)이며, 평균 WER은 기존 Whisper‑fine‑tuned 대비 0.17%p만 상승했다. GPU 메모리 사용량 측정에서는 배치 크기와 시퀀스 길이가 증가할수록 Whisper‑MLA가 Whisper보다 약 50% 적은 메모리를 사용했으며, 배치 64·시퀀스 2048 상황에서 Whisper는 OOM을 일으키는 반면 Whisper‑MLA는 15.4 GB로 정상 동작했다.

이러한 결과는 절대 위치 임베딩을 갖는 음성 인식 모델에서도 MLA가 효과적으로 적용될 수 있음을 증명한다. 특히 디코더 자체 어텐션에만 MLA를 적용하는 전략이 메모리 효율성과 인식 정확도 사이의 최적 균형을 제공한다는 점이 핵심 인사이트이다. 향후 연구에서는 더 높은 차원의 키 보존 전략, 동적 랭크 조정, 그리고 실시간 스트리밍 시나리오에 대한 적용 가능성을 탐색할 여지가 있다.