Optimus: A Robust Defense Framework for Mitigating Toxicity while Fine-Tuning Conversational AI

1. El Problema

La personalización de Modelos de Lenguaje Grandes (LLMs) mediante fine-tuning (ajuste fino) en conjuntos de datos no confiables plantea un riesgo de seguridad crítico: la inyección de comportamientos tóxicos.

• Amenaza: Un adversario puede envenenar el conjunto de datos de entrenamiento (inyectando conversaciones tóxicas) para que el chatbot aprenda a generar respuestas dañinas, ofensivas o sesgadas.
• Desafíos Técnicos:
  1. Desconocimiento del atacante: El defensor no conoce la distribución exacta del lenguaje tóxico en los datos, lo que dificulta crear clasificadores perfectos.
  2. Evolución del lenguaje: Los clasificadores de toxicidad actuales pueden volverse obsoletos o ser sesgados ante nuevas formas de toxicidad.
  3. Calidad vs. Seguridad: Las defensas tradicionales (filtrado agresivo) suelen degradar la utilidad del chatbot (generando falsos positivos o respuestas genéricas) y no siempre eliminan la toxicidad residual si el clasificador falla.

2. Metodología: El Marco Optimus

Optimus es un marco de defensa de extremo a extremo diseñado para integrarse en tuberías de fine-tuning existentes. Su objetivo es mitigar la toxicidad aprendida de datos no confiables sin sacrificar la calidad de la conversación. Se compone de dos etapas principales:

A. Clasificación de Toxicidad (Etapa 1)

En lugar de depender exclusivamente de clasificadores supervisados tradicionales, Optimus propone un esquema sin entrenamiento (training-free) que aprovecha las propiedades de seguridad de los LLMs comerciales ya alineados (como LLaMA-2-Chat).

• Enfoque "Refusal" (Rechazo): Se utiliza un LLM alineado con seguridad para evaluar si es seguro generar la siguiente respuesta en un par contexto-respuesta. Si el modelo predice que es inseguro (basado en la probabilidad de generar un token de "no"), la muestra se marca como tóxica.
• Ventaja: No requiere conjuntos de datos etiquetados y aprovecha la alineación de seguridad inherente de modelos como LLaMA-2.

B. Ajuste Fino y Alineación (Etapa 2)

Optimus reconoce que los clasificadores pueden ser imperfectos o sesgados. Por lo tanto, no se limita a filtrar datos, sino que implementa una estrategia dual:

1. Generación de "Datos de Curación" (Healing Data):

   • Reemplaza las respuestas tóxicas identificadas en el conjunto de datos con respuestas sintéticas no tóxicas.
   • Curación No Contextual (NH): Usa respuestas genéricas (ej. "Lo siento, no estoy seguro...").
   • Curación Contextual (CH): Utiliza un LLM alineado para generar respuestas empáticas y prosociales que sean relevantes para el contexto de la conversación, manteniendo la utilidad del chatbot.

2. Alineación del Modelo mediante DPO (Direct Preference Optimization):

   • Después del fine-tuning inicial con los datos curados, se aplica DPO para alinear el modelo hacia las respuestas deseables.
   • Se crea un conjunto de datos de preferencia donde, para cada contexto tóxico, se compara la respuesta original (tóxica) con la respuesta curada (segura).
   • Ventaja clave: DPO es más eficiente que el RLHF (Aprendizaje por Refuerzo con Retroalimentación Humana) y, crucialmente, es robusto ante clasificadores imperfectos. Incluso si el clasificador detecta solo una fracción de las muestras tóxicas (o tiene sesgos), DPO puede guiar al modelo hacia la distribución de preferencia segura y generalizar a casos no vistos.

3. Contribuciones Clave

1. Clasificador de Toxicidad sin Entrenamiento: Demostraron que los LLMs alineados con seguridad (como LLaMA-2-Chat) pueden adaptarse para detectar toxicidad con un rendimiento superior (hasta un 28.4% más que APIs industriales como OpenAI o Perspective) sin necesidad de datos etiquetados.
2. Resiliencia ante Clasificadores Imperfectos: Optimus es el primer esquema de extremo a extremo capaz de mitigar la toxicidad eficazmente incluso cuando el clasificador de toxicidad tiene un alto nivel de sesgo o degradación (hasta un 85% de reducción en la Recall para subcategorías específicas).
3. Mantenimiento de la Utilidad: A diferencia de los filtros agresivos, Optimus preserva la calidad de la conversación y puede reforzar comportamientos deseables (empatía, prosocialidad) mediante la curación contextual.
4. Robustez ante Ataques Adaptativos: El marco demuestra resistencia contra ataques de jailbreak y ataques adversarios diseñados para evadir tanto la clasificación de toxicidad como la generación de datos de curación.

4. Resultados Experimentales

Los autores evaluaron Optimus en tres modelos base (BART, BlenderBot, LLaMA-2) con datos envenenados en categorías "ofensivas" y "especializadas" (acuerdo tóxico, opiniones sesgadas, ignorancia de riesgos).

• Eficacia en la Mitigación:
  • Optimus redujo la Tasa de Toxicidad de Respuestas (RTR) a niveles cercanos o inferiores a los de un escenario sin ataque, incluso utilizando clasificadores sesgados (como las APIs de OpenAI o Perspective).
  • En el modelo LLaMA-2 (el más desafiante), las defensas basadas solo en filtrado o curación fallaron, pero el pipeline completo de Optimus (con DPO) logró reducir la toxicidad drásticamente.
• Comparación con el Estado del Arte (StarDSS):
  • Optimus superó a StarDSS (el estado del arte actual). StarDSS falló en escenarios con clasificadores imperfectos, mostrando una alta inestabilidad en el entrenamiento y una degradación severa de la utilidad (PPL ~229), mientras que Optimus mantuvo la estabilidad y la calidad.
• Resistencia a Ataques:
  • Ataques Adversarios: Optimus mitigó exitosamente ataques de PromptAttack que intentaban engañar al clasificador.
  • Ataques de Jailbreak: Incluso cuando los atacantes lograron engañar al clasificador de "rechazo" (reduciendo la Recall en un 52%), Optimus mantuvo la RTR baja (7.9%) gracias a la alineación con DPO.
  • Aprendizaje Basado en Diálogo (DBL): En escenarios de reentrenamiento post-despliegue, Optimus redujo la toxicidad a casi 0%, superando a los filtros multinivel existentes.

5. Significancia e Impacto

El trabajo de Optimus es fundamental para la seguridad de la IA conversacional por varias razones:

• Pragmatismo: Reconoce que los clasificadores de toxicidad perfectos no existen y que los datos de entrenamiento a menudo son "sucios". En lugar de depender de una detección perfecta, Optimus utiliza un mecanismo de corrección (DPO + datos curados) que funciona a pesar de los errores de detección.
• Escalabilidad: Al utilizar LLMs existentes para la detección y DPO para la alineación, el marco es aplicable a diversas arquitecturas de modelos y esquemas de personalización (como LoRA) sin requerir infraestructura masiva de anotación humana.
• Calidad de Servicio: Demuestra que es posible hacer que los chatbots sean seguros sin convertirlos en bots rígidos o inutiles, manteniendo la coherencia y la empatía en las respuestas.

En resumen, Optimus establece un nuevo estándar para la defensa en el fine-tuning de LLMs, priorizando la robustez ante la incertidumbre de los datos y la evolución de las amenazas, asegurando que los chatbots personalizados sean tanto útiles como seguros.