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

2026-06-10

AI Summary

This work addresses the challenge of efficiently and safely editing knowledge in masked diffusion language models (MDLMs), for which existing methods are inadequate. The authors propose the first training-free, gradient-free inference-time editing framework that dynamically injects new factual knowledge into a frozen MDLM backbone. By leveraging temporal causal tracing to identify critical intervention coordinates, the method combines closed-form low-rank residual memory with ridge-regularized sparse updates. Requiring only three hyperparameters, the approach generalizes across models without additional GPU memory overhead and scales to up to 400 facts. Evaluated on MDLMs such as LLaDA, it reduces log-probabilities on forgetting sets by an average of 83 nats while keeping fluctuations on retained sets below 1 nat, achieving 4–14× faster inference.

Abstract

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.

Problem

• masked diffusion language models
• knowledge editing
• unlearning
• inference-time editing
• gradient-free

Innovation

• knowledge editing
• masked diffusion language models
• causal tracing
• low-rank update
• inference-time editing