NanoFLUX: Distillation-Driven Compression of Large Text-to-Image Generation Models for Mobile Devices

Resumen Técnico: NanoFLUX

Declaración del Problema

Los modelos de difusión de Texto-a-Imagen (T2I) a gran escala, como la familia FLUX (que escala hasta los 17B de parámetros) y Qwen-Image (20B parámetros), han alcanzado una calidad visual de vanguardia. Sin embargo, su escala masiva crea una barrera significativa para el despliegue en dispositivos, requiriendo infraestructura de GPU costosa o inferencia basada en la nube. Las soluciones existentes para dispositivos suelen depender de modelos más pequeños que carecen de una calidad comparable o emplean técnicas de compresión (poda, cuantización, reducción de pasos) que no abordan completamente la redundancia arquitectónica de los Transformers de Difusión (DiTs) de gran escala y sus codificadores de texto asociados. El desafío radica en comprimir estos modelos masivos para dispositivos móviles preservando la calidad de generación y minimizando la latencia de inferencia.

Metodología: NanoFLosX

Los autores presentan NanoFLUX, un modelo T2I de 2.4B de parámetros destilado del maestro FLUX.1-Schnell de 17B. La compresión se logra mediante un pipeline progresivo impulsado por la destilación que apunta a los dos componentes más pesados: el Transformer de Difusión (DiT) de 12B y el codificador de texto T5-XXL de 5B.

1. Compresión del Transformer (12B $\to$ 2B)

Los autores emplean una estrategia de compresión progresiva de cuatro etapas para reducir el tamaño del DiT:

• Etapa C1: Poda de Cabezales de Atención (12B $\to$ 5B):

  • Análisis: La Descomposición en Valores Singulares (SVD) en las salidas de los cabezales de atención reveló que no todos los cabezales son independientes; muchos contienen información redundante.
  • Acción: El número de cabezales de atención ($H$) se redujo uniformemente de 24 a 16.
  • Destilación: Se introdujo una pérdida a nivel de características para asegurar que los cabezales restantes preserven colectivamente la información de los cabezales podados. El estudiante fue inicializado con los primeros 16 cabezales del maestro.

• Etapa C2: Reducción de la Dimensión del Cabezal (5B $\to$ 3B):

  • Análisis: Se aplicó SVD nuevamente a las características del modelo de 5B.
  • Acción: La dimensión del cabezal ($d_H$) se redujo de 128 a 96, reteniendo los principales componentes singulares. Esto produjo un modelo de 3B de parámetros.

• Etapa C3: Poda de Profundidad mediante Fusión de Bloques (3B $\to$ 2.5B):

  • Análisis: Se calculó la similitud de coseno entrada-salida para los bloques del transformer. Los bloques de Flujo Único (SS) mostraron una alta similitud (>0.85), lo que indica redundancia funcional.
  • Acción: En lugar de eliminar estos bloques, los autores fusionaron una cadena de 15 bloques SS redundantes en un solo bloque promediando sus parámetros. Esto redujo la profundidad del transformer de 38 a 24 bloques.

• Etapa C4: Normalización de Capas Estática (2.5B $\to$ 1.8B):

  • Análisis: Los coeficientes de Normalización de Capas Adaptativa (AdaLN) predichos por el modelo mostraron una baja varianza entre muestras, sugiriendo que el condicionamiento dinámico de las proyecciones de CLIP era redundante.
  • Acción: Los coeficientes dinámicos de AdaLN fueron reemplazados por promedios estáticos precomputados (Static-LN). Esto eliminó aproximadamente 0.7B de parámetros. El modelo fue ajustado brevemente con coeficientes congelados para recuperar la calidad.

2. Submuestreo Progresivo de Tokens

Para abordar la complejidad cuadrática de la autoatención sin sacrificar el detalle:

• Mecanismo: Un módulo de submuestreo basado en ResNet ($D$) y un módulo de sobremuestreo ($U$) se insertaron en el transformer.
• Estrategia Híbrida: El modelo opera bajo un esquema de resolución híbrida. Los pasos tempranos de la difusión (que capturan la estructura global) procesan tokens submuestreados (reduciendo el conteo de tokens en 4x), mientras que los pasos posteriores (que refinan detalles) procesan tokens de resolución completa.
• Entrenamiento: Se utilizó una estrategia de "Submuestreo Progresivo de Tokens" (PTD) para estabilizar el entrenamiento. Los módulos se insertaron y entrenaron de forma incremental, moviendo el punto de submuestreo hacia etapas más tempranas de la red a lo largo de las iteraciones.
• Resultado: 23 de los 43 bloques del transformer operan sobre tokens submuestreados, reduciendo significativamente la latencia mientras mantienen el procesamiento de alta resolución donde es necesario.

3. Destilación del Codificador de Texto (5B $\to$ 330M)

• Objetivo: Destilar el codificador T5-XXL de 5B en un modelo T5-Large de 330M.
• Método: Un proceso de destilación de dos etapas:
  1. Emparejamiento de Embeddings: Se añade un MLP al estudiante para igualar la dimensión de salida del maestro, minimizando el MSE entre los embeddings de los prompts.
  2. Destilación por Bloques: El estudiante es entrenado para igualar los estados ocultos del prompt del maestro en capas específicas del transformer. Crucialmente, el método aprovecha señales visuales de las capas tempranas del denoiser (vía atención conjunta) para guiar al codificador de texto, mitigando la acumulación de errores.
  3. Estabilidad: Los gradientes se desprenden para pasos de tiempo de alto ruido ($t < \hat{t}$) para evitar una supervisión inestable, enfocando la destilación en los pasos de tiempo tardíos y en las tres primeras capas del transformer.

Resultados Clave

• Tamaño del Modelo: El modelo final NanoFLUX es de 2.4B de parámetros (2B DiT + 330M de Codificador de Texto), una reducción de 7$\times$ respecto al maestro FLUX.1-Schnell de 17B.
• Rendimiento:
  • En dispositivos móviles, NanoFLUX genera imágenes de 512 $\times$ 512 en aproximadamente 2.5 segundos.
  • Los benchmarks cuantitativos (One-IG, DPG, GenEval, HPDv3) muestran que NanoFLUX mantiene un rendimiento comparable al modelo maestro, con solo una degradación menor en métricas específicas (por ejemplo, HPSv3 cae de 11.45 a 10.41).
  • Los resultados cualitativos indican que el modelo preserva la calidad de la imagen y la adherencia al prompt de manera comparable al modelo original de 17B.
• Eficiencia: La estrategia de submuestreo progresivo de tokens logra equilibrar con éxito la latencia de inferencia y la calidad de generación, permitiendo que el modelo se ejecute eficientemente en hardware con recursos limitados.

Significancia y Reivindicaciones

El artículo afirma que NanoFLUX es el primer framework que estudia exhaustivamente la compresión de modelos a gran escala para la generación T2I, apuntando específicamente a modelos como FLUX. Su significancia radica en:

1. Viabilidad de la Generación en Dispositivos: Demostrar que la generación T2I de alta calidad es factible en dispositivos móviles sin depender de infraestructura en la nube.
2. Perspectiva Arquitectónica: Proporcionar una vía práctica para escalar hacia abajo grandes modelos generativos mediante la identificación y explotación de redundancias arquitectónicas específicas (cabezales de atención, profundidad de bloques, dinámica de normalización) en lugar de depender únicamente de la cuantización genérica o la reducción de pasos.
3. Generalizabilidad: Ofrecer conocimientos y un pipeline de destilación progresiva que pueden generalizarse a futuros modelos generativos de gran escala.

Los autores posicionan este trabajo como un paso hacia la ampliación de la accesibilidad a la IA generativa de alta calidad para usuarios con recursos computacionales limitados.