OBS-Diff: Accurate Pruning For Diffusion Models in One-Shot

Ce papier présente OBS-Diff, un cadre de pruning one-shot et sans entraînement qui adapte l'algorithme Optimal Brain Surgeon aux modèles de diffusion text-to-image en utilisant une construction de Hessien sensible aux timesteps et une stratégie de pruning séquentiel par groupes pour réduire les coûts de calcul tout en préservant la qualité visuelle.

L'essentiel

🎨 Le Problème : Des Géants Trop Lents

Imaginez que les modèles de diffusion (comme ceux qui créent des images à partir de texte, par exemple "un chat en astronaut") sont comme des géants artistiques. Ils sont incroyablement talentueux et peuvent peindre des chefs-d'œuvre, mais ils sont énormes et lourds.

Pour peindre une seule image, ce géant doit faire des milliers de petits pas (des "itérations de débruitage"). C'est comme si un peintre devait effacer et redessiner sa toile 50 fois avant de montrer le résultat final.

• Le problème : Ces géants prennent trop de temps et d'énergie (mémoire de l'ordinateur) pour fonctionner. Ils sont difficiles à utiliser sur un simple ordinateur portable ou un téléphone.

✂️ La Solution Habituelle : La Taille "Au Couteau"

Pour rendre ces géants plus légers, les chercheurs essaient de les "élaguer" (pruning), c'est-à-dire de retirer les parties inutiles de leur cerveau (les poids du réseau neuronal).

• L'ancienne méthode : C'est comme si on prenait un couteau et qu'on coupait les branches les plus petites d'un arbre au hasard, ou simplement celles qui semblent les plus faibles visuellement.
• Le résultat : Souvent, on coupe une branche qui semblait petite mais qui était en fait vitale pour la structure de l'arbre. L'arbre (le modèle) s'effondre ou produit des images bizarres et floues. De plus, pour réparer les dégâts, il faut souvent réentraîner l'arbre pendant des jours, ce qui est très coûteux.

🚀 OBS-Diff : Le "Chirurgien Optimal" en Une Seule Passe

L'article présente OBS-Diff, une nouvelle méthode qui agit comme un chirurgien de précision capable de réduire le poids du modèle en une seule passe, sans avoir besoin de le réentraîner.

Voici comment cela fonctionne, avec trois analogies clés :

1. Le Timing est Tout (L'Analogie de la Construction)

Dans la création d'une image, le modèle commence par un gros nuage de bruit et le nettoie petit à petit.

• L'erreur classique : Si vous faites une erreur au tout début de la construction d'une maison (les fondations), tout ce qui suit sera bancal. Si vous faites une erreur à la fin (peindre le mur), ce n'est pas grave.
• La méthode OBS-Diff : Elle comprend que les premiers pas (les fondations) sont les plus critiques. Elle utilise une "balance intelligente" qui donne beaucoup plus de poids aux erreurs commises au début du processus. Elle protège donc farouchement les connexions importantes pour les premières étapes, là où les autres méthodes coupent trop facilement.

2. Le "Chirurgien" (OBS) Adapté

Le modèle utilise une vieille technique mathématique appelée "Optimal Brain Surgeon" (OBS), qui calcule exactement quelle connexion supprimer pour faire le moins de dégâts possible.

• Le défi : Cette technique est très lourde à calculer pour des modèles géants.
• L'astuce OBS-Diff : Au lieu de calculer tout d'un coup (ce qui ferait exploser la mémoire), ils découpent le modèle en paquets (des groupes de pièces). Ils traitent un paquet, le réparent, puis passent au suivant. C'est comme si vous rénoviez une maison pièce par pièce plutôt que de tout démonter en même temps.

3. La "Carte au Trésor" (Le Hessian)

Pour savoir quoi couper, le modèle utilise une "carte de sensibilité" (appelée Hessian).

• L'innovation : OBS-Diff crée une carte qui change dynamiquement selon l'étape de la création de l'image. Elle dit : "À l'étape 1, cette connexion est vitale ! À l'étape 20, on peut s'en passer." Cela permet de garder les connexions qui comptent vraiment pour la qualité finale.

🏆 Les Résultats : Plus Léger, Toujours Beau

Les chercheurs ont testé OBS-Diff sur des modèles très puissants (comme Stable Diffusion 3 et Flux).

• Résultat : Ils ont pu enlever jusqu'à 50% ou 70% des connexions du modèle.
• Qualité : Les images générées sont toujours incroyablement belles, nettes et fidèles à la description. Les autres méthodes, elles, produisaient des images floues ou déformées à ce niveau de compression.
• Vitesse : Comme le modèle est plus petit, il génère les images plus vite.
• Coût : Tout cela se fait sans réentraînement. C'est comme si on pouvait alléger une voiture de course en une après-midi, sans avoir besoin de refaire le moteur.

En Résumé

OBS-Diff, c'est comme avoir un architecte génie qui sait exactement quelles briques retirer d'un château de sable géant pour le rendre plus petit et plus rapide, sans qu'il ne s'effondre. Il sait que les briques du bas sont plus importantes que celles du haut, et il travaille si vite qu'il ne faut pas attendre des jours pour voir le résultat.

C'est une avancée majeure pour rendre l'art généré par l'IA accessible à tous, même sur des ordinateurs moins puissants.

Résumé technique

1. Problématique

Les modèles de diffusion text-to-image à grande échelle (comme Stable Diffusion 3, 3.5 ou Flux) offrent des capacités de génération d'images impressionnantes, mais leur coût computationnel et leurs exigences mémoire sont prohibitifs, limitant leur accessibilité.
Bien que le pruning (élagage) soit une méthode éprouvée pour compresser les réseaux de neurones, son application aux modèles de diffusion pose des défis uniques :

• Nature itérative : Contrairement aux modèles de langage (LLM) qui effectuent un seul passage avant, les modèles de diffusion itèrent sur une trajectoire de débruitage. Les erreurs introduites tôt dans le processus se propagent et s'accumulent, dégradant gravement la sortie finale.
• Limites des méthodes existantes : Les méthodes de pruning "one-shot" (en une seule passe) et sans entraînement (training-free) développées pour les LLM (comme SparseGPT ou Wanda) ne peuvent pas être appliquées directement car elles ne tiennent pas compte de la dynamique temporelle du débruitage. De plus, les méthodes existantes pour la diffusion nécessitent souvent un réentraînement coûteux ou sont spécifiques à certaines architectures (ex: U-Net), manquant de généralité pour les nouveaux transformateurs multimodaux (MMDiT).

L'objectif est donc de développer un cadre de pruning général, sans entraînement et en une seule passe, capable de gérer des architectures variées et plusieurs granularités d'élagage.

2. Méthodologie : OBS-Diff

Les auteurs proposent OBS-Diff, un cadre innovant qui revitalise l'algorithme classique Optimal Brain Surgeon (OBS) en l'adaptant aux spécificités des modèles de diffusion.

A. Construction de la Matrice Hessienne Sensible au Pas de Temps (Timestep-Aware Hessian)

C'est l'apport théorique central. L'objectif de pruning est reformulé pour minimiser l'erreur de reconstruction pondérée, où les poids des étapes de débruitage précoces sont augmentés.

• Logique : Les erreurs aux premiers pas ($t$ petit) sont plus dommageables car elles affectent toutes les étapes suivantes.
• Implémentation : Une fonction de pondération logarithmique décroissante ($\alpha_t$) est appliquée. La matrice Hessienne ($H_l$), qui capture les informations d'ordre deux sur l'importance des poids, est calculée comme une somme pondérée sur tous les pas de temps :
  $$H_l = 2 \sum_{t=1}^{T} \alpha_t E[X_{l,t}X_{l,t}^T]$$
  Cela permet d'identifier et de conserver les poids critiques pour les étapes initiales du processus de génération.

B. Stratégie d'Élagage Séquentiel par Groupes (Module Packages)

Pour surmonter le coût prohibitif de la calibration (qui nécessiterait de faire passer tout le jeu de données pour chaque couche individuellement), les auteurs introduisent une stratégie par "Paquets de Modules".

• Principe : Le modèle est divisé en blocs de couches (Basic Units) traités par paquets.
• Avantage : Pour chaque paquet, on collecte les statistiques d'activation sur l'ensemble de la trajectoire de débruitage en une seule fois, puis on élague toutes les couches du paquet simultanément. Cela amortit le coût de collecte de données tout en maintenant une précision de calibration élevée.

C. Adaptabilité aux Granularités

Le cadre OBS-Diff est conçu pour être flexible :

• Non structuré : Élagage de poids individuels.
• Semi-structuré (N:M) : Par exemple, le motif 2:4 (garder 2 poids sur 4), compatible avec le matériel.
• Structuré : Élagage de neurones entiers (FFN) ou de têtes d'attention (MHA). Pour les mécanismes d'attention conjointe (MMDiT) où les têtes sont partagées entre modalités, une fusion des rangs (Reciprocal Rank Fusion) est utilisée pour déterminer l'importance globale d'une tête avant de l'élaguer.

3. Contributions Clés

1. Adaptation de l'OBS aux Diffusion Models : Première application réussie du principe Optimal Brain Surgeon aux modèles de diffusion itératifs, gérant des architectures complexes comme le MMDiT.
2. Pondération Temporelle : Introduction d'une construction de Hessienne sensible au pas de temps, priorisant la préservation des étapes de débruitage critiques pour éviter l'accumulation d'erreurs.
3. Efficacité Computationnelle : Développement d'une stratégie de calibration par paquets qui rend le processus de pruning faisable pour des modèles de plusieurs milliards de paramètres sans fine-tuning.
4. Polyvalence : Support unifié pour le pruning non structuré, semi-structuré et structuré.

4. Résultats Expérimentaux

Les expériences ont été menées sur une gamme de modèles (SD v2.1, SD 3-Medium, SD 3.5-Large, Flux.1-dev) et comparées à des baselines comme Magnitude, Wanda, DSnoT, EcoDiff et Diff-Pruning.

• Qualité de Génération : OBS-Diff surpasse systématiquement les méthodes de référence, en particulier aux taux d'élagage élevés (50% à 70%).
  • Sur SD 3.5-Large à 60% d'élagage, les méthodes baselines s'effondrent (FID > 80, ImageReward négatif), tandis qu'OBS-Diff maintient un FID de ~29 et un ImageReward positif.
  • Les images générées par OBS-Diff conservent la cohérence sémantique et la qualité visuelle, là où les autres méthodes produisent des artefacts graves ou des images détruites.
• Métriques : OBS-Diff obtient les meilleurs scores sur ImageReward (alignement avec les préférences humaines) et CLIP Score dans la majorité des cas, avec une dégradation minime du FID par rapport au modèle dense.
• Efficacité : Le processus de pruning est rapide (moins de 15 minutes pour un modèle de 2B paramètres sur une RTX 4090) et ne nécessite aucun réentraînement.
• Robustesse : Le modèle élagué reste robuste face à des changements de paramètres d'inférence (nombre de pas, solvers, CFG scale) et de prompts hors distribution (OOD).

5. Signification et Impact

OBS-Diff établit un nouvel état de l'art pour le pruning des modèles de diffusion text-to-image.

• Accessibilité : Il permet de déployer des modèles de diffusion massifs sur du matériel moins puissant (ex: cartes grand public) tout en préservant une qualité de génération élevée.
• Généralité : En s'affranchissant de l'architecture spécifique (fonctionnant aussi bien sur U-Net que sur MMDiT) et de l'entraînement, il offre une solution pratique pour la compression de modèles génératifs.
• Fondation Future : La méthodologie ouvre la voie à l'application de techniques d'élagage basées sur le Hessien à d'autres paradigmes de diffusion (génération vidéo, modèles de langage basés sur la diffusion) et à d'autres architectures multimodales.

En résumé, OBS-Diff comble le fossé entre les techniques de compression avancées des LLM et les besoins spécifiques des modèles de diffusion, offrant une solution efficace, rapide et de haute qualité pour la compression de ces modèles.