Technical Summary: Prototype Guided Post-pretraining for Single-Cell Representation Learning

Problem Statement

Single-cell representation learning (SCRL) aims to decode high-dimensional gene expression data to uncover cellular regulatory logic. While recent foundation models inspired by large language models (treating genes as tokens and cells as sentences) have shown promise, they face two fundamental limitations:

1. Long-Tailed Cell-Type Distributions: Single-cell datasets exhibit extreme class imbalance, where rare but biologically critical cell populations (e.g., disease-initiating stem cells) constitute a tiny fraction of samples. This leads to noisy or collapsed embeddings for rare cell types, as models are biased toward abundant populations.
2. Covariate Shifts: Technical artifacts, such as batch effects and varying sequencing depths, create distribution gaps between training and inference data. Existing models often overfit to batch-specific patterns rather than capturing universal regulatory logic, limiting their generalization across different platforms or experimental protocols.

Current fine-tuning strategies often fail to fully resolve these issues because they begin from an unstable latent embedding manifold established during initial pretraining.

Methodology: CellRefine

To address these challenges, the authors introduce CellRefine, a post-pretraining framework designed to operate between the initial pretraining stage and downstream fine-tuning. CellRefine injects biologically grounded inductive biases to reshape the latent embedding space into a more stable and discriminative structure.

The method employs a multi-faceted objective function ($L_{Total}$) combining four components:

1. Prototype-Guided Learning (Marker Gene Programs):

   • Concept: Utilizes known marker gene sets as structural priors. Since full marker sets are often unknown, the method uses curated subsets of markers to create "prototypes" for each cell type.
   • Implementation: Marker genes for a cell type $c_i$ are ordered based on the cell ontology hierarchy (from specific subtypes to broader lineages) using a function $H(c_i)$. These ordered sequences are tokenized and encoded into prototype embeddings ($z_m$).
   • Loss: A prototype-guided regularization loss ($L_{prototype}$) maximizes the cosine similarity between a cell's embedding ($z_c$) and its corresponding prototype embedding, while minimizing similarity to other prototypes. This injects biological structure into the representation.

2. Lineage-Aware Regularization:

   • Concept: Addresses the issue where closely related cell types (sharing a parent lineage in the ontology) remain insufficiently separated in the embedding space.
   • Implementation: A lineage-aware separation loss ($L_{lineage}$) penalizes excessive similarity between cell pairs that share the same parent lineage and ontology level, encouraging the model to distinguish fine-grained biological identities.

3. Gaussian Mixture Variational Encoding (GMVE):

   • Concept: Explicitly models the latent space as a clustered structure to improve representation quality.
   • Implementation: A Gaussian Mixture Variational Encoder is appended to the model. A KL divergence loss ($L_{GMVE}$) regularizes the variational posterior against a learned Gaussian mixture prior, enforcing a clustered latent structure.

4. Masked Language Modeling (MLM):

   • The standard pretraining objective is retained to ensure the model continues to learn general gene expression patterns.

The total objective is:
$$L_{Total} = L_{MLM} + \lambda_1 L_{prototype} + \lambda_2 L_{lineage} + \lambda_3 L_{GMVE}$$

Key Contributions

• Novel Post-Pretraining Stage: Proposes an intermediate adaptation stage (CellRefine) that refines the latent manifold of foundation models before task-specific fine-tuning, rather than relying solely on pretraining followed by direct fine-tuning.
• Biological Priors Integration: Systematically integrates mechanistic biological knowledge (marker gene programs and cell ontologies) directly into the optimization process to guide representation learning.
• Handling Long-Tail and Covariate Shift: Demonstrates that grounding the model in biological priors mitigates the bias toward abundant cell types and improves robustness to technical shifts.

Experimental Results

The authors evaluated CellRefine across three core computational biology tasks on 11 datasets spanning various tissues (blood, pancreas, liver, brain, lung, heart, etc.):

1. Cell Identity Prediction:

   • On-Domain: CellRefine consistently outperformed state-of-the-art baselines (including full fine-tuning, LoRA, and standard MLM alignment) across 10 datasets. It achieved macro-F1 improvements of up to 15% over the strongest baseline (MLM $\to$ FF) on imbalanced datasets like Blood and Pancreas.
   • Out-of-Domain (Zero-Shot): In cross-modal transfer (training on scRNA-seq, testing on spatial transcriptomics), CellRefine showed superior recall@k performance, indicating more robust and generalizable representations.

2. Spatial Transcriptomic Imputation:

   • CellRefine improved the prediction of unmeasured gene expressions, achieving higher Pearson correlation and cosine similarity scores compared to baselines on liver spatial datasets.

3. Gene Perturbation Response Prediction:

   • In predicting gene expression changes under perturbation, CellRefine yielded consistent, albeit smaller, improvements over baselines, achieving the best overall performance among all tested methods.

Ablation Studies:

• Removing any of the three regularization terms ($L_{prototype}$, $L_{lineage}$, $L_{GMVE}$) resulted in performance degradation, confirming that the multi-objective approach is synergistic.
• Sample Efficiency: In few-shot learning scenarios (5–50 shots), CellRefine demonstrated significantly steeper performance gains at low sample counts compared to standard fine-tuning, suggesting it leverages limited labeled data more effectively.

Qualitative Analysis:

• UMAP visualizations of the latent space revealed that CellRefine reorganizes the embedding manifold, making rare "tail" cell populations (e.g., specific T-cell subsets, rare liver cell types) more distinguishable and less entangled with major cell clusters compared to models trained only with MLM.

Significance and Claims

The paper claims that CellRefine represents a promising direction for single-cell representation learning by demonstrating that biologically informed inductive biases can significantly improve the downstream performance of foundation models.

• Clinical Utility: By better capturing rare cell types and improving robustness to covariate shifts, the method enhances the potential for identifying disease mechanisms and therapeutic targets in rare or transitional cell states that are often obscured in standard models.
• Methodological Shift: The work argues against relying solely on data-driven pretraining followed by fine-tuning, advocating instead for intermediate stages that explicitly incorporate domain knowledge (marker genes and ontologies) to stabilize the learning process.
• Limitations: The authors modestly note that the approach relies on curated marker gene programs, which may be incomplete or noisy for novel cell types. Additionally, the performance gains do not necessarily imply improved parameter efficiency.

In summary, CellRefine offers a framework to bridge the gap between large-scale pretraining and the specific, often imbalanced, realities of single-cell biological data, yielding more discriminative and generalizable cell representations.