PISmith: Reinforcement Learning-based Red Teaming for Prompt Injection Defenses

Título: PISmith: Red Teaming Baseado em Aprendizado por Reforço para Defesas contra Injeção de Prompt

1. O Problema

As aplicações de Grandes Modelos de Linguagem (LLMs), especialmente agentes autônomos, são vulneráveis a ataques de injeção de prompt. Neste cenário, um adversário insere instruções maliciosas no contexto de entrada (ex.: conteúdo da web, documentos) para manipular a saída do LLM, desviando-o de sua tarefa original.

Embora diversas defesas tenham sido propostas (baseadas em filtragem ou prevenção), a robustez delas contra ataques adaptativos (onde o atacante ajusta sua estratégia especificamente para contornar a defesa) permanece insuficientemente avaliada. A maioria das avaliações atuais utiliza ataques estáticos ou de busca, que não capturam a capacidade de um atacante aprender e evoluir. Além disso, métodos existentes de red teaming (testes de invasão) baseados em Aprendizado por Reforço (RL) enfrentam o problema da esparsidade extrema de recompensas: como as defesas fortes bloqueiam a maioria das tentativas, o modelo de ataque recebe pouquíssimos sinais de sucesso, levando ao colapso da exploração e ao fracasso no treinamento.

2. Metodologia: PISmith

Os autores propõem o PISmith, um framework de red teaming baseado em RL que treina um LLM de ataque para otimizar prompts injetados em um cenário de "caixa preta" (o atacante só vê a entrada e a saída do modelo defendido).

Desafios Identificados no RL Padrão (GRPO):
Ao aplicar o algoritmo padrão GRPO (Group Relative Policy Optimization) diretamente, os autores observaram dois modos de falha:

1. Colapso de Entropia: O modelo superajusta (overfit) aos poucos prompts que funcionam acidentalmente, reduzindo a diversidade de geração e parando a exploração de novas estratégias.
2. Diluição do Gradiente: Mesmo mantendo a exploração, os poucos sucessos são "diluídos" pela vasta maioria de falhas no cálculo do gradiente, impedindo um aprendizado eficiente.

Soluções Propostas (PISmith):
Para superar esses obstáculos, o PISmith introduz dois mecanismos inovadores:

• Regularização de Entropia Adaptativa (Adaptive Entropy Regularization):

  • Em vez de uma penalidade fixa, o sistema ajusta dinamicamente o incentivo à exploração (entropia) com base na recompensa média atual ($\bar{r}$).
  • Quando a recompensa é baixa (poucos sucessos), o coeficiente de entropia aumenta, forçando o modelo a gerar respostas mais diversas para descobrir novas falhas.
  • Um "teto de entropia" (entropy cap) impede que o modelo gere texto incoerente.
  • À medida que a taxa de sucesso aumenta, o incentivo à exploração diminui, permitindo que o modelo explore (exploit) as estratégias vencedoras.

• Pesagem Dinâmica de Vantagem (Dynamic Advantage Weighting):

  • Para lidar com a raridade dos sucessos, o PISmith amplifica o peso do gradiente das tentativas bem-sucedidas.
  • Um multiplicador dinâmico ($\gamma$) aumenta a contribuição dos poucos rollouts (tentativas) que resultam em recompensa positiva, garantindo que esses sinais raros dirijam a atualização da política, em vez de serem ignorados pela maioria das falhas.

A função de perda final combina a otimização padrão do GRPO com esses dois componentes, permitindo que o modelo aprenda eficientemente mesmo em ambientes de recompensa extremamente esparsa.

3. Contribuições Principais

1. Framework PISmith: Uma nova abordagem de red teaming que supera as limitações do GRPO padrão em cenários de defesa forte, utilizando regularização de entropia adaptativa e pesagem de vantagem dinâmica.
2. Avaliação Abrangente: Testes realizados em 13 benchmarks (incluindo Resposta a Perguntas, Geração Aumentada por Recuperação - RAG, e tarefas de Longo Contexto) contra 8 defesas de última geração.
3. Descoberta Fundamental: Evidência de que as defesas atuais não conseguem manter simultaneamente alta utilidade (desempenho em tarefas benignas) e robustez contra ataques adaptativos. Existe uma tensão fundamental entre utilidade e segurança.
4. Validação em Agentes: O método foi testado com sucesso em ambientes de agentes complexos (InjecAgent e AgentDojo), demonstrando eficácia contra modelos de código aberto e fechados (incluindo GPT-4o-mini e GPT-5-nano).

4. Resultados

• Desempenho de Ataque: O PISmith alcançou uma Taxa de Sucesso de Ataque (ASR) média de 0.87 (ASR@1) e 1.0 (ASR@10) contra o modelo de defesa Meta-SecAlign-8B, superando significativamente todas as outras linhas de base (incluindo RL-Hammer, TAP, PAIR e ataques estáticos).
  • Em tarefas difíceis como HotpotQA-Long, o PISmith atingiu 0.61 de ASR@1, enquanto o melhor concorrente (RL-Hammer) obteve apenas 0.01.
• Generalização: Treinado em apenas 100 amostras de um único dataset, o PISmith generalizou perfeitamente para 12 benchmarks não vistos, incluindo tarefas de longo contexto e RAG.
• Eficiência: O PISmith requer menos consultas ao modelo alvo (queries) do que métodos baseados em busca (como TAP e PAIR), sendo mais eficiente computacionalmente.
• Trade-off Utilidade-Robustez: A análise mostrou que defesas que mantêm alta utilidade são facilmente quebradas pelo PISmith, enquanto defesas robustas tendem a degradar severamente o desempenho do modelo em tarefas legítimas. Nenhuma defesa testada ocupou a região ideal (alta utilidade e baixa taxa de sucesso de ataque).
• Agentes: No benchmark AgentDojo, o PISmith superou todas as estratégias baseadas em templates, alcançando taxas de sucesso superiores a 0.60 em modelos fechados onde ataques estáticos falharam completamente.

5. Significado e Conclusão

O trabalho do PISmith destaca uma lacuna crítica na segurança de LLMs: as defesas atuais são frágeis contra adversários adaptativos inteligentes. O estudo demonstra que a simples aplicação de técnicas de RL padrão é insuficiente devido à esparsidade de recompensas, e que mecanismos adaptativos são necessários para realizar red teaming efetivo.

A principal implicação é que o desenvolvimento de defesas futuras deve focar em resolver o dilema entre utilidade e segurança, pois as abordagens atuais não conseguem oferecer ambos. O PISmith serve como uma ferramenta essencial para estressar e avaliar a robustez de novas defesas antes de sua implantação no mundo real, promovendo o desenvolvimento de sistemas LLM mais seguros e resilientes. O código do projeto está disponível publicamente para a comunidade de pesquisa.