Dynamic Weight Grafting: Localizing Finetuned Factual Knowledge in Transformers

이 논문은 미세조정된 사실적 지식을 사전 훈련된 모델의 가중치에 선택적으로 이식하는 '동적 가중치 이식 (Dynamic Weight Grafting)' 기법을 제안하여, LLM 이 엔티티 처리 시 정보를 '보강'하거나 최종 예측 직전에 '회상'하는 두 가지 서로 다른 경로를 통해 지식을 활용함을 규명했습니다.

핵심

🎬 영화 배우의 새로운 정보를 기억하는 법: AI 의 '동적 접목' 연구

이 논문은 **거대 언어 모델 (LLM)**이 새로운 사실 (예: "새로 선출된 교황은 누구인가?" 또는 "새로 개봉한 영화의 배우는 누구인가?") 을 배웠을 때, 그 정보가 모델의 뇌 (파라미터) 안에 어디에, 어떻게 저장되고 어떻게 꺼내지는지 분석한 연구입니다.

기존의 연구 방법들은 마치 뇌의 특정 부위를 '수술'해서 기능을 멈추게 하는 방식이라, 정보가 어떻게 흐르는지 정확히 파악하기 어려웠습니다. 이 연구팀은 새로운 방법인 **'동적 가중치 접목 (Dynamic Weight Grafting)'**이라는 기술을 개발해 그 비밀을 밝혀냈습니다.

---

🌱 1. 새로운 방법: '동적 가중치 접목'이란?

기존의 '활성화 패치 (Activation Patching)' 방식은 모델이 정보를 처리하는 도중, 특정 시점의 '생각 (활성화)'을 다른 모델의 생각으로 갈아끼우는 것이었습니다. 이는 마치 레고 블록을 조립하는 도중, 이미 붙인 블록을 떼어내고 다른 블록으로 교체하는 것과 같습니다. 그러면 이전 단계에서 어떻게 정보가 쌓였는지 알 수 없게 됩니다.

반면, 이 논문에서 제안한 **'동적 가중치 접목'**은 다음과 같습니다:

• 비유: 두 명의 요리사 (기존 모델과 새로운 정보를 학습한 모델) 가 있다고 가정해 봅시다.
• 기존 방식: 요리사가 요리를 하는 도중, 특정 단계의 '손맛'을 다른 요리사의 손맛으로 바꿔버리는 것입니다.
• 새로운 방식: 요리를 시작할 때, **특정 재료 (가중치)**만 새로운 요리사의 것을 가져와서 사용합니다. 예를 들어, "양파를 다지는 칼은 새로운 요리사의 것을 쓰고, 소금 뿌리는 손은 기존 요리사의 것을 쓴다"는 식으로 세부 부품만 교체합니다.

이렇게 하면 모델의 전체 흐름을 방해하지 않으면서, 어떤 부품이 새로운 정보를 기억하는 데 핵심적인지 정확히 찾아낼 수 있습니다.

---

🔍 2. 발견한 두 가지 기억 경로

연구팀은 이 기술을 통해 AI 가 새로운 사실을 기억하고 꺼낼 때 사용하는 두 가지 주요 경로를 발견했습니다.

① '풍요로움 (Enrichment)' 경로: 처음 만날 때 정보를 채우기

• 상황: 문장 속에서 새로운 배우 이름 (예: "Zendaya") 이 처음 등장할 때입니다.
• 비유: 여행 가방을 채우는 과정입니다. Zendaya 라는 이름을 처음 보자마자, AI 는 그 이름이 가진 새로운 정보 (새로운 영화 정보 등) 를 가방 (잔여 스트림) 에 채워 넣습니다.
• 결과: 이 가방이 충분히 채워지면, 나중에 질문을 받기만 해도 가방에서 정보를 꺼내 답할 수 있습니다.

② '기억 회상 (Recall)' 경로: 마지막에 정보를 꺼내기

• 상황: 문장의 마지막 부분, 정답을 말하기 직전입니다.
• 비유: 마지막에 다시 확인하는 과정입니다. 처음에 가방을 채우지 않았더라도, 정답을 말하기 직전에 "아, 이 배우가 어떤 영화에 나왔지?"라고 마지막 순간에 정보를 찾아내는 메커니즘이 작동합니다.
• 결과: 이 과정은 문장의 마지막 토큰 (단어) 위치에서 일어나며, Attention(주목) 메커니즘과 Feedforward Network(전통적인 신경망 층) 가 협력하여 정보를 꺼냅니다.

---

🧩 3. 핵심 발견: 두 가지 중 하나만 있어도 된다?

흥미로운 점은 모델에 따라 두 경로 중 하나만 있어도 정답을 맞출 수 있다는 것입니다.

• 경우 A: 처음 배우 이름이 나올 때 정보를 잘 채워주면 (풍요로움), 나중에 정답을 말하기 직전에 어떤 처리를 하든 상관없이 정답을 맞춥니다.
• 경우 B: 처음에는 정보를 채우지 않아도, 정답을 말하기 직전에 (기억 회상) 정보를 찾아내는 능력이 뛰어나면 정답을 맞춥니다.

하지만 두 경로 모두를 막아버리면 (가방도 비우고, 마지막 확인도 못 하게 하면), AI 는 완전히 망각한 상태가 되어 정답을 못 맞춥니다. 즉, 이 두 가지 경로가 새로운 정보를 기억하는 데 필수불가결하다는 뜻입니다.

---

🛠️ 4. 어떤 부품이 중요한가?

연구팀은 '기억 회상' 경로가 구체적으로 어떤 부품에서 일어나는지까지 찾아냈습니다.

• Attention (주목): 문장의 처음 (배우 이름) 과 마지막 (정답 직전) 에서 무엇을 주목할지 결정하는 역할을 합니다.
• Feedforward Network (신경망 층): 마지막 단계에서 실제 정보를 꺼내서 정답으로 변환하는 역할을 합니다.
• 비유: Attention 은 도서관의 검색창이고, Feedforward Network 는 책장을 뒤져서 책을 꺼내는 사서입니다. 검색창 (Attention) 이 올바른 책을 찾아주고, 사서 (Feedforward) 가 그 책을 가져와서 정답을 알려주는 것입니다.

---

💡 5. 결론: AI 의 기억 방식을 이해하다

이 연구는 AI 가 새로운 사실을 배울 때 단순히 '기억'만 하는 것이 아니라, 처음 정보를 접할 때 저장하는 방식과 답변을 할 때 다시 찾아내는 방식이라는 두 가지 전략을 동시에 사용한다는 것을 보여줍니다.

• 기존의 오해: AI 는 새로운 정보를 배울 때 무조건 처음부터 끝까지 모든 정보를 저장한다고 생각했습니다.
• 새로운 사실: AI 는 상황에 따라 처음에 저장하거나 마지막에 찾아내거나, 혹은 둘 다 사용하는 유연한 방식을 가집니다.

이처럼 AI 의 내부 작동 원리를 '부품 교체' 방식으로 정밀하게 분석함으로써, 우리는 AI 가 어떻게 학습하고 기억하는지 더 깊이 이해하게 되었습니다. 이는 향후 AI 의 오류를 수정하거나, 새로운 지식을 더 효율적으로 주입하는 데 큰 도움이 될 것입니다.

기술 요약

1. 연구 배경 및 문제 제기 (Problem)

• 배경: LLM 은 사전 학습 (Pre-training) 을 통해 방대한 관계를 저장하지만, 새로운 사실을 파인튜닝할 때 이 정보가 모델 파라미터의 어디에, 어떻게 인코딩되는지는 명확하지 않습니다.
• 기존 방법의 한계: 기존 지식 국소화 (Localization) 연구는 주로 **활성화 패칭 (Activation Patching)**이나 제거 (Ablation) 기법을 사용했습니다.
  • 문제점: 활성화 패칭은 특정 시점의 모델 내부 상태 (Residual Stream) 를 교체하는 방식인데, 이는 이전 단계에서 계산된 정보 (예: 엔티티 토큰이 사실 정보로 '풍부해짐' enriched) 를 덮어쓰게 됩니다. 따라서 모델이 새로운 정보를 **실제로 추출 (Recall)**하는지, 아니면 단순히 이전 단계에서 계산된 정보를 전달하는지 구분하기 어렵습니다.
• 연구 목표: 파인튜닝된 지식이 생성 과정에서 어떻게 인출되는지, 그리고 이를 담당하는 구체적인 메커니즘 (컴포넌트) 과 토큰 위치를 파괴적이지 않으면서 정밀하게 규명하는 새로운 방법론이 필요합니다.

2. 제안 방법: 동적 가중치 이식 (Dynamic Weight Grafting)

저자들은 **동적 가중치 이식 (Dynamic Weight Grafting)**이라는 새로운 분석 기법을 제안했습니다.

• 핵심 개념: 사전 학습된 모델 (Pretrained, $\theta_{pre}$) 과 파인튜닝된 모델 (Finetuned, $\theta_{ft}$) 의 가중치 (Weights) 를 선택적으로 교체하며 생성을 수행합니다.
• 동작 방식:
  • 활성화 패칭이 '상태 (Activation)'를 교체하는 반면, 이 방법은 '메커니즘 (가중치 행렬)'을 교체합니다.
  • 각 토큰 위치 ($t$) 와 모델 구성 요소 ($c$, 예: Attention Head, Feed-Forward Network 등) 에 대해 가중치 마스크 ($\gamma$) 를 적용하여, 특정 위치에서는 파인튜닝된 가중치를, 다른 위치에서는 사전 학습된 가중치를 사용합니다.
  • 수식적 정의: $\tilde{\theta}_m(t) = \begin{cases} \theta^A_c & \text{if } \gamma_c(t) = 0 \\ \theta^B_c & \text{if } \gamma_c(t) = 1 \end{cases}$
  • 이를 통해 모델의 계산 그래프 전체에 걸친 정보 흐름을 방해하지 않으면서, 특정 메커니즘이 파인튜닝된 지식 인출에 필수적인지 (Necessary) 혹은 충분조건인지 (Sufficient) 를 인과적으로 분석할 수 있습니다.

3. 주요 실험 및 결과 (Key Results)

저자들은 Llama3, Pythia, GPT-2 XL, Gemma 등 4 가지 모델을 사용하여 실험을 진행했습니다.

A. 지식 인출의 두 가지 경로 (Two Pathways)

동적 가중치 이식을 통해 파인튜닝된 관계 정보를 인출하는 두 가지 주요 경로가 발견되었습니다.

1. 풍부화 (Enrichment) 경로: 엔티티 토큰 (예: "Zendaya") 이 처음 처리될 때, 해당 토큰의 표현이 새로운 관계 정보로 '풍부화'되는 과정.
2. 회상 (Recall) 경로: 예측 직전 마지막 토큰 위치에서, 풍부화된 엔티티 정보나 직접적인 단서를 바탕으로 관계 사실을 '회상'하여 생성하는 과정.

• 결과:
  • 단일 경로로도 성능 회복: 어떤 경우에는 엔티티 토큰만 이식하거나, 마지막 토큰만 이식하는 것만으로도 파인튜닝된 모델의 성능을 상당 부분 회복할 수 있었습니다.
  • 결합 시 완전 회복: 두 경로를 모두 이식하면 파인튜닝 모델의 성능을 거의 완벽하게 재현했습니다.
  • 필수성 검증: 엔티티와 마지막 토큰을 제외한 나머지 모든 부분을 파인튜닝 모델로 이식하더라도 (Complement grafting), 성능은 사전 학습 모델 수준으로 떨어졌습니다. 이는 이 두 경로가 지식 인출에 **필수적 (Necessary)**임을 의미합니다.

B. 구성 요소 국소화 (Component Localization)

"회상 (Recall)" 경로가 구체적으로 어떤 모델 구성 요소를 통해 이루어지는지 분석했습니다.

• Attention vs. FFN:
  • 첫 번째 엔티티 토큰: 작업 특이적 (Task-specific) Attention 메커니즘이 엔티티 정보를 처리하는 데 필수적입니다.
  • 마지막 토큰 (예측 직전): 관계 특이적 (Relation-specific) 추출은 마지막 레이어의 **Output Projection Matrix (O matrix)**와 **Feed-Forward Networks (FFN)**에서 발생합니다.
• 발견: Attention 메커니즘 전체를 교체하지 않고도, 마지막 레이어의 O matrix 와 FFN 만을 파인튜닝된 가중치로 교체하면 (단, Attention 은 작업 특이적 기능을 유지할 때) 관계 완성 성능을 대부분 회복할 수 있었습니다. 이는 파인튜닝된 지식이 FFN 과 O matrix 를 통해 구체적으로 인출됨을 시사합니다.

C. 일반화 검증

• 템플릿 데이터 vs. 실제 데이터: 합성 데이터 (가짜 영화/배우) 뿐만 아니라, 모델 출시 이후의 실제 위키피디아 기사 (실제 영화) 에 대한 파인튜닝에서도 유사한 패턴 (엔티티 풍부화 + 마지막 토큰 회상) 이 관찰되었습니다.
• 모델 아키텍처 차이: Gemma 와 Llama3 와 같은 최신 모델에서는 "회상" 경로가 특히 강력하게 작용하는 반면, GPT-2 XL 나 Pythia 에서는 "풍부화" 경로가 더 두드러지는 등 아키텍처에 따라 우세한 경로가 다를 수 있음을 발견했습니다.

4. 주요 기여 (Contributions)

1. 새로운 분석 도구 개발: 기존 활성화 패칭의 한계 (과거 정보 덮어쓰기) 를 극복하는 동적 가중치 이식 (Dynamic Weight Grafting) 기법을 제안했습니다. 이는 모델의 파라미터와 메커니즘을 직접 조작하여 인과 관계를 분석할 수 있게 합니다.
2. 지식 인출 메커니즘 규명: LLM 이 파인튜닝된 사실을 인출할 때, **엔티티 토큰에서의 '풍부화'**와 **마지막 토큰에서의 '회상'**이라는 두 가지 병렬 경로를 사용함을 증명했습니다.
3. 구체적 구성 요소 매핑: 지식 인출이 Attention, FFN, O matrix 등 구체적인 구성 요소와 어떤 토큰 위치에서 발생하는지를 정밀하게 매핑했습니다. 특히 마지막 레이어의 FFN 과 O matrix 가 핵심적인 역할을 함을 밝혔습니다.

5. 의의 및 결론 (Significance)

• 해석 가능성의 심화: 단순히 "어디에 지식이 있는가"를 넘어, "지식이 어떻게 인출되는가 (메커니즘)"에 대한 이해를 심화시켰습니다.
• 비파괴적 분석: 기존 방법보다 파괴적이지 않으면서 더 정밀한 국소화를 가능하게 하여, 모델의 내부 작동 원리를 이해하는 데 새로운 패러다임을 제시합니다.
• 향후 연구 방향: 이 기법은 지식 편집 (Knowledge Editing), 안전성 강화 (Safety Tuning), 그리고 모델의 기억 메커니즘을 제어하는 연구에 중요한 통찰을 제공할 것입니다.

요약하자면, 이 논문은 LLM 이 새로운 사실을 학습하고 이를 활용하는 과정을 동적 가중치 이식을 통해 해부학적으로 분석하여, 엔티티 처리 단계의 정보 풍부화와 생성 직전의 정보 회상이라는 두 가지 핵심 메커니즘이 상호 보완적으로 작동함을 증명했습니다.