Modelos de difusão destilados aprendem a esquecer conceitos indesejados com DPO

Modelos text-to-image rápidos, comprimidos para alguns passos de geração em vez de dezenas, tornaram-se o padrão em produtos de baixa latência. Junto com a velocidade, eles herdaram dos modelos base a capacidade de gerar o que não deveriam: pessoas específicas, nudez, objetos protegidos. Os autores de um novo preprint no arXiv (cs.CV) afirmam que os métodos existentes de remoção desses conceitos funcionam mal em modelos destilados e propõem uma variante própria, baseada em Direct Preference Optimization.

Por que o unlearning clássico falha em modelos destilados

A maioria dos métodos de esquecimento orientados a dados faz o ajuste fino dos pesos do modelo com funções-objetivo específicas. Essas funções se apoiam implicitamente na dinâmica da remoção de ruído em vários passos: o modelo limpa a imagem aos poucos, e o sinal de «esquecimento» se distribui ao longo de uma trajetória longa. Nos modelos few-step distilled (FSD), a trajetória se reduz a pouquíssimos passos, a dinâmica de geração é outra, e essa premissa deixa de valer. O resultado, segundo os autores, é um esquecimento ineficaz: o conceito-alvo continua aparecendo.

Há um caminho óbvio para contornar o problema. Dá para fazer o unlearning no modelo base não destilado e, depois, repetir a destilação. Mas isso exige muita computação e tempo, e em muitos cenários o processo é impraticável. Principalmente quando o modelo base é atualizado e a lista de conceitos proibidos cresce.

O que os autores propõem

Os pesquisadores voltaram ao Direct Preference Optimization para modelos de difusão. O DPO padrão e suas variantes para unlearning são formulados em torno do erro de predição do ruído. Em modelos FSD, essa formulação se transfere mal justamente por causa da dinâmica de geração alterada.

No lugar dela, os autores introduzem uma formulação modificada de preference optimization, alinhada de forma explícita às propriedades da geração em poucos passos. A ideia é remover o conceito diretamente no modelo destilado, sem afetar a velocidade e sem destruir as demais habilidades.

Parâmetros do trabalho divulgados no resumo (abstract):

• Objeto: modelos de difusão text-to-image few-step distilled.
• Método: unlearning orientado por preferências, baseado em um DPO reformulado para a geração FSD.
• Tarefas principais: remoção de identidades (identity) e de conteúdo NSFW (nudity).
• Adicionalmente: unlearning no nível de objetos.
• Propriedades declaradas: esquecimento robusto dos conceitos-alvo, preservação do regime de poucos passos, forte retenção das capacidades «permitidas».
• O resumo não traz métricas numéricas; é preciso consultá-las no texto completo do artigo.

O que as equipes que lançam modelos rápidos devem fazer

Para quem coloca modelos destilados em produção, a conclusão prática principal é que não se pode presumir que os métodos de segurança do modelo base se transfiram automaticamente para o destilado. A verificação deve ser feita no próprio modelo FSD e nos números de passos realmente usados no serviço. Se um filtro «esqueceu» o conceito com 50 passos, isso não garante nada para 1–4 passos.

Por enquanto, o trabalho existe apenas como preprint e não tem verificação independente. Antes de adotá-lo, convém esperar o código e os checkpoints, rodar seus próprios conjuntos de prompts de red team e medir a degradação colateral: por exemplo, como a remoção de uma pessoa específica afeta a geração de pessoas com aparência parecida. Também vale comparar o custo da abordagem com o de uma nova destilação no seu hardware: a economia de tempo e de computação é o principal argumento aqui.

Em cenários regulatórios e jurídicos, de pedidos de remoção de biometria e da imagem de uma pessoa à moderação de NSFW, poder corrigir diretamente um modelo já comprimido parece bem mais conveniente do que refazer todo o pipeline. Os autores não dizem, no resumo, o quanto o método resiste a prompts de contorno e a ajuste fino feito por um agente malicioso, e essa questão segue em aberto.