Variance-Tilted Diffusion Models for Diverse Sampling

Resumen Técnico: Modelos de Difusión con Sesgo de Varianza para un Muestreo Diverso

Planteamiento del Problema
Los modelos de difusión se han convertido en el estándar para la modelación generativa en diversos dominios, que van desde la síntesis de imágenes de alta resolución hasta el descubrimiento científico (por ejemplo, el plegamiento de proteínas y el acoplamiento molecular). Sin embargo, el paradigma de muestreo estándar se basa en extraer muestras independientes e idénticamente distribuidas (i.i.d.). Este enfoque suele ser subóptimo para tareas posteriores que requieren un conjunto diverso de candidatos que cubran regiones distintas del manifold de datos, como la generación de ligandos estructuralmente diversos en el diseño de fármacos o la propuesta de trayectorias robustas en la planificación de movimiento robótico.

Los métodos existentes para imponer diversidad, como la guía de partículas, introducen potenciales repulsivos en la dinámica de muestreo. Aunque son efectivos, estos métodos modifican el drift de forma heurística, lo que resulta en una distribución objetivo implícita que es computacionalmente intratable de evaluar o controlar con precisión.

Metodología
Los autores proponen una formulación de "objetivo primero" (target-first) donde la distribución objetivo del lote se define explícitamente antes de derivar el muestreador exacto. La innovación central es una distribución de lote ponderada por la varianza que favorece colecciones de muestras con una gran dispersión empírica tras un mapa de características lineales prescrito.

Para un lote de $n$ muestras $\mathbf{x} = (x^{(1)}, \dots, x^{(n)})$, la distribución objetivo $\pi^A_0(\mathbf{x})$ se define como:
$$\pi^A_0(\mathbf{x}) \propto \text{Var}^A_n(\mathbf{x}) \prod_{i=1}^n p_0(x^{(i)})$$
donde $p_0$ es la distribución de datos original, $A: \mathbb{R}^d \to \mathbb{R}^k$ es un operador lineal que mapea las muestras a un espacio de características específico de la tarea, y $\text{Var}^A_n(\mathbf{x})$ es la varianza empírica de las características $A(x^{(i)})$.

Para muestrear de esta distribución, los autores derivan la transformación $h$ de Doob exacta de la dinámica de difusión independiente. Esto implica calcular la esperanza condicional de la varianza del lote dado el estado ruidoso en el tiempo $t$:
$$h^A_t(\mathbf{x}_t) := \mathbb{E}[\text{Var}^A_n(\mathbf{X}_0) \mid \mathbf{X}_t = \mathbf{x}_t]$$

Utilizando las identidades de Tweedie y una descomposición de sesgo-varianza, los autores derivan una expresión de forma cerrada para la corrección del score $\nabla_{\mathbf{x}_t} \log h^A_t(\mathbf{x}_t)$. Esta corrección consta de dos términos distintos:

1. Término de Interacción (Repulsión): Un término que repele las medias posteriores de denoisado ($\mu_t$) de las partículas entre sí en el espacio de características. Esto se computa eficientemente utilizando productos de vector-Jacobiano (VJP) a través del denoiser.
2. Término de Curvatura: Un término que mueve las partículas hacia regiones de mayor varianza de características, el cual involucra la divergencia del score (específicamente $\text{div}(A^\top A \nabla s_t)$). Este se estima mediante diferencias finitas o estimadores de traza estocásticos (por ejemplo, estimadores de Hutchinson).

El muestreador resultante es un proceso de partículas interactuantes donde el score de cada partícula es la suma del score estándar y la corrección de Doob derivada.

Contribuciones Clave

1. Formulación de Objetivo Primero: El artículo plantea la generación diversa como un problema de muestreo exacto de una distribución objetivo de lote explícita, evitando las fuerzas repulsivas no controladas de los métodos heurísticos.
2. Corrección de Doob de Forma Cerrada: Los autores proporcionan una derivación analítica de la corrección del score para un operador lineal general $A$, descomponiéndola en un término de repulsión para las medias posteriores y un término de curvatura para la varianza de las características.
3. Algoritmo de Muestreo Tratable: Se introduce un muestreador de partículas interactuantes discretizado. El término de interacción principal aprovecha las computaciones eficientes de VJP, mientras que el término de curvatura se gestiona mediante estimación numérica, haciendo que el método sea práctico para los modelos de difusión existentes.

Resultados
Los autores evaluaron el enfoque en la generación de imágenes utilizando Stable Diffusion v1.5 (512x512) con diversos prompts. Compararon tres configuraciones:

• Guía de Clasificador Libre (CFG) estándar.
• Sesgo de Varianza (VT) sin el término de divergencia.
• Sesgo de Varianza + Divergencia (VT + Divergencia).

Los resultados cualitativos (Figuras 1–5) demuestran que los métodos propuestos producen un conjunto de muestras más diverso en comparación con el CFG estándar. Las imágenes generadas cubren un rango más amplio de poses, fondos y estilos visuales, en lugar de producir meras variaciones menores de los modos de alta probabilidad. La inclusión del término de divergencia refina aún más la diversidad, aunque los autores señalan que el método completo puede afectar ligeramente la calidad de la imagen y aumentar el costo computacional.

Significación y Reivindicaciones
El artículo afirma ofrecer una alternativa probabilística a la repulsión de partículas heurística. Al derivar el muestreador exacto para una distribución objetivo ponderada por la varianza, el método proporciona un objetivo probabilístico transparente en lugar de uno implícito. Los autores posicionan esto como un enfoque riguroso para el muestreo diverso que mantiene vías de implementación prácticas.

Los autores mantienen la modestia respecto a las limitaciones, reconociendo que, si bien el método mejora cualitativamente la diversidad, puede degradar ligeramente la calidad de la imagen e incurrir en mayores costos computacionales debido a la estimación de derivadas de segundo orden (divergencia). Sugieren que el trabajo futuro podría centrarse en el enfriamiento (annealing) de la corrección de Doob para mitigar los problemas de calidad y en el desarrollo de estimadores más eficientes para el término de curvatura.