기술 요약: 트랜스포머 및 상태 공간 모델(SSM)의 메커니즘적 평가

문제 정의

상태 공간 모델(SSM)은 언어 모델링을 위한 효율적인 하위 이차(sub-quadratic) 대안을 제공하지만, 표준 벤치마크가 포착하지 못하는 결함, 즉 컨텍스트로부터 정보를 복사하는 검색 작업에서의 미흡함을 보입니다. 이전 연구들은 이러한 성능 격차를 식별하기 위해 연관 회상(Associative Recall, AR)과 같은 합성 작업을 활용해 왔으나, 행동 지표(예: 작업 정확도)에 의존하는 것은 특정 아키텍처가 왜 성공하거나 실패하는지에 대한 통찰을 거의 제공하지 못합니다. 핵심 문제는 서로 다른 아키텍처가 어떻게 인컨텍스트 검색을 구현하는지, 그리고 이들이 견고하고 일반화 가능한 솔루션(예: 유도/induction)을 학습하는지 아니면 취약하고 위치 의존적인 휴리스틱에 의존하는지에 대한 메커니즘적 이해가 부족하다는 점입니다.

방법론

저자들은 모델이 검색 작업을 해결하는 방식을 분석하기 위해 인과적 개입(causal interventions)을 사용하는 메커니즘적 평가 프레임워크를 제안합니다.

1. 작업:

   • 연관 회상 (Associative Recall, AR): 모델에게 일련의 키-값(key-value) 쌍을 제시한 후 쿼리 키를 주어, 그에 대응하는 값을 출력하도록 하는 합성 작업입니다.
   • 연관 트리콜 (Associative Treecall, ATR): 확률적 문맥 자유 문법(PCFG)을 통해 생성된 새로운 계층적 검색 작업입니다. 키와 값이 인접한 AR과 달리, ATR은 잠재적인 트리 구조를 기반으로 비인접 토큰 간의 연관성을 요구하며, 이는 실제 자연어 검색을 더 밀접하게 모방합니다.

2. 메커니즘 지표 (교환 개입, Interchange Interventions):

   • 저자들은 모델 내부 구조에 대한 교환 개입(Geiger et al., 2021; 2024)을 사용합니다.
   • 절차: 특정 입력 토큰(예: 키)을 오염시켜 성능을 저하시킵니다. 그 후 특정 레이어 내의 특정 히든 스테이트(키, 값 또는 쿼리 위치)에서 개입하여, 오염된 표현을 대조적인 입력으로부터 가져온 "깨끗한" 표현으로 교체합니다.
   • 지표: 이 개입이 정답의 가능성을 얼마나 복구하는지 측정합니다. 값(value) 토큰에서의 높은 복구율은 유도(induction) 메커니즘(나중에 검색하기 위해 값에 연관을 저장함)을 나타냅니다. 쿼리(query) 또는 **키(key)**에서의 높은 복구율은 직접 검색(direct retrieval)(쿼리 또는 특정 레이어에서 즉시 연관을 계산함)을 나타냅니다.

3. 평가된 아키텍처:

   • 트랜스포머 (Transformer - Attention)
   • SSM: Mamba, DeltaNet, H3, Hyena
   • 하이브리드 모델: Based, BaseConv, DeltaNet

4. 절제 연구 (Ablation Studies):

   • Mamba와 Based에서 연관 형성에 필수적인지 확인하기 위해 단기 컨볼루션(short convolutions)(커널 크기 $d_{conv}$ 변화)을 체계적으로 제거하거나 수정합니다.

주요 기여 및 결과

1. 유사한 정확도에 대한 상이한 메커니즘

행동 지표는 트랜스포머, Based, Mamba가 AR에서 높은 정확도를 달로는 반면, 다른 SSM(H3, Hyena)은 실패함을 보여줍니다. 그러나 메커니즘 분석은 근본적인 차이를 드러냅니다:

• 트랜스포머 및 Based: 유도 메커니즘을 학습합니다. 이들은 순전파 과정에서 값(value) 토큰에 키-값 연관을 저장하여, 이후 레이어가 답을 검색할 수 있도록 합니다. 이는 견고한 2-레이어 솔루션입니다.
• Mamba 및 DeltaNet: 주로 직접 검색을 사용합니다. 이들은 (종종 레이어 0 또는 1에서) 값 토큰에 연관을 명시적으로 저장하지 않고 쿼리 토큰에서 연관을 계산합니다.
• H3 및 Hyena: 이러한 조건 하에서 작업을 효과적으로 학습하는 데 실패합니다.

2. 단기 컨볼루션의 결정적 역할

저자들은 단기 컨볼루션이 Mamba와 Based에서 연관을 가능하게 하는 핵심 요소임을 식별했습니다:

• Mamba에서: 단기 컨볼루션은 AR에 필수적입니다. 컨볼루션 커널 크기를 1로 줄이거나 제거하면, Mamba의 AR 성능은 거의 0에 가깝게 붕괴됩니다.
• Mamba의 메커니즘: Mamba가 오로지 SSM 상태에만 의존한다는 가정과 달리, 저자들은 Mamba 내부의 단기 컨볼루션이 유도와 유사한 연관 단계(키 정보를 값 또는 다음 키로 이동)를 수행하는 반면, SSM 구성 요소는 검색을 처리한다는 것을 발견했습니다.
• Based에서: 단기 컨볼루션을 암시적 장기 컨볼루션으로 대체하면 성능이 크게 저하되며, 이는 이 특정 작업에 대한 짧은 수용 영역의 필요성을 확인해 줍니다.

3. 계층적 검색(ATR)으로의 일반화

저자들은 메커니즘이 인접한 토큰을 넘어 일반화되는지 테스트하기 위해 ATR을 도입했습니다:

• 행동 결과: Mamba는 단기 컨볼루션 없이도 다른 모든 아키텍처보다 우수하거나 대등한 성능을 ATR에서 보여줍니다.
• 메커니즘의 변화: ATR에서 단기 컨볼루션이 없는 Mamba는 트랜스포머와 유사하게 (값에 저장하는) 2-레이어 유도 메커니즘을 성공적으로 학습합니다.
• 시사점: 이는 Mamba가 유도 능력을 갖추고 있으나, 더 단순한 인접 작업(AR)에서는 단기 컨볼루션을 통해 직접 검색 방식으로 기본 설정된다는 것을 시사합니다. 단기 컨볼루션에 대한 의존성은 작업에 따라 달라집니다. 즉, 이들은 AR에는 필수적이지만, 더 일반적인 유도 메커니즘으로 폴백(fallback)할 수 있는 ATR에서는 반드시 필요한 것은 아닙니다.

의의 및 주장

본 논문은 행동적 정확도만으로는 오해의 소지가 있기 때문에, 아키텍처의 역량을 이해하기 위해 메커니즘적 평가가 필수적이라고 주장합니다.

• 실질적 차이: 유사한 정확도를 가진 아키텍처들(예: 트랜스포머 vs Mamba)은 근본적으로 다른 내부 메커니즘을 사용할 수 있습니다. 트랜스포머와 Based는 일반화 가능한 유도 메커니즘을 학습하는 반면, Mamba의 AR 성공은 로컬하게 유도를 구현하는 특정 아키텍처 구성 요소(단기 컨볼루션)에 크게 의존합니다.
• 작업 의존성: 학습된 메커니즘은 고정되어 있지 않습니다. Mamba는 작업 구조(인접 vs 계층)에 따라 전략을 조정합니다.
• 설계 시사점: 연구 결과는 단기 컨볼루션이 SSM에서 연관 메모리를 가능하게 하는 핵심 동력임을 시사하지만, 그것의 부재가 (ATR과 같은) 작업 구조가 요구할 경우 유도 메커니즘 학습을 가로막는 것은 아님을 보여줍니다.
• 향후 방향: 저자들은 직관과 경험적 벤치마킹을 넘어 모델이 작업을 어떻게 해결하는지 설명하기 위해, 메커니즘적 지표가 아키텍처 분석의 표준 도구로서 채택되어 AI 발전의 "블랙박스"를 열어야 한다고 주장합니다.

결론적으로, 본 논문은 합성 작업이 아키텍처 혁신을 이끌어왔지만, 이를 단순히 행동적 벤치마크로만 취급하는 것은 귀중한 신호를 버리는 것이라고 말합니다. 인과적 개입을 적용함으로써, 저자들은 SSM의 검색 작업에 대한 "성공"이 SSM 공식 자체에 내재된 보편적이고 견고한 유도 능력보다는 특정 작업 민감적 구성 요소에 의존하고 있음을 밝혀냈습니다.