Resumo Técnico: Replay What Matters: Off-Policy Replay para Unlearning Eficiente de LLMs

Declaração do Problema

O unlearning (desaprendizado) de Grandes Modelos de Linguagem (LLMs) visa remover conhecimento perigoso (ex: conteúdo protegido por direitos autorais, informações privadas) de modelos pré-treinados sem o custo proibitivo de um retreinamento completo. Abordagens recentes baseadas em aprendizado por reforço (RL), como o RULE, reformulam o unlearning como o aprendizado de um comportamento de recusa: ensinar o modelo a recusar consultas relacionadas a um "conjunto de esquecimento" (forget set) enquanto preserva respostas em um "conjunto de retenção" (retain set) e um "conjunto de fronteira" (boundary set).

No entanto, os autores identificam uma ineficiência crítica na otimização on-policy da Etapa 2 (especificamente usando GRPO) empregada por métodos como o RULE. Durante o treinamento, surge uma pronunciada disparidade entre casos fáceis e difíceis (hard–easy disparity):

1. Casos Fáceis: Muitos prompts convergem rapidamente para estados de alta recompensa. Continuar a amostrar esses prompts em estágios posteriores do treinamento gera pouco sinal de gradiente informativo (vantagem $\approx 0$), mas consome recursos computacionais significativos.
2. Casos Difíceis: Prompts próximos à fronteira semântica entre o esquecimento e a retenção frequentemente permanecem com baixa recompensa durante todo o treinamento. Em configurações on-policy padrão, as rollouts de baixa recompensa geradas para esses casos difíceis são descartadas após um único uso, apesar de conterem informações valiosas sobre os modos de falha do modelo.

Essa dinâmica leva a uma subotimização da eficiência do treinamento, onde a computação é desperdiçada em casos fáceis já convergidos, enquanto casos difíceis permanecem subotimizados, potencialmente degradando a capacidade do modelo de preservar a utilidade geral (Qualidade de Retenção) ao esquecer alvos específicos.

Metodologia: ReRULE

Para abordar isso, os autores propõem o ReRULE (Replay-enhanced Reinforcement UnLEarning), um mecanismo de aprimoramento de replay off-policy integrado ao framework RULE.

Visão Geral do Framework

O ReRULE mantém a estrutura de duas etapas do RULE:

• Etapa 1 (Rejection Steering): Ajuste Fino Supervisionado (SFT) em um conjunto de dados de recusa para iniciar o aquecimento da capacidade de recusa do modelo.
• Etapa 2 (Refusal Boundary Optimization): Otimização baseada em RL para refinar a fronteira entre consultas de esquecimento e retenção.

O Pipeline ReRULE

A inovação central ocorre na Etapa 2, onde o ReRULE introduz um buffer de replay e treinamento híbrido:

1. Identificação e Armazenamento de Casos Difíceis:
   Durante o início do treinamento GRPO, o sistema monitora a recompensa média dos grupos de rollouts gerados para cada prompt. Se a média de um grupo cair abaixo de um limiar predefinido $\tau$, ele é classificado como um "caso difícil" e armazenado em um buffer de replay $\mathcal{B}$. Este buffer acumula trajetórias de baixa recompensa que representam os modos de falha atuais do modelo.

2. Treinamento Híbrido (Atualizações Off-Policy):
   Uma vez que o buffer contém um número suficiente de casos difíceis ($|\mathcal{B}| \ge B_{min}$) e um período de aquecimento ($t \ge T_{warm}$) é concluído, o ReRULE ativa um regime de treinamento híbrido. Em cada etapa, ele realiza:

   • Atualização On-Policy Padrão: Uma atualização GRPO padrão usando amostras frescas da política atual $\pi_\theta$.
   • Atualização de Replay Off-Policy: Uma atualização de gradiente adicional usando amostras extraídas do buffer de replay $\mathcal{B}$.

3. Amostragem de Importância (Importance Sampling):
   Para corrigir o desvio de distribuição entre a política que gerou as rollouts armazenadas ($\pi_{\theta_{old}}$) e a política atual ($\pi_\theta$), o ReRULE utiliza amostragem de importância. O gradiente off-policy é computado como:
   $$\nabla_\theta J_{off}(\theta) = \mathbb{E}_{g \sim \mathcal{B}} \left[ \mathbb{E}_{y \sim \pi_{\theta_{old}}} \left[ \tilde{\rho}(y|x) \nabla_\theta \log \pi_\theta(y|x) \hat{A}(x, y) \right] \right]$$
   onde $\tilde{\rho}(y|x)$ é um peso de importância clipado e normalizado. Isso permite que o modelo retreine em casos difíceis sem exigir novas e caras amostragens.

Principais Contribuições

1. Identificação de Ineficiência de Treinamento: O artigo demonstra empiricamente que o GRPO on-policy no unlearning de LLMs sofre de retornos decrescentes em casos fáceis e subutilização de casos difíceis, levando a uma "disparidade hard/easy" que dificulta a convergência.
2. Proposta do ReRULE: Os autores introduzem um mecanismo de replay off-policy que armazena rollouts de casos difíceis de baixa recompensa e as reutiliza via atualizações de amostragem de importância. Isso redireciona o foco computacional para casos de fronteira que exigem maior otimização.
3. Análise Teórica: Os autores fornecem uma prova teórica mostrando que o ReRULE produz um limite de convergência mais apertado para casos difíceis em comparação ao RULE puramente on-policy. Ao aumentar o número efetivo de tokens de treinamento na região difícil ($N_{H}^{ReRULE} = N_{H}^{on} + \eta_\rho N_{H}^{off}$), o método reduz o erro de exploração na região difícil.
4. Validação Empírica: Experimentos extensos em três benchmarks (RWKU, MUSE-Books, TOFU) validam a eficácia do método.

Resultados Experimentais

Os autores avaliaram o ReRULE contra baselines incluindo GA, NPO, SimNPO e o RULE original em três conjuntos de dados:

• MUSE-Books (Unlearning de Direitos Autorais): Este cenário, envolvendo conteúdo de Harry Potter, apresenta uma alta proporção de casos difíceis. O ReRULE melhorou significativamente a Qualidade de Retenção de 46.3 (RULE) para 56.2, mantendo uma Qualidade de Esquecimento de alto nível. O tempo de treinamento aumentou apenas 5–11% (de 6.8h para 7.1h).
• RWKU (Conhecimento do Mundo Real): O ReRULE alcançou a melhor Qualidade de Esquecimento e manteve o top-2 desempenho em Naturalidade de Esquecimento e Qualidade de Retenção, superando o RULE em eficiência em casos difíceis.
• TOFU (Unlearning de Entidades): Neste cenário, o conjunto de esquecimento é menor e geralmente mais fácil (menos casos difíceis). O ReRULE mostrou melhorias limitadas em relação ao RULE, o que os autores citam como evidência de que o mecanismo de replay é mais benéfico quando existe uma disparidade pronunciada entre hard/easy.
• Estudos de Ablação:
  • Substituir a seleção de casos difíceis por armazenamento aleatório causou a degradação do desempenho para o nível do RULE padrão, confirmando que o direcionamento para casos difíceis é crucial.
  • Variar o limiar $\tau$ mostrou que configurá-lo muito alto (incluindo casos mais fáceis) reduz a eficiência, enquanto a configuração ideal (0.4 para MUSE) equilibra seletividade e tamanho do buffer.

Significância e Alegações

O artigo afirma que o ReRULE oferece um aprimoramento custo-efetivo para o unlearning baseado em RL. Ao reciclar rollouts "falhas" de casos difíceis, ele melhora a Qualidade de Retenção do modelo desaprendido — uma métrica crítica frequentemente comprometida em tarefas de unlearning — sem exigir retreinamento total ou computação adicional significativa.

Os autores enfatizam que o método é condicional: seus benefícios são mais pronunciados em cenários com uma disparidade significativa entre prompts fáceis e difíceis (ex: MUSE-Books). Em configurações mais simples (ex: TOFU), os ganhos são marginais, apoiando a hipótese de que o mecanismo de replay visa especificamente a ineficiência de casos difíceis subotimizados. O trabalho posiciona o ReRULE como um passo prático para um machine unlearning mais eficiente e robusto, particularmente para preservar a utilidade geral do modelo enquanto remove conhecimento perigoso específico.