Stable-GFlowNet: Toward Diverse and Robust LLM Red-Teaming via Contrastive Trajectory Balance

1. Enunciado del Problema

El red-teaming de Modelos de Lenguaje Grande (LLM) tiene como objetivo descubrir proactivamente vulnerabilidades (salidas tóxicas o dañinas) antes del despliegue. Si bien los métodos de Aprendizaje por Refuerzo (RL) pueden encontrar ataques de alta recompensa, sufren de colapso de modos, generando patrones de ataque repetitivos y estrechos. Por el contrario, las Redes de Flujo Generativo (GFN) ofrecen un marco para la coincidencia de distribuciones, permitiendo teóricamente el descubrimiento de trayectorias de alta recompensa diversas. Sin embargo, aplicar GFNs al red-teaming de LLM enfrenta dos desafíos críticos:

1. Inestabilidad en el Entrenamiento: Los objetivos estándar de GFN (como el Equilibrio de Trayectoria, TB) requieren estimar una función de partición global ($Z$). En el espacio discreto y de alta dimensión de los LLM, estimar $Z$ es numéricamente inestable, lo que lleva a una convergencia deficiente y al colapso de modos.
2. Recompensas Ruidosas: Los clasificadores de toxicidad proporcionan señales densas pero ruidosas. A menudo asignan recompensas no nulas a texto sin sentido o texto Fuera de Distribución (OOD), lo que hace que el modelo "hackee" la recompensa generando texto incoherente en lugar de ataques fluidos y diversos.

2. Metodología: Stable-GFlowNet (S-GFN)

Los autores proponen Stable-GFlowNet (S-GFN), un marco que elimina la necesidad de estimación explícita de la función de partición y filtra las recompensas ruidosas para garantizar un entrenamiento estable y diverso. La tubería consta de tres componentes principales:

A. Equilibrio de Trayectoria Contrastivo (CTB)

Para abordar la inestabilidad de estimar $Z$, S-GFN reemplaza el objetivo de coincidencia de flujo absoluto con un objetivo de comparación por pares.

• Mecanismo: En lugar de minimizar el error entre el flujo de la política y la recompensa relativo a $Z$, CTB minimiza la diferencia al cuadrado entre el log-ratio de las probabilidades de la política y el log-ratio de las recompensas para pares de trayectorias $(y_1, y_2)$:
  $$L_{CTB} = \left( \log \frac{\pi_\theta(y_1)}{\pi_\theta(y_2)} - \log \frac{R(y_1)}{R(y_2)} \right)^2$$
• Garantía Teórica: Los autores demuestran que la política óptima para CTB es idéntica a la del Equilibrio de Trayectoria estándar ($\pi_\theta(y) \propto R(y)$).
• Reducción de Varianza: Al utilizar una muestra como base para otra (similar a Reinforce Leave-One-Out), CTB reduce implícitamente la varianza del gradiente sin necesidad de una red de base externa o estimación de $Z$.

B. Poda de Gradientes Ruidosos (NGP)

Para mitigar el riesgo de amplificación del ruido de recompensa inherente a las comparaciones por pares, S-GFN introduce un filtro basado en saliencia.

• Mecanismo: Los gradientes solo se calculan para pares de trayectorias donde la diferencia absoluta en los log-recompensas supera un umbral $\sigma$:
  $$L_{NGP} = \mathbb{I}(|\log R(y_1) - \log R(y_2)| > \sigma) \cdot L_{CTB}$$
• Efecto: Esto evita que el optimizador actualice basándose en diferencias de recompensa insignificantes causadas por la estocasticidad del LLM víctima o del clasificador de toxicidad, podando eficazmente los gradientes ruidosos.

C. Estabilizador de Fluidez Min-K (MKS)

Para evitar que el modelo converja a mínimos locales de texto sin sentido (que a menudo reciben puntuaciones de toxicidad espurias), S-GFN emplea una restricción de fluidez.

• Mecanismo: En lugar de una penalización global de divergencia KL (que puede sobreajustarse a un modelo de referencia), MKS calcula la estadística Min-K: el log-probabilidad promedio de los $k$ tokens menos probables en una secuencia generada bajo un modelo de referencia.
• Penalización Dura: Si la puntuación Min-K cae por debajo de un umbral $T_{mks}$, la recompensa se recorta a cero (o a un valor negativo grande).
  $$R_{MKS}(y) = \mathbb{I}(M_k(y) \ge T_{mks}) \cdot R(y)$$
• Beneficio: Esto asegura que solo los prompts fluidos y lingüísticamente coherentes se consideren de alta recompensa, permitiendo que el modelo explore un espacio más amplio de ataques válidos sin generar sinsentidos.

3. Contribuciones Clave

1. Marco Stable-GFN: Un marco novedoso de red-teaming que elimina la necesidad de estimación de la función de partición ($Z$), resolviendo la causa principal de la inestabilidad en el entrenamiento de GFNs para LLMs.
2. Equilibrio de Trayectoria Contrastivo (CTB): Una nueva función objetivo que logra la misma distribución óptima que TB pero con una varianza significativamente menor y una mayor estabilidad en el entrenamiento.
3. Robustez al Ruido: La integración de NGP y MKS aborda específicamente los desafíos únicos del red-teaming (clasificadores ruidosos y recompensas de texto sin sentido), previniendo el colapso de modos y el hackeo de recompensas.
4. Universalidad: El artículo demuestra que CTB y NGP son efectivos no solo en red-teaming, sino también en otras tareas de coincidencia de distribuciones como la generación molecular y la navegación en hipergrid.

4. Resultados Experimentales

Los autores evaluaron S-GFN frente a líneas base que incluyen PPO, Jailbreak R1, Rainbow Teaming y GFN estándar (TB) en varios modelos víctimas (Qwen, Llama, Gemma) y configuraciones de defensa.

• Diversidad (Ataques Únicos): S-GFN generó 134 prompts de ataque únicos, una mejora masiva sobre la línea base de GFN estándar (17) y PPO (3). Produjo aproximadamente 7 veces más ataques únicos que la línea base de GFN.
• Tasa de Éxito del Ataque (ASR): S-GFN mantuvo un ASR alto de ~92.5%, comparable a las mejores líneas base (GFN: 93.7%, PPO: 91.7%), demostrando que la diversidad no se logró a costa de la efectividad.
• Transferencia de Ataques Cruzados: Cuando los modelos víctimas fueron ajustados finamente para la seguridad utilizando ataques de S-GFN, se volvieron significativamente más robustos contra ataques de otros métodos (el ASR cayó a cerca del 0% para ataques de GFN). Por el contrario, los ataques de S-GFN permanecieron efectivos contra modelos ajustados finamente con otros métodos, demostrando una generalización superior.
• Ataques de Transferencia: S-GFN transfirió con éxito ataques a modelos víctimas no vistos (por ejemplo, Gemma3, Llama3.2), superando a otros métodos tanto en ASR como en diversidad en escenarios de transferencia.
• Eficiencia: A pesar de las comparaciones por pares de $O(N^2)$, el costo computacional es insignificante porque son operaciones escalares; el costo real de retropropagación permanece en $O(N)$, haciendo que S-GFN sea casi tan eficiente como el GFN estándar.

5. Significado

Este artículo representa un avance significativo en la investigación de seguridad de LLM. Al resolver la inestabilidad en el entrenamiento de GFNs, permite el descubrimiento sistemático de vulnerabilidades diversas y de alto impacto que los métodos tradicionales de RL pasan por alto debido al colapso de modos.

• Para la Seguridad: Proporciona una herramienta más robusta para el red-teaming, permitiendo a los desarrolladores identificar un espectro más amplio de vectores de ataque antes del despliegue.
• Para la Defensa: Los ataques diversos generados por S-GFN conducen a modelos ajustados finamente para la seguridad que son más resilientes a un rango más amplio de ataques futuros.
• Para el Aprendizaje Automático: El Equilibrio de Trayectoria Contrastivo propuesto ofrece una solución generalizable para la coincidencia de distribuciones en entornos de alta dimensión y ruidosos, extendiendo la aplicabilidad de GFNs más allá del red-teaming a campos como el descubrimiento de fármacos y la ciencia de materiales.

En resumen, Stable-GFlowNet cierra con éxito la brecha entre la diversidad teórica de las GFNs y la estabilidad práctica requerida para el red-teaming de LLMs en el mundo real, estableciendo un nuevo estado del arte para la generación de prompts adversarios diversos y efectivos.