Causal Scrubbing a method for rigorously testing interpretability hypothesis [Redwood Research] 이해하기

Original authors: by LawrenceC, Adrià Garriga-alonso, Nicholas Goldowsky-Dill, ryan_greenblatt, jenny, Ansh Radhakrishnan, Buck, Nate Thomas

Translated by: Yong-Min Shin

역자 노트: 해당 글은 어떤 mechanistic 모델 해석이 이미 주어졌을 때, 이에 대한 타당성을 검증하는 방법론을 소개하는 글이다. 즉, 이 글에서는 모델 해석 (여기에서는 그 해석은 검증이 아직 되지 않았기 때문에 가설이라고 부른다)을 어떻게 찾을 수 있었는지가 아닌, 결국 이 해석이 실제 모델을 얼만큼 정확하게 기술하고 있는가에 초점을 맞추고 있다 (물론 뒷 부분에서 이런 방법론이 더 좋은 해석을 찾을 수 있게 도움을 줄 수 있다는 내용도 다루기는 한다). 이를 생각하고 글을 읽으면 도움이 될 것이다.

원문 링크: https://www.lesswrong.com/s/h95ayYYwMebGEYN5y/p/JvZhhzycHu2Yd57RN

원 번역 글 Notion link: Causal Scrubbing a method for rigorously testing interpretability hypothesis [Redwood Research] (요약 번역) - coming soon

요약

이번 글에서는 “causal scrubbing”, 즉 어떠한 mechanistic interpretation의 타당성을 실험하기 위해 고안된 체계적인 방법론을 다룬다. Causal scrubbing의 핵심은 모델의 행동 특성을 보존하는 방식으로 resampling을 하면서 ablation을 진행하는 것이다. 우리는 이 방법론을 가지고 어떻게 작은 언어모델이 induction (귀납적 추론)을 행하고, 구체적인 예시로 어떻게 괄호들로 이루어진 sequence가 주어졌을 때 괄호가 제대로 닫혀있는지 판별할 수 있는지를 규명한다.

Introduction

모든 mechanistic interpretability를 연구하는 사람들이 던져야 하는 질문은 “해당 interpretation (모델에 대한 해석)이 실제 모델이 일으키는 현상을 얼만큼 해석할 수 있는가?”가 될 것이다. 최근에 이루어지는 여러 mechanistic interpretability 연구들을 보면, 각 연구에서 제시하는 해석 (및 이를 검증하는 과정까지 포함된다)이 ad-hoc (거의 case-by-case)인 경향성을 지닌다 [1].

이러한 경향성, 즉 해석 뿐만 아니라 이에 대한 검증을 위한 실험이 제각각이다보니 자연스럽게 mechanistic interpretation 연구의 범위를 넓히는데 근본적인 한계점으로 다가오게 된다. 현재로서는, 주어진 연구 결과를 검증하기 위해서 분석의 대상이 되는 현상 및 이에 대한 해석에 대한 깊은 이해가 밑바탕이 되어야하고, 이를 바탕으로 연구자의 주관적 판단이 이루어져야한다. 우리의 이상적인 목표는 mechanistic interpretability 분야에서도 proterty-based testing 같은 분석을 진행할 수 있게 만드는 것이다 [2] — 해석의 타당성을 연구자의 수고스러운 노동과 주관적 판단이 아닌 자동적으로 판별할 수 있게 하는 방향성이라고 할 수 있다. 또한 절차 지향적인 접근은 mechanistic interpretation의 연구의 스케일 (더 큰 모델, 더욱 섬세함을 요구하는 분석, 더 많은 연구자들 모두 포함된다)을 늘이는데에도 도움이 될 것이다. 이러한 방향성에 도움을 주고자, 우리는 좋은 해석과 그렇지 못한 해석을 구분할 수 있을 만큼 강력하면서도, 더욱 복잡한 해석에도 적용이 될 수 있을 만큼 일반적인 절차를 개발하고자 하였다.

이 연구에서는 causal scrubbing이라는, 특정 neural network이 어떻게 작동하는지에 대한 (엄밀하게 세워진) 가설을 기계적/순차적인 ablation을 통해 검증할 수 있는 프레임워크를 제안한다 [Ablation]. 구체적으로, 주어진 모델이 보이는 특정 행동 패턴에 대한 해석으로 이 모델의 한 부분이 특정 계산을 하고 있다는 (다소 덜 엄밀한) 가설이 주어졌을 때, 이 가설을 모델 안에서 일어나는 (실제) computation graph와 사람이 이해할 수 있는 computation graph 사이의 대응관계로 더욱 formal하게 변환하는 작업을 먼저 진행한다. 그 다음에, output으로부터 recursive하게 (재귀적으로) 가설에 의해 불변하는 부분들을 모두 파악한 다음, 그 부분에 대한 activation들을 (데이터와 가설에 의한 몇 가지 제한 조건을 만족하면서) 최대 entropy 분포로 대체하는 작업을 한다 (이를 scrubbing—즉 ‘닦아내는’—작업이라고 표현한다). 이렇게 scrubbing이 끝난 모델이 처음에 분석하고자 했던 행동 패턴이 여전히 얼만큼 잘 발현되고 있는지를 측정한다. 직관적으로 생각해서, 처음 세운 가설이 모델의 행동 패턴에 대한 적절한 가설이라고 한다면, scrubbing을 하더라도 모델이 그 행동 패턴을 재현하는데에는 문제가 없어야 할 것이다.

주어진 각 상황에 따라 달라졌던 기존의 검증 절차들과 달리, causal scrubbing은 다양한 interpretability hypothesis (해석 혹은 가설 [3]. 이 글에서는 문맥상 큰 문제가 없으면 이를 굳이 구별하지 않는다)에 적용될 수 있는 목표를 지니고 있으며, 이는 mechanistic interpretability 연구들이 제시했던 가설들의 대부분에 해당한다. 우리는 이러한 방법론을 적용한 두 가지 예시를 추후 다룰 것이다: 첫 번째로는 괄호 (parenthesis) 밸런스 확인 작업이고, 두 번째로는 2-layer attention-only 모델에 존재하는 induction head와 관련된다.

우리의 contribution들은 다음과 같다:

우리는 다양한 mechanistic 해석들을 더욱 formal하게 표현할 수 있는 interpretatility hypothesis에 대한 개념을 정의하였다.
우리는 causal scrubbing, 즉 순차적으로 activation들을 모델의 성능에 영향을 미치지 않으면서 교체해나가는 검증 방법론을 제시하였다.
우리는 이 방법론의 두 가지 적용 예시를 보여주어 일반적으로 적용이 가능하다는 것을 보여준다.

(생략 — 원문을 통해 확인하도록 하자)

Setup (셋업)

데이터셋 $D$ : 도메인 (정의역) $X$ 에서 정의됨
(분석 대상이 되는) 함수 $f: X \rightarrow \mathbb{R}$
(분석 대상이 되는) 행동 패턴: $\mathbb{E}_{x \sim D}[f(x)]$

조금 더 구체적으로는, 데이터셋에 label을 포함하고 보통 output (행동 패턴)은 특정 metric에 의한 loss값으로 상정한다. 밑의 그림 (왼쪽)이 이를 그림으로 잘 나타내준다. 오른쪽의 그림은 그의 구체적인 예시로, …[A][B]…[A]와 같은 prompt가 언어 모델의 입력으로 들어갔을 때 다음에 나와야 할 token인 [B]에 대한 loss를 구하는 것을 보여준다.

이 행동 패턴을 기술하는 ‘가설’ $h$ 는 세 가지 요소의 tuple $(G, I, c)$ 로 정의된다:

Computation graph $G$ : 여기서 $G$ 는 모델 $f$ 와 (사실상) 같으며 [Extensionally equal], 모델의 계산이 실제로 어떻게 진행되는지를 기술한다.
Computation graph $I$ : 이는 모델이 아닌 모델에 대한 해석을 기술한다.
대응 관계를 나타내는 함수 $c$ : $G$ 와 $I$ 의 노드끼리 대응해주며, injective graph homomorphism이어야 한다 (즉, $I$ 에 존재하는 모든 edge는 반드시 $G$ 에 존재하며, $G$ 에만 존재하는 edge도 있다. 또한 graph structure가 대응 관계에서 망가지지 않는다).

추가적인 조건으로 $G$ 와 $I$ 모두 input과 output을 나타내주는 노드가 딱 한 개 씩 존재하도록 요구한다.

다음은 가설의 한 예시이다:

위의 예시에서, 가설을 나타내주는 computation graph $I$ 가 핵심적으로 주장하고자 하는 바는 다음과 같다: 모델을 나타내주는 computation graph $G$ 에서 A, B, C, D (파란색 부분)이 실질적으로 계산해주는 것은 $z_1$ 이 3보다 작은지, $z_2$ 가 3보다 큰지 등이라는 것이다.

몇 가지 주목할만한 사항들은 다음과 같다:

앞서 언급했다시피, $I$ 는 가설을 나타내주는 일종의 heuristic이라고 보면 된다. 이 가설은 완벽하지 않을 가능성도 충분히 있고, label $y$ 과도 일치하지 않을 경우도 있다. 어차피 모델을 표현하는 것이 가설의 목적이기 때문에, 후자의 경우 원래 모델도 같은 양상을 보이기만 하면 문제가 없다.
하지만 $c$ 가 각 노드의 activation 값들이 어떻게 $I$ 안에서 표현되는지는 알려주지 않는다. 다만 (개념적으로) 어느 노드에 대응이 되는지만 알려줄 뿐이다.
$c(I)$ 를 $G$ 의 “중요한 노드”들이라고 부를 것이다.
- 각 그래프에서 두 개의 대응되는 노드 쌍 $n_I$ , $n_G$ 를 보자 (즉, $c(n_I)=n_G$ ).
  - 같은 input에 대해 두 그래프 $G$ 와 $I$ 를 돌렸다고 생각하자. 그러면 $n_G$ (activation)이 $n_I$ (input에 대한 간단한 로직)을 표현해준다고 생각할 수 있다.
  - Causal scrubbing 알고리즘은 약간 더 약한 claim에 대한 증명을 할 것이다. 그 claim은 “ $n_I$ 에 들어가는 input에 대한 equivalence class는 $n_G$ 에 들어가는 input에 대한 equivalence class와 같다”이다. 이 claim에 대한 검증이 (틀린 가설을 모두 소거하기엔 힘들지라도) 우리가 mechanistic interpretability 안에서 관심을 가지는 가설들을 검증하기에 충분하다고 생각한다.
우리는 모델을 나타내주는 computation graph $G$ 를 때때로 더 보기 편한 방식으로 표현할 것이다.

‘가설’에 대한 더 심도있는 논의는 appendix를 보도록 하자.

Causal Scrubbing

여기서 우리는 causal scrubbing를 두 가지 다른 방식으로 설명하겠다.

조금 덜 포멀하고 직관적인 버전: Causal scrubbing의 핵심 아이디어를 이해할 수 있도록 많은 그림들과 함께 설명할 것이다.
Causal scrubbing 알고리즘 및 pseudocode.

독자마다 잘 맞는 설명 방식이 다 다를 것이기에, 이 글을 여기저기 옮겨다니면서 훑어보는 것을 추천한다.

우리의 목표는 우리의 해석 (가설이라고 생각해도 될 것이다) $I$ 에 의해 동일(equivalent)하다고 생각되는 activation들을 샘플해가며 $E_{\text{scrubbed}}(h, D)$ 이라는 metric을 정의하는 것이고, 이 값은 $\mathbb{E}_{d \sim D}[f(d)]$ 와 비교될 것이다. 만약 가설 $I$ 가 (합리적인 범주 안에) 정확하다면, causal scrubbing을 통해 activation을 대체하는 작업을 해도 loss가 그대로 유지될 것이고 $E_{\text{scrubbed}}(h, D) \approx \mathbb{E}_{d \in D}f(d)$ 일 것이다. 우리는 이 둘 ( $E_{\text{scrubbed}}(h, D)$ 와 $\mathbb{E}_{d \in D}f(d)$ )의 차이가 가설의 faithfulness를 측정하기에 유의미한 proxy가 되리라 생각한다 — 즉, 모델이 어떤 행동을 취했을 때 그 “진짜 이유”를 얼만큼 보여줄 수 있다는 것이다.

개념적인 설명: 가설은 어떤 activation들을 (마음대로, 혹은 적절한 것으로) 교체해도 된다고 말하고 있을까?

어떤 가설 $h = (G, I, c)$ 가 있다고 가정하자 (그림 참조).

가설(오른쪽)에 의하면 activation A와 B는 각각 ‘ $z_1$ 이 3보다 큰지’와 ‘ $z_2$ 가 3보다 큰지’를 판단하는 부분이다. 그 다음, 여전히 가설에 의하면 activation D는 이 둘 (A 혹은 B)중 하나라도 참인지 확인한다. 게다가 $z_3$ 과 activation C는 중요하지 않은 것으로 보인다 (적어도 y에 대한 log loss를 계산하는 행동에 있어서는 그렇다).

만약 이 가설이 참이라면, 우리는 다음 두 가지 ‘resampling ablation’을 진행해도 모델에 변화가 없을 것이다:

Activation A, B, D에 대해 가설 $I$ 의 논리 안에서 동일한 다른 input에서 나온 activation으로 교체하거나,
Activation C, 혹은 $z_1 \rightarrow B$ 와 같은 부분에 대해서 아무 input에 대한 activation으로 교체하거나 (가설에 의하면 중요하지 않은 부분이라고 했으므로).

이런 intervention (개입)을 더 설명하기 위해, $G$ 의 구조를 “treeified”한 버전으로 다시 그릴 것이다 (즉, 모든 input으로부터 output까지의 경로를 input을 중복해서 다시 그리는 방식으로 펼쳐놓음 — 그림 참조). 이렇게 되었을 때, 모델 중간의 한 부분에 대한 activation을 대체하게 된다면 그 지점 윗쪽에 해당하는 모든 가지들에 대해 activation을 다 교체하는 것과 같게 된다.

Intervention 1: semantically equivalent subtrees

모델에 두 input $x_1 = (5,6,7,\text{True})$ , $x_2 = (8,0,4,\text{True})$ 을 집어넣는다고 해보자. 가설에 의하면, $A’$ 의 결과는 $x_1$ 과 $x_2$ 모두 동일할 것이다. 따라서, 만약 가설이 참이라면, 원래 모델 $G$ 에 대해서도 두 input에 대한 결과가 equivalent해야할 것이다. 즉, $G$ 에 $x_1$ 을 넣게 되었을 때 $A$ 의 activation을 $x_2$ 을 넣었을 때의 것으로 교체해도 된다는 것이다:

따라서 A 및 그 윗 부분에 있는 모든 부분에 대해서 다 activation을 교체하였다 (꼭 $x_2$ 가 아니라 A’ 에 대해서 그 결과가 동일한 모든 $x \in D$ 에 대해 다 가능하다).

사실 우리의 가설에 의하면 이 뿐만 아니라 여러 activation 교체들이 가능하다. 예를 들어, D와 관련해서 이런 교체도 가능할 것이다:

Intervention 2: unimportant inputs

가설이 참이라는 가정하에 가능한 또다른 결의 개입 (intervention)은 중요하지 않은 노드들에 대해 교체를 하는 것이다. 예를 들어, 가설 $h$ 에 [6] 의하면 A 기준으로 유의미한 input은 $z_1$ 이다. 따라서 이 부분에 대해 $z_2$ 및 $z_3$ 를 다른 것으로 바꿔도 모델의 행동에 큰 변화가 없어야 한다 (또한, activation을 바꾸는 것과 그 activation을 일으키는 input을 바꾸는 것은 동일한 효과를 불러일으킬 것이다). 비슷한 논리는 $B$ 노드에서 $z_1$ 와 $z_3$ 에 대해서도 적용될 것이다. 게다가, 가설에 의하면 C는 D에 어떠한 영향도 일으키지 않는다고 하기 때문에, C를 어떠한 값으로 바꿔도 모델에 미치는 영향은 미미할 것이다.

그림으로 이를 표현하자면 다음과 같을 것이다:

참고로 여기서 우리는 3가지의 다른 input으로 3가지 다른 교체 작업을 하고 있다. 그럼에도, 만약 가설이 정확하다면, 모델에게 있어 중요한 정보는 변하지 않았을 것이고 $\text{Treeify}(G)$ 의 결과는 원래의 것과 비슷할 것이다.

Causal scrubbing 알고리즘은 이러한 두 종류의 개입을 여러번 진행하게 된다. 사실은, 이러한 개입을 가설이 허용하는 한 모델을 돌릴 때 마다 가능한 한 최대한으로 하는 것이 목표이다.

The causal scrubbing algorithm

이제 이 가설 검증 알고리즘에 대해 구체적으로 정의를 하겠다. 이 알고리즘은 이전 section에서 설명 되었던, 가설에 의해서 어떤 activation이 교체가 가능한지에 대한 기준을 이용한다.

핵심 아이디어는 이 가설을 일종의 “개입 블랙리스트 (intervention blacklist)” 처럼 해석하는 것이다. 비유적으로, 가설이 마치 중요하다고 지정된 노드 목록을 제외하고는 다른 노드들에 대해서는 아무리 바꿔도 상관없다고 외치고 있는 것과 같다.

한 문장으로 요약하자면: 우리가 모델의 어떤 activation을 계산할 때마다 “가설 $h$ 에 따르면, 이 activation을 교체해도 모델에 영향을 미치지 않는 다른 activation들이 어떤 것들이 있을까?”라고 질문을 던지고, 이에 해당하는 교체를 데이터셋(의 적절한 일부분)으로부터 uniformly random하게 뽑아 진행하고, 이를 재귀적으로 반복하는 것이다.

알고리즘에서는 treeify를 직접적으로 진행하지는 않는다; 하지만 마치 tree 구조를 지닌 것 처럼 경로를 하나씩 탐색하게 된다.

우리는 scrubbed expectation $E_{\text{scrubbed}}(h, D)$ 을 이 알고리즘에 의해 발생하는 (교체에 의해 발생하는) 샘플들에 대한 $f$ 의 행동 기댓값으로 정의한다.

직관적인 버전

(이 부분은 pseudocode를 이해하고 있다면 거의 중복되는 내용이므로 선호하는 방식대로 읽는 것을 추천한다.)

알고리즘의 pseudocode는 다음과 같이 쓸 수 있다:

어떤 기준이 되는 (reference) input $x$ 를 $D$ 로부터 샘플한다.
run_scrub을 재귀적으로 호출하면서 $I$ 의 output으로부터 input까지의 모든 경로를 훑어본다. 각 노드마다, 그 노드 ( $n_I$ 라고 하자)보다 ‘윗 쪽에 존재하는’ 모든 노드들을 모은 subgraph를 생각해볼 수 있다 (즉 이 노드를 계산하는데 영향을 미치는 모든 노드들이 된다). 그리고 이 subgraph들은 실제적으로 $G$ 에도 해당되는 부분이 $c(I)$ 안에 존재할 것이다.
run_scrub(n_I, c, D, x)의 리턴값은 $G$ 안에서의 어떤 activation이다. 구체적으로, 가설에 의해 input이 x일 때 $n_I$ 의 값에 대응되는 $G$ 안에서의 노드가 될 것이다.
- $n_G = c(n_I)$ 라고 하자.
- 만약 $n_G$ 가 input 노드라면 리턴값은 그냥 $x$ 가 될 것이다.
- 그렇지 않다면 $n_G$ 로부터 부모 노드들로 거슬러 올라가면서 input까지 activation들을 계산할 것이다. (가설 단에서) $p_I$ 를 $n_I$ 의 부모 노드, (실제 computation graph $G$ 에서) $p_G$ 를 $n_G$ 의 부모 노드라고 할 때, 다음을 수행한다:
  - $p_G$ 에 대응하는 $p_I$ 가 있는 경우: 그렇다면 가설은 $p_G$ 의 값이 $n_G$ 의 값에 유의미한 영향을 미친다고 생각한다는 것이다. 따라서 가설에 의해 x와 같은 효과를 가지는 (즉 $p_I$ 에서 결과가 같음) 새로운 데이터 new_x를 샘플한다. 그리고 run_scrub를 $p_I$ 에 대해서 재귀적으로 다시 부른다.
  - 대응하는 부모 노드가 없는 “중요하지 않은” 노드인 경우, other_x라는 다른 input을 선택한다. 이 input은 데이터셋으로부터 무작위로 선택되었다. 하지만, 이 경우로 다다를 때마다 다른 데이터를 무작위로 고르기 보다는 한 번 무작위로 골라진 input을 계속 고정해서 사용하도록 한다 [7]. 이 other_x를 사용했을 때의 $p_G$ 값을 따로 기록해둔다.
- 이제 $n_G$ 의 모든 부모 노드들에 대해 activation 값들을 계산하였다 — 그리고 이에 따라 $n_G$ 에 대한 개입을 한 결과를 얻기 위한 input도 얻게 되었으니, 이 input을 통해 모델을 돌리고 함수의 리턴값 얻어낸다.

Pseudocode: 원문 링크

왜 resampling을 통한 ablation을 해야 하는 것일까?

중요하지 않은 노드에서 “중요하지 않다”는 건 무슨 의미인가?

예를 들어 어떤 가설에 의하면 모델의 특정 모듈이 주어진 결과값을 계산하는데 중요하지 않다고 했다고 해보자. 이 “중요하지 않음”을 시험하기 위한 다양한 개입을 생각해볼 수 있다. 예를 들어:

Zero ablation: 해당 모듈의 activation들을 모두 0으로 고정
Mean ablation: 해당 모듈의 activation들을 데이터셋 전체에 대한 empirical mean으로 대체
Resampling ablation: 해당 모듈의 activation을 임의의 다른 input에 대한 activation으로 대체

이들 중 가장 적절한 것이 무엇인지 알아내기 위해, 우리가 ablation을 통해 정확히 무엇을 검증하려고 하는지 다시 한 번 생각해보자.

만약 우리가 검증하려고 하는 것이 “이 모듈의 activation들이 말 그대로 한 번도 사용된 적이 없다”라는 주장이면, 차라리 activation들을 엄청나게 큰 숫자나 심지어 NaN값으로 대체해볼만 하다. 하지만 실제로는, 이렇게 한다면 모델 작동에 문제가 생길것이기에, 이런 주장을 검증하는 것은 아닐 것이다.

조금 더 그럴싸한 주장은 다음과 같을 것이다: “모델은 해당 모듈의 activation (혹은 그로 인한 여러 특성들)에 의해 영향을 받을 것이지만, 그 activation들은 사실 해당 subtask에 대해 그렇게 유의미한 정보들을 담고 있지 않다”. 다른 말로 표현하자면: 해당 모듈의 activation의 분포는 (어쩌면) 모델 행동에 유의미한 영향을 미칠 것이지만, 어떤 특정한 input을 받았는지에 대해 조건부적으로 달라지는 분포의 특성에 의존하지 않는다는 것이다.

그렇기에, 우리는 이 가설을 개입적인 실험 (intervention experiment)으로 바꿨을 때 가장 적절한 방안은 무작위로 샘플된 input을 사용하는 것이라고 생각한다. 이렇게 되면 모듈의 activation들은 이미 계산된 적이 있는 분포일 것이지만, 그 분포가 모델의 최종 결과값에 미치는 연결고리는 닦여졌을 것이다 (scrubbed away).

Zero 및 mean ablation의 문제점

다른 선행연구에서 zero 및 mean ablation이 많이 쓰였지만 [8], 우리가 보이고자 하는 주장들과 결이 맞지 않는 부분들이 있다.

위에서 언급이 되었듯이, 우리가 보이려고 하는 주장은 현재 모델에 있어서 이 모듈의 output 정보가 중요하지 않다이지, 모듈을 아예 삭제해도 모델에 영향을 미치지 않는다가 아니다. 우리는 여러 가설들을 검증할 때 현재 모델 [9]과 가장 가까운 상태에서 검증하려고 하며, 그와 비슷하지만 엄연히 다른 모델로 대체하고 싶지 않으며, zero 및 mean ablation을 사용하면 이 문제를 피할 수 없게 된다. 이 문제는 세 가지의 방향성으로 드러날 수 있을 것이다:

Zero 및 mean ablation은 모델을 예측할 수 없는 데이터 분포 위에서 작동하게 만든다.
Zero 및 mean ablation은 모델 성능에 예측할 수 없는 영향을 미칠 것이다.
Zero 및 mean ablation은 variance를 없애며 모델의 행동 패턴을 왜곡할 것이다.

Results

Causal scrubbing 방법론의 효과를 보여주기 위해, 여기에서 두 가지 task에 대하여 이 알고리즘을 적용해 본다. 그 두 가지 task는 1) 우리가 예전에 이미 해석을 진행한 적인 모델에 적용하여 가설을 검증하는 것, 그리고 2) 2-layer attention-only 모델에서 induction head에 대한 가설 검증 및 개선이다. Causal scrubbing의 효과를 보여주기 위해 여기서는 실험 결론을 요약해서 보여주며, 자세한 내용은 각자의 글에서 확인하기를 바란다.

Parentheses balance checker

여기서는 causal scrubbing을 괄호 (parentheses)가 제대로 닫혀있는지 판별할 수 있게 학습된 작은 transformer 모델에 적용해보고자 한다. 자세한 내용은 해당 글을 참고하기 바란다. 구체적으로, 이 모델의 메커니즘과 관련하여 다음 세 가지 주장을 검증해보고자 한다. (역자 1: 참고로 attention head를 표시할 때 (layer):(head) 방식으로 표기한다. 예를 들어, head 2:3은 두 번째 레이어의 3 번째 head이다.)

(역자 2: Input은 ‘()())()()))’과 같이 생겼다)

주장 1: Output으로 주요 정보를 직접적으로 전달하는 head가 3개 존재한다.

Head 1.0과 2.0은 합동해서 1) 열린 괄호 ‘(’와 닫힌 괄호 ‘)’의 총합이 같은지와 2) 열린 괄호로 시작하는지를 확인한다.
Head 2.1은 nesting depth가 음의 값을 가지는 시점이 존재하는지 확인한다.

주장 1은 다음 가설로 표현될 수 있다:

주장 2: Head 1.0과 2.0은 input의 첫 번째 position에만 의존하며, 이 input은 간접적으로 다음에 의존한다:

Position 1에서 head 0.0의 output. 이 부분은 전체적으로 열린 괄호의 비중이 얼마인지를 계산한다. 이 정보는 residual stream에 특정한 방향을 가지는 벡터로 쓰여(write)진다.
Position 1의 embedding. 이는 sequence가 열린 괄호로 시작하는지에 대한 정보를 담고 있다.

주장 3: Head 2.1은 input의 모든 position으로부터 영향을 받으며, 특히 왼쪽에서 오른쪽으로 읽었을 때 nesting depth가 음의 값을 가지는지에 영향을 받는다.

이 세 주장을 종합하여 그림으로 표현하면 다음과 같을 것이다:

이 주장들을 causal scrubbing으로 검증을 해본 결과, 유의미하지만 완전하지는 않은 해석을 제공하는 것으로 보인다:

검증하는 주장

재현된 모델의 행동 비중 [11]

93%

1 + 2

88%

1 + 3

84%

1 + 2 + 3

72%

예상했던대로, 가설이 점점 구체적이 될 수록 모델을 표현하는 정확도가 낮아진다. 이는 가설이 자세해짐에 따라 교체해야하는 activation들이 많아지고, 교체되는 위치가 모델 깊숙하게 존재한다는 점에서 기인한다.

결과에 의하면 우리의 가설은 다양한 방식으로 미묘하게 틀렸다는 것을 방증한다 (정보가 흐르는 몇몇 경로를 놓치던가 모델이 실제로 사용하는 feature들을 모두 파악하지 못했거나 등).

On induction heads

여기서 우리는 2-layer attention-only 모델에서 ‘induction’ head에 대해 검증하는 실험을 진행한다 [12]. 그 과정에서 가설을 검증하고 손쉽게 가설을 개선할 수 있었다.

일단 첫 번째로 단순한 induction 가설을 세운 후 검증을 진행했다. 간단하게, input이 layer 1의 induction head으로 들어오면 value/key/query 이 세 가지 경로로 갈라지고 각각의 경로에서 중요한 정보가 어디에서 오는지에 대한 가설을 간단하게 세웠다. 여기에서 우리는 value와 query는 input쪽 token embedding으로부터 대부분의 정보가 직접적으로 오고 layer 0을 거치지 않는다는 가설을 세웠다. Key에 대해서는, 직전 token position에만 반응하는 attention layer 0의 output에 의해서 결정된다고 가설을 세웠다.

이 가설을 openwebtext의 subset에 대해 검증하였으며, 이는 induction이 (일어난다는 보장은 없지만) 도움이 될 것이라는 예상에 기반하였다. 이 데이터셋에 대하여, 이 단순한 induction 가설은 기존 성능의 35%밖에 재현하지 못하였다. 이 수치를 증가시키기 위하여 정보가 더 다양한 경로를 통해 흐를 수 있도록 가설을 다음과 같이 수정하였다.

일단, induction head의 attention pattern가 (하나의 token이 아닌) 연속적으로 붙어있는 3개의 token을 비교하는 것이 가능하도록 수정하였다.
그리고, query와 value 값이 현재 position에 해당하는 layer 0의 결과값에도 영향을 받을 수 있도록 수정하였다.
마지막으로 특정 케이스에 대해여 현재 token이 반복적으로 나타날 경우에 대해 반응하는 layer 0 head를 고려하였다.

이렇게 수정한 결과, 모델 성능의 86%까지 재현하는데 성공하였다.

우리는 기존처럼 ad-hoc한 접근법으로 가설을 세우는 방식으로 접근했으면 이 정도로 구제척인 가설을 만들어내는 것이 어려웠을 것이며, 이 가설에 대한 확신 또한 가지는 것이 힘들었을 것이라고 생각한다.

관련한 더 자세한 내용은 appendix 글을 참고하면 되겠다.

Alignment와의 관계에 대하여

Causal scrubbing과 alignment의 연결성 중 가장 먼저 떠오르는 것은 mechanistic interpretation을 검증하는 작업일 것이다. 이와 관련하여, alignment와 관련이 있는 몇 가지 사례들을 생각해볼 수 있을 것이다:

사람이 발견해낸 모델 해석의 검증: Mechanistic interpretability 연구의 규모를 키우는 데에 (causal scrubbing과 같은) 표준화된, 신뢰있는, 그리고 간편한 검증 절차가 있는 것은 많은 도움이 될 것이다; 이는 특히나 혁신적인 AI 모델 개발 전에 interpretatility 프로젝트 같은 것이 있을 때 도움이 될 것으로 보인다.
해석 모델 검출의 자동화: 어떤 경우에서는 연구자가 가설 후보들의 space를 규정하고 optimization 알고리즘을 이용하여 가장 적절한 해석을 찾으려고 할 수 있다. 우리는 이와 유사한 결의 방향성을 가지는 일을 하였고 이를 추후에 더 발전시킬 것이다.
AI-assist된 해석: 이를 이용하여 높은 퀄리티의 이해 가능한 해석을 내놓을 수 있는 모델을 학습할 수 있을 것이다.

이 세 가지 경우에서, 우리는 연구자들이 caual scrubbing에 의해 검증된 해석 (가설)들을 이해하고 있을 것을 요구하였다. 하지만, 큰 모델 (예를 들어, LLM)들 안에서 일어나는 계산들이 단순화가 불가능할 수준이라면 사람이 이해가 가능한 수준의 해석이 아예 불가능할수도 있다. 또한, 만약 사람이 이해 가능한 해석이 존재하더라도, 그런 해석을 찾는 것이 힘들거나 불가능할수도 있다.

Causal scrubbing (그리고 더 나아가 mechanistic interpretability)에 대해 우리가 관심을 가지는 것은 해석에 해당하는 개념들을 우리가 직접 조절할 수 있지만 (인간이던 AI이던) 반드시 이해가 되었음을 전제로 하지 않는 테크닉들이 필요한 application들에 있다:

ELK를 풀기 위한 자동화된 방법론들의 개발 [ELK]: ARC (Alignment Research Center)는 ELK를 해결하는 데에, (굳이 인간이 이해하지 않아도 되는) causal scrubbing으로 찾아낸 해석과 관련된 모듈들을 학습에 이용하는 몇몇 방향성들에 대해 긍정적으로 바라보고 있다.
함정 alignment 찾아내기: 만약 약간 믿을 수 있는 모델과 매우 믿을 수 없는 모델 두 개가 있다고 하자. 두 모델이 비슷한 행동을 한다고 쳤을 때, 매우 믿을 수 없는 모델이 왜 (아마 겉으로 보기에만) align이 되어있는지 직접 모델을 탐구하지 않고 해석을 통해 알아낼 수 있을 것이다.
Relaxed adversarial training에서, 모델의 내부를 바라보고 예상치 못한 행동을 할지를 판단해야하는 상황이 생긴다 (이는 보통 random sampling으로는 찾기 힘든 행동이며, 아주 드물게 training 분포 바깥에서 일어난다). 이 과정에서 내부 모습을 인간이 그 의미를 이해하는 것이 필수로 요구되지는 않는다. 단지 causal scrubbing을 이용해서 밝혀낸 내부 구조를 이용해서 모델의 행동을 예측할 수 있으면 된다.

Limitations

불행히도, causal scrubbing은 우리가 다루고 싶은 모든 interpretabiltiy 가설을 다루는 것은 아니다:

Causal scrubbing에서는 가설이 완벽하게 허용하는 범위 안에서만 activation 교체를 진행할 수 있다.
- 가설에 의하면 어떤 activation을 교체할 수 있는지를 엄격하게 보여주고 있지만, 실제에서는 이 범위 바깥을 벗어나는 가설들도 허용할 때가 있다.
- 예를 들어, 모델의 어떤 부분에서 연속적인 (continuous) 피쳐를 만들어낸다고 가설을 세웠을 때, 이 부분에 대한 activation을 다른 input을 통해 재현하려고 하면 아마 이를 정확하게 재현하는 것은 불가능하고 이와 비슷한 값을 목표로 해야할 것이다. 엄밀히 말하면, 이는 허용이 되지 않는다.
- 이를 해결하는 한 방법으로는 causal scrubbing을 조금 더 일반화해서, 가설 space를 continuous하게 취급을 하는 것이다 (이는 여러가지 장점이 있는데, 예를 들어 이 space 위에서 SGD를 돌릴 수도 있을 것이다). 이에 대해 우리는 기초적인 작업을 진행했지만, 아직 개념적으로 해결해야할 문제들이 남아있다.
Causal scrubbing은 제공된 data 안에서만 검증을 진행할 수 있다. 만약 causal scrubbing을 하다가 현재 input와 비슷한 성질을 가진 다른 input을 사용해야하는데, 그런 data가 존재하지 않는다면 더 이상 검증을 할 수 없다.
데이터셋 안에서 완벽한 상관관계를 보이는 두 변수가 있다면, causal scrubbing은 그 두 변수를 헷갈려하더라도 크게 티가 안날것이다.

또다른 한계점으로는

모델한테 모든 상황에 적용하는 것이 불가능한 heuristic이 있다고 할 때, 그 heuristic을 억제하기 위해 다른 circuit을 사용할수도 있다 (예를 들어, Indirect object identification 논문에 나오는 negative name mover head). 하지만 이런 억제 circuit들은 heuristic들이 적용 가능할 때는 방해가 될 뿐이다. 이런 경우, 억제 circuit들을 무시하게 될 경우, 그 heuristic이 모델 성능에 미치는 영향을 과대평가하게 될 수도 있고 불완전한 해석에 다다르게 될수도 있다.
만약 두 항 간의 상관관계가 높다면, (무작위로 activation을 교체하는 방식으로) 그들을 독립적으로 샘플링하면 합의 variance가 줄어들게 된다. 어떤 경우에는, 이 variance가 모델 성능에 악영향을 미칠 때도 있다 — polysemanticity와 관련된 예시. 이렇게 되면 이런 상관관계가 포함되어있지 않은 가설로 귀결될 가능성이 있고, causal scrubbing에 의해서 더 ‘정확한’ 가설로 보여질 것이다.

이런 예시들은 가설이 충분히 구체적이지 못하고 모델 내부에서 높은 상관관계를 지니는 부분들 (input과 관련된 부분이거나 activation과 관련된 부분이거나)을 포함하지 않아 생기는 일이다. 우리는 이러한 문제점들이 causal scrubbing 자체의 문제라기보단, 이런 해석 방식이 defeasible reasoning의 일종이기에 발생하는 것이라고 본다. 즉, 주어진 증거를 토대로 강력하고 설득력 있는 논증을 펼치지만 새로운 증거가 제시되면 이 논증 또한 변경 및 수정이 될 수 있다.

우리는 이런 문제들이 일종의 adversarial 접근법처럼, 가설에 가장 큰 영향을 미치는 방향으로 약간씩 수정을 가하면서 다시 검증하는 방식으로 점점 더 구체화를 시키는 방식으로 어느 정도 해결될 수 있다고 생각한다. 직관적으로, 이 adversarial한 process는 최종적인 가설이 모델 메커니즘 안에서 주요한 역할을 하는 부분은 모두 포함하기를 바라게 될 것이다 (그렇지 않으면 어떤 adversarial한 가설에 의해 이 부분이 지적을 당하게 됨으로). 이 부분과 관련한 자세한 논의는 이 appendix를 참고하도록 하자.

이런 한계점에도 불구하고, 우리는 causal scrubbing에 큰 기대를 가지고 있다. 간단한 모델들에 causal scrubbing을 적용해 해석을 곧바로 할 수 있었고 더 큰 모델과 더 복잡한 현상들에 대한 적용 가능성에 대해서도 긍정적이다. 현재로서는 추후 mechanistic interpretability 연구를 할 때 causal scrubbing이 중요한 방법론으로써 도움이 될 것이라 예상한다.

각주

[1] (Olsson et al.) (Meng et al.) (Wang et al.) (Nanda et al.)

[2] (역자) 속성기반 테스팅 (wiklpedia). 특정 input-output 쌍에 대한 알고리즘의 특성이 아닌 임의의 input과 그에 따른 output에서 진행한다는 점에서 원 저자들이 비유를 들었다.

[3] Mechanistic interpretabiltiy는 결국 모델을 해석하고자 하는 것이 목표이기 때문에 그 결과물은 당연이 해석이 될 것이고, 그와 동시에 검증을 거쳐야 하니 가설이기도 하다.

[4] Jacobi & Goldberg, Towards Faithfully Interpretable NLP Systems: How should we define and evaluate faithfulness?, ACL 2020

[5] Miao et al., Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism, ICML 2022

[6] (역자) 문맥상 큰 혼란이 없다면 가설을 나타내기 위해 $h$ 와 $I$ 를 혼용한다.

[7] 중요하지 않은 input에 대해서 다른 input으로 대체할 때 어떤 데이터를 샘플해야 하는지에 대한 논의는 appendix에 자세한 논의가 있다.

[8] (역자) Mean ablation은 Lundberg & Lee, A Unified Approach to Interpreting Model Predictions, NeurIPS 2017 등에서도 등장한다.

[9] (역자) 문맥 상 여기에서 모델은 실제 모델 뿐만 아니라 데이터셋 분포까지도 포함하고 있다고 볼 수 있다. 이는 XAI안에서도 out-of-distribution 문제와 연결이 되는데, 간단하게 요약하자면 모델을 해석 할 때 training dataset에서 존재하지도 않는 (변형된) 데이터셋을 이용해서 모델을 해석하려고 하는 시도가 문제가 있다는 내용이다. 관련된 논문으로는 Hase et al., The Out-of-Distribution Problem in Explainability and Search Methods for Feature Importance Explanations, NeurIPS 2021이 있다. 원문에서도 더욱 자세한 discussion이 appendix에 기술되어있다.

[10] (역자) Residual stream에 대해서는 Anthropic의 이 글 혹은 우리의 번역본을 참고하면 된다.

[11] 이는 모델의 원래 loss (0.0003)을 100%, 무작위로 label을 뒤섞었을 때의 loss를 0%로 잡는 방식으로 측정되었다. 이에 대한 자세한 논의는 appendix에 있다.

[12] (역자) Induction head에 대한 자세한 내용은 역시 Anthropic의 이 두 글 (글 1, 글 2) 혹은 우리의 번역본을 참고하면 된다.

[Ablation] Ablation study는 어떤 아키텍쳐를 이루고 있는 여러 부품 (혹은 모듈들)을 하나씩 (혹은 여러 개를 동시에) 제거하거나 무효화하면서 그 반응을 살펴보는 실험을 의미한다. ML 논문에서도 분석에 자주 쓰이는 툴이며, 직역하면 제거나 절개에 가깝다.

[ELK] Eliciting latent knowledge (한글로 직역하면 ‘잠재 지식의 이끌어냄’ 정도가 될 것 같다)은 고도로 발달된 AI 모델이 보유한 숨겨진 또는 암묵적인 지식을 추출하고 전달하는 문제를 말한다. 이는 특히 alignment의 개념에서 인간이 바라보는 방식과 일치하도록 모든 관련 정보를 정확하게 보고할 수 있도록 보장하는 데 중점을 두고 있다.

이해를 돕기 위한 예시로, 보안 목적으로 설계된 가상의 AI 모델에 ELK 문제가 발생할 수 있는 시나리오를 생각해보자. 이 AI 모델은 금고 안을 바라볼 수 있도록 다양한 카메라 및 센서들로부터 정보를 읽어내며, 적절한 학습을 통해 (예를 들어 강화학습이라던지) 외부 침입자가 금고 안으로 들어왔을 때 귀중한 물품을 훔치지 못하도록 조치를 취할 수 있는 의사결정까지 할 수 있는 수준까지 올라왔다고 가정하자. 여기에서 만약 침입자가 카메라나 센서에 접근해서 강제로 정상 시그널을 보내도록 강제했다고 하자. AI 모델 내부적으로는 anomaly detection 등으로 인해 지금이 anomaly 상황 (센서가 보내는 신호를 믿을 수 없음)인 것은 인지하고 있으나 이 정보를 관리자에게 알리거나, output에 반영하지 못하는 경우 문제가 생길 수 있다. 즉, 내부(잠재)적으로는 비정상 상황에 대한 인식이 있으나, 이를 이끌어내도록 설계가 되어있지 않아 인간의 의도와 다르게 행동하는 결과가 있을 수 있다. 자세한 사항은 다음 문서를 참조하도록 하자.

[Extensionally equal] 겉으로 보이는 결과들에 대한 관찰들을 기반으로 두 함수가 같은지 판단하는 방식이다. 여기에서는, input domain $X$ 에 대해서 두 함수 $f$ , $g$ 가 같은 output을 내뱉으면 사실상 두 함수가 같다고 판단하는 것이다. (위키피디아)

[Faithfulness] 원문에서는 [4]로 링크를 걸어두고 있다. Faithfulness은 ‘충실도’ 정도로 직역되며, XAI 문헌에서 혼용되어 자주 나타나는 용어인 fidelity는 ‘(특히 배우자에 대한) 신의/정절’으로 번역되는 것을 보면 그 의도를 짐작할 수 있다 (참고로 이런 용어들의 혼용 문제는 저 논문에서도 지적되고 있다.~~역자도 솔직히 어느 순간부터는 굳이 구분하지 않고 있다~~). XAI에서 faithfulness는, 어떤 딥러닝 모델에 대한 해석이 존재할 때, 그 해석이 실제로 그 모델의 작동원리를 얼만큼 반영할 수 있는가를 측정한다. 사견을 붙이자면, 모델에 대한 해석이라고 이런 성질을 지니는 것이 당연한게 아닌가?라고 생각할 수 있지만, XAI 문헌 전반적으로 ‘생각보다는’ 그렇게 강조되고 있지는 않다. 사실 faithfulness 보다 더 강조되는 것은 accuracy로, 마치 classification 문제에서 모델의 결과물이 ground truth label(class)와 얼만큼 일치하는가를 중요하게 보듯이, 해석의 결과물이 얼마나 ground truth explanation과 일치하는가를 보는 것이다. 하지만 explanation accuracy만 강조하게 되면 큰 함정에 빠질 수 있는데, 실제 모델이 학습이 충분히 이루어지지 않아 성능이 떨어지는데도 ‘해석 모델’의 성능이 너무 좋거나 많은 tuning을 통해 그럴싸한 explanation을 도출할 수 있는 가능성이 있다. 이런 문제는 [5]에서도 일부 지적하고 있다.

[Graph homomorphism] 대략적으로, 두 그래프 $G$ 와 $H$ 사이에서, $G$ 에 존재하는 노드들에서 $H$ 에 존재하는 노드들로 연결 관계를 해치지 않고 대응시킬 수 있다면 graph homomorphism이다. Graph isomorphism을 포함하는 더 넓은 개념이며, $G$ 가 $H$ 보다 ‘크기가 더 커도’ 성립할 수 있다 (예시) (엄밀한 정의).

PreviousIn context learning and Induction heads 이해하기 Next[기계적 해석 기초] 다의성(Polysemanticity)과 Sparse Autoencoder 설명

Last updated 7 months ago

Causal Scrubbing a method for rigorously testing interpretability hypothesis [Redwood Research] 이해하기

Original authors: by LawrenceC, Adrià Garriga-alonso, Nicholas Goldowsky-Dill, ryan_greenblatt, jenny, Ansh Radhakrishnan, Buck, Nate Thomas

Translated by: Yong-Min Shin

원문 링크: https://www.lesswrong.com/s/h95ayYYwMebGEYN5y/p/JvZhhzycHu2Yd57RN

원 번역 글 Notion link: Causal Scrubbing a method for rigorously testing interpretability hypothesis [Redwood Research] (요약 번역) - coming soon

요약

Introduction

우리의 contribution들은 다음과 같다:

우리는 다양한 mechanistic 해석들을 더욱 formal하게 표현할 수 있는 interpretatility hypothesis에 대한 개념을 정의하였다.
우리는 causal scrubbing, 즉 순차적으로 activation들을 모델의 성능에 영향을 미치지 않으면서 교체해나가는 검증 방법론을 제시하였다.
우리는 이 방법론의 두 가지 적용 예시를 보여주어 일반적으로 적용이 가능하다는 것을 보여준다.

(생략 — 원문을 통해 확인하도록 하자)

Setup (셋업)

데이터셋 $D$ : 도메인 (정의역) $X$ 에서 정의됨
(분석 대상이 되는) 함수 $f: X \rightarrow \mathbb{R}$
(분석 대상이 되는) 행동 패턴: $\mathbb{E}_{x \sim D}[f(x)]$

이 행동 패턴을 기술하는 ‘가설’ $h$ 는 세 가지 요소의 tuple $(G, I, c)$ 로 정의된다:

Computation graph $G$ : 여기서 $G$ 는 모델 $f$ 와 (사실상) 같으며 [Extensionally equal], 모델의 계산이 실제로 어떻게 진행되는지를 기술한다.
Computation graph $I$ : 이는 모델이 아닌 모델에 대한 해석을 기술한다.
대응 관계를 나타내는 함수 $c$ : $G$ 와 $I$ 의 노드끼리 대응해주며, injective graph homomorphism이어야 한다 (즉, $I$ 에 존재하는 모든 edge는 반드시 $G$ 에 존재하며, $G$ 에만 존재하는 edge도 있다. 또한 graph structure가 대응 관계에서 망가지지 않는다).

추가적인 조건으로 $G$ 와 $I$ 모두 input과 output을 나타내주는 노드가 딱 한 개 씩 존재하도록 요구한다.

다음은 가설의 한 예시이다:

몇 가지 주목할만한 사항들은 다음과 같다:

앞서 언급했다시피, $I$ 는 가설을 나타내주는 일종의 heuristic이라고 보면 된다. 이 가설은 완벽하지 않을 가능성도 충분히 있고, label $y$ 과도 일치하지 않을 경우도 있다. 어차피 모델을 표현하는 것이 가설의 목적이기 때문에, 후자의 경우 원래 모델도 같은 양상을 보이기만 하면 문제가 없다.
하지만 $c$ 가 각 노드의 activation 값들이 어떻게 $I$ 안에서 표현되는지는 알려주지 않는다. 다만 (개념적으로) 어느 노드에 대응이 되는지만 알려줄 뿐이다.
$c(I)$ 를 $G$ 의 “중요한 노드”들이라고 부를 것이다.
- 각 그래프에서 두 개의 대응되는 노드 쌍 $n_I$ , $n_G$ 를 보자 (즉, $c(n_I)=n_G$ ).
  - 같은 input에 대해 두 그래프 $G$ 와 $I$ 를 돌렸다고 생각하자. 그러면 $n_G$ (activation)이 $n_I$ (input에 대한 간단한 로직)을 표현해준다고 생각할 수 있다.
  - Causal scrubbing 알고리즘은 약간 더 약한 claim에 대한 증명을 할 것이다. 그 claim은 “ $n_I$ 에 들어가는 input에 대한 equivalence class는 $n_G$ 에 들어가는 input에 대한 equivalence class와 같다”이다. 이 claim에 대한 검증이 (틀린 가설을 모두 소거하기엔 힘들지라도) 우리가 mechanistic interpretability 안에서 관심을 가지는 가설들을 검증하기에 충분하다고 생각한다.
우리는 모델을 나타내주는 computation graph $G$ 를 때때로 더 보기 편한 방식으로 표현할 것이다.

‘가설’에 대한 더 심도있는 논의는 appendix를 보도록 하자.

Causal Scrubbing

여기서 우리는 causal scrubbing를 두 가지 다른 방식으로 설명하겠다.

조금 덜 포멀하고 직관적인 버전: Causal scrubbing의 핵심 아이디어를 이해할 수 있도록 많은 그림들과 함께 설명할 것이다.
Causal scrubbing 알고리즘 및 pseudocode.

독자마다 잘 맞는 설명 방식이 다 다를 것이기에, 이 글을 여기저기 옮겨다니면서 훑어보는 것을 추천한다.

개념적인 설명: 가설은 어떤 activation들을 (마음대로, 혹은 적절한 것으로) 교체해도 된다고 말하고 있을까?

어떤 가설 $h = (G, I, c)$ 가 있다고 가정하자 (그림 참조).

만약 이 가설이 참이라면, 우리는 다음 두 가지 ‘resampling ablation’을 진행해도 모델에 변화가 없을 것이다:

Activation A, B, D에 대해 가설 $I$ 의 논리 안에서 동일한 다른 input에서 나온 activation으로 교체하거나,
Activation C, 혹은 $z_1 \rightarrow B$ 와 같은 부분에 대해서 아무 input에 대한 activation으로 교체하거나 (가설에 의하면 중요하지 않은 부분이라고 했으므로).

Intervention 1: semantically equivalent subtrees

사실 우리의 가설에 의하면 이 뿐만 아니라 여러 activation 교체들이 가능하다. 예를 들어, D와 관련해서 이런 교체도 가능할 것이다:

Intervention 2: unimportant inputs

그림으로 이를 표현하자면 다음과 같을 것이다:

The causal scrubbing algorithm

알고리즘에서는 treeify를 직접적으로 진행하지는 않는다; 하지만 마치 tree 구조를 지닌 것 처럼 경로를 하나씩 탐색하게 된다.

직관적인 버전

(이 부분은 pseudocode를 이해하고 있다면 거의 중복되는 내용이므로 선호하는 방식대로 읽는 것을 추천한다.)

알고리즘의 pseudocode는 다음과 같이 쓸 수 있다:

어떤 기준이 되는 (reference) input $x$ 를 $D$ 로부터 샘플한다.
run_scrub을 재귀적으로 호출하면서 $I$ 의 output으로부터 input까지의 모든 경로를 훑어본다. 각 노드마다, 그 노드 ( $n_I$ 라고 하자)보다 ‘윗 쪽에 존재하는’ 모든 노드들을 모은 subgraph를 생각해볼 수 있다 (즉 이 노드를 계산하는데 영향을 미치는 모든 노드들이 된다). 그리고 이 subgraph들은 실제적으로 $G$ 에도 해당되는 부분이 $c(I)$ 안에 존재할 것이다.
run_scrub(n_I, c, D, x)의 리턴값은 $G$ 안에서의 어떤 activation이다. 구체적으로, 가설에 의해 input이 x일 때 $n_I$ 의 값에 대응되는 $G$ 안에서의 노드가 될 것이다.
- $n_G = c(n_I)$ 라고 하자.
- 만약 $n_G$ 가 input 노드라면 리턴값은 그냥 $x$ 가 될 것이다.
- 그렇지 않다면 $n_G$ 로부터 부모 노드들로 거슬러 올라가면서 input까지 activation들을 계산할 것이다. (가설 단에서) $p_I$ 를 $n_I$ 의 부모 노드, (실제 computation graph $G$ 에서) $p_G$ 를 $n_G$ 의 부모 노드라고 할 때, 다음을 수행한다:
  - $p_G$ 에 대응하는 $p_I$ 가 있는 경우: 그렇다면 가설은 $p_G$ 의 값이 $n_G$ 의 값에 유의미한 영향을 미친다고 생각한다는 것이다. 따라서 가설에 의해 x와 같은 효과를 가지는 (즉 $p_I$ 에서 결과가 같음) 새로운 데이터 new_x를 샘플한다. 그리고 run_scrub를 $p_I$ 에 대해서 재귀적으로 다시 부른다.
  - 대응하는 부모 노드가 없는 “중요하지 않은” 노드인 경우, other_x라는 다른 input을 선택한다. 이 input은 데이터셋으로부터 무작위로 선택되었다. 하지만, 이 경우로 다다를 때마다 다른 데이터를 무작위로 고르기 보다는 한 번 무작위로 골라진 input을 계속 고정해서 사용하도록 한다 [7]. 이 other_x를 사용했을 때의 $p_G$ 값을 따로 기록해둔다.
- 이제 $n_G$ 의 모든 부모 노드들에 대해 activation 값들을 계산하였다 — 그리고 이에 따라 $n_G$ 에 대한 개입을 한 결과를 얻기 위한 input도 얻게 되었으니, 이 input을 통해 모델을 돌리고 함수의 리턴값 얻어낸다.

Pseudocode: 원문 링크

왜 resampling을 통한 ablation을 해야 하는 것일까?

중요하지 않은 노드에서 “중요하지 않다”는 건 무슨 의미인가?

Zero ablation: 해당 모듈의 activation들을 모두 0으로 고정
Mean ablation: 해당 모듈의 activation들을 데이터셋 전체에 대한 empirical mean으로 대체
Resampling ablation: 해당 모듈의 activation을 임의의 다른 input에 대한 activation으로 대체

이들 중 가장 적절한 것이 무엇인지 알아내기 위해, 우리가 ablation을 통해 정확히 무엇을 검증하려고 하는지 다시 한 번 생각해보자.

Zero 및 mean ablation의 문제점

다른 선행연구에서 zero 및 mean ablation이 많이 쓰였지만 [8], 우리가 보이고자 하는 주장들과 결이 맞지 않는 부분들이 있다.

Zero 및 mean ablation은 모델을 예측할 수 없는 데이터 분포 위에서 작동하게 만든다.
Zero 및 mean ablation은 모델 성능에 예측할 수 없는 영향을 미칠 것이다.
Zero 및 mean ablation은 variance를 없애며 모델의 행동 패턴을 왜곡할 것이다.

Results

Parentheses balance checker

(역자 2: Input은 ‘()())()()))’과 같이 생겼다)

주장 1: Output으로 주요 정보를 직접적으로 전달하는 head가 3개 존재한다.

Head 1.0과 2.0은 합동해서 1) 열린 괄호 ‘(’와 닫힌 괄호 ‘)’의 총합이 같은지와 2) 열린 괄호로 시작하는지를 확인한다.
Head 2.1은 nesting depth가 음의 값을 가지는 시점이 존재하는지 확인한다.

주장 1은 다음 가설로 표현될 수 있다:

주장 2: Head 1.0과 2.0은 input의 첫 번째 position에만 의존하며, 이 input은 간접적으로 다음에 의존한다:

Position 1에서 head 0.0의 output. 이 부분은 전체적으로 열린 괄호의 비중이 얼마인지를 계산한다. 이 정보는 residual stream에 특정한 방향을 가지는 벡터로 쓰여(write)진다.
Position 1의 embedding. 이는 sequence가 열린 괄호로 시작하는지에 대한 정보를 담고 있다.

이 세 주장을 종합하여 그림으로 표현하면 다음과 같을 것이다:

이 주장들을 causal scrubbing으로 검증을 해본 결과, 유의미하지만 완전하지는 않은 해석을 제공하는 것으로 보인다:

검증하는 주장

재현된 모델의 행동 비중 [11]

93%

1 + 2

88%

1 + 3

84%

1 + 2 + 3

72%

On induction heads

일단, induction head의 attention pattern가 (하나의 token이 아닌) 연속적으로 붙어있는 3개의 token을 비교하는 것이 가능하도록 수정하였다.
그리고, query와 value 값이 현재 position에 해당하는 layer 0의 결과값에도 영향을 받을 수 있도록 수정하였다.
마지막으로 특정 케이스에 대해여 현재 token이 반복적으로 나타날 경우에 대해 반응하는 layer 0 head를 고려하였다.

이렇게 수정한 결과, 모델 성능의 86%까지 재현하는데 성공하였다.

관련한 더 자세한 내용은 appendix 글을 참고하면 되겠다.

Alignment와의 관계에 대하여

사람이 발견해낸 모델 해석의 검증: Mechanistic interpretability 연구의 규모를 키우는 데에 (causal scrubbing과 같은) 표준화된, 신뢰있는, 그리고 간편한 검증 절차가 있는 것은 많은 도움이 될 것이다; 이는 특히나 혁신적인 AI 모델 개발 전에 interpretatility 프로젝트 같은 것이 있을 때 도움이 될 것으로 보인다.
해석 모델 검출의 자동화: 어떤 경우에서는 연구자가 가설 후보들의 space를 규정하고 optimization 알고리즘을 이용하여 가장 적절한 해석을 찾으려고 할 수 있다. 우리는 이와 유사한 결의 방향성을 가지는 일을 하였고 이를 추후에 더 발전시킬 것이다.
AI-assist된 해석: 이를 이용하여 높은 퀄리티의 이해 가능한 해석을 내놓을 수 있는 모델을 학습할 수 있을 것이다.

ELK를 풀기 위한 자동화된 방법론들의 개발 [ELK]: ARC (Alignment Research Center)는 ELK를 해결하는 데에, (굳이 인간이 이해하지 않아도 되는) causal scrubbing으로 찾아낸 해석과 관련된 모듈들을 학습에 이용하는 몇몇 방향성들에 대해 긍정적으로 바라보고 있다.
함정 alignment 찾아내기: 만약 약간 믿을 수 있는 모델과 매우 믿을 수 없는 모델 두 개가 있다고 하자. 두 모델이 비슷한 행동을 한다고 쳤을 때, 매우 믿을 수 없는 모델이 왜 (아마 겉으로 보기에만) align이 되어있는지 직접 모델을 탐구하지 않고 해석을 통해 알아낼 수 있을 것이다.
Relaxed adversarial training에서, 모델의 내부를 바라보고 예상치 못한 행동을 할지를 판단해야하는 상황이 생긴다 (이는 보통 random sampling으로는 찾기 힘든 행동이며, 아주 드물게 training 분포 바깥에서 일어난다). 이 과정에서 내부 모습을 인간이 그 의미를 이해하는 것이 필수로 요구되지는 않는다. 단지 causal scrubbing을 이용해서 밝혀낸 내부 구조를 이용해서 모델의 행동을 예측할 수 있으면 된다.

Limitations

불행히도, causal scrubbing은 우리가 다루고 싶은 모든 interpretabiltiy 가설을 다루는 것은 아니다:

Causal scrubbing에서는 가설이 완벽하게 허용하는 범위 안에서만 activation 교체를 진행할 수 있다.
- 가설에 의하면 어떤 activation을 교체할 수 있는지를 엄격하게 보여주고 있지만, 실제에서는 이 범위 바깥을 벗어나는 가설들도 허용할 때가 있다.
- 예를 들어, 모델의 어떤 부분에서 연속적인 (continuous) 피쳐를 만들어낸다고 가설을 세웠을 때, 이 부분에 대한 activation을 다른 input을 통해 재현하려고 하면 아마 이를 정확하게 재현하는 것은 불가능하고 이와 비슷한 값을 목표로 해야할 것이다. 엄밀히 말하면, 이는 허용이 되지 않는다.
- 이를 해결하는 한 방법으로는 causal scrubbing을 조금 더 일반화해서, 가설 space를 continuous하게 취급을 하는 것이다 (이는 여러가지 장점이 있는데, 예를 들어 이 space 위에서 SGD를 돌릴 수도 있을 것이다). 이에 대해 우리는 기초적인 작업을 진행했지만, 아직 개념적으로 해결해야할 문제들이 남아있다.
Causal scrubbing은 제공된 data 안에서만 검증을 진행할 수 있다. 만약 causal scrubbing을 하다가 현재 input와 비슷한 성질을 가진 다른 input을 사용해야하는데, 그런 data가 존재하지 않는다면 더 이상 검증을 할 수 없다.
데이터셋 안에서 완벽한 상관관계를 보이는 두 변수가 있다면, causal scrubbing은 그 두 변수를 헷갈려하더라도 크게 티가 안날것이다.

또다른 한계점으로는

모델한테 모든 상황에 적용하는 것이 불가능한 heuristic이 있다고 할 때, 그 heuristic을 억제하기 위해 다른 circuit을 사용할수도 있다 (예를 들어, Indirect object identification 논문에 나오는 negative name mover head). 하지만 이런 억제 circuit들은 heuristic들이 적용 가능할 때는 방해가 될 뿐이다. 이런 경우, 억제 circuit들을 무시하게 될 경우, 그 heuristic이 모델 성능에 미치는 영향을 과대평가하게 될 수도 있고 불완전한 해석에 다다르게 될수도 있다.
만약 두 항 간의 상관관계가 높다면, (무작위로 activation을 교체하는 방식으로) 그들을 독립적으로 샘플링하면 합의 variance가 줄어들게 된다. 어떤 경우에는, 이 variance가 모델 성능에 악영향을 미칠 때도 있다 — polysemanticity와 관련된 예시. 이렇게 되면 이런 상관관계가 포함되어있지 않은 가설로 귀결될 가능성이 있고, causal scrubbing에 의해서 더 ‘정확한’ 가설로 보여질 것이다.