기술 요약: SEMA – 토큰 국소화 및 평균화를 통한 확장 가능하고 효율적인 Mamba 스타일 어텐션

1. 문제 정의

본 논문은 컴퓨터 비전 작업에 어텐션 메커니즘을 적용할 때 발생하는 두 가지 주요 과제를 다룹니다:

1. 계산 복잡도: 바닐라 풀 어텐션(softmax 기반)은 입력 크기에 대해 이차적 계산 복잡도($O(n^2)$)를 보이며, 이는 고해상도 이미지에 적용하기에는 매우 부담스럽습니다.
2. 분산 현상 (The Dispersion Phenomenon): 선형 복잡도를 달성하기 위해 설계된 선형 어텐션 변형들은 종종 관련 정보에 집중하지 못하는 문제를 겪습니다. 토큰 수($n$)가 증가함에 따라, 이러한 메커니즘은 모든 키(key)에 거의 동일한 가중치를 할당하여 소프트맥스 어텐션이 가진 변별력을 상실하는 "분산"되는 경향이 있습니다.

Mamba(상태 공간 모델) 및 MILA(Mamba에서 영감을 받은 선형 어텐션)와 같은 최근 연구들이 이러한 문제들을 해결하려고 시도했으나, 저자들은 선형 어텐션 메커니즘이 여전히 표현력 있는 집중 능력이 부족하며, 분산 문제를 해결하기 위한 기존 솔루션들(예: 적응형 온도 조절 또는 차분 어텐션)은 대개 휴리스틱적이거나 이론적으로 여전히 극한 상황에서의 분산 문제에 취약하다고 관찰했습니다.

2. 방법론

2.1 이론적 기초: 일반화된 어텐션과 분산

저자들은 먼저 바닐라 소프트맥스 어텐션과 선형 어텐션을 모두 포괄하는 일반화된 어텐션(Generalized Attention) 프레임워크를 구축합니다. 이들은 연속적인 정규화 함수 $\phi$에 의해 어텐션 계수가 결정되는 $\Phi$-정규화 어텐션 메커니즘을 정의합니다.

핵심 이론적 결과 (분산 정리):
본 논문은 임의의 연속 정규화 함수 $\phi$에 대하여, 토큰 수 $n$이 무한대로 갈 때 어텐션 계수가 필연적으로 **분산(disperse)**된다는 것을 증명합니다. 즉, 특정 키에 할당되는 어텐션 가중치는 0에 수렴하며($\Theta(1/n)$), 결과적으로 모든 토큰에 동일한 가중치를 할당하게 됩니다.

• 이는 바닐라 소프트맥스($\phi(x) = \exp(x)$)와 선형 어텐션($\phi(x) = x$) 모두에 해당합니다.
• 이 증명은 입력 특징 공간의 컴팩트성(compactness)과 관련 함수의 연속성에 기반합니다.
• 저자들은 차분 어텐션(differential attention)이 노이즈를 완화하려고 시도하지만, 이론적으로는 여전히 분산 문제($\Theta(1/n)$)에 직면해 있음을 언급합니다.

2.2 제안된 솔루션: SEMA

분산 특성에 대한 동기 부여를 바탕으로, 저자들은 **SEMA (Scalable and Efficient Mamba-like Attention)**를 설계했습니다. SEMA의 핵심 철학은 전역 어텐션의 점근적 행동(분산)을 수용하면서도 국소적인 집중 능력을 보존하는 것입니다.

SEMA는 두 가지 구성 요소를 결합합니다:

1. 토큰 국소화 (Window Attention): 고주파 특징에 대한 바닐라 어텐션의 "집중" 속성을 유지하기 위해, SEMA는 국소화된 윈도우 어텐션 메커니즘을 사용합니다. 이는 각 쿼리가 사전에 정의된 윈도우 내의 고정된 수의 키와만 상호작용하도록 제한하여, 해당 로컬 컨텍스트 내에서의 분산을 방지합니다.
2. 이론적 평균화 (Global Approximation): 이차적 복잡도를 초래하거나 전체 선형 어텐션의 분산 문제를 겪지 않으면서 전역적인 어텐션 측면을 포착하기 위해, SEMA는 산술 평균(arithmetic averaging) 항을 포함합니다.
   • 점근적 분석을 바탕으로, 저자들은 매우 큰 $n$의 극한에서 전체 어텐션의 전역적 구성 요소가 모든 값 벡터의 단순 평균($\frac{1}{n}\sum v_j$)으로 근사될 수 있다고 주장합니다.
   • 이는 "균질한 혼합(homogeneous mixing)" 또는 저역 통과 필터(low-pass filter) 역할을 하여 전역 정보를 효율적으로 집계합니다.

수식:
윈도우 크기가 $w$인 쿼리 $q$, 키 $K$, 값 $V$에 대한 SEMA 어텐션은 다음과 같이 정의됩니다:
$$\text{SEMA}(Q, K, V) = A_{\Phi}(Q, K, V, w) + \frac{1}{n} \sum_{j=1}^{n} v_j$$
여기서 $A_{\Phi}$는 윈도우 기반 어텐션(소프트맥스 또는 선형 변형 사용)이며, 두 번째 항은 브로드캐스트된 전역 평균입니다.

2.3 아키텍처

SEMA는 Mamba 스타일의 매크로 아키텍처(MILA와 유사)에 통합됩니다. 저자들은 MILA의 선형 어텐션 블록을 SEMA 블록으로 교체했습니다. MILA와 달리, SEMA는 성능을 위해 필요하다고 판단한 값 $V$에 대한 선형 투영(linear projection)을 포함하며, 이는 파라미터 수를 약간 증가시키지만 효율성은 유지합니다.

3. 주요 기여

1. 분산의 이론적 증명: 본 논문은 특정 정규화 함수를 사용하는 것과 관계없이, 토큰 수가 무한히 증가함에 따라 일반화된 어텐션(소프트맥스 및 선형 변형 포함)이 분산된다는 엄밀한 수학적 증명을 제공합니다.
2. SEMA 설계: 로컬 분산을 피하기 위한 토큰 국소화와 전역 어텐션의 측면을 이론적으로 근사하는 산술 평균을 활용하여, 선형 어텐션의 스케일 가능한 대안을 제시하는 새로운 어텐션 메커니즘을 제안합니다.
3. 실험적 검증: 표준 벤치마크, 특히 이미지 해상도가 높아질 때 SEMA가 최근의 Vision Mamba 모델(예: VMamba) 및 선형 어텐션 베이스라인(MILA)보다 우수한 성능을 보임을 입증했습니다.

4. 실험 결과

4.1 ImageNet-1K 분류

• 표준 해상도 (224²): SEMA-T는 유사한 파라미터 수(~26M)를 가진 MILA-T(83.3%) 및 DefMamba(83.5%)를 능가하는 **83.7%**의 Top-1 정확도를 달성했습니다.
• 확장성 (고해상도): 가장 유의미한 이득은 높은 해상도에서 나타납니다.
  • 1024²에서 VMamba-T는 미세 조정 없이 74.7%, 미세 조정 시 83.4%로 성능이 크게 떨어지는 반면, SEMA-T는 미세 조정 없이 75.4%, 미세 조정 시 **84.2%**를 유지하며 견고함을 보여줍니다.
  • 미세 조정이 없는 상태에서 SEMA와 VMamba 사이의 격차는 1024²에서 **7.3%**까지 벌어지며, 이는 SEMA의 우수한 토큰 수 증가 대응 능력을 나타냅니다.

4.2 다운스트림 태스크

• COCO (인스턴스 분할): Mask R-CNN을 사용하여, SEMA-T는 1x 및 3x 학습 스케줄 모두에서 모든 지표(AP, AP50, AP75)에 대해 SOTA(State-of-the-art) 또는 상위 2위의 성능을 기록하며 MILA-T 및 VMamba-T를 능가했습니다.
• ADE20K (시맨틱 분할): SEMA-T는 48.2 mIOU(단일 스케일)를 달성하여 VMamba-T(47.9)를 앞섰으며, 멀티 스케일 설정에서도 경쟁력 있는 성능을 유지했습니다.

4.3 효율성

• 추론 시간: SEMA는 뛰어난 런타임 확장성을 보여줍니다. 1024² 해상도에서 SEMA는 MILA보다 약 33% 빠릅니다. 2048² 해상도에서는 이 속도 향상이 더욱 두드러지며, 이는 평균화 구성 요소의 선형 복잡도 이점을 확인시켜 줍니다.

4.4 절제 연구 (Ablation Studies)

• 값 투영 (Value Projection): 값 투영 레이어를 제거하면 정확도가 약 0.8% 감소하여, 그 중요성을 입증했습니다.
• 평균화 구성 요소: 224² 해상도에서 평균화 항의 기여도는 미미했습니다(~0.2%). 그러나 해상도가 672²로 높아짐에 따라 그 기여도가 크게 증가하여(+2.2%, 미세 조정 없음), 토큰 수가 증가하고 분산 효과가 심화됨에 따라 평균화가 매우 중요하다는 이론적 주장을 검증했습니다.

5. 의의 및 주장

본 논문은 SEMA가 선형 어텐션의 분산 문제에 대한 이론적 근거를 갖춘 솔루션을 제공한다고 주장합니다. 전역 어텐션이 극한 상황에서 분산되는 것이 고유한 특성임을 증명함으로써, 저자들은 산술 평균을 유효한 전역 근사치로 사용하는 정당성을 확보했습니다.

SEMA의 의의는 다음과 같습니다:

• 효과적인 스케일링: 이미지 해상도(및 토큰 수)가 증가함에 따라 성능과 효율성을 유지하며, 이는 현재 Mamba 기반 비전 모델들의 알려진 약점입니다.
• 국소 및 전역의 균형: 윈도우 어텐션의 집중력과 평균화를 통한 전역 컨텍스트를 성공적으로 결합하여, 풀 어텐션의 이차적 비용과 순수 선형 어텐션의 "흐릿한" 집중 문제를 모두 피했습니다.
• 강력한 대안 제공: 특히 표준 이미지 분류보다 토큰 수가 훨씬 많은 경우(예: 의료 영상 등 대규모 입력이 필요한 시나리오)에 유용한, 확장 가능하고 효율적인 비전 아키텍처를 위한 강력한 대안이 됩니다.

저자들은 평균화 구성 요소가 표준 해상도에서는 미미할 수 있지만, 대규모 토큰 영역에서는 결정적인 "안전망"과 성능 향상을 제공하여 SEMA를 확장 가능한 비전 아키텍처를 위한 견고한 선택지로 만든다고 결론지었습니다.