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

1. Declaração do Problema

O red-teaming de Modelos de Linguagem de Grande Escala (LLMs) visa descobrir proativamente vulnerabilidades (saídas tóxicas ou prejudiciais) antes da implantação. Embora métodos de Aprendizado por Reforço (RL) possam encontrar ataques de alta recompensa, eles sofrem de colapso de modo, gerando padrões de ataque repetitivos e estreitos. Por outro lado, Redes de Fluxo Generativo (GFNs) oferecem um framework para correspondência de distribuição, permitindo teoricamente a descoberta de trajetórias de alta recompensa diversas. No entanto, a aplicação de GFNs ao red-teaming de LLMs enfrenta dois desafios críticos:

1. Instabilidade de Treinamento: Objetivos padrão de GFN (como o Equilíbrio de Trajetória, TB) exigem a estimativa de uma função de partição global ($Z$). No espaço discreto e de alta dimensão dos LLMs, estimar $Z$ é numericamente instável, levando a uma convergência pobre e colapso de modo.
2. Recompensas Ruidosas: Classificadores de toxicidade fornecem sinais densos, mas ruidosos. Frequentemente atribuem recompensas não nulas a textos sem sentido (gibberish) ou fora da distribuição (OOD), fazendo com que o modelo "hackeie" a recompensa gerando texto nonsenso em vez de ataques fluídos e diversos.

2. Metodologia: Stable-GFlowNet (S-GFN)

Os autores propõem o Stable-GFlowNet (S-GFN), um framework que elimina a necessidade de estimativa explícita da função de partição e filtra recompensas ruidosas para garantir um treinamento estável e diverso. O pipeline consiste em três componentes principais:

A. Equilíbrio de Trajetória Contrastivo (CTB)

Para abordar a instabilidade da estimativa de $Z$, o S-GFN substitui o objetivo de correspondência de fluxo absoluto por um objetivo de comparação par a par.

• Mecanismo: Em vez de minimizar o erro entre o fluxo da política e a recompensa em relação a $Z$, o CTB minimiza a diferença quadrada entre a razão logarítmica das probabilidades da política e a razão logarítmica das recompensas para pares de trajetórias $(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$$
• Garantia Teórica: Os autores provam que a política ótima para o CTB é idêntica à do Equilíbrio de Trajetória padrão ($\pi_\theta(y) \propto R(y)$).
• Redução de Variância: Ao usar uma amostra como baseline para outra (semelhante ao Reinforce Leave-One-Out), o CTB reduz implicitamente a variância do gradiente sem necessidade de uma rede de baseline externa ou estimativa de $Z$.

B. Poda de Gradientes Ruidosos (NGP)

Para mitigar o risco de amplificação de ruído de recompensa inerente às comparações par a par, o S-GFN introduz um filtro baseado em saliência.

• Mecanismo: Gradientes são computados apenas para pares de trajetórias onde a diferença absoluta nas recompensas logarítmicas excede um limiar $\sigma$:
  $$L_{NGP} = \mathbb{I}(|\log R(y_1) - \log R(y_2)| > \sigma) \cdot L_{CTB}$$
• Efeito: Isso impede que o otimizador atualize com base em diferenças de recompensa insignificantes causadas pela estocasticidade do LLM vítima ou do classificador de toxicidade, podando efetivamente os gradientes ruidosos.

C. Estabilizador de Fluência Min-K (MKS)

Para evitar que o modelo convirja para mínimos locais de texto sem sentido (que frequentemente recebem pontuações de toxicidade espúrias), o S-GFN emprega uma restrição de fluência.

• Mecanismo: Em vez de uma penalidade global de divergência KL (que pode levar ao overfitting em um modelo de referência), o MKS calcula a estatística Min-K: a média das probabilidades logarítmicas dos $k$ tokens menos prováveis em uma sequência gerada sob um modelo de referência.
• Penalidade Rígida: Se a pontuação Min-K cair abaixo de um limiar $T_{mks}$, a recompensa é limitada a zero (ou a um valor negativo grande).
  $$R_{MKS}(y) = \mathbb{I}(M_k(y) \ge T_{mks}) \cdot R(y)$$
• Benefício: Isso garante que apenas prompts fluídos e linguisticamente coerentes sejam considerados de alta recompensa, permitindo que o modelo explore um espaço mais amplo de ataques válidos sem gerar nonsense.

3. Contribuições Principais

1. Framework Stable-GFN: Um novo framework de red-teaming que elimina a necessidade de estimativa da função de partição ($Z$), resolvendo a causa primária da instabilidade de treinamento em GFNs para LLMs.
2. Equilíbrio de Trajetória Contrastivo (CTB): Uma nova função objetivo que alcança a mesma distribuição ótima que o TB, mas com variância significativamente menor e maior estabilidade de treinamento.
3. Robustez ao Ruído: A integração de NGP e MKS aborda especificamente os desafios únicos do red-teaming (classificadores ruidosos e recompensas de texto sem sentido), prevenindo colapso de modo e hacking de recompensa.
4. Universalidade: O artigo demonstra que o CTB e o NGP são eficazes não apenas no red-teaming, mas também em outras tarefas de correspondência de distribuição, como geração molecular e navegação em hipergrid.

4. Resultados Experimentais

Os autores avaliaram o S-GFN contra baselines incluindo PPO, Jailbreak R1, Rainbow Teaming e GFN padrão (TB) em vários modelos vítimas (Qwen, Llama, Gemma) e configurações de defesa.

• Diversidade (Ataques Únicos): O S-GFN gerou 134 prompts de ataque únicos, uma melhoria massiva sobre a baseline padrão de GFN (17) e PPO (3). Produziu aproximadamente 7x mais ataques únicos que a baseline de GFN.
• Taxa de Sucesso do Ataque (ASR): O S-GFN manteve uma ASR alta de ~92,5%, comparável às baselines de melhor desempenho (GFN: 93,7%, PPO: 91,7%), provando que a diversidade não veio à custa da eficácia.
• Transferibilidade Cruzada de Ataques: Quando os modelos vítimas foram ajustados para segurança (safety-fine-tuned) usando ataques do S-GFN, tornaram-se significativamente mais robustos contra ataques de outros métodos (a ASR caiu para perto de 0% para ataques de GFN). Por outro lado, os ataques do S-GFN permaneceram eficazes contra modelos ajustados em outros métodos, demonstrando generalização superior.
• Ataques de Transferência: O S-GFN transferiu com sucesso ataques para modelos vítimas não vistos (por exemplo, Gemma3, Llama3.2), superando outros métodos tanto em ASR quanto em diversidade em cenários de transferência.
• Eficiência: Apesar das comparações par a par $O(N^2)$, o custo computacional é negligenciável porque são operações escalares; o custo real de retropropagação permanece $O(N)$, tornando o S-GFN quase tão eficiente quanto o GFN padrão.

5. Significado

Este artigo representa um avanço significativo na pesquisa de segurança de LLMs. Ao resolver a instabilidade de treinamento de GFNs, permite a descoberta sistemática de vulnerabilidades diversas e de alto impacto que métodos tradicionais de RL perdem devido ao colapso de modo.

• Para Segurança: Fornece uma ferramenta mais robusta para red-teaming, permitindo que desenvolvedores identifiquem um espectro mais amplo de vetores de ataque antes da implantação.
• Para Defesa: Os ataques diversos gerados pelo S-GFN levam a modelos ajustados para segurança que são mais resilientes a uma gama mais ampla de ataques futuros.
• Para Aprendizado de Máquina: O Equilíbrio de Trajetória Contrastivo proposto oferece uma solução generalizável para correspondência de distribuição em ambientes de alta dimensão e ruidosos, estendendo a aplicabilidade de GFNs além do red-teaming para campos como descoberta de fármacos e ciência dos materiais.

Em resumo, o Stable-GFlowNet conecta com sucesso a lacuna entre a diversidade teórica das GFNs e a estabilidade prática exigida para red-teaming de LLMs no mundo real, estabelecendo um novo estado da arte para a geração de prompts adversariais diversos e eficazes.