Closing the Distribution Gap in Adversarial Training for LLMs

Resumen Técnico: Cerrando la Brecha de Distribución en el Entrenamiento Adversarial para LLMs

1. El Problema: La Fragilidad de las Defensas Actuales

A pesar de los avances en el entrenamiento adversarial (AT) para Grandes Modelos de Lenguaje (LLMs), estos modelos siguen siendo vulnerables a ataques simples que caen fuera de la distribución de datos de entrenamiento (out-of-distribution), como:

• Reescribir prompts en tiempo pasado.
• Traducir ataques a idiomas de bajos recursos.
• Reformulaciones naturales que no requieren perturbaciones optimizadas específicas del modelo.

Diagnóstico del Artículo:
Los autores argumentan que esta fragilidad persistente no se debe a una falla en la optimización adversarial (el "bucle interno"), sino a una limitación fundamental en la aproximación de la distribución de datos.

• El AT estándar minimiza el riesgo adversarial empírico sobre un conjunto de datos fijo y pequeño ($D$).
• Esto crea una brecha de distribución entre el riesgo empírico y el riesgo poblacional real ($R_{pop}$).
• Los métodos actuales cubren bien las perturbaciones específicas del modelo dentro de una bola local, pero fallan al cubrir la variedad natural de prompts que podrían desencadenar respuestas dañinas en la distribución real de datos ($q(x, y)$).

2. Metodología: Entrenamiento Adversarial Distribucional (DAT)

Para cerrar esta brecha, los autores proponen DAT (Distributional Adversarial Training), un marco que combina la optimización adversarial continua con la generación de datos basada en la distribución real.

Componentes Clave:

• Surrogados Generativos (Diffusion LLMs):

  • En lugar de usar un conjunto de datos estático, DAT utiliza Modelos de Difusión para LLMs como sustitutos generativos.
  • A diferencia de los modelos autoregresivos estándar que aprenden $q(y|x)$, los modelos de difusión capturan la distribución conjunta $q(x, y)$.
  • Esto permite resolver el problema inverso: muestrear prompts diversos y de alta probabilidad ($x$) condicionados a una respuesta dañina específica ($y$), es decir, $x \sim p_{diff}(x|y)$.

• Muestreo Específico de Datos:

  • El enfoque prioriza prompts que tienen alta probabilidad de elicitar respuestas dañinas bajo la distribución natural de datos, en lugar de perturbaciones raras o sintéticas que solo funcionan en un modelo específico.
  • Se utiliza un proceso de "inpainting" (relleno) en el modelo de difusión para generar variaciones de prompts para cada comportamiento dañino objetivo.

• Objetivo de Optimización Unificado:
  La función de pérdida de DAT combina dos bucles:

  1. Bucle Interno (Ataque Continuo): Se aplica perturbación continua a las representaciones de los tokens ($\delta$) para maximizar la pérdida adversarial ($L_{adv}$), asegurando robustez local.
  2. Bucle Externo (Muestreo de Distribución): Se minimiza la pérdida esperada sobre las muestras generadas por el modelo de difusión, aproximando así el riesgo poblacional real.
  3. Regularización: Se incluye un término de divergencia KL ($L_{KL}$) sobre un conjunto de retención ($D_{ret}$) para evitar que el modelo colapse y mantener su utilidad general.

  $$L_{DAT} = \mathbb{E}_{y, x} \left[ \max_{\delta} L_{adv}(x+\delta, y) \right] + \lambda_{KL} L_{KL}$$

3. Contribuciones Principales

1. Formalización de la Brecha de Robustez:

   • Identifican y formalizan la discrepancia entre minimizar el riesgo empírico (datos finitos) y el riesgo poblacional (distribución real). Distinguen entre vulnerabilidades específicas del modelo y fallos de generalización específicos de los datos.

2. Propuesta de DAT:

   • Introducen un nuevo objetivo de entrenamiento que integra el muestreo generativo de modelos de difusión con la optimización adversarial continua.
   • Proporcionan una justificación teórica (Teorema 3.1) que demuestra que, bajo una suposición de fidelidad del sustituto, optimizar sobre la distribución del modelo de difusión reduce la brecha entre el riesgo empírico y el poblacional.

3. Mejora Empírica de la Robustez:

   • Demuestran que DAT supera significativamente a los métodos anteriores (como CAT, LAT, Circuit Breakers y MixAT) en escenarios de ataques de "peor caso", especialmente contra ataques que explotan la distribución natural del lenguaje.

4. Resultados Experimentales

Los experimentos se realizaron en modelos Llama3-8B y Qwen2.5-14B, evaluando robustez frente a ataques como GCG, PAIR, BoN, Inpainting (basado en difusión) y ataques directos.

• Robustez de Peor Caso (ASR - Tasa de Éxito del Ataque):

  • Los métodos basados en AT estándar (CAT, LAT) y Circuit Breakers fallan estrepitosamente ante ataques de Inpainting (basados en difusión), con tasas de éxito de ataque (ASR) que oscilan entre el 77% y el 94%.
  • DAT reduce drásticamente el ASR. En Llama3-8B, el ASR de "Mejor de Todos" (BoA, el peor escenario entre todos los ataques) baja del 77% (MixAT) al 36%. En Qwen2.5-14B, baja del 75-93% al 18%.
  • DAT es el único método que mantiene una defensa sólida tanto contra ataques de optimización (GCG) como contra ataques de distribución natural (Inpainting).

• Compromiso Robustez-Utilidad (Pareto):

  • DAT logra un compromiso óptimo de Pareto. Mantiene una alta utilidad (medida en XSTest, MMLU, ARC) mientras ofrece una robustez superior.
  • A diferencia de otros métodos que sacrifican utilidad por seguridad, DAT permite controlar el nivel de robustez mediante hiperparámetros sin degradar significativamente la capacidad del modelo para tareas benignas.

• Validación de la Teoría:

  • Ablación de Muestreo: Aumentar el número de muestras de difusión por comportamiento (de 1 a 16) reduce linealmente el ASR de los ataques de Inpainting, validando que una mejor aproximación de la distribución de datos mejora la robustez.
  • Importancia de la Especificidad de Datos: Entrenar con muestras de baja probabilidad (que no elicitan respuestas dañinas en modelos no defendidos) no mejora la robustez. Solo las muestras de alta probabilidad (específicas de la distribución dañina real) son efectivas, confirmando que la calidad de la aproximación de la distribución es crítica.

5. Significado e Impacto

• Cambio de Paradigma: El trabajo desplaza el foco del entrenamiento adversarial desde la mera optimización de perturbaciones locales hacia la aproximación de la distribución de datos subyacente.
• Seguridad en la Vida Real: Al abordar las vulnerabilidades que surgen de variaciones naturales del lenguaje (traducciones, cambios de tiempo verbal), DAT ofrece una defensa más realista y robusta contra la manipulación de LLMs en entornos de producción.
• Viabilidad Teórica y Práctica: Proporciona un marco teórico sólido que conecta la fidelidad de los modelos generativos con la seguridad del modelo, demostrando que los modelos de difusión pueden actuar como herramientas poderosas para la ciberseguridad de la IA, no solo para la generación de contenido.

En conclusión, DAT demuestra que para lograr una robustez real en LLMs, no basta con entrenar contra ejemplos adversarios estáticos; es necesario entrenar contra la distribución completa de posibles ataques naturales, utilizando modelos generativos avanzados para descubrir y cubrir esas vulnerabilidades ocultas.