다중스케일 계층형 주의 메커니즘 MAHA

초록

MAHA는 입력 시퀀스를 학습 가능한 다운샘플링으로 여러 해상도(스케일)로 분해하고, 각 스케일에서 독립적인 어텐션을 계산한다. 이후 스케일별 출력을 볼록 최적화 혹은 내시 균형 게임 이론에 기반한 가중합으로 통합해, 전통적인 MHSA의 O(N²) 복잡도를 크게 낮추면서도 전역·국부 정보를 균형 있게 보존한다. 실험 결과 4096 길이에서 81 % FLOPs 절감과 기존 모델 대비 동등하거나 향상된 성능을 입증했다.

상세 분석

본 논문은 대규모 언어 모델에서 가장 큰 병목인 다중 헤드 자체 주의(MHSA)의 이차 복잡도를 구조적·수학적으로 해결하고자 한다. 첫 번째 핵심은 계층적 다중스케일 분해이다. 입력 텐서 X∈ℝ^{N×d}에 대해 학습 가능한 다운샘플링 연산 D_l(·)을 순차적으로 적용해 L개의 스케일 X_l을 만든다. D_l은 스트라이드 1D 컨볼루션 혹은 어댑티브 풀링으로 구현되며, 압축 비율 r>1에 따라 길이 n_l=⌊n_{l‑1}/r⌋ 로 기하급수적으로 감소한다. 이렇게 하면 높은 스케일일수록 전역 의미를, 낮은 스케일일수록 국부 구문 정보를 효율적으로 캡처한다.

두 번째 혁신은 스케일별 어텐션 계산이다. 각 스케일마다 Q_l, K_l은 개별 파라미터 W_Q^l, W_K^l 로 투사하지만, V는 모든 스케일에 대해 공유값 프로젝션 V_base = XW_V 를 사용한다. 이후 V_l = D_l(V_base) 로 다운샘플링해 스케일별 어텐션 O_l = softmax(Q_lK_l^T/√d_k)·V_l 를 얻는다. 이 설계는 파라미터 수를 크게 줄이며, 값 정보의 일관성을 유지한다.

세 번째 핵심은 수학적으로 정당화된 스케일 통합이다. 저자는 두 가지 접근을 제시한다. ① 볼록 최적화 기반: 가중치 w∈ℝ^{L} 를 ℓ₁ 정규화와 합계 제약(∑w_l=1, w_l≥0) 하에 최소화 문제 min_w ‖∑_l w_l·U_l(O_l) – O*‖₂² + λ‖w‖₁ 을 풀어 가장 정보량이 풍부한 스케일을 자동 선택한다. ② 내시 균형 기반: 각 스케일을 플레이어로 보고, 다른 스케일의 가중치를 고정했을 때 자신의 재구성 오차를 최소화하는 w_l* 를 찾는다. 이 게임은 비협력적이지만, 균형점에서 어느 스케일도 단독으로 개선할 여지가 없어 안정적인 통합을 보장한다. 두 방법 모두 미분 가능 최적화 레이어로 구현돼 역전파가 가능하다.

복합적으로, 하이브리드 팽창 컨볼루션 블록을 어텐션 앞에 삽입해 지역 패턴을 먼저 추출하고, 크로스 스케일 게이팅으로 상위 스케일이 하위 스케일에 정보를 전달한다. 복잡도 분석에서는 전체 연산량이 Ω(N·∑_{l=0}^{L-1} (1/r^l)²·d) ≈ O(N·d·(1‑1/r^{2L})/(1‑1/r²)) 로, r=2일 경우 거의 선형에 근접한다. 실험에서는 4096 토큰 입력 시 표준 어텐션 대비 FLOPs를 81 % 절감했으며, GLUE, PG‑19, WMT14, SQuAD 등 다양한 벤치마크에서 성능 저하 없이 혹은 소폭 향상을 기록했다.

이러한 설계는 다중스케일 의미론적 granularity를 유지하면서도 자원 할당 최적화를 수학적으로 보장한다는 점에서 기존의 희소 어텐션(고정 패턴)이나 단순 계층 어텐션(가중 평균)보다 이론적·실용적 우위를 가진다. 또한, 최적화·게임 이론 레이어를 통해 학습 과정에서 자동으로 스케일 중요도를 조정하므로, 다양한 길이와 도메인에 대한 적응성이 뛰어나다.