Attention to Mamba: A Recipe for Cross-Architecture Distillation

1. 문제 정의 (Problem)

최근 자연어 처리 (NLP) 분야에서 State Space Models (SSM), 특히 Mamba는 Transformer 의 Attention 메커니즘 대비 선형 (linear) 시간 복잡도와 낮은 메모리 사용량으로 인해 주목받고 있습니다. 그러나 다음과 같은 한계가 존재합니다:

• 성능 격차: 대규모 모델에서 Mamba 는 기존 Transformer 의 성능 (Perplexity 및 하위 작업) 을 완전히 따라잡지 못하는 경우가 많습니다.
• 지식 활용의 어려움: 이미 방대한 데이터로 사전 학습된 Transformer 모델의 지식을 Mamba 로 직접 이전 (Distillation) 하려는 시도 (Naïve Distillation) 는 실패하거나 성능이 크게 저하되는 경향이 있습니다. 이는 두 아키텍처 간의 계산 패러다임 (Attention vs. SSM) 이 근본적으로 다르기 때문입니다.
• 기존 해결책의 한계: Attention 과 SSM 블록을 혼합하는 하이브리드 방식은 성능을 개선하지만, 순수 SSM 모델의 효율성 (선형 복잡도) 을 잃게 만듭니다.

핵심 목표: 사전 학습된 Transformer 의 지식을 보존하면서, 순수 Mamba 아키텍처로 효율적으로 이전할 수 있는 효과적인 증류 (Distillation) 레시피를 개발하는 것입니다.

2. 방법론 (Methodology)

저자들은 Transformer 의 Attention 을 Mamba 로 직접 변환하는 대신, **원칙적인 초기화 (Principled Initialization)**를 통해 두 아키텍처 간의 수학적 연결고리를 활용하는 2 단계 증류 전략을 제안합니다. 이를 HedgeMamba라고 명명합니다.

1 단계: Softmax Attention $\rightarrow$ Linear Attention (Hedgehog)

• 목표: 비선형인 Softmax Attention 을 선형 (Linear) 형태로 근사화합니다.
• 기법: Hedgehog 방법론 (Zhang et al., 2024) 을 활용합니다. Mercer 의 정리에 기반하여, Softmax 의 지수 함수를 특정 특징 맵 (Feature Map, $\phi$) 을 통한 내적 (Dot Product) 으로 근사합니다.
• 구현: 학습 가능한 단일 레이어 MLP 를 사용하여 Softmax 의 행동을 모방하는 특징 맵 $\phi_{MLP}$를 학습합니다. 이때 Teacher 모델의 Attention 출력과 Student 의 Linear Attention 출력을 코사인 유사도 (Cosine Similarity) 로 매칭하여 학습합니다.
• 의의: 이 단계는 Attention 의 핵심 정보를 선형 형태로 추출하여 Mamba 가 이해할 수 있는 형태로 변환하는 '다리' 역할을 합니다.

2 단계: Linear Attention $\rightarrow$ Mamba (HedgeMamba)

• 목표: 1 단계에서 학습된 Linear Attention 을 Mamba 의 초기값으로 사용하여 성능을 극대화합니다.
• 매핑 전략:
  • Linear Attention 의 $Q, K, V$ 행렬을 Mamba 의 SSM 파라미터 ($C, B, X$) 에 매핑합니다.
  • 특히, 1 단계에서 학습된 Hedgehog 특징 맵 ($\phi_{MLP}$) 을 적용하여 $K$와 $Q$를 변환합니다.
  • Mamba 의 상태 행렬 $\Lambda$는 단위 행렬 (Identity) 로 초기화하여 1 단계의 선형 Attention 동작을 유지하도록 합니다.
• 추가 구성 요소:
  • 정규화 (Normalization): Attention 점수의 정규화를 위해 SSM 출력에 정규화 인자를 추가합니다.
  • Gate 및 Conv: Mamba 고유의 게이트 (Gate) 브랜치와 짧은 컨볼루션 (Conv) 레이어를 도입하여 모델의 표현력 (Expressivity) 을 높입니다. 초기에는 이들을 항등 연산 (Identity) 으로 설정하여 1 단계 학습을 방해하지 않도록 합니다.
• 파인튜닝: 전체 모델 (임베딩 제외) 을 Cross-Entropy 손실 함수로 파인튜닝하여 최종 성능을 끌어올립니다.

3. 주요 기여 (Key Contributions)

1. 새로운 교차 아키텍처 증류 레시피 제안: Transformer 에서 Mamba 로의 지식 이전을 위해, Attention $\rightarrow$ Linear Attention $\rightarrow$ Mamba 의 2 단계 브리징 전략을 제안했습니다. 이는 단순한 직접 증류보다 훨씬 효과적입니다.
2. 원칙적인 초기화 기법: Mercer 의 정리와 커널 트릭을 활용하여 Transformer 의 지식을 Mamba 파라미터에 체계적으로 초기화하는 방법을 제시했습니다.
3. HedgeMamba 아키텍처: Hedgehog (선형 Attention) 과 Mamba 의 장점을 결합한 하이브리드 레이어를 설계하여, 순수 Mamba 모델이 Transformer 의 성능을 유지하도록 했습니다.
4. 광범위한 실험적 검증: 1B 파라미터 모델 (Pythia-1B) 을 대상으로 10B 토큰 데이터를 사용하여 아키텍처 구성 요소, 토큰 예산 할당, 스케일링 등에 대한 철저한 분석 (Ablation, Scaling, Sensitivity Study) 을 수행했습니다.

4. 실험 결과 (Results)

• Teacher 모델: Pythia-1B (사전 학습 토큰: 300B).
• Student 모델: HedgeMamba (증류 토큰: 10B, Teacher 의 약 2.7% 에 해당).
• Perplexity (PPL):
  • Teacher (Pythia-1B): 13.86
  • Student (HedgeMamba): 14.11
  • 기존 Baseline (Hedgehog 단독): 14.89
  • 결과: 제안한 방법은 Teacher 의 성능을 매우 근접하게 유지하며, 기존 선형 Attention 기반 증류보다 우월한 성능을 보입니다.
• 하위 작업 (Downstream Tasks): ARC, BoolQ, Lambada 등 다양한 벤치마크에서 Teacher 모델과 유사한 성능을 달성했습니다.
• Ablation Study:
  • 토큰 할당: 1 단계 (Linear Attention 학습) 에 10%, 2 단계 (Mamba 파인튜닝) 에 90% 의 토큰을 할당하는 것이 최적 (10/90 split) 임을 확인했습니다. 1 단계가 없으면 성능이 급격히 떨어집니다.
  • 구성 요소: Gate 브랜치와 Conv 레이어를 추가할수록 성능이 향상됨을 확인했습니다.

5. 의의 및 결론 (Significance)

이 논문은 Mamba 와 같은 차세대 SSM 아키텍처가 기존 Transformer 의 지식을 효과적으로 흡수하여, 재학습 없이도 높은 성능과 효율성을 동시에 달성할 수 있음을 증명했습니다.

• 효율성: Transformer 의 $O(L^2)$ 복잡도를 Mamba 의 $O(L)$ 복잡도로 낮추면서도, Teacher 모델의 성능을 거의 잃지 않습니다.
• 실용성: 이미 존재하는 대규모 사전 학습 Transformer 모델을 활용하여, 더 가볍고 빠른 SSM 모델을 빠르게 구축할 수 있는 구체적인 방법론을 제공합니다.
• 미래 전망: 단순한 모델 변환을 넘어, 서로 다른 신경망 아키텍처 간의 지식 이전을 위한 수학적 연결고리를 찾는 새로운 패러다임을 제시했습니다.

요약하자면, 이 연구는 **"Attention 을 Mamba 로 변환할 때, 중간에 '선형화된 Attention'을 거치는 2 단계 증류 과정과 원칙적인 초기화가 핵심이다"**라는 통찰을 제공하며, Mamba 의 실용적 도입을 가속화할 수 있는 중요한 기여를 했습니다.