Stable-GFlowNet: Rumo a um Red-Teaming de LLM Diverso e Robusto via Equilíbrio de Trajetória Contrastivo

Resumo

O Teste de Adversário (Red-Teaming) de Modelos de Linguagem de Grande Porte (LLM), que identifica proativamente as vulnerabilidades desses modelos, é um processo essencial para garantir a segurança. Encontrar ataques eficazes e diversificados durante o teste de adversário é importante, mas alcançar ambos é um desafio. As Redes de Fluxo Generativo (Generative Flow Networks - GFNs), que realizam correspondência de distribuição, são métodos promissores, mas são notórios pela instabilidade no treinamento e pelo colapso modal. Especificamente, recompensas instáveis no teste de adversário aceleram o colapso modal. Propomos a Stable-GFN (S-GFN), que elimina a estimativa da função de partição Z na GFN e reduz a instabilidade do treinamento. A S-GFN evita a estimativa de Z por meio de comparações pareadas e emprega uma metodologia robusta de mascaramento contra recompensas ruidosas. Adicionalmente, propomos um estabilizador de fluência para evitar que o modelo fique preso em ótimos locais que produzem textos sem sentido. A S-GFN proporciona um treinamento mais estável, mantendo a política ótima da GFN. Demonstramos o desempenho de ataque avassalador e a diversidade da S-GFN em várias configurações.

English

Large Language Model (LLM) Red-Teaming, which proactively identifies vulnerabilities of LLMs, is an essential process for ensuring safety. Finding effective and diverse attacks in red-teaming is important, but achieving both is challenging. Generative Flow Networks (GFNs) that perform distribution matching are a promising methods, but they are notorious for training instability and mode collapse. In particular, unstable rewards in red-teaming accelerate mode collapse. We propose Stable-GFN (S-GFN), which eliminates partition function Z estimation in GFN and reduces training instability. S-GFN avoids Z-estimation through pairwise comparisons and employs a robust masking methodology against noisy rewards. Additionally, we propose a fluency stabilizer to prevent the model from getting stuck in local optima that produce gibberish. S-GFN provides more stable training while maintaining the optimal policy of GFN. We demonstrate the overwhelming attack performance and diversity of S-GFN across various settings.