TimeROME-DLM: 시간적 인과 추론 및 저랭크 추론 기반 지식 편집 기법 - 마스크된 확산 언어 모델 (Masked Diffusion Language Models)을 위한 방법

TimeROME-DLM: Temporal Causal Tracing and Low-Rank Inference-Time Knowledge Editing for Masked Diffusion Language Models

마스크된 확산 언어 모델(MDLM)은 LLaDA와 같이 현재 autoregressive (AR) LLM에 필적하는 성능을 보이지만, 기존의 지식 편집 및 망각 기법(ROME, MEMIT 등)들은 AR 트랜스포머를 대상으로 하며, 반복적인 디노이징 과정에서 오류가 발생하거나, 수십 GB의 추가 VRAM을 필요로 하는 역전파 연산을 수행하여 MDLM의 성능 저하를 야기합니다. 본 논문에서는 MDLM에 대한 첫 번째의 학습-불필요(training-free)하고, 기울기-불필요(gradient-free)하며, 추론-시간(inference-time) 지식 편집 프레임워크인 TimeROME-DLM을 제안합니다. 이는 두 가지 구성 요소로 이루어집니다. 첫째, 각 사실에 대해 특정 좌표를 식별하고, 이후 디노이징 단계에서 객체 예측에 가장 큰 영향을 미치는 좌표를 찾는 시간적 간접 효과(TIE) 기반 인과 추론 프로토콜입니다. 둘째, 모든 망각될 사실에 대한 subject key와 target delta 값을 집계하고, 단일의 ridge-regularized 업데이트를 해당 좌표에 적용하여 메모리를 효율적으로 관리하는 closed-form, low-rank residual edit memory입니다. 핵심 파라미터는 고정되어 있으며, 세 개의 하이퍼파라미터(alpha, lambda, q)만이 작은 검증 데이터셋에서 조정됩니다. TOFU forget01 데이터셋과 TOFU로 미세 조정된 LLaDA-8B-Base 모델을 사용하여 TimeROME-DLM은 망각 집합의 로그 확률을 약 83 nats만큼 감소시켰습니다. 동일한 구성이 LLaDA-8B-Instruct, Dream-7B, MMaDA-8B, DiffuLLaMA-7B 및 LLaDA-MoE-1.4B 모델에도 적용 가능합니다. TimeROME-DLM은 유지 집합의 로그 확률을 거의 일정하게 유지(유틸리티 안전 지점에서 약 1 nat 이내)하면서 50개의 사실을 순차적으로 추가하고, 최적화된 학습 기반 모델 대비 4배에서 14배 빠른 속도를 제공하며, 400개 사실까지 확장 가능합니다. TimeROME-DLM은 AR LLM과 MDLM 간의 지식 편집 격차를 상당한 계산 비용 절감으로 해소합니다.

Masked diffusion language models (MDLMs) such as LLaDA now rival autoregressive (AR) LLMs, but every existing knowledge-editing and unlearning method (ROME, MEMIT, etc.) targets AR transformers and either makes assumptions that fail under iterative denoising, or requires gradient updates whose backward-pass activations cost tens of GB of extra VRAM and which collapse MDLMs at standard learning rates. We introduce TimeROME-DLM, the first training-free, gradient-free, inference-time knowledge-editing framework for MDLMs. It couples two components: a Temporal Indirect Effect (TIE) causal-tracing protocol that identifies, for each fact, the coordinate whose intervention most strongly drives the object prediction at later denoising steps; and a closed-form, low-rank residual edit memory that aggregates subject keys and target deltas across all forget facts and applies a single ridge-regularised update at that coordinate at every diffusion forward, with sparsification to limit utility spillover. Backbone weights stay frozen; only three hyperparameters (alpha, lambda, q) are tuned on a small validation split. On TOFU forget01 with TOFU-finetuned LLaDA-8B-Base, TimeROME-DLM cuts forget-set log-probability by roughly 83 nats. The same configuration transfers to LLaDA-8B-Instruct, Dream-7B, MMaDA-8B, DiffuLLaMA-7B, and LLaDA-MoE-1.4B. It keeps retain-set log-probability nearly flat (within ~1 nat at the utility-safe operating point) across 50 sequentially inserted facts, delivers a four- to fourteen-fold wall-clock speedup with zero additional VRAM over the strongest converged training-time baseline, and scales sub-linearly to 400 facts. TimeROME-DLM closes the locate-then-edit gap between AR LLMs and MDLMs at a fraction of the computational cost.