Resumen Técnico: STEP-OPD

Declaración del Problema

La destilación on-policy (OPD) ha surgido como un enfoque viable para consolidar múltiples modelos de generación de imágenes especializados en tareas en un único modelo estudiante mediante la supervisión de sus propias trayectorias de eliminación de ruido. Sin embargo, los métodos de OPD existentes sufren dos limitaciones primarias:

1. Techo de Emparejamiento con el Maestro (Teacher-Matching Ceiling): Los métodos actuales optimizan al estudiante únicamente para que coincida con la velocidad de salida del maestro en los estados visitados. Esto trata al maestro como el límite superior absoluto del objetivo de optimización. Como ilustra el artículo, la pérdida estándar de emparejamiento de velocidad al cuadrado se minimiza exactamente cuando la predicción del estudiante es igual a la del maestro; cualquier intento de "ir más allá" del maestro aumenta la pérdida, lo que impide que el estudiante alcance capacidades superiores a las del maestro.
2. Dinámicas Internas Sub-Constreñidas: La supervisión de la velocidad a nivel de salida proporciona solo una señal de punto final. Dado que la velocidad final agrega los efectos de todos los bloques de la red, las desviaciones en las actualizaciones de las representaciones a través de diferentes capas pueden cancelarse entre sí en la predicción final. En consecuencia, un estudiante puede lograr una baja pérdida de salida mediante compensación sin aprender fielmente el procesamiento estructurado de tokens visuales e información de condicionamiento del maestro. Esto debilita la transferencia de capacidades que requieren un desarrollo progresivo a través de las capas, como la vinculación de atributos, la composición espacial y el renderizado de texto.

Metodología: STEP-OPD

Los autores proponen STEP-OPD (Extrapolación de Salida Paso a Paso y Alineación de Cambio de Representación), un marco que extiende el objetivo de aprendizaje del estudiante más allá del maestro e introduce restricciones explícitas sobre la evolución de la representación interna. El método consta de dos componentes principales:

1. Extrapolación de Salida Específica de la Tarea

En lugar de utilizar la velocidad del maestro ($v_{T_k}$) como el objetivo final, STEP-OPD utiliza la diferencia de velocidad entre el maestro específico de la tarea y el modelo base compartido ($v_B$) como una dirección para un aprendizaje adicional.

• Mecanismo: La velocidad de objetivo extrapolada $\hat{v}_{k,i}$ se construye añadiendo una versión escalada de la diferencia maestro-base a la velocidad del maestro:
  $$\hat{v}_{k,i} = v_{T_k,i} + \alpha_k (v_{T_k,i} - v_{B,i})$$
  donde $\alpha_k \geq 0$ es un coeficiente de extrapolación específico de la tarea.
• Efecto: Cuando $\alpha_k > 0$, el objetivo de aprendizaje se desplaza más allá del maestro a lo largo de la dirección de mejora introducida por el post-entrenamiento. Esto transforma al maestro de un punto final fijo en una guía direccional para una optimización adicional.
• Estrategia de Entrenamiento: Para evitar la inestabilidad, el coeficiente $\alpha_k$ se activa gradualmente de 0 a un valor máximo durante un período de calentamiento (warm-up).

2. Alineación de Cambio de Representación Intermedia (RCA)

Para abordar las dinámicas internas sub-constreñidas, STEP-OPD alinea la evolución de los estados ocultos entre el estudiante y el maestro a través de los bloques de la red.

• Mecanismo: En lugar de emparejar los estados ocultos absolutos, el método alinea la dirección y magnitud de los cambios en las representaciones ocultas entre bloques de Transformer adyacentes. Para un par de bloques $p=(l_0, l_1)$, el cambio se define como $\Delta h^p_{M,i} = h^{l_1}_{M,i} - h^{l_0}_{M,i}$.
• Función de Pérdida: La pérdida de alineación consta de dos partes:
  1. Pérdida de Dirección: Minimiza el MSE entre los cambios normalizados por RMS por token del estudiante y del maestro.
  2. Pérdida de Magnitud: Minimiza la diferencia al cuadrado entre las magnitudes RMS escaladas logarítmicamente de los cambios.
• Integración: Esta supervisión interna se aplica directamente a los cambios del maestro (no extrapolados), asegurando que el estudiante aprenda la dinámica de transformación específica del maestro mientras que el objetivo de salida tiene permitido exceder al maestro.

El objetivo de entrenamiento final combina la pérdida de salida extrapolada ($L_{out}$) y la pérdida de estado oculto ($L_{hidden}$):
$$L = L_{out} + \lambda_h L_{hidden}$$

Contribuciones Clave

1. Propuesta de Marco de Trabajo: Introducción de STEP-OPD, un marco de destilación on-policy que permite el aprendizaje "más allá del maestro" mientras transfiere la dinámica de la representación interna.
2. Extrapolación de Salida: Una técnica novedosa que construye objetivos de aprendizaje más allá del maestro aprovechando la diferencia de velocidad entre el modelo base y los maestros específicos de cada tarea.
3. Alineación de Cambio de Representación: Un método para supervisar la computación interna del estudiante mediante el emparejamiento de la dirección y magnitud de las actualizaciones de la representación entre capas, abordando las limitaciones de la supervisión solo de salida.
4. Validación Empírica: Experimentos exhaustivos que demuestran que el estudiante unificado supera tanto a los baselines de OPD estándar como a los correspondientes maestros de una sola tarea.

Resultados Experimentales

Los experimentos se realizaron en Stable Diffusion 3.5 Medium a través de tres grupos de capacidades: alineación composicional (GenEval), renderizado de texto (OCR) y preferencia humana (PickScore, Aesthetic, HPSv2.1, ImageReward).

• Ganancias de Rendimiento: Aplicado a DiffusionOPD, STEP-OPD aumentó la puntuación general de GenEval de 0.927 a 0.961 y la puntuación de OCR de 0.941 a 0.946. También mejoró todas las métricas de preferencia humana.
• Superar a los Maestros: El estudiante unificado de STEP-OPD superó a los correspondientes maestros de una sola tarea en las tres categorías de capacidad (por ejemplo, 0.961 frente a 0.935 en GenEval), demostrando una consolidación de capacidades y un rendimiento más allá del maestro efectivos.
• Estudios de Ablación:
  • Alineación de Cambio de Representación (RCA): Por sí sola, RCA mejoró GenEval (0.927 $\to$ 0.959) y OCR (0.941 $\to$ 0.944), confirmando que la supervisión interna es crucial para las tareas composicionales y de texto.
  • Extrapolación de Salida (OE): Por sí sola, OE proporcionó ganancias significativas en las métricas basadas en preferencias y mejoró aún más GenEval/OCR, validando la hipótesis de "más allá del maestro".
  • Combinado: La combinación de ambos componentes produjo los resultados más sólidos.
• Análisis Cualitativo: Las comparaciones visuales mostraron que STEP-OPD produjo una adherencia al prompt más precisa, un mejor renderizado de texto y una mayor calidad visual en comparación con los baselines. Específicamente, RCA redujo las estructuras internas espurias, mientras que OE mejoró los detalles finos y las texturas de los materiales.

Significado y Reivindicaciones

El artículo afirma que STEP-OPD cambia fundamentalmente el paradigma de la destilación on-policy del emparejamiento del maestro a la extrapolación guiada por el maestro. Al desacoplar el objetivo de salida de la velocidad exacta del maestro y supervisar explícitamente la dinámica de transformación interna, el método permite que un único modelo estudiante consolide múltiples capacidades especializadas y supere el rendimiento de los maestros individuales que generaron los datos de entrenamiento. Los autores sostienen que este enfoque resuelve las limitaciones de la supervisión de solo salida y proporciona un mecanismo robusto para transferir capacidades generativas complejas y dependientes de las capas.