Learning What to Forget: Improving LLM Unlearning via Learned Token-Level Importance

Resumo Técnico: Aprender o que Esquecer: Melhorando o Desaprendizado de LLMs via Importância de Tokens Aprendida

Declaração do Problema
O desaprendizado de máquina (machine unlearning) visa remover conhecimento direcionado (ex: dados pessoais, textos protegidos por direitos autorais, conteúdo nocivo) de Grandes Modelos de Linguagem (LLMs) treinados sem a necessidade de um novo treinamento do zero. Um desafio crítico em LLMs autorregressivos é que nem todos os tokens dentro de uma "amostra para esquecer" (os dados a serem removidos) são igualmente relevantes para a informação alvo. Muitos tokens são estruturais (ex: palavras funcionais, pontuação, padrões sintáticos) e são necessários para as capacidades linguísticas gerais. As abordagens existentes de desaprendizado com pesos por token tentam aplicar seletivamente a perda de esquecimento a tokens específicos, mas dependem de anotações externas, modelos auxiliares, heurísticas baseas em probabilidade ou regras específicas de tarefas para estimar a relevância do token. Esses métodos frequentemente falam em distinguir efetivamente o conteúdo específico do esquecimento dos tokens estruturais, levando ou a um esquecimento insuficiente ou a uma degradação desnecessária da utilidade geral.

Metodologia: Desaprendizado com Pesos de Token Alternados (ATWU)
Os autores propõem que a especificidade de esquecimento ao nível do token deve ser caracterizada pela interação entre os próprios objetivos de retenção e esquecimento, em vez de sinais externos. Eles formalizam isso como um problema de otimização conjunta sobre os parâmetros do modelo ($\theta$) e os pesos dos tokens ($z$).

1. Formulação Teórica: O artigo define uma métrica de "conflito de retenção", $\kappa(A)$, que mede o custo de suprimir um conjunto de tokens $A$ enquanto mantém a otimalidade no conjunto de retenção. Sob uma suposição de "separação de conflito de retenção", os autores provam que o conjunto oráculo de tokens específicos de esquecimento pode ser recuperado como o minimizador de um objetivo conjunto que equilibra a perda de retenção e uma perda de esquecimento ponderada. Este resultado teórico sugere que os pesos dos tokens podem ser aprendidos diretamente do objetivo de desaprendizado sem supervisão externa.
2. Algoritmo (ATWU): Para tornar isso tratável, os autores introduzem o Desaprendizado com Pesos de Token Alternados (ATWU).
   • Pontuador (Scorer): Em vez de aprender pesos independentes para cada token, o ATWU utiliza um pontuador linear leve, $g_w$, parametrizado por um vetor $w$ no espaço do estado oculto do modelo. A pontuação para um token é $\hat{z}_{x,t} = \sigma(\langle w, h_\theta(x_t) \rangle)$, onde $h_\theta$ é a representação oculta.
   • Otimização Conjunta: O framework otimiza um objetivo relaxado e contínuo que inclui a perda de retenção, uma entropia negativa de entropia cruzada saturada (SATGA+) para o esquecimento, um regularizador de entropia para empurrar as pontuações em direção a valores binários e uma restrição de orçamento para controlar o número de tokens selecionados.
   • Atualizações Alternadas: Os parâmetros do modelo $\theta$ e os parâmetros do pontuador $w$ são atualizados alternadamente. Quando $\theta$ é atualizado, as pontuações são desconectadas (detached); quando $w$ é atualizado, $\theta$ é congelado. Isso estabiliza o treinamento ao desacoplar o ciclo de feedback entre a geometria mutável do modelo e a seleção de tokens.

Principais Contribuições

• Caracterização do Conflito de Retenção: O artigo fornece uma caracterização teórica da especificidade de esquecimento ao nível do token baseada na tensão entre os objetivos de retenção e esquecimento, provando que uma otimização conjunta pode recuperar o suporte de esquecimento oráculo sob condições naturais de separação.
• Aprendizado Não Supervisionado de Pesos de Token: O ATWU é um framework leve que aprende a especificidade de esquecimento do token conjuntamente com os parâmetros do modelo usando apenas as representações internas do modelo. Ele não requer supervisão de token externa, modelos auxiliares ou anotações caras.
• Integração de Perda Adaptativa: O método introduz o SATGA+, uma variante pontuada por score da entropia negativa de entropia cruzada saturada, onde o score de token aprendido é injetado no expoente de saturação. Isso permite atualizações mais suaves e conscientes da incerteza durante as fases iniciais de otimização.

Resultados
Os autores avaliam o ATWU em dois benchmarks: TOFU (biografias fictícias de autores sintéticos) e RWKU (conhecimento de figuras públicas do mundo real).

• Desempenho: O ATWU atinge o estado da arte em Qualidade de Desaprendizado (UQ), uma métrica que equilibra qualidade de esquecimento (FQ) e degradação de retenção (RD). No TOFU, atinge um UQ de 91,7 (excedendo o segundo colocado em 5,8 pp) e no RWKU, 58,7 (excedendo o segundo colocado em 6,8 pp).
• Preservação de Utilidade: Ao contrário de métodos de desaprendizado agressivos que colapsam a qualidade de geração, o ATWU preserva a utilidade geral (medida por MMLU, repetitividade e taxas de vitória pareadas) próxima ao checkpoint original.
• Qualidade do Score do Token: Os scores de token aprendidos alinham-se significativamente melhor com os intervalos (spans) de esquecimento específicos da verdade fundamental (medido por AUROC) em comparação com heurísticas baseadas em probabilidade, abordagens de modelos auxiliares e outros baselines de ponderação de tokens.
• Versatilidade: O framework é agnóstico à perda; quando combinado com outros esqueimentos de perda (DPO, NPO, SIMNPO), o ATWU melhora consistentemente o equilíbrio entre esquecimento e retenção em relação aos seus equivalentes não ponderados.

Significância e Alegações
O artigo alega que o "conflito de retenção" fornece um critério intrínseco e eficaz para identificar o que os modelos de linguagem devem esquecer. Ao mudar o desaprendizado ponderado por token de um pipeline de heurísticas externas para um problema de otimização conjunta de ponta a ponta, o ATWU demonstra que:

1. O próprio objetivo de desaprendizado contém sinal latente suficiente para recuperar os tokens específicos de esquecimento sem supervisão externa.
2. O gargalo no desaprendizado seletivo não é meramente encontrar os tokens certos, mas o codisponível principiado do mecanismo de pontuação e do objetivo de esquecimento.
3. Um desaprendizado escalável, preservando a privacidade e altamente direcionado é alcançável sem depender de anotações externas, que frequentemente introduzem riscos de privacidade e problemas de escalabilidade.

Os autores observam limitações, incluindo a dependência de um hiperparâmetro para o orçamento de seleção de tokens e a suposição de uma separação estrita entre tokens específicos de esquecimento e estruturais, que pode ser apenas aproximada na linguagem natural. Além disso, as avaliações são atualmente limitadas a benchmarks específicos e execuções únicas devido aos custos computacionais.