Reinforcement Learning for Neural Model Editing

Resumo Técnico: Aprendizado por Reforço para Edição de Modelos Neurais

Definição do Problema
Modificar redes neurais pré-treinadas para alcançar objetivos específicos — como apagamento de conceitos, mitigação de viés ou incorporação de novos conhecimentos — tradicionalmente requer algoritmos especializados e projetados manualmente. Esses métodos costem ser demorados para serem desenhados e carecem de generalização entre diferentes tarefas. Este artigo investiga se o processo de edição de modelos neurais pode ser automatizado ao formulá-lo como um problema de Aprendizado por Reforço (RL). A questão central de pesquisa é se um agente de RL pode aprender uma política para modificar pesos de redes neurais usando apenas feedback de recompensa, sem acesso a gradientes ou conjuntos de dados construídos manualmente para a tarefa de edição específica.

Metodologia
Os autores propõem um framework exploratório onde um agente interage com uma rede neural ao atualizar iterativamente seus pesos com base em um sinal de recompensa. O framework consiste em dois ambientes distintos:

1. MaskWorld: O agente modifica os pesos escalando-os multiplicativamente. A ação $a_t$ escala os pesos de uma camada específica $l$ de tal forma que $W^{(l)}_t = W^{(l)} \odot a_t$, onde $a_t$ é um vetor de constantes entre 0.0 e 1.0.
2. ShiftWorld: O agente modifica os pesos aplicando atualizações aditivas. A ação $a_t$ desloca os pesos de tal forma que $W^{(l)}_t = W^{(l)} + a_t$, onde $a_t$ contém constantes entre -0.05 e 0.05.

Para abordar a intratabilidade computacional de prever vetores de ação de alta dimensão (ex: $m \times n$ para uma matriz de pesos), os autores empregam um mecanismo inspirado em Low-Rank Adaptation (LoRA) para a tarefa de desaprendizado de máquina (machine unlearning). Em vez de prever a matriz de atualização de pesos completa, o agente prevê duas matrizes de baixo posto, $A$ e $B$, reduzindo a dimensionalidade do espaço de ação e estabilizando a otimização da política.

Função de Recompensa
A função de recompensa $R$ é uma combinação ponderada de dois componentes:
$$R = \alpha U(w) + \beta O(w)$$

• Termo de Utilidade ($U(w)$): Incentiva a preservação do desempenho geral do modelo (ex: acurácia em um conjunto de "retenção" ou dataset geral).
• Termo de Objetivo de Edição ($O(w)$): Quantifica o sucesso do objetivo de edição específico (ex: minimizar a acurácia em um conjunto de "esquecimento" ou reduzir métricas de viés).
• Coeficientes ($\alpha, \beta$): Controlam o equilíbrio entre manter a utilidade e atingir o objetivo de edição. Notavelmente, $\beta$ é definido como um valor negativo em tarefas de desaprendizado para penalizar o desempenho no conjunto de esquecimento.

O agente é treinado usando Proximal Policy Optimization (PPO) com uma rede de política MLP. Crucialmente, o agente não recebe gradientes; ele aprende exclusivamente através de tentativa e erro baseado no desempenho do modelo modificado em lotes (batches) amostrados.

Principais Contribuições e Resultados Experimentais
O framework foi avaliado em duas tarefas distintas: desaprendizado de máquina em classificação de imagens e mitigação de viés em classificação de texto.

• Desaprendizado de Máquina (MNIST): O objetivo era "desaprender" a classe de dígito 7 enquanto mantinha o desempenho nas classes 0–6 e -8–9.

  • Resultados: As políticas aprendidas reduziram a acurácia no conjunto de esquecimento (classe 7) para quase 0% (0.0000 para MaskWorld, 0.0006 para ShiftWorld).
  • Preservação de Utilidade: Contrário às expectativas de degradação de desempenho, os modelos modificados melhoraram ligeiramente a acurácia no conjunto de retenção (de 92.84% para ~93.5%) em comparação ao modelo original.
  • Observação: Os agentes alcançaram esses resultados acessando apenas 15% dos dados de treinamento durante a otimização, demonstrando forte generalização para o conjunto de teste.

• Mitigação de Viés (Jigsaw Toxic Comments): O objetivo era reduzir correlações espúrias (ex: associar a palavra "black" à toxicidade) enquanto mantinha a utilidade de classificação geral.

  • Resultados: Os modelos modificados melhoraram a acurácia em um conjunto de teste sintético sem viés em aproximadamente 6–7% (ex: ShiftWorld melhorou de 55.36% para 62.52%).
  • Preservação de Utilidade: A acurácia do dataset geral permaneceu estável, diminuindo apenas ~1% para MaskWorld, indicando que o agente conseguiu equilibrar a redução de viés com a preservação da utilidade.

Significância e Alegações
O artigo não alega superar algoritmos especializados de última geração nessas tarefas específicas. Em vez disso, sua principal significância reside em demonstrar uma mudança de paradigma: a edição de modelos neurais pode ser tratada como um problema de RL aprendível, em vez de um problema projetado manualmente.

Os autores alegam que suas descobertas sugerem que o RL fornece um framework geral e extensível para aprender políticas de edição a partir de feedback de recompensa. Essa abordagem permite a descoberta de modificações de pesos que alcançam objetivos direcionados (como desaprendizado ou redução de viés) sem exigir o design algorítmico específico para cada novo objetivo. O trabalho serve como uma prova de conceito para motivar pesquisas futuras em abordagens aprendidas para a edição de modelos em áreas como justiça (fairness), robustez e privacidade.

Limitações
Os autores reconhecem que o framework atual enfrenta desafios de escalabilidade inerentes ao RL. Embora o mecanismo inspirado em LoRA ajude a lidar com camadas de alta dimensão, otimizar totalmente todas as camadas de um modelo grande permanece inviável devido à instabilidade causada pelo enorme espaço de ação. Sugere-se que trabalhos futuros explorem configurações multi-agente e abordagens de teoria dos jogos para endereçar essas limitações.