Élagage de Tokens pour un Meilleur Caching : Accélération 9× sur Stable Diffusion Gratuitement

Informations Fondamentales

• ID de l'article : 2501.00375
• Titre : Token Pruning for Caching Better: 9× Acceleration on Stable Diffusion for Free
• Auteurs : Evelyn Zhang, Bang Xiao, Jiayi Tang, Qianli Ma, Chang Zou, Xuefei Ning, Xuming Hu, Linfeng Zhang
• Classification : cs.CV (Vision par Ordinateur), cs.LG (Apprentissage Automatique)
• Date de publication : 31 décembre 2024
• Lien de l'article : https://arxiv.org/abs/2501.00375
• Lien du code : github.com/EvelynZhang-epiclab/DaTo

Résumé

Stable Diffusion a remporté un succès remarquable dans le domaine de la génération d'images à partir de texte, mais son mécanisme itératif de débruitage entraîne des coûts de calcul élevés et une génération lente. Bien que des méthodes telles que le caching de caractéristiques aient attiré l'attention en raison de leur efficacité et de leur simplicité, la simple réutilisation des caractéristiques calculées aux étapes temporelles précédentes conduit à des caractéristiques similaires entre étapes temporelles adjacentes, réduisant la dynamique des caractéristiques au fil du temps et affectant finalement la qualité des images générées. Cet article propose une méthode d'élagage de tokens consciente de la dynamique (DaTo) pour résoudre les limitations du caching de caractéristiques. DaTo élague sélectivement les tokens ayant une dynamique faible, permettant uniquement aux tokens hautement dynamiques de participer aux couches d'auto-attention, étendant ainsi la dynamique des caractéristiques entre étapes temporelles. Lorsqu'elle est appliquée à Stable Diffusion sur ImageNet, cette méthode réalise une accélération 9×, tandis que le FID diminue de 0,33 ; une accélération 7× est observée sur COCO-30k, avec une diminution significative du FID de 2,17.

Contexte de Recherche et Motivation

Contexte du Problème

Les modèles de diffusion ont réalisé des progrès majeurs dans le domaine de la modélisation générative, avec des applications largement répandues dans la génération d'images à partir de texte, la génération vidéo et autres tâches. Cependant, le mécanisme itératif de débruitage des modèles de diffusion entraîne des coûts de calcul énormes et une génération lente, limitant leurs applications plus larges.

Limitations des Méthodes Existantes

Les méthodes actuelles pour accélérer les modèles de diffusion incluent principalement :

1. Réduction du nombre d'étapes d'échantillonnage : comme les échantillonneurs rapides DDIM
2. Réduction du coût de calcul à chaque étape : incluant la distillation de connaissances, l'élagage structurel, la quantification, l'élagage de tokens et le caching de caractéristiques

Parmi ceux-ci, le caching de caractéristiques est populaire en raison de son efficacité et de sa simplicité, stockant les caractéristiques calculées aux étapes temporelles précédentes et les réutilisant aux étapes temporelles suivantes. Cependant, la réutilisation de caractéristiques force les caractéristiques à différentes étapes temporelles à avoir des valeurs similaires, réduisant la dynamique des caractéristiques le long des étapes temporelles, endommageant le processus de diffusion original et réduisant ainsi la qualité de génération.

Motivation de la Recherche

L'article observe expérimentalement que, comparé au Stable Diffusion original, le modèle utilisant le caching de caractéristiques montre une différence de caractéristiques significativement réduite entre étapes temporelles adjacentes. Cela soulève une question clé : Est-il possible d'effectuer le caching de caractéristiques tout en maintenant la dynamique correcte des caractéristiques ?

Contributions Principales

1. Proposition de la méthode d'élagage de tokens consciente de la dynamique (DaTo) : en élaguant les tokens dont la dynamique est réduite par le caching de caractéristiques à différentes étapes temporelles et en les récupérant avec des tokens hautement dynamiques, évitant la dégradation de la qualité de génération causée par le caching de caractéristiques.
2. Conception d'une stratégie de recherche évolutive : proposition d'une méthode de recherche par évolution pour trouver les stratégies optimales de caching de caractéristiques et d'élagage de tokens, libérant pleinement le potentiel de DaTo.
3. Réalisation d'améliorations de performance significatives : les expériences étendues sur Stable Diffusion et SDXL montrent qu'une accélération jusqu'à 9× peut être obtenue sur Stable Diffusion sans entraînement ni données supplémentaires, avec une perte de qualité de génération minimale.

Explication Détaillée de la Méthode

Définition de la Tâche

La tâche de cet article est d'accélérer significativement le processus d'inférence du modèle Stable Diffusion tout en maintenant la qualité de génération d'images. L'entrée est une invite textuelle, la sortie est l'image de haute qualité correspondante, avec la contrainte qu'aucun réentraînement du modèle n'est nécessaire.

Architecture du Modèle

1. Élagage de Tokens Conscient de la Dynamique (DaTo)

Sélection de Tokens de Base :

• Score de Différence de Bruit Temporel : Pour la t-ième étape temporelle, calculer la différence absolue des sorties des deux étapes temporelles adjacentes précédentes :

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

• Sélection de Tokens Basée sur les Patches : Diviser l'image en patches non chevauchants s×s, sélectionner le token ayant le DiffScore le plus élevé dans chaque patch comme token de base.

Alignement CFG : Pour traiter la guidance sans classificateur (CFG), copier les positions des tokens de base de la génération conditionnelle à la génération inconditionnelle :

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

Sélection de Tokens à Élaguer : Sélectionner les K tokens les plus similaires aux tokens de base en fonction de la similarité cosinus pour l'élagage :

X_prune = arg topK max Cosine_Similarity(X_i, X_j)

Récupération de Tokens Élaguées : Récupérer les tokens élaguées en copiant directement leur token de base le plus similaire.

2. Caching de Caractéristiques Conscient des Étapes Temporelles

Élagage de l'Espace de Recherche :

• Profondeur de cache d limitée à {0, 1, 1/2}
• Ratio d'élagage r limité à {0.3, 0.4, 0.5, 0.6, 0.7}

Algorithme de Recherche Évolutive : Utiliser l'algorithme d'optimisation multi-objectif NSGA-II, avec les objectifs d'optimisation incluant :

• Latence d'inférence
• Qualité de génération (FID)

Le processus de recherche inclut les opérations évolutives standard telles que la sélection, le croisement et la mutation, obtenant finalement la stratégie optimale consciente des étapes F(t).

Points d'Innovation Technique

1. Mécanisme de Récupération de Dynamique : En élaguant sélectivement les tokens de faible dynamique et en les récupérant avec des tokens de haute dynamique, récupérer avec succès la distribution de dynamique des caractéristiques endommagée par le caching de caractéristiques.
2. Cadre Unifié de Caching-Élagage : Combiner le caching de caractéristiques et l'élagage de tokens dans un cadre indépendant de l'entraînement, réalisant la réutilisation d'informations au niveau temporel et au niveau des tokens.
3. Recherche de Stratégie Adaptative : Pour les différentes caractéristiques de redondance à différentes étapes temporelles, proposer une méthode pour rechercher automatiquement la profondeur de cache optimale et le ratio d'élagage.

Configuration Expérimentale

Ensembles de Données

• ImageNet-1k : Génération de 2000 images 512×512 (2 par classe)
• COCO-30k : Génération de 30000 images (1 par titre)
• Ensemble de validation MS COCO : Pour l'évaluation SDXL, génération de 5k images 1024×1024

Métriques d'Évaluation

• FID (Fréchet Inception Distance) : Mesurer la qualité de génération
• CLIP Score : Évaluer l'alignement texte-image
• Inception Score : Évaluation de la qualité d'image
• Latence et Ratio d'Accélération : Évaluation de l'efficacité

Méthodes de Comparaison

• DDIM/DPM : Échantillonneurs rapides
• ToMeSD : Méthode de fusion de tokens
• DeepCache : Méthode de caching de caractéristiques
• DeepCache & ToMeSD : Méthode de combinaison naïve

Détails d'Implémentation

• Utiliser l'algorithme évolutif NSGA-II, taille de population 20, exécution 100 générations
• CFG scale : 7.5 (SD v1.5), 9.0 (SD v2), 7.0 (SDXL)
• Nombre d'étapes d'échantillonnage : 50 étapes PLMS
• Test sur un seul GPU 4090

Résultats Expérimentaux

Résultats Principaux

Stable Diffusion v1.5 (ImageNet) :

• Configuration e1 : Accélération 9.01×, FID diminue de 27.64 à 27.31
• Supérieur aux méthodes de comparaison dans toutes les configurations

Stable Diffusion v2 (ImageNet) :

• Configuration e2 : Accélération 7.25×, FID de 28.20
• Comparé au modèle original, FID diminue de 29.8 à 28.20

Ensemble de Données COCO-30k :

• SD v1.5 : Accélération 7×, FID diminue de 12.15 à 9.98 (diminution de 2.17)
• SD v2 : Accélération 7.25×, FID de 13.68 à 13.88

SDXL (MS COCO) :

• Accélération 2.32×, FID diminue de 24.25 à 23.10
• Significativement supérieur à DeepCache (1.75×) et DeepCache&ToMeSD (1.78×)

Expériences d'Ablation

Efficacité du DiffScore : Avec différents paramètres de caching et ratios d'élagage, l'utilisation du DiffScore améliore constamment les scores FID, prouvant l'efficacité du score de différence de bruit temporel.

Impact de l'Alignement CFG : Avec l'augmentation du ratio d'élagage, les gains apportés par la configuration d'alignement CFG augmentent progressivement, avec des améliorations FID allant de 13 à 30 points à des ratios d'élagage élevés (0.7).

Analyse de Cas

Les résultats de comparaison visuelle montrent que DaTo excelle dans plusieurs aspects :

• Fidélité du Contenu : Hautement similaire au contenu de l'image originale
• Préservation des Détails : Maintien des textures fines dans les scènes à haute détail
• Adaptation de Style : Équilibre entre préservation du contenu et précision du style dans les tâches image-à-image
• Alignement d'Invite : Génération précise de tous les éléments dans les invites textuelles complexes

Découvertes Expérimentales

1. Récupération de la Dynamique des Caractéristiques : DaTo récupère avec succès la distribution de différence de caractéristiques à un niveau proche du Stable Diffusion original
2. Effet de Codage Creux : L'élagage modéré de tokens et le caching de caractéristiques peuvent améliorer les performances du modèle en se concentrant sur les caractéristiques clés
3. Généralisation de la Stratégie : La stratégie recherchée sur SD v1.5 fonctionne bien sur SDXL et d'autres ensembles de données

Travaux Connexes

Modèles de Diffusion Efficaces

• Réduction des Étapes d'Échantillonnage : DDIM, modèles de cohérence, etc.
• Compression de Réseau : Quantification, élagage, distillation, etc.
• Optimisation d'Architecture : Améliorations U-Net, optimisations Transformer, etc.

Stratégies de Réduction de Tokens

• Méthodes Apprises : DynamicViT, A-ViT, etc. utilisant des modèles auxiliaires pour le tri et l'élagage
• Méthodes Heuristiques : Token Pooling, Token Merging et autres méthodes indépendantes de l'entraînement
• Applications aux Modèles de Diffusion : ToMeSD, AT-EDM et autres adaptations pour les tâches de génération

Mécanismes de Caching

• Caching U-Net : DeepCache exploitant la redondance temporelle pour le caching de caractéristiques
• Caching DiT : Stratégie de caching Δ-DiT pour Diffusion Transformer
• Défis d'Optimisation : Équilibre entre amélioration d'efficacité et maintien de la qualité de génération

Conclusion et Discussion

Conclusions Principales

1. DaTo résout avec succès le problème de perte de dynamique des caractéristiques causée par le caching de caractéristiques
2. La stratégie adaptative obtenue par recherche évolutive est significativement supérieure aux configurations fixes
3. Cette méthode réalise une accélération et une amélioration de qualité significatives sur plusieurs modèles et ensembles de données

Limitations

1. Coût de Recherche : Bien que ≤20 heures GPU soit acceptable, cela nécessite toujours des ressources de calcul supplémentaires
2. Dépendance Matérielle : Les améliorations de performance peuvent varier selon la configuration matérielle
3. Limitations dans les Configurations Extrêmes : Des ratios d'élagage trop élevés ou une fréquence de mise à jour de cache trop basse endommagent les performances

Directions Futures

1. Apprentissage de Stratégie Adaptative : Développer des stratégies de caching et d'élagage plus intelligentes et adaptatives
2. Adaptation à d'Autres Architectures : Étendre à plus d'architectures de modèles de diffusion
3. Analyse Théorique : Approfondir la compréhension du rôle des principes de codage creux dans les modèles de diffusion

Évaluation Approfondie

Avantages

1. Forte Innovativité : Première résolution systématique du problème de perte de dynamique du caching de caractéristiques
2. Méthode Pratique : Indépendante de l'entraînement, facile à déployer et intégrer
3. Expérimentation Complète : Évaluation exhaustive sur plusieurs modèles et ensembles de données
4. Support Théorique : Fournit une explication théorique du codage creux
5. Convivialité Open Source : Fournit une implémentation de code complète

Insuffisances

1. Analyse Théorique Insuffisante : L'explication théorique de pourquoi cette méthode améliore le FID est relativement simple
2. Dépendance à l'Algorithme de Recherche : Nécessite une recherche évolutive pour trouver la stratégie optimale, augmentant la complexité d'utilisation
3. Métriques d'Évaluation Limitées : Dépend principalement du FID pour l'évaluation, manquant de métriques de qualité plus diversifiées
4. Absence d'Étude Utilisateur : Pas d'évaluation humaine pour vérifier la qualité de génération

Impact

1. Valeur Académique : Fournit de nouvelles perspectives et méthodes pour l'accélération des modèles de diffusion
2. Valeur Pratique : Peut être directement appliqué aux modèles Stable Diffusion existants
3. Reproductibilité : Fournit des détails d'implémentation détaillés et du code open source
4. Caractère Inspirant : Fournit un exemple d'application de l'optimisation au niveau des tokens dans les modèles génératifs

Scénarios Applicables

1. Environnements aux Ressources Limitées : Appareils mobiles, scénarios d'informatique en périphérie
2. Applications en Temps Réel : Applications interactives nécessitant une génération rapide
3. Génération en Masse : Tâches de génération d'images à grande échelle
4. Prototypes de Recherche : Projets de recherche nécessitant une itération rapide

Références

L'article cite 47 travaux connexes, couvrant plusieurs domaines connexes tels que les modèles de diffusion, la réduction de tokens et les mécanismes de caching, fournissant une base théorique solide et des repères de comparaison pour cette recherche.

Évaluation Globale : Ceci est un article de haute qualité en vision par ordinateur, proposant une solution innovante à l'important problème de l'accélération des modèles de diffusion. La conception de la méthode est ingénieuse, l'évaluation expérimentale est complète et la valeur pratique est remarquable. Bien qu'il y ait une certaine insuffisance dans la profondeur de l'analyse théorique, ses contributions pratiques et son impact méritent d'être reconnus.