활성화 서브스페이스 병목을 통한 SSM 해석 및 조정

초록

본 논문은 Mamba 계열 상태공간 모델(SSM)의 내부 동작을 활성화 서브스페이스 관점에서 분석하고, 정보 흐름을 제한하는 “병목”을 찾아낸다. 식별된 병목에 대해 테스트 시 스칼라 배율을 곱하는 간단한 스티어링을 적용하면, 5가지 SSM과 6가지 벤치마크에서 평균 8.27%의 성능 향상을 달성한다. 또한, 병목을 구조적으로 해소한 Stable‑Mamba를 설계·재학습하여 장기 컨텍스트 성능을 추가로 개선한다.

상세 분석

이 연구는 메커니즘 해석 기법을 SSM에 적용한 최초 사례 중 하나로, 특히 Mamba 모델의 숨겨진 상태를 Sparse Autoencoder(SAE)와 사전 학습된 사전(dictionary) 학습을 통해 고차원 서브스페이스로 분해한다. 저자들은 각 레이어의 파라미터를 Stochastic Parameter Decomposition(SPD)으로 분석해 엔트로피 급증 지점을 병목으로 정의하고, Δ‑parameter가 제어하는 시간‑스텝 업데이트와 강하게 상관된 “Delta‑Sensitive” 서브스페이스를 추출한다. 이러한 서브스페이스는 입력에 따라 활성화 변동이 큰데, 이는 모델이 새로운 토큰 정보를 충분히 반영하지 못하고 기존 상태에 과도하게 의존한다는 신호이다.

실험에서는 768개의 서브스페이스 중 668개가 높은 변동성을 보였으며, 이들을 0으로 차단했을 때 성능이 급격히 떨어지는 것을 확인함으로써 병목의 기능적 중요성을 입증했다. 이후 스티어링 단계에서는 해당 서브스페이스에 스칼라 팩터(주로 5배)를 곱해 활성화를 증폭시켰다. 이 간단한 연산은 추론 시에만 적용되며, 별도의 파인튜닝이나 데이터 재학습이 필요 없다는 점에서 실용성이 크다.

또한, 병목을 구조적으로 해소하기 위해 Stable‑Mamba를 설계했다. 여기서는 다중 시간 스케일(state dynamics), 희소 글로벌 컨텍스트 주입, 학습 가능한 게이팅, 적응형 압축 길이, 스케일된 잔차 연결 등 7가지 모듈을 추가했으며, 전체 파라미터 수는 256개만 증가했다. 이러한 변경은 모델의 메모리·연산 비용에 거의 영향을 주지 않으면서, 장기 컨텍스트(LRA, RULER, LongContext V2)에서 기존 Mamba 대비 10~15% 정도의 정확도 향상을 보였다.

전체적으로 논문은 (1) SSM 내부를 서브스페이스 단위로 정량·정성 분석, (2) 병목 서브스페이스를 식별·증폭하는 테스트‑타임 스티어링, (3) 최소 구조 변경으로 병목을 근본적으로 제거하는 두 단계 접근법을 제시한다. 이는 SSM이 기존 트랜스포머 대비 효율성을 유지하면서도, 해석 가능성과 조정 가능성을 확보할 수 있음을 증명한다.