Variance-Tilted Diffusion Models for Diverse Sampling

Resumo Técnico: Modelos de Difusão com Inclinação de Variância para Amostragem Diversa

Problema
Os modelos de difusão tornaram-se o padrão para modelagem generativa em diversos domínios, que vão desde a síntese de imagens de alta resolução até a descoberta científica (ex: dobramento de proteínas e acoplamento molecular). No entanto, o paradigma de amostragem padrão baseia-se na extração de amostras independentes e identicamente distribuídas (i.i.d.). Esta abordagem é frequentemente subótima para tarefas de jusante (downstream) que exigem um conjunto diversificado de candidatos que cubram regiões distintas da variedade (manifold) de dados, como a geração de ligantes estruturalmente diversos no design de fármacos ou propostas de trajetórias robustas no planejamento de movimento robótico.

Métodos existentes para impor diversidade, como a orientação por partículas (particle guidance), introduzem potenciais repulsivos na dinâmica de amostragem. Embora eficazes, esses métodos modificam o drift de forma heurística, resultando em uma distribuição alvo implícita que é computacionalmente intratável de avaliar ou controlar com precisão.

Metodologia
Os autores propõem uma formulação "alvo-primeiro" (target-first), onde a distribuição alvo do lote é explicitamente definida antes da derivação do amostrador exato. A inovação central é uma distribuição de lote ponderada pela variância que favorece coleções de amostras com grande dispersão empírica após um mapeamento de características lineares prescrito.

Para um lote de $n$ amostras $\mathbf{x} = (x^{(1)}, \dots, x^{(n)})$, a distribuição alvo $\pi^A_0(\mathbf{x})$ é definida como:
$$\pi^A_0(\mathbf{x}) \propto \text{Var}^A_n(\mathbf{x}) \prod_{i=1}^n p_0(x^{(i)})$$
onde $p_0$ é a distribuição original dos dados, $A: \mathbb{R}^d \to \mathbb{R}^k$ é um operador linear que mapeia as amostras para um espaço de características específico da tarefa, e $\text{Var}^A_n(\mathbf{x})$ é a variância empírica das características $A(x^{(i)})$.

Para amostrar desta distribuição, os autores derivam a transformação de Doob $h$ exata da dinâmica de difusão independente. Isso envolve calcular a esperança condicional da variância do lote dado o estado ruidoso no tempo $t$:
$$h^A_t(\mathbf{x}_t) := \mathbb{E}[\text{Var}^A_n(\mathbf{X}_0) \mid \mathbf{X}_t = \mathbf{x}_t]$$

Utilizando as identidades de Tweedie e uma decomposição de viés-variância, os autores derivam uma expressão de forma fechada para a correção do score $\nabla_{\mathbf{x}_t} \log h^A_t(\mathbf{x}_t)$. Esta correção consiste em dois termos distintos:

1. Termo de Interação (Repulsão): Um termo que repele as médias posteriores de denoising ($\mu_t$) das partículas umas das outras no espaço de características. Isso é computado eficientemente usando produtos de vetor-Jacobiano (VJP) através do denoiser.
2. Termo de Curvatura: Um termo que move as partículas em direção a regiões de maior variância de características, envolvendo a divergência do score (especificamente $\text{div}(A^\top A \nabla s_t)$). Este é estimado via diferenças finitas ou estimadores de traço estocástico (ex: estimadores de Hutchinson).

O amostrador resultante é um processo de partículas interagentes onde o score de cada partícula é a soma do score padrão e da correção de Doob derivada.

Principais Contribuições

1. Formulação Alvo-Primeiro: O artigo enquadra a geração diversa como um problema de amostragem exata de uma distribuição alvo de lote explícita, evitando as forças repulsivas não controladas de métodos heurísticos.
2. Correção de Doob em Forma Fechada: Os autores fornecem uma derivação analítica da correção de score para um operador linear geral $A$, decompondo-a em um termo de repulsão para as médias posteriores e um termo de curvatura para a variância de características.
3. Algoritmo de Amostragem Tratável: Um amostrador de partículas interagentes discretizado é introduzido. O principal termo de interação aproveita computações eficientes de VJP, enquanto o termo de curvatura é tratado via estimativa numérica, tornando o método prático para modelos de difusão existentes.

Resultos
Os autores avaliaram a abordagem na geração de imagens usando Stable Diffusion v1.5 (512x512) com vários prompts. Eles compararam três configurações:

• Orientação de Classificador Livre Padrão (CFG - Classifier-Free Guidance).
• Inclinação de Variância (VT - Variance-Tilted) sem o termo de divergência.
• Variância Inclinada + Divergência (VT + Divergência).

Resultados qualitativos (Figuras 1–5) demonstram que os métodos propostos produzem um conjunto de amostras mais diverso em comparação ao CFG padrão. As imagens geradas cobrem uma gama mais ampla de poses, fundos e estilos visuais, em vez de produzir apenas variações menores de modos de alta probabilidade. A inclusão do termo de divergência refina ainda mais a diversidade, embora os autores notem que o método completo pode impactar levemente negativamente a qualidade da imagem e aumentar o custo computacional.

Significância e Alegações
O artigo afirma oferecer uma alternativa probabilística à repulsão de partículas heurística. Ao derivar o amostrador exato para uma distribuição alvo ponderada pela variância, o método fornece um objetivo probabilístico transparente em vez de um implícito. Os autores posicionam isso como uma abordagem rigorosa para amostragem diversa que mantém caminhos práticos de implementação.

Os autores mantêm-se modestos quanto às limitações, reconhecendo que, embora o método melhore qualitativamente a diversidade, ele pode degradar levemente a qualidade da imagem e incorrer em custos computacionais mais altos devido à estimativa de derivadas de segunda ordem (divergência). Eles sugerem que trabalhos futuros poderiam focar no recozimento (annealing) da correção de Doob para mitigar problemas de qualidade e no desenvolvimento de estimadores mais eficientes para o termo de curvatura.