DualOptim+: Bridging Shared and Decoupled Optimizer States for Better Machine Unlearning in Large Language Models

Resumen Técnico de DualOptim+

Enunciado del Problema
El Olvido Mecánico (MU) en Modelos de Lenguaje Grandes (LLM) busca eliminar la influencia de datos de entrenamiento específicos (el conjunto de olvido) preservando al mismo tiempo la utilidad general del modelo en los datos restantes (el conjunto de retención). Los enfoques existentes enfrentan un desafío significativo al equilibrar estos objetivos conflictivos.

• Actualización Conjunta: La mayoría de los métodos minimizan la suma de las pérdidas de olvido y retención simultáneamente. Esto a menudo conduce a una degradación significativa en la utilidad del modelo debido a un olvido excesivo.
• Actualización Alternada: Alternar entre los objetivos de olvido y retención mejora el rendimiento, pero sufre de inestabilidad y sensibilidad a los hiperparámetros.
• Optimizadores Desacoplados (DualOptim): Un enfoque previo, DualOptim, desacopla los estados del optimizador (por ejemplo, términos de momento) para cada objetivo para evitar el entrelazamiento de gradientes. Aunque es efectivo en visión por computadora, DualOptim produce solo mejoras marginales en LLM. Los autores observan que en el olvido de LLM, los gradientes para el olvido y la retención no siempre están en conflicto directo; a menudo comparten representaciones comunes, especialmente en las fases posteriores del entrenamiento. Los estados completamente desacoplados no aprovechan estas características compartidas, mientras que los estados completamente compartidos no logran manejar los conflictos.

Metodología: DualOptim+
El artículo propone DualOptim+, un marco de optimización plug-and-play diseñado para cerrar la brecha entre los estados del optimador completamente compartidos y completamente desacoplados. La innovación central es la descomposición del estado del optimizador en dos componentes:

1. Estado Base Compartido ($B$): Este estado captura las representaciones comunes compartidas por ambos objetivos de olvido y retención. Se actualiza conjuntamente utilizando gradientes de ambos objetivos ($\nabla_\theta L_f$ y $\nabla_\theta L_r$).
2. Estados Delta Desacoplados ($\Delta_f, \Delta_r$): Estos estados preservan los residuos específicos del objetivo. Se actualizan basándose en la diferencia entre el gradiente del objetivo específico y el estado base corregido por sesgo.
   • $\Delta_f \leftarrow \beta \Delta_f + (1-\beta)(\nabla_\theta L_f - \hat{B})$
   • $\Delta_r \leftarrow \beta \Delta_r + (1-\beta)(\nabla_\theta L_r - \hat{B})$

Actualización de Parámetros:
El término de momento final utilizado para las actualizaciones de parámetros se reconstruye combinando el estado base corregido por sesgo con el estado delta corregido por sesgo respectivo ($\hat{B} + \hat{\Delta}_f$ o $\hat{B} + \hat{\Delta}_r$).

Insight Teórico Clave:
Los análisis teóricos y numéricos demuestran que DualOptim+ funciona como un intermedio adaptativo.

• Cuando los gradientes están positivamente correlacionados (las características compartidas dominan), los estados delta convergen a cero y el método se comporta como un esquema Alternado (utilizando la base compartida).
• Cuando los gradientes están negativamente correlacionados (objetivos conflictivos), el estado base converge a cero y el método se comporta como un esquema DualOptim (utilizando estados desacoplados).
  Esta adaptabilidad permite que el optimizador se ajuste dinámicamente en función del conflicto direccional entre los gradientes de olvido y retención.

DualOptim+ 8bit
Para abordar la sobrecarga de memoria asociada con el almacenamiento de estados adicionales del optimizador (3 veces la memoria de un optimizador Adam estándar), los autores introducen DualOptim+ 8bit. Esta variante cuantiza los estados del optimizador a una precisión de 8 bits, reduciendo significativamente el consumo de memoria (a aproximadamente 3/4 de un Adam estándar) sin comprometer el rendimiento.

Contribuciones Clave

1. Marco Novel: La propuesta de DualOptim+, que introduce un estado base compartido y estados delta desacoplados para conectar eficazmente los estados del optimizador compartidos y desacoplados.
2. Diseño Plug-and-Play: Un marco compatible con cualquier optimizador que almacene estados (por ejemplo, AdamW, Muon) y aplicable a la optimización multiobjetivo.
3. Validación Empírica: Experimentos extensos en conjuntos de datos ficticios (TOFU), escenarios del mundo real (olvido en Llama-3), alineación de seguridad y tareas de aprendizaje multi-tarea.
4. Eficiencia: La introducción de la variante cuantizada de 8 bits que mitiga la sobrecarga de memoria.

Resultados Experimentales
El artículo evalúa DualOptim+ frente a las líneas base Joint, Alternate y DualOptim en varios LLM (Phi-1.5, Llama-2, Llama-3) y tareas:

• Escenarios Ficticios (TOFU): DualOptim+ logra consistentemente el mejor Rendimiento General (OVR), definido como un equilibrio entre Eficacia de Olvido (FE) y Utilidad del Modelo (MU). Supera a DualOptim, particularmente en el olvido dirigido, al aprovechar mejor las características compartidas.
• Escenarios del Mundo Real: En el olvido de individuos del mundo real de Llama-3, DualOptim+ mantiene una utilidad superior en tareas posteriores (ARC-c, MMLU, etc.) en comparación con las líneas base, mientras que DualOptim a veces tiene un rendimiento inferior al enfoque Alternado más simple.
• Alineación de Seguridad: En una tarea para olvidar conocimientos inseguros mientras se mantienen las capacidades de seguimiento de instrucciones, DualOptim+ logra el mejor equilibrio seguridad-utilidad y la tasa más baja de rechazo excesivo en la prueba XSTest.
• Aprendizaje Multi-tarea: A diferencia del olvido, donde los conflictos de gradientes son altos, en el aprendizaje multi-tarea donde los conflictos son menos pronunciados, DualOptim+ continúa ofreciendo aumentos sustanciales de rendimiento, mientras que DualOptim muestra ganancias negativas.
• Eficiencia de Memoria: Las variantes de 8 bits logran una paridad casi total en el uso de memoria con las líneas base estándar mientras mantienen los beneficios de rendimiento del método de precisión completa.

Significado y Afirmaciones
Los autores afirman que DualOptim+ proporciona un equilibrio más estable y efectivo entre la eliminación de conocimientos y la preservación de la utilidad que los métodos existentes. Posicionan el método no solo como una herramienta de olvido, sino como un marco de optimización generalizable aplicable a escenarios más amplios que involucran múltiples objetivos de aprendizaje conflictivos, como la alineación de LLM y el aprendizaje multiobjetivo. El artículo enfatiza que la naturaleza adaptativa de conectar estados compartidos y desacoplados es crucial para manejar las relaciones dinámicas de gradientes encontradas en el olvido de LLM.