Technische Samenvatting: Prototype-geleide Post-pretraining voor Single-Cell Representatieleren

Probleemstelling

Single-Cell Representatieleren (SCRL) heeft als doel het decoderen van hoogdimensionale genexpressiedata om cellulaire regulerende logica bloot te leggen. Hoewel recente fundamentele modellen, geinspireerd op grote taalmodellen (waarbij genen als tokens en cellen als zinnen worden behandeld), veelbelovend zijn, staan ze voor twee fundamentele beperkingen:

1. Langstaartverdelingen van Celtypen: Single-cell datasets vertonen extreme onbalans in klassen, waarbij zeldzame maar biologisch kritieke celpopulaties (bijvoorbeeld ziekte-initiërende stamcellen) slechts een tiny fractie van de samples uitmaken. Dit leidt tot ruis of ingestorte embeddins voor zeldzame celtypen, omdat modellen vertekend zijn ten gunste van overvloedige populaties.
2. Covariatenverschuivingen: Technische artefacten, zoals batch-effecten en variërende sequentieerdieptes, creëren distributieverschillen tussen trainings- en inferentiegegevens. Bestaande modellen overfitten vaak op batchespecifieke patronen in plaats van universele regulerende logica te vangen, wat hun generalisatievermogen over verschillende platforms of experimentele protocollen beperkt.

Huidige fine-tuningstrategieën slagen er vaak niet in deze problemen volledig op te lossen omdat ze beginnen vanuit een onstabiel latent embeddingsmanifold dat tijdens de initiële pretraining is vastgesteld.

Methodologie: CellRefine

Om deze uitdagingen aan te pakken, introduceren de auteurs CellRefine, een post-pretraining framework dat is ontworpen om te opereren tussen de initiële pretrainingsfase en downstream fine-tuning. CellRefine injecteert biologisch onderbouwde inductieve bias om de latente embeddingsruimte te herschikken naar een stabielere en discriminerendere structuur.

De methode maakt gebruik van een multifaceted doelfunctie ($L_{Total}$) die vier componenten combineert:

1. Prototype-geleerd leren (Marker Gen Programma's):

   • Concept: Gebruikt bekende sets van marker-genen als structurele priors. Omdat volledige marker sets vaak onbekend zijn, gebruikt de methode samengestelde subsets van markers om "prototypes" voor elk celtype te creëren.
   • Implementatie: Marker-genen voor een celtype $c_i$ worden gesorteerd op basis van de cel-ontologiehiërarchie (van specifieke subtypen tot bredere lijnen) met behulp van een functie $H(c_i)$. Deze gesorteerde sequenties worden getokeniseerd en gecodeerd tot prototype-embeddings ($z_m$).
   • Verlies: Een prototype-geleide regularisatieverlies ($L_{prototype}$) maximaliseert de cosinus-similariteit tussen de embedding van een cel ($z_c$) en de bijbehorende prototype-embedding, terwijl de similariteit met andere prototypes wordt geminimaliseerd. Dit injecteert biologische structuur in de representatie.

2. Lijnbewust Regularisatie:

   • Concept: Adresseert het probleem waarbij nauw verwante celtypen (die een ouderlijke lijn delen in de ontologie) onvoldoende gescheiden blijven in de embeddingsruimte.
   • Implementatie: Een lijnbewust scheidingsverlies ($L_{lineage}$) straft excessieve similariteit tussen celparen die dezelfde ouderlijke lijn en ontologieniveau delen, en moedigt het model aan om fijnmazige biologische identiteiten te onderscheiden.

3. Gaussian Mixture Variational Encoding (GMVE):

   • Concept: Modelleert de latente ruimte expliciet als een gegroepeerde structuur om de kwaliteit van de representatie te verbeteren.
   • Implementatie: Een Gaussian Mixture Variational Encoder wordt aan het model toegevoegd. Een KL-divergentieverlies ($L_{GMVE}$) regulariseert de variational posterior tegen een geleerde Gaussian mixture prior, waardoor een gegroepeerde latente structuur wordt afgedwongen.

4. Masked Language Modeling (MLM):

   • Het standaard pretrainingsdoel wordt behouden om ervoor te zorgen dat het model blijft leren over algemene genexpressiepatronen.

Het totale doel is:
$$L_{Total} = L_{MLM} + \lambda_1 L_{prototype} + \lambda_2 L_{lineage} + \lambda_3 L_{GMVE}$$

Belangrijkste Bijdragen

• Nieuwe Post-Pretraining Fase: Stelt een intermediaire adaptatiefase (CellRefine) voor die de latente manifold van fundamentele modellen verfijnt voordat taakspecifieke fine-tuning plaatsvindt, in plaats van uitsluitend te vertrouwen op pretraining gevolgd door directe fine-tuning.
• Integratie van Biologische Priors: Integreert systematisch mechanistische biologische kennis (marker gen programma's en cel-ontologieën) direct in het optimalisatieproces om representatieleren te sturen.
• Omgaan met Langstaart en Covariatenverschuiving: Toont aan dat het grondvesten van het model in biologische priors de vertekening ten gunste van overvloedige celtypen vermindert en de robuustheid tegen technische verschuivingen verbetert.

Experimentele Resultaten

De auteurs hebben CellRefine geëvalueerd over drie kernberekeningsbiologische taken op 11 datasets die diverse weefsels bestrijken (bloed, pancreas, lever, hersenen, long, hart, enz.):

1. Celidentiteit Voorspelling:

   • On-Domain: CellRefine presteerde consistent beter dan state-of-the-art baselines (inclusief volledige fine-tuning, LoRA en standaard MLM-alignement) over 10 datasets. Het behaalde macro-F1 verbeteringen tot 15% ten opzichte van de sterkste baseline (MLM $\to$ FF) op onbalans datasets zoals Bloed en Pancreas.
   • Out-of-Domain (Zero-Shot): Bij cross-modale transfer (trainen op scRNA-seq, testen op ruimtelijke transcriptomics) toonde CellRefine superieure recall@k-prestaties, wat wijst op robuustere en generaliseerbaardere representaties.

2. Ruimtelijke Transcriptomische Imputatie:

   • CellRefine verbeterde de voorspelling van niet-gemeten genexpressies, met hogere Pearson-correlatie en cosinus-similariteit scores in vergelijking met baselines op lever-ruimtelijke datasets.

3. Voorspelling van Gen Perturbatie Respons:

   • Bij het voorspellen van veranderingen in genexpressie onder perturbatie leverde CellRefine consistente, zij het kleinere, verbeteringen op ten opzichte van baselines, en behaalde de beste algehele prestatie onder alle geteste methoden.

Ablatie Studies:

• Het verwijderen van een van de drie regularisatietermen ($L_{prototype}$, $L_{lineage}$, $L_{GMVE}$) resulteerde in prestatiedaling, wat bevestigt dat de multi-objectieve aanpak synergetisch is.
• Streekefficiëntie: In few-shot learning scenario's (5–50 shots) toonde CellRefine aanzienlijk steilere prestatiewinsten bij lage aantallen samples in vergelijking met standaard fine-tuning, wat suggereert dat het beperkte gelabelde data effectiever benut.

Kwalitatieve Analyse:

• UMAP-visualisaties van de latente ruimte onthulden dat CellRefine de embeddingsmanifold herschikt, waardoor zeldzame "staart"-celpopulaties (bijvoorbeeld specifieke T-cel subsets, zeldzame leverceltypen) beter onderscheidbaar worden en minder verstrengeld met grote celclusters in vergelijking met modellen die uitsluitend met MLM zijn getraind.

Betekenis en Claims

Het artikel claimt dat CellRefine een veelbelovende richting voor single-cell representatieleren vertegenwoordigt door aan te tonen dat biologisch onderbouwde inductieve bias de downstream prestaties van fundamentele modellen aanzienlijk kan verbeteren.

• Klinische Nut: Door zeldzame celtypen beter te vangen en de robuustheid tegen covariatenverschuivingen te verbeteren, verhoogt de methode het potentieel voor het identificeren van ziektemechanismen en therapeutische targets in zeldzame of overgangscelstaten die vaak onzichtbaar blijven in standaardmodellen.
• Methodologische Verschuiving: Het werk pleit tegen het uitsluitend vertrouwen op datagedreven pretraining gevolgd door fine-tuning, en pleit in plaats daarvan voor intermediaire fasen die domeinkennis (marker-genen en ontologieën) expliciet integreren om het leerproces te stabiliseren.
• Beperkingen: De auteurs merken bescheiden op dat de aanpak afhankelijk is van samengestelde marker gen programma's, die onvolledig of ruisachtig kunnen zijn voor nieuwe celtypen. Bovendien impliceren de prestatiewinsten niet noodzakelijk een verbeterde parameter-efficiëntie.

Kortom, CellRefine biedt een raamwerk om de kloof te overbruggen tussen grootschalige pretraining en de specifieke, vaak onbalans, realiteiten van single-cell biologische data, wat resulteert in meer discriminerende en generaliseerbare celrepresentaties.