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

Technische Samenvatting: Van Punt-schattingen naar Distributies: GMM-pooling voor MIL bij de voorspelling van vroeggeboorte

Probleemformulering
Vroeggeboorte (PTB), gedefinieerd als een bevalling vóór de 37 weken zwangerschap, is een belangrijke oorzaak van neonatale morbiditeit en mortaliteit. Hoewel klinische richtlijnen aanbevelen om meerdere transvaginale echografiebeelden (TVUS) per patiënt te verkrijgen om de door de operator afhankelijke variabiliteit in de meting van de lengte van de baarmoederhals op te vangen, vertrouwen bestaande machine learning-benaderingen voor PTB-voorspelling doorgaans op één geselecteerd beeld per patiënt. Deze praktijk verwerpt informatieve intra-examen variabiliteit (bijv. verschillen in weergave en acquisitiecondities) en beperkt de voorspellende prestaties.

Om dit aan te pakken, formuleren de auteurs PTB-voorspelling als een Multiple Instance Learning (MIL) probleem. In dit kader komt elke patiënt overeen met een "bag" van variabele groottes aan TVUS-beelden ($X = \{x_1, \dots, x_N\}$) geassocieerd met een enkel bag-niveau label ($Y$) dat de uitkomst van de bevalling aangeeft (PTB versus termijn). De kernuitdaging is het aggregeren van deze variabele aantallen instance-niveau kenmerken naar een representatie met een vaste dimensie, terwijl de informatieve variabiliteit binnen de bag behouden blijft, in plaats van deze te reduceren tot een enkele punt-schatting.

Methodologie: GMM Pooling
Het artikel stelt een Gaussian Mixture Model (GMM) pooling module voor om standaard MIL-aggregators (zoals max-, gemiddelde- of attention-pooling) te vervangen, die een bag samenvatten als een enkele punt-schatting in de feature-ruimte. De voorgestelde methode modelleert de binnen-bag feature-distributie expliciet om heterogeniteit te vangen. De architectuur bestaat uit drie stadia:

1. Feature Extractie: Een gedeelde encoder (geïmplementeerd als een U-Net backbone voor de PTB-taak) brengt elk beeld in kaart naar een instance feature-vector. De encoder wordt gezamenlijk getraind voor zowel de segmentatie van de baarmoederhalskanaal (als een hulpdoelstelling) als voor bag-niveau classificatie om anatomisch betekenisvolle representaties te leren.
2. GMM Parameter Schatting: In plaats van eenvoudige middeling, schat de methode een GMM over de instance features. Dit omvat:
   • Soft Instance-to-Component Assignments: Een lichtgewicht MLP ($f_r$) genereert een waarschijnlijkheidsverdeling ($r_{i,k}$) over $M$ mengcomponenten voor elke instance.
   • Instance Importance Weighting: Een apart netwerk ($f_w$) berekent instance-gewichten ($w_i$) via softmax.
   • Attention Combinatie: De uiteindelijke attention-weging $a_{i,k}$ combineert de instance-belangrijkheid en component-verantwoordelijkheid ($a_{i,k} = w_i \cdot r_{i,k}$).
   • Parameter Leren: Met behulp van deze attention-gewichten worden de mengcoëfficiënten ($\pi_k$), gemiddelden ($\mu_k$) en diagonale varianties ($\sigma^2_k$) van de GMM end-to-end geschat. De auteurs kiezen voor diagonale covariantie om computationele stabiliteit en schaalbaarheid te garanderen.
3. Distributie Embedding: Om een embedding met een vaste lengte te produceren, evalueert de methode de log-densiteit van de geleerde GMM op een set van $P$ leerbare probe-vectoren. De resulterende log-densiteitswaarden vormen de uiteindelijke bag-embedding die wordt gebruikt voor voorspelling.

Belangrijkste Bijdragen
De auteurs identificeren drie primaire bijdragen:

1. GMM Pooling Module: Een nieuwe MIL-poolingmechanisme die informatieve binnen-bag feature-distributies behoudt, waarmee de informatie-bottleneck inherent aan punt-schatting pooling wordt aangepakt.
2. Differentiabele Distributie Embedding: Een methode om een embedding met een vaste lengte te construeren via log-densiteit evaluaties op leerbare probe-vectoren, waardoor de distributionele representatie end-to-end geleerd kan worden.
3. Uitgebreide Empirische Evaluatie: Een rigoureuze evaluatie op een private hoog-risico PTB-cohort en een publieke lymfeklier metastase benchmark, die aantoont dat MIL beter presteert dan instance-niveau voorspelling en dat GMM pooling competitieve prestaties biedt over zowel classificatie- als regressietaken.

Experimentele Resultaten
De methode werd geëvalueerd op twee datasets:

• Private PTB Cohort: 182 hoog-risico zwangere patiënten (44 PTB gevallen) met 1–43 beelden per patiënt.
  • MIL verbeterde aanzienlijk ten opzichte van instance-gebaseerde modellen, waarbij de PR-AUC steeg van 0,44 naar 0,56 (max pooling) en 0,56 (GMM pooling).
  • GMM pooling behaalde een PR-AUC van $0,56 \pm 0,03$, vergelijkbaar met max ($0,57 \pm 0,05$) en attention ($0,56 \pm 0,02$) pooling. Echter, GMM vertoonde een lagere standaarddeviatie over verschillende runs, wat wijst op verbeterde stabiliteit.
• Publieke Lymfeklier Metastase Benchmark: 933 trainings-, 668 validatie- en 736 test-samples.
  • GMM pooling bereikte state-of-the-art prestaties met een F1-score van $0,91 \pm 0,01$ en een ROC-AUC van $0,89 \pm 0,01$ voor classificatie.
  • Voor regressie (metastatisch pixelpercentage) behaalde het een Mean Absolute Error (MAE) van $0,18 \pm 0,02$, competitief met de best presterende methoden.

Significantie en Claims
Het artikel claimt dat het modelleren van feature-distributies via GMM pooling een robuust alternatief biedt voor standaard punt-schatting pooling. Terwijl standaard aggregators de binnen-bag variabiliteit reduceren, legt GMM pooling de gezamenlijke distributionele structuur over de feature-dimensies vast. De auteurs merken op dat hoewel de aanpak specifieke distributie-gerelateerde hyperparameters introduceert (aantal componenten, probes, temperatuur) die dataset-afhankelijke afstemming vereisen, de resultaten over twee verschillende medische beeldvormingstaken wijzen op een goede generalisatie van het distributionele modelleren. Het werk benadrukt dat het benutten van alle beschikbare beelden per patiënt via MIL cruciaal is voor PTB-voorspelling, met name in hoog-risico cohorten waar de klassenseparabiliteit laag is.