Fusion de Tokens Adaptatifs en Cache : Réduction Dynamique de Tokens et Élimination des Calculs Redondants dans les Modèles de Diffusion

Informations Fondamentales

• ID de l'article : 2501.00946
• Titre : Cached Adaptive Token Merging: Dynamic Token Reduction and Redundant Computation Elimination in Diffusion Model
• Auteurs : Omid Saghatchian, Atiyeh Gh. Moghadam, Ahmad Nickabadi (Université de Technologie Amirkabir)
• Classification : cs.CV (Vision par Ordinateur)
• Date de Publication : 1er janvier 2025 (prépublication arXiv)
• Lien de l'article : https://arxiv.org/abs/2501.00946
• Lien du code : https://github.com/omidiu/ca_tome

Résumé

Les modèles de diffusion sont devenus une méthode prometteuse pour générer des images de haute qualité et de haute dimension. Cependant, ces modèles sont entravés par des coûts de calcul élevés et une vitesse d'inférence lente, en partie en raison de la complexité de calcul quadratique du mécanisme d'auto-attention par rapport à la taille d'entrée. Cet article propose la méthode de fusion de tokens adaptatifs en cache (CA-ToMe), qui résout ce problème en calculant la similarité entre les tokens et en fusionnant les tokens dont la similarité dépasse un paramètre de seuil t. En raison des motifs répétitifs observés dans les étapes adjacentes et des variations de fréquence de similarité, la méthode améliore les approches de fusion de tokens en implémentant un seuil adaptatif et en ajoutant un mécanisme de cache. Les résultats expérimentaux montrent que la méthode, en tant que méthode d'accélération sans entraînement, réalise une accélération de 1,24× dans le processus de débruitage tout en maintenant le même score FID que les méthodes existantes.

Contexte de Recherche et Motivation

Définition du Problème

Les modèles de diffusion excellent dans les tâches de génération d'images, mais font face à de graves problèmes d'efficacité de calcul :

1. Coûts de calcul élevés : La complexité quadratique du mécanisme d'auto-attention entraîne une vitesse d'inférence lente
2. Processus de débruitage en série : Impossible de paralléliser, chaque étape de débruitage nécessite des calculs répétés
3. Calculs redondants : Il existe une quantité importante de calculs répétés entre les étapes temporelles adjacentes

Importance du Problème

• La latence élevée des modèles de diffusion limite leur utilisation dans les applications nécessitant une inférence rapide
• Les coûts de calcul élevés rendent le déploiement des modèles difficile, en particulier dans les environnements aux ressources limitées
• Les méthodes d'accélération existantes nécessitent soit un réentraînement, soit entraînent des pertes de qualité significatives

Limitations des Méthodes Existantes

1. Les méthodes réduisant le nombre d'étapes d'échantillonnage nécessitent généralement un réentraînement ou l'utilisation de solveurs numériques complexes
2. Les méthodes d'élagage de tokens entraînent une perte d'information et une dégradation des performances
3. La fusion de tokens traditionnelle (ToMe) utilise un taux de fusion fixe, incapable de s'adapter aux variations de distribution de similarité entre différentes étapes temporelles et couches

Motivation de la Recherche

Basée sur l'observation de deux phénomènes clés :

1. Il existe des variations significatives dans la distribution de similarité des tokens entre différentes étapes temporelles et couches
2. Les paires de tokens entre les étapes d'inférence adjacentes présentent une similarité élevée

Contributions Principales

1. Proposition d'un mécanisme de seuil adaptatif : Ajuste dynamiquement la stratégie de fusion selon la distribution de similarité des tokens, remplaçant le taux de fusion fixe
2. Conception d'un mécanisme de cache : Exploite la similarité entre les étapes adjacentes, met en cache les paires de tokens pour réduire les calculs répétés
3. Implémentation d'une accélération sans entraînement : La méthode peut être appliquée directement aux modèles pré-entraînés sans réentraînement
4. Réalisation d'un meilleur compromis qualité-vitesse : Par rapport à la méthode ToMe de base, réalise une inférence plus rapide tout en maintenant la qualité de l'image

Explication Détaillée de la Méthode

Définition de la Tâche

Entrée : Séquence de tokens dans le processus de débruitage du modèle de diffusion Sortie : Processus d'inférence accéléré avec fusion adaptative et optimisation de cache Contraintes : Maintenir une baisse non significative de la qualité de l'image générée

Architecture du Modèle

1. Mécanisme de Fusion de Tokens Adaptatifs

La méthode ToMe traditionnelle utilise un ratio fixe r pour la fusion de tokens, tandis que CA-ToMe introduit un seuil de similarité t :

Idée Centrale :

• Diviser l'image en régions de foulée de taille sx × sy
• Sélectionner le token du coin supérieur gauche de chaque région de foulée comme token cible
• Calculer la similarité cosinus entre les tokens sources et les tokens cibles
• Fusionner uniquement les paires de tokens dont la similarité dépasse le seuil t

Analyse des Avantages :

• Scénario A : Lorsque la plupart des tokens ont une similarité faible, un taux de fusion fixe force la fusion de tokens non similaires, entraînant une perte d'information. Le seuil adaptatif garantit que seuls les tokens hautement similaires sont fusionnés
• Scénario B : Lorsque la plupart des tokens sont hautement similaires (comme au début du débruitage), un taux de fusion fixe limite la quantité de fusion. Le seuil adaptatif permet de fusionner davantage de tokens, améliorant l'efficacité

2. Conception du Mécanisme de Cache

Basée sur l'analyse de la distance de Jaccard révélant une similarité élevée des paires de tokens entre les étapes adjacentes :

$JaccardDistance(A_n, A_{n+1}) = 1 - \frac{|A_n \cap A_{n+1}|}{|A_n \cup A_{n+1}|}$

où An représente l'ensemble de toutes les paires source-cible de tokens à l'étape n.

Stratégie d'Implémentation :

• Définir des points de contrôle (checkpoints), calculer les matrices de similarité uniquement à des étapes temporelles spécifiques
• Réutiliser les paires de tokens calculées précédemment aux étapes non-checkpoint
• Réduire significativement les frais généraux de calcul répété des matrices de similarité

Points d'Innovation Technique

1. Adaptabilité Dynamique : Ajuste automatiquement la stratégie de fusion selon la distribution de similarité, évitant les limitations des paramètres fixes
2. Optimisation de la Dimension Temporelle : Exploite la redondance entre les étapes temporelles, réduisant la quantité de calcul via le cache
3. Application Sélective au Niveau des Couches : Applique spécifiquement l'optimisation aux couches supérieures du U-Net (D1 et U1) qui sont intensives en calcul
4. Pas de Réentraînement Nécessaire : En tant que méthode d'accélération plug-and-play, peut être appliquée directement aux modèles existants

Configuration Expérimentale

Ensembles de Données

• Ensemble de données ImageNet-1k : Génération de 2000 images de résolution 512×512 (2 images par classe, 1000 classes au total)
• Ensemble de Validation : Utilisation de 5000 images de validation ImageNet-1k pour calculer le score FID
• Modèle de Prompt : « A high-quality photograph of a classname. »

Indicateurs d'Évaluation

1. FID (Distance d'Inception de Fréchet) : Indicateur principal pour mesurer la qualité des images générées
2. Temps d'Inférence : Temps moyen pour générer 2000 images
3. PSNR : Rapport Signal-Bruit de Crête, mesure la qualité de reconstruction au niveau des pixels
4. SSIM : Indice de Similarité Structurelle, évalue la cohérence spatiale et structurelle

Méthodes de Comparaison

• Baseline : Stable Diffusion v1.5 original
• ToMe : Méthode traditionnelle de fusion de tokens (r=50%)

Détails d'Implémentation

• Matériel : GPU Tesla V100S
• Étapes de Diffusion : 50 étapes d'échantillonnage PLMS
• Échelle CFG : 7,5
• Taille de Foulée : Fixée à 2×2
• Couches d'Application : Appliquée uniquement aux couches D1 et U1 du U-Net

Résultats Expérimentaux

Résultats Principaux

Découvertes Clés :

• CA-ToMe réalise la vitesse d'inférence la plus rapide (6,09s)
• Le score FID (34,05) surpasse ToMe (34,16) et se rapproche du baseline (33,66)
• Atteint le meilleur équilibre entre vitesse et qualité

Analyse des Paramètres de Seuil

Observations :

• Les variations dans la plage de seuil 0,4-0,6 sont faibles, car la plupart des tokens ont une similarité ≥0,6
• Le seuil 0,7 offre le meilleur compromis qualité-vitesse
• Les seuils plus élevés améliorent la qualité mais réduisent la vitesse

Comparaison des Configurations de Cache

CONFIG 3 offre les meilleures performances, cohérent avec l'analyse de la distance de Jaccard, avec plus de checkpoints définis aux étapes 8, 11, 13 et aux étapes finales.

Études d'Ablation

En comparant les contributions de différents composants :

1. Seuil Adaptatif Uniquement : Améliore la qualité de l'image par rapport au taux de fusion fixe
2. Mécanisme de Cache Uniquement : Réduit significativement le temps de calcul
3. CA-ToMe Complet : La combinaison des deux techniques réalise les meilleures performances

Travaux Connexes

Méthodes d'Accélération des Modèles de Diffusion

1. Réduction du nombre d'étapes d'échantillonnage :
   • Méthodes de distillation de connaissances 26,51,28
   • Échantillonnage implicite 32
   • Solveurs d'équations différentielles avancés 52,33
   • La plupart nécessitent un réentraînement
2. Réduction du calcul par étape :
   • Méthodes de quantification 31,36
   • Réduction de tokens 21,40,41,43,44
   • Techniques de cache 24,37,38,39
   • Plug-and-play, sans réentraînement nécessaire

Techniques de Réduction de Tokens

1. Élagage de Tokens : Suppression directe des tokens non importants, pouvant entraîner une perte d'information
2. Fusion de Tokens : Fusion des tokens similaires, préservant l'intégrité de l'information
   • ToMe 21 : Utilise un taux de fusion fixe
   • CA-ToMe de cet article : Seuil adaptatif + mécanisme de cache

Techniques de Cache

Les méthodes de cache existantes ciblent différents composants :

• Cache d'attention croisée 38
• Cache d'encodeur U-Net 39
• Cache de caractéristiques avancées 24

Cet article est le premier à appliquer le cache au calcul de similarité de la fusion de tokens.

Conclusion et Discussion

Conclusions Principales

1. Le seuil adaptatif résout efficacement les limitations du taux de fusion fixe, ajustant dynamiquement la stratégie de fusion selon la distribution de similarité
2. Le mécanisme de cache exploite la redondance entre les étapes temporelles, réduisant significativement les calculs répétés
3. La méthode CA-ToMe réalise une accélération de 1,24×, tout en maintenant voire en améliorant légèrement la qualité de l'image
4. La caractéristique sans entraînement confère à la méthode une bonne praticité et scalabilité

Limitations

1. Ajustement des paramètres de seuil : Nécessite d'ajuster le seuil optimal pour différents modèles et tâches
2. Portée d'application limitée : Principalement destinée aux modèles de diffusion avec architecture U-Net
3. Frais généraux de cache : Nécessite de la mémoire supplémentaire pour stocker les informations des paires de tokens en cache
4. Limitation au niveau des couches : Appliquée uniquement aux couches supérieures, peut manquer les opportunités d'optimisation dans d'autres couches

Directions Futures

1. Apprentissage automatique du seuil : Développer des méthodes pour déterminer automatiquement le seuil optimal
2. Extension à d'autres architectures : S'adapter aux nouvelles architectures de modèles de diffusion comme DiT
3. Stratégies de cache plus raffinées : Mécanismes de cache adaptatifs basés sur le contenu
4. Optimisation matérielle : Implémentations optimisées pour des matériels spécifiques

Évaluation Approfondie

Points Forts

1. Innovation Forte : Introduit la pensée adaptative dans la fusion de tokens, formant une solution complète en combinaison avec le mécanisme de cache
2. Valeur Pratique Élevée : Les caractéristiques sans entraînement et plug-and-play la rendent facile à déployer
3. Expériences Complètes : Des études d'ablation complètes et une analyse paramétrique soutiennent l'efficacité de la méthode
4. Fondations Théoriques Solides : L'analyse de similarité basée sur la distance de Jaccard fournit un soutien théorique au mécanisme de cache

Insuffisances

1. Analyse Théorique Insuffisante : Manque de guidance théorique pour la sélection du seuil adaptatif
2. Portée Expérimentale Limitée : Validation uniquement sur ImageNet, manque d'évaluation sur d'autres ensembles de données et tâches
3. Méthodes de Comparaison Limitées : Comparaison principalement avec ToMe, manque de comparaison avec d'autres méthodes d'accélération
4. Évaluation de Qualité Unique : Dépend principalement de la métrique FID, manque d'évaluation humaine et d'autres indicateurs de qualité

Impact

1. Contribution Académique : Fournit de nouvelles perspectives et méthodes pour l'accélération des modèles de diffusion
2. Valeur Pratique : Peut être appliquée directement aux modèles de diffusion existants, avec des perspectives d'application larges
3. Reproductibilité : Fournit une implémentation de code complète, facilitant la reproduction et l'extension
4. Nature Inspirante : Les idées d'adaptabilité et de cache peuvent inspirer davantage de recherches connexes

Scénarios d'Application

1. Environnements aux Ressources Limitées : Appareils mobiles, informatique en périphérie, etc.
2. Applications en Temps Réel : Applications interactives nécessitant une génération d'images rapide
3. Déploiement à Grande Échelle : Réduction des coûts de calcul des serveurs et de la latence
4. Prototypes de Recherche : Fourniture de composants de base pour d'autres techniques d'accélération

Références

Cet article cite 54 références connexes, incluant principalement :

• Théories fondamentales des modèles de diffusion 1,2,3
• Applications de génération d'images 4,5,18,19,20
• Techniques d'accélération 24,25,26,27,28
• Méthodes de traitement de tokens 21,40,41,43,44
• Techniques de cache 24,37,38,39

Évaluation Globale : Ceci est un travail ayant une valeur pratique dans le domaine de l'accélération des modèles de diffusion. Grâce à la combinaison ingénieuse d'un seuil adaptatif et d'un mécanisme de cache, il réalise une amélioration significative de la vitesse tout en maintenant la qualité de l'image. Bien qu'il y ait de la place pour l'amélioration dans l'analyse théorique et la portée expérimentale, sa nature sans entraînement et ses bons résultats expérimentaux lui confèrent une valeur pratique et un impact élevés.