Steering Diffusion Models via Class-Contrastive Influence for Few-Shot Medical Classification

Résumé Technique : Pilotage des Modèles de Diffusion par l'Influence Contrastive de Classe pour la Classification Médicale à Faible Échantillonnage (Few-Shot)

Énoncé du Problème
Dans la classification d'images médicales à faible échantillonnage (few-shot), les données étiquetées sont rares, ce qui rend l'augmentation de données essentielle. Bien que les modèles de diffusion prêts à l'emploi puissent générer des échantillons d'entraînement synthétiques, les approches existantes se concentrent principalement sur l'amélioration du réalisme, de la diversité ou de l'adaptation au domaine de ces images. Les auteurs soutiennent que ces métriques ne garantissent pas l'utilité pour la tâche de classification en aval. Il existe un écart significatif où des images générées d'une qualité visuelle comparable produisent des performances de classification très différentes. Le problème central identifié est l'absence d'un critère fondé sur des principes pour mesurer et optimiser l'« utilité pour la tâche » des échantillons synthétiques. Le papier pose comme limitation centrale des pratiques actuelles non pas la qualité de la génération seule, mais l'incapacité à générer des échantillons qui aident spécifiquement le classificateur à affiner sa frontière de décision.

Méthodologie
La solution proposée, l'Influence Contrastive de Classe (C2I - Class-Contrastive Influence), déplace l'objectif d'optimisation de la qualité de l'image vers l'influence basée sur le gradient des échantillons synthétiques sur le classificateur.

1. Critère d'Influence Contrastive de Classe (C2I) :

   • La méthode exploite l'estimation de l'influence basée sur le gradient pour quantifier comment un échantillon d'entraînement affecte la perte de validation.
   • Les auteurs observent que les scores d'influence bruts, agnostiques à la classe, ne sont pas corrélés à la performance de classification car les ensembles de validation contiennent des signaux conflictuels provenant de différentes classes.
   • Au lieu de cela, ils définissent la C2I basée sur l'écart d'influence : un échantillon est jugé utile si son gradient de perte s'aligne positivement avec les gradients de validation de la même classe et s'oppose (s'aligne négativement avec) aux gradients des autres classes.
   • Pour la classification binaire, la C2I est formulée comme la différence au carré entre la moyenne de l'influence sur les ensembles de validation positifs et négatifs, normalisée par la variance. Pour les tâches multi-classes, elle est définie comme le score softmax de l'influence moyenne sur la classe correspondante par rapport à toutes les classes.

2. Connexion Théorique avec la Difficulté des Échantillons (Sample Hardness) :

   • L'analyse théorique (Théorème 1) démontre que la maximisation de la C2I dans un cadre de régression logistique pousse le générateur d'échantillons vers la moyenne globale de l'ensemble de validation.
   • Puisque la moyenne globale de deux clusters de classes se situe généralement près de la frontière de décision, la maximisation de la C2I sélectionne intrinsèquement des exemples « difficiles » (hard examples) — des échantillons qui sont difficiles à classifier et qui se situent proches de la frontière.
   • L'entraînement sur ces exemples difficiles, proches de la frontière, aide à affiner la frontière de décision, améliorant ainsi la généralisation et la robustesse.

3. Affinage par Apprentissage par Renforcement (RL) :

   • Pour opérationnaliser cela, les auteurs proposent un schéma d'affinage (fine-tuning) pour les modèles de diffusion Text-to-Image (T2I) pré-entraînés (spécifiquement Stable Diffusion 2.1).
   • Fonction de Récompense : Le score C2I sert de signal de récompense dans un cadre d'apprentissage par renforcement. Le générateur est affiné pour maximiser l'espérance de la récompense C2I.
   • Optimisation : Les auteurs utilisent la Denoising Diffusion Policy Optimization (DDPO) pour mettre à jour les paramètres du modèle de diffusion (spécifiquement les poids LoRA) sur la base de la récompense C2I. Cela oriente le processus de génération vers les régions de la variété de données (data manifold) informatives pour la classe, sans nécessiter de changements architecturaux dans le générateur ou le classificateur.
   • Flux de Travail : Un classificateur (ViT-B/16) est d'abord affiné sur l'ensemble d'entraînement à faible échantillonnage. Ses gradients sur un ensemble de validation tenu à l'écart sont utilisés pour calculer la récompense C2I pour les lots générés, qui guident ensuite l'affinage par RL du modèle de diffusion.

Contributions Clés

• Formalisation de l'Efficacité : Le papier formalise l'efficacité des échantillons via la C2I, démontrant que les échantillons à haute C2I correspondent à des exemples difficiles et proches de la frontière plutôt qu'à de simples images réalistes.
• Schéma de Génération Basé sur le RL : Il introduit un cadre de RL « plug-and-play » qui utilise la C2I comme récompense pour orienter les modèles de diffusion standards vers la génération d'échantillons utiles pour la tâche et informatifs pour la classe.
• Validation Empirique : La méthode est validée sur plusieurs benchmarks d'imagerie médicale à faible échantillonnage (BreastMNIST, DermaMNIST, PneumoniaMNIST), montrant des améliorations constantes par rapport à l'augmentation standard et aux autres méthodes de diffusion basées sur l'augmentation.

Résultats Expérimentaux
La méthode a été évaluée sur trois benchmarks MedMNIST dans des configurations de faible échantillonnage (16 ou жды 32 échantillons étiquetés par classe).

• Performance : Le générateur guidé par la C2I a systématiquement surpassé les bases de comparaison, incluant les transformations standards (RandAugment, RandomErasing, Mixup) et d'autres méthodes d'augmentation basées sur la diffusion (DataDream, Dataset Expansion, DistDiff).
  • En moyenne, la méthode proposée a obtenu l'AUC la plus élevée à travers différents backbones (ViT-B/16 et ResNet-18).
  • Notamment, la méthode a amélioré les performances sur ResNet-18, indiquant que les échantillons synthétisés sont largement informatifs et transférables entre les architectures de modèles.
• Robustesse : Les modèles entraînés avec l'augmentation guidée par la C2I ont démontré une robustesse supérieure sous diverses conditions de bruit (sel et poivre, flou gaussien, compression JPEG), suggérant la génération d'une frontière de décision plus stable.
• Corrélation de Difficulté : Les expériences ont confirmé qu'au fur et à mesure que l'affinage par RL progresse, la récompense C2I augmente et les caractéristiques des images générées se rapprochent des centroïdes de caractéristiques de la validation, confirmant la génération d'exemples plus difficiles.
• Ablation : L'augmentation du nombre d'époques d'entraînement RL et du nombre d'images synthétisées a généralement amélioré l'AUC en aval, tandis que certaines bases de comparaison (par exemple, Dataset Expansion) ont montré une dégradation des performances avec l'augmentation des données synthétiques.

Signification et Revendications
Le papier affirme que l'augmentation synthétique est plus efficace lorsqu'elle est guidée par l'utilité pour la tâche plutôt que par la qualité de l'image seule. En alignant les données synthétiques avec les gradients de validation, la C2I transforme les modèles de diffusion génériques en générateurs ciblés qui renforcent les frontières de décision là où les données étiquetées sont les plus rares.

Les auteurs reconnaissent des limites, notant que leur méthode introduit un surcoût de calcul supplémentaire par rapport à l'utilisation des données d'entraînement originales, principalement en raison du processus d'affinage par RL (nécessitant environ 5 heures GPU sur un A100). Cependant, ils affirment que ce coût est justifié par les gains constants en précision et en robustesse dans les régimes de faibles données. L'étude suggère un changement de paradigme dans l'utilisation des modèles de diffusion pour l'augmentation : optimiser la génération pour son influence descendante sur le classificateur.