TMPDiff: Temporal Mixed-Precision for Diffusion Models

Résumé Technique : TMPDiff

1. Problématique

Les modèles de diffusion, en particulier pour la génération d'images à partir de texte (Text-to-Image), souffrent d'une latence d'inférence élevée due à leur processus itératif de débruitage. Bien que la quantification (réduction de la précision des nombres, par exemple de 16 bits à 4 bits) soit une méthode courante pour accélérer l'inférence, elle dégrade généralement la qualité de l'image générée.

Le problème central identifié par les auteurs est que les méthodes de quantification existantes appliquent une précision fixe (uniforme) sur l'ensemble des étapes de débruitage. Or, la sensibilité aux erreurs de quantification varie considérablement selon l'étape temporelle du processus. Les approches actuelles ignorent cette dimension temporelle, laissant un axe d'optimisation inexploité : l'attribution dynamique de différentes précisions numériques à différentes étapes de la diffusion.

2. Méthodologie

Les auteurs proposent TMPDiff, un cadre de travail pour la précision mixte temporelle. L'approche repose sur trois piliers fondamentaux :

• Modèle d'Erreur Additive :
  Les auteurs émettent l'hypothèse que les erreurs de quantification s'accumulent de manière additive à travers les étapes de débruitage. Ils dérivent un modèle mathématique (basé sur une approximation de Taylor et la récursion de l'erreur) montrant que l'erreur finale $\delta_0$ peut être exprimée comme une somme pondérée des erreurs locales injectées à chaque étape $t$, pondérée par le schéma de précision choisi.
  $$\boldsymbol{\delta}_0(Z) \approx \sum_{t=1}^{T} \left( \prod_{j=1}^{t-1} \mathbf{A}_j \right) B_t \boldsymbol{\epsilon}_t (1 - z_t)$$
  Où $z_t$ indique si l'étape $t$ est en pleine précision ($z_t=1$) ou quantifiée ($z_t=0$).

• Validation Empirique :
  Cette hypothèse d'additivité est validée expérimentalement sur plusieurs modèles et jeux de données. Les résultats montrent une forte corrélation linéaire entre les gains de précision d'une seule étape (upcasting) et la réduction globale de l'erreur. Cela permet de prédire la performance d'un schéma complexe en se basant uniquement sur les mesures d'erreur d'étapes individuelles.

• Algorithme de Recherche Adaptative (Bissection) :
  Trouver le schéma de précision optimal ($Z^*$) pour un budget de latence donné est un problème de recherche combinatoire exponentiel ($2^T$). Grâce au modèle d'erreur additive, ce problème devient gérable.

  • Les auteurs développent un algorithme de bissection adaptative. Au lieu d'évaluer toutes les étapes, ils échantillonnent d'abord les étapes clés (début, milieu, fin) et identifient les segments où le gain de précision est le plus élevé.
  • L'algorithme affine ces segments de manière itérative pour reconstruire la courbe de gain d'erreur avec un nombre minimal d'évaluations (complexité linéaire par rapport au nombre d'échantillons, évitant la recherche exhaustive).
  • Une fois les gains d'erreur estimés pour chaque étape, l'algorithme sélectionne de manière gourmande (greedy) les $K$ étapes les plus sensibles pour les exécuter en pleine précision, tandis que le reste est quantifié.

3. Contributions Clés

1. Modèle d'Erreur Additive et Validation : Proposition et validation empirique d'un modèle montrant que les erreurs de quantification s'accumulent de manière additive dans le latent final, permettant une modélisation analytique du problème.
2. Framework TMPDiff : Introduction d'un système de planification de précision temporelle qui utilise un algorithme de bissection adaptative pour identifier la répartition optimale des précisions en temps linéaire, transformant un problème de recherche exponentielle en une procédure efficace.
3. Généralité de la Sensibilité Temporelle : Démonstration que le score de sensibilité par étape capture une propriété fondamentale du processus de débruitage, applicable non seulement à la quantification, mais aussi au mélange de modèles de différentes tailles (ex: utiliser un petit modèle pour les étapes peu sensibles et un grand modèle pour les étapes critiques).

4. Résultats Expérimentaux

Les expériences ont été menées sur quatre modèles de pointe (FLUX.1-dev, PixArt-Σ, Sana, StableDiffusion XL) et trois jeux de données (COCO, DCI, MJHQ).

• Amélioration de la Qualité : À vitesse d'inférence égale, TMPDiff surpasse systématiquement les baselines à précision uniforme, offrant une amélioration de 10 % à 20 % des métriques de qualité perceptuelle (SSIM, LPIPS, PSNR).
• Performance sur FLUX.1-dev : Le modèle atteint 90 % de SSIM par rapport au modèle en pleine précision tout en offrant un accélération de 2,5× par rapport à une inférence standard en 16 bits.
• Efficacité Latence-Qualité : TMPDiff remplit efficacement le compromis (trade-off) entre la qualité et la vitesse. Par exemple, en utilisant 80 % d'étapes quantifiées, il améliore la qualité de 10 points de pourcentage (SSIM) avec un impact minimal sur la latence.
• Mixité de Modèles : Dans une expérience de mélange de modèles (Sana 1.6B et 600M), TMPDiff a surpassé les méthodes heuristiques et de recherche évolutionnaire existantes, prouvant que la sensibilité temporelle est un critère robuste au-delà de la simple quantification.

5. Signification et Impact

Ce travail représente une avancée significative dans l'optimisation des modèles de diffusion génératifs :

• Nouvel Axe d'Optimisation : Il déplace le paradigme de la quantification statique (spatiale) vers une quantification dynamique temporelle, exploitant la nature itérative des modèles de diffusion.
• Efficacité Computationnelle : En réduisant le problème de recherche de schéma optimal de exponentiel à linéaire, TMPDiff rend l'optimisation de précision réalisable même pour des modèles massifs (plusieurs milliards de paramètres).
• Déploiement Pratique : La méthode permet de réduire la latence d'inférence de 60 % avec une dégradation de qualité minimale, rendant les modèles de diffusion de pointe plus accessibles pour des applications en temps réel ou sur des dispositifs aux ressources limitées.

En conclusion, TMPDiff démontre que l'adaptation dynamique de la précision numérique au cours du temps de débruitage est une stratégie puissante pour concilier haute fidélité d'image et efficacité computationnelle.