Resumen Técnico: Hacia un control más general de los modelos de difusión mediante la Guía de Jeffrey

1. Planteamiento del problema

Los modelos de difusión ofrecen una flexibilidad significativa en las tareas generativas, permitiendo el control de la salida durante el muestreo mediante mecanismos de guía sin necesidad de reentrenar la red neuronal subyacente. Si bien las técnicas de guía estándar (por ejemplo, la guía de clasificador) gestionan eficazmente el muestreo condicional para clases específicas, a menudo fallan al abordar objetivos más complejos donde la distribución objetivo es implícita o se define únicamente mediante heurísticas.

Los enfoques actuales para tareas como la equidad (fairness), el control de conceptos y la generación compositiva dependen frecuentemente de modificaciones heurísticas de la dinámica de muestreo. Estos métodos inclinan la distribución original, pero carecen de una distribución objetivo explícitamente definida, lo que hace que el objetivo real sea difícil de caracterizar o interpretar. Existe la necesidad de un marco conceptual que pueda actualizar las distribuciones marginales hacia una distribución objetivo prescrita preservando la estructura condicional del modelo, permitiendo así el control sobre propiedades estadísticas complejas (por ejemplo, la independencia de atributos o el ajuste de la distribución de embeddings) que el condicionamiento estándar no puede expresar.

2. Metodología: Guía de Jeffrey

El artículo propone la guía de Jeffrey, un marco que aprovecha la regla de condicionamiento de Jeffrey para actualizar los modelos de difusión. A diferencia de los enfoques basados en Bayes que condicionan sobre eventos específicos (condicionamiento estándar), la regla de Jeffrey permite condicionar sobre predicados complejos, tales como el ajuste de una distribución marginal específica.

Fundamento Teórico

Dado un modelo conjunto $p(x, y)$ donde $x$ representa el espacio de la imagen e $y$ representa los atributos/características, el objetivo es actualizar la marginal $p(y)$ hacia una distribución objetivo $p^\star(y)$ preservando la estructura condicional $p(x|y)$. Bajo la regla de Jeffrey, la distribución conjunta actualizada se define como:
$$p^\star(x, y) = p(x|y)p^\star(y)$$
Esta actualización minimiza la perturbación de la distribución conjunta original (una proyección de información) mientras satisface la restricción marginal.

El artículo deriva una formulación de reponderación para la marginal sobre $x$:
$$p^\star(x) = p(x) \int p(y|x)\rho(y) dy$$
donde $\rho(y) = p^\star(y)/p(y)$ es la razón de densidad. En los casos donde las características se derivan de forma determinista de las imágenes ($y = u(x)$), esto se simplifica a:
$$p^\star(x) = p(x)\rho(u(x))$$

Implementación en Modelos de Difusión

Para aplicar esto a los modelos de difusión basados en puntuación (score-based), los autores derivan una función de puntuación guiada. La puntuación estándar $\nabla_{x_t} \log p_t(x_t)$ se aumenta con un término de corrección basado en la razón de densidad:
$$\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$$
donde $\hat{x}_0(x_t, t)$ es una estimación de la muestra limpia obtenida mediante la fórmula de Tweedie.

Para mapeos de características deterministas (por ejemplo, espacios de embedding), la guía se simplifica a:
$$\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)))$$

El marco introduce un hiperparámetro $\lambda$ para escalar el término de corrección, controlando el compromiso entre la calidad de la muestra y la fidelidad a la distribución objetivo. La razón de densidad $\rho$ se estima utilizando un clasificador binario entrenado para distinguir entre muestras de la distribución objetivo y la distribución actual del modelo.

3. Principales Contribuciones

El artículo realiza tres contribuciones primarias:

1. Marco de Guía de Jeffrey: Un método de guía novedoso y "plug-and-play" que generaliza la guía de clasificador estándar. Utiliza la regla de Jeffrey para ajustar las distribuciones de salida generadas a distribuciones objetivo arbitrarias, permitiendo el control sobre objetivos que antes eran inaccesibles para la guía estándar.
2. Ajuste de la Distribución de Embeddings: Una prueba de concepto que demuestra la capacidad de ajustar las distribuciones de salida del modelo a una distribución objetivo en un espacio de embedding (específicamente, embeddings de Inception). Esta aplicación reduce significativamente la Distancia de Fréchet Inception (FID) sin requerir reentrenamiento.
3. Control de la Distribución de Atributos para la Equidad: Una demostración del control de las distribuciones de atributos de las imágenes para lograr objetivos de equidad. El método logra con éxito la paridad de género (ajuste de distribución) y desacopla atributos (por ejemplo, "Masculino" y "Joven") en rostos generados, tareas que la guía de clasificador estándar no puede abordar directamente.

4. Resultados Experimentales

Los autores evaluaron la guía de Jeffrey en los conjuntos de datos CIFAR-10, FFHQ y CelebA-HQ.

Guía de Embedding (Reducción de FID)

• Configuración: El método tuvo como objetivo la distribución de embeddings de Inception de los datos de entrenamiento.
• Resultados: En CIFAR-10, el FID mejoró de $\approx 5.88$ a $\approx 2.55$ (conjunto de entrenamiento) y de $\approx 9.76$ a $\approx 6.71$ (conjunto de prueba). En FFHQ, el FID bajó de $\approx 20.13$ a $\approx 12.91$.
• Observación: Los autores señalan que, aunque el FID disminuyó drásticamente, los cambios visuales en las imágenes generadas fueron mínimos, lo que sugiere limitaciones en el FID como métrica puramente perceptual. Los mejores resultados se obtuvieron aplicando la guía solo en los pasos finales de eliminación de ruido ($\delta=10$).

Aplicaciones de Equidad (Fairness)

• Paridad de Género (Ajuste de Distribución): En CelebA-HQ, el método ajustó la proporción de atributos "Masculino" de $\approx 0.296$ a $\approx 0.522$ (apuntando a 0.5), reduciendo la Discrepancia de Equidad (FD) de $0.288$a$0.032$. Esto se logró con una degradación mínima en el FID ($21.61$ frente al $22.30$ de la línea base).
• Desacoplamiento de Atributos: El método desacopló con éxito los atributos "Masculino" y "Joven", reduciendo el coeficiente de correlación de Pearson de $-0.541$ a casi cero ($0.005$) con $\lambda=3.0$. Esto se logró reduciendo el modo conjunto dominante (muestras de mujeres jóvenes) mientras se mantenían las distribuciones marginales relativamente estables, con un ligero aumento del FID de aproximadamente un punto.
• Comparación: En comparación con una línea base de muestreo ancestral combinada con guía estándar, la guía de Jeffrey logró métricas de equidad similares con costos de FID más bajos y modificó menos muestras, ya que equilibra las direcciones de clase de forma continua a lo largo de la trayectoria de muestreo en lugar de empujar las muestras enteras hacia una sola clase.

5. Significado y Reivindicaciones

El artículo posiciona la guía de Jeffrey como una extensión fundamentada del control de los modelos de difusión. Su importancia radica en:

• Generalización: Recontextualiza la guía de clasificador estándar como un caso especial del condicionamiento de Jeffrey, unificando así estos enfoques bajo un marco teórico más amplio.
• Expresividad: Permite la especificación de distribuciones objetivo explícitas para objetivos complejos (como la independencia de atributos) que anteriormente solo eran alcanzables mediante métodos heurísticos o implícitos.
• Practicidad: El método no requiere reentrenamiento del modelo de difusión, solo la adición de un término de corrección durante el muestreo y el entrenamiento de un estimador de razón de densidad ligero.

Los autores mantienen un tono modesto, calificando su trabajo como una "prueba de concepto". Reconocen que la implementación actual depende de un clasificador para estimar las razones de densidad, de manera similar a cómo funcionaba la guía de clasificador temprana, y expresan la esperanza de que el trabajo futuro evolucione hacia variantes libres de clasificador (classifier-free). También advierten que, si bien el método ofrece un nuevo poder expresivo para aplicaciones beneficiosas (equidad, mitigación de la memorización), también conlleva el riesgo de facultar a actores malintencionados para generar contenidos nefastos.