Mamba2 상태 차원 절감을 위한 GHOST 프레임워크

초록

GHOST는 전방 통계만을 이용해 Mamba2의 상태 차원을 50%까지 줄이면서 퍼플렉시티를 1 정도만 상승시키는 구조적 프루닝 방법이다. 제어 이론의 controllability와 observability를 데이터 기반으로 추정해 “phantom state”와 “corporeal state”를 구분하고, 그룹 단위 임계값 적용으로 메모리 대역폭을 크게 절감한다.

상세 분석

본 논문은 Mamba2가 도입한 128 차원의 상태 공간이 메모리 대역폭을 포화시켜 추론 속도를 저해한다는 문제를 제기한다. 기존의 비구조적 스파스화는 가중치만을 희소하게 만들고, 활성화는 여전히 밀집돼 실제 메모리 절감 효과가 없으며, magnitude 기반 프루닝은 정적 가중치 크기와 실제 사용량 사이의 상관관계가 낮아 중요한 “phantom state”를 놓치고 비활성 “corporeal state”를 남긴다. 또한 gradient 기반 구조적 프루닝은 대규모 모델에서 메모리와 연산 비용이 prohibitive 하다.

GHOST는 이러한 한계를 극복하기 위해 두 가지 핵심 아이디어를 도입한다. 첫째, controllability를 hidden state의 공분산(즉, 상태가 입력에 의해 얼마나 활성화되는가)으로, observability를 출력 에너지 Hessian(즉, 현재 상태가 미래 출력에 미치는 영향)으로 근사한다. 이 두 지표의 곱을 각 상태 채널의 saliency score로 정의하고, 이를 시간과 배치에 걸쳐 평균해 안정적인 추정값을 만든다. 둘째, Mamba2의 Grouped Query Attention(GQA) 구조를 활용해 그룹별로 score를 풀링하고, 전체 그룹·채널 풀에서 목표 sparsity에 맞는 전역 임계값을 설정한다. 이렇게 하면 복잡한 동역학을 학습하는 그룹은 더 많은 상태를 보존하고, 중복된 그룹은 aggressive하게 차원을 축소한다.

수식적으로는 Hankel singular value와 유사한 형태의 S = √(∑_t P·Q) 를 사용해 score를 정규화하고, pruning mask M을 binary 형태로 생성한다. pruning 후에는 두 번째 forward pass를 통해 남은 채널에 대한 활성화를 재계산하고, 다음 레이어에 전달함으로써 distribution shift를 최소화한다. 알고리즘 복잡도는 각 레이어당 O(|D_cal|·L·G·K·P·N) 로, 기존 inference와 동일한 수준이며 추가 메모리는 O(G·N) 에 불과하다.

실험에서는 130M~2.7B 파라미터 모델에 대해 50% 상태 차원 감소, WikiText‑2 퍼플렉시티 1.0 상승, 그리고 H100 GPU에서 메모리 사용량을 절반 이하로 감소시킨 결과를 보고한다. Magnitude, Random, Taylor pruning 등과 비교했을 때 GHOST는 gradient 기반 방법과 동등하거나 더 나은 정확도를 유지하면서 훨씬 적은 연산·메모리 비용을 요구한다. 또한 다양한 시퀀스 길이와 OOD 데이터에서도 안정적인 성능을 보이며, 실제 배포 환경에서 bandwidth 제한을 크게 완화할 수 있음을 시사한다.