Synthetic Data Generation for Long-Tail Medical Image Classification: A Case Study in Skin Lesions

Résumé Technique : Génération de Données Synthétiques pour la Classification d'Images Médicales à Distribution Longue Traîne

Énoncé du Problème
Les distributions de classes à longue traîne sont omniprésentes dans les jeux de données médicaux multi-classes, posant des défis majeurs pour les modèles d'apprentissage profond. Dans ces scénarios, quelques classes dominantes « tête » contiennent la majorité des échantillons, tandis que de nombreuses classes « queue » cliniquement critiques (correspondant souvent à des maladies rares et à haut risque) sont gravement sous-représentées. Les modèles entraînés sur de telles données déséquilibrées atteignent généralement une précision globale compétitive mais échouent à généraliser aux classes de queue. Les solutions existantes — catégorisées en conception de fonctions de perte, améliorations de l'architecture du modèle et rééquilibrage des classes — n'offrent que des améliorations marginales. Ces méthodes reposent souvent sur un réglage minutieux des hyperparamètres, présentent une généralisation limitée à travers différents modèles et peinent à se mettre à l'échelle en raison du caractère déterministe de l'augmentation de données manuelle. De plus, bien que les modèles de diffusion aient fait leurs preuves dans la génération d'images naturelles, leur application à l'imagerie médicale spécifique à un domaine et à longue traîne reste sous-exploitée, en particulier concernant la génération d'échantillons diversifiés mais cliniquement significatifs.

Méthodologie
Les auteurs proposent un pipeline d'augmentation de données synthétiques piloté par un modèle de diffusion, adapté à la classification d'images médicales à longue traîne (LTMIC), démontré spécifiquement sur la classification des lésions cutanées. L'approche se compose de trois étapes principales :

1. Architecture du Modèle de Diffusion de Réparation (Inpainting) :
   Le cœur de la méthode est un modèle de diffusion de réparation novateur. Contrairement aux modèles d'incrustation de texte standards, inadaptés aux descriptions médicales hautement spécialisées, les auteurs entraînent un réseau d'incrustation de classe simple entièrement connecté. Le modèle reçoit quatre entrées : l'image dermatoscopique originale ($x$), un masque binaire de zone d'intérêt ($b$), l'arrière-plan de l'image dermatoscopique ($d$) et l'étiquette de classe ($y$). L'image et l'arrière-plan sont encodés dans un espace latent via un Auto-Encodeur Variationnel (VAE) pré-entraîné. Le masque binaire est redimensionné pour correspondre aux dimensions latentes. Ces composantes sont concaténées et alimentées dans un UNet paramétré par $\theta_d$ pour estimer le bruit ajouté lors du processus de diffusion vers l'avant. Un réseau d'incrustation de classe paramétré par $\theta_c$ fournit la conditionnement par étiquette.

2. Affinement avec Sélection d'Échantillons :
   Pour garantir que le modèle de diffusion apprenne la distribution de chaque classe de manière égale, une stratégie de suréchantillonnage est employée lors de l'affinement. Les auteurs s'appuient sur des travaux antérieurs suggérant que les échantillons « approximativement dans la distribution » (AID) (ceux légèrement éloignés du centre de la distribution) sont les plus importants pour l'entraînement des modèles génératifs. Par conséquent, une méthode d'autodétection d'anomalies médicales est utilisée pour calculer les scores d'anomalie pour chaque échantillon. Des poids sont attribués en fonction de la distance d'un échantillon par rapport au score d'anomalie médian de sa classe ; les échantillons plus proches de la région AID se voient attribuer des poids plus élevés et sont répétés pour équilibrer le jeu de données d'affinement à travers toutes les classes.

3. Génération Synthétique et Filtrage OOD :
   Le modèle affiné génère des images synthétiques en prenant les images d'arrière-plan et les étiquettes de classe du jeu de données d'entraînement original comme entrée. En raison du caractère stochastique des modèles de diffusion, le jeu de données synthétique initial ($D_{syn}$) contient des échantillons « bruyants », y compris des images hors distribution (OOD) qui peuvent ne pas être cliniquement valides. Pour remédier à cela, un mécanisme de filtrage OOD médical est appliqué à $D_{syn}$ pour produire un jeu de données synthétique propre ($D_{clean}^{syn}$). Le jeu de données d'entraînement final est formé en combinant le jeu de données à longue traîne original ($D_{LT}$) avec $D_{clean}^{syn}$. Un hyperparamètre $\gamma$ contrôle le pourcentage d'échantillons propres conservés à partir du jeu de données synthétique pour chaque classe.

Contributions Clés
L'article expose trois contributions principales :

1. Pipeline Novateur : Un pipeline de génération de données synthétiques intégrant un modèle de diffusion de réparation avec un filtrage OOD, spécifiquement conçu pour résoudre le problème LTMIC.
2. Architecture Spécialisée : Une nouvelle architecture de modèle de diffusion de réparation adaptée à la génération d'images médicales, utilisant un réseau d'incrustation de classe dédié et un conditionnement par arrière-plan plutôt que des incrustations de texte génériques.
3. Stratégie d'Échantillonnage : Une stratégie d'échantillonnage novatrice basée sur OOD pour l'affinement des modèles de diffusion sur des jeux de données médicaux à longue traîne, qui privilégie les échantillons AID pour améliorer la qualité des données générées pour les classes de queue.

Résultats Expérimentaux
La méthode a été évaluée sur le jeu de données de classification des lésions cutanées ISIC2019, une référence contenant 25 331 images réparties sur 8 classes avec un déséquilibre sévère (allant de 12 875 échantillons pour le Nævus à seulement 62 pour la Lésion Vasculaire). Les expériences ont utilisé une validation croisée à 5 plis et un classificateur simple basé sur CNN (sans ensemble ni données externes) pour isoler l'impact de la méthode de génération de données.

• Performance Globale : L'approche proposée a atteint une Précision Multiclasse Équilibrée (BMA) de 0,802, surpassant significativement les méthodes actuelles de l'État de l'Art (SOTA), y compris les ensembles de modèles multiples (MRE) et d'autres approches basées sur la diffusion (par exemple, CBDM, Dreambooth). Notamment, elle a dépassé la méthode MRE d'environ 4 % en BMA sans nécessiter de données externes.
• Amélioration des Classes de Queue : Les gains les plus significatifs ont été observés dans les classes de queue. Pour la classe avec le moins d'échantillons (Dermatofibrome, DF), la méthode a produit une amélioration de la précision de plus de 28 % par rapport à la ligne de base sans données synthétiques.
• Impact du Filtrage OOD : L'étude a révélé que le filtrage OOD est crucial pour équilibrer les performances. Bien que les données synthétiques sans filtrage aient considérablement amélioré les classes de queue, elles ont introduit du bruit qui a légèrement dégradé les performances des classes de tête. Le processus de filtrage a restauré les performances des classes de tête tout en maintenant les gains dans les classes de queue.
• Sensibilité aux Hyperparamètres : Une étude d'ablation sur $\gamma$ (le taux de rétention des échantillons synthétiques propres) a révélé que les performances s'améliorent nettement avec même une petite quantité de données synthétiques, mais plafonnent ou diminuent lorsque $\gamma$ approche 1 (aucun filtrage). La plage optimale pour $\gamma$ a été identifiée entre 0,2 et 0,6.

Signification et Revendications
Les auteurs affirment que leurs résultats démontrent l'efficacité de l'augmentation basée sur la diffusion pour atténuer le déséquilibre à longue traîne et renforcer la robustesse de la classification médicale. En générant des échantillons synthétiques diversifiés, réalistes et cliniquement significatifs, la méthode répond à la rareté des données des maladies rares sans dépendre d'ensembles de modèles complexes ou de fonctions de perte spécialisées sensibles aux hyperparamètres. L'article conclut que, bien que l'approche actuelle traite la génération de données et la classification comme des processus indépendants, les travaux futurs se concentreront sur l'optimisation conjointe de ces composants pour améliorer davantage les performances LTMIC. La méthode est présentée comme agnostique au modèle, capable d'intégration avec n'importe quelle architecture de classification.