Résumé technique : Locus – Régularisation de l'attention guidée par l'anatomie pour la classification d'images médicales

1. Énoncé du problème

Les modèles de classification d'images médicales sont généralement entraînés en utilisant uniquement des étiquettes au niveau de l'image, sans supervision spatiale explicite. Par conséquent, ces modèles apprennent souvent à s'appuyer sur des corrélations statistiques avec des régions anatomiquement non pertinentes, telles que des artefacts d'imagerie, des tissus environnants ou des indices contextuels spécifiques au jeu de données, plutôt que sur les structures anatomiques diagnostiquement pertinentes. Ce comportement limite l'interprétabilité et la généralisabilité des modèles.

Bien que des approches précédentes aient tenté de remédier à cela en incorporant des informations de forme anatomique (par exemple, des masques de segmentation), elles font face à des limitations significatives :

• Annotation manuelle : S'appuyer sur des masques délimités par des experts est coûteux et difficile à mettre à l'échelle.
• Entraînement conjoint : L'entraînement d'une branche de segmentation parallèlement à un classificateur introduit des défis d'optimisation, tels que des gradients concurrents et des prédictions bruitées qui peuvent corrompre le guidage.
• Pseudo-masques basés sur la CAM : Bien qu'ils évitent les coûts d'annotation, ces méthodes activent souvent uniquement les régions discriminantes corrélées à l'étiquette plutôt que la structure anatomique pertinente complète.
• Modèles de segmentation dédiés : L'entraînement d'un nouveau modèle de segmentation à partir de zéro réintroduit la charge d'annotation et la surcharge de calcul.

2. Méthodologie : Le cadre Locus

Les auteurs proposent Locus, un cadre qui exploite des modèles de fondation de segmentation (SFM) préentraînés (par exemple, Medical SAM, BiomedParse) pour guider l'attention du classificateur sans nécessiter d'entraînement de segmentation dédié ni d'annotations manuelles.

2.1 Régularisation de l'attention centrée sur l'anatomie

Au lieu d'imposer un alignement rigide pixel par pixel entre la carte d'attention du classificateur et le masque anatomique (ce qui peut être trop restrictif), Locus introduit un terme de régularisation à charnière ($R$).

• Mécanisme : Le cadre calcule la moyenne pondérée de l'énergie d'attention dans le premier plan anatomique ($A_{fg}$) et l'arrière-plan ($A_{bg}$) dérivés du masque du SFM.
• Fonction de perte : Le régulariseur est défini par $R(M, A) = \max(0, A_{bg} - A_{fg})$.
• Logique : Ce terme pénalise le classificateur uniquement lorsque l'attention de l'arrière-plan domine celle du premier plan. Si l'attention du premier plan est déjà suffisante ($A_{fg} \geq A_{bg}$), le régulariseur reste inactif. Cette approche évite de supprimer excessivement l'information contextuelle potentiellement pertinente pour le diagnostic tout en garantissant que le modèle ne repose pas principalement sur le bruit de l'arrière-plan.

2.2 Prompting anatomique par patchs

Le prompting standard des SFM sur des images complètes échoue souvent à capturer toutes les structures anatomiques pertinentes, particulièrement lorsqu'elles sont petites ou disjointes (par exemple, des lésions éparpillées ou des organes bilatéraux), en ne retournant souvent que la région la plus dominante.

• Stratégie : Locus partitionne l'image d'entrée en une grille de patchs non chevauchants. Chaque patch est prompté indépendamment avec le même descripteur textuel dérivé de la classe (par exemple, "poumon", "tumeur").
• Agrégation : Les masques par patch résultants sont agrégés via une union pour former un masque anatomique à pleine résolution ($M$).
• Objectif : Cette stratégie maximise la couverture anatomique pour garantir que toutes les régions pertinentes contribuent à la supervision de l'attention, même si le masque résultant est grossier. Les auteurs soulignent qu'une segmentation précise au pixel près n'est pas requise ; l'objectif est un ancrage anatomique suffisant pour guider l'attention.

2.3 Perte du réseau

La perte totale du réseau combine la perte de tâche standard (entropie croisée) avec le régulariseur centré sur l'anatomie :
$$L_{net} = L_{task} + \lambda \cdot R(M, A)$$

• Stratégie d'entraînement : Une stratégie de chauffe (warm-up) est employée où le réseau est entraîné avec $L_{task}$ seul pendant quelques époques avant que la force de la régularisation ($\lambda$) ne soit progressivement augmentée.
• Flexibilité : Le cadre est agnostique à la méthode spécifique d'extraction de l'attention (par exemple, Grad-CAM, Grad-CAM++, CAM vanilla, Gradient Norm) et au choix du SFM.

3. Contributions clés

1. Nouveau cadre de régularisation : Propose un régulariseur d'attention centré sur l'anatomie qui équilibre de manière adaptative l'attention entre le premier plan et l'arrière-plan, en pénalisant le classificateur uniquement lorsque l'attention de l'arrière-plan domine, éliminant ainsi le besoin d'un entraînement de segmentation dédié.
2. Adaptabilité "Plug-and-Play" : Le cadre est facilement adaptable à divers backbones de classificateurs (CNN et Transformers) et aux SFM sans modifier l'architecture du classificateur ou affiner le modèle de fondation.
3. Validation exhaustive : Validé sur huit jeux de données d'imagerie médicale diversifiés (dermoscopie, rayons X, histopathologie, échographie et IRM cardiaque), démontrant des gains de performance constants et une meilleure interprétabilité.

4. Résultats expérimentaux

Les auteurs ont évalué Locus par rapport à trois lignes de base : l'entropie croisée standard (CE), l'entrée masquée (utilisant les masques SFM pour rejeter l'arrière-plan) et MGA (alignement d'attention basé sur Dice nécessitant des annotations d'experts).

• Performance de classification : Locus a obtenu les scores de précision et de F1 les plus élevés dans presque tous les réglages sur six jeux de données de référence (HAM10000, PneumoniaMNIST, BTXRD, PanNuke, ChestX-ray8 et un jeu de données d'IRM cardiaque cine) et a systématiquement surpassé l'entraînement non contraint et les lignes de base d'entrée masquée.
• Qualité de l'attention : Les visualisations (Grad-CAM) ont montré que Locus produisait des cartes d'attention plus focalisées et anatomiquement cohérentes. Notamment, pour les architectures sujettes à la dominance de l'arrière-plan (ex: ConvNeXt, ViT), Locus a considérablement affiné l'attention sur les structures pertinentes (ex: champs pulmonaires, lésions) par rapport à l'entraînement CE standard.
• Robustesse : Sous des perturbations au moment du test (luminosité, contraste, flou gaussien, bruit), les modèles entraînés avec Locus ont maintenu une précision plus élevée que les lignes de base CE, suggérant que la régularisation décourage la dépendance à des caractéristiques spécieuses sensibles aux distorsions.
• Généralisabilité : Les bénéfices du régulariseur ont été constants à travers différentes méthodes d'extraction d'attention (Grad-CAM, Grad-CAM++, Vanilla CAM, Gradient Norm) et différents backbones de réseaux (ResNet, ConvNeXt, ViT).

5. Signification et affirmations

L'article affirme que Locus offre une solution pratique à la nature de "boîte noire" des classificateurs d'images médicales en les orientant vers des structures anatomiques diagnostiquement pertinentes sans le coût prohibitif de l'annotation manuelle ou la complexité d'un entraînement conjoint de segmentation-classification.

En exploitant les capacités de localisation "prêtes à l'emploi" des modèles de fondation et en employant un régulariseur à charnière flexible, le cadre permet de réaliser :

• Une meilleure interprétabilité : Les modèles se concentrent sur les régions anatomiques correctes (ex: le myocarde en IRM cardiaque) plutôt que sur les artefacts ou l'arrière-plan.
• Une généralisabilité accrue : Meilleure performance à travers diverses modalités et robustesse face aux perturbations d'image.
• Efficacité : Il évite la surcharge de calcul liée à l'entraînement d'un modèle de segmentation dédié tout en fournissant un guidage anatomique fiable.

Les auteurs concluent que cette approche représente une étape significative vers la création d'une IA médicale plus fiable et cliniquement interprétable, particulièrement dans les scénarios où les annotations au niveau du pixel ne sont pas disponibles. Des travaux futurs sont suggérés pour explorer le pondération sensible à l'incertitude et l'apprentissage de bout en bout des configurations de patchs.