Cross-Space Distillation: Teaching One-Step Students with Modern Diffusion Teachers

Resumen Técnico: Destilación de Espacio Cruzado (Cross-Space Distillation)

Planteamiento del Problema
Los modelos de difusión de un solo paso modernos logran una generación de alta calidad mediante la destilación de pasos de tiempo basada en la distribución. Sin embargo, estos métodos dependen de una restrictiva restricción de Espacio Compartido: el Maestro y el Estudiante deben habitar el mismo espacio latente, compartiendo una resolución latente y una parametrización de VAE idénticas. Este supuesto impide la transferencia de conocimiento desde Maestros modernos de alta capacidad (por ejemplo, Stable Diffusion 3.5, Flux, Kolors) que operan a altas resoluciones (por ejemplo, 1024×1024) y utilizan VAEs avanzados, hacia Estudiantes compactos y aptos para el despliegue (por ejemplo, Stable Diffusion 1.5) que operan a resoluciones más bajas (por ejemplo, 512×512) con diferentes VAEs.

Los objetivos de destilación estándar fallan en este escenario porque los tensores latentes ($z_S$ y $z_T$) son inherentemente incompatibles debido a:

1. Resolución Cruzada: Formas de rejilla espacial desajustadas (por ejemplo, $h \times w$ frente a $H \times W$).
2. VAE Cruzado: Diferentes parametrizaciones de autoencoder que resultan en espacios latentes desalineados.
3. Arquitectura/Mecanismo Cruzado: Diferencias entre backbones UNet y DiT, o entre formulaciones de difusión estándar y de coincidencia de flujo (flow matching).

Las alternativas existentes, como la poda de redes (network pruning) o la construcción de nuevos backbones compactos desde cero, enfrentan desafíos que incluyen umbrales de escasez rígidos, requisitos de reentrenamiento extensos y la pérdida de compatibilidad con el ecosistema.

Metodología: Destilación de Espacio Cruzado y el Puente (The Bridge)
Los autores formalizan este escenario como Destilación de Espacio Cruzado e introducen Bridge ($B_\phi$), un módulo ligero diseñado para mapear los latentes del Estudiante al espacio latente del Maestro sin modificar el backbone del Estudiante.

• Arquitectura: El Bridge se factoriza en dos etapas para manejar el desajuste de manera eficiente:

  1. Prior Espacial (Congelado): Los primeros $n$ bloques de decodificación del decodificador VAE del Estudiante ($D_S^{(n)}$) se mantienen congelados. Estos bloques actúan como un upsampler espacial preentrenado, expandiendo el latente del Estudiante $z_S$ desde baja resolución a la alta resolución del Maestro ($H \times W$). Esto aprovecha el conocimiento espacial existente del Estudiante sin añadir parámetros entrenables para el upsampling.
  2. Proyector Aprendible ($g_\phi$): Una red compacta y aprendible (implementada como SwinIR) mapea las características upsampled del espacio de características del Estudiante al espacio de representación latente del Maestro, produciendo un latente compatible con el Maestro $\hat{z}_T$.

• Objetivos de Entrenamiento: El Bridge se entrena utilizando un conjunto de datos sintético de latentes emparejados $(z_S, z_T)$ derivados de la misma imagen codificada por ambos VAEs. La optimización minimiza una suma ponderada de dos pérdidas:

  1. Reconstrucción de Latente ($L_{rec}$): Una pérdida $\ell_1$ entre el latente predicho $\hat{z}_T$ y el latente real del Maestro $z_T$. Los autores señalan que la $\ell_1$ proporciona mejor estabilidad y preservación de la estructura que la $\ell_2$ en este contexto.
  2. Fidelidad de Atención ($L_{attn}$): Una pérdida de divergencia KL inversa que alinea las distribuciones de autoatención del denoiser del Maestro (cuando se le alimenta con $\hat{z}_T$) frente a $z_T$. Esto asegura que el latente mapeado induzca las mismas respuestas estructurales internas (dependencias de largo alcance) que el latente real del Maestro, preservando detalles finos que las pérdidas de píxel podrían omitir.

• Destilación e Inferencia: Una vez entrenado, el Bridge se adjunta al Estudiante. Durante la destilación, la salida del Estudiante se pasa a través del Bridge para generar $\hat{z}_T$, permitiendo que se apliquen directamente los objetivos de la distribución (por ejemplo, Destilación de Puntuación Variacional, Destilación Adversaria) contra el Maestro. Crucialmente, el backbone del Estudiante permanece sin cambios.

Contribuciones Clave

1. Formalización de la Destilación de Espacio Cruzado: El artículo define un nuevo régimen donde la destilación ocurre a través de resoluciones latentes y espacios VAE desajustados, abordando una brecha crítica en la transferencia de conocimiento desde modelos maestros modernos de alta resolución hacia modelos legados compactos.
2. El Módulo Bridge: Una interfaz ligera y modular que permite los objetivos de destilación estándar sin alterar la arquitectura del Estudiante. Combina un prior espacial congelado con un proyector aprendible y utiliza un objetivo de Fidelidad de Atención para una alineación estable.
3. Efectividad Práctica: El método permite que los Estudiantes compactos (por ejemplo, SD 1.5) hereden los priors visuales de diversos Maestros modernos (SDXL, Kolors, PixArt-$\Sigma$, Flux, SD 3.5) manteniendo la inferencia de un solo paso, baja latencia y una amplia compatibilidad con el ecosistema.
4. Capacidad de Fusión de Modelos: Debido a que el método destila diversos Maestros en Estudiantes con el mismo backbone y espacio latente, los checkpoints resultantes pueden fusionarse trivialmente mediante el promedio de parámetros para crear un único modelo que agregue el conocimiento de múltiples fuentes de Maestros.

Resultados
Los experimentos demuestran mejoras sustanciales en la calidad de generación para Estudiantes de un solo paso compactos:

• Ganancias Cuantitativas: Destilar SD 3.5 Medium (Maestro) en SD 1.5 (Estudiante) mediante Bridge mejora la puntuación HPSv3 de 5.37 a 9.42 e ImageReward de -0.29 a 0.62. Fusionar checkpoints de cinco Maestros diferentes aumenta aún más el HPSv3 de SD 1.5 a 10.53.
• Mejoras Cualitativas: Las comparaciones visuales muestran que los estudiantes habilitados por Bridge producen imágenes con texturas más claras, mayor consistencia estructural y menos artefactos en comparación con las líneas base inicializadas o los modelos podados.
• Eficiencia: El Bridge añade una sobrecarga insignificante (aprox. 5M de parámetros) y permite la inferencia de un solo paso. También permite la mejora de resolución en tiempo de inferencia, donde una muestra de un Estudiante de 512×512 puede mapearse a 1024×1024 a través del Bridge y ser decodificada con el decodificador del Maestro sin reentrenar al Estudiante.
• Comparación con la Poda: La destilación habilitada por Bridge supera a la poda de redes (por ejemplo, podar SD 3.5 al 30% de escasez), la cual tiende a producir desenfoque y artefactos.

Significancia
El artículo afirma que la Destilación de Espacio Cruzado desbloquea el potencial de transferir las capacidades generativas de Maestros masivos y modernos hacia backbones eficientes y desplegables sin el costo de entrenar nuevas arquitecturas desde cero. Al tratar la alineación de representaciones latentes heterogéneas como un problema de interfaz reutilizable, este trabajo permite:

• Compatibilidad de Ecosistema: Modelos compactos como SD 1.5 pueden actualizarse con calidad de vanguardia mientras mantienen la compatibilidad con herramientas y hardware existentes.
• Unificación de Conocimiento: La capacidad de fusionar el conocimiento de múltiples Maestros arquitectónicamente distintos (por ejemplo, Flux y SD 3.5) en un único modelo ligero de tipo Estudiante.
• Modularidad: El Bridge sirve como un componente plug-and-play que puede reutilizarse tanto para la destilación como para la mejora de resolución en tiempo de inferencia, desacoplando la capacidad de generación del Estudiante de sus restricciones de representación latente.

Los autores posicionan este trabajo como un paso hacia la exploración más amplia de la destilación a través de espacios latentes heterogéneos, aplicable más allá de la generación de texto a imagen, hacia la generación de video y multimodal.