Poda de Tokens para Almacenamiento en Caché Mejorado: Aceleración de 9 Veces en Stable Diffusion de Forma Gratuita

Información Básica

• ID del Artículo: 2501.00375
• Título: Token Pruning for Caching Better: 9× Acceleration on Stable Diffusion for Free
• Autores: Evelyn Zhang, Bang Xiao, Jiayi Tang, Qianli Ma, Chang Zou, Xuefei Ning, Xuming Hu, Linfeng Zhang
• Clasificación: cs.CV (Visión por Computadora), cs.LG (Aprendizaje Automático)
• Fecha de Publicación: 31 de diciembre de 2024
• Enlace del Artículo: https://arxiv.org/abs/2501.00375
• Enlace del Código: github.com/EvelynZhang-epiclab/DaTo

Resumen

Stable Diffusion ha logrado un éxito significativo en el campo de la generación de imágenes a partir de texto, pero su mecanismo iterativo de desruido conlleva altos costos computacionales y velocidades de generación lentas. Aunque métodos como el almacenamiento en caché de características han recibido atención por su efectividad y simplicidad, la simple reutilización de características calculadas en pasos de tiempo anteriores hace que las características en pasos de tiempo adyacentes se vuelvan similares, reduciendo la dinamicidad de las características a lo largo del tiempo y afectando finalmente la calidad de las imágenes generadas. Este artículo propone un método de poda de tokens consciente de la dinámica (DaTo) para abordar las limitaciones del almacenamiento en caché de características. DaTo poda selectivamente tokens con menor dinamicidad, permitiendo que solo tokens altamente dinámicos participen en capas de autoatención, extendiendo así la dinamicidad de características entre pasos de tiempo. Cuando se aplica a Stable Diffusion en ImageNet, el método logra una aceleración de 9×, mientras que el FID se reduce en 0.33; en COCO-30k se observa una aceleración de 7×, con una reducción significativa de FID de 2.17.

Antecedentes de Investigación y Motivación

Contexto del Problema

Los modelos de difusión han logrado avances significativos en el campo de la modelación generativa, con aplicaciones generalizadas en generación de imágenes a partir de texto, generación de video y otras tareas. Sin embargo, el mecanismo iterativo de desruido de los modelos de difusión conlleva enormes costos computacionales y velocidades de generación lentas, limitando sus aplicaciones más amplias.

Limitaciones de Métodos Existentes

Los métodos actuales para acelerar modelos de difusión incluyen principalmente:

1. Reducción del número de pasos de muestreo: Como muestreadores rápidos tipo DDIM
2. Reducción del costo computacional por paso: Incluyendo destilación de conocimiento, poda estructural, cuantificación, poda de tokens y almacenamiento en caché de características

Entre estos, el almacenamiento en caché de características es ampliamente popular por su efectividad y simplicidad, almacenando características calculadas en pasos de tiempo anteriores y reutilizándolas en pasos de tiempo posteriores. Sin embargo, la reutilización de características obliga a que las características en diferentes pasos de tiempo tengan valores similares, reduciendo la dinamicidad de las características a lo largo de los pasos de tiempo, dañando el proceso de difusión original y reduciendo así la calidad de generación.

Motivación de la Investigación

El artículo descubre a través de experimentos que, en comparación con Stable Diffusion original, los modelos que utilizan almacenamiento en caché de características muestran diferencias de características significativamente reducidas entre pasos de tiempo adyacentes. Esto plantea una pregunta clave: ¿Es posible mantener la dinamicidad correcta de características mientras se realiza el almacenamiento en caché de características?

Contribuciones Principales

1. Propone el método de poda de tokens consciente de la dinámica (DaTo): Al podar tokens cuya dinamicidad ha sido reducida por el almacenamiento en caché de características en diferentes pasos de tiempo, y recuperarlos con tokens de gran dinamicidad, evita la degradación de la calidad de generación causada por el almacenamiento en caché de características.
2. Diseña una estrategia de búsqueda evolutiva: Propone buscar a través de métodos evolutivos la estrategia óptima de almacenamiento en caché de características y poda de tokens, liberando completamente el potencial de DaTo.
3. Logra mejoras de rendimiento significativas: Experimentos extensos en Stable Diffusion y SDXL demuestran que sin entrenamiento ni datos adicionales, se puede lograr una aceleración de hasta 9× en Stable Diffusion sin pérdida de calidad de generación.

Explicación Detallada del Método

Definición de la Tarea

La tarea de este artículo es acelerar significativamente el proceso de inferencia del modelo Stable Diffusion mientras se mantiene la calidad de generación de imágenes. La entrada es un mensaje de texto, la salida es la imagen correspondiente de alta calidad, y la restricción es que no se requiere reentrenamiento del modelo.

Arquitectura del Modelo

1. Poda de Tokens Consciente de la Dinámica (DaTo)

Selección Básica de Tokens:

• Puntuación de Diferencia de Ruido Temporal: Para el paso de tiempo t, se calcula la diferencia absoluta de los resultados de los dos pasos de tiempo adyacentes anteriores:

  DiffScore = (1/C) * Σ|f_up_0(x_{t+2})_c - f_up_0(x_{t+1})_c|
  

• Selección de Tokens Basada en Parches: Se divide la imagen en parches no superpuestos de s×s, seleccionando en cada parche el token con el DiffScore más alto como token base.

Alineación CFG: Para manejar la guía libre de clasificador (CFG), se copian las posiciones de tokens base de la generación condicional a la generación incondicional:

X_base,i,j[k] = X_base,i,j[k - B/2], k ∈ {B/2, B/2+1, ..., B-1}

Selección de Tokens a Podar: Se seleccionan los K tokens más similares a los tokens base basándose en similitud de coseno para podar:

X_prune = arg topK max Cosine_Similarity(X_i, X_j)

Recuperación de Tokens Podados: Los tokens podados se recuperan copiando directamente su token base más similar.

2. Almacenamiento en Caché de Características Consciente del Paso de Tiempo

Poda del Espacio de Búsqueda:

• Profundidad de caché d limitada a {0, 1, 1/2}
• Relación de poda r limitada a {0.3, 0.4, 0.5, 0.6, 0.7}

Algoritmo de Búsqueda Evolutiva: Se utiliza el algoritmo de optimización multiobjetivo NSGA-II, con objetivos de optimización que incluyen:

• Latencia de inferencia
• Calidad de generación (FID)

El proceso de búsqueda incluye operaciones evolutivas estándar como selección, cruce y mutación, obteniendo finalmente la estrategia óptima consciente de pasos F(t).

Puntos de Innovación Técnica

1. Mecanismo de Recuperación de Dinamicidad: Al podar selectivamente tokens de baja dinamicidad y recuperarlos con tokens de alta dinamicidad, se recupera exitosamente la distribución de dinamicidad de características dañada por el almacenamiento en caché de características.
2. Marco Unificado de Caché-Poda: Se combina el almacenamiento en caché de características y la poda de tokens en un marco independiente del entrenamiento, realizando reutilización de información a nivel de tiempo y token.
3. Búsqueda de Estrategia Adaptativa: Para las diferentes características de redundancia en diferentes pasos de tiempo, se propone un método para buscar automáticamente la profundidad de caché óptima y la relación de poda.

Configuración Experimental

Conjuntos de Datos

• ImageNet-1k: Generación de 2000 imágenes de 512×512 (2 por clase)
• COCO-30k: Generación de 30000 imágenes (1 por título)
• Conjunto de Validación MS COCO: Para evaluación de SDXL, generación de 5k imágenes de 1024×1024

Métricas de Evaluación

• FID (Distancia de Inception de Fréchet): Mide la calidad de generación
• Puntuación CLIP: Evalúa la alineación texto-imagen
• Puntuación de Inception: Evaluación de calidad de imagen
• Latencia y Relación de Aceleración: Evaluación de eficiencia

Métodos de Comparación

• DDIM/DPM: Muestreadores rápidos
• ToMeSD: Método de fusión de tokens
• DeepCache: Método de almacenamiento en caché de características
• DeepCache & ToMeSD: Método de combinación ingenua

Detalles de Implementación

• Se utiliza algoritmo evolutivo NSGA-II, tamaño de población 20, ejecución de 100 generaciones
• Escala CFG: 7.5 (SD v1.5), 9.0 (SD v2), 7.0 (SDXL)
• Pasos de muestreo: 50 pasos PLMS
• Pruebas en una sola GPU 4090

Resultados Experimentales

Resultados Principales

Stable Diffusion v1.5 (ImageNet):

• Configuración e1: Aceleración de 9.01×, FID de 27.64 a 27.31
• Supera a todos los métodos de comparación en todas las configuraciones

Stable Diffusion v2 (ImageNet):

• Configuración e2: Aceleración de 7.25×, FID de 28.20
• Comparado con el modelo original, FID de 29.8 a 28.20

Conjunto de Datos COCO-30k:

• SD v1.5: Aceleración de 7×, FID de 12.15 a 9.98 (reducción de 2.17)
• SD v2: Aceleración de 7.25×, FID de 13.68 a 13.88

SDXL (MS COCO):

• Aceleración de 2.32×, FID de 24.25 a 23.10
• Significativamente superior a DeepCache (1.75×) y DeepCache&ToMeSD (1.78×)

Experimentos de Ablación

Efectividad de DiffScore: En diferentes configuraciones de caché y relaciones de poda, el uso de DiffScore mejora consistentemente las puntuaciones de FID, demostrando la efectividad de la puntuación de diferencia de ruido temporal.

Impacto de la Alineación CFG: Con el aumento de la relación de poda, los beneficios de la configuración de alineación CFG aumentan gradualmente, con mejoras de FID que van desde 13 puntos hasta 30 puntos en relaciones de poda altas (0.7).

Análisis de Casos

Los resultados de comparación visual muestran que DaTo se desempeña excepcionalmente en múltiples aspectos:

• Fidelidad de Contenido: Altamente similar al contenido de la imagen original
• Preservación de Detalles: Mantiene texturas finas en escenas de alto detalle
• Adaptación de Estilo: Equilibra la preservación de contenido con la precisión de estilo en tareas de imagen a imagen
• Alineación de Mensajes: Genera con precisión todos los elementos en mensajes de texto complejos

Hallazgos Experimentales

1. Recuperación de Dinamicidad de Características: DaTo recupera exitosamente la distribución de diferencias de características a un nivel cercano al Stable Diffusion original
2. Efecto de Codificación Dispersa: La poda de tokens moderada y el almacenamiento en caché de características pueden mejorar el rendimiento del modelo enfocándose en características clave
3. Generalización de Estrategia: La estrategia buscada en SD v1.5 se desempeña bien en SDXL y otros conjuntos de datos

Trabajo Relacionado

Modelos de Difusión Eficientes

• Reducción de Pasos de Muestreo: DDIM, modelos de consistencia, etc.
• Compresión de Red: Cuantificación, poda, destilación, etc.
• Optimización de Arquitectura: Mejoras de U-Net, optimización de Transformer, etc.

Estrategias de Reducción de Tokens

• Métodos de Aprendizaje: DynamicViT, A-ViT y otros que utilizan modelos auxiliares para clasificación y poda
• Métodos Heurísticos: Token Pooling, Token Merging y otros métodos independientes del entrenamiento
• Aplicación en Modelos de Difusión: ToMeSD, AT-EDM y otras adaptaciones para tareas generativas

Mecanismos de Caché

• Caché de U-Net: DeepCache aprovecha la redundancia temporal para almacenar características en caché
• Caché de DiT: Estrategia de caché Δ-DiT para Diffusion Transformer
• Desafíos de Optimización: Equilibrio entre mejora de eficiencia y mantenimiento de calidad de generación

Conclusiones y Discusión

Conclusiones Principales

1. DaTo resuelve exitosamente el problema de pérdida de dinamicidad de características causado por el almacenamiento en caché de características
2. La estrategia adaptativa obtenida a través de búsqueda evolutiva supera significativamente las configuraciones fijas
3. El método logra aceleración significativa y mejora de calidad en múltiples modelos y conjuntos de datos

Limitaciones

1. Costo de Búsqueda: Aunque ≤20 horas de GPU es aceptable, aún requiere recursos computacionales adicionales
2. Dependencia de Hardware: Las mejoras de rendimiento pueden variar según la configuración de hardware
3. Limitaciones en Configuraciones Extremas: Relaciones de poda muy altas o frecuencias de actualización de caché muy bajas dañarán el rendimiento

Direcciones Futuras

1. Aprendizaje de Estrategia Adaptativa: Desarrollar estrategias de caché y poda más inteligentes y adaptativas
2. Adaptación a Otras Arquitecturas: Extender a más arquitecturas de modelos de difusión
3. Análisis Teórico: Comprender más profundamente el papel de los principios de codificación dispersa en modelos de difusión

Evaluación Profunda

Fortalezas

1. Innovación Fuerte: Primera solución sistemática al problema de pérdida de dinamicidad en almacenamiento en caché de características
2. Método Práctico: Independiente del entrenamiento, fácil de desplegar e integrar
3. Evaluación Completa: Evaluación integral en múltiples modelos y conjuntos de datos
4. Apoyo Teórico: Proporciona explicación teórica de codificación dispersa
5. Amigable con Código Abierto: Proporciona implementación de código completa

Deficiencias

1. Análisis Teórico Insuficiente: La explicación teórica de por qué el método mejora el FID es relativamente simple
2. Dependencia del Algoritmo de Búsqueda: Requiere búsqueda evolutiva para encontrar estrategias óptimas, aumentando la complejidad de uso
3. Métricas de Evaluación Limitadas: Depende principalmente de FID para evaluación, carece de métricas de calidad más diversificadas
4. Ausencia de Investigación de Usuarios: No se realizó evaluación humana para verificar la calidad de generación

Impacto

1. Valor Académico: Proporciona nuevas ideas y métodos para aceleración de modelos de difusión
2. Valor Práctico: Puede aplicarse directamente a modelos Stable Diffusion existentes
3. Reproducibilidad: Proporciona detalles de implementación detallados y código de código abierto
4. Inspiración: Proporciona un ejemplo de aplicación de optimización a nivel de token en modelos generativos

Escenarios Aplicables

1. Entornos con Recursos Limitados: Dispositivos móviles, escenarios de computación perimetral
2. Aplicaciones en Tiempo Real: Aplicaciones interactivas que requieren generación rápida
3. Generación por Lotes: Tareas de generación de imágenes a gran escala
4. Prototipos de Investigación: Proyectos de investigación que requieren iteración rápida

Referencias

El artículo cita 46 referencias relacionadas, cubriendo múltiples campos relacionados como modelos de difusión, reducción de tokens y mecanismos de caché, proporcionando una base teórica sólida y puntos de referencia de comparación para esta investigación.

Evaluación General: Este es un artículo de alta calidad en visión por computadora que propone una solución innovadora a un problema importante en la aceleración de modelos de difusión. El diseño del método es ingenioso, la evaluación experimental es completa y el valor práctico es destacado. Aunque hay cierta insuficiencia en la profundidad del análisis teórico, sus contribuciones prácticas e impacto merecen reconocimiento.