PeLAP-A: Adaptive Latent Pruning for Lightweight Latent Diffusion Models

Résumé Technique : PeLAP-A (Adaptive Latent Pruning for Diffusion)

Énoncé du Problème
Les modèles de diffusion latente (LDM) produisent des images de haute qualité en opérant dans un espace latent compressé généré par un auto-encodeur variationnel (VAE). Cependant, il reste une question ouverte de savoir si tous les canaux latents contribuent de manière égale au processus de diffusion ou si une redondance significative existe. Alors que les travaux existants se concentrent sur la compression de l'architecture du modèle de diffusion elle-même (par exemple, l'élagage des poids) ou sur le processus d'échantillonnage, le contenu informationnel de la représentation latente elle-même reste largement inexploré. Les auteurs étudient si des canaux latents redondants peuvent être identifiés et supprimés sélectivement sans dégrader la qualité de la génération.

Méthodologie : PeLAP-A
Le papier propose PeLAP-A (Adaptive Latent Pruning for Diffusion), un cadre léger qui augmente le pipeline standard de diffusion latente avec un prédicteur d'importance par canal, apprenable et adaptatif à l'entrée.

• Architecture : Un perceptron multicouche (MLP) à deux couches contenant seulement 292 paramètres est inséré entre l'encodeur du VAE et l'UNet de débruitage.
• Mécanisme :
  1. Le tenseur latent $z$ provenant de l'encodeur du VAE subit un pooling moyen global pour produire un résumé par canal.
  2. Ce résumé est traité par le MLP et passé à travers une activation sigmoïde pour générer un masque souple $m \in (0, 1)^C$.
  3. Le masque est diffusé spatialement et appliqué élément par élément au tenseur latent ($z' = z \odot m$) avant d'entrer dans l'UNet.
  4. La dernière couche linéaire du prédicteur est initialisée avec un biais de +3,0, garantissant que tous les canaux commencent presque entièrement ouverts ($\sigma(3) \approx 0,95$) pour éviter une suppression prématurée.
• Objectif d'entraînement : Le système est entraîné de bout en bout sur CIFAR-10 en utilisant une fonction de perte conjointe :
  $$L = L_{vae} + L_{diff} + \lambda \cdot \frac{1}{C}\sum_{c=1}^{C} m_c$$
  où $L_{vae}$ est la perte de reconstruction et de divergence KL, $L_{diff}$ est l'erreur quadratique moyenne de la prédiction du bruit de diffusion, et le troisième terme est un régulariseur de parcimonie encourageant la valeur moyenne du masque vers zéro.

Résultats Clés
Les expériences sur CIFAR-10 ont révélé un phénomène surprenant nommé « effondrement de la parcimonie » (sparsity collapse) :

• Dynamique de l'effondrement : Sous une régularisation de parcimonie agressive ($\lambda = 0,01$), le prédicteur d'importance pousse tous les canaux latents vers des valeurs proches de zéro dès les deux premiers époques d'entraînement.
• Performance de débruitage : Malgré la suppression complète de l'information latente, l'UNet de débruitage du modèle PeLAP-A a obtenu des métriques de validation meilleures que la ligne de base non élaguée :
  • Perte de diffusion : 0,0236 (PeLAP-A) contre 0,0240 (Baseline).
  • MSE de reconstruction du VAE : 22,59 (PeLAP-A) contre 24,67 (Baseline).
• Échec génératif : Bien que la voie de débruitage se soit améliorée, la qualité de la sortie générative s'est effondrée. Le score FID s'est considérablement dégradé, passant de 278,1 (Baseline) à 362,6 (PeLAP-A). L'analyse qualitative a montré qu'avec des latents proches de zéro, le décodeur produisait des sorties grises uniformes plutôt que des images structurées.
• Structure résiduelle : Même après l'effondrement, les magnitudes du masque conservaient un ordonnancement sensible à la classe (par exemple, les classes « automobile » avaient des activations moyennes légèrement plus élevées que les classes « avion »), indiquant que le prédicteur encodait des informations pertinentes à la classe même à des magnitudes sous le seuil.

Signification et Revendications
Le papier positionne PeLAP-A comme une étude exploratoire sur la dynamique de la parcimonie dans l'entraînement de la diffusion latente. Ses principales revendications et contributions sont :

1. Robustesse des UNets : Les résultats démontrent que les UNets de débruitage sont remarquablement robustes à la suppression des canaux latents. Ils peuvent apprendre des priors inconditionnels efficaces sur des espaces latents proches de zéro, suggérant que les besoins informationnels pour le processus de débruitage sont bien inférieurs à ce que la dimensionnalité de l'espace latent implique.
2. Limitation de la perte de diffusion : La divergence entre l'amélioration de la perte de diffusion et la dégradation du FID souligne que la perte de diffusion seule est un indicateur insuffisant pour la qualité de la génération lorsque l'espace latent s'est effondré. Les métriques perceptuelles comme le FID sont essentielles pour évaluer l'élagage latent.
3. Phénomène d'effondrement de la parcimonie : Les auteurs identent une transition de phase abrupte où la pression de parcimonie l'emporte sur l'apprentissage de gradients informatifs tôt dans l'entraînement. Ils émettent l'hypothèse que $\lambda = 0,01$ dépasse un seuil critique pour cette architecture spécifique, causant le système à se stabiliser dans un minimum local où l'UNet s'adapte pour débruiter des tenseurs proches de zéro plutôt que des données structurées.
4. Portée modeste : Les auteurs déclarent explicitement que la méthode, telle qu'implémentée, n'a pas atteint l'objectif de conserver sélectivement les canaux informatifs tout en supprimant les autres. Le travail sert de fondation pour comprendre les seuils critiques de la parcimonie dans les modèles génératifs et motive des travaux futurs sur les calendriers de chauffe (warm-up), la relaxation discrète ou l'entraînement sur des espaces latents plus larges pour parvenir à un élagage partiel.