Learning What to Forget: Improving LLM Unlearning via Learned Token-Level Importance

Resumen Técnico: Aprender qué olvidar: Mejora del desaprendizaje de LLM mediante la importancia aprendida a nivel de token

Planteamiento del problema
El desaprendizaje de máquinas (machine unlearning) tiene como objetivo eliminar conocimiento específico (por ejemplo, datos personales, texto protegido por derechos de autor, contenido dañino) de los Grandes Modelos de Lenguaje (LLM) entrenados sin necesidad de reentrenarlos desde cero. Un desafío crítico en los LLM autorregresivos es que no todos los tokens dentro de una "muestra para olvidar" (los datos a eliminar) son igualmente relevantes para la información objetivo. Muchos tokens son estructurales (por ejemplo, palabras funcionales, puntuación, patrones sintácticos) y son necesarios para las capacidades lingüísticas generales. Los enfoques actuales de desaprendizaje ponderado por tokens intentan aplicar selectivamente la pérdida de olvido a tokens específicos, pero dependen de anotaciones externas, modelos auxiliares, heurísticas basadas en probabilidades o reglas específicas de la tarea para estimar la relevancia de los tokens. Estos métodos suelen fallar al distinguir eficazmente entre el contenido específico del olvido y los tokens estructurales, lo que provoca un olvido insuficiente o una degradación innecesaria de la utilidad general.

Metodología: Desaprendizaje de Pesos de Tokens Alterno (ATWU)
Los autores proponen que la especificidad del olvido a nivel de token debe caracterizarse por la interacción entre los propios objetivos de retención y olvido, en lugar de por señales externas. Formalizan esto como un problema de optimización conjunta sobre los parámetros del modelo ($\theta$) y los pesos de los tokens ($z$).

1. Formulación Teórica: El artículo define una métrica de "conflicto de retención", $\kappa(A)$, que mide el costo de suprimir un conjunto de tokens $A$ mientras se mantiene la optimalidad en el conjunto de retención. Bajo un supuesto de "separación de conflicto de retención", los autores demuestran que el conjunto oráculo de los tokens específicos del olvido puede recuperarse como el minimizador de un objetivo conjunto que equilibra la pérdida de retención y una pérdida de olvido ponderada. Este resultado teórico sugiere que los pesos de los tokens pueden aprenderse directamente del objetivo de desaprendizaje sin supervisión externa.
2. Algoritmo (ATWU): Para hacer esto tratable, los autores introducen el Desaprendizaje de Pesos de Tokens Alterno (ATWU).
   • Calificador (Scorer): En lugar de aprender pesos independientes para cada token, ATWU utiliza un calificador lineal ligero, $g_w$, parametrizado por un vector $w$ en el espacio de estados ocultos del modelo. La puntuación de un token es $\hat{z}_{x,t} = \sigma(\langle w, h_\theta(x_t) \rangle)$, donde $h_\theta$ es la representación oculta.
   • Optimización Conjunta: El marco optimiza un objetivo relajado y continuo que incluye la pérdida de retención, una pérdida de entropía negativa saturada (SATGA+) para el olvido, un regularizador de entropía para empujar las puntuaciones hacia valores binarios y una restricción de presupuesto para controlar el número de tokens seleccionados.
   • Actualizaciones Alternas: Los parámetros del modelo $\theta$ y los parámetros del calificador $w$ se actualizan de forma alterna. Cuando se actualiza $\theta$, las puntuaciones se desvinculan (detached); cuando se actualiza $w$, $\theta$ se congela. Esto estabiliza el entrenamiento al desacoplar el bucle de retroalimentación entre la cambiante geometría del modelo y la selección de tokens.

Contribuciones Clave

• Caracterización del Conflicto de Retención: El artículo proporciona una caracterización teórica de la especificidad del olvido a nivel de token basada en la tensión entre los objetivos de retención y olvido, demostrando que una optimización conjunta puede recuperar el soporte de olvido oráculo bajo condiciones de separación natural.
• Aprendizaje No Supervisado de Pesos de Tokens: ATWU es un marco ligero que aprende la especificidad del olvido de los tokens de forma conjunta con los parámetros del modelo utilizando únicamente las representaciones internas del modelo. No requiere supervisión externa a nivel de token, modelos auxiliares ni anotaciones costosas.
• Integración de Pérdida Adaptativa: El método introduce SATGA+, una variante del error de entropía negativa saturada modulada por puntuación donde la puntuación de token aprendida se inyecta en el exponente de saturación. Esto permite actualizaciones más suaves y conscientes de la incertidumbre durante las fases tempranas de la optimación.

Resultados
Los autores evalúan ATWU en dos entornos de prueba: TOFU (biografías ficticias de autores sintéticos) y RWKU (conocimiento de figuras públicas del mundo real).

• Desempeño: ATWU alcanza el estado del arte en Calidad de Desaprendizaje (UQ), una métrica que equilibra la calidad del olvido (FQ) y la degradación de la retención (RD). En TOFU, logra una UQ de 91.7 (superando al segundo mejor por 5.8 pp) y en RWKU, 58.7 (superando al segundo mejor por 6.8 pp).
• Preservación de la Utilidad: A diferencia de los métodos de desaprendizaje agresivos que colapsan la calidad de generación, ATWU preserva la utilidad general (medida por MMLU, repetitividad y tasas de victoria por pares) cerca del checkpoint original.
• Calidad de la Puntuación de Tokens: Las puntuaciones de tokens aprendidas se alinean significativamente mejor con los segmentos de verdad de base de los fragmentos específicos del olvido (medido por AUROC) en comparación con las heurísticas basadas en probabilidad, los enfoques de modelos auxiliares y otros entornos de base de ponderación de tokens.
• Versatilidad: El marco es agnóstico a la pérdida; cuando se combina con otras pérdidas de olvido (DPO, NPO, SIMNPO), ATWU mejora consistentemente el equilibrio olvido-retención sobre sus contrapartes no ponderadas.

Significado y Reivindicaciones
El artículo afirma que el "conflicto de retención" proporciona un criterio intrínseco y efectivo para identificar qué debe olvidar los modelos de lenguaje. Al desplazar el desaprendizaje ponderado por tokens de un flujo de trabajo de heurísticas externas a un problema de optimización conjunta de extremo a extremo, ATWU demuestra que:

1. El propio objetivo de desaprendizaje contiene suficiente señal latente para recuperar los tokens específicos del olvido sin supervisión externa.
2. El cuello de botella en el desaprendizaje selectivo no es simplemente encontrar los tokens correctos, sino el diseño conjunto y principista del mecanismo de puntuación y el objetivo de olvido.
3. Es posible lograr un desaprendizaje escalable, respetuoso con la privacidad y altamente dirigido sin depender de anotaciones externas, que a menudo introducen riesgos de privacidad y problemas de escalabilidad.

Los autores señalan limitaciones, incluyendo la dependencia de un hiperparámetro para el presupuesto de selección de tokens y el supuesto de una separación estricta entre los tokens específicos del olvido y los estructurales, que puede cumplirse solo aproximadamente en el lenguaje natural. Además, las evaluaciones están limitadas actualmente a entornos específicos y ejecuciones únicas debido a los costos computacionales.