Knowledge Editing in Masked Diffusion Language Models

Resumo Técnico: Edição de Conhecimento em Modelos de Linguagem de Difusão Mascarada

Definição do Problema
A edição de conhecimento visa atualizar ou corrigir conhecimentos factuais específicos dentro de Grandes Modelos de Linguagem (LLMs) sem o retreinamento de todo o modelo. O paradigma dominante, "localizar-então-editar" (ex: ROME, MEMIT), foi desenvolvido exclusivamente para Modelos Autorregressivos (ARMs). Esses métodos baseiam-se na rastreabilidade causal para identificar que a recuperação de fatos está localizada nos MLPs (Perceptrons de Camada Multicamadas) das camadas iniciais a intermediárias no último token do sujeito. Eles então modificam os pesos nesse local para associar um novo objeto ao sujeito.

No entanto, os Modelos de Difusão Mascarada (MDMs) representam um paradigma distinto. Ao contrário dos ARMs, que preveem tokens sequencialmente da esquerda para a direita, os MDMs modelam o texto bidirecionalmente e geram texto através de um processo iterativo de denoising (remoção de ruído), recuperando tokens mascarados a partir de um estado totalmente mascarado. Permanece uma questão em aberto se as suposições subjacentes ao "localizar-então-editar" — especificamente a localização dos fatos e a eficácia da edição nesse local — se sustentam para os MDMs, dado que seus objetivos de treinamento e processos de geração são fundamentalmente diferentes.

Metodologia
Os autores transferem o framework de localizar-então-editar (especificamente o algoritmo MEMIT) para MDMs sem alterar seu algoritmo central, aplicando apenas ajustes mínimos para acomodar o mecanismo de predição mascarada do MDM.

1. Configuração Experimental: O estudo compara dois pares correspondentes de modelos:

   • ARMs: LLaMA-3-8B e Qwen2.5-7B.
   • MDMs: LLaDA-8B (treinado do zero) e Dream-v0-7B (inicializado a partir dos pesos do Qwen2.5).
   • Datasets: COUNTERFACT (benchmark padrão, majoritariamente com alvos de um único token) e KAMEL (adaptado para análise de alvos de múltiplos tokens).

2. Rastreabilidade Causal (Causal Tracing): Os autores realizam a rastreabilidade causal para determinar onde os fatos estão causalmente localizados nos MDMs. Eles corrompem os embeddings dos tokens do sujeito e medem a recuperação das predições corretas ao restaurar os estados ocultos em camadas e posições de tokens específicas.

3. Avaliação da Edição: Eles medem o desempenho da edição (eficácia, generalização, especificidade) através de diferentes posições de tokens e intervalos de camadas. Crucialmente, eles avaliam o desempenho em vários comprimentos de alvo (token único vs. múltiplos tokens) e configurações de geração para identificar falhas específicas do processo de geração do MDM.

4. Correção Proposta: Ao identificar uma degradação de desempenho nos MDMs para alvos de múltiplos tokens, os autores propõem uma correção. A edição padrão otimiza a atualização de peso ($\delta$) para um único estado de condicionamento (o estado totalmente mascarado da resposta). A correção proposta amplia a otimização para incluir estados intermediários parcialmente desmascarados que o modelo percorre durante o denoising iterativo.

Principais Descobertas

1. Localização Compartilhada: A suposição central do localizar-então-editar transfere-se entre paradigmas. Apesar do treinamento bidirecional, os MDMs localizam a recuperação factual no mesmo local que os ARMs: o MLP do último token do sujeito nas camadas iniciais a intermediárias. A rastreabilidade causal confirma que o "sítio inicial" (recuperação do fato) e o "sítio tardio" (leitura da predição) existem em ambos os tipos de modelos.

2. Resultados Divergentes (A Falha de Múltiplos Tokens): Embora a localização seja compartilhada, o resultado não é.

   • Alvos de Token Único: A edição apresenta desempenho comparável entre ARMs e MDMs.
   • Alvos de Múltiplos Tokens: O desempenho da edição nos MDMs degrada-se sistematicamente à medida que o comprimento do alvo aumenta. Por exemplo, para alvos de 4 tokens, a eficácia do LLaDA cai de 0,92 (1 token) para ~0,27, enquanto o LLaMA permanece alto (0,89).
   • Causa Raiz: A falha decorre de um descompasso entre o estado de otimização e o processo de geração. O MEMIT otimiza a edição para o estado totalmente mascarado. Nos ARMs, a geração procede da esquerda para a direita, coincidindo com este estado. Nos MDMs, a geração ocorre através de uma cadeia de estados intermediários onde alguns tokens estão desmascarados e outros permanecem mascarados. A edição, otimizada apenas para o estado totalmente mascarado, falha ao não considerar o contexto variável criado por esses tokens intermediários desmascarados.

3. Restauração via Otimização de Estado Intermediário: Ao modificar o objetivo de otimização para ajustar o resíduo $\delta$ em todos os estados intermediários parcialmente desmascarados (ciclando através das contagens de tokens revelados), os autores restauram substancialmente o desempenho da edição de múltiplos tokens. Para o LLaDA, a eficácia em alvos de 4 tokens sobe de 0,27 para 0,73.

Significância e Alegações
O artigo afirma fornecer a primeira análise de edição de conhecimento do tipo localizar-então-editar em MDMs. Suas principais contribuições são:

• Validação da Localização: Demonstrar que a localização causal dos fatos no MLP do token do sujeito é uma propriedade agnóstica de paradigma, válida tanto para modelos de linguagem autorregressivos quanto para modelos de difusão.
• Identificação de uma Falha Induzida pela Geração: Revelar que os benchmarks padrão (dominados por alvos de um único token) mascaram uma falha sistemática do localizar-então-editar em MDMs para fatos de múltiplos tokens. Esta falha é atribuída ao descompasso entre a otimização de estado único e o denoising iterativo.
• Uma Correção Simples: Propor um método para otimizar edições para estados intermediários, o que efetivamente preenche a lacuna de desempenho para a edição de múltiplos tokens em MDMs sem alterar o local da edição ou a estrutura do algoritmo.

Os autores concluem que portar métodos de edição entre arquiteturas requer o ajuste não apenas do local da edição, mas também do processo de geração através do qual o fato editado é realizado. Eles observam que, embora sua correção trate a causa principal da falha, uma lacuna de desempenho permanece nos alvos mais longos, sugerindo que trabalhos futuros podem precisar considerar a ordem específica em que os MDMs desmascaram os tokens.