Causal Evidence of Stack Representations in Modeling Counter Languages Using Transformers

기술 요약: 트랜스포머를 이용한 카운터 언어 모델링에서의 스택 표현에 대한 인과적 증거

문제 정의
기계론적 해석 가능성(mechanistic interpretability) 연구가 형식 언어(formal languages)를 학습한 언어 모델(LM) 내의 잠재적 특징들을 성공적으로 식별해 왔으나, 이러한 특징들의 인과적 역할을 확립하는 데에는 여전히 중요한 공백이 존재한다. Tiwari 등(2025)의 선행 연구는 카운터 언어(예: Dyck-1, Shuffle-k)로 학습된 트랜스포머가 기본 스택 구조와 일치하는 표현을 학습한다는 것을 보여주었다. 그러나 이러한 연구들은 이 표현들이 단순히 작업과 상관관계가 있는 것인지, 아니면 모델의 계산을 위해 인과적으로 필수적인 것인지까지는 결정하지 못했다. 본 논문은 다음 질문을 다룬다: 학습된 스카크 표현이 트랜스포머의 다음 토큰 예측 능력에 인과적으로 관여하는가?

연구 방법론
본 연구는 학습, 프로빙(probing), 그리고 인과적 개입(causal intervention)을 포함하는 3단계 실험 파이프라인을 채택한다:

1. 모델 학습: 인코더 전용(encoder-only) 트랜스포머를 Shuffle-k 언어($k \in \{2, 4, 6, 8\}$)에 대한 다음 토큰 예측을 위해 학습시킨다. Shuffle-k는 $k$개의 독립적인 Dyck-1 언어가 인터리빙(interleaving)된 것으로, 모델이 $k$개의 서로 다른 괄호 스택 깊이를 동시에 추적해야 함을 요구한다. 모델 아키텍처는 단일 레이어, 4개의 어텐션 헤드, 임베딩 차원 32, 은닉 차원 64로 구성된다. 모델은 25번째 에포크(epoch)까지 완벽한 검증 정확도를 달성하며, 이를 통해 이후의 모든 분석은 완전히 수렴된 모델을 대상으로 수행됨을 보장한다.
2. 스택 깊이 프로빙: 모델의 은닉 상태($h^{(l)}_t$)로부터 특정 스택의 깊이($c_i(t)$)를 예측하도록 선형 프로브(linear probes)를 학습시킨 시킨다. 프로브가 허위 상관관계가 아닌 실제 구조적 정보를 포착하도록 하기 위해, 변형된 레이블(perturbed labels)로 학습된 **제어 프로브(control probe)**를 함께 사용한다. 지표인 **선택도(Selectivity)**는 실험용 프로브와 제어 프로브 간의 정확도 차이로 정의된다. 높은 선택도는 스택 깊이가 은닉 공간에 진정으로 인코딩되어 있음을 나타낸다.
3. 인과적 개입 (절제/Ablation):
   • 방향 추출: 프로브의 가중치 행렬 $W$를 특이값 분해(SVD)를 통해 분해한다. 우측 상위 특이 벡터 $w$를 스택 깊이 표현에 대응하는 주요 방향으로 식별한다.
   • 절제(Ablation): 추론 과정 중, $w$와 정렬된 은닉 상태의 성분을 $h' = h - \alpha(w \cdot h)w$ 연산을 사용하여 제거한다. 여기서 $\alpha \in [0, 1]$는 절제 강도를 나타낸다.
   • 대조군: 무작위 방향들도 유사하게 절제되어 베이스라인 역할을 수행한다.
   • 평가: 절제된 상태를 디코더 헤드에 직접 통과시켜(트랜스포머의 내부 레이어를 다시 실행하지 않고) **위치 정확도(Positional Accuracy, $A_{pos}$)**와 **순차 정확도(Sequential Accuracy, $A_{seq}$)**에 미치는 영향을 측정한다.

주요 결과

• 표현의 존재성: 선형 프로브는 높은 선택도(예: $k=6$, Stack 2의 경우 85.8%)와 함께 스택 깊이 정보를 성공적으로 추출하며, 이는 모델이 스택 구조와 일치하는 표현을 학습했음을 확인시켜 준다.
• 인과적 필요성: 주요 스택 방향($w$)을 표적 절제했을 때 모델 성능이 급격히 붕괴된다. 절제 강도 $\alpha$가 0에서 1로 증가함에 따라, 순차 정확도는 0%에 가깝게 떨어진다.
• 특이성: 무작위 방향을 절제하는 것은 정확도에 유의미한 영향을 미치지 않으며, 이는 성능 붕괴가 벡터 조작에 의한 일반적인 현상이 아니라 스택 표현의 제거에 특이적인 것임을 확인해 준다.
• 지표 민감도: 연구 결과 위치 정확도는 다소 하락하지만 순차 정확도가 더 급격히 붕괴됨을 관찰하였다. 이는 오류가 누적되는 성질 때문이며, 단 하나의 잘못된 토큰이 전체 문자열을 무효화하기 때문이다.

의의 및 주장
본 논문은 카운터 언어로 학습된 트랜스포머 내의 스택 표현이 단순한 창발적 부산물이 아니라, 모델의 내부 계산과 다음 토큰 예측에 인과적으로 필수적이라는 것에 대한 강력한 경험적 증거를 제공한다고 주장한다.

특정 학습된 표현 방향의 제거가 모델의 거의 완전한 실패를 초래한다는 것을 입증함으로써, 저자들은 트랜스포머가 형식 언어에 요구되는 논리적 추론 과제를 해결하기 위해 이러한 스택 형태의 구조를 명시적으로 활용하고 있다고 주장한다. 이는 프로빙(probing)을 통한 상관관계 분석을 넘어 인과적 기계론적 이해로 분야를 확장하며, 모델의 "블랙박스" 동작이 특정적이고 식별 가능한 계산 프리미티브(computational primitives)에 의존하고 있음을 입증한다.

향후 연구 방향
저자들은 향-후 연구에서 다음을 탐구할 것을 제안한다:

• 절제되거나 편집된 상태를 피드포워드 네트워크(FFN)를 통과시켜 이후의 출력 변화를 관찰하는 방식의 미드-네트워크(mid-network) 개입.
• 다양한 아키텍처가 이러한 표현의 창발, 중복성 및 인과성에 미치는 역할 조사.
• 트랜스포머가 카운터와 같은 다른 창발적 계산 프리미티브를 어떻게 처리하는지에 대한 이해.