[LLM 해석가능성] ROME: 피드포워드 네트워크(FFN)의 메모리 편집하기

0. Introduction

고정된 메모리를 편집할 수 있을까?

이전 글에서 FFN이 key-value memory로 기능한다는 것을 확인했습니다. $W_1$ 의 행이 key, $W_2$ 의 열이 value로, 훈련 후 고정된 사실 저장소 역할을 합니다.

이전 글의 마지막에 두 가지 질문을 남겼습니다:

핵심 질문

어디에?"에펠탑은 파리에 있다"라는 사실이 48개 층의 FFN 중 어느 층에 저장되어 있는가?

편집?찾은 위치의 가중치를 수정하여, 모델이 "에펠탑은 로마에 있다"고 믿게 만들 수 있는가?

Meng et al.의 논문 "Locating and Editing Factual Associations in GPT" (NeurIPS 2022)은 이 두 질문에 각각 Causal Tracing과 Rank-One Model Editing (ROME)으로 답합니다.

이 글의 여정

1. Causal Tracing: 사실의 위치 찾기

2. Tracing 결과: 두 개의 결정적 site

3. ROME: 사실 편집하기

4. 편집 결과 평가

5. 정리 + 다음 단계

1. Causal Tracing 방법론

인과적 개입으로 사실의 위치를 찾다

사실을 knowledge tuple $t = (s, r, o)$ 로 표현합니다: subject, relation, object. 예: (Space Needle, located in, Seattle).

프롬프트 "The Space Needle is in downtown ___"을 넣으면 모델이 "Seattle"을 예측합니다. 모델 내부의 어떤 구성요소가 이 사실을 '기억'하는 데 결정적인 역할을 하는지 찾고 싶습니다. 이를 위해 세 번 모델을 실행합니다.

Step 1: Clean Run

프롬프트를 정상적으로 넣고, 모든 층의 모든 hidden state를 기록합니다:

\{h_i^{(l)}\}

. 이것이 정답을 알고 있는 상태의 스냅샷입니다.

Step 2: Corrupted Run

Subject 토큰("Space", "Needle")의 임베딩에 가우시안 노이즈를 추가합니다. 모델이 subject를 제대로 인식하지 못해 오답을 출력합니다.

h_i^{(0)} := h_i^{(0)} + \epsilon, \quad \epsilon \sim \mathcal{N}(0, \nu)

Step 3: Corrupted-with-Restoration

Corrupted 상태에서 모델을 돌리되, 특정 층, 특정 토큰 위치의 hidden state만 Clean Run의 값으로 교체합니다. 이 하나의 복원만으로 정답이 돌아오면 → 그 위치가 인과적으로 결정적입니다.

왜 '인과적'인가?

단순히 gradient가 크다거나 activation이 크다는 것으로는 해당 구성요소가 정말로 결과에 영향을 미치는지 알 수 없습니다 (상관관계 ≠ 인과관계). Causal Tracing은 직접 개입(intervention)하여 결과가 바뀌는지를 확인하므로, 인과적 결론을 내릴 수 있습니다.

\text{IE}(h_i^{(l)}) = \mathbb{P}_{*, \text{clean } h_i^{(l)}}[o] - \mathbb{P}_*[o]

Corrupted Run에서 subject 토큰에만 노이즈를 추가하고, 나머지 토큰('is', 'in', 'downtown')은 그대로 둡니다. 왜 subject만 오염시킬까요?

우리가 찾고 싶은 것은 subject에 대한 사실 지식이 어디에 저장되어 있는가입니다. 'is in downtown'은 relation을 나타내는 부분으로, 어떤 종류의 사실을 물어보는지를 알려줍니다. Subject만 오염시키면, 모델이 "something is in downtown ___"이라는 구조는 이해하되 무엇에 대한 사실인지를 모르는 상태가 됩니다. 이 상태에서 특정 위치를 복원했을 때 정답이 돌아오면, 그 위치에 subject의 정체성과 관련된 사실이 저장되어 있다는 것입니다.

2. Causal Tracing 결과 해석

두 개의 결정적 site

Meng et al.은 이 실험을 1000개의 사실 진술에 대해 수행하고, 결과를 평균하여 Average Indirect Effect (AIE) 히트맵을 만들었습니다.

핵심 발견 #1 — Early Site

중간 층(GPT-2 XL에서 약 15~18번째 층)에서, subject의 마지막 토큰 위치에서 강한 인과 효과가 나타납니다. 이 효과를 MLP 기여분과 Attention 기여분으로 분리하면, MLP가 지배적입니다.

의미: 중간 층의 MLP가 subject를 처리하면서 해당 subject에 대한 저장된 사실을 recall하고 있습니다. 이것은 FFN이 key-value memory로 기능한다는 Geva et al.의 해석과 정확히 부합합니다.

핵심 발견 #2 — Late Site

마지막 토큰 위치의 높은 층에서도 강한 인과 효과가 나타나는데, 여기서는 Attention이 지배적입니다.

의미: 높은 층의 Attention이 MLP가 recall한 사실 정보를 마지막 토큰 위치로 복사하여 출력에 사용합니다.

전체 메커니즘 가설

① 중간 층 MLP가 subject 토큰을 처리하며 사실을 recall → residual stream에 쓰기② 높은 층 Attention이 이 정보를 마지막 토큰으로 복사 → 출력에 반영이것이 Meng et al.이 제안하는 Localized Factual Association Hypothesis입니다.

autoregressive 모델에서 정보는 왼쪽에서 오른쪽으로만 흐르기 때문에, subject의 마지막 토큰이 subject 전체의 정보를 집약하는 위치가 됩니다.

Early site에서 MLP가 지배적이고, late site에서 Attention이 지배적이라는 것은, Transformer 블록의 Attention → FFN 순서와 어떤 관련이 있을까요?

각 층에서 Attention이 먼저 문맥 정보를 모으고, 그 다음 FFN이 모인 정보를 바탕으로 지식을 recall합니다. 중간 층에서 이 recall이 일어나고, 높은 층에서는 recall된 정보가 이미 residual stream에 있으므로, Attention이 그것을 최종 출력 위치로 운반하는 역할에 집중하는 것입니다.

3. ROME의 개념적 원리

Rank-One Model Editing

Causal Tracing으로 중간 층 MLP가 사실 recall의 핵심이라는 것을 확인했습니다. 이제 그 MLP의 가중치를 직접 수정하여 새로운 사실을 삽입해봅시다.

이전 포스트를 복기하면: $W_2$ (논문에서는 $W_{proj}$ )가 value 행렬입니다. 임의의 선형 변환 $W$ 는 key 집합과 value 집합에 대해 $WK \approx V$ 를 만족하는 연관 메모리(associative memory)로 볼 수 있습니다. 여기에 새로운 key-value 쌍 $(k^*, v^*)$ 를 하나 추가하고 싶습니다.

Step 1: Key 선택 (k*)

"무엇에 대한 사실인가?"
Subject("Space Needle")를 나타내는 벡터를 key로 사용합니다. 구체적으로: subject를 포함하는 여러 문장을 모델에 넣고, 중간 층 MLP 입력에서의 activation을 평균합니다.

직관: "다양한 문맥에서 'Space Needle'이 나타날 때의 공통적인 내부 표현"

Step 2: Value 최적화 (v*)

"어떤 사실로 바꿀 것인가?"
새로운 object("Paris")를 예측하게 만드는 벡터

v^*

를 최적화로 찾습니다. 이 벡터를 MLP 출력으로 삽입했을 때, 모델이 'Paris'를 높은 확률로 예측하도록 합니다.

추가 제약: subject의 본질(essence)이 변하지 않도록 KL divergence 항을 추가합니다. 예: "Space Needle은 건축물이다"라는 속성은 유지하면서 위치만 바꿉니다.

Step 3: Rank-One Update

새로운 사실을 삽입합니다.

\hat{W} = W + \Lambda (C^{-1} k^*)^T

여기서

\Lambda = (v^* - W k^*)

C = KK^T

는 기존 key들의 공분산입니다.

핵심: rank-one update = 행렬 전체를 바꾸는 것이 아니라, 한 쌍의 벡터 외적(outer product)을 더하는 것입니다. 비유하자면 "백과사전에서 한 항목만 수정하는 것이지, 책 전체를 다시 쓰는 것이 아닙니다."

왜 'Rank-One'인가?

행렬에 rank-1 행렬(

\Lambda u^T

, 두 벡터의 외적)을 더하면, 행렬의 한 방향만 변경됩니다.

k^*

방향으로 들어오는 입력에 대해서만 출력이 바뀌고, 다른 방향의 입력에 대해서는 원래 행렬과 거의 동일하게 동작합니다. 이것이 specificity(다른 사실은 건드리지 않음)를 보장하는 수학적 원리입니다.

ROME은 48층 중 하나의 층만 수정합니다. 논문에서는 몇 번째 층을 선택했을까요? 그리고 그 근거는 무엇일까요?

GPT-2 XL에서 18번째 층을 선택했습니다. 이것은 Causal Tracing에서 MLP의 인과 효과가 가장 강했던 early site의 중심 층입니다. 즉, Causal Tracing의 결과가 ROME의 편집 위치를 직접적으로 결정합니다.

4. 편집 결과 평가

진짜로 사실이 바뀐 것인가?

"사실을 편집했다"고 주장하려면, 세 가지를 확인해야 합니다:

Efficacy (효과) — 편집한 프롬프트에 대해 새로운 답을 내는가?

Generalization (일반화) — 다른 표현으로 물어봐도 새로운 답을 내는가?

Specificity (특이성) — 관련 없는 다른 사실은 그대로 유지되는가?

왜 ROME만 세 가지를 동시에 달성하는가?

ROME은 모델이 사실을 저장하는 정확한 메커니즘(FFN key-value memory)을 이해하고, 그 메커니즘에 맞춰 편집하기 때문입니다. Fine-Tuning은 메커니즘을 모른 채 출력만 맞추려 하므로 부작용이 생기고, Hypernetwork 방법은 표면적 패턴만 학습하므로 일반화가 안 됩니다.

Fine-Tuning이 specificity를 희생하는 이유를 FFN-as-memory 관점에서 설명할 수 있을까요?

Fine-Tuning은 loss를 줄이기 위해

W_{proj}

의 많은 방향을 동시에 변경합니다. 이것은 메모리의 여러 슬롯을 동시에 수정하는 것과 같아서, 편집하려는 사실과 관련 없는 사실까지 영향을 받습니다. 반면 ROME은 rank-one update로 딱 한 방향만 바꾸므로,

k^*

방향과 무관한 key들에 대해서는 행렬이 거의 동일하게 동작합니다.

5. 요약

핵심 정리와 다음 여정

Causal Tracing

Clean → Corrupt → Restore로 인과적 중요도 측정

두 개의 Site

Early site = 중간층 MLP (사실 recall), Late site = 높은층 Attention (정보 복사)

ROME

FFN의 W_proj에 rank-one update로 사실 편집

평가 기준

Efficacy + Generalization + Specificity를 동시에 달성

두 포스트를 관통하는 서사

Geva et al.은 FFN이 key-value memory라는 구조적 해석을 제시했고,

Meng et al.은 이 구조에서 사실의 위치를 인과적으로 특정하고 편집까지 수행함으로써,

FFN이 정말로 사실적 지식의 저장소로 기능한다는 것을 실험적으로 검증했습니다.

구조적 해석 → 인과적 검증 → 직접 편집이라는 흐름은, interpretability 연구의 전형적인 방법론적 패턴입니다.

다음 단계 미리보기

Causal Tracing은 '이 위치가 중요하다'는 것을 보여주지만, 정보가 어떤 경로를 통해 흐르는지는 알려주지 않습니다.

Week 3의 Path Patching은 이 한계를 넘어서, 특정 구성요소 간의 경로별 인과 효과를 측정합니다.