An Exact Instrument for State Usage in Selective State-Space Models, and the Input-Driven Migration It Reveals

Mamba 모델을 거대하고 똑똑한 오케스트라라고 상상해 보세요. 이 오케스트라의 모든 레이어 내부에는 16개의 작은 단음 악기(모드라고 불림)로 구성된 뱅크가 있습니다. 과거의 방식에서는 지휘자(모델)가 음악이 어떻게 흘러가든 상관없이, 예를 들어 가장 좋은 8개의 악기를 고정된 세트로 정해놓고 계속 사용한다고 가정했습니다.

하지만 이 논문은 그 커튼을 걷어내고 충격적인 비밀을 밝혀냈습니다. 지휘자는 고정된 세트를 선택하지 않습니다.

대신, 지휘자는 즉흥 연주의 명수입니다. 모델이 읽는 매 단어(또는 "토큰")마다, 지휘자는 멜로디를 이끌어가는 데 실제로 필요한 8개의 악기가 무엇인지 즉각적으로 다시 결정합니다. 때로는 플루트가 필요하고, 때로는 바이올린이, 때로는 드럼이 필요합니다. "중요한" 악기들은 입력값에 따라 **이동(migrate)**합니다. 만약 오케스트라에게 고정된 악기 세트를 고수하도록 강요한다면(즉, "정적"인 선택을 한다면), 당신은 행진곡단용 악보를 들고 재즈 솔로를 연주하라고 요구하는 것과 같습니다. 작동은 하겠지만, 실제 연주에 비하면 형편없이 들릴 것입니다.

마법의 도구: "정확한 악기(Exact Instrument)"

저자들은 이를 어떻게 알아냈을까요? 그들은 수학적인 "정확한 악기"를 구축했습니다.

오케스트라의 내부 구조가 특수한 종류의 "대각선(diagonal)" 설정(악기들이 서로 간섭하지 않는 구조)이기 때문에, 저자들은 출력을 각 악기의 기여도로 완벽하게 분해할 수 있었습니다. 그들은 "그람 텐서(Gram tensor)"(매우 정밀한 점수판이라고 생각하면 됩니다)를 만들었는데, 이것은 특정 악기 그룹을 제외했을 때 발생하는 오차가 정확히 얼마인지 알려줍니다.

그들은 이 도구를 실제 모델과 대조 테스트했고, 상대 오차가 2.3 × 10⁻⁷ 수준임을 발견했습니다. 이는 지구에서 달까지의 거리를 측정하면서 오차를 인간의 머리카락 굵기보다 작게 만드는 것과 같습니다. 이것은 추정치가 아니라 정밀한 측정값입니다.

핵심 발견: "이동 격차(Migration Gap)"

이 도구를 사용하여 저자들은 아주 작은 모델(130M 파라미터)부터 거대한 모델(배포된 Falcon-Mamba와 같은 7B 파라미터)에 이르기까지 다양한 모델을 조사했습니다.

그 결과, 가장 활발한 레이어에서 고정된 악기 세트를 사용하는 것이 변화하는 세트를 사용하는 것보다 두 배 더 많은 오차를 발생시킨다는 것을 발견했습니다.

• 통계: 가장 큰 영향을 받는 레이어에서 "이동 격차"(고정된 세트와 변화하는 세트 사이의 오차 비율)는 0.44에서 0.57 사이였습니다.
• 의미: 모델이 특정 순간마다 최적의 악기를 선택하도록 허용하면(즉, "입력 스케줄링 오라클"을 사용하면), 한 번 정해진 목록을 고수할 때보다 오차를 절반으로 줄일 수 있습니다.

이 현상은 테스트한 모든 모델에서 나타났습니다: Mamba-1 제품군, 7B Falcon-Mamba, 그리고 Mamba-2까지도 마찬가지였습니다.

무엇이 이동(Migration)을 일으키는가? (이유)

저자들은 질문했습니다. 모델의 어느 부분이 이 모든 전환을 주도하는가?
Mamba 레이어에는 세 가지 주요 신호가 있습니다:

1. 쓰기 맵 (Write Map, $B_t$): 어떤 악기에 입력 신호를 보낼지 결정합니다.
2. 읽기 출력 (Readout, $C_t$): 어떤 악기가 들리게 할지 결정합니다.
3. 타임스텝 ($\Delta$): 흔히 "선택성(selectivity)" 조절 노브라고 생각되는 부분입니다.

그들은 "신호 동결(frozen-signal)" 실험을 수행했습니다. 각 신호를 하나씩 평균값으로 고정하여, 이동 현상이 멈추는지 확인했습니다.

• 결과: **쓰기 맵 ($B_t$)**을 고정했을 때, 이동 현상이 사라졌습니다. 모델은 더 이상 악기를 교체하지 않았습니다.
• 놀라운 점: **타임스텝 ($\Delta$)**을 고정했을 때, 이동 현상은 그대로 유지되었습니다.

결론: 많은 사람이 선택성의 핵심이라고 생각했던 타임스텝($\Delta$) 신호는 이동 신호에 거의 영향을 주지 않았습니다. 진짜 주인공은 **쓰기 맵 ($B_t$)**입니다. 쓰기 맵은 토큰마다 어떤 악기가 연주될지를 결정하는 문지기 역할을 합니다.

결과: 이를 활용할 수 있는가?

저자들은 이 지식을 사용하여 모델을 가지치기(pruning)하여 용량을 줄이는 시도를 했습니다.

• 정적 가지치기 (Static Pruning): 평균적인 활동량을 기준으로 최적의 악기 8개를 골라 영구적으로 유지합니다.
• 입력 스케줄링 가지치기 (Input-Scheduled Pruning): 현재 문장을 살펴보고, 지금 이 순간 활성화된 8개의 악기를 측정하여 그 악기들만 남깁니다.

결과:
상태 예산의 절반(16개 모드 중 8개만 유지)을 사용할 때, 입력 스케줄링 방식은 원본 모델의 가공되지 않은 정확도와 대등하거나(경우에 따라서는 약간 더 나은) 성능을 보였습니다.

• 130M 모델의 경우, 스케줄링 방식의 퍼플렉시티(perplexity)는 11.84였고, 가공되지 않은 모델은 12.38이었습니다.
• 7B Falcon-Mamba의 경우, 스케줄링 방식은 4.44로, 가공되지 않은 모델의 4.48을 앞질렀습니다.

하지만 중요한 주의사항이 있습니다: 논문은 이 "스케줄링" 방식이 **두 번의 패스(two-pass oracle)**를 거치는 방식임을 명시합니다. 즉, 어떤 악기를 남길지 결정하기 위해 전체 윈도우를 한 번 읽은 다음, 두 번째 패스를 통해 출력을 생성합니다. 이는 실제 배포 환경에서 계산량이나 메모리를 절약해주지는 못한다는 의미입니다(데이터를 두 번 읽어야 하기 때문입니다).

저자들은 이 결과가 **실현 가능한 여지(realizable headroom)**를 보여주는 것이지, 실제 배포 시의 효율성 승리를 의미하는 것은 아니라고 설명합니다. 이는 만약 우리가 첫 번째 패스 없이도 적절한 악기를 예측할 수 있는 저렴하고 빠른 예측기를 만들 수 있다면, 아주 작고 효율적인 모델이 거대 모델과 대등해질 수 있는 잠재력이 존재함을 증명하는 것입니다. 현재의 방식은 그 가능성의 천장을 보여주는 것뿐입니다.

이 논문이 배제한 것들

논문은 무엇이 효과가 없는지 매우 명확하게 밝히고 있습니다:

• 정적 순위 지정 (Static Rankings): 평균 활동량이나 "Hankel 에너지"(GHOST나 LAST 같은 방식)를 기반으로 고정된 모드 세트를 뽑는 방법은 훨씬 성능이 떨어집니다. 이들은 움직이는 목표를 따라잡지 못합니다.
• 타임스텝 활동성 (Timestep Activity): 타임스텝의 활성도를 기준으로 가지치기를 하는 것은 막다른 길입니다. 타임스텝은 이동을 유도하지 않기 때문입니다.
• 단순한 예측기 (Simple Predictors): 문장의 앞부분 절반만 사용하거나 "도메인"(코드냐 산문이냐 등)을 기반으로 마스크를 예측하는 시도를 했습니다. 이러한 간단한 트릭들은 잠재적 이득의 **2~6%**만을 회복할 수 있었습니다. "중요한" 악기 세트는 너무 빠르게 변하기 때문에(수백 토큰 내에서), 최대의 이득을 얻으려면 점수를 매기는 특정 토큰들을 직접 측정해야 합니다.

결-론

이 논문은 훈련된 선택적 상태 공간 모델(selective state-space models)이 입력에 따라 내부 자원을 끊임없이 재할당하는 역동적이고 살아있는 존재임을 입증합니다. "쓰기 맵"이 바로 이 이동의 지휘자입니다. 비록 우리가 아직 "두 번의 패스"라는 비용 없이 실시간으로 악기를 교체하는 모델을 배포할 수는 없지만, 이 연구는 우리에게 영역의 정확한 지도를 제공합니다. 이는 현재의 "정적" 가지치기 방식들이 엄청난 성능을 낭비하고 있으며, 앞으로의 길은 이러한 이동을 즉석에서 예측할 수 있는 스케줄러를 구축하는 데 있음을 보여줍니다.