MMD Guidance: Training-Free Distribution Adaptation for Diffusion Models via Maximum Mean Discrepancy Guidance

Resumen Técnico: Guía MMD para la Adaptación de Distribución sin Entrenamiento

Planteamiento del Problema
Los modelos de difusión preentrenados sirven como potentes prioris generativos para la generación de muestras incondicionales y condicionales. Sin embargo, sus salidas suelen desviarse de las distribuciones objetivo específicas del usuario, un problema crítico en tareas de adaptación de dominio donde solo se dispone de un número limitado de ejemplos de referencia (por ejemplo, 50–100). El reentrenamiento o el ajuste fino (fine-tuning) de estos modelos suele ser computacionalmente inviable o poco práctico debido a las restricciones de memoria y a la falta de grandes conjuntos de datos objetivo. Los métodos de guía existentes en tiempo de inferencia, como la guía de clasificador (classifier guidance) o la guía libre de clasificador (classifier-free guidance), ajustan las trayectorias de muestreo optimizando objetivos subrogados (por ejemplo, probabilidades de clase o fuerza de condicionamiento) en lugar de alinear directamente la distribución generada con una distribución de referencia objetivo. En consecuencia, estos métodos no pueden hacer coincidir explícitamente una distribución objetivo arbitraria definida por un pequeño conjunto de muestras de referencia.

Metodología
El artículo propone la Guía MMD, un mecanismo sin entrenamiento que aumenta el proceso de difusión inversa con los gradientes de la Discrepancia de la Media Máxima (MMD, por sus siglas en inglés) entre las muestras generadas y un conjunto de datos de referencia.

1. Mecanismo Central: El método modifica la actualización estándar del muestreo de difusión inversa en cada paso de tiempo $t$ restando el gradiente del MMD empírico al cuadrado con respecto al estado de la muestra actual. La regla de actualización se formula como:
   $$z_{t-1} = \text{sampler}(z_t, t, \epsilon_\theta) - \lambda_t \nabla_{z_t} \widehat{\text{MMD}}^2(\hat{P}_t, \hat{Q})$$
   donde $\hat{P}_t$ es la distribución empírica del lote generado en el paso de tiempo $t$, $\hat{Q}$ es la distribución empírica de las muestras de referencia, y $\lambda_t$ controla la fuerza de la guía.

2. Elección de la Divergencia: Los autores seleccionan MMD sobre las divergencias-$f$ (como KL) y las distancias de Wasserstein porque MMD proporciona estimaciones insesgadas y de baja varianza a partir de datos limitados sin sufrir la maldición de la dimensionalidad. Además, su formulación basada en kernels permite un cálculo de gradiente eficiente y en forma cerrada.

3. Implementación en el Espacio Latente: Para asegurar la eficiencia computacional, la guía se aplica directamente en el espacio latente de los Modelos de Difusión Latente (LDM). Las imágenes de referencia se codifican una sola vez en representaciones latentes utilizando el codificador VAE. Esto reduce la dimensionalidad del cálculo de MMD y mejora la relación señal-ruido al filtrar el ruido a nivel de píxel.

4. Adaptación Consciente del Prompt: Para la generación condicional, el marco se extiende a la adaptación consciente del prompt mediante un kernel de producto. Este kernel, $k_{\otimes}([p, z], [p', z']) = k_{\text{text}}(p, p') \cdot k_{\text{image}}(z, z')$, mide conjuntamente la similitud tanto en el espacio semántico (prompt) como en el espacio visual (latente). Esto permite que la guía alinee las muestras generadas tanto con el prompt textual como con el estilo visual de los pares de referencia sin reentrenar el denoiser.

5. Dinámica de Gradientes: El gradiente del MMD al cuadrado consiste en dos términos: un término intra-lote que crea repulsión entre las muestras generadas (promoviendo la diversidad) y un término transversal que crea atracción hacia las muestras de referencia (fomentando el ajuste de la distribución). El análisis teórico proporciona garantías de concentración para el gradiente del término transversal, asegurando que estime de manera fiable el gradiente de la población incluso con muestras de referencia limitadas.

Contribuciones Clave

• Guía MMD: Introducción de un método sin entrenamiento y en tiempo de inferencia para adaptar modelos de difusión a distribuciones objetivo especificadas por conjuntos de referencia limitados.
• Extensión Consciente del Prompt: Una extensión práctica utilizando kernels de producto para permitir el alineamiento conjunto con prompts de texto y estilos visuales en la generación condicional.
• Implementación Eficiente: Una implementación escalable que opera directamente en el espacio latente de los LDMs, reduciendo significativamente la carga computacional en comparación con la guía en el espacio de píxeles.
• Validación Empírica: Experimentos exhaustivos que demuestran que la Guía MMD logra un alineamiento distributivo superior preservando la fidelidad de la muestra.

Resultados Experimentales
Los autores evaluaron la Guía MMD en Modelos de Mezcla Gaussiana (GMM) sintéticos y en benchmarks de imágenes del mundo real (FFHQ, CelebA-HQ y subconjuntos de estilos específicos de salidas de Stable Diffusion).

• Datos Sintéticos: En GMMs de 8 y 25 modalidades, la Guía MMD recuperó con éxito las proporciones de la mezcla objetivo y corrigió distribuciones desequilibradas, superando a las líneas base en Distancia de Fréchet (FD) y Distancia de Kernel (KD).
• Imágenes del Mundo Real: En experimentos con FFHQ (por ejemplo, generar rostros con estilos específicos de gafas) y CelebA-HQ, la Guía MMD alcanzó las puntuaciones más altas de fidelidad y cobertura distributiva en comparación con el muestreo sin guía, la guía de clasificador e incluso las líneas base entrenadas por el usuario que requerían ajuste fino.
• Generación Condicional: En entornos condicionados por texto (usando Stable Diffusion XL y PixArt), el método adaptó con éxito el estilo visual de las imágenes generadas (por ejemplo, "gatos en estilo anime" o "coches en estilo Pixar") para que coincidieran con los conjuntos de referencia mientras se adhería a los prompts de texto.
• Eficiencia: El análisis de tiempo de ejecución indica que el coste computacional adicional de la guía MMD en el espacio latente es insignificante en comparación con el muestreo sin guía, incluso para modelos a gran escala como SDXL.
• Robustez: El método demostró robustez ante la variación en el número de muestras de referencia (convergiendo eficazmente con tan solo 50–100 muestras) y diferentes brechas de dominio.

Significancia y Reivindicaciones
El artículo afirma que la Guía MMD ofrece una solución práctica para desplegar modelos generativos en escenarios donde los recursos computacionales para el reentrenamiento son limitados o donde los datos específicos del usuario son escasos. Al operar enteramente en tiempo de inferencia sin modificar los parámetros del modelo, el método permite un despliegue inmediato y preserva la privacidad (ya que no se utilizan datos del usuario para el entrenamiento). Los autores enfatizan que el enfoque cierra eficazmente la brecha entre las capacidades de los modelos fundacionales preentrenados y los requisitos específicos del usuario, logrando el alineamiento distributivo mientras mantiene la diversidad y la calidad de las muestras generadas. El trabajo posiciona a la Guía MMD como una herramienta complementaria a las técnicas existentes de ajuste fino e ingeniería de prompts, particularmente valiosa en entornos con restricciones de privacidad o de recursos.