MambAttention 일반화 가능한 단일채널 음성 향상을 위한 Mamba와 멀티헤드 어텐션 결합

초록

본 논문은 Mamba 기반 시계열 모델에 시간·주파수 멀티헤드 어텐션(MHA)을 공유 가중치로 결합한 하이브리드 구조 MambAttention를 제안한다. 새로운 고난도 데이터셋 VB‑DemandEx로 학습한 뒤, DNS 2020(무반향) 및 EARS‑WHAM v2와 같은 외부 코퍼스에서 기존 LSTM, xLSTM, Mamba, Conformer 및 최신 diffusion 모델을 능가하는 일반화 성능을 보였다. 가중치 공유와 MHA를 Mamba 블록 앞에 배치하는 것이 성능 향상의 핵심이며, LSTM·xLSTM에 동일 모듈을 추가해도 개선 효과가 확인되었다.

상세 분석

MambAttention은 기존 시계열 모델이 갖는 장점인 선형 시간 복잡도와 장기 의존성 학습 능력을 유지하면서, 전통적인 트랜스포머식 어텐션이 제공하는 전역 컨텍스트 통합을 도입한다. 구체적으로 각 레이어는 시간‑Mamba(T‑Mamba)와 주파수‑Mamba(F‑Mamba) 블록 사이에 시간‑MHA(T‑MHA)와 주파수‑MHA(F‑MHA)를 삽입한다. 여기서 T‑MHA와 F‑MHA는 동일 파라미터 집합을 공유하도록 설계돼, 시간축과 주파수축 모두에서 동일한 어텐션 매핑을 학습한다. 이 공유 메커니즘은 두 축의 스펙트로그램 특성을 동시에 정렬시켜, 노이즈 유형이나 SNR이 크게 변하는 상황에서도 모델이 데이터셋 불변적인 표현을 학습하도록 유도한다.

Mamba 자체는 구조화된 상태공간 모델(SSM)을 기반으로 하며, 입력‑종속적인 A, B, C, Δ 파라미터를 통해 채널별 선형 동역학을 구현한다. 논문에서는 Mamba의 수식(7)~(10)을 그대로 차용하면서, 입력 차원을 K(주파수)와 시간 T에 대해 독립적으로 적용한다. 이렇게 하면 각 시간 프레임마다 복잡도 O(N·K)인 선형 연산만 수행되므로, 길이가 수천 프레임에 이르는 오디오에도 메모리·연산 부담이 적다.

어텐션 모듈은 스케일드 닷‑프로덕트 방식으로 구현되며, Q, K, V를 각각 1×d_k, 1×d_k, 1×d_v 형태로 변환한다. 중요한 점은 MHA를 Mamba 블록 앞에 배치했다는 점이다. 실험 결과, 어텐션이 먼저 전역 정보를 제공하고, 그 뒤에 Mamba가 지역적인 동역학을 정교화함으로써 일반화 성능이 크게 향상됨을 확인했다. 반대로 Mamba 뒤에 어텐션을 두면 성능 저하가 발생한다.

데이터 측면에서는 기존 VoiceBank+Demand를 확장한 VB‑DemandEx를 제안한다. 이 데이터셋은 30 % 이상의 새로운 잡음 종류(예: 기계음, 교통소음, 실내 대화 잡음 등)와 -5 dB 이하의 저 SNR 구간을 포함한다. 따라서 모델이 훈련 단계에서 다양한 잡음·음성 비율을 경험하게 되어, 실제 환경에서의 강인성이 크게 증가한다.

평가에서는 두 개의 외부 코퍼스(DNS 2020 무반향, EARS‑WHAM v2)를 사용해 인‑도메인·아웃‑오브‑도메인 모두에서 성능을 측정했다. PESQ, STOI, SI‑SDR, DNS‑MOS 등 6가지 지표에서 MambAttention은 동일 파라미터 규모의 LSTM, xLSTM, Mamba, Conformer 대비 평균 0.120.25 dB(또는 0.030.07 MOS)씩 우수했다. 특히 diffusion 기반 생성 모델과 비교했을 때, 추론 속도는 3배 이상 빠르면서도 품질은 동등하거나 약간 앞섰다.

Ablation 연구에서는 (1) T‑MHA와 F‑MHA 가중치 공유 여부, (2) 어텐션 순서(앞/뒤), (3) 레이어 수와 헤드 수, (4) Δ 파라미터 고정 여부 등을 변형했다. 공유 가중치를 사용하지 않으면 아웃‑오브‑도메인 PESQ가 평균 0.08 dB 감소하고, 어텐션을 뒤에 두면 전체 지표가 5 % 이상 악화되는 등, 제안 설계가 일반화에 결정적임을 입증했다.

또한 LSTM·xLSTM에 동일한 공유 MHA 모듈을 삽입한 실험에서도 성능이 0.04~0.09 dB 상승했지만, MambAttention 자체가 여전히 최고 수준을 유지했다. 이는 Mamba의 선형 복합성 및 입력‑종속 동역학이 어텐션과 시너지 효과를 내는 구조적 장점임을 시사한다.

시각화 측면에서는 t‑SNE를 이용해 각 모델의 은닉 표현을 2‑D에 투사했으며, MambAttention은 잡음 종류와 SNR에 관계없이 클러스터가 명확히 구분되는 반면, LSTM·Conformer은 잡음에 따라 크게 섞이는 모습을 보였다. 이는 MHA가 전역적인 스펙트럼 정렬을 수행해 데이터 불변성을 강화한다는 가설을 뒷받침한다.

마지막으로 대규모 DNS 2020 데이터(≈500 h)로 학습했을 때, MambAttention은 파라미터 수 대비 성능 상승률이 가장 높았으며, 학습 데이터가 늘어날수록 과적합 현상이 거의 관찰되지 않았다. 이는 제안 모델이 데이터 규모에 대해 높은 확장성을 가지고 있음을 의미한다.

요약하면, MambAttention은 (1) 선형 시간 복잡도의 SSM 기반 Mamba, (2) 전역 컨텍스트를 제공하는 공유 멀티헤드 어텐션, (3) 어려운 잡음·저 SNR 환경을 반영한 VB‑DemandEx 데이터셋이라는 세 축을 결합해, 기존 최첨단 모델들을 능가하는 일반화 성능을 달성한 혁신적인 음성 향상 프레임워크이다.