MedSteer: Counterfactual Endoscopic Synthesis via Training-Free Activation Steering

MedSteer is een trainingsvrij framework dat actiesturing toepast op endoscopische diffusiemodellen om perfecte contrafactuele beeldparen te genereren waarbij alleen het pathologische concept verandert terwijl de anatomische structuur volledig behouden blijft, wat leidt tot superieure resultaten in data-augmentatie en polypdetectie.

De kern

Stel je voor dat je een zeer slimme, digitale chef-kok hebt die foto's van binnen in het menselijk lichaam kan maken (endoscopie). Deze chef-kok heet een "diffusiemodel". Hij is fantastisch in het bedenken van nieuwe plaatjes, maar hij heeft een groot probleem: als je hem vraagt om een foto van een gezonde darm te maken en daarna vraagt om een foto van een ziek darm (met een poliep), maakt hij twee hele andere foto's.

Het is alsof je een huis bouwt, het afbreekt en dan een compleet nieuw huis bouwt. De muren, de vloer en de tuin zijn allemaal anders. Voor artsen is dit nutteloos, want ze willen weten: "Wat is er precies anders tussen een gezonde en een zieke darm, terwijl de rest exact hetzelfde blijft?"

Hier komt MedSteer om de hoek kijken. Het is een nieuwe, slimme methode om deze "digitale chef-kok" te sturen zonder hem opnieuw te hoeven leren.

De Magische Knop: "Sturen zonder te Herstarten"

Hoe werkt het? Stel je voor dat je een auto rijdt op een weg.

• De oude manier (Re-prompting): Je wilt van richting veranderen. De oude methode is alsof je de motor uitschakelt, de auto terugrijdt naar de start, en dan opnieuw start in een andere richting. Je komt op een heel andere plek uit.
• De tweede oude manier (Inversie): Je probeert de auto terug te draaien op zijn eigen sporen. Maar omdat je niet perfect kunt terugrijden, schuif je een beetje opzij. De weg wordt een beetje vervormd.
• De MedSteer-methode: Je rijdt gewoon door. Maar je hebt een speciale stuurknop die je precies op het moment dat je een bocht neemt, een klein beetje draait. De auto blijft op dezelfde weg, de banden zijn hetzelfde, het landschap is hetzelfde, maar je komt nu uit bij een andere bestemming (bijvoorbeeld: van "ziek" naar "gezond").

Hoe werkt die stuurknop precies?

1. Het vinden van de "Ziekte-Vector":
   De onderzoekers laten de computer eerst kijken naar duizenden voorbeelden van "gezond" en "ziek". Ze zoeken in de hersenen van de computer (de cross-attention lagen) naar het exacte verschil tussen die twee. Ze maken een soort magisch kompas (de vector) dat alleen naar "ziek" wijst. Dit kompas is heel specifiek: het zegt "dit is de ziekte", maar negeert de rest van de foto (de achtergrond, de textuur).

2. Het Sturen (Steering):
   Wanneer de computer nu een nieuwe foto maakt, gebruikt hij datzelfde kompas.

   • Als hij een foto maakt van een gezonde darm, en je wilt een zieke maken, dan duwt hij de computer een klein beetje in de richting van het "ziek"-kompas.
   • Het mooie is: omdat hij de exacte weg van de computer volgt (dezelfde ruis, dezelfde start), blijft alles wat niet met de ziekte te maken heeft, 100% hetzelfde. De vorm van de darm, de belichting, de achtergrond – niets verandert, behalve de ziekte zelf.

3. De "Verf-Verwijderaar":
   Een heel cool voorbeeld in het artikel is het verwijderen van verf (die artsen gebruiken om poliepen beter te zien).

   • Soms is een poliep geverfd, soms niet. De oude methoden konden de verf niet weg halen zonder de poliep zelf te vervormen.
   • MedSteer ziet de verf als een apart "concept". Het pakt het kompas voor "verf" en draait het tegenovergestelde. Resultaat: De poliep blijft precies zo, maar de verf is weg alsof hij er nooit was.

Waarom is dit zo belangrijk?

In de medische wereld is het lastig om genoeg foto's te vinden van zeldzame ziektes om AI te trainen. MedSteer kan oneindig veel "tegenhangers" (counterfactuals) maken.

• Het maakt een foto van een gezonde darm.
• Het maakt er direct een zieke versie van, waarbij alleen de ziekte verandert.
• Dit helpt artsen en computers om beter te leren zien wat een ziekte echt is, zonder dat ze zich laten afleiden door de achtergrond.

Samengevat in één zin:

MedSteer is als een magische pen die je kunt gebruiken om een ziekte in een foto te "tekenen" of "weg te wissen", terwijl de rest van de foto ongeschonden en perfect hetzelfde blijft, zonder dat je de hele foto opnieuw hoeft te schilderen.

Dit maakt het makkelijker om betere AI-tools te bouwen die artsen helpen sneller en accurater ziektes te herkennen.

Technische samenvatting

Probleemstelling

In de analyse van endoscopische beelden is het een uitdaging om pathologiedetectoren te trainen die reageren op ziektespecifieke kenmerken in plaats van op verwarrende anatomische factoren. Een natuurlijke oplossing is het genereren van paren van zieke en gezonde beelden met identieke anatomie, zodat de detector alleen het doelconcept leert. Echter, bestaande methoden hebben ernstige tekortkomingen:

1. Opnieuw prompten (Re-prompting): Het genereren van een nieuw beeld met een andere tekstprompt (bijv. van "poliep" naar "normaal") herstart de volledige generatietrajectorie. Dit resulteert in een compleet ander beeld met variaties in anatomie, textuur en achtergrond, waardoor geen echte controlegroep ontstaat.
2. Inversie-gebaseerde bewerking: Methoden die beginnen met een bronafbeelding en Denoising Diffusion Implicit Models (DDIM) inversie gebruiken, introduceren reconstructiefouten. Dit leidt tot "structuurdrift", waarbij de niet-doelgebieden (zoals de achtergrond) niet perfect behouden blijven. Bestaande oplossingen vereisen vaak maskeraanduidingen of fine-tuning, wat problematisch is in het medische domein waar gelabelde data schaars is.

Er is dus behoefte aan een methode die counterfactuele paren kan genereren waarbij de enige variatie het doelconcept is, terwijl alle andere structurele elementen exact behouden blijven, zonder dat er bronafbeeldingen, maskers of extra training nodig zijn.

Methodologie: MedSteer

MedSteer is een training-vrij (training-free) framework dat gebruikmaakt van activatiesturing (activation steering) binnen een bevroren Diffusion Transformer (DiT). De methode bestaat uit drie hoofdfasen:

1. Backbone en Interventiepunt:

   • Het model is gebaseerd op PixArt-α (een DiT).
   • De interventie vindt plaats in de Cross-Attention (CA) lagen. Dit is het enige kanaal waar tekstuele semantiek de beeldtokens bereikt. Door hierin te interveniëren, kan een klinisch concept onderdrukt worden zonder de anatomie te beïnvloeden die elders in het netwerk is gecodeerd.

2. Schatting van Pathologie-Vectoren (Offline):

   • Er worden contrastieve promptparen samengesteld (bijv. "een geverfde opgetilde poliep" vs. "een poliep").
   • Voor deze paren worden $Z$ forward-passes uitgevoerd met willekeurige zaden (seeds) door het bevroren model.
   • De CA-features ($h_{l,t}$) worden gemiddeld over alle zaden en ruimtelijke tokens.
   • De pathologie-vector ($v_{l,t}$) wordt berekend als de genormaliseerde verschilvector tussen de positieve en negatieve promptfeatures: $v_{l,t} = \frac{\bar{h}^{pos}_{l,t} - \bar{h}^{neg}_{l,t}}{\|\bar{h}^{pos}_{l,t} - \bar{h}^{neg}_{l,t}\|_2}$.
   • Deze vector is concept-specifiek en wordt één keer offline berekend, waarna hij hergebruikt kan worden zonder extra kosten.

3. Ruimtelijk Selectieve Pathologie-Sturing (SSPS) tijdens Inferentie:

   • Beide beelden (het origineel en het counterfactuele) worden gegenereerd vanuit dezelfde ruisseed, wat garandeert dat ze door hetzelfde traject reizen.
   • In plaats van een globale aftrekking, gebruikt MedSteer een Cosine-Similarity Gate (CSG). Voor elke token wordt de uitlijning met de pathologie-vector berekend via het inproduct: $\sigma_{l,t} = \max(\langle h_{l,t}, v_{l,t} \rangle, 0)$.
   • Alleen tokens die positief uitgelijnd zijn met het pathologie-concept worden aangepast. De activatie wordt bijgewerkt door het uitgelijnde component af te trekken: $h'_{l,t} = h_{l,t} - \alpha \cdot \sigma_{l,t} \cdot v_{l,t}$.
   • Hierdoor worden orthogonale componenten (anatomie, textuur, perspectief) volledig behouden.

Belangrijkste Bijdragen

1. Nieuwe Stuurmechanisme: Een op cosinus-similariteit gebaseerd mechanisme dat klinische concepten stuurt door alleen het concept-uitgelijnde component van cross-attention tokens te verwijderen via een per-token poort. Dit biedt ingebouwde ruimtelijke interpreteerbaarheid.
2. Inversie-vrije Generatie: Het eerste framework dat counterfactuele paren "vanaf nul" genereert zonder DDIM-inversie, bronafbeeldingen of maskers. Door het delen van de ruisseed wordt structuurbehoud gegarandeerd door constructie.
3. Toepassing in het Medische Domein: De methode werkt met verstrengelde medische klasselabels en valideert de nuttigheid op downstream klinische taken (detectie), in plaats van alleen op perceptuele metrics.

Resultaten

De evaluatie vond plaats op de datasets Kvasir v3 en HyperKvasir met drie experimenten:

1. Downstream Poliep-detectie:

   • MedSteer verbeterde de prestaties van poliep-detectoren (ViT en ConvNeXt) aanzienlijk.
   • De ViT AUC steeg naar 0.9755 met MedSteer, vergeleken met 0.9083 voor opnieuw prompten en 0.9518 voor de beste inversie-baselines (PnP). Dit bevestigt dat de behouden structuur cruciaal is voor het leren van robuuste pathologiekens.

2. Counterfactuele Generatie:

   • MedSteer behaalde hoge "flip rates" (het percentage beelden dat succesvol van klasse verandert): 0.800 (Poliep $\to$ Normaal), 0.925 (Colitis $\to$ Normaal) en 0.950 (Esophagitis $\to$ Normaal).
   • Het overtrof baselines (PnP, h-Edit) in zowel concept-verandering als behoud van de achtergrondstructuur (gemeten via Bg-LPIPS, Bg-SSIM, Bg-PSNR). Inversiemethoden leden onder reconstructiefouten.

3. Verfijning van Kleurstof (Dye Disentanglement):

   • In een scenario waar tekst geen onderscheid kan maken tussen een poliep en de kleurstof (indigo carmine), slaagde MedSteer erin de kleurstof te verwijderen terwijl de poliepstructuur behouden bleef.
   • De Dye Detection Rate (DDR) was 0.250 voor MedSteer (d.w.z. 75% verwijdering), significant beter dan h-Edit (0.900) en PnP (0.800).

Betekenis en Conclusie

MedSteer lost het fundamentele probleem op van "structuurdrift" in medische beeldgeneratie. Door te vertrouwen op activatiesturing in plaats van inversie, garandeert het dat alleen het doelconcept verandert, wat essentieel is voor het trainen van betrouwbare diagnostische AI-modellen.

• Efficiëntie: Geen fine-tuning of extra annotaties nodig.
• Interpreteerbaarheid: De per-token scores bieden een visuele kaart van waar en wanneer het model het concept aanpast.
• Toekomst: De auteurs plannen uitbreiding naar 3D-volumetrische data en video-endoscopie.

Samenvattend biedt MedSteer een robuuste, schaalbare oplossing voor het genereren van hoogwaardige, causale trainingsdata voor medische AI, waarbij de integriteit van de anatomische structuur volledig wordt bewaard.