Knowledge Editing in Masked Diffusion Language Models

Resumen Técnico: Edición de Conocimiento en Modelos de Lenguaje de Difusión con Máscaras

Declaración del Problema
La edición de conocimiento tiene como objetivo actualizar o corregir conocimientos fácticos específicos dentro de los Modelos de Lenguaje Grandes (LLMs) sin reentrenar el modelo completo. El paradigma dominante, "localizar-luego-editar" (por ejemplo, ROME, MEMIT), se desarrolló exclusivamente para Modelos Autorregresivos (ARMs). Estos métodos dependen del rastreo causal para identificar que la recuperación de hechos se localiza en las capas MLP (Perceptrón Multicapa) de la etapa inicial a media de la última posición del token del sujeto. Posteriormente, modifican los pesos en esta ubicación para asociar un nuevo objeto con el sujeto.

Sin embargo, los Modelos de Difusión con Máscaras (MDMs) representan un paradigma distinto. A diferencia de los ARMs, que predicen tokens secuencialmente de izquierda a derecha, los MDMs modelan el texto de forma bidireccional y generan texto mediante un proceso de eliminación de ruido iterativo, recuperando los tokens enmascarados desde un estado totalmente enmascarado. Sigue siendo una cuestión abierta si los supuestos que subyacen al "localizar-luego-editar" —específicamente la localización de los hechos y la eficacia de la edición en esa ubicación— se mantienen para los MDMs, dado sus objetivos de entrenamiento y procesos de generación fundamentalmente diferentes.

Metodología
Los autores transfieren el marco de "localizar-luego-editar" (específicamente el algoritmo MEMIT) a los MDMs sin alterar su algoritmo central, aplicando únicamente ajustes mínimos para adaptarse al mecanismo de predicción enmascarada del MDM.

1. Configuración Experimental: El estudio compara dos pares de modelos emparejados:

   • ARMs: LLaMA-3-8B y Qwen2.5-7B.
   • MDMs: LLaDA-8B (entrenado desde cero) y Dream-v0-7B (inicializado con pesos de Qwen2.5).
   • Datasets: COUNTERFACT (benchmark estándar, mayoritariamente de objetivos de un solo token) y KAMEL (adaptado para el análisis de objetivos de múltiples tokens).

2. Rastreo Causal: Los autores realizan un rastreo causal para determinar dónde se localizan causalmente los hechos en los MDMs. Corrompen las incrustaciones (embeddings) de los tokens del sujeto y miden la recuperación de las predicciones correctas mediante la restauración de los estados ocultos en capas y posiciones de tokens específicas.

3. Evaluación de la Edición: Miden el rendimiento de la edición (eficacia, generalización, especificidad) a través de diferentes posiciones de tokens y rangos de capas. Crucialmente, evalúan el rendimiento a través de variaciones en la longitud del objetivo (un solo token frente a múltiples tokens) y configuraciones de generación para identificar fallos específicos del proceso de generación de los MDMs.

4. Corrección Propuesta: Tras identificar una degradación del rendimiento en los MDMs para objetivos de múltiples tokens, los autores proponen una corrección. La edición estándar optimiza la actualización del peso ($\delta$) para un único estado de condicionamiento (el estado de la respuesta totalmente enmascarada). La corrección propuesta amplía la optimización para incluir estados intermedios parcialmente desenmascarados que el modelo atraviesa durante la eliminación de ruido iterativa.

Hallazgos Clave

1. Localización Compartida: El supuesto central de "localizar-luego-editar" se transfiere entre paradigmas. A pesar del entrenamiento bidireccional, los MDMs localizan la recuperación fáctica en el mismo sitio que los ARMs: el MLP de la etapa inicial a media del último token del sujeto. El rastreo causal confirma que el "sitio temprano" (recuperación de hechos) y el "sitio tardío" (lectura de la predicción) existen en ambos tipos de modelos.

2. Resultados Divergentes (El Fallo de Múltiples Tokens): Aunque la ubicación es compartida, el resultado no lo es.

   • Objetivos de un Solo Token: La edición funciona de manera comparable en ARMs y MDMs.
   • Objetivos de Múltiples Tokens: El rendimiento de la edición en los MDMs se degrada sistemáticamente a medida que aumenta la longitud del objetivo. Por ejemplo, para objetivos de 4 tokens, la eficacia de LLaDA cae de 0.92 (1 token) a ~0.27, mientras que LLaMA se mantiene alta (0.89).
   • Causa Raíz: El fallo surge de un desajuste entre el estado de optimización y el proceso de generación. MEMIT optimiza la edición para el estado totalmente enmascarado. En los ARMs, la generación procede de izquierda a derecha, coincidiendo con este estado. En los MDMs, la generación procede a través de una cadena de estados intermedios donde algunos tokens están desenmascarados y otros permanecen enmascarados. La edición, optimizada solo para el estado totalmente enmascarado, no tiene en cuenta el contexto cambiante creado por estos tokens intermedios parcialmente desenmascarados.

3. Restauración mediante Optimización de Estados Intermedios: Al modificar el objetivo de optimización para ajustar el residuo $\delta$ a través de todos los estados intermedios parcialmente desenmascarados (ciclando a través de conteos de tokens revelados), los autores restauran sustancialmente el rendimiento de la edición de múltiples tokens. Para LLaDA, la eficacia en objetivos de 4 tokens aumenta de 0.27 a 0.73.

Significancia y Reivindicaciones
El artículo afirma proporcionar el primer análisis de la edición de conocimiento de "localizar-luego-editar" en MDMs. Sus principales contribuciones son:

• Validación de la Localización: Demostrar que la localización causal de los hechos en el MLP del token del sujeto es una propiedad agnóstica al paradigma, válida tanto para modelos de lenguaje autorregresivos como para los de difusión.
• Identificación de un Fallo Inducido por la Generación: Revelar que los benchmarks estándar (dominados por objetivos de un solo token) enmascaran un fallo sistemático de "localizar-luego-editar" en los MDMs para hechos de múltiples tokens. Este fallo se atribuye al desajuste entre la optimización de un solo estado y la eliminación de ruido iterativa.
• Una Corrección Simple: Proponer un método para optimizar las ediciones para estados intermedios, lo que efectivamente cierra la brecha de rendimiento para la edición de múltiples tokens en los MDMs sin cambiar la ubicación de la edición ni la estructura del algoritmo subyacente.

Los autores concluyen que portar métodos de edición entre arquitecturas requiere igualar no solo el sitio de la edición, sino también el proceso de generación a través del cual se realiza el hecho editado. Señalan que, aunque su corrección aborda la causa principal del fallo, persiste una brecha de rendimiento en los objetivos más largos, lo que sugiere que trabajos futuros podrían necesitar tener en cuenta el orden específico en el que los MDMs desenmascaran los tokens.