DIAMOND: Directed Inference for Artifact Mitigation in Flow Matching Models

Resumen Técnico: DIAMOND – Inferencia Dirigida para la Mitigación de Artefactos en Modelos de Flow Matching

Declaración del Problema
A pesar de la alta fidelidad de los modelos recientes de texto a imagen, particularmente aquellos basados en arquitectes de Rectified Flow (p. ej., FLUX.2) y Modelos de Difusión (p. ej., SDXL), las imágenes generadas presentan frecuentemente artefactos estructurales y anatómicos. Estos incluyen manos distorsionadas, extremidades adicionales, rostros malformados y texto incorrecto. Las estrategias de mitigación existentes sufren limitaciones significativas:

• Refinamiento post-hoc: Los métodos que corrigen las imágenes después de la generación a menudo no logran intervenir durante el proceso central de formación.
• Modificaciones invasivas: Las técnicas que requieren el ajuste fino (fine-tuning) del modelo o pesos de LoRA son computacionalmente costosas, difíciles de escalar y carecen de pesos disponibles públicamente para arquitecturas modernas.
• Limitaciones heurísticas: Los métodos de detección en tiempo real suelen depender de umbrales heurísticos o reemplazos aleatorios de latentes, lo que conduce a una localización inexacta o correcciones subóptimas.

Metodología
Los autores proponen DIAMOND (Inferencia Dirigida para la Mitigación de Artefactos en Modelos de Flow Matching), un marco de trabajo zero-shot y libre de entrenamiento que corrige artefactos durante la trayectoria de inferencia sin modificar los pesos del modelo. El método opera bajo el principio de reconstruir una estimación de la muestra limpia en cada paso de tiempo para guiar la generación lejos de los estados latentes propensos a artefactos.

1. Estimación de Latente Limpio:
   A diferencia de los enfoques estándar que operan directamente sobre latentes ruidosos ($x_t$), DIAMOND estima el latente de los datos limpios ($\hat{x}_{0,t}$) en cada paso.

   • Para modelos de Rectified Flow (que predicen la velocidad $v_\theta$): $\hat{x}_{0,t} = x_t - t \cdot v_\theta(x_t, t)$.
   • Para modelos de Difusión (que predicen el ruido $\epsilon_\theta$): $\hat{x}_{0,t} = x_t - \sigma_t \cdot \epsilon_\theta(x_t, t)$.
     Este latente estimado se decodifica al espacio de píxeles ($\hat{I}_t = D(\hat{x}_{0,t})$) para asegurar que el detector de artefactos opere en su dominio óptimo (datos limpios) en lugar de en estados intermedios ruidosos.

2. Detección de Artefactos y Pérdida:
   Un Detector de Artefactos (AD) diferenciable y pre-entrenado (p. ej., de DiffDoctor) analiza $\hat{I}_t$ para generar una máscara de probabilidad de artefactos $M_t$. Se calcula una pérdida de artefacto ($L_a$) como el error cuadrático medio entre esta máscara y una máscara de cero (que representa una imagen libre de artefactos):
   $$L_a = \frac{1}{HW} \sum_{i,j} (M_{i,j}^t)^2$$

3. Corrección de Trayectoria Basada en Gradiente:
   El método calcula el gradiente de la pérdida con respecto al latente actual $x_t$ ($\nabla_{x_t} L_a$). Para asegurar la estabilidad y evitar que la magnitud de la corrección sea dominada por el tamaño de la región del artefacto, el gradiente es normalizado:
   $$\delta_t = \lambda_t \frac{\nabla_{x_t} L_a}{\|\nabla_{x_t} L_a\|^2 + \epsilon}$$
   Este vector de desplazamiento $\delta_t$ se resta del paso de actualización del solver estándar, desviando efectivamente la trayectoria lejos de los artefactos.

4. Programación Dinámica ($\lambda_t$):
   Un esquema de ley de potencia controla la intensidad de la corrección ($\lambda_t$). Basándose en la observación de que los artefactos estructurales aparecen en las etapas tempranas, dominadas por el ruido, mientras que los detalles de alta frecuencia se refinan más tarde, el esquema aplica una corrección fuerte al inicio del proceso de generación y decae hacia el final.

5. Preservación de Identidad (Opcional):
   Para evitar que la corrección altere la estructura global de la imagen o la alineación con el prompt, se puede aplicar un término de regularización opcional ($L_{rec}$). Este término minimiza el Error Absoluto Medio (MAE) entre la imagen corregida y la imagen base en las regiones que no contienen artefactos.

Contribuciones Clave

• Marco Libre de Entrenamiento: DIAMOND es el primer método zero-shot que integra la guía geométrica de la trayectoria directamente en el proceso de inferencia de modelos de Rectified Flow y Difusión sin requerir entrenamiento adicional o modificaciones de pesos.
• Estimación sobre la Marcha (On-the-Fly): El método introduce un mecanismo novedoso para estimar latentes limpios ($\hat{x}_{0,t}$) durante la inferencia, permitiendo una detección de artefactos fiable en pasos intermedios donde los detectores estándar suelen fallar.
• Generalización: Se demuestra que el enfoque funciona eficazmente tanto en modelos de Rectified Flow (FLUX.1, FLUX.2) como en modelos de Difusión estándar (SDXL).
• Normalización de Gradiente: El artículo identifica e implementa la normalización de gradiente para desacoplar la fuerza de la actualización del tamaño espacial del artefacto, asegurando una corrección consistente independientemente de la extensión del mismo.

Resultados Experimentales
Las evaluaciones cuantitativas se realizaron en tres conjuntos de datos (animales, palabras, personas) utilizando modelos como FLUX.1 [dev], FLUX.1 [schnell], FLUX.2 [dev] y SDXL.

• Reducción de Artefactos: En el conjunto de datos de "palabras" con FLUX.1 [dev], DIAMOND redujo la Frecuencia Media de Artefactos de una línea base del 100% a un 9.85%, superando significativamente a las líneas base como DiffDoctor (46.97%) y HPSv2 (56.82%).
• Ratio de Píxeles: El Ratio de Píxeles de Artefacto (APR) se redujo a niveles cercanos a cero (p. ej., 0.009% en el conjunto de datos de palabras) en comparación con las líneas base.
• Preservación de la Calidad: El método mantuvo puntuaciones CLIP-T (alineación texto-imagen) e ImageReward (preferencia humana) elevadas, lo que indica que la eliminación de artefactos no degradó la consistencia semántica ni la calidad visual.
• Preservación de Identidad: Los estudios de ablación mostraron que, si bien el DIAMOND sin regularizar elimina los artefactos de manera efectiva, añadir el término $L_{rec}$ (con $\alpha=0.1$) equilibra de forma óptima la supresión de artefactos con la preservación de la identidad del modelo base.

Significancia y Reivindicaciones
Los autores afirman que DIAMOND proporciona un camino robusto y zero-shot hacia la síntesis de imágenes de alta fidelidad y libres de artefactos. Al intervenir directamente durante el proceso de muestreo mediante la corrección de la trayectoria, el método supera los costos computacionales del ajuste fino y las imprecisiones de las heurísticas post-hoc. El artículo enfatiza que este enfoque permite una supervisión espacial de grano fino y una corrección precisa de artefactos localizados (como manos o texto) manteniendo la estructura global y la adherencia al prompt de la imagen generada. El método se presenta como una solución práctica para aplicaciones profesionales donde la corrección anatómica es crítica, como en la visual de comercio electrónico y marketing.