Unlearning What Matters: Token-Level Attribution for Precise Language Model Unlearning

1. Declaração do Problema

Os Grandes Modelos de Linguagem (LLMs) frequentemente memorizam dados sensíveis, protegidos por direitos autorais ou prejudiciais, tornando necessário o aprendizado desmembrado (machine unlearning) para remover conhecimento específico sem retreinar do zero.

• Limitação Atual: Os métodos de desmembramento existentes operam no nível de sequência, aplicando atualizações de gradiente uniformes em todos os tokens de uma sequência.
• O Problema Central: O conhecimento não é codificado de forma uniforme; apenas um subconjunto esparsos de tokens (por exemplo, entidades específicas, datas, fatos) carrega a informação-alvo, enquanto outros servem a papéis estruturais ou contextuais.
• Consequências das Atualizações Uniformes:
  1. Ruído de Gradiente: Atualizar tokens não críticos introduz ruído, degradando a relação sinal-ruído (SNR) do processo de desmembramento.
  2. Degradação da Utilidade: Atualizações indiscriminadas prejudicam o desempenho do modelo em conhecimento não relacionado (danos colaterais).
  3. Esquecimento Subótimo: O modelo falha em remover efetivamente o conhecimento-alvo porque o sinal de gradiente é diluído por tokens irrelevantes.

2. Metodologia: TokenUnlearn

Os autores propõem o TokenUnlearn, um framework que desloca o desmembramento do nível de sequência para o nível de token, identificando e selecionando seletivamente tokens "críticos para o conhecimento".

A. Mecanismo de Atribuição de Token

Para identificar quais tokens codificam o conhecimento a ser esquecido, o framework calcula uma pontuação de importância composta ($\phi_i$) para cada token usando dois sinais:

1. Atribuição Consciente do Conhecimento (Mascaramento):
   • Mecanismo: O modelo gera previsões na entrada original ($z_{orig}$) e em uma entrada onde as porções relevantes para o conhecimento (especificamente substantivos no prompt) são mascaradas ($z_{mask}$).
   • Métrica: O deslocamento absoluto do log-probabilidade ($\Delta^{unlearn}_i$) é calculado. Um grande deslocamento indica que a previsão do token depende fortemente do conhecimento-alvo.
2. Sinal Complementar Consciente de Entropia:
   • Mecanismo: Calcula a entropia preditiva ($H_i$) do token.
   • Racional: Alta entropia pode indicar pontos de decisão onde o modelo está incerto sobre continuições dependentes de conhecimento (por exemplo, escolher entre múltiplos fatos plausíveis).
   • Combinação: A pontuação final é uma soma ponderada: $\phi_i = \alpha \cdot \bar{\Delta}^{unlearn}_i + (1-\alpha) \cdot \bar{H}_i$.

B. Estratégias de Otimização Seletiva de Token

Uma vez geradas as pontuações de importância, o framework aplica objetivos de desmembramento apenas aos tokens de alta importância usando duas estratégias:

1. Seleção Rígida (Hard Selection):
   • Seleciona a fração top-$r$ de tokens (por exemplo, top 20%) com base em $\phi_i$.
   • Aplica a perda de desmembramento exclusivamente a esses tokens; outros tokens recebem zero atualizações de gradiente.
2. Ponderação Suave (Soft Weighting):
   • Modula a contribuição do gradiente de cada token com base em uma pontuação de importância normalizada (usando um softmax com temperatura $\tau$).
   • Permite uma interpolação suave entre atualizações uniformes e atualizações totalmente seletivas.

C. Integração com Métodos Existentes

O framework é agnóstico ao modelo e integra-se a objetivos padrão de desmembramento baseados em gradiente, incluindo:

• Ascensão de Gradiente (GA)
• Ascensão de Gradiente Ponderada (WGA)
• Otimização de Preferência Negativa (NPO)
• Desmembramento por Misdirecionamento de Representação (RMU)
• Retenção: Todos os métodos são combinados com regularização de divergência KL no "conjunto de retenção" para preservar capacidades gerais.

3. Análise Teórica

O artigo fornece uma justificativa teórica para por que a seleção no nível de token melhora o desmembramento:

• Relação Sinal-Ruído (SNR): Os autores modelam o desmembramento como movimento ao longo de um subespaço de baixa dimensionalidade ($U$) de parâmetros.
• Redução de Ruído: Atualizações uniformes acumulam ruído de todos os $T$ tokens. A seleção no nível de token concentra o peso do gradiente no subconjunto $K$ de tokens críticos para o conhecimento.
• Resultado: A SNR do estimador de gradiente melhora proporcionalmente ao inverso da taxa de seleção ($T/|K|$). Ao suprimir atualizações em tokens não críticos, o método reduz o ruído de gradiente e preserva o conhecimento não relacionado de forma mais eficaz.

4. Resultados Experimentais

Os autores avaliaram o TokenUnlearn em dois benchmarks através de três arquiteturas de modelo (Phi-1.5, Llama-2-7B, Qwen-3-8B).

Benchmarks

1. TOFU (Task of Fictitious Unlearning): Desmembramento de conhecimento granular (removendo fatos sobre autores fictícios).
2. WMDP (Weapons of Mass Destruction Proxy): Alinhamento de segurança (removendo conhecimento perigoso de bio/cibersegurança).

Principais Descobertas

• Esquecimento Superior: Variantes no nível de token (T-* e S-*) superaram consistentemente as linhas de base no nível de sequência na eficácia do esquecimento (menor precisão nos conjuntos de esquecimento).
  • Exemplo: No TOFU com Qwen-3-8B, o T-WGA melhorou significativamente a eficácia do esquecimento em comparação com o WGA padrão.
• Melhor Preservação de Utilidade: Métodos no nível de token mantiveram maior utilidade do modelo (pontuações de retenção) em dados não-alvo.
  • Exemplo: No Qwen-3-8B, a pontuação de retenção melhorou em até 19,0% em comparação com as linhas de base no nível de sequência, enquanto alcançava melhor esquecimento.
• Generalizabilidade: A abordagem funcionou consistentemente em diferentes tamanhos e arquiteturas de modelos.
• Estudos de Ablação:
  • Taxa de Seleção ($r$): Uma taxa de 20% (top 20% dos tokens) forneceu o compromisso ótimo entre esquecimento e retenção.
  • Peso de Entropia: Combinar sinais de entropia (peso aproximado de 30%) forneceu benefícios complementares, capturando pontos de decisão incertos que o mascaramento sozinho poderia perder.

5. Contribuições Principais

1. Framework de Atribuição no Nível de Token: Introduziu um método principiado para identificar tokens críticos para o conhecimento usando mascaramento contrafactual e sinais de entropia sem exigir modelos auxiliares.
2. Estratégias de Desmembramento Seletivo: Propôs estratégias de "Seleção Rígida" e "Ponderação Suave" que estendem algoritmos de desmembramento existentes para operar no nível de token.
3. Insight Teórico: Provou que atualizações seletivas de token melhoram a SNR do gradiente, explicando por que atualizações uniformes falham em otimizar os objetivos duais de esquecimento e retenção.
4. Validação Empírica: Demonstrou melhorias consistentes em três modelos e dois benchmarks distintos (conhecimento e segurança), alcançando até 32,6% de melhoria no esquecimento e 19,0% de melhoria na preservação de utilidade.

6. Significado e Impacto

• Privacidade e Regulamentação: Permite conformidade mais precisa com o "Direito ao Esquecimento" (GDPR/CCPA), removendo rastros de dados específicos sem degradar a utilidade geral do modelo.
• Segurança: Oferece um mecanismo mais eficaz para remover capacidades perigosas (por exemplo, ameaças de biosegurança) de modelos pré-treinados.
• Eficiência: Desafia a suposição de que o desmembramento requer atualizações uniformes de sequência, sugerindo que o conhecimento em LLMs é esparsos e pode ser visado com precisão.
• Limitações: O método introduz sobrecarga computacional devido à passagem forward adicional para mascaramento. É um método aproximado e não garante a remoção de 100% dos rastros de conhecimento (rastros residuais podem existir).

Em conclusão, o TokenUnlearn representa uma mudança de paradigma no aprendizado desmembrado, movendo-se de atualizações de sequência de granularidade grosseira para atribuição de token de granularidade fina, resolvendo assim a tensão fundamental entre esquecimento eficaz e preservação de utilidade.