Resumen Técnico: CoReDiT

Enunciado del Problema

Los Transformadores de Difusión (DiT) han establecido un rendimiento de vanguardia en la generación de imágenes y videos al modelar eficazmente las dependencias de largo alcance mediante mecanismos de atención. Sin embargo, su costo computacional escala cuadráticamente con el número de tokens, lo que hace que la síntesis de alta resolución y la generación de video multiframe sean prohibitivamente costosas para plataformas móviles con recursos limitados. Si bien una porción significativa de los tokens en el espacio latente corresponde a regiones visualmente redundantes (por ejemplo, fondos uniformes), los métodos existentes de poda de tokens enfrentan tres desafíos principales al aplicarse a la inferencia de difusión:

1. Localización Eficiente: Identificar tokens de baja saliencia sin incurrir en la sobrecarga cuadrática de calcular puntuaciones de atención global.
2. Preservación Semántica: Mantener la continuidad visual y prevenir artefactos (por ejemplo, rupturas de textura, distorsiones de límites) cuando se omiten tokens, lo cual es crítico para tareas generativas a diferencia de las discriminativas.
3. Programación Adaptativa: Determinar una estrategia de poda que se adapte a los niveles variables de redundancia a través de diferentes bloques del transformador y pasos de eliminación de ruido.

Metodología: CoReDiT

CoReDiT es un marco estructurado de poda de tokens diseñado para acelerar los DiT preentrenados mientras preserva la fidelidad de salida. Opera a través de tres componentes principales:

1. Selección de Tokens Basada en Coherencia Espacial

En lugar de depender de puntuaciones de atención global (que son $O(N^2)$), CoReDiT utiliza una puntuación de Coherencia Espacial (CE) para estimar la redundancia local en tiempo lineal ($O(N)$).

• Mecanismo: La red de tokens se divide en cuadrículas no superpuestas. Para cada token, la puntuación de CE se calcula como la media de la similitud del coseno entre la incrustación normalizada del token y la incrustación normalizada media de su cuadrícula.
• Lógica: Los tokens con puntuaciones de CE altas se consideran redundantes porque están bien explicados por sus vecinos espaciales.
• Selección: Se seleccionan los $K$ tokens superiores con las puntuaciones de CE más altas para su omisión ($X_S$), mientras que los tokens restantes ($X_R$) proceden a la capa de Atención Múltiple de Cabeza (MHSA).

2. Reconstrucción de Tokens Guiada por Coherencia

Para prevenir las discontinuidades espaciales causadas por omitir tokens, CoReDiT reconstruye las salidas de atención para los tokens omitidos ($Y_S$) utilizando las salidas de los vecinos retenidos ($Y_R$).

• Aproximación: En lugar de calcular similitudes por pares para la reconstrucción (lo que añade sobrecarga), el método aproxima los pesos de similitud utilizando las puntuaciones de CE precalculadas de los vecinos.
• Subcuadrículas Locales: Para evitar un suavizado excesivo, la reconstrucción se realiza dentro de subcuadrículas más pequeñas.
• Micro-diseños:
  • Tamaños de Cuadrícula Alternos: Las dimensiones de la cuadrícula alternan entre bloques consecutivos para desplazar los artefactos de los límites y permitir la propagación de información a través de los bordes de la cuadrícula.
  • Retención con Paso $m$: Un patrón de paso determinista asegura que al menos un token se retenga en cada subcuadrícula, proporcionando anclajes necesarios para la reconstrucción.

3. Programación de Poda Progresiva y Adaptativa al Bloque

Reconociendo que la redundancia varía entre los bloques del transformador y los pasos de eliminación de ruido, CoReDiT emplea una estrategia de poda progresiva durante el ajuste fino ligero.

• Adaptabilidad al Bloque: El presupuesto de poda se asigna a los bloques con la mayor redundancia medida (estimada por la suma de las puntuaciones de CE de los tokens candidatos).
• Decaimiento por Paso de Tiempo: La poda es más agresiva en los primeros pasos de eliminación de ruido (donde domina el ruido y la estructura es gruesa) y se reduce en los pasos posteriores (donde se refinan los detalles finos). Se utiliza un programa de decaimiento de dos fases para equilibrar la eficiencia con la complejidad de gestionar grafos de computación específicos por paso.
• Entrenamiento: El modelo se ajusta fino utilizando una pérdida de destilación que combina la pérdida de salida del modelo original y una pérdida de bloque podado para recuperar el rendimiento.

Contribuciones Clave

1. Estimación de Redundancia en Tiempo Lineal: La introducción de una puntuación de coherencia espacial que identifica tokens redundantes en tiempo $O(N)$, evitando el cuello de botella computacional del análisis de atención global.
2. Preservación de la Red Densa: Un mecanismo de reconstrucción que sintetiza las salidas de atención omitidas a partir de los vecinos, asegurando que la red de tokens permanezca densa para las capas posteriores y previniendo artefactos visuales.
3. Estrategia de Poda Adaptativa: Un programa progresivo y adaptativo al bloque que asigna dinámicamente presupuestos de poda basados en estadísticas de redundancia empíricas a través de la profundidad y el tiempo, facilitando una adaptación estable.

Resultados Experimentales

CoReDiT se evaluó en arquitecturas base de DiT de vanguardia, incluyendo PixArt-α, PixArt-Σ y MagicDrive-V2.

• Ganancias de Eficiencia:
  • Logró una reducción de hasta el 55% en los FLOPs de atención auto-atención.
  • Aceleraciones de Inferencia: 1.33× en GPUs en la nube (Nvidia H100) y 1.72× en NPUs móviles (Qualcomm Snapdragon 8 Elite).
  • Margen de Memoria: Permitió la generación a resoluciones de hasta 1920 en dispositivos móviles, donde el modelo base encontró errores de memoria insuficiente (OOM) a 1600.
• Preservación de la Calidad:
  • En PixArt-α-1024 con un 40% de poda, CoReDiT mantuvo un FID de 28.7 (vs. 27.3 del modelo base) y puntuaciones competitivas de CLIP/IS. Con destilación, la brecha de FID se redujo a 27.4.
  • En PixArt-Σ-2048, una tasa de poda promedio del 23% redujo la latencia de atención auto-atención en un 22% con solo una caída modesta en las métricas de calidad.
  • En MagicDrive-V2 (video), una reducción del 39% en los FLOPs de atención auto-atención preservó la calidad por cuadro (PSNR, LPIPS, SSIM) y la alineación condicional (mAP, mIoU), aunque se notó una ligera caída en la consistencia temporal (FVD).
• Estudios de Ablación:
  • La selección aleatoria de tokens resultó en una degradación catastrófica de la calidad (FID 644.5), validando la necesidad de una selección basada en coherencia.
  • Desactivar la reconstrucción condujo a una pérdida significativa de calidad (FID 30.4 vs. 28.5), confirmando la importancia del paso de reconstrucción.
  • La poda uniforme seguida de ajuste fino tuvo un rendimiento peor que el programa adaptativo progresivo propuesto.

Significado y Afirmaciones

El artículo afirma que CoReDiT ofrece una solución práctica para desplegar modelos de difusión de alta fidelidad en dispositivos de borde al abordar los desafíos específicos de la poda de tokens generativa. Su significado radica en:

• Escalabilidad: Permitir la síntesis de mayor resolución y la generación de video en hardware móvil al reducir tanto la huella computacional como la de memoria.
• Compatibilidad Arquitectónica: Mantener la compatibilidad con las arquitecturas DiT existentes sin requerir cambios estructurales o módulos de enrutamiento complejos.
• Compensación Eficiencia-Calidad: Demostrar que la omisión estructurada y guiada por coherencia puede reducir drásticamente los costos de inferencia (hasta un 55% de FLOPs de atención auto-atención) mientras se mantiene una calidad perceptual comparable a la del modelo completo.

Los autores señalan que, si bien el enfoque actual se centra en la redundancia espacial, el trabajo futuro tiene como objetivo extender la eficiencia a los componentes FFN/MLP e incorporar objetivos de coherencia entre cuadros para mejorar aún más la consistencia temporal del video.