HybridStitch: Pixel and Timestep Level Model Stitching for Diffusion Acceleration

Le papier présente HybridStitch, une nouvelle méthode d'accélération pour la génération d'images par diffusion qui combine un grand et un petit modèle en traitant la création comme une édition, en utilisant le petit modèle pour les zones simples et le grand pour les zones complexes, permettant ainsi un gain de vitesse de 1,83 fois sur Stable Diffusion 3.

L'essentiel

Imaginez que vous commandez un tableau magnifique à un artiste très célèbre, mais très lent (le Grand Modèle). Il prend des heures pour peindre chaque détail, du ciel aux visages. Maintenant, imaginez un assistant rapide mais un peu moins talentueux (le Petit Modèle).

Le problème, c'est que si vous demandez au Grand Modèle de tout faire, cela prend trop de temps. Si vous demandez à l'assistant de tout faire, c'est trop rapide mais le résultat est moche.

Les méthodes précédentes disaient : "Ok, on laisse le Grand Modèle faire les 40 premières minutes, puis on donne le pinceau à l'assistant pour le reste." C'est mieux, mais pas parfait. Pourquoi ? Parce que dans une image, certaines parties sont faciles (un ciel bleu uni) et d'autres sont très complexes (les yeux d'un chien, les plis d'un vêtement).

C'est ici qu'intervient HybridStitch, la nouvelle méthode présentée dans cet article. Voici comment ça marche, avec des images simples :

1. L'idée principale : Le chef d'orchestre et les zones de travail

Au lieu de changer de modèle pour toute l'image d'un coup, HybridStitch agit comme un chef d'orchestre très intelligent qui divise le travail pixel par pixel.

• Le Grand Modèle (l'Expert) : Il ne travaille que sur les zones difficiles. C'est comme si un architecte ne venait vérifier que les fondations et les poutres complexes d'une maison, laissant le reste aux ouvriers.
• Le Petit Modèle (l'Assistant Rapide) : Il travaille sur toute l'image, mais surtout sur les zones faciles (le ciel, l'herbe, les murs lisses). Il fait le gros œuvre rapidement.

2. Le processus en trois étapes (La "Couture")

Imaginez que vous dessinez une image étape par étape, en partant d'un brouillard flou jusqu'à une image nette.

• Étape 1 : Le Brouillard Initial
  Au tout début, tout est flou. L'Expert (Grand Modèle) travaille sur toute l'image pour poser les bases, comme un architecte qui trace les murs principaux.

• Étape 2 : La Division du Travail (Le cœur de la méthode)
  C'est là que la magie opère. Le système regarde l'image et se dit : "Tiens, le ciel est déjà assez joli, l'assistant peut s'en occuper. Mais les yeux du chat sont encore flous, l'expert doit les finir."

  • L'Assistant (Petit Modèle) prend le relais sur tout l'image pour avancer vite.
  • L'Expert (Grand Modèle) intervient uniquement sur les zones difficiles (les yeux du chat, les détails complexes).
  • L'astuce secrète (Le Cache KV) : Pour que l'Expert ne perde pas le fil et ne fasse pas des erreurs de contexte (comme dessiner un chat avec un nez de chien), le système lui "colle" les souvenirs des étapes précédentes. C'est comme si l'Expert avait une photo de l'étape précédente collée sur son bureau pour savoir où il en était, même s'il ne travaille que sur une petite partie du dessin.

• Étape 3 : Le Finissage
  Une fois que les zones difficiles sont assez nettes, même l'Expert peut arrêter de travailler. L'Assistant prend alors le relais sur tout l'image pour polir les derniers détails jusqu'à la fin.

3. Pourquoi c'est génial ? (Les résultats)

• Vitesse Éclair : Parce que l'Expert ne travaille que sur une petite partie de l'image la plupart du temps, le processus est beaucoup plus rapide. L'article montre une accélération de 1,83 fois par rapport à l'utilisation du Grand Modèle seul. C'est comme passer d'une voiture de ville à une Formule 1.
• Qualité Préservée : Même si on va plus vite, l'image reste belle. Pourquoi ? Parce que les zones complexes (les plus importantes pour la beauté de l'image) ont toujours eu l'attention de l'Expert. On ne sacrifie pas la qualité pour la vitesse.
• Pas besoin de réapprendre : Cette méthode ne nécessite pas de réentraîner les modèles. C'est comme un nouveau logiciel qui s'installe sur vos outils existants pour les rendre plus efficaces.

En résumé

HybridStitch, c'est comme avoir un atelier de peinture où le Maître Peintre ne touche au pinceau que là où c'est vraiment nécessaire (les visages, les détails), tandis qu'un apprenti rapide remplit le reste du tableau (le fond, le ciel). Le résultat ? Une œuvre d'art magnifique, produite en moitié moins de temps.

C'est une solution intelligente qui comprend que toutes les parties d'une image ne demandent pas le même effort, et elle adapte le travail en conséquence.

Résumé technique

1. Problème Identifié

Les modèles de diffusion pour la génération d'images à partir de texte (Text-to-Image ou T2I) ont démontré une qualité exceptionnelle, mais souffrent d'une surcharge computationnelle importante, en particulier pour les modèles massifs (plusieurs milliards de paramètres comme Stable Diffusion 3 ou 3.5). Cette latence élevée constitue un obstacle majeur pour les applications sensibles au temps de réponse.

Les travaux antérieurs visant à accélérer l'inférence ont proposé des méthodes de "mélange de modèles" (Mixture of Models), où l'on alterne entre un grand modèle (haute qualité) et un petit modèle (rapide) au cours des étapes de débruitage. Cependant, ces approches existantes (comme T-Stitch ou SRDiffusion) opèrent à une granularité d'image entière : elles basculent d'un modèle à l'autre pour l'image complète à un instant donné.

Le problème central réside dans l'hétérogénéité des besoins computationnels au sein d'une même image et d'une même étape de temps (timestep). Certaines régions (ex: arrière-plan) sont faciles à générer et pourraient être traitées par un petit modèle plus tôt, tandis que d'autres (ex: objets complexes) nécessitent la précision du grand modèle plus longtemps. Les méthodes actuelles, en ignorant cette diversité au niveau des pixels, subissent soit une dégradation de la qualité (si le basculement est trop précoce), soit une latence inutilement élevée (si le basculement est retardé jusqu'à ce que toute l'image soit prête).

2. Méthodologie : HybridStitch

HybridStitch propose un nouveau paradigme de génération qui traite la création d'image comme un processus d'édition, en effectuant un basculement de modèle au niveau du pixel et du timestep.

A. Concept Fondamental

L'idée centrale est de diviser l'image en deux régions dynamiques à chaque étape de débruitage :

1. Régions complexes : Continuent d'être raffinées par le grand modèle.
2. Régions simples : Transitionnent vers le petit modèle pour construire un croquis global et assurer la cohérence.

B. Flux de Travail (3 Étapes)

L'algorithme procède en trois phases distinctes :

1. Phase 1 (Initialisation) : Le grand modèle traite l'ensemble du bruit gaussien pour établir la structure globale de l'image.
2. Phase 2 (Hybride) :
   • Le petit modèle traite l'image entière pour générer une ébauche (draft) rapide.
   • Le grand modèle ne traite que les pixels masqués (les régions difficiles à générer).
   • Astuce technique (KV Cache) : Pour éviter que le grand modèle, en ne voyant qu'une partie de l'image, ne perde le contexte global, HybridStitch réutilise les caches Key-Value (KV) de l'étape précédente pour combler les parties manquantes lors du calcul de l'attention. Cela garantit la cohérence sémantique.
   • La prédiction finale de l'étape est une combinaison : les régions masquées utilisent la sortie du grand modèle, le reste utilise celle du petit modèle.
3. Phase 3 (Finalisation) : Une fois que la différence entre les sorties des deux modèles devient négligeable pour l'ensemble de l'image, le grand modèle est désactivé et seul le petit modèle termine les étapes restantes.

C. Stratégie de Basculement et Masquage

• Détection de différence : Le système calcule la différence (distance L1) entre les états latents de deux étapes de débruitage consécutives.
• Sélection de Masque (Top-K) : À chaque étape, les pixels présentant la plus grande différence (indiquant une instabilité ou un changement significatif) sont sélectionnés pour former le masque. Ce masque est mis à jour dynamiquement à chaque étape.
• Critères de transition : Le passage d'une phase à l'autre est déclenché lorsque la différence moyenne ($D_t$) tombe en dessous de seuils prédéfinis.

3. Contributions Clés

1. Granularité Pixel et Timestep : Première méthode à permettre un basculement de modèle non pas sur l'image entière, mais sur des régions spécifiques (pixels) à chaque étape de temps.
2. Technique sans entraînement (Train-free) : HybridStitch ne nécessite pas de réentraînement des modèles ; il s'agit d'une technique d'accélération d'inférence pure.
3. Mécanisme de Contexte KV : Introduction d'une méthode ingénieuse utilisant le cache KV pour permettre au grand modèle de traiter uniquement des sous-ensembles de pixels tout en maintenant une cohérence globale, évitant ainsi les artefacts de contexte.
4. Paradigme "Édition" : Redéfinition de la génération comme un processus où le grand modèle "édite" et affine les zones complexes tandis que le petit modèle gère le reste.

4. Résultats Expérimentaux

Les évaluations ont été menées sur Stable Diffusion 3 (Grand modèle : SD 3.5 Large, Petit modèle : SD 3 Medium) sur le dataset COCO (5000 échantillons).

• Accélération : HybridStitch atteint un speedup de 1,83x par rapport à l'utilisation exclusive du grand modèle.
  • Comparé à SRDiffusion (l'état de l'art précédent), HybridStitch est 18,06 % plus rapide tout en maintenant une qualité égale ou supérieure.
  • Comparé à T-Stitch, le gain est encore plus significatif.
• Qualité :
  • FID (Fréchet Inception Distance) : HybridStitch obtient un score de 30,43 (pour la configuration 30 %), surpassant T-Stitch (31,87) et SRDiffusion (31,67). Un FID plus bas indique une meilleure qualité.
  • CLIP Score : Reste stable autour de 31,84, indiquant une bonne adhérence sémantique au prompt.
  • LPIPS : Mesure la similarité perceptuelle avec le modèle original ; HybridStitch montre une très faible dégradation.
• Robustesse : Les résultats sont stables sur différentes tailles de masques (10 % à 40 %) et sur différents types de GPU (RTX 6000, H100, A100). L'avantage est particulièrement marqué sur les plateformes moins puissantes.

5. Signification et Impact

HybridStitch représente une avancée significative dans l'optimisation des modèles de diffusion génératifs. En reconnaissant que la complexité de génération n'est pas uniforme à travers une image, la méthode permet de démocratiser l'accès aux grands modèles en réduisant drastiquement leur coût computationnel sans sacrifier la qualité visuelle.

Son approche "region-aware" (consciente des régions) ouvre la voie à des stratégies d'inférence plus intelligentes et adaptatives, dépassant les limitations des méthodes de basculement global. Cela rend possible le déploiement de modèles T2I de haute qualité sur du matériel moins puissant ou dans des contextes nécessitant une faible latence, tout en conservant la richesse des modèles à grande échelle pour les détails critiques.