[논문 리뷰] Tracr: Compiled Transformers as a Laboratory for Interpretability

David Lindner, János Kramár외 4명|arXiv (Cornell University)|2023. 01. 12.

Explainable Artificial Intelligence (XAI)인용 수 6

한 줄 요약

Tracr는 인간이 이해할 수 있는 RASP 프로그램을 표준 디코더 중심 트랜스포머 모델로 컴파일하여, 알려진이고 해석 가능한 구조를 가지게 하여 해석 가능성 연구에서 통제 가능한 실험을 가능하게 한다. 지능형 계산 메커니즘이 확정된 모델을 생성함으로써 Tracr는 슈퍼포지션과 같은 현상 및 기울기 기반 기여도, 인과 추적과 같은 해석 가능성 방법의 검증을 가능하게 한다.

ABSTRACT

We show how to "compile" human-readable programs into standard decoder-only transformer models. Our compiler, Tracr, generates models with known structure. This structure can be used to design experiments. For example, we use it to study "superposition" in transformers that execute multi-step algorithms. Additionally, the known structure of Tracr-compiled models can serve as ground-truth for evaluating interpretability methods. Commonly, because the "programs" learned by transformers are unknown it is unclear whether an interpretation succeeded. We demonstrate our approach by implementing and examining programs including computing token frequencies, sorting, and parenthesis checking. We provide an open-source implementation of Tracr at https://github.com/google-deepmind/tracr.

연구 동기 및 목표

• 해석 가능성 연구에서 지능형 설명의 부재 문제를 해결하기 위해, 알려진이고 인간이 읽을 수 있는 계산 메커니즘을 가진 트랜스포머 모델을 만드는 것.
• 진단 도구의 재현 가능하고 체계적인 테스트 환경을 제공하여, 사전에 정확한 계산이 알려져 있는 환경을 조성하는 것.
• 다단계 알고리즘을 실행하는 모델에서 슈퍼포지션과 같은 신경망 현상에 대한 통제 가능한 조사를 가능하게 하는 것.
• 컴파일된 추적 가능한 모델을 통해 트랜스포머 메커니즘을 이해하는 데 도움이 되는 교육적 및 실험적 플랫폼을 제공하는 것.
• 알고리즘이 알려진 행동을 보이는 모델을 사용하여 해석 가능성 방법의 평가를 지원하는 것 — 예를 들어 분류자 프로브, 기울기 기반 기여도, 인과 추적 등.

제안 방법

• Tracr는 트랜스포머 계산을 위한 도메인 특화 언어인 RASP를 표준 디코더 중심 트랜스포머 아키텍처로 컴파일한다. 이 모델은 알려진된 구조적 구성 요소를 가진다.
• 컴파일러는 각 레이어의 기능 — 예를 들어 지표 계산, 토큰 선택, 값 집계 — 이 고수준 프로그램 연산에 직접 매핑되도록 모델을 생성한다.
• 잔차 스트림은 중심 메모리 메커니즘으로 사용되며, 어텐션과 MLP 레이어는 특정 계산 단계(예: 선택, 집계)를 명시적으로 구현하도록 설계된다.
• 이 방법은 트랜스포머 회로 관점에 기반하며, 쿼리-키 및 값-출력 행렬을 통해 어텐션을 매개변수화하여 모듈화되고 해석 가능한 설계를 가능하게 한다.
• 모델은 경사 하강법을 통해 훈련되며, 내부 표현 분석을 통해 빈도 수세기, 정렬, 괄호 검사와 같은 알고리즘의 정확한 구현 여부를 검증한다.
• 슈퍼포지션은 훈련된 모델을 압축하여, 특징이 낮은 차원의 부분공간에서 어떻게 잘리거나 표현되는지를 관찰함으로써 연구된다.

실험 결과

연구 질문

• RQ1 RASP 프로그램에서 컴파일된 트랜스포머 모델은 해석 가능성 방법 평가를 위한 지능형 기준 기준으로 사용될 수 있는가?
• RQ2 압축된 트랜스포머 모델은 슈퍼포지션에서 특징을 어떻게 표현하는가? 이론 예측과 같이 불필요한 특징이 제거되는가?
• RQ3 기울기 기반 기여도나 인과 추적과 같은 해석 가능성 도구가 Tracr가 생성한 모델의 알려진 계산 구조를 얼마나 정확히 복원할 수 있는가?
• RQ4 정렬이나 괄호 검사와 같은 다단계 알고리즘은 인과적 어텐션 및 MLP 패턴을 가진 트랜스포머로 신뢰성 있게 컴파일될 수 있는가?
• RQ5 순위-1 예측자 제약 조건이 존재할 때, 어텐션 기반 선택기에서 복잡한 논리 연산(예: OR, AND)을 구현하는 데에는 어떤 제한이 있는가?

주요 결과

• Tracr는 지능형 내부 메커니즘을 가진 트랜스포머 모델로 RASP 프로그램을 성공적으로 컴파일하여, 토큰 빈도 수세기, 정렬, 괄호 검사와 같은 알고리즘을 정확히 실행한다.
• 컴파일된 모델은 명확하고 해석 가능한 어텐션 및 MLP 패턴을 보인다: 예를 들어, 한 어텐션 헤드는 선택을 수행하고 다른 하나는 집계를 수행하여 의도된 프로그램 논리와 일치한다.
• 기울기 하강법을 통해 압축된 후, Tracr 모델은 불필요한 특징을 제거하고 중요도가 낮은 특징을 슈퍼포지션에 표현하며, Elhage 등(2022b)의 예측을 확인한다.
• 이 모델들은 기울기 기반 기여도 및 인과 추적과 같은 해석 가능성 도구의 검증을 가능하게 하며, 비교를 위한 알려진 지능형 메커니즘을 제공한다.
• 선택기에서 불리안 연산자인 OR는 특정 조건(예: 공유 s-ops)에서만 구현 가능하며, 일반적인 경우는 순위-1 예측자 제약 조건에 의해 제한된다.
• 이 접근법은 복잡한 알고리즘을 예측 가능하고 모듈화된 구성 요소를 가진 트랜스포머로 컴파일할 수 있음을 보여주며, 연구 및 교육적 용도 모두에 기여한다.