Reinforcement Learning for Neural Model Editing

Resumen Técnico: Aprendizaje por Refuerzo para la Edición de Modelos Neuronales

Planteamiento del Problema
Modificar redes neuronales preentrenadas para lograr objetivos específicos —como la eliminación de conceptos, la mitigación de sesgos o la incorporación de nuevo conocimiento— requiere tradicionalmente algoritmos especializados diseñados manualmente. Estos métodos suelen ser lentos de diseñar y carecen de generalizabilidad entre diferentes tareas. Este artículo investiga si el proceso de edición de modelos neuronales puede automatizarse formulándolo como un problema de Aprendizaje por Refuerzo (RL). La pregunta central de investigación es si un agente de RL puede aprender una política para modificar los pesos de una red neuronal utilizando únicamente retroalimentación de recompensa, sin acceso a gradientes o conjuntos de datos construidos manualmente para la tarea de edición específica.

Metodología
Los autores proponen un marco exploratorio donde un agente interactúa con una red neuronal actualizando iterativamente sus pesos basándose en una señal de recompensa. El marco consta de dos entornos distintos:

1. MaskWorld: El agente modifica los pesos escalándolos multiplicativamente. La acción $a_t$ escala los pesos de una capa específica $l$ tal que $W^{(l)}_t = W^{(l)} \odot a_t$, donde $a_t$ es un vector de constantes entre 0.0 y 1.0.
2. ShiftWorld: El agente modifica los pesos aplicando actualizaciones aditivas. La acción $a_t$ desplaza los pesos de tal manera que $W^{(l)}_t = W^{(l)} + a_t$, donde $a_t$ contiene constantes entre -0.05 y 0.05.

Para abordar la intratabilidad computacional de predecir vectores de acción de alta dimensión (por ejemplo, $m \times n$ para una matriz de pesos), los autores emplean un mecanismo inspirado en la Adaptación de Bajo Rango (LoRA) para la tarea de desaprendizaje de máquinas (machine unlearning). En lugar de predecir la matriz completa de actualización de pesos, el agente predice dos matrices de bajo rango, $A$ y $B$, reduciendo la dimensionalidad del espacio de acción y estabilizando la optimización de la política.

Función de Recompensa
La función de recompensa $R$ es una combinación ponderada de dos componentes:
$$R = \alpha U(w) + \beta O(w)$$

• Término de Utilidad ($U(w)$): Fomenta la preservación del rendimiento general del modelo (por ejemplo, la precisión en un conjunto de "retención" o un conjunto de datos general).
• Término de Objetivo de Edición ($O(w)$): Cuantifica el éxito del objetivo de edición específico (por ejemplo, minimizar la precisión en un conjunto de "olvido" o reducir las métricas de sesgo).
• Coeficientes ($\alpha, \beta$): Controlan el equilibrio entre mantener la utilidad y alcanzar el objetivo de edición. Notablemente, $\beta$ se establece con un valor negativo en las tareas de desaprendizaje para penalizar el rendimiento en el conjunto de olvido.

El agente se entrena utilizando Optimización de Política Próxima (PPO) con una red de política MLP. Crucialmente, el agente no recibe gradientes; aprende únicamente mediante ensayo y error basado en el rendimiento del modelo modificado en lotes (batches) muestreados.

Contribuciones Clave y Resultados Experimentales
El marco fue evaluado en dos tareas distintas: desaprendizaje de máquinas en clasificación de imágenes y mitigación de sesgos en clasificación de texto.

• Desaprendizaje de Máquinas (MNIST): El objetivo era "desaprender" la clase de dígito 7 mientras se mantenía el rendimiento en las clases 0–6 y 8–9.

  • Resultados: Las políticas aprendidas redujeron la precisión en el conjunto de olvido (clase 7) a casi 0% (0.0000 para MaskWorld, 0.0006 para ShiftWorld).
  • Preservación de la Utilidad: Contrario a las expectativas de degradación del rendimiento, los modelos modificados mejoraron ligeramente la precisión en el conjunto de retención (del 92.84% a ~93.5%) en comparación con el modelo original.
  • Observación: Los agentes lograron estos resultados accediendo solo al 15% de los datos de entrenamiento, demostrando una fuerte generalización al conjunto de prueba.

• Mitigación de Sesgos (Jigsaw Toxic Comments): El objetivo era reducir las correlaciones espurias (por ejemplo, asociar la palabra "negro" con toxicidad) mientras se mantiene la utilidad de clasificación general.

  • Resultados: Los modelos modificados mejoraron la precisión en un conjunto de prueba sintético sin sesgos en aproximadamente un 6–7% (por ejemplo, ShiftWorld mejoró del 55.36% al 62.52%).
  • Preservación de la Utilidad: La precisión del conjunto de datos general se mantuvo estable, disminuyendo solo ~1% para MaskWorld, lo que indica que el agente equilibró con éxito la reducción del sesgo con la preservación de la utilidad.

Significancia y Reivindicaciones
El artículo no pretende superar a los algoritmos especializados de última generación en estas tareas específicas. En cambio, su principal significancia radica en demostrar un cambio de paradigma: la edición de modelos neuronales puede plantearse como un problema de RL aprendible en lugar de uno diseñado manualmente.

Los autores afirman que sus hallazgos sugieren que el RL proporciona un marco general y extensible para aprender políticas de edición a partir de la retroalimentación de recompensa. Este enfoque permite el descubrimiento de modificaciones de pesos que logran objetivos específicos (como el desaprendizaje o la mitigación de sesgos) sin requerir un diseño algorítmico específico para cada nuevo objetivo. El trabajo sirve como una prueba de concepto para motivar la investigación futura en enfoques aprendidos para la edición de modelos en áreas como la equidad, la robustez y la privacidad.

Limitaciones
Los autores reconocen que el marco actual enfrenta desafíos de escalabilidad inherentes al RL. Si bien el mecanismo inspirado en LoRA ayuda a manejar capas de alta dimensión, optimizar completamente todas las capas de un modelo grande sigue siendo inviable debido a la inestabilidad causada por el enorme espacio de acción. Se sugiere que el trabajo futuro explore entornos multi-agente y enfoques de teoría de juegos para abordar estas limitaciones.