Résumé Technique : Alignement Contrastif Conditionnel-Non Conditionnel (CCUA) pour les Modèles de Diffusion à Longue Traîne

Énoncé du Problème

Les données d'entraînement pour la synthèse d'images conditionnelle par classe présentent fréquemment une distribution à longue traîne, caractérisée par une rareté d'images pour les classes de la « queue » par rapport aux classes de la « tête ». Ce déséquilibre provoque un effondrement de mode (mode collapse) et une réduction de la diversité dans les modèles de diffusion générant des images pour ces classes sous-représentées. Les méthodes existantes, incluant les modèles de diffusion probabilistes de débruitage standard (DDPM) et les modèles de diffusion basés sur le score (SBDM), échouent souvent à générer des images de haute fidélité et diversifiées pour ces classes sous-représentées sans compromettre la qualité de la génération des classes de tête. Bien que l'apprentissage contrastif se soit avéré efficace pour la reconnaissance à longue traîne et la génération basée sur les GAN, son application aux modèles de diffusion pour la génération avec classes déséquilibrées nécessite une adaptation novatrice pour répondre aux mécanismes spécifiques du processus de débruitage.

Méthodologie : Alignement Contrastif Conditionnel–Non Conditionnel (CCUA)

Les auteurs proposent CCUA, un cadre comprenant deux fonctions de perte synergiques conçues pour améliorer la diversité et la fidélité des images des classes de la queue tout en maintenant la qualité des classes de tête. La méthode est compatible avec les architectures basées sur U-Net et les Transformers de Diffusion (DiT/SiT).

1. Perte Contrastive Non Supervisée (UCL) pour la Génération Non Conditionnelle

Pour remédier à l'effondrement de mode, les auteurs introduisent une perte contrastive non supervisée appliquée aux latents de la génération non conditionnelle.

• Mécanisme : Le réseau d'estimation du bruit est divisé en un encodeur de latent $e(\cdot)$ et un décodeur $d(\cdot)$. L'encodeur projette les images bruitées $x_t$ vers des latents de faible dimension $h$.
• Formulation de la perte : La perte utilise une formulation de type InfoNCE avec uniquement des échantillons négatifs. Au sein d'un mini-lot, chaque image sert d'ancre, et toutes les autres images du lot servent d'échantillons négatifs. Aucun couple positif n'est créé via l'augmentation de données ; l'ancre est traitée comme son propre vecteur positif.
• Objectif : Cela crée une « force répulsive » dans l'espace latent, maximisant la distance entre les images synthétiques indépendamment de leur classe. Cela force le modèle à diversifier les latents non conditionnels, empêchant ainsi l'effondrement vers un vecteur constant ou le surapprentissage de l'échantillonnage limité.
• Justification Théorique : Appliquer directement une perte contrastive aux latents conditionnels peut conduire à une solution triviale où le modèle maximise la diversité inter-classes (séparation des classes) tout en ignorant la diversité intra-classe. En appliant la perte aux latents non conditionnels, le modèle est contraint de diversifier tous les échantillons, ce qui est ensuite distillé dans la génération conditionnelle via la seconde perte.

2. Perte d'Alignement (AL) pour l'Alignement Conditionnel–Non Conditionnel

Le second composant aligne la génération conditionnelle avec la génération non conditionnelle diversifiée.

• Mécanisme : La perte minimise la divergence KL entre la distribution conditionnelle $p_\theta(x_{t-1}|x_t, c)$ et la distribution non conditionnelle $p_\theta(x_{t-1}|x_t)$. En pratique, cela est simplifié en une perte MSE entre le bruit estimé $\epsilon_\theta(x_t, t, c)$ et $\epsilon_\theta(x_t, t, \emptyset)$.
• Pondération Temporelle : Crucialement, cet alignement est pondéré plus lourdement aux instants de temps élevés (les premières étapes du processus inverse).
• Motivation : Les auteurs observent que lors des premières étapes de débruitage, les composantes de basse fréquence des images de différentes classes (et entre la génération conditionnelle et non conditionnelle) sont très similaires. En alignant ces étapes initiales, le modèle facilite le partage de connaissances des classes de tête (données abondantes) vers les classes de la queue (données rares). Cela permet aux classes de la queue de bénéficier de la riche diversité apprise par les classes de tête lors des premières phases de génération.

3. Synergie

Le cadre combine la perte DDPM standard avec la perte CCUA proposée :
$$\mathcal{L}_{total} = \mathcal{L}_{ddpm} + \alpha \cdot \mathcal{L}_{ucl} + \gamma \cdot \mathcal{L}_{al}$$
La synergie réside dans le fait que l'UCL diversifie l'espace latent non conditionnel, et l'AL aligne l'espace latent conditionnel sur cet espace non conditionnel diversifié. Cette diversification indirecte des latents conditionnels est empiriquement démontrée comme étant plus efficace que l'application directe d'une perte contrastive aux latents conditionnels.

Contributions Clés

1. Cadre CCUA : Une méthode novatrice pour les modèles de diffusion à classes déséquilibrées qui combine l'apprentissage contrastif non supervisé sur les latents non conditionnels avec une perte d'alignement conditionnel-non conditionnel.
2. Stratégie d'Alignement : La proposition d'aligner la génération conditionnelle et non conditionnelle spécifiquement aux instants de temps élevés (étapes de débruitage initiales) pour exploiter la similitude des composantes de basse fréquence et permettre le transfert de connaissances des classes de tête vers les classes de la queue.
3. Régularisation Contrastive Non Supervisée : L'utilisation d'une perte contrastive non supervisée avec uniquement des échantillons négatifs sur les latents non conditionnels afin de maximiser la diversité intra-classe et inter-classe sans nécessiter de labels de classe pour les paires contrastives.
4. Agnosticisme de l'Architecture : La méthode est démontrée comme étant efficace tant sur les architectures basées sur U-Net que sur les Diffusion Transformers (SiT).

Résultats Expérimentaux

La méthode a été évaluée sur les jeux de données ImageNet-LT (256×256), TinyImageNet-LT (64×64), Places-LT (64×64) et CIFAR-LT.

• Performance Quantitative :
  • Sur ImageNet-LT, CCUA surpasse significativement les bases de comparaison incluant SiT, CBDM, REPA et la Dispersive Loss.
  • À 900k étapes d'entraînement, CCUA a atteint un FID de 15,1 (contre 19,9 pour SiT) et un FIDtail de 22,5 (contre 33,9 pour SiT), représentant une amélioration de ~11,4 points de la fidélité des classes de la queue.
  • La méthode a également amélioré les métriques de Recall (diversité) et de Precision (fidélité) de manière générale.
  • Sur TinyImageNet-LT, CCUA a obtenu le meilleur FID (12,0) et FIDtail (20,8), surpassant DDPM, CBDM et OCLT.
• Résultats Qualitatifs : Les visualisations montrent que CCUA génère des images avec une plus grande diversité et fidélité pour les classes de la queue (ex: 'red wine', 'espresso', 'window shade') par rapport au SiT standard, qui produit souvent des images répétitives ou de faible qualité pour ces classes.
• Études d'Ablation :
  • La combinaison des deux pertes est nécessaire ; l'utilisation d'une seule donne des résultats sous-optimaux.
  • L'application de la perte d'alignement aux instants de temps élevés (pondération linéaire) est supérieure à une pondération constante ou quadratique.
  • L'utilisation de représentations latentes provenant du bloc N/4 du transformer offre généralement le meilleur compromis entre le FID global et le FID spatial/Recall.

Signification et Revendications

L'article affirme que CCUA traite efficacement le problème de la génération à longue traîne en adaptant les principes des GAN à longue traîne (spécifiquement l'entraînement non conditionnel à des résolutions plus basses/étapes précoces) au paradigme de la diffusion. En découplant la tâche de diversification (gérée par l'apprentissage contrastif non conditionnel) de la tâche de conditionnement (gérée par l'alignement), la méthode évite les pièges de la régularisation contrastive conditionnelle directe.

Les auteurs soulignent que leur approche :

• Améliore la génération des classes de la queue sans dégrader la qualité des classes de tête.
• Est facile à implémenter dans les pipelines DDPM et SBDM existants.
• Surpasse les travaux contemporains comme Dispersive Loss et DiffROP, particulièrement dans les scénarios de déséquilibre de classe extrême (ex: ImageNet-LT avec un facteur de déséquilibre de 0,001).
• Fournit une justification théorique et empirique sur la raison pour laquelle l'alignement de la génération conditionnelle et non conditionnelle lors des premières étapes du processus inverse est bénéfique pour le partage de connaissances dans les contextes de longue traîne.