Halfway to 3D: Ensembling 2.5D and 3D Models for Robust COVID-19 CT Diagnosis

Probleemstelling

Het paper adresseert de uitdagingen bij het ontwikkelen van robuuste deep learning-modellen voor de diagnose van COVID-19 en andere longziekten op basis van thorax-CT-scans. Hoewel CT-scans gedetailleerde volumetrische informatie bieden, ondervinden modellen in de praktijk grote problemen door:

1. Domeinverschuivingen (Domain Shifts): Klinische datasets worden vaak verzameld van meerdere instellingen met verschillende scanners, opnameprotocollen en patiëntenpopulaties. Dit leidt tot een daling in de generalisatiecapaciteit van modellen.
2. Fairness (Eerlijkheid): Modellen vertonen vaak prestatieverschillen tussen verschillende demografische groepen (bijv. geslacht), wat een groot risico vormt voor betrouwbare medische AI.
3. Complementaire informatie: Bestaande methoden focussen vaak óf op 2D-slices (die fijne details missen) óf op volledige 3D-volumes (die computationeel duur zijn en soms grootschalige patronen over het hoofd zien).

Het doel is een framework te ontwikkelen dat robuust is tegen heterogene databronnen en eerlijke prestaties levert over verschillende populaties heen.

Methodologie

De auteurs stellen een hybride deep learning-framework voor dat 2.5D (multi-view slices) en 3D (volledige volumetrische) representaties combineert via een ensemble-strategie.

1. Data Preprocessing

• Ruwe axiale CT-slices worden gereconstrueerd tot volledige 3D-volumes (128 × 128 × 128).
• Toepassing van een 3D Gaussische denoising-filter en mask-scherping voor anatomische duidelijkheid.
• Normalisatie van voxelintensiteiten naar het bereik [0, 255].

2. De 3D-Branche (Volumetrisch)

• Architectuur: Een ResNet-18 aangepast voor single-channel CT-input.
• Trainingsstrategie:
  • Fase 1 (Domain Generalization): Pre-training met Variance Risk Extrapolation (VREx). Deze techniek dwingt het model om domein-invariante representaties te leren door de variantie van de verliesfuncties over verschillende domeinen (bronnen) te minimaliseren.
  • Fase 2 (Fine-tuning): Supervised Contrastive Learning wordt toegepast om de scheidbaarheid van features te verbeteren, gecombineerd met MixUp augmentatie voor gladdere beslissingsgrenzen.

3. De 2.5D-Branche (Multi-View)

• Extractie: Uit elk gereconstrueerd volume worden slices geëxtraheerd uit drie orthogonale vlakken: axiaal, coronaal en sagittaal.
• Architectuur: Een Vision Transformer (ViT) backbone gebaseerd op DINOv3 (zelftoezichtend voorvertrouwd model).
• Verwerking: Slices worden onafhankelijk verwerkt en de resulterende embeddings worden op feature-niveau gefuseerd. Dit benut de sterke visuele representaties van DINOv3 voor medische beeldvorming.

4. Ensemble Integratie

• De voorspellingen van beide takken worden gecombineerd via logit-level ensembling (gemiddelde van de logit-scores).
• Dit combineert de globale context van het 3D-model met de hoge resolutie, slice-gebaseerde details van het 2.5D-model.

Belangrijkste Bijdragen

1. Hybride Framework: Een innovatieve combinatie van multi-view slice-representaties (2.5D) en volumetrische modellering (3D) voor CT-diagnose.
2. Robuuste Training: Implementatie van VREx en supervised contrastive learning in de 3D-tak om domein-generalisatie te verbeteren.
3. Gebruik van Pre-trained Vision Transformers: Toepassing van DINOv3 voor het 2.5D-pipeline, wat bewijst dat grote, voorvertrouwde visuele modellen effectief kunnen worden getransfereerd naar medische taken.
4. Validatie op Eerlijkheid en Robuustheid: Demonstratie dat ensemble-strategieën de robuustheid en eerlijkheid (over geslachtsgroepen) verbeteren in multi-source medische taken.

Resultaten

De methode werd geëvalueerd op de PHAROS-AIF-MIH benchmark, die twee taken omvatte:

• Taak 1: Binair COVID-19 Detectie
  • Het ensemble bereikte een Macro F1-score van 0.751 op de testset, wat leidde tot een tweede plaats in de Multi-Source COVID-19 Detection Challenge.
  • Het ensemble presteerde beter dan de individuele modellen, vooral in het overbruggen van prestatiedalingen bij specifieke ziekenhuisbronnen (bijv. H2).
• Taak 2: Multi-class Diagnose (Gezond, Adenocarcinoom, Plaveiselcelcarcinoom, COVID-19)
  • Het ensemble behaalde een Macro F1-score van 0.633, wat een derde plaats opleverde in de Fair Disease Diagnosis Challenge.
  • Fairness: Er was een merkbare prestatiekloof tussen geslachten. Het 0.7-gewogen ensemble verbeterde de prestaties voor mannelijke patiënten aanzienlijk, wat leidde tot een beter evenwicht tussen geslachten.
  • Opmerkelijk: Voor de multi-class taak presteerde het 2.5D-model (DINOv3) vaak beter dan het 3D-model of het ensemble, wat suggereert dat slice-gebaseerde features voldoende zijn voor complexe pathologie-gradering.

Betekenis en Conclusie

Het paper toont aan dat er geen "one-size-fits-all" oplossing is voor medische beeldanalyse.

• Complementaire Krachten: 3D-modellen zijn superieur in het vastleggen van globale context voor binair onderscheid (COVID vs. niet-COVID), terwijl 2.5D-modellen (vooral met DINOv3) beter zijn in het onderscheiden van complexe, multi-class pathologieën door hun vermogen om fijne slice-details te analyseren.
• Praktische Toepassing: De ensemble-strategie biedt een praktische route naar klinisch bruikbare tools die bestand zijn tegen de domeinverschuivingen die inherent zijn aan real-world ziekenhuisomgevingen.
• Fairness: Door het combineren van modellen en het aanpassen van ensemble-weights, kan de systematische bias ten opzichte van bepaalde demografische groepen worden verminderd.

Kortom, dit werk benadrukt dat het integreren van verschillende representatiedimensies (2.5D en 3D) essentieel is voor het creëren van robuuste, eerlijke en nauwkeurige AI-systemen voor medische diagnostiek.