Towards Fair and Robust Volumetric CT Classification via KL-Regularised Group Distributionally Robust Optimisation

Probleemstelling

De automatische diagnose van longziekten (zoals COVID-19 en longkanker) op basis van thorax-CT-scans stuit in klinische toepassingen op twee fundamentele uitdagingen:

1. Distributieverandering (Domain Shift): Scans variëren sterk tussen verschillende ziekenhuizen door verschillen in scannerhardware, reconstructiekernen en beeldvormingsprotocollen. Modellen getraind op één locatie generaliseren vaak slecht naar andere locaties.
2. Demografische Ongelijkheid: Classificatiemodellen presteren vaak slechter op ondervertegenwoordigde subgroepen (bijv. specifieke geslachts- en ziektecategorieën-combinaties) omdat standaard trainingsdoelstellingen de gemiddelde prestatie optimaliseren ten koste van minderheidsgroepen.

Deze problemen zijn verergerd in kleine, multi-centrale datasets waar bestaande methoden voor domeingeneralisatie vaak falen.

Methodologie

Het auteurs stellen een framework voor dat twee complementaire taken aanpakt:

• Taak 1: Binair COVID-19-classificatie over vier verschillende opnamelocaties.
• Taak 2: Vier-klassen longpathologie-herkenning met gender-gerechtigheidseisen.

Het framework bestaat uit drie kerncomponenten:

1. Architectuur: Slice Encoder en Volumetric Aggregator

Om de hoge rekentijd van volledige 3D-volumen te vermijden en toch ruimtelijke context te behouden, gebruiken de auteurs een hybride aanpak:

• Slice Encoder: Elke individuele axiale slice wordt onafhankelijk verwerkt door een gedeelde MobileViT-XXS backbone. Dit is een lichtgewicht model (ongeveer 1,3 miljoen parameters) dat convolutieblokken combineert met lokale Transformer-blokken. De invoer wordt aangepast voor single-channel (grayscale) CT-data.
• SliceTransformer Aggregator: De embeddings van 64 slices worden samengevoegd tot een sequentie en verwerkt door een twee-laags Transformer-encoder. Deze laag gebruikt self-attention om te leren welke anatomische niveaus de sterkste diagnostische signalen bevatten, in plaats van alle slices gelijk te behandelen (zoals bij eenvoudige pooling).
• Classificatiehead: Een lineaire classifier produceert de uiteindelijke voorspellingen.

2. Trainingsdoel: KL-Geregulariseerde Group DRO

In plaats van Empirical Risk Minimization (ERM), die de gemiddelde fout minimaliseert, optimaliseren de auteurs voor de worst-case groepsfout via Group Distributionally Robust Optimisation (Group DRO).

• Groepsdefinitie:
  • Voor Taak 1: Groepen zijn de vier opnamelocaties.
  • Voor Taak 2: Groepen zijn fijne combinaties van Geslacht × Klasse (bijv. Vrouwelijk + Plaveiselcelcarcinoom). Dit is cruciaal voor zeer ondervertegenwoordigde groepen (slechts 5 trainingssamples).
• KL-Regularisatie: Standaard Group DRO kan leiden tot "weight collapse", waarbij de gewichten volledig naar één moeilijke groep springen, wat de prestaties op andere groepen schaadt. De auteurs introduceren een Kullback-Leibler (KL) straffing naar een uniforme verdeling.
  • De loss-functie is: $L = \sum w_g \ell_g + \alpha \cdot D_{KL}(w \parallel u)$.
  • De parameter $\alpha$ balanceert tussen het beschermen van de slechtst presterende groepen en het behouden van stabiele gemiddelde prestaties.

Belangrijkste Bijdragen

1. Efficiënte Volumetrische Architectuur: Een framework dat MobileViT-XXS combineert met een SliceTransformer, wat concurrerende prestaties levert met minder dan 2 miljoen parameters, ideaal voor kleine medische datasets.
2. KL-Geregulariseerde Group DRO: Een trainingsdoel dat adaptief gewichten toewijst aan ondervertegenwoordigde opnamelocaties en subgroepen. De KL-term voorkomt instabiliteit en zorgt voor een evenwicht tussen robuustheid en gemiddelde prestaties.
3. Fijne Groepsgranulariteit: Voor de fairnesstaken worden groepen gedefinieerd op het niveau van geslacht × ziekteklasse, wat directe aandacht geeft aan extreem zeldzame combinaties die anders zouden worden overschaduwd door meerderheidsgroepen.
4. Systematische Analyse: Een uitgebreide ablatiestudie van de regularisatiesterkte ($\alpha$) en het aggregator-ontwerp, wat waardevolle inzichten biedt voor praktici.

Resultaten

De resultaten tonen aanzienlijke verbeteringen ten opzichte van bestaande methoden en uitdaging-baselines:

• Taak 1 (COVID-19):

  • De beste configuratie ($\alpha = 0.5$) bereikte een F1-score van 0,835.
  • Dit is een verbetering van +5,9% ten opzichte van de beste eerder gepubliceerde uitdaging-deelname (0,776).
  • Opmerking: De prestaties op "Centre 2" bleven laag (F1 ≈ 0), wat wijst op een fundamentele beperking van reweighting-methoden bij volledig ongezette domeinen die niet in de trainingsdata voorkomen.

• Taak 2 (Longpathologie & Fairness):

  • De methode bereikte een gemiddelde per-geslacht macro F1 van 0,815.
  • Dit is +11,1% beter dan de beste uitdaging-deelname.
  • Specifiek voor de ondervertegenwoordigde groep Vrouwelijk Plaveiselcelcarcinoom steeg de F1-score met +17,4% ten opzichte van de Focal Loss-baseline.
  • De waarde $\alpha = 0.5$ bleek optimaal om de kloof tussen geslachten te minimaliseren zonder de prestaties van de meerderheidsgroep te veel te laten dalen.

Betekenis en Conclusie

Dit werk demonstreert dat het combineren van lichtgewicht transformer-architecturen met een KL-geregulariseerde Group DRO-objectief een effectieve strategie is voor het oplossen van zowel domeinverschuiving als demografische ongelijkheid in medische beeldvorming.

De studie benadrukt dat:

1. Fijne granulariteit in groepsdefinitie essentieel is om extreme ondervertegenwoordiging aan te pakken.
2. Regularisatie (KL) noodzakelijk is om de stabiliteit van Group DRO te garanderen op kleine datasets.
3. Hoewel gewichtshervorming effectief is voor bekende verschuivingen, zijn aanvullende strategieën (zoals testtijd-adaptatie) nodig voor domeinen die volledig afwezig zijn in de trainingsdata.

De auteurs zien modelcompressie voor edge-apparaten (zoals draagbare scanners) als een natuurlijke volgende stap voor de implementatie in hulpbronnenbeperkte klinische omgevingen.