Knowledge Editing in Masked Diffusion Language Models

Technische Samenvatting: Kennisbewerking in Masked Diffusion Taalmodellen

Probleemstelling
Kennisbewerking (knowledge editing) beoogt specifieke feitelijke kennis binnen Large Language Models (LLMs) bij te werken of te corrigeren zonder het volledige model opnieuw te trainen. Het dominante paradigma, "locate-then-edit" (bijv. ROME, MEMIT), is exclusief ontwikkeld voor Autoregressieve Modellen (ARMs). Deze methoden vertrouwen op causale tracing om te identificeren dat feitelijke herinnering gelokaliseerd is in de vroege-tot-midden-laag Multi-Layer Perceptrons (MLPs) bij het laatste token van het onderwerp. Vervolgens modificeren ze de gewichten op deze locatie om een nieuw object aan het onderwerp te koppelen.

Echter, Masked Diffusion Models (MDMs) vertegenwoordigen een verschillend paradigma. In tegenstelling tot ARMs, die tokens sequentieel van links naar rechts voorspellen, modelleren MDMs tekst bidirectioneel en genereren ze tekst door middel van iteratieve denoising, waarbij gemaskeerde tokens worden hersteld vanuit een volledig gemaskeerde staat. Het blijft een open vraag of de aannames die ten grondslag liggen aan "locate-then-edit"—specifiek de lokalisatie van feiten en de effectiviteit van bewerking op die locatie—standhouden voor MDMs, gezien hun fundamenteel andere trainingsdoelstellingen en generatieprocessen.

Methodologie
De auteurs dragen het "locate-then-edit" raamwerk (specifiek het MEMIT-algoritme) over naar MDMs zonder het kernalgoritme te wijzigen, waarbij slechts minimale aanpassingen worden gedaan om rekening te houden met het gemaskeerde voorspellingsmechanisme van de MDM.

1. Experimentele Opzet: De studie vergelijkt twee bij elkaar passende paren modellen:

   • ARMs: LLaMA-3-8B en Qwen2.5-7B.
   • MDMs: LLaDA-8B (vanaf nul getraind) en Dream-v0-7B (geïnitialiseerd vanaf Qwen2.5-gewichten).
   • Datasets: COUNTERFACT (standaard benchmark, voornamelijk single-token targets) en KAMEL (aangepast voor multi-token target analyse).

2. Causale Tracing: De auteurs voeren causale tracing uit om te bepalen waar feiten causaal gelokaliseerd zijn in MDMs. Ze corrumperen de embeddings van het onderwerp-token en meten het herstel van correcte voorspellingen door de verborgen toestanden (hidden states) te herstellen op specifieke lagen en tokenposities.

3. Evaluatie van Bewerking: Ze meten de prestaties van de bewerking (effectiviteit, generalisatie, specificiteit) over verschillende tokenposities en laagbereiken. Cruciaal is dat ze de prestaties evalueren over variërende target-lengtes (single-token versus multi-token) en generatie-instellingen om breakdowns te identificeren die specifiek zijn voor het MDM-generatieproces.

4. Voorgestelde Correctie: Na het identificeren van een prestatieverslechtering in MDMs voor multi-token targets, stellen de auteurs een correctie voor. De standaard bewerking optimaliseert de gewichtsupdate ($\delta$) voor een enkele conditionerende staat (de volledig gemaskeerde antwoordstaat). De voorgestelde correctie breidt de optimalisatie uit naar intermediaire gedeeltelijk ongemaskeerde staten die het model doorloopt tijdens iteratieve denoising.

Belangrijkste Bevindingen

1. Gedeelde Lokalisatie: De kernaanname van "locate-then-edit" wordt overgedragen tussen paradigma's. Ondanks bidirectionele training lokaliseren MDMs feitelijke herinnering naar dezelfde locatie als ARMs: de vroege-tot-midden-laag MLP bij het laatste onderwerp-token. Causale tracing bevestigt dat zowel de "early site" (feitelijke herinnering) als de "late site" (voorspellings-readout) bestaan in beide modeltypen.

2. Divergerende Resultaten (De Multi-Token Breakdown): Hoewel de locatie gedeeld wordt, is de uitkomst dat niet.

   • Single-Token Targets: Bewerking presteert vergelijkbaar in ARMs en MDMs.
   • Multi-Token Targets: De prestaties van de bewerking in MDMs verslechteren systematisch naarmate de target-lengte toeneemt. Bijvoorbeeld, bij 4-token targets daalt de effectiviteit van LLaDA van 0,92 (1-token) naar ~0,27, terwijl LLaMA hoog blijft (0,89).
   • Oorzaak: De mislukking komt voort uit een mismatch tussen de optimalisatiestatus en het generatieproces. MEMIT optimaliseert de bewerking voor de volledig gemaskeerde staat. In ARMs verloopt de generatie van links naar rechts, wat overeenkomt met deze staat. In MDMs verloopt de generatie via een keten van intermediaire staten waarbij sommige tokens ongemaskeerd zijn en andere gemaskeerd blijven. De bewerking, geoptimaliseerd voor enkel de volledig gemaskeerde staat, houdt geen rekening met de verschuivende context die door deze intermediaire ongemaskeerde tokens wordt gecreëerd.

3. Herstel via Intermediaire Staat Optimalisatie: Door de optimalisatiedoelstelling aan te passen om de residu $\delta$ over alle intermediaire gedeeltelijk ongemaskeerde staten te fitten (doorheen cycli van het aantal onthulde tokens), herstellen de auteurs de multi-token bewerkingsprestaties aanzienlijk. Voor LLaDA stijgt de effectiviteit op 4-token targets van 0,27 naar 0,73.

Betekenis en Claims
Het artikel claimt de eerste analyse te bieden van "locate-then-edit" kennisbewerking in MDMs. De primaire bijdragen zijn:

• Validatie van Lokalisatie: Het aantonen dat de causale lokalisatie van feiten naar de onderwerp-token MLP een paradigma-agnostische eigenschap is, die standhoudt voor zowel autoregressieve als diffusie-gebaseerde taalmodellen.
• Identificatie van een Generatie-geïnduceerde Fout: Het onthullen dat standaard benchmarks (gedomineerd door single-token targets) een systematische fout van "locate-then-edit" in MDMs voor multi-token feiten maskeren. Deze fout wordt toegeschreven aan de mismatch tussen single-state optimalisatie en iteratieve denoising.
• Een Eenvoudige Correctie: Het voorstellen van een methode om bewerkingen te optimaliseren voor intermediaire staten, wat het prestatieverschil voor multi-token bewerking in MDMs effectief overbrugt zonder de onderliggende locatie of algoritmische structuur te veranderen.

De auteurs concluderen dat het porten van bewerkingsmethoden naar andere architecturen vereist dat niet alleen de locatie van de bewerking, maar ook het generatieproces waarbij het bewerkte feit wordt gerealiseerd, wordt afgestemd. Ze merken op dat hoewel hun correctie de primaire oorzaak van de fout aanpakt, er een prestatiekloof blijft bestaan bij de langste targets, wat suggereert dat toekomstig werk mogelijk rekening moet houden met de specifieke volgorde waarin MDMs tokens onmaskeren.