Beyond Generative Priors: Minority Sampling with JEPA-Guided Diffusion

Aquí se presenta un resumen técnico detallado del artículo "Beyond Generative Priors: Minority Sampling with JEPA-Guided Diffusion".

Planteamiento del Problema

El muestreo de minorías tiene como objetivo generar instancias de baja densidad en una variedad de datos, lo cual es crítico para aplicaciones como el diagnóstico médico, la detección de anomalías y la IA creativa. Sin embargo, los enfoques existentes definen "minoría" en relación con el prior generativo aprendido por el propio modelo (es decir, la distribución de los datos de entrenamiento). Esto crea una limitación fundamental: las muestras identificadas como "raras" a menudo son meramente artefactos del conjunto de entrenamiento específico o de los sesgos inductivos del modelo, en lugar de instancias genuinamente atípicas en el mundo real. Por ejemplo, un modelo entrenado en conjuntos de datos estándar podría considerar un perro sobre un fondo blanco como "raro" simplemente porque está subrepresentado en los datos de entrenamiento, mientras que un avión furtivo podría ser genuinamente raro en el mundo real pero no capturarse como tal si el prior del modelo es estrecho. El artículo argumenta que los métodos actuales no logran alinearse con los priors del mundo real, que son necesarios para capturar la rareza semántica verdadera.

Metodología: Guía JEPA

Los autores proponen la Guía JEPA, un marco que desplaza el muestreo de minorías de una perspectiva centrada en el generador a una centrada en el mundo. Esto se logra guiando los modelos de difusión utilizando la densidad implícita de una Arquitectura de Predicción de Incrustación Conjunta (JEPA), una clase de modelos del mundo (por ejemplo, DINOv2) conocida por codificar representaciones amplias y semánticamente ricas.

Mecanismo Central:

1. Definición Centrada en el Mundo: En lugar de minimizar la probabilidad bajo la densidad del modelo de difusión $p_\theta(x)$, el método busca minimizar la probabilidad bajo la densidad inducida por JEPA $p_\phi(x)$. Las muestras de minoría se definen como aquellas que residen en regiones de baja densidad del espacio de representación de JEPA.
2. JEPA-SCORE: El método utiliza JEPA-SCORE, una medida de densidad derivada del Jacobiano de un codificador JEPA preentrenado $f_\phi$. Específicamente, es la suma de los logaritmos de los valores singulares del Jacobiano $J_f(x)$. Los valores más bajos de JEPA-SCORE indican menor densidad (mayor atipicidad) en el prior del mundo.
3. Aproximación mediante SVD Aleatorizado: Calcular el Jacobiano completo y su Descomposición en Valores Singulares (SVD) en cada paso de difusión es computacionalmente prohibitivo. Los autores emplean SVD Aleatorizado para aproximar eficientemente los $k$ valores singulares superiores del Jacobiano comprimido. Proporcionan una cota superior teórica sobre el error de aproximación, descomponiéndolo en el error de estimación del SVD aleatorizado y el error de truncamiento.
4. Teorema de la Envoltura para la Eficiencia: Para evitar la retropropagación a través del bucle de optimización interno del SVD aleatorizado (que depende de la entrada $x_t$), los autores invocan el Teorema de la Envoltura. Esto les permite tratar la matriz de proyección óptima $Q^*$ como una constante durante el cálculo del gradiente, reduciendo significativamente el uso de memoria y la sobrecarga computacional mientras se preservan los gradientes de primer orden correctos.
5. Guía Diferida: Para abordar la brecha de dominio entre los latentes intermedios ruidosos del modelo de difusión y las entradas limpias esperadas por el codificador JEPA, los autores introducen la guía diferida. La guía JEPA solo se aplica después de un paso de tiempo intermedio $\tau T$ (por ejemplo, $t=0.8T$), asegurando que las estimaciones desruidadas $\hat{x}_{0|t}$ sean lo suficientemente claras para que el codificador JEPA las evalúe eficazmente. Esto también permite que el método se extienda a la generación condicional (por ejemplo, texto a imagen), ya que la condición se establece temprano en el proceso antes de que comience la guía centrada en el mundo.

Contribuciones Clave

• Muestreo de Minorías Centrada en el Mundo: Un cambio de paradigma que redefine las muestras de minoría basándose en priors del mundo real (a través de JEPAs) en lugar de los priors estrechos del modelo generativo.
• Marco de Guía JEPA: Un método de muestreo de difusión principista que dirige las trayectorias hacia regiones de baja densidad bajo priors inducidos por JEPA.
• Cálculo Eficiente: El desarrollo de estrategias de aproximación principistas (SVD Aleatorizado) y la aplicación del Teorema de la Envoltura para hacer que la guía JEPA sea computacionalmente práctica, acompañada de cotas de error teóricas.
• Integración sin Entrenamiento: El enfoque es completamente libre de entrenamiento, requiriendo solo modelos de difusión y JEPA preentrenados, y puede integrarse en el momento de la inferencia sin modificar los modelos subyacentes.

Resultados Experimentales

Los autores evaluaron la guía JEPA en tareas de generación incondicional, condicional por clase y de texto a imagen (utilizando CelebA, ImageNet, Stable Diffusion 1.5 y SDXL-Lightning).

• Atipicidad: La guía JEPA logra consistentemente el JEPA-SCORE más bajo (indicando mayor atipicidad bajo priores del mundo) en comparación con líneas base como Sehwag et al. (2022), MG, SGMS y MinorityPrompt.
• Fidelidad y Validez: A pesar de apuntar a regiones raras, el método mantiene una alta fidelidad. En tareas de texto a imagen, logra puntuaciones competitivas de alineación con el texto (CLIPScore, PickScore, ImageReward) mientras genera muestras semánticamente distintas (por ejemplo, personal militar anciano, jugadores de béisbol afroamericanos) que otros métodos no logran producir.
• Comparación Cualitativa: Las visualizaciones muestran que los métodos centrados en el generador a menudo producen muestras con contextos atípicos (por ejemplo, un águila cerca de una grúa), whereas la guía JEPA produce muestras con semántica o comportamientos atípicos (por ejemplo, águilas apareándose, formas animales inusuales), alineándose mejor con la rareza del mundo real.
• Utilidad en Tareas Posteriores: En una tarea de aumento de datos para el entrenamiento de clasificadores, el aumento con muestras de minoría guiadas por JEPA arrojó un rendimiento superior o comparable al de otros métodos utilizando menos muestras, lo que sugiere que estas muestras proporcionan señales de entrenamiento más informativas.

Significado y Afirmaciones

El artículo afirma que la guía JEPA ofrece una solución robusta a la limitación de que los muestreadores de minorías existentes estén atrapados dentro de la propia distribución del modelo generativo. Al aprovechar los modelos del mundo (JEPAs), el método genera muestras que son genuinamente raras en el mundo real, no solo raras en relación con un conjunto de datos específico.

Los autores enfatizan que este enfoque:

1. Desacopla la noción de minoría del prior generativo.
2. Cierra la brecha entre el modelado generativo y el modelado del mundo, un componente clave en la búsqueda de la inteligencia artificial general.
3. Proporciona una herramienta práctica y libre de entrenamiento para aplicaciones que requieren datos diversos y atípicos, como pruebas de robustez, mejora de la equidad y generación de contenido creativo.

Los autores también reconocen una posible preocupación de doble uso: invertir la dirección de la guía podría reforzar los sesgos existentes en los priores del mundo, y alientan un uso responsable. Señalan que, aunque el método introduce una sobrecarga computacional en comparación con los muestreadores estándar, sigue siendo más rápido que algunas líneas base complejas (como MinorityPrompt) y ofrece un compromiso razonable para la ganancia significativa en atipicidad centrada en el mundo.