MANATEE: Inference-Time Lightweight Diffusion Based Safety Defense for LLMs

Resumo Técnico: MANATEE

1. O Problema

A defesa de Grandes Modelos de Linguagem (LLMs) contra ataques de jailbreak (contorno de segurança) e comportamentos adversariais permanece um desafio crítico. As soluções existentes apresentam limitações significativas:

• Classificadores Binários: Frequentemente falham quando o input adversarial cai fora da fronteira de decisão aprendida, exigindo retreinamento constante e não garantindo generalização.
• Ajuste Fino (Fine-tuning) Repetido: É computacionalmente caro e pode degradar as capacidades do modelo ou falhar em eliminar comportamentos enganosos (como "agentes adormecidos" ou sleeper agents).
• Vulnerabilidade: Benchmarks como o JailbreakBench mostram que as defesas mais avançadas ainda são vulneráveis a ataques de otimização (ex: GCG) e de priming.

O objetivo central é desenvolver uma defesa robusta que não dependa de dados de treinamento maliciosos, não altere a arquitetura do modelo base e preserve a utilidade do modelo em entradas benignas.

2. Metodologia: MANATEE

O MANATEE (MAnifold-based ANomaly Detection and Steering via Diffusion) é uma defesa executada no momento da inferência (inference-time). Ele opera no espaço de representações (espaço latente) do modelo, sem modificar os pesos do LLM base.

A abordagem baseia-se na premissa de que, em vez de perguntar "esta representação é prejudicial?", deve-se perguntar "como projetar esta representação no manifold de representações benignas?".

Fases Principais:

1. Definição do Manifold Benigno (Treinamento):

   • O sistema extrai os estados ocultos da camada final de gerações benignas do modelo.
   • Esses vetores são padronizados (média zero, variância unitária).
   • Um modelo de difusão (DDPM - Denoising Diffusion Probabilistic Models) é treinado para aprender a função de pontuação (score function) e a densidade desses estados benignos. O objetivo é modelar a distribuição de probabilidade das representações seguras.

2. Detecção de Anomalias (Inferência):

   • Para uma nova entrada, o estado oculto final ($h$) é extraído e padronizado.
   • Calcula-se uma pontuação de anomalia ($s(h)$) baseada na magnitude do ruído previsto pelo denoiser em um passo de tempo fixo ($t_{check}$).
   • Estados dentro da distribuição (benignos) geram pontuações baixas; estados fora da distribuição (ataques/ backdoors) geram pontuações altas.

3. Mecanismo de Resposta Condicional:

   • Se $s(h) > \tau$ (Limiar): O sistema dispara uma recusa automática (o modelo recusa a resposta).
   • Se $s(h) \le \tau$: O sistema aplica direcionamento condicional (steering).
     • O estado oculto anômalo é injetado com ruído e submetido a um processo de difusão reversa (denoising).
     • Isso projeta iterativamente a representação de volta para regiões de alta densidade do manifold benigno.
     • O estado "purificado" é então passado para a cabeça de linguagem (LM head) para gerar uma resposta segura.

3. Contribuições Chave

• Defesa Leve e Plug-and-Play: Não requer dados de treinamento maliciosos, nem alterações na arquitetura do LLM base. Funciona como uma camada externa de segurança.
• Abordagem Baseada em Manifold: Reformula a segurança como um problema de estimativa de densidade, utilizando modelos de difusão para detectar e corrigir desvios no espaço latente.
• Transferibilidade: O método foi demonstrado como eficaz entre diferentes famílias de modelos (Mistral, Llama, Gemma) sem necessidade de retreinamento do defensor.
• Preservação de Utilidade: Ao contrário de filtros rígidos que podem bloquear entradas legítimas, o MANATEE tenta corrigir a representação antes de recusar, mantendo a capacidade do modelo em tarefas benignas.

4. Resultados Experimentais

Os experimentos foram conduzidos em três modelos principais: Mistral-7B-Instruct, Llama-3.1-8B-Instruct e Gemma-2-9B-it. Foram testados contra conjuntos de dados de backdoor (MAD, Sleeper Agent) e jailbreak (JailbreakBench).

• Redução da Taxa de Sucesso do Ataque (ASR):
  • MAD (Mechanistic Anomaly Detection): Redução média de 58,7%.
  • JBB (JailbreakBench): Redução média de 77,6%.
  • ASA (Anthropic Sleeper Agent): Redução de 100% em todos os modelos testados (Mistral, Gemma, Llama).
  • Em casos específicos (ex: Mistral no JBB), a ASR caiu de 98% para 0%.
• Preservação de Utilidade: As entradas benignas permaneceram praticamente inalteradas, indicando que o modelo de difusão foca especificamente nas representações desviadas (maliciosas) e não degrada o desempenho em tarefas normais.
• Qualidade da Resposta: Exemplos qualitativos mostram que, enquanto o modelo backdoored fornecia instruções para atividades ilegais (ex: falsificar prescrições, contrabando), o MANATEE redirecionava a resposta para uma recusa segura ("Não posso ajudar com isso").

5. Significado e Impacto

O MANATEE representa um avanço significativo na segurança de LLMs ao demonstrar que a segurança e a capacidade não precisam ser trade-offs.

• Eficiência: Oferece uma defesa computacionalmente mais viável do que o retreinamento contínuo de modelos.
• Robustez: Ao tratar comportamentos inseguros como trajetórias fora da distribuição (OOD) e corrigi-los via difusão, oferece uma defesa mais robusta contra ataques que tentam explorar fronteiras de decisão rígidas.
• Futuro: O trabalho sugere que mecanismos baseados em difusão podem ser expandidos para raciocínio em espaço latente contínuo (como Continuous Chain of Thought), abrindo caminho para defesas mais sofisticadas e adaptáveis.

Em suma, o MANATEE propõe uma mudança de paradigma: em vez de apenas classificar e rejeitar, o sistema corrige a representação interna do modelo, garantindo que a saída permaneça alinhada com o comportamento seguro aprendido.