Knowledge Editing in Masked Diffusion Language Models

기술 요약: 마스크드 디퓨전 언어 모델에서의 지식 편집

문제 정의
지식 편집(Knowledge editing)은 전체 모델을 재학습시키지 않고 대규모 언어 모델(LLM) 내의 특정 사실적 지식을 업데이트하거나 수정하는 것을 목표로 한다. 지배적인 패러다임인 "위치 파악 후 편집(locate-then-edit)"(예: ROME, MEMIT)은 오직 자기회귀 모델(Autoregressive Models, ARMs)만을 위해 개발되었다. 이러한 방법들은 인과적 추적(causal tracing)에 의존하여, 사실적 회상이 주어(subject)의 마지막 토큰에 위치한 초기-중기 레이어의 다층 퍼셉트론(MLP)에 국한된다는 점을 식별한다. 그 후, 이 위치의 가중치를 수정하여 새로운 목적어(object)를 주어와 연관시킨다.

그러나 마스크드 디퓨전 모델(Masked Diffusion Models, MDMs)은 이와 구별되는 패러다임을 나타낸다. 왼쪽에서 오른쪽으로 순차적으로 토큰을 예측하는 ARM과 달리, MDM은 텍스트를 양방향으로 모델링하며, 완전히 마스킹된 상태로부터 반복적인 디노이징(denoising)을 통해 마스킹된 토큰을 복구하며 텍스트를 생성한다. 따라서 사실의 국소성(localization of facts)과 해당 위치에서의 편집 효능이라는 'locate-then-edit'의 전제 조건이, 근본적으로 다른 훈련 목표와 생성 과정을 가진 MDM에서도 유효한지는 미해결 과제로 남아 있다.

방법론
저자들은 핵심 알고리즘을 변경하지 않고 MDM의 마스크드 예측 메커니즘에 맞게 최소한의 조정만을 가하여, locate-then-edit 프레임워크(특히 MEMIT 알고리즘)를 MDM으로 전이시킨다.

1. 실험 설정: 연구는 일치하는 두 쌍의 모델을 비교한다:

   • ARM: LLaMA-3-8B 및 Qwen2.5-7B.
   • MDM: LLaDA-8B (처음부터 학습됨) 및 Dream-v0-7B (Qwen2.5 가중치로 초기화됨).
   • 데이터셋: COUNTERFACT (주로 단일 토큰 타겟인 표준 벤치마크) 및 KAMEL (멀티 토큰 타겟 분석을 위해 조정됨).

2. 인과적 추적(Causal Tracing): 저자들은 MDM에서 사실이 인과적으로 어디에 국소화되어 있는지 결정하기 위해 인과적 추적을 수행한다. 이들은 주어 토큰 임베딩을 오염시키고, 특정 레이어와 토큰 위치의 은닉 상태(hidden states)를 복구함으로써 올바른 예측의 회복 여부를 측정한다.

3. 편집 평가: 저자들은 다양한 토큰 위치와 레이어 범위에 걸쳐 편집 성능(효능, 일반화, 특이성)을 측정한다. 결정적으로, MDM의 생성 과정 특유의 결함을 식별하기 위해 다양한 타겟 길이(단일 토큰 vs 멀티 토큰)와 생성 설정에 따른 성능을 평가한다.

4. 제안된 교정: 멀티 토큰 타겟에 대한 MDM의 성능 저하를 확인한 후, 저자들은 교정안을 제안한다. 표준 편집은 단일 조건 상태(완전히 마스킹된 정답)에 대한 가중치 업데이트($\delta$)를 최적화한다. 제 제안된 교정은 최적화 범위를 모델이 반복적인 디노이징 과정 중에 거치는 중간 단계의 부분적으로 언마스킹된 상태들까지 확장한다.

주요 결과

1. 공유된 국소성: 'locate-then-edit'의 핵심 가정은 패러다임을 가리지 않고 전이된다. 양방향 훈련에도 불구하고, MDM은 ARM과 동일한 위치, 즉 주어 토큰의 초기-중기 레이어 MLP에 사실적 회상을 국소화한다. 인과적 추적은 "초기 사이트"(사실 회상)와 "후기 사이트"(예측 출력)가 두 모델 유형 모두에 존재함을 확인해 준다.

2. 엇갈린 결과 (멀티 토큰의 붕괴): 국소성은 공유되지만, 결과는 그렇지 않다.

   • 단일 토큰 타겟: 편집 성능은 ARM과 MDM에서 유사하게 나타난다.
   • 멀티 토큰 타겟: MDM의 편집 성능은 타겟 길이가 길어짐에 따라 체계적으로 저하된다. 예를 들어, 4-토큰 타겟의 경우 LLaMA는 높은 수준(~0.89)을 유지하는 반면, LLaDA의 효능은 ~0.92(1-토큰)에서 ~0.27(4-토큰)로 급격히 떨어진다.
   • 근본 원인: 이러한 실패는 최적화 상태와 생성 과정 사이의 불일치에서 기인한다. MEMIT은 완전히 마스킹된 상태를 위해 편집을 최적화한다. ARM에서는 생성이 왼쪽에서 오른쪽으로 진행되어 이 상태와 일치한다. 그러나 MDM에서는 일부 토큰은 언마스킹되고 다른 토큰은 마스킹된 상태인 중간 상태들의 체인을 통해 생성이 진행된다. 완전히 마스킹된 상태만을 위해 최적화된 편집은 이러한 중간 언마스킹 토큰들에 의해 생성되는 변화하는 컨텍스트를 고려하지 못한다.

3. 중간 상태 최적화를 통한 복구: 모든 중간 단계의 부분적으로 언마스킹된 상태(드러난 토큰의 수를 순환하며)에 대해 잔차 $\delta$를 적합시키도록 최적화 목표를 수정함으로써, 저자들은 멀티 토큰 편집 성능을 실질적으로 복구했다. LLaDA의 경우, 4-토큰 타겟에 대한 효능이 0.27에서 0.73으로 상승했다.

의의 및 주장
본 논문은 MDM에서의 locate-then-edit 지식 편집에 대한 첫 번째 분석을 제공한다고 주장한다. 주요 기여는 다음과 같다:

• 국소성의 검증: 사실의 인과적 국소화가 주어 토큰 MLP에 존재한다는 것이 자기회귀 모델과 디퓨전 기반 언어 모델 모두에 적용되는 패러다임 불변의 속성임을 입증하였다.
• 생성 유도형 실패의 식별: 표준 벤치마크(단일 토큰 타겟 위주)가 멀티 토큰 사실에 대한 MDM의 체계적인 locate-then-edit 실패를 가리고 있음을 밝혀냈다. 이 실패는 단일 상태 최적화와 반복적 디노이징 사이의 불일치에 기인한다.
• 간단한 교정법: 편집을 중간 상태에 맞게 최적화하는 방법을 제안하여, 근본적인 편집 위치나 알고리즘 구조를 변경하지 않고도 MDM의 멀티 토큰 편집을 위한 성능 격차를 효과적으로 해소하였다.

저자들은 편집 방법을 아키텍처 간에 이식할 때는 편집의 *위치(site)*뿐만 아니라, 편집된 사실이 실현되는 생성 과정(generation process) 또한 일치시켜야 한다고 결론짓는다. 또한, 제안된 교정이 주요 실패 원인을 해결하긴 했으나 가장 긴 타겟에서는 여전히 성능 격차가 존재함을 언급하며, 향후 연구에서 MDM이 토큰을 언마스킹하는 특정 순서를 고려해야 할 수도 있음을 시사하였다.