DerMAE: Improving skin lesion classification through conditioned latent diffusion and MAE distillation

L'article présente DerMAE, une méthode qui améliore la classification des lésions cutanées en générant des images synthétiques via des modèles de diffusion conditionnés et en utilisant l'auto-entraînement MAE suivi d'une distillation de connaissances pour déployer des modèles légers et performants sur des appareils mobiles.

L'essentiel

🩺 Le Problème : Un Médecin qui voit trop de "Bons" et pas assez de "Mauvais"

Imaginez que vous essayez d'apprendre à un jeune médecin (l'intelligence artificielle) à reconnaître les cancers de la peau. Le problème, c'est que dans les livres d'images qu'il utilise pour étudier (les données), il y a 90 % de grains de beauté inoffensifs et seulement 10 % de cancers dangereux.

C'est comme si vous appreniez à un enfant à reconnaître les lions en lui montrant 90 photos de chats et seulement 10 photos de lions. À la fin, l'enfant pensera que tout animal poilu est un chat ! Il va rater les vrais lions (les cancers) parce qu'il n'en a jamais assez vu. C'est ce qu'on appelle le déséquilibre des classes.

🎨 La Solution : Trois Astuces Magiques

Les chercheurs de l'Université Fédérale de Pernambouc (au Brésil) ont inventé une méthode en trois étapes pour résoudre ce problème.

1. L'Imprimeur 3D de Peau (Génération Synthétique)

Au lieu de chercher désespérément plus de photos de vrais cancers (qui sont rares), ils ont créé un imprimeur 3D magique (un modèle de diffusion).

• Comment ça marche ? C'est comme un artiste qui regarde une photo de cancer, la brouille avec du bruit, puis apprend à la reconstruire parfaitement.
• L'astuce : Ils ont demandé à l'imprimeur de créer 600 000 nouvelles photos de cancers et de grains de beauté.
• Le résultat : Soudain, le jeune médecin a un tas de photos équilibrées. Il peut étudier autant de lions que de chats. De plus, ils ont ajouté un "filtre de beauté" (une perte perceptuelle) pour s'assurer que les faux cancers ressemblent vraiment à la vraie peau, et pas juste à des taches floues.

2. L'Entraînement de l'Expert (Pré-entraînement MAE)

Maintenant qu'ils ont beaucoup de photos, ils doivent entraîner un super-héros (un très gros modèle d'IA appelé ViT-Huge).

• Le jeu : Imaginez que vous cachez 75 % d'une photo de peau avec un cache (c'est le "Masque"). Le super-héros doit regarder les petits bouts qui restent et deviner à quoi ressemble le reste de la peau.
• Pourquoi ? Cela force l'IA à comprendre la structure globale de la peau (les formes, les contours) plutôt que de juste mémoriser des détails. C'est comme apprendre à reconnaître un visage en ne voyant qu'un œil ou une bouche.
• Grâce aux 600 000 photos générées, ce super-héros devient un expert infaillible.

3. Le Professeur et l'Élève (Distillation de Connaissances)

Voici le hic : le super-héros (ViT-Huge) est trop lourd ! Il est comme un éléphant dans une boutique de porcelaine : il est trop gros et consomme trop d'énergie pour tenir dans le téléphone portable d'un médecin ou d'une application mobile.

• La solution : Ils utilisent la distillation de connaissances.
• L'analogie : Imaginez que le super-héros (le Professeur) a toutes les connaissances du monde. Il ne peut pas voyager léger, alors il prend un petit élève (un modèle plus simple, comme un ViT-Base ou EfficientNet) et lui transmet tout son savoir.
• Le petit élève ne regarde pas seulement les réponses correctes, il essaie de penser comme le Professeur. Il apprend à voir les nuances que le Professeur voit.
• Le résultat final : On obtient un petit modèle, léger et rapide, qui tient dans un téléphone, mais qui a l'intelligence du géant.

🚀 Pourquoi c'est génial ?

1. Plus juste : Grâce aux photos générées, l'IA ne rate plus les cancers rares. Elle est plus précise.
2. Plus rapide : Grâce à la distillation, on peut utiliser cette IA intelligente directement sur un téléphone portable, même sans connexion internet. C'est parfait pour les zones reculées où il n'y a pas de dermatologue.
3. Économique : On n'a pas besoin de milliers de vrais patients pour entraîner le système, juste un peu de données réelles et beaucoup de données "fabriquées" intelligentes.

En résumé

Les chercheurs ont dit : "Si nous n'avons pas assez de vrais cancers pour apprendre à l'IA, nous allons en fabriquer de faux mais réalistes. Ensuite, nous allons entraîner un géant sur ces images, et enfin, nous allons lui apprendre à transmettre son génie à un petit modèle qui pourra voyager dans votre poche pour aider les médecins."

C'est une victoire pour la médecine de précision, rendue accessible à tous, partout dans le monde.

Résumé technique

1. Problématique

La classification des lésions cutanées souffre de deux obstacles majeurs qui limitent l'efficacité des modèles d'apprentissage profond :

• Déséquilibre sévère des classes : Les ensembles de données dermatologiques (comme HAM10000) sont dominés par des lésions bénignes (environ 90 %), tandis que les cas malins sont sous-représentés. Cela entraîne des frontières de décision biaisées favorisant la classe majoritaire.
• Limites des données et des architectures : Les jeux de données médicaux sont souvent petits et non standardisés. De plus, les architectures puissantes comme les Vision Transformers (ViT), bien que performantes grâce à l'attention, nécessitent généralement de très grandes quantités de données pour apprendre des représentations stables, ce qui les rend peu efficaces sur des datasets dermatologiques déséquilibrés.
• Contraintes de déploiement : Les modèles lourds (comme les ViT-Huge) sont difficiles à déployer dans des environnements cliniques réels ou sur des appareils mobiles en raison de leurs exigences computationnelles et mémoire élevées.

2. Méthodologie

Les auteurs proposent un pipeline en trois étapes, nommé DERMAE, pour surmonter ces limitations :

A. Génération de Données Synthétiques (Diffusion Latente Conditionnée)

Pour pallier le déséquilibre des classes, l'équipe utilise un modèle de Denoising Diffusion Probabilistic Model (DDPM) conditionné par la classe.

• Stratégie : Au lieu d'une génération unconditionnelle, le modèle utilise des embeddings de classe lors de l'inférence pour générer spécifiquement des images de lésions malignes, équilibrant ainsi artificiellement le dataset.
• Optimisation : L'entraînement combine une perte de type Mean Squared Error (MSE) avec une perte perceptuelle (perceptual loss). Cette dernière est cruciale pour améliorer la fidélité visuelle et réduire les artefacts de basse fréquence souvent générés par la MSE seule.
• Données : Le modèle est entraîné sur le dataset HAM10000 (10 000 images, 8 catégories) et génère des centaines de milliers d'échantillons synthétiques.

B. Pré-entraînement Auto-supervisé (MAE)

Pour exploiter ces données synthétiques massives, les auteurs emploient une stratégie de Masked Autoencoder (MAE) avec un backbone ViT-Huge (ViT-H/16).

• Principe : Le modèle apprend à reconstruire des patches masqués (75 % d'occultation) de l'image.
• Objectif : Cela force le modèle à apprendre des représentations latentes robustes et sémantiquement riches (structures globales des lésions) plutôt que de simples corrélations de bas niveau, en s'entraînant sur le vaste corpus de données synthétiques générées.

C. Distillation de Connaissances (Knowledge Distillation)

Afin de rendre le système utilisable sur des appareils mobiles, les représentations apprises par le grand modèle "enseignant" (ViT-H pré-entraîné en MAE) sont transférées à un modèle "élève" plus léger.

• Architecture : Un ViT-Base (ViT-B/16) ou un EfficientNet-B0 sert d'élève.
• Processus : L'élève est entraîné avec une perte de distillation combinant la vraisemblance des étiquettes réelles (hard labels) et la distribution de probabilité du modèle enseignant (soft labels), via une divergence de Kullback-Leibler.
• Affinage : Le modèle final est affiné (fine-tuned) avec un mélange de données réelles et synthétiques.

3. Contributions Clés

1. Intégration Diffusion-MAE : Une approche novatrice combinant la génération de données conditionnelle (pour l'équilibrage) et le pré-entraînement auto-supervisé (pour l'apprentissage de caractéristiques robustes).
2. Optimisation pour le Edge Computing : Démonstration qu'il est possible de transférer la puissance des grands modèles ViT vers des architectures légères (EfficientNet, ViT-B) via la distillation, rendant la classification précise possible sur mobile.
3. Amélioration de la Fidélité Visuelle : L'intégration d'une perte perceptuelle dans le modèle de diffusion améliore significativement la qualité et la diversité des lésions synthétiques générées.

4. Résultats Expérimentaux

Les expériences ont été menées sur le dataset HAM10000, comparant les approches proposées à des baselines (ViT, EfficientNet) et à l'état de l'art (Derm-t2im).

• Performance de Classification :
  • Les modèles ViT baselines (sans pré-entraînement) ont des performances médiocres en raison du manque de données et du déséquilibre.
  • L'ajout de pré-entraînement MAE sur données synthétiques améliore drastiquement les performances. Par exemple, le modèle MAE + ViT-H + synth atteint un score F1 de 0,8911 en classification catégorielle et 0,8715 en classification binaire (bénin/malin), surpassant largement les baselines.
  • La distillation permet de conserver une grande partie de cette performance dans des modèles plus petits. Le modèle MAE ViT-H distillé vers EfficientNet-B0 + synth atteint un score F1 binaire de 0,9010, surpassant même le modèle ViT-Huge non distillé dans certaines configurations.
• Efficacité Computationnelle :
  • Le modèle ViT-Huge nécessite environ 55-60 GFLOPs par inférence (impraticable pour le mobile).
  • Le modèle distillé (EfficientNet-B0) ne nécessite que ~0,6 GFLOPs et ~5M de paramètres, permettant une inférence efficace sur appareil tout en maintenant une haute précision.

5. Signification et Impact

Ce travail démontre qu'il est possible de surmonter le problème du déséquilibre des données en dermatologie sans sacrifier l'efficacité computationnelle.

• Clinique : La méthode permet de créer des systèmes de diagnostic assisté par IA robustes, capables de détecter des lésions malignes rares avec précision.
• Déploiement : En combinant génération synthétique, pré-entraînement auto-supervisé et distillation, l'approche rend possible le déploiement de modèles de pointe sur des dispositifs mobiles (télédermatologie), facilitant l'accès aux soins spécialisés dans des zones sous-équipées.
• Généralité : Le cadre proposé (Synthèse conditionnelle + MAE + Distillation) offre une voie prometteuse pour d'autres tâches médicales souffrant de pénurie de données étiquetées et de déséquilibre de classes.