Resumo Técnico: Rumo a um Controle Mais Geral de Modelos de Difusão Usando a Orientação de Jeffrey

1. Declaração do Problema

Os modelos de difusão oferecem uma flexibilidade significativa em tarefas generativas, permitindo o controle da saída durante a amostragem por meio de mecanismos de orientação sem a necessidade de retreinar a rede neural subjacente. Embora as técnicas de orientação padrão (por exemplo, orientação por classificador) lidem efetivamente com a amostragem condicional para classes específicas, elas frequentemente falham ao abordar objetivos mais complexos onde a distribuição alvo é implícita ou definida apenas por heurísticas.

Abordagens atuais para tarefas como equidade (fairness), controle de conceitos e geração composicional frequentemente dependem de modificações heurísticas na dinâmica de amostragem. Esses métodos inclinam a distribuição original, mas carecem de uma distribuição alvo explicitamente definida, tornando o objetivo real difícil de caracterizar ou interpretar. Há uma necessidade de um framework principiado que possa atualizar distribuições marginais em direção a um alvo prescrito enquanto preserva a estrutura condicional do modelo, permitindo assim o controle sobre propriedades estatísticas complexas (por exemplo, independência de atributos ou correspondência de distribuição de embedding) que a condicionamento padrão não consegue expressar.

2. Metodologia: Orientação de Jeffrey

O artigo propõe a orientação de Jeffrey (Jeffrey guidance), um framework que utiliza a regra de condicionamento de Jeffrey para atualizar modelos de difusão. Diferente das abordagens baseadas em Bayes que condicionam em eventos específicos (condicionamento padrão), a regra de Jeffrey permite condicionar em predicados complexos, como a correspondência de uma distribuição marginal específica.

Fundamentação Teórica

Dado um modelo conjunto $p(x, y)$ onde $x$ representa o espaço da imagem e $y$ representa atributos/características, o objetivo é atualizar a marginal $p(y)$ para uma distribuição alvo $p^\star(y)$ enquanto preserva a estrutura condicional $p(x|y)$. Sob a regra de Jeffrey, a distribuição conjunta atualizada é definida como:
$$p^\star(x, y) = p(x|y)p^\star(y)$$
Esta atualização minimiza a perturbação à distribuição conjunta original (uma projeção de informação) enquanto satisfaz a restrição marginal.

O artigo deriva uma formulação de reponderação para a marginal sobre $x$:
$$p^\star(x) = p(x) \int p(y|x)\rho(y) dy$$
onde $\rho(y) = p^\star(y)/p(y)$ é a razão de densidade. Em casos onde as características são derivadas deterministicamente de imagens ($y = u(x)$), isso se simplifica para:
$$p^\star(x) = p(x)\rho(u(x))$$

Implementação em Modelos de Difusão

Para aplicar isso a modelos de difusão baseados em score, os autores derivam uma função de score orientada. O score padrão $\nabla_{x_t} \log p_t(x_t)$ é aumentado com um termo de correção baseado na razão de densidade:
$$\nabla_{x_t} \log \tilde{p}^\star_t(x_t) = \nabla_{x_t} \log p_t(x_t) + \nabla_{x_t} \log \int p(y|\hat{x}_0(x_t, t))\rho(y) dy$$
onde $\hat{x}_0(x_t, t)$ é uma estimativa da amostra limpa obtida via a fórmula de Tweedie.

Para mapeamentos de características determinísticos (ex: espaços de embedding), a orientação simplifica-se para:
$$\nabla_{x_t} \log \tilde{p}^\star_t(x_t) = \nabla_{x_t} \log p_t(x_t) + \nabla_{x_t} \log \rho(u(\hat{x}_0(x_t, t)))$$

O framework introduz um hiperparâmetro $\lambda$ para escalar o termo de correção, controlando o equilíbrio entre a qualidade da amostra e a fidelidade à distribuição alvo. A razão de densidade $\rho$ é estimada usando um classificador binário treinado para distinguir entre amostras da distribuição alvo e a distribuição atual do modelo.

3. Principais Contribuições

O artigo apresenta três contribuições primárias:

1. Framework de Orientação de Jeffrey: Um novo método de orientação "plug-and-play" que generaliza a orientação por classificador padrão. Ele usa a regra de Jeffrey para fazer a correspondência das distribuições de saída geradas com distribuições alvo arbitrárias, permitindo o controle sobre objetivos anteriormente inacessíveis à orientação padrão.
2. Correspondência de Distribuição de Embedding: Uma prova de conceito demonstrando a capacidade de fazer a correspondência das distribuições de saída do modelo com uma distribuição alvo em um espaço de embedding (especificamente embeddings Inception). Esta aplicação reduz significativamente a Distância de Fréchet Inception (FID) sem exigir retreinamento.
3. Controle de Distribuição de Atributos para Equidade: Uma demonstração de controle das distribuições de atributos de imagem para alcançar objetivos de equidade. O método aplica com sucesso a paridade de gênero (correspondência de distribuição) e decorrelaciona atributos (ex: "Masculino" e "Jovem") em faces geradas, tarefas que a orientação por classificador padrão não consegue abordar diretamente.

4. Resultados Experimentais

Os autores avaliaram a orientação de Jeffrey nos datasets CIFAR-10, FFHQ e CelebA-HQ.

Orientação de Embedding (Redução de FID)

• Configuração: O método visou a distribuição de embedding Inception dos dados de treinamento.
• Resultados: No CIFAR-10, o FID melhorou de $\approx 5.88$ para $\approx 2.55$ (conjunto de treinamento) e de $\approx 9.76$ para $\approx 6.71$ (conjunto de teste). No FFHQ, o FID caiu de $\approx 20.13$ para $\approx 12.91$.
• Observação: Os autores observam que, embora o FID tenha diminuído drasticamente, as mudanças visuais nas imagens geradas foram mínimas, sugerindo limitações no FID como uma métrica puramente perceptual. Os melhores resultados foram alcançados aplicando a orientação apenas nos passos finais de denoising ($\delta=10$).

Aplicações de Equidade (Fairness)

• Paridade de Gênero (Correspondência de Distribuição): No CelebA-HQ, o método ajustou a proporção de atributos "Masculino" de $\approx 0.296$ para $\approx 0.522$ (visando 0.5), reduzindo a Discrepância de Equidade (FD) de $0.288$ para $0.032$. Isso foi alcançado com degradação mínima no FID ($21.61$ vs. baseline $22.30$).
• Decorrelação de Atributos: O método decorrelacionou com sucesso os atributos "Masculino" e "Jovem", reduzindo o coeficiente de correlação de Pearson de $-0.541$ para próximo de zero ($0.005$) em $\lambda=3.0$. Isso foi alcançado reduzindo o modo conjunto dominante (amostras de mulheres jovens) enquanto mantinha as distribuições marginais relativamente estáveis, com um aumento menor de aproximadamente um ponto no FID.
• Comparação: Comparado a um baseline de amostragem ancestral combinada com orientação padrão, a orientação de Jeffrey alcançou métricas de equidade semelhantes com custos de FID menores e modificou menos amostras, pois equilibra as direções de classe continuamente ao longo da trajetória de amostragem em vez de empurrar amostras inteiramente para uma única classe.

5. Significância e Alegações

O artigo posiciona a orientação de Jeffrey como uma extensão principiada do controle de modelos de difusão. Sua significância reside em:

• Generalização: Recontextualiza a orientação por classificador padrão como um caso especial de condicionamento de Jeffrey, unificando essas abordagens sob um framework teórico mais amplo.
• Expressividade: Permite a especificação de distribuições alvo explícitas para objetivos complexos (como independência de atributos) que antes só eram alcançáveis através de métodos heurísticos ou implícitos.
• Praticidade: O método não requer retreinamento do modelo de difusão, exigindo apenas a adição de um termo de correção durante a amostragem e o treinamento de um estimador de razão de densidade leve.

Os autores mantêm um tom modesto, enquadrando o trabalho como uma "prova de conceito". Eles reconhecem que a implementação atual depende de um classificador para estimar razões de densidade, de forma semelhante ao modo como a orientação por classificador precoce funcionava, e expressam esperança de que trabalhos futuros evoluam para variantes sem classificador (classifier-free). Eles também observam que, embora o método ofereça novo poder expressivo para aplicações benéficas (equidade, mitigação de memorização), ele também carrega o risco de capacitar atores maliciosos a gerar outputs nefastos.