🤖 machine learning

Interpreting and Steering State-Space Models via Activation Subspace Bottlenecks

본 논문은 테스트 시 활성화 부분공간 병목 현상을 식별하여 작업별 튜닝 없이 다양한 벤치마크에서 성능을 크게 향상시키고, 동시에 개선된 장문맥 기능을 갖춘 재학습된 "Stable-Mamba" 아키텍처의 생성을 가능하게 함으로써 상태-공간 모델의 해석 가능성과 제어 가능성을 진전시킵니다.

원저자: Vamshi Sunku Mohan, Kaustubh Gupta, Aneesha Das, Chandan Singh

게시일 2026-05-22

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Vamshi Sunku Mohan, Kaustubh Gupta, Aneesha Das, Chandan Singh

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

State-Space Model(SSM)인 인기 있는 Mamba를 초효율적인 택배 서비스로 상상해 보세요. 모든 단어를 읽을 때마다决策을 내리기 위해 과거에 읽은 모든 단어를 담은 거대하고 무거워지는 배낭을 운반하는 기존 모델(Transformer)과 달리, Mamba 는 단일하고 간결한"노트북"을 운반합니다. 이 노트북은 읽는 과정에서 업데이트되어 과거의 무게에 짓눌리지 않고 정보를 놀라울 정도로 빠르게 처리할 수 있게 합니다.

하지만 이 논문 연구자들은 이 택배 서비스가 빠르기는 하지만 뇌에 교통 체증이 있다는 사실을 발견했습니다.

문제: "병목 현상"

저자들은 정보가 Mamba 의 레이어 (기차 노선의 여러 역으로 생각하세요) 를 통과할 때, 모든 것이 20 번째 레이어 주변의 특정하고 좁은 역으로 몰려든다는 사실을 발견했습니다.

비유: 10 차선 고속도로의 교통이 갑자기 단일 차선 터널로 강제되는 상황을 상상해 보세요. 차 (데이터) 가 빠르게 움직이고 있더라도 서로 으깨지며 밀집하게 됩니다. 터널이 너무 좁아 택배가 중요한 세부 사항을 잊어버리거나 혼란에 빠지게 만들며, 특히"패키지"(텍스트) 가 매우 길 때 두드러집니다.
증거: 연구자들은 모델 내부를 살펴보기 위해 특별한"X 선"도구 (확률적 매개변수 분해) 를 사용했습니다. 그들은 이 특정 레이어에서"엔트로피"(혼란 또는 혼동의 척도) 가 급격히 치솟는 것을 목격했습니다. 마치 터널 바로 앞에서 차량이 대량으로 정차하는 것을 보는 것과 같습니다.

해결책: 교통"조종"

전체 택배 시스템을 재건하는 대신, 팀은 **사후 조종 (Post-hoc Steering)**이라는 교묘하고 비침습적인 해결책을 시도했습니다.

비유: 가속 페달이 걸린 차를 운전한다고 상상해 보세요. 엔진을 교체하는 대신, 딱 필요한 순간에 페달을 막대기로 살짝 밀어 조금 더 힘을 실어주는 것입니다.
실행 방법: 그들은 교통 체증을 일으키는 특정"부분 공간"(내부 경로) 을 식별했습니다. 그런 다음 모델이 사고하는 과정에서 해당 경로의 활동을 특정 숫자 (예: 5 배) 로 곱했습니다.
결과: 이 작은 밀어냄이 교통 체증을 해소했습니다. 모델을 재학습시키거나 새로운 교훈을 가르칠 필요가 없었습니다. 내부 신호를"조종"하기만 했을 뿐, 질문 답변이나 헤이 stack 에서 바늘 찾기 등 6 가지 다른 작업에서 모델의 성능을 평균 **8.27%**만큼 향상시켰습니다. 이는 7 가지 다른 모델 버전에서 작동하여 교통 체증이 설계의 보편적 결함임을 입증했습니다.

장기적 해결책:"Stable-Mamba"

"조종"이 임시 교통 경찰처럼 작동하는 동안, 연구자들은 아키텍처 수준에서 문제를 해결하기 위해 Stable-Mamba라는 모델의 새로운 버전을 구축하기도 했습니다.

비유: 차를 살짝 밀어주는 대신, 고속도로 자체를 재설계한 것입니다.
- 다중 시간 척도: 모든 교통에 하나의 속도를 적용하는 대신, 빠르고 중간, 느린 차선을 추가하여 짧은 기억과 긴 기억이 공존할 수 있게 했습니다.
- 전역 주입: 택배가 큰 그림을 잃지 않도록, 전체 텍스트의 요약을 지역 교통에 가끔씩 떨어뜨리는"헬리콥터 시점"을 추가했습니다.
- 개선된 게이트: 다른 모든 것을 차단하지 않고 올바른 정보를 통과시키는 더 똑똑한 신호등을 설치했습니다.
결과: 이 새로운 모델인 Stable-Mamba는 처음부터 학습되었습니다. 이는 단순히 교통 체증을 해결한 것이 아니라 정보 흐름 전체를 매끄럽게 만들었습니다. 특히 매우 긴 텍스트에서"조종"된 버전보다 더 나은 성능을 보였으며, 모델에 256 개의 새로운 매개변수라는 아주 작은 양의 추가 무게만 더했습니다.

왜 이것이 중요한가

이 논문은 특정 약점 (병목 현상) 을 찾아내어 이러한 복잡한 AI"블랙박스"를 이해할 수 있음을 보여줍니다. 모델이 어디서 어려움을 겪고 있는지 알면 다음과 같이 대응할 수 있습니다.

조종하기: 재학습 없이 테스트 시간에 빠르고 무료한 부스트를 제공합니다.
설계 수정하기: 처음부터 병목 현상이 없는 더 나은 버전을 구축합니다.

저자들은 이 접근 방식이 AI 를 더 투명하고 효율적으로 만들어, 이러한 강력한 모델이 붕괴 없이 긴 대화와 복잡한 작업을 처리할 수 있게 한다고 강조합니다. 이는 처음에 인기를 끌게 만든 속도와 효율성을 유지하면서 가능합니다.

기술적 요약: 활성화 부분공간 병목 현상을 통한 상태-공간 모델 해석 및 조향

1. 문제 제기

상태 - 공간 모델 (SSM), 특히 Mamba 아키텍처는 고정된 크기의 숨겨진 상태를 활용하여 성장하는 키 - 값 캐시 대신 선형 시간 훈련과 상수 시간 추론을 제공함으로써 언어 모델링을 위한 효율적인 대안으로 부상했습니다. 그러나 효율성에도 불구하고 SSM 은 해석 가능성과 조향성 (steerability) 에서 상당한 도전에 직면해 있습니다. 명시적인 어텐션 헤드와 뉴런 표현을 갖추고 있어 해석 가능성 연구가 광범위하게 진행된 트랜스포머와 달리, SSM 은 이러한 명시적 구성 요소가 결여된 재귀적 구조에 의존합니다. 이러한 불투명성은 문맥 내 학습 (in-context learning) 과 긴 문맥 검색 성능과 같은 실패 모드를 이해하는 것을 방해합니다. furthermore, 기존 해석 가능성 연구는 SSM 에 실질적인 개선을 거의 가져오지 못했으며, 현재 사용 중인 조향 방법들은 대부분 트랜스포머 아키텍처에 맞춰져 있습니다.

2. 방법론

저자들은 Mamba 모델 내의 **활성화 부분공간 병목 현상 (Activation Subspace Bottlenecks)**을 식별하는 데 초점을 맞춘 기계적 조사를 제안합니다. 방법론은 세 가지 주요 구성 요소로 이루어집니다:

A. 활성화 부분공간 병목 현상 식별

저자들은 기능적 병목 현상을 찾기 위해 다단계 파이프라인을 사용합니다:

희소 오토인코더 (SAEs): SAE 는 희소 잠재 표현을 사용하여 숨겨진 활성화를 재구성하도록 훈련되어, 지배적인 활성화 수준 특징을 식별하기 위한 압축 병목으로 작용합니다.
어텐션 매핑 근사: Mamba 에는 명시적인 어텐션 헤드가 없으므로, 저자들은 암시적 어텐션 매핑 절차를 사용하여 Mamba 믹서의 SSM 레이어로부터 어텐션 스타일 행렬 ( $A$ ) 을 구성합니다. 이를 통해 토큰 수준의 중요도 벡터를 유도하고 숨겨진 상태의 가중 합으로 활성화 부분공간을 정의할 수 있습니다.
확률적 파라미터 분해 (SPD): SPD 는 모델 파라미터를 희소하게 사용되는 벡터로 분해하는 데 사용됩니다. 엔트로피, 분산, 유효 랭크, 그리고 제거 후 KL 발산과 같은 지표들이 레이어 단위로 계산되어 병목 현상을 식별합니다.
델타 민감 부분공간: 저자들은 $\Delta$ 파라미터 (이산화된 시간 단계 및 입력 의존적 업데이트를 제어) 에 반응하여 높은 분산을 보이는 부분공간을 식별합니다. 이러한 부분공간은 재귀적 역학에 중요하다고 간주되어 조향 대상으로 선정됩니다.

B. 사후 조향 (Post-Hoc Steering)

병목 현상이 식별되면, 저자들은 미세 조정 (fine-tuning) 이 필요 없는 추론 시 개입을 도입합니다:

대상 선정: 파이프라인은 엔트로피 급증과 높은 제거 후 KL 발산을 보이는 특정 레이어 (예: Mamba-130M 의 20 번째 레이어) 를 식별합니다. 이러한 레이어 내에서 다음 토큰 예측 정확도에 미치는 영향이 가장 큰 델타 민감 부분공간이 제거 실험을 기준으로 선정됩니다.
개입: 식별된 부분공간의 활성화는 추론 중 스칼라 조향 인자 (예: 고영향 부분공간의 경우 5 배, 중간 영향의 경우 2 배) 로 곱해집니다. 이는 식별된 병목 현상을 통한 신호 흐름을 증폭시킵니다.

C. Stable-Mamba 아키텍처

식별된 병목 현상이 훈련의 산물이 아닌 구조적 한계임을 검증하기 위해, 저자들은 해석 가능성 통찰을 통합한 수정된 아키텍처인 Stable-Mamba를 제안합니다:

다중 시간 척도 상태 역학: 병렬 SSM 분기 (빠른, 중간, 느린) 를 통해 단거리, 중거리, 장거리 처리를 동시에 수행할 수 있습니다.
앙상블 출력: 적응형 시간 해상도를 가능하게 하기 위해 선형 출력 프로젝션 대신 가중치 앙상블을 사용합니다.
희소 전역 문맥 주입: 희소 어텐션이 지역 상태를 안정화하기 위해 전역 요약을 지역 상태에 주입합니다.
학습된 게이트 및 적응형 압축: 앙상블 게이트는 극단적인 희소성을 줄이고, 적응형 압축은 용량을 동적으로 조절합니다.
그래디언트 스케일링 및 스케일된 잔차: 이러한 수정 사항들은 깊은 SSM 체인을 안정화하고 그래디언트 흐름을 개선합니다.

3. 주요 결과

본 연구는 7 개의 SSM 과 6 개의 다양한 벤치마크 (TriviaQA, SQuAD, MuSiQue, IFEval, RULER, DROP 포함) 에 걸쳐 제안된 방법을 평가했습니다.

조향 성능: 사후 조향 개입은 작업별 튜닝 없이 7 개의 서로 다른 SSM 에서 평균 **8.27%**만큼 성능을 향상시켰습니다. 특히, "The Pile"의 샘플을 사용하여 식별된 개입은 다양한 벤치마크로 일반화되었습니다.
긴 문맥 향상: 긴 문맥 벤치마크 (RULER, Long Range Arena, LongBench v2) 에서 조향된 모델은 유의미한 개선을 보였습니다. 예를 들어, Stable-Mamba 는 Vanilla Mamba 에 비해 Long Range Arena 의 PathFinder 작업에서 35 퍼센트 포인트의 향상을 달성했습니다.
병목 현상 검증:
- 엔트로피 분석: Vanilla Mamba 는 20 번째 레이어에서 급격한 엔트로피 급증을 보여 라우팅 병목 현상을 나타냅니다. 조향과 Stable-Mamba 아키텍처 모두 이 엔트로피를 완화하여 정보 흐름을 복원합니다.
- 제거 연구: Vanilla Mamba 에서 20 번째 레이어를 제거하면 실제로 성능이 향상됩니다 (77.0% 정확도 대 63.5% 기준). 이는 해당 레이어가 제한적인 병목으로 작용함을 확인시켜 줍니다. 반면, 조향되거나 Stable 모델에서 동일한 레이어를 제거하면 성능이 저하되어 병목 현상이 완화되었음을 나타냅니다.
- 파라미터 민감도: 20 번째 레이어의 $dt\_proj.bias$ 파라미터는 훈련 중 사실상 동결되는 "마스터 게이트" 역할을 하며, 예측의 45% 를 제어하고 제거 시 치명적인 정보 손실 (KL 발산 813) 을 초래합니다.
효율성: Stable-Mamba 는 256 개의 파라미터만 추가합니다. 토큰당 FLOPs 는 약 2.8 배 증가하지만, 병렬화로 인해 Vanilla Mamba 에 비해 추론 지연 시간은 약 9% 만 증가하고 메모리 오버헤드는 무시할 수 있습니다.

4. 중요성 및 주장

본 논문은 기계적 해석 가능성이 SSM 에 대한 실질적인 성능 향상으로 이어질 수 있음을 입증한다고 주장합니다. 모델 실패를 특정 활성화 부분공간과 아키텍처 구성 요소로 국한시킴으로써, 저자들은 다음을 보여줍니다:

해석 가능성은 실행 가능합니다: 병목 현상을 식별하면 재훈련 없이 성능을 향상시키는 표적 개입 (조향) 이 가능합니다.
구조적 한계가 존재합니다: SSM 은 긴 문맥 추론과 정보 흐름을 방해하는 특정 아키텍처 제약 (예: 단일 시간 척도 처리, 동결된 게이트 파라미터) 으로 고통받습니다.
최소한의 수정이 큰 성과를 냅니다: 해석 가능성에 기반한 작은 아키텍처 변경 (Stable-Mamba) 은 이러한 한계를 해결하여 더 큰 모델과同等하거나 그 이상의 성능을 달성하면서도 효율성을 유지합니다.

저자들은 일반화에 대해 겸손한 태도를 유지하며, 현재 파이프라인은 Mamba 에 초점을 맞추고 있으며 동일한 접근 방식이 다른 아키텍처나 매우 큰 파라미터 크기로 일반화되는지는 명확하지 않다고 지적합니다. 또한 평가된 작업 집합은 표준적이며 복잡한 추론 작업으로의 일반화 가능성은 추가 조사가 필요하다고 인정합니다.