Resumen Técnico: Aprendizaje Multi-Objetivo para Modelos de Difusión bajo Aprendizaje Semi-Supervisado

1. Enunciado del Problema

Los modelos de difusión se despliegan cada vez más como generadores condicionales para diversas tareas, como la síntesis de texto a imagen en dominios de prompts heterogéneos o el control robótico en múltiples entornos. Estos escenarios constituyen naturalmente un problema de Aprendizaje Multi-Objetivo (MOL), donde un único modelo debe lograr un rendimiento óptimo de Pareto a través de múltiples distribuciones objetivo.

Un desafío central en este contexto es el costo estadístico de entrenar un modelo "generalista". Para manejar múltiples objetivos, la clase de modelos generalistas ($\mathcal{F}$) típicamente requiere una capacidad sustancialmente mayor que las clases de modelos "especialistas" ($\mathcal{H}_k$) necesarias para tareas individuales. En configuraciones totalmente supervisadas, la complejidad de muestras para entrenar al generalista escala con su propia complejidad ($C_\mathcal{F}$), lo cual a menudo es prohibitivamente alto.

Además, los despliegues en el mundo real a menudo enfrentan un régimen de datos asimétrico:

• Los datos emparejados $(x, y)$ (por ejemplo, pares texto-imagen curados o demostraciones de robots expertos) son escasos y costosos de adquirir.
• Los datos de condición no emparejados $y$ (por ejemplo, prompts de texto crudos u observaciones de estado) son abundantes.

El artículo aborda la siguiente pregunta: ¿Podemos aprender un modelo de difusión óptimo de Pareto de manera eficiente en muestras aprovechando datos de condición débilmente supervisados (no etiquetados)?

2. Metodología

Los autores proponen un marco de entrenamiento semi-supervisado en dos etapas que reconcilia la necesidad de un generalista de alta capacidad con la escasez de datos emparejados.

Etapa 1: Entrenamiento de Especialistas

Para cada una de las $K$ objetivos (distribuciones objetivo $P_k$), se entrena una función de puntuación especialista ligera $h_k \in \mathcal{H}_k$ utilizando únicamente los datos emparejados limitados disponibles $\{(x_i^k, y_i^k)\}_{i=1}^n$.

• Se asume que las clases de especialistas $\mathcal{H}_k$ tienen una complejidad significativamente menor ($C_{\mathcal{H}_k} \ll C_\mathcal{F}$) que la clase generalista.
• El entrenamiento minimiza el objetivo de coincidencia de puntuación de eliminación de ruido en los datos emparejados.

Etapa 2: Entrenamiento del Generalista mediante Muestreo Pseudo

Se entrena un modelo generalista grande $f \in \mathcal{F}$ para aproximar el conjunto óptimo de Pareto a través de todas las objetivos.

• Generación de Pseudo-Muestras: Los especialistas entrenados $\{h_k\}$ se utilizan para generar datos pseudo-emparejados. Para cada condición no etiquetada abundante $y_i^k \sim P_Y^k$, el especialista genera un pseudo-objetivo $\tilde{x}_i^k \sim P_{h_k}(\cdot|y_i^k)$.
• Objetivo: El generalista se entrena sobre una combinación de los datos etiquetados originales y los datos pseudo-emparejados generados. El objetivo de entrenamiento minimiza una escalarización $S$ de las pérdidas de coincidencia de puntuación a través de las $K$ objetivos:
  $$\hat{f}_S = \arg\min_{f \in \mathcal{F}} S(\tilde{L}_1(f), \dots, \tilde{L}_K(f))$$
  donde $\tilde{L}_k$ es la pérdida empírica calculada sobre las pseudo-muestras (con un término de corrección para tener en cuenta el sesgo del especialista).

Extensión a la Toma de Decisiones Secuencial

El marco se extiende a políticas de difusión en aprendizaje por imitación. Aquí, la "condición" es el estado $s$, y el "objetivo" es la acción $a$. Los autores abordan el desplazamiento de distribución inherente a las simulaciones en política: dado que las distribuciones de estado experto son difíciles de muestrear directamente, la segunda etapa utiliza trayectorias generadas por las políticas de especialistas aprendidas bajo la dinámica del entorno para recopilar datos pseudo para la política generalista.

3. Contribuciones Teóricas Clave

Límites de Complejidad de Muestras

El artículo establece límites de generalización no asintóticos para el enfoque semi-supervisado propuesto.

• Resultado Principal (Teorema 3.2): El número de muestras emparejadas (etiquetadas) requeridas para lograr un error objetivo $\epsilon$ depende únicamente de la complejidad de las clases de especialistas ($C_{\mathcal{H}_k}$), no de la clase generalista ($C_\mathcal{F}$).
  • Muestras etiquetadas requeridas: $\tilde{O}\left(\frac{K C_\mathcal{H}}{\epsilon^2}\right)$.
  • Condiciones no etiquetadas requeridas: $\tilde{O}\left(\frac{K C_\mathcal{F}}{\epsilon^2}\right)$.
• Comparación: En contraste, el entrenamiento totalmente supervisado convencional del generalista requeriría $\tilde{O}\left(\frac{K C_\mathcal{F}}{\epsilon^2}\right)$ muestras etiquetadas. Esto demuestra una mejora significativa en la eficiencia de las muestras cuando $C_\mathcal{F} \gg C_\mathcal{H}$.
• Escalarizaciones Lineales (Teorema 3.5): Para escalarizaciones lineales (suma ponderada) y clases de funciones convexas, los autores derivan tasas más precisas, refinando aún más el límite sobre el error de aproximación.

Estimación de Distribución y Garantías de Política

• Estimación de Distribución: El artículo cierra la brecha entre el error de coincidencia de puntuación y la estimación de distribución (distancia de variación total). Demuestra que el generalista aprendido logra un rendimiento óptimo de Pareto en la aproximación de las distribuciones condicionales objetivo.
• Políticas de Difusión (Teorema 3.6): Los autores proporcionan la primera garantía teórica sobre la brecha de sub-optimidad de las políticas de difusión en el aprendizaje por imitación bajo configuraciones semi-supervisadas. Muestran que el marco maneja con éxito el desplazamiento de distribución causado por las simulaciones en política, manteniendo al mismo tiempo los beneficios de eficiencia de muestras del pipeline de especialista a generalista.

4. Resultados Experimentales

Los autores validan su teoría en dos dominios distintos:

1. Manipulación Robótica (StackCube & PickCube):

   • Configuración: Entrenamiento de políticas de difusión para tareas bajo variación de aleatorización de dominio (escena, iluminación, pose de la cámara).
   • Resultado: El método MOL semi-supervisado supera significativamente a una línea base solo con etiquetas, particularmente en configuraciones Fuera de Distribución (OOD) (por ejemplo, +16% en la tasa de éxito en StackCube Nivel 2, +14% en el Nivel 3). Esto sugiere que las simulaciones pseudo proporcionan una supervisión crucial para manejar los desplazamientos de distribución.

2. Restauración de Imágenes (Inpainting de CelebA-HQ):

   • Configuración: Inpainting de rostros con tres tipos diferentes de máscaras (líneas alternadas, región facial, trazo ancho).
   • Resultado: El enfoque semi-supervisado mejoró el PSNR, SSIM y LPIPS en todos los tipos de máscaras en comparación con la línea base solo con etiquetas, demostrando la efectividad del método en tareas de generación de imágenes.

5. Significado y Afirmaciones

El artículo afirma proporcionar un marco estadístico principiado para el aprendizaje multi-objetivo en modelos de difusión bajo regímenes semi-supervisados. Su significado principal radica en:

• Puente entre Teoría y Práctica: Formaliza el pipeline de destilación "de especialista a generalista", ofreciendo una justificación teórica de por qué este enfoque es más eficiente en muestras que el aprendizaje multi-tarea directo cuando los datos etiquetados son escasos.
• Eficiencia de Muestras: Demuestra que el costo estadístico de aprender un generalista complejo puede desacoplarse de su propia capacidad, dependiendo en cambio de la complejidad de especialistas ligeros entrenados con datos limitados.
• Manejo del Desplazamiento de Distribución: Extiende estas garantías a la toma de decisiones secuenciales (aprendizaje por imitación), abordando un desafío práctico crítico en robótica.

Los autores mantienen una postura modesta, reconociendo que sus garantías dependen de una asunción de realizabilidad (que las funciones de puntuación verdaderas se encuentran dentro de las clases de especialistas) y que actualmente no abordan la generalización Fuera de Distribución (OOD) más allá de los niveles específicos de aleatorización de dominio probados. Consideran este trabajo como un paso fundamental hacia la comprensión de las propiedades estadísticas del MOL semi-supervisado para modelos generativos.