1. 문제 정의 (Problem Statement)

Mamba2 모델은 시계열 모델링 능력을 향상시키기 위해 상태 차원 (State Dimension, $N$) 을 기존 Mamba1 의 16 에서 128 로 8 배 확장했습니다. 그러나 이 확장은 추론 시 메모리 대역폭 (Memory Bandwidth) 포화와 캐시 지역성 저하를 초래하여 심각한 병목 현상을 야기합니다. 예를 들어, 13 억 파라미터 모델의 경우 재귀 상태 크기가 약 12MB 에서 100MB 로 급증합니다.

기존의 가지치기 (Pruning) 방법들은 이 문제를 해결하는 데 한계가 있습니다:

• 비구조적 (Unstructured) 가지치기: 가중치만 희소화할 뿐, 활성화 (Activation) 는 밀집된 상태로 남아 메모리 대역폭 절감 효과가 없습니다.
• 정적 크기 기반 (Static Magnitude) 가지치기: 고정된 가중치 노름 (Norm) 만을 기준으로 중요도를 판단하여, 실제 동적 에너지 (Dynamic Energy) 가 높은 '유령 상태 (Phantom States)'를 잘못 제거하거나, 가중치는 크지만 실제 활용도가 낮은 '구체적 상태 (Corporeal States)'를 유지하는 오류를 범합니다.
• 기반한 (Gradient-based) 방법: 2 차 최적화 (Taylor pruning 등) 는 높은 정확도를 제공하지만, 전체 계산 그래프에 대한 기울기 (Gradient) 와 헤시안 (Hessian) 을 계산해야 하므로 VRAM 사용량이 과도하고 (13 억 파라미터 모델 기준 약 45GB), 계산 비용이 매우 높습니다.

2. 방법론: GHOST (Methodology)

저자들은 GHOST (Grouped Hidden-state Output-aware Selection and Truncation) 라는 새로운 구조적 가지치기 프레임워크를 제안합니다. 이는 제어 이론의 균형 축소 (Balanced Truncation) 원리를 Mamba2 의 순방향 전달 (Forward-pass) 통계만 사용하여 근사화한 것입니다.

핵심 원리

GHOST 는 상태 채널의 중요도를 평가하기 위해 제어 가능성 (Controllability) 과 관측 가능성 (Observability) 을 결합합니다.

1. 제어 가능성 (Controllability): 입력 히스토리가 현재 상태에 얼마나 큰 영향을 미치는지 측정합니다.
   • 측정 지표: 은닉 상태 (Hidden State, $H_t$) 의 공분산 (Covariance) 또는 제곱합.
2. 관측 가능성 (Observability): 현재 상태가 미래 출력에 얼마나 큰 영향을 미치는지 측정합니다.
   • 측정 지표: 출력 에너지 ($y_{SSM}$) 에 대한 헤시안 (Hessian) 근사치. 구체적으로는 출력 행렬 $C'_t$ 의 제곱합 $(C'_t)^T C'_t$ 를 사용합니다.

알고리즘 흐름

• 점수 산정: 각 상태 채널 $i$ 에 대해 점수 $S_i = \sqrt{P \cdot Q}$ 를 계산합니다. 여기서 $P$ 는 제어 가능성 (은닉 상태 에너지), $Q$ 는 관측 가능성 (출력 영향력) 입니다. 이는 균형 축소에서의 Hankel 특이값 (Hankel Singular Values) 에 해당합니다.
• 그룹화 및 임계값 적용: Mamba2 의 그룹화된 쿼리 어텐션 (GQA) 구조를 활용하여, 각 레이어 내의 그룹 간 점수를 통합 (Pooling) 합니다. 목표 희소도 ($\kappa$) 에 따라 전역 임계값을 설정하여, 복잡한 동역학을 모델링하는 그룹은 상태를 더 많이 유지하고, 중복된 그룹은 과감히 가지치기합니다.
• 순차적 업데이트: 레이어별로 순차적으로 처리하며, 가지치기된 레이어의 활성화 값을 다음 레이어의 보정 (Calibration) 에 사용하여 분포 이동 (Distribution Shift) 을 완화합니다.
• 비용: 기울기 (Gradient) 나 헤시안 (Hessian) 계산이 필요 없으며, 순방향 전달 2 회만 수행하면 됩니다.

3. 주요 기여 (Key Contributions)

1. GHOST 프레임워크 제안: Mamba2 의 상태 차원 ($N$) 을 줄이기 위해 제어 가능성과 관측 가능성을 기반으로 한 데이터 기반 구조적 가지치기 방법을 최초로 도입했습니다.
2. 유령 상태 (Phantom States) 식별: 기존 크기 기반 가지치기가 놓치는 '가중치는 작지만 활동이 높은 상태'를 정확히 식별하고 보존하며, '가중치는 크지만 활동이 없는 상태'를 제거함으로써 시스템의 동적 특성을 유지합니다.
3. 효율성과 성능의 균형: 기울기 계산 없이 Taylor 가지치기 (2 차 방법) 와 유사한 성능을 내면서도, VRAM 사용량을 45GB 에서 15GB 수준으로 획기적으로 줄였습니다.

4. 실험 결과 (Results)

130M 에서 2.7B 파라미터까지 다양한 크기의 Mamba2 모델을 대상으로 실험을 수행했습니다.

• 희소성 vs. 성능 (Sparsity Sweep):
  • WikiText-2 에서 50% 상태 차원 감소 시, perplexity 는 약 1 포인트 증가 (13.17 $\to$ 14.23) 하는 선형적인 성능 저하를 보였습니다.
  • 70% 이상의 고희소성 영역에서도 GHOST 는 안정적인 성능을 유지한 반면, Taylor 가지치기는 70% 이상에서 성능이 급격히 붕괴 (Perplexity 폭발) 했습니다.
• 시퀀스 길이 일반화 (Sequence Length Generalization):
  • 짧은 컨텍스트 (128 토큰) 로 보정하고 긴 컨텍스트 (2048 토큰) 로 평가했을 때, GHOST 는 밀집 모델과 유사한 성능 저하 곡선을 보였습니다. 반면, Taylor 는 긴 시퀀스에서 성능이 극도로 나빠졌습니다 (Perplexity 1613).
• 확장 법칙 (Scaling Laws):
  • 작은 모델 (130M, 370M) 에서 Taylor 는 기울기 불안정성으로 인해 완전히 실패한 반면, GHOST 는 모든 규모에서 견고한 성능을 발휘했습니다.
• Zero-Shot 및 OOD (Out-of-Distribution) 성능:
  • Lambada, PIQA, ARC 등 다양한 벤치마크에서 GHOST 는 Taylor 와 유사하거나 더 나은 성능을 보였습니다. 특히 코드 생성 (HumanEval) 과 수학 문제 (MMLU) 와 같은 OOD 작업에서도 기존 방법들보다 우수한 일반화 능력을 입증했습니다.
• 자원 효율성:
  • Taylor 가지치기는 40GB VRAM 을 초과하는 45GB 이상의 메모리가 필요했으나, GHOST 는 15GB 만으로도 실행 가능했습니다.

5. 의의 및 결론 (Significance)

이 논문은 대규모 Mamba2 모델의 추론 병목 현상을 해결하기 위한 실용적인 솔루션을 제시합니다.

• 이론적 통찰: 단순한 가중치 크기가 아닌, 시스템의 동적 에너지 (Dynamic Energy) 와 입력 - 출력 간 관계를 기반으로 중요도를 판단함으로써, 기존 가지치기 방법의 '맹점 (Blind Spot)'을 해결했습니다.
• 실용적 가치: 기울기 계산 없이도 고품질의 모델 압축이 가능하므로, 제한된 하드웨어 (예: 소비자용 GPU) 에서도 대규모 SSM(State Space Model) 을 배포할 수 있는 길을 열었습니다.
• 미래 지향성: 제어 이론과 딥러닝 압축을 결합한 접근법은 다른 아키텍처로 확장될 수 있는 가능성을 보여주며, 메모리 대역폭 효율성을 극대화하여 Mamba2 의 실제 적용을 가속화할 것으로 기대됩니다.

결론적으로, GHOST 는 Mamba2 의 상태 차원을 줄이면서도 모델의 핵심 동역학을 보존하는 가장 효율적이고 강력한 가지치기 방법으로 입증되었습니다.