SHOT-CCR: Biologically guided adversarial training for test-time adaptation in cellular morphology

Het paper introduceert SHOT-CCR, een biologisch onderlegd testtijd-adaptatiestelsel dat door middel van gradiëntomkering batch-effecten elimineert en zo de classificatie van genetische verstoringen in celmorfologie aanzienlijk verbetert op verschillende datasets.

De kern

De "Shot-CCR" methode: Hoe we AI helpen om cellen te herkennen, ondanks "ruis" in het lab

Stel je voor dat je een detective bent die miljoenen foto's van cellen moet analyseren om te zien welke medicijnen werken. Elke foto toont een cel die is blootgesteld aan een specifiek genetisch experiment (een "perturbatie"). Je AI-model is getraind om deze foto's te bekijken en te zeggen: "Ah, dit is een cel met een gebroken gen A!"

Het probleem? De "batch-effecten".

Het Probleem: De Ruis van het Lab

In de echte wereld worden deze foto's niet allemaal op dezelfde dag, met dezelfde microscopen of door dezelfde mensen gemaakt. Soms is het licht net iets anders, soms zijn de cellen iets voller of leger, en soms is de chemische oplossing net iets anders.

Dit is alsof je een detective bent die getraind is om een verdachte te herkennen aan zijn gezicht. Maar als de verdachte een dag later wordt gefotografeerd, draagt hij een andere hoed, staat hij in een andere kamer met ander licht, en heeft hij misschien een andere houding. Als je AI te sterk leert op de "hoed" of de "kamer", vergeet hij het gezicht. In de wetenschap noemen we dit batch-effecten: technische ruis die de echte biologische signalen (het gezicht) verbergt.

In het verleden faalden AI-modellen vaak als ze naar nieuwe foto's keken die in een ander "batch" (een andere experimentele ronde) waren gemaakt, zelfs als ze de oude foto's perfect herkenden.

De Oplossing: SHOT-CCR

De auteurs van dit papier hebben een slimme nieuwe methode bedacht, genaamd SHOT-CCR. Laten we dit uitleggen met een paar analogieën.

1. De "Cellen-teller" als Slechte Gewoonte

Een van de grootste technische verschillen tussen batches is simpelweg hoeveel cellen er op de foto zitten.

• In Batch A zitten misschien 100 cellen per foto.
• In Batch B zitten er misschien 200.

De AI is slim, maar ook lui. Hij merkt op: "Oh, als er 200 cellen zijn, is het waarschijnlijk Batch B, en dus is dit waarschijnlijk een bepaald type gen." De AI begint dus te gokken op het aantal cellen in plaats van op het uiterlijk van de cellen. Dit is alsof de detective denkt: "Als de verdachte in een grote kamer staat, is hij schuldig," in plaats van naar zijn gezicht te kijken.

De CCR-methode (Cell Count Reversal) is hier de oplossing.
Stel je voor dat je de AI een "anti-leer" geeft. We zeggen tegen de AI: "Probeer te voorspellen hoeveel cellen er op de foto zitten, maar straf je eigen succes als je het goed hebt!"

• We laten de AI proberen het aantal cellen te raden.
• Maar zodra hij het goed raadt, keren we de straf om. In plaats van beloning, krijgt hij een boete.
• Hierdoor leert de AI: "Ik mag het aantal cellen niet gebruiken om mijn beslissing te maken." Hij wordt gedwongen om echt naar de biologische details (het gezicht van de verdachte) te kijken en de "ruis" (het aantal cellen) te negeren.

2. SHOT: De "Zelflerende" Detective

Naast het negeren van het aantal cellen, gebruiken ze een techniek genaamd SHOT (Test-Time Adaptation).
Stel je voor dat je detective (het AI-model) al getraind is in een oude zaak. Nu komt er een nieuwe zaak met nieuwe foto's. In plaats van de detective opnieuw jarenlang te laten trainen (wat tijd en data kost), laten we hem ter plekke even snel oefenen op de nieuwe foto's.

• De AI kijkt naar de nieuwe foto's.
• Hij probeert zijn eigen voorspellingen te verbeteren door te zeggen: "Ik ben het zekerste over deze groep, dus ik ga mijn instellingen een beetje aanpassen om nog zekerder te zijn."
• Dit gebeurt zonder dat iemand de antwoorden (labels) hoeft te geven. De AI leert zichzelf aan de hand van de nieuwe situatie.

Wat is het Resultaat?

De auteurs hebben deze methode getest op twee enorme datasets met miljoenen celafbeeldingen (RxRx1 en JUMP-CP).

• De Prestatie: Hun nieuwe methode (SHOT-CCR) scoort 91,6% nauwkeurigheid. De vorige beste methode (de "standaard") haalde maar 87,1%. Dat klinkt als een klein verschil, maar in de wereld van AI en medicijnontwikkeling is dat een enorme sprong voorwaarts.
• De Sterkste Winst: De methode werkt het beste bij de cellen die het moeilijkst waren (de U2OS cellen). Hier was de "ruis" het grootst, en juist daar hielp het negeren van het cel-aantal het meest.
• Biologische Betekenis: Ze keken ook na of de AI nu echt "slimmer" was geworden over de biologie. Ja! De AI kon nu beter specifieke biologische processen herkennen die eerder verborgen zaten onder de technische ruis.

Waarom is dit belangrijk?

Medicijnontwikkeling is duur en langzaam. Als AI-modellen beter kunnen omgaan met verschillende experimenten zonder opnieuw getraind te hoeven worden, kunnen wetenschappers sneller nieuwe medicijnen vinden.

Samengevat in één zin:
De auteurs hebben een slimme truc bedacht om AI te leren "niet te letten op hoeveel er op de foto zitten" (het aantal cellen), zodat de AI zich kan focussen op wat er echt belangrijk is: hoe de cellen eruitzien en welke medicijnen werken, ongeacht in welk lab de foto is gemaakt.

Technische samenvatting

Probleemstelling

In het veld van de drugontdekking worden steeds vaker grote datasets gebruikt van "Cell Painting" (hoge-inhoud screening), waarbij cellen worden gefotografeerd na genetische of chemische verstoringen. Een groot obstakel bij het gebruik van deze data voor AI-modellen zijn batch-effecten. Technische verschillen tussen experimentele batches (bijv. verschillende tijdstippen, laboratoria of apparatuur) leiden tot artefacten die biologische signalen verstoren. Hierdoor falen modellen vaak om te generaliseren naar nieuwe, ongezette batches, zelfs als ze binnen de trainingsbatches goed presteren.

Specifiek blijkt dat modellen vaak onterecht vertrouwen op celtelling (het aantal cellen in een beeld) als een proxy voor de experimentele batch, in plaats van op de daadwerkelijke morfologische veranderingen veroorzaakt door de genetische verstoring. Dit is vooral problematisch bij datasets met verschillende celtypen, waar de celdichtheid per batch sterk kan variëren.

Methodologie: SHOT-CCR

De auteurs stellen een nieuw raamwerk voor genaamd SHOT-CCR (SHOT Cell Count Reversal), dat Test-Time Adaptation (TTA) combineert met biologisch gerichte adversarial training.

1. Architectuur:

   • Het model gebruikt een voorgeprogrammeerde backbone (voornamelijk DenseNet-161) die is aangepast voor Cell Painting-afbeeldingen (5 of 6 kanalen).
   • Tijdens training wordt een extra regressiehoofd toegevoegd dat het aantal cellen in het beeld voorspelt.

2. Biologisch Gesteunde Adversarial Training (CCR):

   • In plaats van te proberen alle batch-effecten te verwijderen (wat vaak leidt tot het verliezen van nuttige biologische signalen), richt de methode zich specifiek op celtelling als een bekende confounder.
   • Er wordt een Gradient Reversal Layer (GRL) gebruikt op de cel-telling regressie. Dit dwingt de feature extractor om representaties te leren die onafhankelijk zijn van het aantal cellen, maar wel behoudend voor de genetische verstoring.
   • Belangrijk: De auteurs gebruiken aparte hyperparameters (leer snelheid en $\alpha$) voor deze laag om te voorkomen dat het model volledig blind wordt voor celtelling (wat schadelijk kan zijn) of juist te veel afhankelijk blijft. Het doel is een "gedeeltelijke invariantie" (partial invariance).

3. Test-Time Adaptation (TTA):

   • Tijdens de inferentie (testtijd) wordt het model aangepast aan de nieuwe batch zonder toegang tot labels.
   • De methode gebruikt SHOT (Source Hypothesis Transfer), een techniek die entropy-minimalisatie, diversiteitsverlies en pseudo-labeling combineert.
   • De classifier (hypothesis) wordt bevroren, terwijl de feature extractor wordt geoptimaliseerd om de onzekerheid van de voorspellingen te minimaliseren op de nieuwe data.

Kernbijdragen

1. Biologisch onderbouwde TTA: Uitbreiding van computer-vision TTA-technieken naar Cell Painting-data, specifiek gericht op het oplossen van batch-effecten door biologische priors (celtelling) te gebruiken.
2. Cell Count Adversarial Training: Introductie van een nieuw mechanisme dat het netwerk dissuadeert om te veel te vertrouwen op celtelling, wat superieur is aan generieke batch-effect reversie.
3. Uitgebreide Validatie: Demonstration van consistente verbeteringen over twee grote datasets (RxRx1 en JUMP-CP) en vier verschillende celtypen, wat een nieuw benchmark stelt voor morfologische batch-correctie.

Resultaten

De prestaties zijn gemeten aan de hand van de classificatie-accuraatheid van genetische verstoringen (siRNA en CRISPR).

• RxRx1 Dataset (siRNA):

  • De huidige state-of-the-art (AdaBN) haalde 87,1% nauwkeurigheid.
  • SHOT-CCR bereikte 91,6% nauwkeurigheid, een verbetering van 4,5%.
  • De verbetering was het grootst bij het U2OS-celtype (van 68,2% naar 76,2%), dat eerder als het moeilijkst te voorspellen werd beschouwd vanwege de beperkte trainingsdata en grote variatie in celtelling.

• JUMP-CP Dataset (CRISPR):

  • Voor deze dataset (484 CRISPR-perturbaties) verbeterde de methode de nauwkeurigheid met 15,7% ten opzichte van de baseline.
  • De voordelen van CCR waren hier kleiner omdat de celtelling-distributie over de batches in deze dataset homogener was, wat bevestigt dat de methode het meest effectief is waar er grote variatie in celtelling is.

• Ablatie Studies:

  • Generieke batch-identiteit reversie (zonder biologische focus) bleek schadelijk voor de prestaties.
  • De combinatie van SHOT en CCR was superieur aan het gebruik van alleen SHOT of alleen CCR.
  • Het verwijderen van batches met extreme afwijkingen in celtelling uit de trainingsdata verbeterde de prestaties van basismodellen, wat aantoont dat celtelling een kritieke factor is bij het opzetten van train/test splits.

Significantie en Conclusie

Het paper toont aan dat het selectief verwijderen van specifieke, biologisch gemotiveerde confounders (zoals celtelling) via adversarial training, in combinatie met test-time adaptatie, een krachtige strategie is om batch-effecten in celbiologie te overwinnen.

• Biologische Validatie: Een gen-enrichment analyse toonde aan dat de verbeterde classificatie vooral plaatsvond bij genen die betrokken zijn bij subtielere morfologische veranderingen (zoals RNA-helicases en endosomen), die eerder vaak werden overschaduwd door batch-ruis.
• Praktische Impact: De methode maakt het mogelijk om modellen te gebruiken die zijn getraind op oude datasets, om nieuwe experimenten (met andere batches of celtypen) te analyseren zonder dat er nieuwe labels nodig zijn.
• Toekomst: De auteurs raden aan om deze aanpak uit te breiden naar andere soorten batch-effecten (zoals rij- en kolom-effecten in platen) en om te testen op transformer-architecturen.

Kortom, SHOT-CCR biedt een robuuste oplossing voor een van de grootste uitdagingen in high-content screening: het generaliseren van AI-modellen over verschillende experimentele condities door gebruik te maken van biologische kennis in het trainingsproces.