Unlearning for One-Step Generative Models via Unbalanced Optimal Transport

Each language version is independently generated for its own context, not a direct translation.

Imagine que você tem um chef de cozinha genial (o modelo de IA) que aprendeu a cozinhar milhões de pratos diferentes, desde pizza até sushi. Esse chef é tão rápido que consegue criar um prato completo em uma única piscada de olhos (os chamados "modelos de um passo").

No entanto, há um problema: o chef aprendeu a fazer alguns pratos que não deveriam ser servidos. Talvez sejam pratos com ingredientes proibidos por lei, ou receitas que violam direitos autorais. A solução óbvia seria demitir o chef e contratar um novo que nunca viu essas receitas. Mas isso é caro e demorado.

Aqui entra a ideia de "Esquecimento de Máquina" (Machine Unlearning): queremos ensinar o chef a esquecer especificamente aquele prato proibido, sem que ele perca a habilidade de fazer todos os outros milhões de pratos deliciosos.

O Problema: Por que os métodos antigos falham?

Antes, os chefs de IA eram lentos. Eles criavam um prato desenhando-o primeiro com um rabisco, depois ajustando o rabisco, depois polindo, e assim por diante, em dezenas de passos. Para fazer o chef "esquecer" algo, os cientistas tentavam mudar a receita em cada um desses passos intermediários.

Mas o nosso novo chef (o modelo de um passo) não faz rabiscos. Ele pula direto do "nada" para o "prato pronto". Não existem passos intermediários para mexer! Se tentarmos usar as técnicas antigas, é como tentar consertar um carro de Fórmula 1 enquanto ele está voando: não funciona e pode explodir o motor (destruir a qualidade das imagens).

A Solução: O "Transporte Desbalanceado" (UOT)

Os autores deste paper propõem uma solução inteligente chamada UOT-Unlearn. Para entender, vamos usar uma analogia de mudança de casa.

O Cenário: Imagine que a memória do chef é uma casa cheia de caixas. A caixa "Gato" (o que queremos esquecer) está cheia de brinquedos de gato. As outras caixas têm brinquedos de cachorro, pássaro, etc.
O Método Antigo (Otimização de Transporte Clássica): Era como tentar mover exatamente a mesma quantidade de brinquedos para cada caixa, sem perder nada. Se você tirasse os brinquedos de gato, teria que jogar fora tudo, deixando a casa vazia ou bagunçada. O resultado era uma casa com buracos ou brinquedos estranhos misturados.
O Método Novo (Transporte Desbalanceado - UOT): O UOT é mais flexível. Ele diz: "Ok, vamos tirar os brinquedos de gato. Mas, em vez de jogar fora, vamos redistribuir esses brinquedos para as caixas de cachorro e pássaro de forma natural".

O segredo do UOT é que ele permite que a "massa" (a probabilidade de gerar uma imagem) saia da caixa proibida e se espalhe suavemente pelas caixas permitidas, mantendo a casa organizada.

Como funciona na prática?

O método usa duas regras principais (custos):

A Regra do "Não Faça Isso" (Custo de Esquecimento): O sistema cria um "ímã" invisível para o conceito proibido (por exemplo, um ímã para a palavra "Gato"). Se o chef tentar gerar algo parecido com um gato, o sistema aplica uma multa pesada. Isso força o chef a se afastar desse conceito.
A Regra da "Qualidade" (Custo de Fidelidade): Ao mesmo tempo, o sistema diz: "Mas não invente coisas estranhas! Mantenha a qualidade dos outros pratos". Se o chef tentar gerar um "Gato-Cachorro-Raio", ele é punido. Ele é obrigado a transformar a ideia de "Gato" em algo que se pareça com "Cachorro" ou "Pássaro", mas que ainda seja um animal bonito e realista.

O Resultado Mágico

Ao usar essa técnica, os pesquisadores conseguiram:

Apagar o conceito proibido: O chef quase nunca mais gera a imagem que você quer esquecer (90% a 95% de sucesso).
Manter a qualidade: As imagens que ele gera continuam lindas e realistas. Ele não começa a gerar "barulhos" ou imagens distorcidas.
Não precisar de dados reais: O método é tão inteligente que consegue fazer isso usando apenas imagens que o próprio chef gera, sem precisar que os humanos mostrem fotos reais dos outros pratos para ele estudar.

Resumo em uma frase

É como ensinar um gênio da culinária a esquecer uma receita proibida sem que ele perca a habilidade de cozinhar, fazendo com que ele transforme os ingredientes dessa receita proibida em novos pratos deliciosos, em vez de simplesmente jogar tudo no lixo e estragar a cozinha.

Isso é crucial para o futuro, pois modelos de IA estão ficando cada vez mais rápidos. Se não tivermos uma maneira rápida e segura de "apagar" coisas ruins desses modelos rápidos, eles podem espalhar conteúdo prejudicial em velocidade da luz. O UOT-Unlearn é o freio de segurança inteligente para essa nova geração de IAs.

Each language version is independently generated for its own context, not a direct translation.

1. Problema e Contexto

O avanço recente em modelos generativos de um passo (one-step), como Modelos de Fluxo (Flow Maps) e Modelos de Consistência (Consistency Models), permitiu a geração de imagens de alta qualidade em uma única passagem de rede neural, eliminando a necessidade das centenas de iterações de denoising exigidas pelos modelos de difusão tradicionais. Embora isso resolva o gargalo de velocidade de inferência, cria um novo desafio de segurança: a Machine Unlearning (Esquecimento de Máquina).

O Desafio: Garantir que modelos generativos rápidos não produzam conteúdo indesejado (ex: NSFW, materiais com direitos autorais) sem a necessidade de retreinamento completo, que é computacionalmente proibitivo.
A Lacuna: As técnicas existentes de esquecimento foram desenvolvidas para modelos de difusão de múltiplos passos, que permitem modificações em etapas intermediárias de denoising. Essas técnicas são incompatíveis com modelos de um passo, pois estes mapeiam o ruído diretamente para os dados sem etapas intermediárias acessíveis para ajustes graduais.
A Necessidade: É urgente desenvolver um framework de esquecimento que funcione nativamente na arquitetura de "passada única" (single forward pass) sem depender de dados reais de retenção (retain data) durante a otimização.

2. Metodologia: UOT-Unlearn

Os autores propõem o UOT-Unlearn, um framework plug-and-play baseado no Transporte Ótimo Desbalanceado (Unbalanced Optimal Transport - UOT).

Conceito Central

Ao contrário do Transporte Ótimo (OT) clássico, que exige um casamento perfeito das margens das distribuições (massa de probabilidade total deve ser preservada), o UOT relaxa essas restrições. Ele minimiza um compromisso (trade-off) entre o custo de transporte e o desvio das margens (penalizado por divergências f-divergence).

Formulação do Problema

O processo de esquecimento é formulado como um problema de transporte onde:

Distribuição Fonte ( $\mu$ ): A distribuição do modelo pré-treinado ( $p_{pre}$ ).
Distribuição Alvo ( $\nu$ ): A distribuição de dados original ( $p_{data}$ ).
Objetivo: Redistribuir a massa de probabilidade associada à classe a ser esquecida (forget class) para as classes restantes, minimizando o custo de transporte, mas permitindo um desvio nas margens para "remover" a classe indesejada.

Mecanismo de Custo de Esquecimento

A inovação chave é o design de uma função de custo ( $c_{ul}$ ) que atua em duas frentes:

Custo de Esquecimento (Active Expulsion): Para amostras geradas que caem na região da classe a ser esquecida (definida por uma distância coseno em relação a um "centroide" pré-computado da classe), aplica-se uma penalidade pesada. Isso força o modelo a afastar a probabilidade dessa região.
Custo de Retenção (Fidelity & Transport): Para amostras fora da região de esquecimento, aplica-se um custo $L_2$ quadrado entre a saída do modelo atual e a saída do modelo pré-treinado. Isso preserva a fidelidade das classes restantes e atua como o termo de transporte.

Otimização sem Dados Reais

Um diferencial crucial é que o método não requer dados reais de retenção durante o treinamento.

Utiliza apenas amostras sintéticas geradas pelo próprio modelo pré-treinado.
Utiliza um centroide pré-computado ( $\mu_f$ ) da classe a ser esquecida (calculado uma única vez com um pequeno conjunto de dados).
Aproxima a distribuição alvo ( $\nu$ ) usando a própria distribuição do modelo pré-treinado, permitindo uma otimização totalmente "data-free" (sem dados reais) após a fase inicial.

A função objetivo é derivada da formulação semi-dual do UOT, onde um potencial dual ( $v_\phi$ ) e o mapeamento de transporte (o próprio gerador $G_\theta$ ) são otimizados alternadamente.

3. Principais Contribuições

Primeiro Framework para Modelos de Um Passo: Introdução do UOT-Unlearn, o primeiro método de esquecimento de classe projetado especificamente para arquiteturas generativas de um passo (como CTM e MeanFlow).
Formulação via UOT: Desenvolvimento de um objetivo baseado em UOT que permite a redistribuição suave da massa de probabilidade da classe esquecida para as classes restantes, evitando o colapso em ruído ou amostras de baixa qualidade.
Independência de Dados Reais: O método opera sem acesso a dados reais de retenção durante a fase de otimização, superando a dependência de dados de grandes conjuntos de treinamento presentes em métodos anteriores.
Mecanismo de Redistribuição Estruturada: Em vez de apenas suprimir a classe (o que pode degradar a qualidade geral), o método "move" a probabilidade para regiões semânticas válidas das classes retidas.

4. Resultados Experimentais

Os autores avaliaram o método nos conjuntos de dados CIFAR-10 e ImageNet-256, utilizando arquiteturas como CTM e MeanFlow.

Métricas:
- PUL (Percentage of Unlearning): Mede a eficácia na remoção da classe (redução na frequência de geração).
- u-FID (Unlearned FID): Mede a qualidade das imagens geradas para as classes retidas (quanto menor, melhor).
Desempenho:
- O UOT-Unlearn superou consistentemente os baselines (Gradient Ascent, Selective Amnesia, SalUn, VDU) em todos os cenários.
- Alta Eficácia: Alcançou PUL superior (ex: >90% em alguns casos no CIFAR-10), indicando remoção quase total da classe alvo.
- Alta Fidelidade: Mantiveu um u-FID muito baixo, preservando a qualidade e a diversidade das classes restantes. Em contraste, métodos como Gradient Ascent (GA) frequentemente causaram distorções severas na distribuição (u-FID alto) ao tentar remover a classe.
- Visualização: Em dados sintéticos 2D, o método mostrou uma redistribuição suave da massa de probabilidade da classe esquecida para as classes retidas, enquanto métodos de base distorciam a distribuição para regiões inválidas.
- ImageNet-256: Mesmo em alta resolução e cenários condicionais, o método conseguiu remover conceitos (ex: "Goldfish") mantendo a integridade estrutural das classes aquáticas retidas, com um PUL de 85% e u-FID de 20, comparado a um u-FID de quase 80 para o baseline GA.

5. Significado e Impacto

Este trabalho é significativo por preencher uma lacuna crítica na segurança de IA generativa. À medida que os modelos de geração de imagem se tornam mais rápidos (um passo), o risco de disseminação de conteúdo prejudicial aumenta. O UOT-Unlearn oferece uma solução viável e eficiente para "desaprender" conceitos indesejados sem o custo proibitivo de retreinar o modelo do zero ou sem a necessidade de manter grandes bancos de dados de treinamento.

Ao reformular o esquecimento como um problema de transporte de probabilidade desbalanceado, os autores demonstram que é possível equilibrar a força de remoção e a consistência da distribuição de forma matematicamente fundamentada. Isso abre caminho para a aplicação segura de modelos generativos ultra-rápidos em ambientes de produção onde a conformidade e a segurança são obrigatórias.