Token-UNet: A New Case for Transformers Integration in Efficient and Interpretable 3D UNets for Brain Imaging Segmentation

Dit artikel introduceert Token-UNet, een efficiënt en interpreteerbaar 3D-segmentatiemodel voor hersenbeeldvorming dat convolutie en transformatormodules combineert om de rekenkosten en het geheugengebruik aanzienlijk te verlagen ten opzichte van SwinUNETR, terwijl de prestaties worden verbeterd.

De kern

Token-UNet: De Slimme Vertaler voor Medische Beelden

Stel je voor dat je een gigantische, driedimensionale puzzel hebt: een MRI-scan van een hersentumor. Voor een computer is dit een enorme berg data, vol met miljoenen kleine blokjes (voxels). Om deze puzzel op te lossen en de tumor precies af te bakenen, gebruiken artsen en onderzoekers kunstmatige intelligentie (AI).

Maar hier zit het probleem: de meest geavanceerde AI-modellen, die we "Transformers" noemen, zijn als olifanten in een porseleinen winkel. Ze zijn extreem slim en kunnen het hele plaatje tegelijk zien, maar ze zijn ook zo zwaar en hongerig naar rekenkracht dat ze alleen werken in superduurzame datacenters. De meeste ziekenhuizen en kleine onderzoekslabs hebben die krachtige computers niet. Ze zitten vast met "normale" computers die deze modellen niet aankunnen.

De Oplossing: Token-UNet

In dit artikel presenteren de auteurs Token-UNet. Dit is een nieuwe manier om die slimme AI-modellen te maken, maar dan zo licht en efficiënt dat ze op een gewone laptop of een standaard ziekenhuiscomputer kunnen draaien.

Hier is hoe het werkt, vertaald in alledaagse termen:

1. Het Oude Probleem: Te Veel Detail

Stel je voor dat je een boek moet samenvatten. De oude, zware AI-modellen (zoals SwinUNETR) proberen elk woord in het boek te lezen, te analyseren en te onthouden voordat ze een conclusie trekken. Bij een 3D-hersenscan betekent dit dat de computer elke millimeter van de hersenen apart moet bekijken. Dit kost enorm veel tijd en geheugen. Het is alsof je een heel bos probeert te tellen, boom voor boom, terwijl je eigenlijk alleen maar wilt weten waar de brand is.

2. De Nieuwe Methode: De "Token" Vertaler

Token-UNet gebruikt een slimme truc met twee nieuwe onderdelen: TokenLearner en TokenFuser.

• TokenLearner (De Slimme Samenvatter):
  In plaats van elk woord apart te lezen, kijkt deze module naar het hele verhaal en zegt: "Oké, dit stukje tekst gaat over de brand, dit stukje over de wind, en dit stukje is gewoon ruis."
  De computer groepeert miljoenen kleine beeldpuntjes samen in slechts een handjevol "tokens" (samenvattingen).

  • Analogie: In plaats van 10.000 losse foto's van een stad te bekijken, maakt de AI één collage van de belangrijkste gebouwen. Hij zegt: "Hier is het ziekenhuis, hier is de brand, hier is de rivier." Hij negeert de onbelangrijke details. Dit vermindert de hoeveelheid werk met wel 90%!

• De Transformer (De Denker):
  Nu, met deze kleine, samengevatte collage (de tokens), kan de slimme "Transformer"-AI aan de slag. Omdat er maar een paar tokens zijn, kan deze AI razendsnel en zonder zware computers nadenken over de relaties tussen de brand en de rest van de stad. Hij is nog steeds super slim, maar nu werkt hij op een fiets in plaats van op een vrachtwagen.

• TokenFuser (De Terugbrenger):
  Nadat de AI heeft nagedacht en de tumor heeft gevonden, moet hij dit resultaat weer terugvertalen naar het oorspronkelijke, gedetailleerde beeld. De TokenFuser neemt de conclusies van de tokens en "plakt" ze weer terug op de juiste plekken in de 3D-scan, zodat de arts precies ziet waar de tumor zit.

3. Waarom is dit zo belangrijk?

• Toegankelijkheid: Omdat Token-UNet zo weinig rekenkracht nodig heeft, kunnen nu ook kleine ziekenhuizen en labs in ontwikkelingslanden de beste AI-tools gebruiken. Je hebt geen miljoenen euro's aan supercomputers meer nodig.
• Snelheid: Het model is veel sneller. Wat eerder uren duurde, gaat nu in minuten.
• Betrouwbaarheid (Interpreteerbaarheid): Dit is misschien wel het coolste deel. Omdat de AI werkt met "tokens" die specifieke delen van de hersenen vertegenwoordigen, kunnen we zien waar de AI naar kijkt.
  • Voorbeeld: Als de AI een tumor detecteert, kunnen we een kaartje zien waarop de AI aangeeft: "Ik kijk hier naartoe omdat ik hier een donkere vlek zie die lijkt op een tumor." Dit helpt artsen om de AI te vertrouwen, in plaats van dat het een "black box" is die alleen een ja/nee antwoord geeft.

Het Resultaat

De onderzoekers hebben getest of hun nieuwe model net zo goed werkt als de zware, dure modellen. Het antwoord is verrassend: Ja!
Token-UNet is net zo nauwkeurig in het vinden van hersentumoren (zelfs iets beter dan de zware modellen), maar het is:

• 3 keer lichter qua geheugen.
• 10 keer sneller in gebruik.
• 3 keer minder zwaar qua parameters (de "hersenen" van de AI).

Conclusie
Token-UNet bewijst dat je niet altijd de zwaarste, duurste technologie nodig hebt om de beste resultaten te behalen. Door slimme vertalers (TokenLearner) te gebruiken die de informatie samenvatten voordat de "supercomputer" eraan te pas komt, maken we medische AI democratischer. Het stelt de wereld in staat om sneller, goedkoper en transparanter te werken aan het vinden van ziektes, zonder dat je een datacenter in je kelder nodig hebt.

Technische samenvatting

Probleemstelling

De toepassing van geavanceerde Deep Learning-modellen, specifiek Transformers, voor 3D medische beeldsegmentatie (zoals hersentumoren in MRI-scans) stuit op aanzienlijke computatiebeperkingen.

• Rekenkracht en Geheugen: Bestaande state-of-the-art modellen, zoals SwinUNETR, gebruiken Transformers als encoder. Deze verwerken de invoer als "tokens" (kleine subvolumes van 8³ voxels). Omdat de aandachtsmechanisme (Self-Attention) kwadratisch schaalt met het aantal tokens, en het aantal tokens kubisch toeneemt met de 3D-resolutie, worden deze modellen extreem rekenintensief.
• Toegankelijkheid: Dit vereist zware hardware (meerdere GPU's) die vaak niet beschikbaar is in standaard ziekenhuizen of academische laboratoria. Hierdoor kunnen veel instellingen geen gebruik maken van de beste beschikbare tools voor patiëntenzorg.
• Trade-off: Er is een duidelijke spanning tussen de hoge prestaties van Transformer-architecturen en de noodzaak voor efficiëntie, interpreteerbaarheid en het gebruik van beperkte hardware.

Methodologie: Token-UNet

De auteurs stellen Token-UNet voor, een hybride architectuur die de sterktes van convolutie (voor lokale patronen) en attention (voor globale interacties) combineert, maar met een fundamenteel andere aanpak voor tokenisatie.

Kerncomponenten:

1. Convolutionele Encoder/Decoder: In plaats van de invoer direct te tokeniseren (zoals bij ViT of SwinUNETR), behoudt Token-UNet een conventionele convolutie-encoder (gebaseerd op een verbeterde UNet-variant, genaamd UNet). Deze verwerkt lokale ruimtelijke patronen efficiënt.
2. TokenLearner: Dit is het cruciale innovatieve blok dat tussen de encoder en decoder wordt geplaatst.
   • Het neemt de 3D feature maps van de encoder en reduceert deze tot een vast, klein aantal tokens ($N=8$ in dit paper), ongeacht de oorspronkelijke invoergrootte.
   • Het gebruikt een Multi-Layer Perceptron (MLP) om voor elke voxel te bepalen tot welke van de $N$ semantische klassen deze behoort.
   • Hierdoor worden $N$ spatial attention maps gegenereerd. Tokens worden gepoold op basis van deze scores, waardoor alleen semantisch relevante informatie wordt behouden.
   • Dit doorbreekt de koppeling tussen invoerresolutie en het aantal tokens, wat de complexiteit van de daaropvolgende Transformer drastisch verlaagt.
3. Transformer Encoder: Een kleine Transformer (4 blokken) verwerkt de gegenereerde tokens. Omdat het aantal tokens vast en klein is, is de Self-Attention-berekening zeer efficiënt en niet afhankelijk van de beeldgrootte.
4. TokenFuser: Dit module zet de verwerkte tokens terug om naar de oorspronkelijke 3D-ruimtelijke resolutie.
   • Het genereert nieuwe spatial attention masks en mixt de tokens lineair met de feature maps.
   • Dit zorgt voor een naadloze overgang naar de decoder voor de uiteindelijke segmentatie.

Architectuurvarianten:
De auteurs testen verschillende configuraties: een basis UNet, een verbeterde UNet** (met additieve skip-connections in plaats van concatenatie), en Token-UNet met en zonder de Transformer-blokken, om de bijdrage van elk component te isoleren.

Belangrijkste Bijdragen

• Efficiënte Integratie van Transformers: Het paper toont aan dat Transformers kunnen worden ingebouwd in UNet-achtige structuren zonder de rekenkosten van volledige 3D-tokenisatie, door gebruik te maken van TokenLearner en TokenFuser.
• Decoupling van Resolutie en Tokens: Het aantal tokens is vastgesteld (bijv. 8), waardoor de modelcomplexiteit niet meer exponentieel toeneemt met de beeldgrootte.
• Interpreteerbaarheid: De TokenLearner module genereert natuurlijke attention maps. Deze visualiseren welke delen van de hersenen (bijv. tumorcore, randen, ventrikels) het model als relevant beschouwt, wat essentieel is voor medische betrouwbaarheid.
• Democratisering van AI: Het model is ontworpen om te draaien op standaard hardware (één GPU), waardoor geavanceerde 3D-segmentatie toegankelijk wordt voor kleinere onderzoeksgroepen en klinieken.

Resultaten

De modellen werden getest op de FeTS 2022 / BraTS dataset (1251 patiënten met glioblastoma) met 5-voudige kruisvalidatie.

• Prestaties (Dice Score):
  • Token-UNet: 87.21% ± 0.35%
  • SwinUNETR (SOTA): 86.75% ± 0.19%
  • Conclusie: Token-UNet presteert statistisch significant beter dan SwinUNETR, ondanks een veel kleinere architectuur.
• Efficiëntie:
  • Parameters: Het zwaarste Token-UNet model heeft slechts 35% van het aantal parameters van SwinUNETR.
  • Geheugen (Memory Footprint): Verminderd tot 33% van SwinUNETR.
  • Inferentie-tijd: Verminderd tot 10% van de tijd die SwinUNETR nodig heeft.
• Conclusie: Het toevoegen van TokenLearner en TokenFuser aan een UNet levert een grotere prestatieverbetering op dan het toevoegen van de Transformer zelf, wat aangeeft dat de tokenisatiestrategie de belangrijkste factor is.

Betekenis en Impact

Dit werk is een belangrijke stap in de richting van efficiënte foundation models voor de biomedische beeldvorming.

1. Toegankelijkheid: Het bewijst dat top-presterende medische AI niet afhankelijk hoeft te zijn van "elite" hardware. Dit stelt meer onderzoekers in staat om nieuwe methoden te ontwikkelen, te fine-tunen en te valideren.
2. Interpreteerbaarheid: De gegenereerde attention maps bieden inzicht in het besluitvormingsproces van het model, wat cruciaal is voor het vertrouwen van artsen en het analyseren van foutgevallen.
3. Toekomstperspectief: De methode opent de deur voor het gebruik van self-supervised learning en pre-training op grote datasets, zelfs met beperkte middelen, omdat de tokenisatie de rekenlast beheersbaar houdt.

Kortom, Token-UNet herdefinieert de balans tussen convolutie en attention, en biedt een schaalbare, interpreteerbare en hardware-vriendelijke oplossing voor complexe 3D medische segmentatietaken.