One Mask to Rule Them All: On Hidden Facts after Editing and How to Find Them

기술적 요약: "모든 것을 지배하는 하나의 마스크: 편집 후 숨겨진 사실과 이를 찾는 방법"

문제 제기
지식 편집 (KE) 방법, 특히 ROME 와 MEMIT 과 같은 '위치 파악 및 편집 (locate-and-edit)' 접근법은 MLP 가중치를 직접 수정함으로써 트랜스포머 모델 내 사실적 연관성을 업데이트하는 것을 목표로 합니다. 이러한 방법들은 MLP 계층이 특정 사실을 덮어쓸 수 있는 표적 업데이트가 가능한 연관 기억체로 기능한다는 가설 하에 작동합니다. 그러나 이러한 편집의 내부 메커니즘은 아직 충분히 탐구되지 않았습니다. 개념적 난제가 존재합니다: 트랜스포머 내의 사실적 지식이 중복된 경로 ( '히dra 효과') 에 분산되어 있다면, 어떻게 단일 계층이나 소수의 계층 범위를 수정하는 것이 지식 검색을 성공적으로 대체할 수 있을까요? 현재 평가는 오직 출력 행동 (즉, 모델이 새로운 사실을 출력하는가?) 에만 의존하며, 지식이 실제로 덮어쓰였는지 아니면 단순히 억제되었는지에 대한 질문은 해결되지 않은 채 남아 있습니다.

방법론
저자들은 ROME 와 MEMIT 이 원래 지식을 덮어쓰는 것이 아니라 하위 Attention 을 장악하는 증폭된 신호를 주입하여 그 검색을 억제한다고 가정합니다. 이를 검증하기 위해 다양한 사실에 걸쳐 편집을 유지하는 데 중요한 공통 기능적 부분 공간 (functional subspace) 을 격리하는 것을 제안합니다.

1. 공유 마스크 학습: 핵심 방법은 여러 가지 다양한 편집의 편집된 가중치 행렬 ($\hat{W}$) 에 대해 이진 마스크 ($K$) 를 동시에 학습하는 것입니다. 마스크는 학습 가능한 매개변수 $\Theta$로 파라미터화되며, 조건부 순전파 (conditional forward pass) 동안 요소별 ($\hat{W} \odot K$) 로 적용됩니다.
2. 최적화 목적 함수: 마스크는 결합된 손실 함수를 통해 세 가지 제약을 만족하도록 학습됩니다:
   • 복원: 가지치기된 모델 ($M_p$) 은 편집된 객체 ($o^*$) 보다 원래 객체 ($o$) 를 선호해야 합니다.
   • 최소성: 마스크는 가능한 가장 작은 가중치 부분 집합을 가지치기해야 합니다 (희소성 제약).
   • 행동 보존: 가지치기된 모델의 일반 언어 능력 (원래 모델과의 KL 발산으로 측정) 은 유지되어야 합니다.
3. 실험 설정: 본 연구는 세 가지 모델 (GPT-2 XL, LLaMA-3.2, Qwen2.5) 과 두 가지 편집 방법 (단일 편집용 ROME, 배치 편집용 MEMIT) 에 걸쳐 CounterFact 데이터셋을 활용합니다. 마스크는 편집 집합에서 학습되고, 훈련 집합과 홀드아웃 테스트 집합을 포함한 분포 외 (OOD) 관계에 대해 평가됩니다.

주요 결과

• 보편적 역전: 단일 희소 이진 마스크 (일반적으로 가중치의 10% 미만을 가지치기) 가 훈련 집합에서 80% 이상의 편집을 성공적으로 역전시키고, 보지 못한 테스트 편집에서도 70% 이상을 역전시킵니다. 이 일반화는 서로 다른 모델과 편집 방법 전반에 걸쳐 유효하며, 공유된 구조적 메커니즘을 시사합니다.
• 메커니즘 식별: 잔여 스트림 분해는 마스크가 후속 계층에서의 과도한 어텐션 (overattention) 을 제거함으로써 편집을 역전시킨다는 것을 보여줍니다. 편집된 모델은 편집 계층 하류에서 어텐션 기여도가 급증하는 것을 보이지만, 가지치기된 모델은 원래 모델과 일치하도록 어텐션 수준을 복원합니다. 결정적으로, 원래 지식을 계속 인코딩하는 MLP 경로는 마스크의 영향을 거의 받지 않습니다.
• 메커니즘의 필요성: 학습된 마스크를 편집 과정 도중 주입 (식별된 부분 공간 차단) 하면 편집 성공률이 98% 에서 38% 로 떨어집니다. 이는 식별된 가중치 부분 공간이 역전을 위해 충분할 뿐만 아니라 편집이 성공하기 위해 구조적으로 필수적임을 확인시켜 줍니다.
• 능력 복원: 마스크를 적용하면 ROME 편집으로 인해 저하된 일반 언어 모델링 능력 (퍼플렉시티) 을 회복하며, 종종 퍼플렉시티 급증을 수백만 수준에서 기준선 수준으로 되돌립니다.

의의 및 주장
본 논문은 '위치 파악 및 편집' 방법에 대한 이해를 근본적으로 재정의한다고 주장합니다:

1. 억제 vs 덮어쓰기: ROME 와 MEMIT 은 저장된 지식을 덮어쓰는 것이 아니라 어텐션 메커니즘을 장악하여 그 검색을 억제합니다. 원래 사실은 MLP 경로에 인코딩되어 있지만, 편집된 신호에 대한 불균형적인 어텐션으로 인해 표면화되지 못합니다.
2. 공유 기능적 부분 공간: 사실별 가중치 변화에도 불구하고, 다양한 편집은 편집을 유지하기 위해 공통된 소수의 가중치 부분 집합에 의존합니다. 이 공유 메커니즘은 이러한 방법들이 관련 사실에 변화를 전파하지 못하는 이유 (파급 효과 부재) 와 역전에 취약한 이유를 설명합니다.
3. 방어 및 탐지: 공유 메커니즘의 존재는 편집이 공통된 서명을 남긴다는 것을 의미합니다. 단일 희소 마스크는 원치 않는 편집을 역전시킬 뿐만 아니라 필요한 부분 공간을 차단함으로써 미래의 편집에 대해 모델을 '잠금'할 수도 있습니다. 이는 악의적 편집으로부터 오픈 가중치 모델을 방어하고, 공유된 구조적 발자국을 기반으로 편집된 사실을 탐지하는 실용적인 경로를 제공합니다.

저자들은 이러한 방법을 동기부여하는 연관 기억 가설이 불완전하다고 결론지었습니다. 사실은 MLP 에 저장될 수 있지만, 검색은 어텐션에 크게 의존하며, 현재 편집 방법들은 이를 업데이트하는 것이 아니라 악용합니다. 이러한 한계는 '위치 파악 및 편집' 방법이 달성할 수 있는 상한선을 시사합니다: 그들은 검색을 억제할 수는 있지만 근본적인 지식 그래프를 실제로 수정할 수는 없습니다.