기술 요약: 태스크 구조가 시퀀스 모델의 계층별 상태 인코딩을 역전시킨다

문제 정의

시퀀스 모델에 대한 기계론적 해석 가능성(mechanistic interpretability) 연구는 종종 계층별 상태 인코딩을 고정된 아키텍처 특성으로 취급한다. 순환 모델(예: LSTM, GRU)과 상태 공간 모델(SSM, 예: Mamba)은 후기 계층에 읽기 가능한 상태를 집중시키는 반면, 어텐션 기반 모델(Transformer)은 깊이에 따라 상태 축적을 분산시킨다는 것이 지배적인 가설이다. 본 논문은 이러한 인코딩 전략이 안정적인 아키텍처 불변량(architectural invariant)이라는 관점에 이의를 제기한다. 저자들은 동일한 아키텍처라도 태스크의 계산 구조가 변하면 계층별 상태 인코딩 프로필이 역전된다고 주장한다. 또한, 기존 문헌은 이러한 패턴의 두 가지 잠재적 동인인 대수적 구조(특히 교환성)와 계산적 구조(접두사 업데이트 vs 스택 추적)를 혼동하는 경향이 있다. 저자들은 이 두 요인을 분리하여, 상태 인코딩이 아키텍처만의 고유한 속성인지 아니면 아키텍처와 태스크 사이의 상호작용 결과인지를 규명하고자 한다.

방법론

본 연구는 Transformer, Mamba, Mamba-2, LSTM, GRU의 다섯 가지 아키텍처 패밀리를 대상으로 통제된 실험 세트를 적용한다.

태스크:
저자들은 세 가지 형식적 상태 추적 태스크와 한 가지 준실제(semi-real) 벤치마크를 활용한다:

1. Parity (패리티): 교환 가능한 단일 비트 접두사 업데이트 태스크.
2. $S_3$ Permutation Composition (순열 합성): 비교환적인 6개 클래스 접두사 업데이트 태스크. 이는 결정적인 변별 도구 역할을 한다. 만약 대수적 교환성이 인코딩 프로필을 결정한다면 $S_3$는 Dyck와 유사하게 동작해야 하며, 계산적 구조가 결정한다면 $S_3$는 Parity와 유사하게 동작해야 한다.
3. Dyck-$k$: 최대 깊이 10의 유계 깊이 스택 추적(괄호 매칭) 태스크.
4. Code-Depth: CodeSearchNet의 파이썬 코드에서 추출한 토큰 단위 중첩 깊이 태스크로, 준실제 검증용으로 사용된다.

실험 설정:

• 형식적 태스크: 4개 계층, $d_{model}=128$인 4개 계층 모델을 사용한다. 아키텍처는 가능한 경우 파라미터 수를 약 830K로 맞춘다. 훈련은 시퀀스 길이 $L=40$에서 수행되며, 평가는 $L \in \{40, 60, 80, 100, 120\}$에서 이루어진다.
• 사전 학습 규모: Mamba-130M(24개 계층) 및 Pythia-160M/410M(12/24개 계층)에 대해 미세 조정(fine-tuning)을 수행한다.
• 측정 지표:
  • 선형 프로빙(Linear Probing): 동결된 은닉 상태(frozen hidden states)에 로지스틱 회귀 분류기를 훈련시켜 각 계층에서의 상태 가독성을 측정한다.
  • 인과적 개입(Causal Interventions): 제로-어블레이션(zero-ablation, 프로브 방향을 투영하여 제거) 및 서브스페이스 어블레이션(subspace ablation, 상위 PCA 또는 특이 방향을 제거)을 통해 기능적 필요성을 테스트한다.
  • OOD 일반화: $L=40$에서 훈련된 어블레이션을 보지 못한 길이($L=80, 100, 120$)에 적용한다.
  • 컴포넌트 귀속(Component Attribution): 순수 PyTorch 백엔드 환경에서 Mamba의 Conv1D 컴포넌트를 어블레이션하여 태스크별 기여도를 격리한다.
  • 액티베이션 패칭(Activation Patching): 사전 학습된 Mamba에서, 클린 입력의 중간 위치 액티베이션을 오염된 입력으로 패칭하여 로짓(logit) 회복력을 측정한다.

주요 기여

1. 태스크 의존적 역전: 본 논문은 "후기 집중형" 대 "분산 축적형" 패턴이 아키텍처 불변량이 아님을 입증한다. 이 패턴은 태스크에 따라 뒤바뀐다:
   • 접두사 업데이트 태스크 (Parity, $S_3$): 순환 모델과 SSM은 최종 계층에 가독성 있는 상태를 집중시키는 반면, Transformer는 여러 계층에 걸쳐 점진적으로 구축한다.
   • 유계 스택 추서 태스크 (Dyck): 패턴이 역전된다. Transformer는 즉각적으로(초기 계층) 높은 가독성을 달로하는 반면, 순환 모델과 SSM은 점진적으로 상태를 구축해 나간다.
2. 대수와 계산의 분리: 비교환적인 $S_3$ 태스크를 도입함으로써, 저자들은 인코딩 프로필이 대수적 구조(교환성)가 아닌 계산적 구조(접두사 업데이트 vs 스택)를 추적한다는 것을 보여준다. $S_3$는 모든 아키텍처에서 Parity와 함께 그룹화되었으며, 이는 교환성이 인코딩 프로필의 주요 동인이라는 가설을 반박한다.
3. 가독성-인과성 디커플링(Decoupling): 상태가 가장 선형적으로 읽기 쉬운 계층(최적의 프로브 계층)이 항상 태스크에 인과적으로 필요한 계층과 일치하는 것은 아님을 밝혀냈다. 이러한 디커플링은 표현이 더 중복적이고 분산되어 있는 사전 학습된 모델에서 특히 두드러진다.
4. 컴포넌트 수준의 태스크 의존성: 단일 아키텍처(Mamba) 내에서도 특정 컴포넌트의 중요성은 태스크에 따라 달라진다. 예를 들어, Conv1D 컴포넌트는 $S_3$와 Parity에는 결정적이지만 Dyck 성능에는 유의미한 영향을 미치지 않는데, 이는 아키텍처의 단일한 "메커니즘"이라는 개념에 도전한다.

주요 결과

• 행동 수준: 순환 모델 베이스라인은 Parity와 $S_3$에서 완벽하게 일반화하며, Dyck에서는 SSM과 Transformer보다 우수한 성능을 보인다. $S_3$의 성능 순위(Recurrent > SSM > Transformer)는 Dyck이 아닌 Parity의 양상을 따른다.
• 프로빙 프로필:
  • Parity/$S_3$: Mamba와 LSTM은 최종 계층($4$개 계층 모델의 경우$L_3$)에서 프로브 정확도가 급격히 상승한다. Transformer는 $L_1$에서 $L_3$까지 점진적인 증가를 보인다.
  • Dyck: Transformer는 $L_0$ 또는 $L_1$에서 높은 정확도에 도달한다. Mamba와 LSTM은 점진적인 증가를 보이며 후기 계층에서야 포화 상태에 도달한다.
  • $S_3$ 확인: $S_3$는 Parity 프로필(순환 모델/SSM의 후기 집중, Transformer의 점진적 구축)을 따르며, 이를 통해 계산적 구조 가설을 확인한다.
• 인과적 개입:
  • 4개 계층의 형식적 모델에서, 최적의 프로브 방향을 어블레이션하면 상당한 정확도 저하가 발생하며, 이 방향들은 더 긴 시퀀스 길이(OOD)에서도 여전히 필수적이다.
  • 사전 학습된 모델: 분리 현상이 나타난다.
    • Pythia (Dyck): 강력한 중간 계층 병목(bottleneck)이 존재한다 ($160\text{M}$ 모델에서 $L_6$~$L_7$ 어블레이션 시 정확도가 약 81% 하락). 최적의 프로브 계층($L_{11}$)은 가장 인과적으로 중요한 계층이 아니다.
    • Mamba (Parity/Dyck/$S_3$): 최종 계층은 가독성이 매우 높지만, 단일 방향 어블레이션은 미미한 저하(<0.2%)만을 일으킨다. 그러나 최종 계층의 중간 위치 액티베이션 패칭은 로짓 차이의 약 97~98%를 회복시키는데, 이는 상태가 충분히 존재하지만 단일 선형 프로브에 포착되지 않고 여러 방향에 걸쳐 분산되어 있음을 시사한다.
• Conv1D 귀속: Mamba에서 Conv1D를 제거하면 $S_3$(-12.64 pp)와 Parity(-6.83 pp)에서 큰 폭의 하락이 발생하지만, Dyck에서는 유의미한 하락이 없다(+4.04 pp, 유의하지 않음). 이는 컴포넌트의 중요성이 태스크 의존적임을 나타낸다.

의의 및 주장

본 논문은 "아키텍처 X가 어떻게 상태를 인코딩하는가?"라는 질문은 태스크의 계산적 구조를 명시하지 않고는 불완전하다고 주장한다. 저자들의 주장은 다음과 같다:

1. 기계론적 시그니처는 결합된 속성이다: 상태 인코딩 전략은 고정된 아키텍처 특성이 아니라, 모델의 귀납적 편향(inductive bias)과 태스크의 계산적 요구사항(접두사 업데이트 vs 스택) 사이의 상호작용에서 발생하는 결과물이다.
2. 방법론적 주의: 한 가지 태스크 영역(예: 접두사 업데이트)에서 얻은 기계론적 발견을 다른 영역(예: 유계 스택 추적)으로 확장하는 것은 태스크 특이적 시그니처를 아키텍처 불변량으로 오인할 위험이 있다.
3. 가독성 vs 인과성: 프로빙은 정보가 선형적으로 이용 가능한 곳을 식별하지만, 이것이 반드시 계산이 병목되는 지점과 일치하는 것은 아니다. 특히 중복성이 높은 대규모 사전 학습 모델에서 더욱 그러하다.
4. 대수 vs 계산: 이러한 형식적 태스크의 맥락에서, 계산적 구조(순차적 접두사 업데이트 vs 계층적 스택)가 교환성과 같은 대수적 성질보다 계층별 인코딩 프로필을 더 잘 예측하는 지표이다.

저자들은 이러한 발견이 형식적 태스크, 사전 학습된 규모(최대 410M 파라미터), 그리고 준실제 코드 벤치마크 전반에 걸쳐 유효함을 입증하며, 시퀀스 모델링에서 아키텍처와 태스크 구조 사이의 견고한 상호작용을 시사한다고 결론짓는다.