Modular Energy Steering for Safe Text-to-Image Generation with Foundation Models

1. O Problema

Os modelos generativos de texto para imagem (T2I) enfrentam desafios críticos de segurança, como a geração de conteúdo não seguro para o trabalho (NSFW), violação de direitos autorais e síntese de imagens de indivíduos reais sem consentimento. As abordagens de segurança existentes apresentam limitações significativas:

• Filtros de Entrada/Saída: Sofrem com altas taxas de falsos positivos e são vulneráveis a prompts adversariais.
• Ajuste Fino (Fine-tuning) ou Esquecimento (Unlearning): Requerem conjuntos de dados curados e atualização de pesos do modelo, o que frequentemente degrada a qualidade da geração e a capacidade geral do modelo.
• Direcionamento (Steering) Existente: Métodos de direcionamento em tempo de inferência que não alteram pesos muitas vezes dependem de conjuntos de dados curados ou operam apenas no nível de embeddings de prompts, resultando em desempenho subótimo.

O objetivo é desenvolver um mecanismo de segurança que seja treinamento-free (sem re-treinamento), livre de dados (sem necessidade de conjuntos de dados específicos de domínio) e que preserve a alta qualidade da geração.

2. Metodologia

Os autores propõem um framework de Direcionamento Modular Baseado em Energia (Modular Energy Steering) que opera em tempo de inferência. A ideia central é utilizar modelos fundacionais de visão e linguagem (como CLIP ou VLMs) congelados como estimadores de "energia semântica" para guiar o processo de geração.

Conceitos Chave:

• Estimadores de Energia Semântica: O framework trata a segurança como um problema de amostragem baseada em energia. Em vez de modificar os pesos do gerador, o método injeta um sinal de feedback derivado de modelos fundacionais pré-treinados (CLIP ou VLMs) para penalizar a presença de conceitos indesejados.
• Estimativa de Latente Limpo ($\hat{x}_{0|t}$): Em cada passo de amostragem do processo de difusão (ou flow-matching), o modelo estima uma versão "limpa" da imagem a partir do estado latente ruidoso atual. Essa estimativa é passada para o modelo fundacional (CLIP/VLM).
• Cálculo do Gradiente de Segurança:
  • Usando CLIP: Define-se um conjunto de hipóteses positivas (conceitos proibidos, ex: "nu", "sexual") e negativas (conceitos benignos). Calcula-se a diferença de log-likelihood (usando a função Log-Mean-Exponential) entre os conjuntos para obter um gradiente que empurra a geração para longe dos conceitos proibidos.
  • Usando VLMs: O VLM é consultado com uma pergunta binária (ex: "Esta imagem contém nudez? Sim/Não"). Os logits dos tokens de resposta ("Sim" vs "Não") são usados para calcular o gradiente de segurança via retropropagação através do codificador visual.
• Equação de Direcionamento: O gradiente de segurança é subtraído do campo de pontuação (score field) do modelo gerador, ponderado por uma força de direcionamento ($\lambda$). Isso modifica a trajetória de geração sem alterar os pesos originais do modelo.

Algoritmo:

O processo segue um ciclo de amostragem onde, a cada passo, o latente é atualizado, uma estimativa limpa é gerada, o modelo fundacional calcula a probabilidade de um conceito indesejado, e um gradiente de correção é aplicado se essa probabilidade exceder um limiar.

3. Principais Contribuições

1. Formulação Baseada em Energia: O controle de segurança em tempo de inferência é formulado matematicamente como um problema de amostragem baseada em energia, permitindo uma penalização principista de conceitos indesejados.
2. Modelos Fundacionais como Estimadores: Demonstra-se que modelos pré-treinados (CLIP e VLMs) podem ser reutilizados como estimadores de energia semântica eficazes, sem necessidade de treinamento adicional ou dados específicos.
3. Direcionamento Multi-Alvo Escalável e Livre de Dados: O método requer apenas uma "lista negra" textual definida pelo usuário. Ele escala naturalmente para suprimir múltiplos conceitos simultaneamente (ex: várias celebridades ou categorias) sem interferência entre eles e sem re-treinamento.
4. Compatibilidade Universal: O framework é compatível com diferentes arquiteturas de geração contínua, incluindo modelos baseados em Difusão (Diffusion Models) e Flow-Matching (como SD-v3 e SD-v3.5).

4. Resultados Experimentais

Os autores avaliaram o método em cenários de supressão de NSFW (nudez/sexualidade) e restrição de Identidade (celebridades).

• Robustez contra Ataques (Red-teaming):
  • O método alcançou desempenho state-of-the-art (SOTA) em benchmarks de ataques adversariais (P4D, Ring-A-Bell, MMA-Diffusion, UnlearnDiffAtk).
  • Reduziu drasticamente a Taxa de Sucesso do Ataque (ASR) para níveis muito baixos (ex: ~0.01 a 0.11 dependendo do dataset), superando ou igualando métodos que exigem ajuste fino ou datasets externos.
• Preservação de Qualidade:
  • A aplicação do direcionamento não degradou significativamente a qualidade da geração em prompts benignos. Métricas como FID e CLIP Score no dataset COCO-30K permaneceram comparáveis ao modelo original.
• Direcionamento de Identidade:
  • O método conseguiu suprimir com precisão a geração de identidades específicas (ex: Elon Musk, Taylor Swift) enquanto preservava a geração de outras identidades não-alvo, demonstrando alta seletividade.
  • Funcionou eficazmente em direcionamento múltiplo (até 18 identidades simultâneas).
• Compatibilidade com Modelos Avançados:
  • O método foi validado com sucesso no Stable Diffusion v3 e v3.5 (baseados em Flow-Matching e arquitetura DiT), provando sua generalização além dos modelos baseados em UNet tradicionais.
• Custo Computacional:
  • Adiciona um atraso moderado (aprox. 2.7s a 6.3s por imagem dependendo do modelo fundacional usado), mas evita os custos massivos de treinamento e curadoria de dados.

5. Significado e Impacto

Este trabalho oferece uma abordagem prática e escalável para a segurança em IA generativa. Ao transformar modelos fundacionais em "juízes" em tempo de execução, o framework elimina a necessidade de:

• Coletar e curar grandes conjuntos de dados de segurança.
• Realizar fine-tuning ou unlearning que podem corromper a capacidade criativa do modelo.
• Manter múltiplos modelos especializados para diferentes restrições.

A solução é modular: novos conceitos proibidos podem ser adicionados instantaneamente apenas atualizando a lista de texto (blacklist) e o prompt de consulta, tornando-a ideal para ambientes de produção onde as regras de segurança evoluem rapidamente (ex: novas leis de privacidade, direitos de celebridades). O trabalho estabelece um novo paradigma onde a segurança é tratada como um guia de energia semântica, em vez de uma restrição rígida de pesos.