Knowledge Editing in Masked Diffusion Language Models

Résumé Technique : Édition de Connaissances dans les Modèles de Langage de Diffusion Masquée

Énoncé du Problème
L'édition de connaissances vise à mettre à jour ou à corriger des connaissances factuelles spécifiques au sein des Grands Modèles de Langage (LLM) sans réentraîner l'intégralité du modèle. Le paradigme dominant, « localiser puis éditer » (ex. ROME, MEMIT), a été développé exclusivement pour les Modèles Autorégressifs (ARM). Ces méthodes reposent sur le traçage causal pour identifier que le rappel de faits est localisé dans les couches MLP (Perceptrons Multicouches) précoces à intermédiaires, au niveau du dernier jeton (token) du sujet. Elles modifient ensuite les poids à cet emplacement pour associer un nouvel objet au sujet.

Cependant, les Modèles de Diffusion Masquée (MDM) représentent un paradigme distinct. Contrairement aux ARM, qui prédisent les jetons séquentiellement de gauche à droite, les MDM modélisent le texte de manière bidirectionnelle et génèrent du texte par un processus de débruitage itératif, récupérant les jetons masqués à partir d'un état entièrement masqué. Il reste une question ouverte de savoir si les hypothèses sous-jacentes à la méthode « localiser puis éditer » — spécifiquement la localisation des faits et l'efficacité de l'édition à cet emplacement — sont vérifiées pour les MDM, étant donné leurs objectifs d'entraînement et leurs processus de génération fondamentalement différents.

Méthodologie
Les auteurs transfèrent le cadre de travail « localiser puis éditer » (spécifiquement l'algorithme MEMIT) aux MDM sans en modifier l'algorithme central, en appliquant seulement des ajustements minimaux pour s'adapter au mécanisme de prédiction masquée des MDM.

1. Configuration Expérimentale : L'étude compare deux paires de modèles appariés :

   • ARM : LLaMA-3-8B et Qwen2.5-7B.
   • MDM : LLaDA-8B (entraîné de zéro) et Dream-v0-7B (initialisé à partir des poids de Qwen2.5).
   • Jeux de données : COUNTERFACT (benchmark standard, principalement des cibles à jeton unique) et KAMEL (adapté pour l'analyse de cibles à multi-jetons).

2. Traçage Causal : Les auteurs effectuent un traçage causal pour déterminer où les faits sont causalement localisés dans les MDM. Ils corrompent les plongements (embeddings) des jetons du sujet et mesurent la récupération des prédictions correctes en restaurant les états cachés à des couches et des positions de jetons spécifiques.

3. Évaluation de l'Édition : Ils mesurent la performance de l'édition (efficacité, généralisation, spécificité) à travers différents emplacements de jetons et plages de couches. Crucialement, ils évaluent la performance selon des longueurs de cible variables (jeton unique vs multi-jetons) et des configurations de génération pour identifier les ruptures spécifiques au processus de génération des MDM.

4. Correction Proposée : Suite à l'identification d'une dégradation de la performance dans les MDM pour les cibles multi-jetons, les auteurs proposent une correction. L'édition standard optimise la mise à jour du poids ($\delta$) pour un état de conditionnement unique (la réponse entièrement masquée). La correction proposée élargit l'optimisation pour inclure des états intermédiaires partiellement démasqués que le modèle parcourt durant le débruitage itératif.

Principaux Résultats

1. Localisation Partagée : L'hypothèse centrale de « localiser puis éditer » se transfère entre les paradigmes. Malgré l'entraînement bidirectionnel, les MDM localisent le rappel factuel au même endroit que les ARM : le MLP du jeton du sujet dans les couches précoces à intermédiaires. Le traçage causal confirme que le « site précoce » (rappel de faits) et le « site tardif » (lecture de la prédiction) existent dans les deux types de modèles.

2. Résultats Divergents (La Rupture Multi-Jeton) : Bien que l'emplacement soit partagé, le résultat ne l'est pas.

   • Cibles à Jeton Unique : L'édition est comparable entre les ARM et les MDM.
   • Cibles Multi-Jetons : La performance de l'édition dans les MDM se dégrade systématiquement à mesure que la longueur de la cible augmente. Par exemple, pour des cibles de 4 jetons, l'efficacité de LLaDA chute d'environ 0,92 (1 jeton) à 0,27, tandis que LLaMA reste élevé (~0,89).
   • Cause Racine : L'échec provient d'une inadéquation entre l'état d'optimisation et le processus de génération. MEMIT optimise l'édition pour l'état entièrement masqué. Dans les ARM, la génération progresse de gauche à droite, correspondant à cet état. Dans les MDM, la génération progresse à travers une chaîne d'états intermédiaires où certains jetons sont démasqués et d'autres restent masqués. L'édition, optimisée uniquement pour l'état entièrement masqué, ne parvient pas à prendre en compte le contexte changeant créé par ces jetons intermédiaires démasqués.

3. Restauration via l'Optimisation d'États Intermédiaires : En modifiant l'objectif d'optimisation pour ajuster le résidu $\delta$ à travers tous les états intermédiaires partiellement démasqués (en cyclant à travers le nombre de jetons révélés), les auteurs restaurent substantiellement la performance d'édition multi-jetons. Pour LLaDA, l'efficacité sur les cibles de 4 jetons passe de 0,27 à 0,73.

Signification et Revendications
L'article affirme fournir la première analyse de l'édition de connaissances de type « localiser puis éditer » dans les MDM. Ses principales contributions sont :

• Validation de la Localisation : Démontrer que la localisation causale des faits au MLP du jeton du sujet est une propriété indépendante du paradigme, s'appliquant tant aux modèles de langage autorégressifs qu'aux modèles de diffusion.
• Identification d'un Échec Induit par la Génération : Révéler que les benchmarks standards (dominés par des cibles à jeton unique) masquent un échec systématique de « localiser puis éditer » dans les MDM pour les faits multi-jetons. Cet échec est attribué au décalage entre l'optimisation à état unique et le débruitage itératif.
• Une Correction Simple : Proposer une méthode pour optimiser les éditions pour les états intermédiaires, ce qui comble efficacement l'écart de performance pour l'édition multi-jetons dans les MDM sans changer l'emplacement de l'édition ni la structure de l'algorithme.

Les auteurs concluent que porter des méthodes d'édition d'une architecture à l'autre nécessite de faire correspondre non seulement le site de l'édition, mais aussi le processus de génération par lequel le fait édité est réalisé. Ils notent que bien que leur correction traite la cause principale de l'échec, un écart de performance subsiste pour les cibles les plus longues, suggérant que les travaux futurs devront peut-être prendre en compte l'ordre spécifique dans lequel les MDM démasquent les jetons.