Resumen Técnico: Replay What Matters: Off-Policy Replay for Efficient LLM Reinforcement Unlearning

Declaración del Problema

El desaprendizaje de Grandes Modelos de Lenguaje (LLM) busca eliminar el conocimiento peligroso (por ejemplo, contenido protegido por derechos de autor, información privada) de los modelos preentrenados sin el costo prohibitivo de un reentrenamiento completo. Los enfoques recientes basados en aprendizaje por refuerzo (RL), como RULE, reformulan el desaprendizaje como el aprendizaje de un comportamiento de rechazo: enseñar al modelo a rechazar consultas relacionadas con un "conjunto de olvido" (forget set) mientras preserva las respuestas en un "conjunto de retención" (retain set) y un "conjunto de frontera" (boundary set).

Sin embargo, los autores identifican una ineficiencia crítica en la optimización on-policy de la Etapa 2 (específicamente utilizando GRPO) empleada por métodos como RULE. Durante el entrenamiento, surge una pronunciada disparidad entre casos fáciles y difíciles (hard–easy disparity):

1. Casos Fáciles: Muchos prompts convergen rápidamente a estados de alta recompensa. Continuar muestreando estos prompts en etapas posteriores del entrenamiento genera poca señal de gradiente informativa (ventaja $\approx 0$) pero consume recursos computacionales significativos.
2. Casos Difíciles: Los prompts cerca de la frontera semántica entre el olvido y la retención a menudo permanecen con recompensas bajas durante todo el entrenamiento. En entornos on-policy estándar, los rollouts de baja recompensa generados para estos casos difíciles se descartan tras un único uso, a pesar de que contienen información valiosa sobre los modos de falla del modelo.

Esta dinámica conduce a una eficiencia de entrenamiento subóptima, donde la computación se desperdicia en casos fáciles ya convergidos mientras que los casos difíciles permanecen suboptimizados, lo que potencialmente degrada la capacidad del modelo para preservar la utilidad general (Calidad de Retención).

Metodología: ReRULE

Para abordar esto, los autores proponen ReRULE (Replay-enhanced Reinforcement UnLEarning), un mecanismo de mejora de replay off-policy integrado en el marco de trabajo de RULE.

Visión General del Marco de Trabajo

ReRULE mantiene la estructura de dos etapas de RULE:

• Etapa 1 (Rejection Steering): Ajuste Fino Supervisado (SFT) sobre un conjunto de datos de rechazo para iniciar el calentamiento de la capacidad del modelo para rechazar.
• Etapa 2 (Refusal Boundary Optimization): Optimización basada en RL para refinar la frontera entre las consultas de olvido y de retención.

El Pipeline de ReRULE

La innovación central ocurre en la Etapa 2, donde ReRULE introduce un buffer de replay y un entrenamiento híbrido:

1. Identificación y Almacenamiento de Casos Difíciles:
   Durante el entrenamiento temprano de GRPO, el sistema monitorea la recompensa media de los grupos de rollouts generados para cada prompt. Si la recompensa media de un grupo cae por debajo de un umbral predefinido $\tau$, se clasifica como un "caso difícil" y se almacena en un buffer de replay $\mathcal{B}$. Este buffer acumula trayectorias de baja recompensa que representan los modos de falla actuales del modelo.

2. Entrenamiento Híbrido (Actualizaciones Off-Policy):
   Una vez que el buffer contiene un número suficiente de casos difíciles ($|\mathcal{B}| \ge B_{min}$) y se completa un periodo de calentamiento ($t \ge T_{warm}$), ReRULE activa un régimen de entrenamiento híbrido. En cada paso, realiza:

   • Actualización On-Policy Estándar: Una actualización GRPO estándar utilizando muestras frescas de la política actual $\pi_\theta$.
   • Actualización de Replay Off-Policy: Una actualización de gradiente adicional utilizando muestras extraídas del buffer de replay $\mathcal{B}$.

3. Muestreo de Importancia (Importance Sampling):
   Para corregir el cambio de distribución entre la política que generó los rollouts almacenados ($\pi_{\theta_{old}}$) y la política actual ($\pi_\theta$), ReRULE emplea muestreo de importancia. El gradiente off-policy se calcula 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]$$
   donde $\tilde{\rho}(y|x)$ es un peso de importancia recortado y normalizado. Esto permite al modelo reentrenar en casos difíciles sin requerir un nuevo y costoso muestreo.

Contribuciones Clave

1. Identificación de la Ineficiencia de Entrenamiento: El artículo demuestra empíricamente que el GRPO on-policy en el desaprendizaje de LLM sufre de rendimientos decrecientes en casos fáciles y una subutilización de los casos difíciles, lo que conduce a una "disparidad entre casos fáciles y difíciles" que obstaculiza la convergencia.
2. Propuesta de ReRULE: Los autores introducen un mecanismo de replay off-policy que almacena rollouts de casos difíciles de baja recompensa y los reutiliza mediante actualizaciones de muestreo de importancia. Esto redirige el enfoque computacional hacia los casos de frontera que requieren mayor optimización.
3. Análisis Teórico: Los autores proporcionan una prueba teórica que muestra que ReRULE produce un límite de convergencia más ajustado para los casos difíciles en comparación con el RULE puramente on-policy. Al aumentar el número efectivo de tokens de entrenamiento en la región difícil ($N_{H}^{ReRULE} = N_{H}^{on} + \eta_\rho N_{H}^{off}$), el método reduce el error de exploración en la región difícil.
4. Validación Empírica: Experimentos extensos en tres benchmarks (RWKU, MUSE-Books, TOFU) validan la efectividad del método.

Resultados Experimentales

Los autores evaluaron ReRULE frente a baselines que incluyen GA, NPO, SimNPO y el RULE original a través de tres conjuntos de datos:

• MUSE-Books (Desaprendizaje de Contenido Protegido): Este entorno, que involucra contenido de Harry Potter, presenta una alta proporción de casos difíciles. ReRULE mejoró significativamente la Calidad de Retención de 46.3 (RULE) a 56.2, manteniendo una Calidad de Olvido de primer nivel. El tiempo de entrenamiento aumentó solo un 5–11% (de 6.8h a 7.1h).
• RWKU (Conocimiento del Mundo Real): ReRULE logró la mejor Calidad de Olvido y mantuvo un rendimiento top-2 en Naturalidad de Olvido y Calidad de Retención, superando a RULE en eficiencia en casos difíciles.
• TOFU (Desaprendizaje de Entidades): En este entorno, el conjunto de olvido es más pequeño y generalmente más fácil (menos casos difíciles). ReRULE mostró una mejora limitada respecto a RULE, lo que los autores citan como evidencia de que el mecanismo de replay es más beneficioso cuando existe una disparidad pronunciada entre casos fáciles y difíciles.
• Estudios de Ablación:
  • Reemplazar la selección de casos difíciles con almacenamiento aleatorio causó que el rendimiento degradara al nivel de RULE estándar, confirmando que apuntar a casos difíciles es crucial.
  • Variar el umbral $\tau$ mostró que establecerlo demasiado alto (incluyendo casos más fáciles) reduce la eficiencia, mientras que la configuración óptima (0.4 para MUSE) equilibra la selectividad y el tamaño del buffer.

Significado y Reivindicaciones

El artículo afirma que ReRULE ofrece una mejora rentable para el desaprendizaje basado en RL. Al reciclar rollouts "fallidos" de casos difíciles, mejora la Calidad de Retención del modelo desaprendido —una métrica crítica que a menudo se ve comprometida en tareas de desaprendizaje— sin requerir un reentrenamiento completo ni un cómputo adicional significativo.

Los autores enfatizan que el método es condicional: sus beneficios son más pronunciados en escenarios con una disparidad significativa entre prompts fáciles y difíciles (por ejemplo, MUSE-Books). En entornos más simples (por ejemplo, TOFU), las ganancias son marginales, lo que respalda la hipótesis de que el mecanismo de replay apunta específicamente a la ineficiencia de los casos difíciles suboptimizados. El trabajo posiciona a ReRULE como un paso práctico hacia un desaprendizaje de máquina más eficiente y robusto, particularmente para preservar la utilidad general del modelo mientras se elimina conocimiento peligroso específico.