Mask-Guided Attention Regulation for Anatomically Consistent Counterfactual CXR Synthesis

Cet article présente un cadre d'inférence régulant l'attention pour synthétiser des radiographies thoraciques contrefactuelles anatomiquement cohérentes, en combinant une régularisation guidée par des masques pour préserver la structure et un module orienté pathologie pour contrôler précisément la localisation et l'étendue des lésions.

L'essentiel

Imaginez que vous êtes un photographe médical très talentueux. Votre travail consiste à prendre une photo de poumons sains d'un patient et à dire : « Et si ce patient avait une petite infection ici ? » ou « Et si son cœur était un peu plus gros ? ».

C'est ce qu'on appelle la génération contrefactuelle : créer une image médicale réaliste montrant une maladie qui n'existe pas encore, tout en gardant le reste du corps (les côtes, la forme des poumons) exactement comme il était. C'est très utile pour entraîner les intelligences artificielles à reconnaître des maladies sans avoir besoin de photos de vrais patients malades.

Cependant, les outils actuels (basés sur une technologie appelée "modèles de diffusion") ont deux gros problèmes, un peu comme un artiste qui a la main trop lourde :

1. Le problème du "dérapage" : Quand on demande à l'IA d'ajouter une maladie, elle a tendance à modifier tout le reste de l'image. Si vous demandez d'ajouter une tache sur le poumon droit, l'IA pourrait déformer le cœur ou les côtes. C'est comme si vous essayiez de repeindre une fenêtre, mais la peinture coule sur tout le mur et change la couleur du parquet.
2. Le problème du "chuchotement" : Les maladies sont souvent de petites taches subtiles. L'IA a du mal à les entendre. Elle les ignore ou les étale partout, comme si quelqu'un chuchotait une instruction dans un stade de foot bruyant. Le résultat est une maladie floue ou mal placée.

La solution de l'équipe : Le "Régulateur de Focus"

Les chercheurs de l'Université de Tianjin ont créé une nouvelle méthode qui agit comme un régulateur de focus intelligent pour l'IA, sans avoir besoin de la réentraîner à chaque fois (ce qui serait long et coûteux). Ils utilisent deux astuces principales :

1. Le "Bouclier Anatomique" (Pour ne pas abîmer le reste)

Imaginez que vous avez un pochoir en carton (un masque) qui recouvre tout le corps du patient, sauf la zone précise où vous voulez ajouter la maladie.

• Comment ça marche : Pendant que l'IA dessine l'image, ce "pochoir" lui dit : « Tu as le droit de toucher uniquement cette zone. Le reste est protégé ! ».
• L'analogie : C'est comme si vous jouiez à un jeu vidéo où vous ne pouvez toucher que les ennemis dans une petite zone, et que le reste du décor est invulnérable. Cela empêche l'IA de "déformer" les côtes ou le cœur quand elle modifie les poumons.

2. Le "Mégaphone pour la Maladie" (Pour bien voir la tache)

Maintenant, pour s'assurer que la maladie est bien dessinée au bon endroit, l'IA utilise un "mégaphone".

• Comment ça marche : L'IA sait où elle doit regarder grâce à une carte (le masque). Elle amplifie le signal pour cette zone précise, comme si elle mettait un projecteur sur la tache à créer.
• L'ajustement fin : Si l'IA commence à étaler la maladie un peu trop, un petit mécanisme de correction (comme un correcteur orthographique automatique) intervient pour recentrer le dessin exactement là où il faut, sans en faire trop.

Le résultat ?

Grâce à cette méthode, l'IA peut :

• Créer des images de poumons malades qui ressemblent à de vraies photos médicales.
• Garder la forme exacte du corps du patient (les côtes ne bougent pas, le cœur reste à sa place).
• Placer la maladie exactement là où le médecin l'a demandée, ni plus, ni moins.

En résumé : C'est comme donner à un artiste une règle et un pinceau spécial. La règle l'empêche de salir le reste de la toile, et le pinceau spécial lui permet de peindre la maladie avec une précision chirurgicale. Cela permet de créer des milliers d'exemples de maladies virtuelles pour aider les médecins et les ordinateurs à mieux comprendre et diagnostiquer les problèmes réels, le tout sans toucher aux données sensibles des patients.

Résumé technique

1. Problématique et Contexte

La génération d'images médicales contrefactuelles vise à simuler des scénarios cliniques "et si" (par exemple, l'apparition ou la disparition d'une pathologie) tout en préservant la structure anatomique spécifique du patient. Dans le contexte des radiographies thoraciques (CXR), cela permet d'introduire des anomalies pulmonaires localisées sans perturber les structures stables comme la forme des poumons, les côtes ou les contours cardiaques.

Bien que les modèles de diffusion aient ouvert de nouvelles perspectives, leur application à l'édition médicale rencontre deux défis majeurs :

• L'instabilité structurelle : Les structures anatomiques globales tendent à se stabiliser tôt dans le processus de débruitage et dominent la génération via l'attention globale. Lorsque des prompts pathologiques sont introduits, cette propagation globale peut déformer des régions non ciblées, entraînant une perte de cohérence anatomique.
• L'instabilité de l'expression pathologique : Les lésions médicales sont souvent subtiles, localisées et hétérogènes. Cela génère des signaux de conditionnement faibles et bruyants, conduisant à une localisation imprécise des lésions ou à une étendue incontrôlée de la modification.

Les méthodes existantes nécessitent souvent un réentraînement spécifique au domaine ou des branches de contrôle apprenables, ce qui augmente les coûts de maintenance et limite la généralisation.

2. Méthodologie

Les auteurs proposent un cadre de régulation de l'attention au moment de l'inférence (inference-time), évitant ainsi le réentraînement coûteux du modèle. La méthode s'appuie sur un modèle de diffusion conditionnel (basé sur Stable Diffusion v1.5) et intègre deux modules clés :

A. Régularisation de l'attention consciente de l'anatomie (Anatomy-aware Attention Regularization)

• Objectif : Empêcher la propagation de sémantiques structurelles vers des zones sensibles aux pathologies et limiter les distorsions non intentionnelles.
• Mécanisme : Un module de "gating" (porte) applique un masque d'organe ($M_{anat}$) aux cartes d'auto-attention ($S_t$) du réseau UNet.
• Formule : $S_{anat}^t = S_t \odot (M_{anat} \downarrow q)$, où $\odot$ est la multiplication élément par élément et $q$ est la résolution spatiale.
• Effet : Cela confine les interactions structurelles aux régions d'intérêt (ROI) anatomiques valides, préservant ainsi la cohérence globale de l'image.

B. Régulation de l'attention guidée par la pathologie (Pathology-guided Attention Regulation)

Ce module vise à renforcer la génération de lésions dans les zones cibles et à contrôler leur étendue. Il opère en deux étapes :

1. Repondération de l'attention croisée :
   • Un prior spatial ($\Omega$) est dérivé du masque d'organe et des indices textuels (ex: "poumon droit").
   • Pendant les premières étapes de débruitage ($t < \mu T$), les poids de l'attention croisée pour les tokens liés à la pathologie sont amplifiés dans la zone cible : $A_{t,k}^{path} = A_{t,k} \odot (1 + \eta \Omega)$.
2. Correction légère de l'espace latent :
   • Pour éviter la diffusion de l'attention vers des zones non ciblées, une métrique de concentration différentiable est calculée pour mesurer l'alignement entre l'attribution des tokens pathologiques et la ROI.
   • Une énergie de pathologie ($L_{path}$) est définie et utilisée pour effectuer une mise à jour du gradient sur le latent intermédiaire ($z_t$) : $\hat{z}_t \leftarrow z_t - \alpha_t \nabla_{z_t} L_{path}(t)$.
   • Cela affine la trajectoire de débruitage pour une localisation plus précise de la lésion.

3. Contributions Clés

1. Cadre d'inférence sans réentraînement : Une approche qui impose des contraintes sur le processus d'échantillonnage du diffusion au moment de l'inférence, réduisant les coûts de maintenance et améliorant la généralisation entre différents domaines et dispositifs.
2. Régulation conjointe de l'attention : Combinaison innovante d'une régularisation de l'auto-attention (pour l'anatomie) et d'une régulation de l'attention croisée (pour la pathologie), permettant de dissocier efficacement les variations pathologiques de la cohérence anatomique.
3. Contrôle précis de la localisation et de l'étendue : Grâce à la correction latente guidée par l'énergie d'attention, la méthode permet un contrôle fin de la taille et de la position des lésions générées.

4. Résultats Expérimentaux

Les expériences ont été menées sur les ensembles de données MIMIC-CXR-JPG et ChexpertPlus (images 512x512).

• Comparaison qualitative (Fig. 2) : La méthode proposée génère des images où les modifications pathologiques sont plus précises et mieux confinées aux zones cibles que les méthodes de pointe (SD-inpainting, PIE, BiomedJourney, ProgEmu). Les régions non ciblées (fond, autres organes) restent stables.
• Comparaison quantitative (Tableau 1) :
  • Confiance (Conf) : 0.709 (le plus élevé), indiquant une meilleure adéquation avec le prompt cible.
  • CLIP-I : 0.870 (le plus élevé), montrant une forte similarité sémantique avec l'image originale.
  • LPIPS et FID : Des scores compétitifs (0.18 et 29.0 respectivement) prouvent que la qualité visuelle et la fidélité à la distribution réelle ne sont pas sacrifiées au profit du contrôle.
• Étude d'ablation (Tableau 2) :
  • La configuration complète surpasse toutes les variantes partielles.
  • La suppression de la régularisation de l'auto-attention dégrade la cohérence structurelle (SSIM baisse).
  • La suppression de la régulation de l'attention croisée réduit la précision de la pathologie (Conf et CLIP-I baissent).
  • La correction latente apporte une stabilité supplémentaire, affinant le résultat final.

5. Signification et Impact

Ce travail démontre qu'il est possible de réaliser une édition médicale contrefactuelle fiable et contrôlable sans réentraîner massivement les modèles de diffusion. En résolvant le compromis entre la stabilité structurelle et la génération de pathologies localisées, la méthode offre un outil puissant pour :

• L'interprétabilité des modèles : Comprendre comment les modèles de diagnostic réagissent à des changements pathologiques spécifiques.
• L'augmentation de données : Générer des données synthétiques variées et réalistes pour entraîner des modèles de détection de maladies, en particulier pour des cas rares.
• La modélisation de la progression des maladies : Simuler l'évolution d'une pathologie dans le temps pour des études cliniques.

En résumé, cette approche propose une solution élégante et efficace pour l'édition d'images médicales, en utilisant des priors anatomiques pour guider intelligemment les mécanismes d'attention des modèles de diffusion.