DC-ViT: Modulating Spatial and Channel Interactions for Multi-Channel Images

Probleemstelling

Het trainen en evalueren van visuele encoders op Multi-Channel Imaging (MCI) data is fundamenteel uitdagend vanwege de heterogeniteit in kanaalconfiguraties. In tegenstelling tot standaard RGB-afbeeldingen (drie vaste kanalen), kunnen MCI-data (zoals in microscopie of aardobservatie) een willekeurig aantal kanalen bevatten met verschillende semantische betekenissen afhankelijk van de staining-protocollen, sensortypes of opnamemodi.

Bestaande Multi-Channel Vision Transformers (MC-ViTs) lossen dit op door tokens van alle kanalen in één lange sequentie te combineren en gezamenlijk te verwerken via gedeelde self-attention lagen. Dit leidt echter tot twee problemen:

1. Feature Verdunning (Feature Dilution): De onbeperkte interactie tussen tokens van verschillende kanalen kan leiden tot het "verwateren" van kanaalspecifieke semantiek. Kanalen in MCI-data vertegenwoordigen vaak fundamenteel verschillende concepten (bijv. celkern vs. mitochondriën), die niet zomaar gemengd mogen worden.
2. Gebrek aan Controle: Bestaande methoden hebben geen expliciete controle over hoe cross-kanaal informatie wordt geïntegreerd. Hoewel volledige isolatie van kanalen ook suboptimaal is (omdat kanalen vaak complementaire context bieden), is de huidige "unrestricted token mixing" schadelijk voor de representatiekracht.

Methodologie: DC-ViT

De auteurs stellen Decoupled Vision Transformer (DC-ViT) voor, een architectuur die de informatie-uitwisseling tussen kanalen expliciet reguleert in plaats van deze volledig te laten of volledig te blokkeren. De methode bestaat uit twee kerncomponenten:

1. Decoupled Self-Attention (DSA)

DSA deconstrueert de standaard self-attention mechanisme in twee complementaire paden:

• Ruimtelijke Updates (Spatial): Deze modus modelleert de structuur binnen een enkel kanaal. Tokens van hetzelfde kanaal wisselen informatie uit over hun ruimtelijke locaties, wat de kanaalspecifieke kenmerken versterkt.
• Kanaal-Updates (Channel-wise): Deze modus integreert adaptief informatie tussen kanalen op vaste ruimtelijke locaties.
• Selectieve Interactie: In tegenstelling tot MC-ViTs waar alle tokens in één grote sequentie zitten, wordt DSA toegepast op een tensor met dimensies $C \times N \times D$ (Kanaal $\times$ Tokens $\times$ Dimensie). Kanaal-attention wordt niet op elke laag toegepast, maar alleen op een geselecteerde subset van lagen $\mathcal{M}$.
  • In lagen waar $\mathcal{M}$ niet geldt, wordt alleen ruimtelijke attention uitgevoerd (kanalen blijven gescheiden).
  • In lagen waar $\mathcal{M}$ geldt, wordt een gewogen som genomen van kanaal- en ruimtelijke attention, waarbij $\alpha$ een leerbare schaalparameter is.
  • Dit creëert een "bottleneck" die zorgt dat de encoder eerst sterke, onafhankelijke kanaalrepresentaties leert, voordat er selectieve cross-kanaal interactie plaatsvindt.

2. Decoupled Aggregation (DAG)

Bestaande MC-ViTs gebruiken vaak een cls-token of globale pooling over alle tokens, wat de unieke bijdrage van elk kanaal negeert. DAG introduceert een hiërarchische aggregatie:

1. Ruimtelijke Pooling: Eerst wordt er gepoold binnen elk kanaal (over de $N$ tokens), wat resulteert in een kanaalspecifieke embedding.
2. Kanaal Pooling: Vervolgens worden deze kanaal-embeddings geaggregeerd tot de uiteindelijke representatie.
   Deze tweede stap maakt het mogelijk om kanaalbelang (channel importance) te leren, wat cruciaal is voor datasets waar sommige kanalen meer informatief zijn dan anderen voor de specifieke taak.

Belangrijkste Bijdragen

1. Decoupled Self-Attention (DSA): Een nieuw attention-mechanisme dat ruimtelijke en kanaalinteracties expliciet scheidt om feature-verwatering te voorkomen terwijl nuttige contextdeling mogelijk blijft.
2. Decoupled Aggregation (DAG): Een hiërarchisch pooling-schema dat rekening houdt met de heterogene aard van MCI-data door eerst per kanaal te aggregeren en daarna de kanalen te combineren.
3. Architectonische Generalisatie: DC-ViT is ontworpen als een plug-in vervanging voor MC-ViT blokken en werkt effectief zowel bij training vanaf nul als bij transfer learning van pre-getrainde encoders.

Resultaten

De auteurs evalueren DC-ViT op drie standaard MCI-benchmarks: CHAMMI (microscopie), JUMP-CP (chemische perturbaties) en So2Sat (aardobservatie).

• Prestatieverbetering: DC-ViT presteert consistent beter dan bestaande MC-ViT baselines (zoals ChAdaViT, ChannelViT, DiChaViT en ChaMAEViT) op bijna alle metrieken.
• Transfer Learning: Wanneer DC-ViT wordt geïnitieerd met pre-getrainde encoders (die vaak op single-channel data zijn getraind), behaalt het betere resultaten dan de originele MC-ViT architecturen. Dit suggereert dat DC-ViT beter kan omgaan met de overgang van single-channel naar multi-channel data.
• Ablatie Studies:
  • Het verwijderen van DSA of DAG leidt tot een significante daling in prestaties, wat aantoont dat beide componenten noodzakelijk zijn.
  • Efficiëntie: DSA heeft een lineaire schaalbaarheid met het aantal kanalen ($O(CN^2 + NC^2)$), terwijl standaard MC-ViT (MSA) kwadratisch schaalt ($O(C^2N^2)$). Dit maakt DC-ViT aanzienlijk sneller en efficiënter bij datasets met veel kanalen.
  • Lageselectie: De beste resultaten worden behaald wanneer cross-kanaal interactie plaatsvindt in de middelste lagen van de encoder, niet in de allereerste of allerlaatste lagen.

Betekenis en Conclusie

DC-ViT biedt een elegante oplossing voor het fundamentele probleem van heterogene multi-kanaal data in deep learning. Door de interactie tussen kanalen te "ontkoppelen" en te reguleren, behoudt het model de unieke semantiek van individuele kanalen (cruciaal voor wetenschappelijke interpretatie) terwijl het toch profiteert van complementaire context.

De studie onderstreept dat het niet voldoende is om kanalen simpelweg samen te voegen; de structuur van de interactie is even belangrijk als de data zelf. DC-ViT stelt een nieuwe standaard voor het modelleren van MCI-data en biedt een robuustere basis voor downstream taken zoals classificatie en segmentatie in biomedische en aardobservatiedomeinen.