Resumen Técnico: Guía Amortizada para la Inpainting de Imágenes con Modelos de Difusión Preentrenados (AID)

Enunciado del Problema

La inpainting de imágenes es una tarea de generación condicional cuyo objetivo es reconstruir las regiones faltantes de una imagen manteniendo la consistencia con los píxeles visibles. Aunque los modelos de difusión preentrenados han avanzado significativamente esta capacidad, los enfoques existentes enfrentan una dicotomía en eficiencia y adaptabilidad:

1. Entrenamiento Específico de la Tarea: Métodos como LaMa, MAT y BrushNet entrenan modelos dedicados o módulos de adaptación grandes. Aunque esto amortiza el costo de inferencia, requiere un tiempo de entrenamiento sustancial y añade una complejidad significativa al modelo.
2. Optimización por Instancia: Métodos como RePaint, DPS y MCG mantienen el núcleo de difusión preentrenado fijo, pero lo adaptan por separado para cada imagen enmascarada en el momento de la prueba. Esto preserva el modelo preentrenado, pero incurre en costos de optimización repetidos durante la implementación, tratando cada imagen como un nuevo problema de inferencia.

Existe un vacío para una solución de "punto medio" que amortice la inferencia sin entrenar un modelo grande específico de la tarea ni realizar optimización por instancia en el momento de la prueba.

Metodología: Inpainting Amortizado con Difusión (AID)

Los autores proponen AID, un marco que mantiene el núcleo de difusión preentrenado fijo y entrena un módulo de guía pequeño y reutilizable de forma offline. Este módulo se reutiliza luego en todas las imágenes enmascaradas en el momento de la prueba sin optimización adicional.

1. Formulación de Control

El problema se formula como un problema de control determinista.

• Dinámica: El proceso de difusión inverso está gobernado por una ecuación diferencial ordinaria (ODE) de flujo de probabilidad impulsada por una función de puntuación preentrenada fija $\hat{S}$.
• Guía: Se añade un término de control dependiente de la tarea $u(t, x; \xi)$ a la deriva, donde $\xi = (M, y)$ representa la máscara observable y los píxeles visibles. El control no depende de la imagen limpia oculta $x^\dagger$.
• Objetivo: El objetivo es minimizar una pérdida terminal $\Psi$ (que equilibra la fidelidad a los píxeles visibles y la región faltante) más un costo en ejecución que penaliza la magnitud del campo de guía. Esto conduce a una ecuación de Hamilton-Jacobi-Bellman (HJB) donde el control óptimo $u^*$ es determinista.

2. Puente Teórico: De Determinista a Estocástico

Resolver directamente la ecuación HJB determinista de alta dimensión es intratable. Los autores introducen un problema de control aleatorizado auxiliar utilizando políticas gaussianas para hacer el problema aprendible mediante Aprendizaje por Refuerzo de Tiempo Continuo (CT-RL).

• Problema Auxiliar: En lugar de un control determinista, el sistema utiliza una política gaussiana $\pi^{(\lambda)} \sim \mathcal{N}(\mu, \sigma^2)$, donde la media $\mu$ es la ley de guía candidata.
• Puente Preservador del Optimizador: El artículo demuestra dos teoremas críticos:
  1. Suficiencia: La guía determinista óptima $u^*$ puede elevarse a una política gaussiana óptima estableciendo la media de la política en $u^*$.
  2. Necesidad: Cualquier política gaussiana óptima para el problema auxiliar debe tener una media que recupere exactamente la guía determinista óptima $u^*$.
• Significado: Este puente asegura que aprender una política estocástica (lo que permite actualizaciones basadas en gradientes) no altera la guía determinista objetivo utilizada en la implementación. La varianza de la política gaussiana es fija y sirve únicamente como dispositivo de aprendizaje, no como un regularizador de entropía.

3. Algoritmo Actor-Crítico Amortizado

Basándose en el puente teórico, los autores desarrollan un algoritmo actor-crítico totalmente impulsado por datos:

• Parametrización: El actor (media de guía $\hat{\mu}_\phi$) y el crítico (función de valor $\hat{V}_\theta$) están parametrizados por redes neuronales.
• Aprendizaje: El algoritmo utiliza condiciones de ortogonalidad de martingala derivadas de la ecuación HJB auxiliar. Actualiza al actor y al crítico utilizando aproximación estocástica basada en trayectorias muestreadas de la distribución de la tarea.
• Implementación: Una vez entrenado, el módulo de guía se implementa de forma determinista. Para una nueva imagen enmascarada, el sistema utiliza la guía media aprendida $\hat{\mu}_\phi$ dentro del solucionador de tiempo inverso, descartando el ruido de la política.

Contribuciones Clave

1. Metodología: AID introduce un marco de guía controlada que amortiza la inferencia entrenando un módulo de guía ligero de forma offline, el cual se reutiliza en todas las tareas sin optimización por instancia.
2. Teoría: El artículo establece un riguroso "puente de equivalencia de políticas" entre el problema de guía determinista y una formulación gaussiana auxiliar. Demuestra que el problema auxiliar es preservador del optimizador, justificando el uso de métodos actor-crítico estocásticos para aprender un campo de control determinista.
3. Eficiencia: El enfoque añade menos del 1% de parámetros entrenables en relación con el núcleo preentrenado, mejorando significativamente el equilibrio entre calidad y velocidad.

Resultados Empíricos

Los autores evaluaron AID en AFHQv2, FFHQ (utilizando la tubería EDM en espacio de píxeles) e ImageNet (utilizando la tubería EDM2 en espacio latente).

• Rendimiento frente a Líneas Base de Núcleo Fijo: AID supera consistentemente a fuertes líneas base de núcleo fijo (Sin Guía, Reemplazo, MCG, DPS) en múltiples métricas (PSNR, SSIM, LPIPS) y tipos de máscara (libre, centro, franja).
• Rendimiento frente a RePaint: AID supera a RePaint (una fuerte línea base de optimización iterativa) en calidad utilizando solo ~10% del presupuesto de muestreo (por ejemplo, 35 pasos frente a 350 pasos para RePaint).
• Rendimiento frente a LatentPaint: En ImageNet, AID supera a LatentPaint (un método de adaptación en espacio latente reutilizable) en todas las métricas utilizando un módulo de guía con menos del 45% de los parámetros entrenables.
• Latencia: AID soporta un punto de operación de baja latencia (12 pasos) que reduce aún más el tiempo de inferencia mientras mantiene un rendimiento superior a las líneas base que no son RePaint y sigue siendo competitivo con RePaint.
• Generalización: El mismo módulo de guía entrenado se transfiere eficazmente a tipos de máscara no vistos y diferentes presupuestos de muestreo sin reentrenamiento.

Significado y Afirmaciones

El artículo afirma que AID es el primer marco de guía controlada amortizada para la inpainting de imágenes que:

1. Mantiene el núcleo de difusión preentrenado congelado.
2. Entrena únicamente un módulo de guía pequeño y reutilizable.
3. Proporciona un puente exacto preservador del optimizador desde el aprendizaje actor-crítico estocástico hasta el campo de guía determinista utilizado en la implementación.

Los autores posicionan este trabajo como un paso hacia el uso de marcos basados en control y RL para proporcionar métodos con fundamento teórico para estudiar y mejorar los modelos de difusión. Enfatizan que el RL de tiempo continuo sustenta tanto el análisis riguroso como el diseño práctico del algoritmo, ofreciendo una alternativa principista a la optimización heurística por instancia o al reentrenamiento pesado específico de la tarea.