LLM : 빠른 논문 리뷰 : Locating and Editing Factual Associations in GPT

Paper Study Note: Locating and Editing Factual Associations in GPT

---

Terminology (용어 설명)

• Causal Tracing (Causal Mediation Analysis): 모델 내부의 특정 hidden state activation이 최종 factual prediction에 미치는 인과적 영향(Indirect Effect)을 측정하기 위해, 정상 실행(clean run), 노이즈를 주입한 실행(corrupted run), 특정 state만 정상으로 복구한 실행(corrupted-with-restoration run)의 세 가지 경로를 비교하는 기법.
• Early Site & Late Site: Causal Tracing 결과에서 도출된 두 가지 핵심 연산 위치. Early site는 subject의 마지막 token이 처리되는 중간(middle) layer의 MLP 모듈 위치이며, Late site는 최종 출력 직전 마지막 token이 처리되는 후반(late) layer의 Attention 모듈 위치임.
• Linear Associative Memory: 행렬 W가 key 벡터 K와 value 벡터 V 사이의 선형 매핑(WK ≈ V)을 통해 정보를 저장하는 key-value 메모리로 작동한다는 관점.
• ROME (Rank-One Model Editing): Transformer의 중간 layer MLP projection weight를 Linear Associative Memory로 간주하고, 제약 조건이 있는 최소제곱법(constrained least-squares)을 풀어 rank-one 행렬 업데이트만으로 새로운 사실(k*, v*)을 직접 삽입하는 모델 편집 방법론.
• Essence Drift: 모델의 지식을 편집할 때 목표로 한 사실만 바뀌는 것이 아니라, 해당 subject의 근본적인 정체성이나 본질(essence)까지 훼손되거나 왜곡되는 현상 (예: 게임 이름을 마이크로소프트가 만들었다고 편집했더니 아예 게임을 오피스 소프트웨어로 묘사함).
• Bleedover: 특정 subject의 사실을 수정했을 때, 그와 유사한 latent space에 위치한 다른 연관 subject들까지 잘못된 새로운 object로 예측하게 되는 부작용(Specificity 저하 현상).
• Regurgitation: 모델이 새로운 사실의 의미를 깊이 이해하고 일반화(Generalization)한 것이 아니라, 단순히 프롬프트에 대해 목표 단어만 기계적으로 뱉어내는 표면적 암기 현상.

---

Purpose of the Paper

• 기존 연구의 한계 극복: 기존 지식 위치 파악 및 model editing 연구들은 주로 양방향 문맥을 보는 masked language models (BERT 등)에 집중되어 있었으며, 단방향(unidirectional) attention과 generation 능력을 갖춘 대규모 autoregressive GPT 모델들이 내부적으로 어떻게 factual association을 저장하고 인출하는지는 거의 규명되지 않았음.
• 방법론적 차별점: 기존 probing 방식은 네트워크의 실제 행동과 분리되어 단순 상관관계에 의존했고, gradient 기반 saliency 기법(Integrated Gradients 등)은 국소적 특성에만 민감해 거시적인 정보 흐름을 파악하지 못했음. 또한 기존 fine-tuning이나 hypernetwork 기반 editing 기법들은 Generalization(과적합으로 인한 실패)과 Specificity(언더피팅 및 Bleedover 발생) 중 하나를 희생해야만 했음.
• 새로운 접근 제시: 인과 추론(Causal Mediation Analysis)을 통해 GPT 내 지식 인출의 정확한 메커니즘을 규명하고, 이를 바탕으로 개별 뉴런 단위가 아닌 전체 행렬 단위의 Linear Associative Memory 관점을 도입하여 Generalization과 Specificity를 동시에 달성하는 직접적인 parameter 조작(ROME)이 가능함을 증명하고자 함.

---

Key Contributions & Novelty

• Localized Factual Association Hypothesis 정립 및 인과적 규명
  • Contribution: Causal Tracing을 통해 autoregressive 모델에서 사실 인출이 일어나는 구체적인 3차원 위치(MLP 모듈, 특정 중간 layer, subject의 마지막 token)를 정확히 찾아냄.
  • Novelty: 중간 layer의 MLP가 subject를 인코딩한 입력을 받아 암기된 속성을 인출(recall)하고, 후반 layer의 self-attention이 이 정보를 마지막 token으로 복사(copy)한다는 명확한 분업 구조를 인과적으로 최초 입증함.
• ROME (Rank-One Model Editing) 알고리즘 제안
  • Contribution: 중간 layer MLP의 두 번째 가중치 행렬(W_proj)에 대해 closed-form 대수적 해법(W_hat = W + Lambda * (C^-1 * k*)^T)을 도출하여 단일 사실(s, r, o*)을 즉각적으로 삽입하는 기법 개발.
  • Novelty: 기존 연구(Geva et al.)가 MLP를 개별 뉴런 단위의 key-value로 보았던 것과 달리, 전체 행렬을 Linear Associative Memory로 모델링함. 또한 무작위 prefix context를 활용한 key(k*) 평균화 및 KL divergence로 Essence Drift를 억제하는 value(v*) 최적화 공식을 고안하여 기존 데이터 손실 없이 정교한 편집을 달성함.
• 고난이도 벤치마크 데이터셋 COUNTERFACT 구축
  • Contribution: 21,919개의 다양한 counterfactual 기록을 포함한 새로운 평가 데이터셋 COUNTERFACT 및 세부 평가지표 제안.
  • Novelty: 모델이 원래 높은 확률로 예측하던 쉬운 사실만 수정하던 기존 벤치마크(zsRE 등)의 맹점을 극복하기 위해, 발생 확률이 매우 낮은 거짓 사실(counterfactual)을 주입하도록 설계함. 특히 대상 subject와 latent space 상에서 가까운 이웃 엔티티들을 WikiData SPARQL로 수집해 Bleedover를 정밀 측정하고, 생성 텍스트의 Fluency와 Consistency까지 평가하는 다각도 검증 체계를 마련함.

---

Experimental Highlights

• Experimental Setup
  • Models: GPT-2 XL (1.5B, 주 실험 모델), GPT-J (6B), GPT-NeoX (20B), GPT-2 Medium (345M), GPT-2 Large (774M).
  • Datasets: zsRE (10,000 records), COUNTERFACT (21,919 records).
  • Baselines: Fine-Tuning (FT), Constrained Fine-Tuning (FT+L), Knowledge Neurons (KN, gradient 기반 뉴런 편집), Knowledge Editor (KE, hypernetwork), MEND (gradient 분해 기반 hypernetwork).
  • Metrics: Efficacy (ES, EM), Generalization / Paraphrase (PS, PM), Specificity / Neighborhood (NS, NM), 전체 조화평균 Score (S), Fluency (GE, n-gram 엔트로피), Consistency (RS, 생성 텍스트와 참조 텍스트 간 TF-IDF 코사인 유사도).
• Key Results
  • Causal Tracing 검증 (GPT-2 XL, 1000개 사실): Clean run과 Corrupted run의 평균 확률 차이인 Average Total Effect (ATE)는 18.6%로 나타남. Subject의 마지막 token의 15번째 layer hidden state 하나만 복구해도 Average Indirect Effect (AIE)가 8.7%에 달함. Early site에서는 MLP의 기여도(AIE 6.6%)가 Attention(AIE 1.6%)을 압도하며, MLP 경로를 차단(severing)할 경우 초반 layer들의 인과적 효과가 완전히 사라짐을 통해 중간 layer MLP의 필수적 역할을 증명함.
  • zsRE Benchmark 성능: ROME은 GPT-2 XL에서 Efficacy 99.8%, Paraphrase 88.1%, Specificity 24.2%를 기록함. 별도의 대규모 학습 시간이 필요한 hypernetwork(KE, MEND)나 fine-tuning 기법들과 대등하거나 더 뛰어난 성능을 추가 학습 오버헤드 없이 달성함.
  • COUNTERFACT Benchmark 압도적 SOTA 달성:
    • GPT-2 XL (7,500 records): ROME은 종합 Score (S) 89.2를 기록하여 FT (65.1), FT+L (66.9), MEND (57.9), KE (52.2), KN (35.6)을 큰 폭으로 압도함. 특히 Generalization (PS 96.4%)과 Specificity (NS 75.4%)를 동시에 높게 유지한 유일한 모델임 (FT는 NS가 40.4%로 폭락하며 Bleedover 발생, FT+L은 PS가 48.7%로 과적합 발생).
    • GPT-J (6B, 2,000 records): ROME은 종합 Score (S) 91.5 (ES 99.9%, PS 99.1%, NS 78.9%, RS 43.0)를 달성하며 모델 크기가 커져도 일관되게 작동함을 입증함 (MEND는 63.2, FT+L은 68.7, FT는 25.5에 그침).
  • Attention Intervention (AttnEdit) 비교 실험: 후반 layer(layer 33)의 Attention 가중치를 수정한 AttnEdit는 훈련 프롬프트에 대한 Efficacy는 높지만, Paraphrase나 생성 프롬프트에서는 원래 사실을 그대로 말하는 실패(단순 Regurgitation)를 보임. 반면 중간 layer(layer 18) MLP를 수정한 ROME은 완벽한 일반화를 보여주어 가설을 재검증함.
  • Human Evaluation: 15명의 평가자가 50개 시나리오(150개 평가)를 검증한 결과, ROME이 생성한 텍스트가 FT+L 대비 삽입된 사실과 일치(Consistent)한다고 평가받을 확률이 1.8배 높았음. 다만 Fluency 측면에서는 FT+L보다 1.3배 덜 자연스럽다는 평가를 받아 미세한 유창성 저하가 확인됨.

---

Limitations and Future Work

• 단일 사실 편집의 한계와 확장성 (Single-Fact Editing vs. Scalability)
  • Limitation: ROME은 지식 저장 메커니즘을 규명하기 위한 도구로 설계되어 한 번에 하나의 사실만 편집할 수 있으며, 대규모 모델 학습이나 대량의 지식을 동시에 업데이트하는 실용적 용도로는 적합하지 않음.
  • Importance & Future Work: 실제 환경에서는 수천 개 이상의 지식을 동시에 업데이트해야 하므로, ROME의 rank-one 업데이트 아이디어를 확장하여 다수의 메모리를 동시에 대량 편집할 수 있는 후속 연구(MEMIT 등)로 발전시킬 수 있음.
• 지식의 방향성 (Directional Associations)
  • Limitation: ROME이 편집하는 사실 관계는 단방향적임. 예를 들어 "시애틀의 랜드마크는 스페이스 니들이다"와 "스페이스 니들은 시애틀의 랜드마크이다"는 모델 내부에서 별개로 저장되므로, 양방향 지식을 모두 바꾸려면 두 번의 편집이 필요함.
  • Importance & Future Work: 인간과 달리 언어 모델이 대칭적 관계를 통합하여 표상하지 못함을 보여주며, 향후 대칭적/역방향 관계까지 한 번에 일반화하여 수정할 수 있는 표상 구조 연구가 필요함.
• 사실적 연관성에 국한된 분석 (Scope of Knowledge Types)
  • Limitation: 본 연구는 (Subject, Relation, Object) 형태의 단순 사실적 연관성(factual association)만 다루었으며, 논리적(logical), 공간적(spatial), 수치적(numerical) 지식에 대해서는 탐구하지 못함. 또한 사실이 성공적으로 수정되더라도 모델이 증거 없는 그럴듯한 거짓 사실을 추측(guessing/hallucination)하는 현상이 여전히 발생함.
  • Importance & Future Work: 언어 모델을 신뢰할 수 있는 지식 소스로 활용하기 위해서는 단순 팩트 외에 복잡한 추론 및 논리 구조가 네트워크 어디에서 어떻게 연산되는지 규명하고 제어하는 연구로 확장되어야 함.

---

Overall Summary

이 논문은 Causal Tracing을 통해 autoregressive GPT 모델 내부에서 사실적 지식이 'subject의 마지막 token을 처리하는 중간 layer의 MLP 모듈'에 집중적으로 저장 및 인출된다는 것을 인과적으로 규명했습니다. 나아가 이 MLP 가중치를 Linear Associative Memory로 모델링하여, 단 한 번의 rank-one 행렬 업데이트만으로 새로운 사실을 주입하는 ROME 알고리즘과 이를 엄밀히 평가하기 위한 COUNTERFACT 벤치마크를 제시했습니다. 이 연구는 거대 언어 모델의 블랙박스 내부에서 지식이 처리되는 위치(where)와 방식(how)을 명확히 밝혀냈을 뿐만 아니라, 모델의 전체 재학습 없이도 내부 계산 메커니즘을 직접 조작하여 지식을 정교하게 수정할 수 있다는 Mechanistic Interpretability 및 Model Editing 분야의 핵심적인 이정표를 세웠습니다.

---

쉬운 설명 (Intuitive Analogy)

이 논문의 핵심 아이디어는 거대 언어 모델을 **"도서관의 사서 시스템"**으로 비유할 수 있습니다.
질문이 들어왔을 때 모델 내부의 모든 부품이 다 같이 정답을 찾는 것이 아니라, **중간 층에 있는 특정 부서(중간 layer MLP)**가 질문 속 주어의 마지막 글자를 보고 백과사전 서랍에서 정답 데이터를 꺼내는 '핵심 기억 창고(Key-Value Memory)' 역할을 하고, **후반 층의 부서(후반 layer Attention)**는 단지 그 꺼내온 데이터를 문장 맨 끝으로 **'복사-붙여넣기'**만 한다는 사실을 밝혀낸 것입니다.
따라서 모델에게 새로운 지식을 가르칠 때 모델 전체를 다시 공부시킬(Fine-tuning) 필요 없이, 그 핵심 기억 창고의 서랍 속 색인 카드 딱 한 장만 수학 공식으로 정교하게 바꿔 끼워 넣으면(ROME), 다른 지식은 전혀 건드리지 않고도 원하는 사실만 완벽하게 바꿀 수 있다는 것을 보여준 연구입니다.

 

 

 

 

Causal Tracing (Causal Mediation Analysis)
델 내부의 특정 Activation이 최종 예측에 실제로 얼마나 영향을 미치는지, 해당 Activation을 인위적으로 변경하여 측정하는 기법

Clean Run: 정상 입력을 넣고 모델이 Seattle을 예측할 때의 모든 Activation 저장.



Corrupted Run: Space Needle에 해당하는 토큰의 Embedding에 Gaussian Noise를 추가하여 정답 예측을 방해.



Restoration Run: Corrupted 상태에서 특정 Layer·Token의 Activation 하나를 Clean Run의 값으로 복원.



Indirect Effect (IE): 해당 Activation을 복원했을 때 Seattle의 예측 확률이 얼마나 회복되는지 측정.


Early Site & Late Site
Early Site (초기 지점): 모델이 사실 정보를 회상하는 과정에서 중요한 영향을 미치는 Subject의 마지막 Token에 위치한 중간 Layer의 MLP 영역.
Late Site (후기 지점): 최종 정답을 예측하기 직전, 입력 문장의 마지막 Token에 위치한 상위 Layer의 Attention 영역.


ROME 논문에서의 의미
예시: "The Space Needle is in downtown ___" → Seattle
1. Early Site: 중간 Layer의 MLP가 Space Needle이라는 Subject와 관련된 사실 정보를 회상하는 역할.
2. Late Site: 상위 Layer의 Attention이 앞서 처리된 정보를 마지막 Token인 downtown으로 전달하여 정답 예측에 활용하는 역할.


ROME (Rank-One Model Editing):
LLM 내부의 특정 사실 지식을 수정하기 위해, 중간 Layer의 MLP 가중치에 Rank-One Update를 적용하는 Model Editing 기법.


ROME의 작동 원리
예시: "Space Needle은 Seattle에 있다" → "Space Needle은 Paris에 있다"로 변경
1. 위치 탐색 (Causal Tracing): 사실 정보의 회상에 중요한 중간 Layer의 MLP를 식별.
2. Key 추출 (\(k^*\)): Space Needle이라는 Subject에 대응하는 MLP 내부 Activation을 추출.
3. Value 최적화 (\(v^*\)): 모델이 Seattle 대신 Paris를 예측하도록 만드는 새로운 Value 벡터를 최적화.
\(k^*\): 실제 Forward Pass에서 추출한 MLP 내부 Activation

\(v^*\): 원하는 정답을 출력하도록 인위적으로 최적화한 MLP 출력 벡터


4. Rank-One Update: 기존 MLP의 가중치를 수정하여 \(k^*\)가 입력되면 \(v^*\)가 출력되도록 변경.


Essence Drift
Essence Drift는 모델의 특정 사실 하나만 수정하려고 했는데, 그 과정에서 해당 Subject 자체에 대한 정체성이나 본질적인 표현까지 같이 변해버리는 현상입니다.


Bleedover
Bleedover는 특정 Subject의 사실을 수정했는데, 그 수정 효과가 다른 Subject의 지식까지 번지는 현상입니다.

Regurgitation
Regurgitation은 모델이 새로운 사실을 실제 지식처럼 내부적으로 통합한 것이 아니라, 편집에 사용된 Target Object나 단어를 표면적으로 뱉어내는 데 그치는 현상입니다.


사실 회상에는 특정 위치가 중요하다.
Causal Tracing을 해보니, Subject의 마지막 Token을 처리하는 중간 Layer MLP가 factual prediction에 강한 인과적 영향을 보였습니다. 즉 Space Needle 같은 Subject를 봤을 때 관련 사실을 꺼내는 핵심 계산이 이 근처에서 일어난다는 주장입니다.     2202.05262v5

그 MLP는 단순히 정보를 전달하는 게 아니라, factual association을 저장하는 역할을 할 수 있다.
저자들은 중간 MLP를 일종의 associative memory로 봅니다. Subject representation이 들어오면 그 Subject와 연결된 속성 정보를 꺼내는 구조라는 가설입니다.     2202.05262v5

그 가설이 맞다면, 그 MLP의 weight를 직접 바꿨을 때 사실 하나를 바꿀 수 있어야 한다.
그래서 ROME을 만든 겁니다. 실제로 특정 factual association을 바꿨을 때, 단순히 target word만 뱉는 게 아니라 paraphrase에도 어느 정도 일반화되고 다른 Subject는 비교적 유지되는 결과를 보여줍니다. 저자들은 이걸 근거로 **“우리가 어디에 지식이 저장되는지뿐 아니라, 어느 정도 어떻게 저장되는지도 이해했다”**고 주장합니다.