💻 computer science

CoReDiT: Spatial Coherence-Guided Token Pruning and Reconstruction for Efficient Diffusion Transformers

CoReDiT est un cadre efficace de pruning de tokens pour les Transformers de diffusion qui exploite la cohérence spatiale pour identifier et sauter les tokens redondants tout en reconstruisant leurs sorties via des tokens voisins, permettant ainsi des accélérations computationnelles et des économies de mémoire significatives sans compromettre la qualité visuelle.

Auteurs originaux : Zhuojin Li, Hsin-Pai Cheng, Hong Cai, Shizhong Han, Fatih Porikli

Publié 2026-05-15

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Zhuojin Li, Hsin-Pai Cheng, Hong Cai, Shizhong Han, Fatih Porikli

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayiez de peindre une fresque massive et incroyablement détaillée sur un mur. Vous avez une équipe d'artistes (les « tokens ») debout devant le mur, et chaque artiste doit parler à tous les autres pour décider exactement quelle couleur appliquer sur leur emplacement spécifique. Cela garantit que l'image est parfaite, mais c'est épuisant. Si vous avez 10 000 artistes, ils doivent tenir 100 millions de conversations juste pour s'accorder sur le prochain coup de pinceau. C'est ainsi que fonctionnent les générateurs d'images IA actuels (appelés Diffusion Transformers) : ils sont brillants, mais incroyablement lents et voraces en puissance de calcul.

L'article présente une nouvelle méthode appelée CoReDiT pour accélérer ce processus sans gâcher la peinture. Voici comment cela fonctionne, décomposé en trois étapes simples :

1. Le détecteur d'« artiste ennuyé » (Cohérence spatiale)

Dans une peinture typique, certaines zones sont très animées et détaillées (comme un visage ou un arbre), tandis que d'autres sont de simples arrière-plans ennuyeux (comme un ciel bleu ou un mur lisse).

Le problème : L'ancienne IA traite tous les artistes de la même manière, obligeant ceux qui peignent le ciel bleu ennuyeux à avoir les mêmes conversations intenses que ceux qui peignent le visage.
La solution CoReDiT : CoReDiT observe les artistes et demande : « Qui se tient à côté de quelqu'un qui leur ressemble exactement ? » Si un artiste se trouve dans une zone ennuyeuse et uniforme, ses voisins sont pratiquement des clones. CoReDiT attribue à ces artistes « ennuyés » un score de cohérence élevé.
L'action : Il dit aux artistes à score élevé (redondants) : « Vous n'avez pas besoin de parler à tout le monde pour l'instant. Faites une pause. » Cela saute les conversations coûteuses pour les parties ennuyeuses de l'image, économisant ainsi un temps considérable.

2. L'artiste « remplissage de trous » (Reconstruction)

Si vous disiez simplement aux artistes « ennuyés » de rentrer chez eux et d'arrêter de travailler, vous vous retrouveriez avec des trous dans votre fresque. L'image paraîtrait brisée et patchée.

Le problème : Supprimer simplement des tokens (artistes) crée des artefacts visuels.
La solution CoReDiT : Au lieu de laisser un trou, CoReDiT utilise les artistes qui sont restés pour combler les vides. Parce que les artistes « ennuyés » étaient si similaires à leurs voisins, les artistes restants peuvent facilement deviner ce que les absents auraient peint.
L'action : C'est comme un jeu de « téléphone arabe » où les artistes restants chuchotent la bonne couleur aux emplacements vides. Cela garantit que l'image finale est lisse et complète, sans lacunes visibles, même si moins d'artistes ont réellement fourni l'effort principal.

3. Le gestionnaire intelligent (Élagage progressif)

Toutes les étapes du processus de peinture ne sont pas identiques.

Le problème : Tout au début de la création d'une image, celle-ci n'est que du bruit statique (comme la neige à la télévision), donc presque tout se ressemble et est redondant. Mais tout à la fin, lorsque les détails sont affinés, chaque coup de pinceau compte. Si vous éliminez trop d'artistes trop tôt, l'IA se confond. Si vous en éliminez trop à la fin, l'image devient floue.
La solution CoReDiT : CoReDiT agit comme un gestionnaire intelligent qui apprend l'emploi du temps parfait. Il commence par laisser l'IA éliminer de nombreux artistes redondants lorsque l'image n'est que du bruit. À mesure que l'image devient plus claire et plus détaillée, le gestionnaire cesse progressivement d'éliminer autant d'artistes, veillant à ce que les détails critiques reçoivent toute l'attention nécessaire.
L'action : Il détermine automatiquement quelles parties du cerveau de l'IA (blocs) et quelles étapes du processus (pas de temps) peuvent supporter le plus d'éliminations, ajustant le plan à la volée pour maintenir une qualité élevée.

Les résultats

En utilisant cette méthode, les auteurs ont constaté que :

Vitesse : L'IA fonctionne 1,33 fois plus vite sur des ordinateurs cloud puissants et 1,72 fois plus vite sur les puces de téléphones mobiles.
Efficacité : Elle réduit le travail mathématique lourd (calculs) jusqu'à 55 %.
Qualité : Les images restent tout aussi bonnes que la version lente, sans fissures évidentes ni artefacts étranges.
Bonus : Parce qu'elle utilise moins de mémoire, elle peut même générer des images de résolution plus élevée sur des téléphones qui planteraient normalement par manque de mémoire.

En bref, CoReDiT est comme embaucher un contremaître intelligent pour un projet artistique qui sait exactement quand laisser l'équipe se reposer et quand la faire travailler dur, aboutissant à un chef-d'œuvre réalisé beaucoup plus rapidement.

Résumé Technique : CoReDiT

Énoncé du Problème

Les Transformers de Diffusion (DiT) ont établi des performances de pointe dans la génération d'images et de vidéos en modélisant efficacement les dépendances à long terme grâce aux mécanismes d'attention. Cependant, leur coût computationnel évolue de manière quadratique avec le nombre de tokens, rendant la synthèse haute résolution et la génération de vidéos multi-images prohibitivement coûteuses pour les plateformes mobiles aux ressources limitées. Bien qu'une part significative des tokens dans l'espace latent corresponde à des régions visuellement redondantes (par exemple, des arrière-plans uniformes), les méthodes existantes de pruning de tokens font face à trois défis principaux lorsqu'elles sont appliquées à l'inférence de diffusion :

Localisation Efficace : Identifier les tokens à faible saillance sans encourir la surcharge quadratique du calcul des scores d'attention globale.
Préservation Sémantique : Maintenir la continuité visuelle et prévenir les artefacts (par exemple, ruptures de texture, distorsions de frontières) lors du saut de tokens, ce qui est critique pour les tâches génératives contrairement aux tâches discriminatives.
Planification Adaptative : Déterminer une stratégie de pruning qui s'adapte aux niveaux de redondance variables à travers différents blocs de transformateurs et étapes de débruitage.

Méthodologie : CoReDiT

CoReDiT est un cadre structuré de pruning de tokens conçu pour accélérer les DiT pré-entraînés tout en préservant la fidélité de la sortie. Il fonctionne via trois composants principaux :

1. Sélection de Tokens Basée sur la Cohérence Spatiale

Au lieu de s'appuyer sur des scores d'attention globale (qui sont de $O(N^2)$ ), CoReDiT utilise un score de Cohérence Spatiale (SC) pour estimer la redondance locale en temps linéaire ( $O(N)$ ).

Mécanisme : Le réseau de tokens est partitionné en grilles non chevauchantes. Pour chaque token, le score SC est calculé comme la moyenne de la similarité cosinus entre l'embedding normalisé du token et l'embedding normalisé moyen de sa grille.
Logique : Les tokens avec des scores SC élevés sont jugés redondants car ils sont bien expliqués par leurs voisins spatiaux.
Sélection : Les $K$ tokens ayant les scores SC les plus élevés sont sélectionnés pour être sautés ( $X_S$ ), tandis que les tokens restants ( $X_R$ ) procèdent à la couche d'Attention Multi-Têtes (MHSA).

2. Reconstruction de Tokens Guidée par la Cohérence

Pour prévenir les discontinuités spatiales causées par le saut de tokens, CoReDiT reconstruit les sorties d'attention pour les tokens sautés ( $Y_S$ ) en utilisant les sorties des voisins conservés ( $Y_R$ ).

Approximation : Plutôt que de calculer des similarités par paires pour la reconstruction (ce qui ajouterait une surcharge), la méthode approxime les poids de similarité en utilisant les scores SC pré-calculés des voisins.
Sous-grilles Locales : Pour éviter un lissage excessif, la reconstruction est effectuée au sein de sous-grilles plus petites.
Micro-Designs :
- Tailles de Grilles Alternées : Les dimensions de grille alternent entre les blocs consécutifs pour décaler les artefacts de frontière et permettre la propagation d'information à travers les bords de grille.
- Conservation par Pas m (m-Stride) : Un motif de pas déterministe garantit qu'au moins un token est conservé dans chaque sous-grille, fournissant les ancres nécessaires à la reconstruction.

3. Planification de Pruning Progressive et Adaptative aux Blocs

Reconnaissant que la redondance varie à travers les blocs de transformateurs et les étapes de débruitage, CoReDiT emploie une stratégie de pruning progressive lors du fine-tuning léger.

Adaptativité aux Blocs : Le budget de pruning est alloué aux blocs présentant la redondance mesurée la plus élevée (estimée par la somme des scores SC des tokens candidats).
Décroissance par Pas de Temps : Le pruning est plus agressif dans les premières étapes de débruitage (où le bruit domine et la structure est grossière) et réduit dans les étapes ultérieures (où les détails fins sont affinés). Un calendrier de décroissance en deux phases est utilisé pour équilibrer l'efficacité avec la complexité de la gestion des graphes de calcul spécifiques aux étapes.
Entraînement : Le modèle est fine-tuné en utilisant une fonction de perte de distillation qui combine la perte de sortie du modèle original et une perte de bloc pruné pour récupérer les performances.

Contributions Clés

Estimation de Redondance en Temps Linéaire : L'introduction d'un score de cohérence spatiale qui identifie les tokens redondants en temps $O(N)$ , évitant le goulot d'étranglement computationnel de l'analyse d'attention globale.
Préservation du Réseau Dense : Un mécanisme de reconstruction qui synthétise les sorties d'attention sautées à partir des voisins, assurant que le réseau de tokens reste dense pour les couches suivantes et empêchant les artefacts visuels.
Stratégie de Pruning Adaptative : Un calendrier progressif et adaptatif aux blocs qui alloue dynamiquement les budgets de pruning basés sur les statistiques de redondance empiriques à travers la profondeur et le temps, facilitant une adaptation stable.

Résultats Expérimentaux

CoReDiT a été évalué sur des architectures de DiT de pointe, notamment PixArt-α, PixArt-Σ et MagicDrive-V2.

Gains d'Efficacité :
- Réduction allant jusqu'à 55 % des FLOPs d'attention auto.
- Accélérations d'Inférence : 1,33× sur les GPU cloud (Nvidia H100) et 1,72× sur les NPU mobiles (Qualcomm Snapdragon 8 Elite).
- Marge Mémoire : Permet la génération à des résolutions allant jusqu'à 1920 sur des appareils mobiles, là où le modèle de base rencontrait des erreurs de Mémoire Insuffisante (OOM) à 1600.
Préservation de la Qualité :
- Sur PixArt-α-1024 avec 40 % de pruning, CoReDiT a maintenu un FID de 28,7 (contre 27,3 pour la base) et des scores CLIP/IS compétitifs. Avec la distillation, l'écart de FID s'est réduit à 27,4.
- Sur PixArt-Σ-2048, un taux de pruning moyen de 23 % a réduit la latence d'attention auto de 22 % avec seulement une baisse modeste des métriques de qualité.
- Sur MagicDrive-V2 (vidéo), une réduction de 39 % des FLOPs d'attention auto a préservé la qualité par image (PSNR, LPIPS, SSIM) et l'alignement conditionnel (mAP, mIoU), bien qu'une légère baisse de la cohérence temporelle (FVD) ait été notée.
Études d'Ablation :
- La sélection aléatoire de tokens a entraîné une dégradation catastrophique de la qualité (FID 644,5), validant la nécessité d'une sélection basée sur la cohérence.
- La désactivation de la reconstruction a conduit à une perte de qualité significative (FID 30,4 contre 28,5), confirmant l'importance de l'étape de reconstruction.
- Le pruning uniforme suivi d'un fine-tuning a donné de moins bons résultats que le calendrier adaptatif progressif proposé.

Importance et Revendications

L'article revendique que CoReDiT offre une solution pratique pour le déploiement de modèles de diffusion haute fidélité sur des appareils périphériques en abordant les défis spécifiques du pruning de tokens génératifs. Son importance réside dans :

Évolutivité : Permettre la synthèse haute résolution et la génération vidéo sur du matériel mobile en réduisant à la fois l'empreinte de calcul et la mémoire.
Compatibilité Architecturale : Maintenir la compatibilité avec les architectures DiT existantes sans nécessiter de modifications structurelles ou de modules de routage complexes.
Compromis Efficacité-Qualité : Démontrer que le saut structuré et guidé par la cohérence peut réduire drastiquement les coûts d'inférence (jusqu'à 55 % des FLOPs d'attention auto) tout en maintenant une qualité perceptuelle comparable au modèle complet.

Les auteurs notent que, bien que l'approche actuelle se concentre sur la redondance spatiale, les travaux futurs visent à étendre l'efficacité aux composants FFN/MLP et à intégrer des objectifs de cohérence inter-images pour améliorer davantage la cohérence temporelle des vidéos.