Sparse Autoencoders Reveal Interpretable Features in Single-Cell Foundation Models

Deze studie toont aan dat het trainen van sparse autoencoders op verborgen representaties van single-cell foundation-modellen interpreteerbare biologische en technische kenmerken onthult, waardoor het mogelijk wordt om deze modellen te sturen en ongewenste technische effecten te verminderen.

De kern

De "X-ray" voor de digitale cel: Hoe we de gedachten van AI ontrafelen

Stel je voor dat je een enorme, superintelligente robot hebt die miljoenen foto's van menselijke cellen heeft gezien. Deze robot, een Single-Cell Foundation Model (scFM), kan je vertellen wat voor soort cel je voor je hebt, hoe die reageert op medicijnen, of zelfs voorspellen wat er gebeurt als je een gen uitschakelt. Het is een wonder van de moderne biologie.

Maar er is een groot probleem: niemand weet precies hoe de robot denkt. Het is een "zwarte doos". Je geeft het een cel, en het geeft een antwoord, maar je ziet niet welke interne gedachten of regels het gebruikt om tot dat antwoord te komen. Het is alsof je een auto hebt die perfect rijdt, maar niemand weet hoe de motor werkt of welke knoppen er in het dashboard zitten.

Dit artikel beschrijft hoe de onderzoekers een nieuwe techniek hebben gebruikt om die zwarte doos open te maken. Ze noemen deze techniek Sparse Autoencoders (SAE).

1. De "Lego-blokken" van de robot

Stel je voor dat de robot een cel beschrijft door een enorme muur van Lego-blokken te bouwen. In het verleden dachten we dat de robot één groot, complex blok per cel gebruikte. Maar de onderzoekers ontdekten dat de robot eigenlijk duizenden kleine, specifieke Lego-blokjes gebruikt.

Met hun nieuwe "X-ray" (de SAE) konden ze zien dat deze blokjes twee soorten zijn:

• De "Gen-blokjes": Deze herkennen specifieke eigenschappen van individuele genen. Bijvoorbeeld: "Ah, dit blokje gaat af als een gen veel wordt uitgedrukt" of "Dit blokje herkent alleen mitochondriale genen (de energiecentrales van de cel)."
• De "Cel-blokjes": Deze kijken naar het hele plaatje. Ze zeggen: "Deze combinatie van genen betekent dat we een B-cel van het immuunsysteem hebben," of "Dit is een cel uit een longpatiënt."

De verrassing: De robot leert deze blokjes al tijdens zijn "schooltijd" (pre-training), zelfs voordat hij gespecialiseerde taken krijgt. Hij bouwt al een heel rijk begrip van de biologie op, net zoals een kind al woorden leert voordat het een proefschrift schrijft.

2. De "Stoorzenders" in de radio

Er is echter een addertje onder het gras. De robot is niet alleen slim, hij is ook een beetje een paranoïde luisteraar. Hij hoort niet alleen de biologische signalen (de muziek), maar ook de ruis van de radio (de technische achtergrond).

• De biologische muziek: De echte celtypen en ziektes.
• De technische ruis: Hoe de cel is gemeten. Was het met een bepaalde machine (bijv. "SMARTer")? Van welke persoon kwam het monster? In welk laboratorium?

De onderzoekers ontdekten dat de robot soms meer aandacht besteedt aan de machine die de cel heeft gemeten dan aan de cel zelf. Hij leert bijvoorbeeld: "Als de genen lang zijn en weinig GC-content hebben, dan komt dit waarschijnlijk van de 'SMARTer'-machine." Dit is een stoorzender die de resultaten kan verstoren.

3. Het "Afknijpen" van de gedachten

Dit is waar het echt cool wordt. Omdat ze nu weten welke Lego-blokjes (features) voor welke gedachten staan, kunnen ze de robot besturen.

Stel je voor dat je een radio hebt met een knop voor "Stoorzender". De onderzoekers hebben een knop gevonden die ze kunnen draaien om de "stoorzenders" (de technische ruis van verschillende laboratoria) uit te schakelen.

• Ze hebben de robot dwongen om die specifieke "machine-blokjes" te negeren.
• Het resultaat: De robot zag plotseling dat een cel uit laboratorium A en een cel uit laboratorium B precies hetzelfde waren, omdat de technische ruis was verwijderd. De biologische muziek klonk nu helder en duidelijk.

Ze noemen dit steering (sturen). Het is alsof je een autopilot in een vliegtuig kunt sturen om de windvlagen (batch-effecten) te negeren, zodat het vliegtuig rechtuit blijft vliegen, zonder dat je de hele motor moet vervangen.

4. Waarom is dit belangrijk?

Voorheen waren deze enorme AI-modellen voor cellen onbegrijpelijk en soms onbetrouwbaar omdat ze verward raakten door technische verschillen in de data.

Dit artikel laat zien dat:

1. Deze modellen slimmer zijn dan we dachten: Ze hebben al een diep begrip van biologie, zelfs zonder specifieke training.
2. Ze verwarrend zijn: Ze leren ook onzin (technische ruis) mee.
3. We ze nu kunnen repareren: Door de "gedachten" van de robot te zien en te manipuleren, kunnen we ze schoner en betrouwbaarder maken.

Kortom: De onderzoekers hebben een bril opgezet voor de AI, waardoor we kunnen zien wat hij denkt. En door te weten wat hij denkt, kunnen we hem vertellen: "Hé, ignoreer die ruis van de meetmachine, focus op de echte cel!" Dit maakt de toekomst van medische AI veel transparanter en veiliger.

Technische samenvatting

Probleemstelling

Single-cell foundation models (scFMs), zoals scGPT, scFoundation en Geneformer, hebben veelbelovende toepassingen gevonden in de computationele biologie, waaronder celtype-annotatie, data-integratie en het voorspellen van effecten van celverstoringen (bijv. medicijnrespons). Desondanks functioneren deze modellen grotendeels als "black boxes". Hun interne mechanismen zijn slecht begrepen, wat de interpretatie van hun voorspellingen bemoeilijkt.

Er zijn drie kernproblemen die dit artikel aanpakt:

1. Gebrek aan transparantie: Het is onduidelijk welke biologische of technische signalen de modellen leren en hoe deze worden gecodeerd in de interne representaties.
2. Technische artefacten: Modellen kunnen gevoelig zijn voor technische variatie (bijv. batch-effecten, sequentieertechnologieën) die de biologische signalen kunnen overstemmen.
3. Beperkte controle: Zonder inzicht in de interne features is het moeilijk om de output van het model te sturen om ongewenste effecten te verminderen zonder de biologische integriteit te verliezen.

Methodologie

De auteurs gebruiken Sparse Autoencoders (SAE's) als mechanistische interpretatietechniek om de interne representaties van scFMs te ontleden.

1. Doelmodellen: De studie richt zich op drie wijdverspreide scFMs:

   • scGPT: Een generatief model dat veelvuldig wordt gebruikt.
   • scFoundation: Een pre-trained model.
   • Geneformer: Een model dat vaak wordt gebruikt voor transfer learning.
   • Er werden zowel pre-trained als fijn-gestemd (fine-tuned) varianten geanalyseerd.

2. SAE Training:

   • De SAE's werden getraind op de intermediate token-representaties (residual stream activations) tijdens de forward pass van de transformer-modellen, in plaats van alleen op de uiteindelijke cel-embeddings. Dit vangt rijkere biologische informatie voordat deze wordt gecomprimeerd.
   • Er werd gebruikgemaakt van BatchTopK SAE's, die bleken superieur te zijn in het scheiden van monosemantische (één betekenis dragende) features.
   • De training vond plaats op diverse datasets, waaronder het CellXGene Census (37 miljoen cellen), een COVID-19 cohort en tissue-specifieke datasets (long, pancreas, immuunsysteem).

3. Feature Analyse:

   • Cel-niveau associaties: Activaties werden geaggregeerd naar cel-niveau (via max-pooling) en vergeleken met ground-truth labels (celtype, ziektestatus, batch) met behulp van Adjusted Mutual Information (AMI) en F1-scores.
   • Functionele verrijking: Genen met sterke feature-activaties werden getest op verrijking voor biologische processen (Gene Ontology) en celtype-specifieke markers (PanglaoDB).
   • Sturing (Steering): Om de causaliteit van features te testen, werden specifieke features (bijv. die correleren met batch-effecten) tijdens de inferentie "geclamped" (op een negatieve waarde gezet) om te zien of dit de modeloutput systematisch verandert.

Belangrijkste Bijdragen

1. Ontdekking van interpretabele features: Het bewijzen dat pre-trained scFMs complexe en betekenisvolle biologische concepten leren, zelfs zonder taak-specifiek fine-tuning.
2. Architecturale invloed: Het in kaart brengen van hoe verschillende trainingsprotocollen en architecturen (bijv. binning vs. positionele encoding) de manier waarop informatie wordt gecodeerd, beïnvloeden.
3. Technische vs. Biologische variatie: Het karakteriseren van hoe modellen technische variatie (batch-effecten) naast biologische informatie coderen, waarbij vaak study-specifieke patronen worden gevonden in plaats van universele celtype-features.
4. Functionele sturing: Het demonstreren dat SAE-afgeleide features functioneel gerelateerd zijn aan modelgedrag en kunnen worden gemanipuleerd om batch-effecten te verminderen terwijl het biologische signaal behouden blijft.
5. Open Source: Het vrijgeven van een codebase voor het trainen van SAE's op scFMs.

Resultaten

1. Twee dimensies van features:
De SAE's onthulden dat scFMs informatie organiseren langs twee assen:

• Gen-specifieke features: Deze coderen eigenschappen van individuele genen (expressieniveau, genfamilie, moleculaire functie) onafhankelijk van de celcontext. Bijvoorbeeld, features die specifiek reageren op ribosomale genen of mitochondriale genen.
• Cel-specifieke features: Deze vangen eigenschappen van de hele cel op, gedistribueerd over meerdere gen-token. Ze ontstaan door contextuele informatie die via self-attention wordt verspreid. Deze features coderen celtypen, ziektestatus en technische batches.

2. Complexe coderingsstrategieën:
Modellen gebruiken soms onverwachte strategieën om celidentiteit te coderen:

• Negatieve codering: Een feature activeert voor T-cellen maar niet voor B-cellen, wat fungeert als een "afwezigheid van niet-B-cel signalen".
• Proxy-codering: Een feature activeert op specifieke ribosomale genen die B-cellen onderscheiden van andere cellen, zonder dat deze genen direct biologisch gekoppeld zijn aan B-celfunctie.

3. Technische bias en generalisatie:

• Pre-trained modellen coderen sterke technische signalen (zoals sequentieertechnologieën zoals "SMARTer" vs. "inDrop").
• Celtype-features tonen vaak variabele generalisatie: een feature die een celtype in één studie herkent, activeert soms zwak of niet in dezelfde celtype in een andere studie. Dit suggereert dat modellen study-specifieke representaties leren in plaats van geconsolideerde biologische concepten.

4. Sturing (Steering) voor Batch Correctie:

• Door de top 20-50 batch-gerelateerde features te onderdrukken (clampen op -2) tijdens de inferentie, kon de batch-integratie significant worden verbeterd.
• Op het pancreas-dataset overtrof de gestuurde fine-tuned scGPT zowel de native DAR-batchcorrectie van scGPT als de ongestuurde versie, met behoud van biologische conservatie.
• Dit bewijst dat deze features causaal bijdragen aan de technische variatie in de modeloutput.

Betekenis en Conclusie

De studie toont aan dat scFMs tijdens pre-training een genuanceerd, maar gefragmenteerd begrip van celbiologie ontwikkelen. Hoewel ze rijke biologische features leren, splitsen ze deze vaak op in study-specifieke contexten en coderen ze sterke technische artefacten.

De belangrijkste implicaties zijn:

• Interpretbaarheid: SAE's bieden een krachtig middel om de "black box" van scFMs te openen en te begrijpen wat deze modellen werkelijk leren.
• Controle: Het is mogelijk om specifieke features te manipuleren om ongewenste bias (zoals batch-effecten) te verwijderen zonder het model opnieuw te hoeven trainen. Dit opent de deur voor gerichte interventies in foundation models.
• Toekomstige richting: De resultaten benadrukken dat architecturale keuzes en trainingsstrategieën (zoals hoe expressie wordt gebinned of geschaald) een grote invloed hebben op de kwaliteit en interpretatie van de geleerde representaties. Verdere interpretabiliteitsstudies zijn essentieel om betrouwbare en controleerbare scFMs te ontwikkelen.

Kortom, dit werk levert een fundamentele stap voorwaarts in het begrijpen en beheersen van single-cell foundation models door hun interne mechanismen zichtbaar en manipuleerbaar te maken.