Masked Diffusion Language Models에서의 지식 편집 (Knowledge Editing)

지식 편집 (Knowledge editing)은 언어 모델 내의 사실적 지식을 업데이트하거나 수정하는 것을 목표로 합니다. 널리 사용되는 접근 방식인 '위치 파악 후 편집 (locate-then-edit)'은 두 단계로 이를 수행합니다. 먼저 모델 내에서 특정 사실의 위치를 파악한 다음, 해당 위치의 가중치 (weights)를 편집합니다. 현재까지 이러한 방법들은 오직 자기회귀 모델 (Autoregressive Models, ARMs)에서만 개발되어 왔습니다. 텍스트를 양방향으로 모델링하고 다음 토큰 예측 (next-token prediction) 대신 반복적인 노이즈 제거 (iterative denoising)를 통해 생성하는 마스크 확산 모델 (Masked Diffusion Models, MDMs)에서도 이러한 기본 가정이 유효한지는 여전히 미해결 과제로 남아 있습니다. 본 연구에서는 '위치 파악 후 편집' 방식을 MDM으로 전이시키고, 동일한 규모의 두 가지 MDM (LLaDA, Dream)과 두 가지 ARM (LLaMA, Qwen)을 비교함으로써 이 문제를 다룹니다. 우리의 핵심 발견은 두 부분으로 나뉩니다. 첫째, 편집이 적용되는 위치는 패러다임 간에 전이됩니다. 인과적 추적 (causal tracing) 결과, 두 모델 모두 마지막 주어 토큰 (subject token)에 위치한 동일한 초기-중기 레이어의 MLP를 강조하며, 그곳에서 편집이 가장 효과적입니다. 둘째, 이러한 공유된 위치가 공유된 결과를 보장하지는 않습니다. 단일 토큰 편집은 두 모델 모두에서 성공하지만, 대상이 길어질수록 ARM과 달리 MDM에서는 편집 성능이 체계적으로 저하됩니다. 이러한 실패는 편집된 사실이 생성되는 방식에서 기인합니다. 다중 토큰 대상을 생성하려면 편집이 최적화되지 않은 부분적으로 마스크가 해제된 중간 상태 (partially unmasked intermediate states)를 거쳐야 하기 때문입니다. 이러한 진단을 바탕으로, 우리는 이러한 상태들에 대해 편집을 최적화하는 간단한 교정 방법을 도입하여 다중 토큰 성능을 실질적으로 복구합니다.