HEIMDALL: Disentangling tokenizer design for robust transfer in single-cell foundation models

Deze paper introduceert HEIMDALL, een unificerend kader dat aantoont dat het zorgvuldig ontwerpen van tokenizers, met name op het gebied van genidentiteit, expressiecodering en volgorde, cruciaal is voor de robuuste generalisatie van foundation modellen voor single-cell RNA-sequencing onder distributieveranderingen.

De kern

Stel je voor dat je een superintelligente robot wilt bouwen die alle cellen in het menselijk lichaam kan begrijpen. Deze robot moet kunnen voorspellen waarom iemand ziek wordt, welke medicijnen werken, of hoe een cel verandert. Om dit te doen, moet je de robot eerst "leren" spreken.

Maar hier zit het probleem: een cel is geen zin in een boek. Een cel is een enorme lijst van duizenden genen, elk met een bepaald activiteitsniveau. Hoe vertaal je die chaotische lijst naar iets wat een computer (een AI-model) begrijpt?

In de wetenschap noemen we dit tokenisatie: het omzetten van data in "woorden" (tokens) die een AI kan lezen.

Dit nieuwe onderzoek, genaamd HEIMDALL, zegt: "Wacht even, we zijn al te lang gefocust op hoe slim de robot is (het brein), maar we vergeten hoe goed we de taal hebben vertaald (de vertaler)."

Hier is een simpele uitleg van wat ze hebben ontdekt, met een paar creatieve vergelijkingen:

1. Het probleem: De "Vertaler" is verward

Stel je voor dat je een recept wilt geven aan een kok.

• Model A zegt: "Neem 2 koppen bloem, 1 ei, en meng ze in willekeurige volgorde."
• Model B zegt: "Neem 2 koppen bloem, 1 ei, en meng ze eerst de bloem, dan het ei."
• Model C zegt: "Schrijf de ingrediënten op in het alfabet."

Tot nu toe hebben wetenschappers gekeken naar welke kok het beste baktaartjes maakt. Maar ze hebben nooit gekeken of het misschien wel de volgorde of de manier van schrijven was die het verschil maakte. Sommige modellen werken goed als je ze in dezelfde keuken gebruikt, maar als je ze naar een andere keuken (een ander weefsel of een andere soort) stuurt, gaan ze volledig in de war.

2. De oplossing: HEIMDALL (De Ontleedmachine)

De onderzoekers hebben HEIMDALL bedacht. Dit is geen nieuwe AI, maar een modulair raamwerk. Het is alsof ze de vertaler uit elkaar hebben gehaald en in drie losse onderdelen hebben verdeeld, zodat ze elk onderdeel apart kunnen testen:

1. De Naamplaatjes (Gene Identity): Hoe noemen we het gen? Is het een willekeurige code, of gebruiken we de biologische naam?
2. De Hoeveelheid (Expression Encoding): Hoe geven we aan hoeveel er van het gen is? Is het een getal, een categorie, of een kleur?
3. De Volgorde (Ordering): In welke volgorde zetten we de genen op de lijst? Alphabetisch? Van meest naar minst actief? Of willekeurig?

3. De Grote Ontdekkingen

Door deze onderdelen te mixen en te matchen (zoals Lego-blokken), ontdekten ze drie belangrijke dingen:

• In de "thuiskeuken" maakt het niet uit: Als je de AI test op dezelfde data waar hij mee is getraind, maakt de keuze van de vertaler weinig uit. Alle modellen doen het ongeveer even goed.

• In de "vreemde keuken" is het cruciaal: Zodra je de AI naar een nieuw weefsel (bijv. van darm naar hersenen), een andere soort (van mens naar muis) of een andere lijst van genen stuurt, breekt het systeem als de vertaler niet goed is.

  • Vergelijking: Het is alsof je een kok die alleen Italiaans spreekt naar Frankrijk stuurt. Als je hem niet vertelt hoe de Franse ingrediënten heten of in welke volgorde je ze moet gebruiken, maakt hij een ramp.

• Er is geen "perfecte" vertaler: Er is niet één manier om alles te vertalen die voor elke situatie werkt.

  • Voor het overstappen van mens naar muis werkt een vertaler die kijkt naar de DNA-sequentie (de bouwtekening) het beste.
  • Voor het voorspellen van ziektes in nieuwe weefsels werkt een vertaler die de genen sorteert op activiteit (van hoog naar laag) het beste.

4. Waarom is dit belangrijk?

Voorheen dachten onderzoekers: "Als we maar een groter en slimmer brein bouwen, werkt het wel."
Dit papier zegt: "Nee, het gaat om de taal."

Als je een AI wilt bouwen die echt robuust is en kan helpen bij het vinden van nieuwe medicijnen voor verschillende ziektes of soorten, moet je eerst de vertaler (de tokenisatie) op orde hebben. Je moet kiezen welke "biologische prikkels" (zoals de volgorde van genen of hun activiteit) je aan de AI geeft.

Samenvattend in één zin:

HEIMDALL laat zien dat voor slimme AI in de biologie, het niet gaat om hoe groot het brein is, maar om hoe goed je de taal van de cellen hebt vertaald naar een vorm die het brein begrijpt, vooral als je die AI naar een nieuwe omgeving stuurt.

Het is alsof je een wereldreiziger bent: je kunt de beste gids hebben (het brein), maar als je de kaart verkeerd leest (de tokenisatie), beland je in de verkeerde stad.

Technische samenvatting

Probleemstelling

Single-cell RNA-sequencing (scRNA-seq) foundation modellen (scFMs) beloven krachtige, universele tools te zijn voor biologisch onderzoek, zoals celtype-annotatie en voorspelling van perturbaties. Echter, hun prestaties zijn vaak inconsistent, vooral bij toepassing op nieuwe biologische contexten (bijv. nieuwe weefsels, soorten of genpanelen).

Een onderbelichte oorzaak van deze beperkingen is tokenisatie: het proces waarbij het continue, ongeordende setje van genexpressiewaarden van een cel wordt omgezet in een sequentie van "tokens" die door een model (zoals een Transformer) kan worden verwerkt.

• In tegenstelling tot tekst of afbeeldingen, ontbreekt er een canonieke tokenisatiestandaard voor single-cell data.
• Bestaande scFMs gebruiken verschillende, vaak heuristische tokenisatieschema's die verschillende biologische aannames verbergen.
• Bestaande benchmarks vergelijken volledige modellen, waardoor het onmogelijk is om te bepalen of prestatieverschillen het gevolg zijn van de modelarchitectuur, de trainingsdata of juist de tokenisatie.
• Er is geen principieel kader om te begrijpen welke tokenisatiecomponenten essentieel zijn voor robuuste generalisatie onder distributieveranderingen (distribution shift).

Methodologie: Het HEIMDALL Framework

De auteurs introduceren HEIMDALL, een modulair framework om tokenisatiestrategieën systematisch te ontleden, te evalueren en opnieuw te ontwerpen.

1. Modulaire Decompositie
HEIMDALL breekt elke tokenisator op in drie functionele modules, die verder worden onderverdeeld:

• $F_G$ (Gene Identity Encoding): Codeert de identiteit van een gen. Opties zijn: willekeurige initialisatie, vooraf getrainde embeddings (ESM2 op aminozuursequenties, Gene2vec op co-expressie, GenePT op tekstbeschrijvingen, HyenaDNA op DNA-sequenties).
• $F_E$ (Expression Encoding): Codeert de expressiewaarde van een gen. Opties zijn: No-op (constant), continue mapping (MLP), kwantiel-binning, integer-binning, of autobinning.
• $F_C$ (Cell Construction): Integreert $F_G$ en $F_E$ tot een celrepresentatie. Dit wordt onderverdeeld in:
  • ORDER: Bepaalt de volgorde van de tokens (bijv. gesorteerd op expressie, chromosomaal, of willekeurig).
  • SEQUENCE: Selecteert welke genen worden opgenomen en construeert de sequentie (bijv. truncatie of gewenste steekproef).
  • REDUCE: Combineert de embeddings (bijv. optelling of identiteit).

2. Experimenteel Ontwerp

• Re-implementatie: De auteurs hebben de tokenizers van vijf toonaangevende scFMs (scGPT, Geneformer, scFoundation, scBERT, UCE) herbouwd binnen dit gemeenschappelijke HEIMDALL-framework.
• Gecontroleerde Vergelijking: Alle modellen worden getraind van scratch (zonder pre-training) met een minimale Transformer-achterkant en identieke hyperparameters. Hierdoor worden verschillen in prestaties puur toegeschreven aan de tokenisatie.
• Ablatiestudies: Door modules uit te wisselen (bijv. de $F_G$ van model A gebruiken met de $F_E$ van model B), kunnen de auteurs de bijdrage van elk component isoleren.

3. Benchmarktaken
Het framework wordt getest op vier uitdagende transfer learning-taken:

1. Cross-tissue generalisatie: Trainen op darmweefsel, testen op hersenen.
2. Cross-species generalisatie: Trainen op menselijke cellen, testen op muiscellen (zonder extra finetuning).
3. Gen-panel generalisatie: Trainen op een groot genpanel, testen op een klein, overlappend panel (zoals bij ruimtelijke transcriptomics).
4. Reverse perturbation prediction: Voorspellen welke gen-knockout leidt tot een specifieke celtoestand.

Belangrijkste Resultaten

1. Tokenisatie is cruciaal bij distributieverandering

• In situaties waar train- en testdata uit dezelfde verdeling komen (in-distribution), heeft de keuze van de tokenisator weinig invloed op de prestaties.
• Onder distribution shift (nieuwe weefsels, soorten, of genpanels) wordt de tokenisatie echter de bepalende factor. Een slecht gekozen tokenisatie kan leiden tot falen, terwijl een goede tokenisatie robuuste transfer mogelijk maakt.

2. Drie kritieke ontwerpaspecten
De studie identificeert dat robuuste transfer afhangt van drie specifieke asen:

• Genidentiteit ($F_G$): Hoe genen worden vertegenwoordigd.
• Expressie-codering ($F_E$): Hoe expressiewaarden worden verwerkt.
• Ordening ($ORDER$): De volgorde van de tokens in de input.

3. Specifieke bevindingen per taak

• Cross-tissue: De prestaties worden voornamelijk gedreven door de ORDER-module. Geneformer-tok presteert het beste omdat het genen sorteert op basis van expressie, wat impliciete biologische informatie toevoegt, zelfs zonder expliciete expressie-codering ($F_E$). Pre-training had hier weinig extra effect.
• Cross-species:
  • Zonder aanpassing faalt de meeste modellen omdat ze een soort-specifiek vocabulaire hebben.
  • UCE-tok presteert het beste zonder orthologie-mapping omdat het ESM2 (proteïne-sequenties) gebruikt voor genidentiteit, wat soort-onafhankelijk is.
  • Wanneer orthologie-mapping wordt toegepast, presteren modellen met sterke $F_E$ en $ORDER$ componenten (zoals scBERT-tok) beter dan UCE.
  • Conclusie: Voor niet-modelsoorten zonder betrouwbare orthologie-mapping is een sequentie-gebaseerde $F_G$ (zoals ESM2) essentieel.
• Gen-panel generalisatie (Spatial Transcriptomics):
  • scBERT-tok presteert het beste. De ablatie-studies tonen aan dat dit komt door de Gene2vec-embeddings ($F_G$), die co-expressiepatronen beter vastleggen dan sequentie-embeddings.
  • Ook een goede expressie-codering ($F_E$) en ordening zijn belangrijk.
• Reverse Perturbation:
  • De toevoeging van expliciete expressie-informatie via de $F_E$ en $ORDER$ modules is cruciaal. UCE-tok, dat standaard geen expressie-codering gebruikt, presteert slecht. Zodra er een continue of gebinned expressie-codering wordt toegevoegd, verbetert de prestatie drastisch.
  • De beste resultaten worden behaald door hybride tokenizers te maken die de beste onderdelen van verschillende modellen combineren (bijv. $F_E$ van scBERT en $ORDER$ van Geneformer).

Bijdragen en Significance

Technische Bijdragen:

1. HEIMDALL Framework: Een open-source, modulair platform dat tokenisatie ontzint in herbruikbare componenten, wat eerlijke vergelijkingen en systematische ablaties mogelijk maakt.
2. Ontmaskering van Tokenisatie: Het paper toont aan dat tokenisatie geen statische, vooraf bepaalde keuze is, maar een kritieke ontwerpaspect dat de inductieve bias van het model bepaalt.
3. Hybride Tokenizers: Het bewijs dat er geen "universeel beste" tokenisator bestaat, maar dat de beste strategie afhankelijk is van de specifieke distributieverandering. Het combineren van componenten (hybride tokenizers) kan superieur zijn aan bestaande, monolithische modellen.

Wetenschappelijke en Praktische Impact:

• Richtinggevend voor Ontwikkeling: Het biedt ontwikkelaars van scFMs een blauwdruk voor het ontwerpen van robuustere modellen door expliciet na te denken over genidentiteit, expressie-codering en ordening.
• Gebruikersgids: Het helpt gebruikers bij het kiezen van het juiste model voor hun specifieke toepassing (bijv. "Gebruik een model met ESM2-embeddings als je werkt met niet-modelsoorten").
• Toekomstperspectief: Het legt de basis voor "virtuele cellen" waarbij tokenisatie zal moeten evolueren om multimodale data (genomisch, epigenomisch, proteomisch) coherent te integreren.

Samenvattend stelt HEIMDALL dat de robuustheid van single-cell foundation modellen niet alleen afhangt van de schaal of architectuur van het model, maar fundamenteel wordt bepaald door hoe biologische informatie via de tokenisatie-interface aan het model wordt gepresenteerd.