[논문 리뷰] Interpreting Attention Layer Outputs with Sparse Autoencoders

Connor Kissane, Robert Krzyzanowski외 3명|arXiv (Cornell University)|2024. 06. 25.

Ferroelectric and Negative Capacitance Devices인용 수 4

한 줄 요약

이 논문은 트랜스포머의 어텐션 레이어 출력에 스퍼스 자동에코더(SAEs)를 적용하여, 장거리 컨텍스트, 단거리 컨텍스트, 인덕션 패턴과 관련된 희박하고 해석 가능한 특징을 추출함을 보여준다. 저자들은 입력 토큰으로의 특징 기원을 추적하기 위해 재귀적 직접 특징 기여도(RDFA)를 도입하여 어텐션 메커니즘과 회로 행동의 세부 분석을 가능하게 하며, 이는 여분의 인덕션 헤드의 역할을 규명하고 간접 목적어 식별 회로를 검증하는 데 기여한다.

ABSTRACT

Decomposing model activations into interpretable components is a key open problem in mechanistic interpretability. Sparse autoencoders (SAEs) are a popular method for decomposing the internal activations of trained transformers into sparse, interpretable features, and have been applied to MLP layers and the residual stream. In this work we train SAEs on attention layer outputs and show that also here SAEs find a sparse, interpretable decomposition. We demonstrate this on transformers from several model families and up to 2B parameters. We perform a qualitative study of the features computed by attention layers, and find multiple families: long-range context, short-range context and induction features. We qualitatively study the role of every head in GPT-2 Small, and estimate that at least 90% of the heads are polysemantic, i.e. have multiple unrelated roles. Further, we show that Sparse Autoencoders are a useful tool that enable researchers to explain model behavior in greater detail than prior work. For example, we explore the mystery of why models have so many seemingly redundant induction heads, use SAEs to motivate the hypothesis that some are long-prefix whereas others are short-prefix, and confirm this with more rigorous analysis. We use our SAEs to analyze the computation performed by the Indirect Object Identification circuit (Wang et al.), validating that the SAEs find causally meaningful intermediate variables, and deepening our understanding of the semantics of the circuit. We open-source the trained SAEs and a tool for exploring arbitrary prompts through the lens of Attention Output SAEs.

연구 동기 및 목표

• 스퍼스 자동에코더(SAEs)가 어텐션 레이어 출력을 희박하고 해석 가능한 특징들로 분해할 수 있음을 입증하는 것.
• 재귀적 직접 특징 기여도(RDFA)를 개발하고 적용하여 SAE 특징 활성화가 입력 토큰과 상游 구성 요소로 돌아가 기여도를 추적하는 것.
• 다중의미 어텐션 헤드와 보이기 쉬운 여분의 인덕션 헤드의 功能 다양성에 관한 열린 질문을 해결하는 것.
• SAEs가 알려진 회로, 예를 들어 간접 목적어 식별 회로에서 인과적으로 의미 있는 중간 변수를 식별함을 검증하는 것.
• 커뮤니티가 어텐션 메커니즘을 해석하는 데 사용할 수 있도록 오픈소스 도구와 SAEs를 제공하는 것.

제안 방법

• 모든 헤드의 사전선형 어텐션 출력(z_cat)을 연결하여 트랜스포머 레이어에서 SAEs를 훈련시켜 고차원적 활성화를 재구성하는 것.
• 가중치 기반 헤드 기여도를 사용하여 학습된 SAE 특징을 특정 어텐션 헤드와 연관지켜 다의미성 문제를 완화하는 것.
• 재귀적 직접 특징 기여도(RDFA)를 구현하여 SAE 특징 활성화를 상游 구성 요소, 즉 잔류 스트림 특징과 입력 토큰의 기여도로 재귀적으로 분해하는 것.
• 트랜스포머의 선형 구조를 활용하여 z_cat의 기여도를 어텐션 패턴, 값 투영, 잔류 스트림 특징으로 선형 분해하는 것.
• 어텐션 출력에 대한 SAEs와 잔류 스트림 SAEs를 결합하여 최종 SAE 특징에서 입력 토큰으로의 종단 간 기여도 추적을 가능하게 하는 것.
• 모든 프롬프트에서 특징 회로를 시각화하고 탐색할 수 있도록 https://robertzk.github.io/circuit-explorer 에 도구를 배포하는 것.

Figure 1 : Overview. We train Sparse Autoencoders (SAEs) on $\textbf{z}_{\text{cat}}$ , the attention layer outputs pre-linear, concatenated across all heads. The SAEs extract linear directions that correspond to concepts in the model, giving us insight into what attention layers learn in practice.

실험 결과

연구 질문

• RQ1 어텐션 레이어는 어떤 개념을 학습하며, 이를 의미 있는 해석 가능한 특징들로 분해할 수 있는가?
• RQ2 어텐션 헤드는 모델 행동에 어떻게 기여하며, 다의미성 또는 여분의 구성에서 어떤 역할을 하는가?
• RQ3 어텐션 출력에 대한 SAEs가 왜 많은 보이기 쉬운 여분의 인덕션 헤드가 존재하는지에 대한 수수께끼를 해결하는 데 도움이 되는가?
• RQ4 어텐션 레이어의 SAE 특징은 간접 목적어 식별 회로와 같은 알려진 회로에서 인과적으로 의미 있는 중간 변수와 대응하는가?
• RQ5 재귀적 기여도 기법, 예를 들어 RDFA는 특정 입력 토큰과 상游 구성 요소로의 특징 기원을 추적할 수 있는가?

주요 결과

• SAEs는 어텐션 레이어 출력을 희박하고 해석 가능한 특징들로 성공적으로 분해하며, 이는 장거리 컨텍스트, 단거리 컨텍스트, 인덕션 패턴을 포함한다.
• GPT-2 Small의 어 schemes 중 최소 90퍼센트 이상이 다의미성을 띠며, 맥락에 따라 서로 관련 없는 기능을 수행한다.
• RDFA는 보이기 쉬운 여분의 인덕션 헤드가 장접두 또는 단접두 인덕션에 특화되어 있음을 드러내어 핵심 열린 질문을 해결한다.
• 간접 목적어 식별 회로는 SAE 특징을 통해 분석할 경우, 인과적으로 의미 있는 중간 변수를 포착하므로 더 잘 이해될 수 있다.
• 재귀적 기여도로 확인된 lin, 'Dave'에 대한 정보는 'is' 토큰에서 'isn’t'으로 어텐션 연결을 통해 라우팅된다.
• 저자들은 훈련된 SAEs와 회로 탐색 도구를 배포하여 프롬프트 간 어텐션 특징을 커뮤니티가 탐색할 수 있도록 한다.

(a)