Resumen Técnico: Poda de Subpalabras Tóxicas para la Generación de Respuestas en Diálogos en Modelos de Lenguaje Grandes

Declaración del Problema

El avance rápido de los Modelos de Lenguaje Grandes (LLM) ha introducido riesgos significativos en cuanto a la generación de contenido dañino o tóxico. Si bien la investigación de seguridad existente se centra principalmente en defender modelos "seguros" contra ataques de jailbreak o prompts tóxicos, persiste una brecha crítica para los modelos que ya son tóxicos. Este escenario surge cuando desarrolladores individuales crean involuntariamente modelos tóxicos o cuando atacantes obtienen acceso a los pesos del modelo para ajustarlos finamente con contenido dañino.

Las estrategias actuales de mitigación enfrentan varias limitaciones:

• Alineación de Seguridad: El ajuste fino de modelos preentrenados para la seguridad es computacionalmente costoso y a menudo inaccesible para desarrolladores individuales.
• Clasificadores Externos: Utilizar clasificadores de toxicidad externos para detectar y bloquear la salida tóxica es imperfecto. Los atacantes pueden eludirlos mediante prompts de inyección maliciosos, y el bloqueo continuo puede volver inútil al modelo al impedir cualquier salida.
• Desintoxicación Compleja: Los métodos existentes para desintoxicar LLMs tóxicos suelen ser demasiado complejos para un despliegue rápido.
• Vulnerabilidad a la Inyección de Prompts: Incluso los modelos bien alineados pueden ser activados para generar toxicidad mediante inyección de prompts.

El artículo postula que se necesita una solución ligera, en tiempo de inferencia, para podar las capacidades tóxicas directamente del proceso de decodificación del modelo sin requerir reentrenamiento costoso ni actualizaciones de pesos.

Metodología: Poda de Subpalabras Tóxicas (ToxPrune)

Los autores proponen ToxPrune, un algoritmo novedoso que poda subpalabras (subtokens) asociadas a palabras tóxicas del vocabulario de un LLM entrenado durante el proceso de decodificación.

Mecanismo Central

1. Lista Tóxica Personalizada: El método acepta una lista definida por el usuario de palabras tóxicas.
2. Tokenización de Subpalabras: Estas palabras se tokenizan en subpalabras utilizando el tokenizador del modelo (por ejemplo, BPE). Los autores señalan que las palabras tóxicas y no tóxicas pueden compartir subpalabras; ToxPrune maneja esto podando las subpalabras específicas derivadas de la lista tóxica.
3. Modificación de la Distribución: Durante la inferencia, se modifica la distribución de probabilidad del siguiente token. En lugar de muestrear del vocabulario completo $V$, el modelo muestrea de un conjunto restringido $V \setminus V_{toxic}$, donde $V_{toxic}$ representa las subpalabras tóxicas podadas.
   • Formalmente, para el muestreo top-$k$, la distribución cambia de $P(\bar{x}_i | \dots) \in V^{(k)}$ a $P(\bar{x}_i | \dots) \in V^{(k)} \setminus V^{(k)}_{toxic}$.
   • Para la Búsqueda en Haz (Beam Search), la probabilidad de secuencias que contienen palabras podadas se establece en 0, eliminándolas efectivamente de la lista de candidatos.
4. Expansión y Truncamiento:
   • Lista Negra Parafraseada: Para aumentar la cobertura, la lista inicial de palabras tóxicas puede expandirse utilizando un generador de paráfrasis para incluir variaciones de términos tóxicos.
   • Lista Blanca Truncada: Se puede utilizar una lista blanca predefinida para preservar palabras clave específicas eliminando las subpalabras tóxicas que se superponen con palabras seguras, asegurando que los términos no tóxicos críticos sigan siendo utilizables.

Implementación

ToxPrune opera completamente en la etapa de inferencia. No requiere actualizar pesos del modelo, reentrenar ni utilizar clasificadores externos. Se integra con estrategias de muestreo estándar (Top-$k$, Top-$p$, Búsqueda en Haz) filtrando simplemente el vocabulario candidato antes de la selección.

Contribuciones Clave

1. Algoritmo Novel: La introducción de ToxPrune, un método que elimina subpalabras tóxicas para afectar directamente el proceso de decodificación de la generación de diálogos.
2. Validación Dual: El método se valida tanto en un LLM tóxico conocido (NSFW-3B) como en un LLM oficialmente alineado y no tóxico (Llama-3.1-6B Base).
3. Evidencia Empírica: Estudios de caso extensos y evaluaciones automáticas demuestran que ToxPrune mitiga efectivamente la generación tóxica mientras mantiene o mejora la calidad del diálogo.

Resultados Experimentales

Conjuntos de Datos y Modelos

• Conjuntos de Datos: Los experimentos se realizaron en PERSONACHAT-ORI y PERSONACHAT-REV (una versión revisada manualmente, más difícil), que son conjuntos de datos de diálogos de múltiples turnos condicionados a personas.
• Modelos:
  • NSFW-3B: Un LLM tóxico de código abierto entrenado para generar palabras obscenas.
  • Llama-3.1-6B: La versión base oficial de Llama-3.1 con alineación de seguridad.
• Líneas Base: Las comparaciones incluyeron un enfoque basado en paráfrasis y un Clasificador de Toxicidad (TC) que utiliza GPT-4 para la regeneración.

Rendimiento en LLMs Tóxicos (NSFW-3B)

• Reducción de Toxicidad: ToxPrune redujo drásticamente las puntuaciones de toxicidad (evaluadas por GPT-4) de 0.89 (línea base) a 0.13 (con poda del 100%).
• Mejora de la Calidad: Contrario a la expectativa de que la poda degradaría el rendimiento, la mayoría de las métricas mejoraron:
  • BLEU: Mientras que BLEU-1 experimentó una ligera caída, BLEU-2 a BLEU-4, y el BLEU ponderado aumentaron.
  • ROUGE: Todas las métricas ROUGE (R-1, R-2, R-L) mejoraron.
  • Diversidad: Las puntuaciones Distinct-1 y Distinct-2 aumentaron significativamente (por ejemplo, Distinct-2 de 0.712 a 0.800).
• Comparación: Las líneas base que utilizaban paráfrasis o clasificadores de toxicidad no lograron mitigar la toxicidad de manera efectiva porque el modelo subyacente NSFW-3B regeneraba consistentemente contenido tóxico.

Rendimiento en LLMs No Tóxicos (Llama-3.1-6B)

• Seguridad: La toxicidad se mantuvo en 0.00.
• Diversidad: ToxPrune sorprendentemente mejoró las métricas de diversidad (Distinct-1 y Distinct-2) en el modelo limpio, lo que sugiere que enmascarar subpalabras frecuentes compartidas entre contextos tóxicos y no tóxicos aplanan la distribución de palabras, lo que lleva a salidas más variadas.
• Calidad: Las métricas BLEU y ROUGE mostraron ligeras variaciones, con algunas mejoras en ROUGE, lo que indica que la poda no dañó las capacidades generativas centrales del modelo.

Evaluación Humana

• LLM Tóxico: Los anotadores humanos prefirieron significativamente ToxPrune sobre la línea base NSFW-3B en cuanto a Adecuación, Informatividad, Interesante y Humanidad. La toxicidad de la línea base fue eliminada en gran medida.
• LLM No Tóxico: ToxPrune no degradó el rendimiento en Llama-3.1-6B y mostró una mejora estadísticamente significativa en Informatividad (diversidad).

Significado y Afirmaciones

El artículo afirma que ToxPrune ofrece un enfoque simple, efectivo y novedoso para la seguridad de la IA con ventajas únicas:

1. Remediación de Modelos Tóxicos: Puede rescatar modelos que han sido comprometidos o entrenados intencionalmente para ser tóxicos, restaurando su capacidad para generar diálogos significativos y no tóxicos sin reentrenamiento.
2. Mejora de Modelos Seguros: Puede mejorar la diversidad del diálogo en modelos ya seguros, desafiando la noción de que podar el vocabulario es inherentemente dañino para las tareas de generación.
3. Robustez Frente a Ataques: Al eliminar los tokens tóxicos de la lista de muestreo por completo, ToxPrune previene la generación tóxica incluso bajo ataques de inyección de prompts, ya que el modelo físicamente no puede emitir los tokens podados.
4. Eficiencia: A diferencia de la alineación de seguridad o la edición de modelos, ToxPrune no requiere actualizaciones del modelo, es computacionalmente barato y permite la personalización dinámica de la lista de palabras tóxicas.

Los autores concluyen que, si bien el trabajo anterior sugería que podar tokens BPE perjudica tareas como la traducción automática, este estudio demuestra su utilidad sorprendente en el dominio de la seguridad de la IA y la generación de diálogos.