From Point Estimates to Distributions: GMM Pooling for MIL in Preterm Birth Prediction

Résumé Technique : Des estimations ponctuelles aux distributions : le GMM Pooling pour l'apprentissage par instances multiples (MIL) dans la prédiction de l'accouchement prématuré

Formulation du problème
L'accouchement prématuré (AP), défini comme un accouchement avant 37 semaines de gestation, est une cause importante de morbidité et de mortalité néonatales. Bien que les directives cliniques recommandent d'acquérir plusieurs images par échographie transvaginale (TVUS) par patient pour tenir compte de la variabilité dépendante de l'opérateur dans la mesure de la longueur cervicale, les approches actuelles d'apprentissage automatique pour la prédiction de l'AP reposent généralement sur une seule image sélectionnée par patient. Cette pratique néglige la variabilité informative intra-examen (ex. : différences de vue et conditions d'acquisition) et limite les performances prédictives.

Pour y remédier, les auteurs formulent la prédiction de l'AP comme un problème d'apprentissage par instances multiples (Multiple Instance Learning - MIL). Dans ce cadre, chaque patient correspond à un « sac » (bag) de variables quantités d'images TVUS ($X = \{x_1, \dots, x_N\}$) associé à un étiquette unique au niveau du sac ($Y$) indiquant le résultat de l'accouchement (AP vs terme). Le défi central consiste à agréger ces nombres variables de caractéristiques au niveau des instances en une représentation de dimension fixe, tout en préservant la variabilité informative au sein du sac, plutôt que de la réduire à une estimation ponctuelle unique.

Méthodologie : GMM Pooling
Le papier propose un module de GMM pooling (regroupement par modèle de mélange gaussien) pour remplacer les agrégateurs MIL standards (tels que le max, la moyenne ou le pooling par attention), qui résument un sac en une estimation ponctuelle unique dans l'espace des caractéristiques. La méthode proposée modélise explicitement la distribution des caractéristiques au sein du sac pour capturer l'hétérogénéité. L'architecture comprend trois étapes :

1. Extraction de caractéristiques : Un encodeur partagé (implémenté comme un squelette U-Net pour la tâche d'AP) projette chaque image en un vecteur de caractéristiques d'instance. L'encodeur est entraîné conjointement pour la segmentation du canal cervical (comme objectif auxiliaire) et la classification au niveau du sac afin d'apprendre des représentations anatomiquement significatives.
2. Estimation des paramètres du GMM : Au lieu d'une simple moyenne, la méthode estime un GMM sur les caractéristiques des instances. Cela implique :
   • Assignations souples de l'instance aux composantes : Un MLP léger ($f_r$) génère une distribution de probabilité ($r_{i,k}$) sur $M$ composantes de mélange pour chaque instance.
   • Pondération de l'importance de l'instance : Un réseau distinct ($f_w$) calcule les poids des instances ($w_i$) via une fonction softmax.
   • Combinaison par attention : Le poids d'attention final $a_{i,k}$ combine l'importance de l'instance et la responsabilité de la composante ($a_{i,k} = w_i \cdot r_{i,k}$).
   • Apprentissage des paramètres : En utilisant ces poids d'attention, les coefficients de mélange ($\pi_k$), les moyennes ($\mu_k$) et les variances diagonales ($\sigma^2_k$) du GMM sont estimés de bout en bout. Les auteurs optent pour une covariance diagonale afin de garantir la stabilité computationnelle et l'extensibilité.
3. Plongement de distribution (Distribution Embedding) : Pour produire une représentation de sac de longueur fixe, la méthode évalue la log-densité du GMM appris à un ensemble de vecteurs de sondage (probes) apprenables. Les valeurs de log-densité résultantes forment l'embedding final du sac utilisé pour la prédiction.

Contributions clés
Les auteurs identifient trois contributions principales :

1. Module de GMM Pooling : Un nouveau mécanisme de pooling MIL qui préserve les distributions de caractéristiques informatives au sein du sac, répondant au goulot d'étranglement informationnel inhérent au pooling par estimation ponctuelle.
2. Plongement de distribution différentiable : Une méthode pour construire un embedding de longueur fixe via des évaluations de log-densité à des vecteurs de sondage apprenables, permettant au la représentation distributionnelle d'être apprise de bout en bout.
3. Évaluation empirique complète : Une évaluation rigoureuse sur une cohorte privée de haut risque d'AP et sur un benchmark public de métastases ganglionnaires, démontrant que le MIL améliore les performances par rapport à la prédiction au niveau de l'instance et que le GMM pooling offre des performances compétitives à travers les tâches de classification et de régression.

Résultats expérimentaux
La méthode a été évaluée sur deux ensembles de données :

• Cohorte privée d'AP : 182 patientes enceintes à haut risque (44 cas d'AP) avec 1 à 43 images par patiente.
  • Le MIL a significativement amélioré les modèles basés sur l'instance, faisant passer la PR-AUC de 0,44 à 0,56 (max pooling) et de 0,44 à 0,56 (GMM pooling).
  • Le GMM pooling a atteint une PR-AUC de $0,56 \pm 0,03$, comparable au max ($0,57 \pm 0,05$) et au pooling par attention ($0,56 \pm 0,02$). Cependant, le GMM a présenté un écart-type plus faible entre les exécutions, suggérant une meilleure stabilité.
• Benchmark public de métastases ganglionnaires : 933 échantillons d'entraînement, 668 de validation et 736 de test.
  • Le GMM pooling a atteint des performances de pointe avec un score F1 de $0,91 \pm 0,01$ et un ROC-AUC de $0,89 \pm 0,01$ pour la classification.
  • Pour la régression (pourcentage de pixels métastatiques), il a atteint une erreur absolue moyenne (MAE) de $0,18 \pm 0,02$, compétitive avec les meilleures méthodes.

Signification et affirmations
Le papier affirme que la modélisation des distributions de caractéristiques via le GMM pooling offre une alternative robuste au pooling standard par estimation ponctuelle. Alors que les agrégateurs standards réduisent la variabilité intra-sac, le GMM pooling capture la structure distributionnelle conjointe à travers les dimensions de caractéristiques. Les auteurs notent que bien que l'approche introduise des hyperparamètres spécifiques à la distribution (nombre de composantes, sondes, température) nécessitant un réglage dépendant du jeu de données, les résultats sur deux tâches d'imagerie médicale distinctes suggèrent que la modélisation distributionnelle se généralise bien. Ce travail souligne qu'exploiter toutes les images disponibles par patient via le MIL est crucial pour la prédiction de l'AP, particulièrement dans les cohortes à haut risque où la séparabilité des classes est faible.