scTrilemma: Balancing Identity, Invariance, and Fidelity in Single-Cell Representation Learning

Yunhak Oh; Yoonho Lee; Junseok Lee; Namkyeong Lee; Sang-Yeon Hwang; Yinhua Piao; Hyomin Kim; Seonghwan Kim

cs.LGSep 30, 2026

scTrilemma: Balancing Identity, Invariance, and Fidelity in Single-Cell Representation Learning

Authors: Yunhak Oh, Yoonho Lee, Junseok Lee, Namkyeong Lee, Sang-Yeon Hwang, Yinhua Piao, Hyomin Kim, Seonghwan Kim, +4 more

Organizations: KAIST · Work done during an internship at HITS · HITS

Abstract

Single-cell RNA-seq representation learning is fundamentally label-free: cell identities, states, and contexts are not fixed training targets, so what constitutes signal or nuisance is analysis-dependent. A single representation must therefore preserve biological identity and state, remain robust to nuisance context, and retain the gene-level variation needed for expression analysis, three demands we call the representation trilemma. To tackle this problem, we introduce scTrilemma, a latent-bottleneck VAE that routes expression-derived variation to the embedding, the decoder, or the prior rather than forcing all of it through one embedding. It gates gene tokens by expression, routes the cell representation through the decoder, and conditions the prior on unlabeled pseudo-bulk context, under a single reconstruction objective and without target annotations or auxiliary representation losses. In release-based zero-shot evaluation on successive CZ CELLxGENE Census releases, scTrilemma leads all three demands at once and preserves biological-state, differential-expression, and pathway structure across multiple disease settings. Latent interventions further show that context can be removed at almost no cost to the other demands, leaving identity against fidelity as the remaining tension. Code is publicly available at https://github.com/yunhak0/scTrilemma.

Figures & tables

Figure 1: The representation trilemma ; edge labels name the operation each pair contends over. Biological Identity preserves cell-type and state structure. Context Invariance prevents nuisance context from defining cellular similarity. Expression Fidelity retains gene-level variation needed for biological analysis.

Figure 2: scTrilemma architecture. (a) Overall routing architecture trained with ZINB reconstruction and KL regularization. (b) E-Gate applies an expression-derived feature-wise gate before latent compression. (c) C-Route mean-pools posterior means for the evaluated cell representation and uses the pooled latent sample to modulate decoder queries, while full sampled latent tokens remain reconstruction memory. (d) PB-Cond softly assigns dataset–donor pseudo-bulks to fixed pretraining centroids to parameterize the conditional prior.

	Biological identity					Context invariance
Method	NMI	ARI	ASW	cLISI	Iso. Label	BRAS	iLISI	PCR
scVI	$0.640$ $(0.138)$	$\mathit{0.476}$ $(0.155)$	$\mathit{0.570}$ $(0.049)$	$0.973$ $(0.050)$	$\mathit{0.598}$ $(0.072)$	$0.730$ $(0.093)$	$\mathit{0.218}$ $(0.174)$	$0.195$ $(0.135)$
Geneformer	$0.558$ $(0.166)$	$0.377$ $(0.156)$	$0.545$ $(0.037)$	$\mathbf{0.980}$ $(0.046)$	$0.567$ $(0.062)$	$0.732$ $(0.106)$	$0.167$ $(0.169)$	$0.112$ $(0.115)$
scGPT	$\mathit{0.645}$ $(0.141)$	$0.470$ $(0.158)$	$\mathbf{0.575}$ $(0.052)$	$0.977$ $(0.047)$	$\mathbf{0.600}$ $(0.081)$	$0.728$ $(0.091)$	$0.196$ $(0.173)$	$0.215$ $(0.107)$
CellPLM	$0.595$ $(0.142)$	$0.404$ $(0.150)$	$0.559$ $(0.077)$	$0.970$ $(0.049)$	$0.598$ $(0.127)$	$0.677$ $(0.087)$	$\mathbf{0.236}$ $(0.178)$	$\mathit{0.284}$ $(0.175)$
scPRINT $\dagger$	$0.405$ $(0.154)$	$0.240$ $(0.139)$	$0.510$ $(0.017)$	$0.914$ $(0.073)$	$0.523$ $(0.035)$	$\mathit{0.748}$ $(0.105)$	$0.138$ $(0.135)$	$0.025$ $(0.063)$

Table 1: Release-based zero-shot embedding benchmark. Values are 20-repeat means (SD) across datasets. Bold and italics mark the best and second-best values; stars mark gains over the strongest baseline by paired Wilcoxon test ( ∗ :

p<

0.05, ∗∗∗ :

p<

0.001).

\dagger

Metadata-supervised pretraining.

Metric	scVI	Geneformer †	scGPT †	CellPLM	scPRINT ‡	scTrilemma
Biological identity and state
Cell-type BA $-$ chance	$0.6535$	$0.6112$	$\mathit{0.6667}$	$0.6041$	$0.4475$	$\mathbf{0.6697}$
Neighbor purity	$\mathit{0.6817}$	$0.6006$	$\mathbf{0.7041}$	$0.6459$	$0.4327$	$0.6795$
Disease-state BA	$0.5583$	$\mathit{0.5785}$	$0.5744$	$0.5507$	$0.5432$	$\mathbf{0.5856}$
Context invariance
Residual donor invariance	$0.8889$	$0.8864$	$0.8778$	$\mathit{0.9161}$	$0.9016$	$\mathbf{0.9336}$

Table 2: Matched evaluation of the three representation demands on the common cohorts. Values are equal-cohort means; bold and italics mark the best and second-best values. Expression metrics are reported only for models with a full expression output ( † : no decoder weights in the checkpoint; ‡ : metadata-supervised pretraining).

Figure 3: Component ablation of the full model. Diamonds show the mean score change without each component (w/o) relative to the full model, standardized by its across-dataset standard deviation. Negative values indicate degradation, MSE is sign-flipped, and

r

denotes Pearson correlation. For C-Route, the mean-pooled readout is retained and only decoder-query routing is removed. Asterisks indicate significance after BH-FDR correction of paired two-sided Wilcoxon tests (

{}^{*}q<0.05

;

{}^{**}q<0.01

Figure 4: Demand-targeted interventions on the trained model. Each panel pushes one demand on the posterior-mean latent at inference (centroid shrinkage, donor centering, or latent refinement at strength

\alpha=0.25

0.5

0.75

) and plots the change of the other two demands against the change of the pushed metric; fidelity is shown as reconstruction Pearson and as MSE with its sign reversed. Points are dataset means with SEM.

Appendix figures & tables33 assets

Supplementary material from the paper’s appendix.

Appendix

Cohort	Dataset ID prefix	Fixed context	Disease/state (vs. normal)	Evaluation support †
Kidney	867757c1	10x multiome, kidney	obstructive nephropathy	13 / 2,243 / 520 / 1,524
Liver	e3ed2ba4	BD Rhapsody WTA, liver	metastatic colorectal carcinoma	11 / 2,073 / 440 / 2,135
Brain	203025fe	10x multiome, dorsolateral prefrontal cortex	cognitive disorder	7 / 830 / 280 / 1,821
Tendon	acd544d0	10x 3’ v3, quadriceps femoris tendon	injury	4 / 854 / 160 / 1,555
Colon	829a3cd1	10x 3’ v3, sigmoid colon	colorectal cancer	6 / 959 / 240 / 1,639

Appendix

Table 3: Metadata-defined cohorts and evaluation support for Table 2 (RQ2–4). † Evaluation-support entries report eligible cell types, common-pool cells, cells per repeat, and common genes, in that order.

Dataset ID prefix	Disease (vs. normal)	Fixed context	Evaluation support †
16023185	Colon adenocarcinoma	10x 3 ′ v2, colorectum	43 / 38 / 12,000
f14bc322	Interstitial lung disease	10x 5 ′ v2, lung	93 / 33 / 12,000
d68a8b48	Bronchopulmonary dysplasia	10x 3 ′ v3, middle lobe of right lung	21 / 35 / 12,000
19053a82	Ulcerative colitis	10x 3 ′ v2, colon	11 / 40 / 12,000

Appendix

Table 4: Metadata-defined non-blood disease-state cohorts used in RQ2. Cohorts were fixed from metadata and cell-count support before model evaluation. Eligible cell types contain at least 20 cells from at least two donors in each state. † Evaluation-support entries report donors, eligible cell types, and sampled cells, in that order.

Candidate collection	Shared donors	Eligible strata	Evaluation support †	Outcome
RPE/choroid atlas	4	24	4 / 9	Selected
Anterior-segment atlas	8	3	2 / 3	Lower matched support
Ocular-surface atlas	0	0	0 / 0	No shared donor

Appendix

Table 5: Metadata-defined screening of matched scRNA-seq/snRNA-seq candidates for RQ3. Eligible-stratum counts are evaluated in the fixed cross-method evaluation pool after matching donor, annotated cell type, and tissue, with at least 50 observations from each modality. † Evaluation-support entries report represented donors and cell types, in that order.

Donor	Eligible cell types	Strata	Cells per repeat
BCM_22_0500	endothelial cell of venule, fenestrated endothelial cell, fibroblast, macrophage, melanocyte	5	878
BCM_22_0698	endothelial cell of venule, fenestrated endothelial cell, fibroblast, macrophage, melanocyte, pericyte	6	1,068
BCM_22_0769	T cell, endothelial cell of venule, fibroblast, macrophage, melanocyte, retinal pigment epithelial cell	6	1,120
BCM_22_0784	T cell, endothelial cell of venule, fibroblast, macrophage, melanocyte, monocyte, retinal pigment epithelial cell	7	1,290
Total	9 unique cell types	24	4,356

Appendix

Table 6: Composition of the 24 eligible donor–cell-type strata used in RQ3.

Dataset ID prefix	Disease (vs. normal)	Evaluation support †
30c2a6fd	Open-angle glaucoma	15 / 6 / 1,230
84cfa5aa	Intestinal failure-associated liver disease	14 / 6 / 2,045
867757c1	Obstructive nephropathy	13 / 6 / 1,524
e3ed2ba4	Metastatic colorectal carcinoma	11 / 6 / 2,327
203025fe	Alzheimer disease	8 / 6 / 1,877
829a3cd1	Colorectal cancer	8 / 6 / 1,811

Appendix

Table 7: Disease cohorts retained by the donor-supported screen of all 89 held-out datasets for the supplementary RQ4 expression-fidelity analysis. All seven eligible disease-labeled datasets are included and sampled to 2,500 cells each. † Evaluation support reports the numbers of supported cell types, selected contrasts, and common genes, in that order. Common genes are the shared gene list of Appendix A.2.5 after the 2% detection filter.

Package	Version
torch	2.8.0+cu128
lightning / pytorch-lightning	2.6.0
hydra-core / omegaconf	1.3.2 / 2.3.0
numpy / scipy / pandas	2.3.5 / 1.16.3 / 2.3.3
anndata / scanpy	0.12.7 / 1.11.5
scikit-learn	1.8.0

Appendix

Table 8: Main Python package versions in the training environment.

Model	Params	Emb. dim	Batch size	Median time (s)	Cells/s	Peak GPU memory
scVI	8.0M	50	512	0.041	49,225.2	0.09GB
scGPT	51.9M	512	64	2.141	934.0	1.52GB
CellPLM	65.8M	512	2,000 †	0.396	5,056.6	0.88GB
Geneformer	316.3M	1152	16	32.349	61.8	3.15GB
scTrilemma	58.3M	512	256	6.433	310.9	19.79GB

Appendix

Table 9: Direct embedding-extraction runtime and memory profile. All rows use a fixed 2,000-cell sample from the same held-out dataset. Throughput is computed from the median runtime. † CellPLM processes the sample as a single batch, as in its benchmark pipeline.

Category	Setting	Value
Training corpus	CELLxGENE Census release	20250130
Zero-shot evaluation release	CELLxGENE Census release	20251108
Held-out zero-shot benchmark	Release-based benchmark	89 newly appearing datasets
Counterfactual subset	Held-out blood/protocol subset	10 blood datasets from the 20251108 release
Random seed	Seed	42
Batch size	Cells per batch	256 per GPU (1,024 total across 4 GPUs)

Appendix

Table 10: Data, batching, and preprocessing settings.

Category	Setting	Value
Optimizer	AdamW	learning rate $2\times 10^{-4}$ , weight decay $10^{-4}$ , betas $(0.9,0.999)$
Scheduler	Cosine with linear warmup	1,250 warmup steps (5% of 25,000), then cosine decay to step 25,000
KL loss	KL coefficient, warmup	$5\times 10^{-4}$ , 2,500 steps (10% of the 25,000-step recipe)
Trainer	Precision, devices, strategy	bf16-mixed; 4 GPUs; DDP
Trainer schedule	Steps, validation interval	25,000, 5,000
Gradient clipping	Clip value	1.0

Appendix

Table 11: Optimization, loss, and trainer settings.

Component	Setting	Role
Model family	Latent-bottleneck VAE with cross-attention encoder and decoder	Compresses gene-expression tokens into latent cell tokens before reconstruction.
Trainable parameters	58,347,460	Counted from the active final model configuration.
Gene vocabulary	61,890 Ensembl genes	Gene IDs are indexed by the 20250130 CELLxGENE vocabulary.
Token width	$d_{\mathrm{model}}=512$	Shared width for gene tokens, latent tokens, attention blocks, and decoder states.
Gene embedding	Learnable, 512-dimensional	Provides the base gene identity token.
Expression-gated gene encoding	Multiplicative expression gate	Gates gene embeddings by measured expression before compression.

Appendix

Table 12: Architecture settings for the final scTrilemma model.

Component	Setting	Role
Grouping unit	Dataset-donor group	Defines the collection-level unit for pseudo-bulk profiles.
Pseudo-bulk profile	Per-cell CP10K normalization, group mean, then elementwise $\log(1+\cdot)$	Produces the group-level expression signature used for conditioning.
Centroid construction	k-means with $M=32$ , random seed 42, and five random initializations	Learns expression-derived collection archetypes from pseudo-bulk profiles.
Soft assignment	Euclidean distance to centroids with distance standardization	Produces a 32-dimensional soft pseudo-bulk code.
Zero-shot-release assignment	Fixed 20250130 centroids; no centroid refit on 20251108	Assigns unseen release groups by soft assignment to the precomputed centroids.
Prior conditioning	Diagonal Gaussian prior	Conditions the VAE prior using the pseudo-bulk code.

Appendix

Table 13: Pseudo-bulk-derived conditioning settings.

Evaluation	Metrics	Reported on	Direction †
RQ1 biological identity	NMI, ARI, ASW-label, cLISI, isolated labels	Cell embedding	$\uparrow$
RQ1 context invariance	BRAS, iLISI, PCR comparison	Cell embedding	$\uparrow$
RQ2 biological identity and state	Cell-type BA minus chance, neighbor purity, disease-state BA	Cell embedding	$\uparrow$
RQ3 context invariance	Residual donor invariance, conditional donor BRAS, local donor mixing	Cell embedding	$\uparrow$
RQ4 expression fidelity	logFC Spearman, DEG Jaccard and sign, pathway Jaccard and score Spearman	Reconstructed expression	$\uparrow$
Supplementary identity and state	Program-neighbor Spearman, purity, accuracy, balanced accuracy	Embedding neighborhoods and expression	$\uparrow$

Appendix

Table 14: Summary of the evaluation metrics defined in Appendix C . † Direction indicates whether higher (

\uparrow

) or lower (

\downarrow

) values are better.

Figure 5: Dataset characteristics associated with RQ1 score differences. Each point is one held-out dataset, and

\Delta

is scTrilemma minus the fixed external comparator. Panels A and B relate annotated cell-type count to NMI and isolated-label silhouette; Panel C relates raw context burden to BRAS. Red lines show least-squares trends on the displayed scale, dashed lines mark equal performance, and

\rho

denotes Spearman correlation.

Marker program	Share	scVI	Geneformer	scGPT	CellPLM	scPRINT	scTrilemma
CD4 T cell	19.3	0.8832	0.8983	0.8869	0.8820	0.8575	0.9070
CD8/cytotoxic T	15.0	0.8411	0.8581	0.8256	0.8032	0.7853	0.8468
B cell	13.5	0.5883	0.5850	0.5730	0.5752	0.5511	0.5885
NK cell	11.7	0.7811	0.7819	0.7713	0.7637	0.7686	0.7811
Classical monocyte	5.8	0.6446	0.6518	0.6554	0.6457	0.6679	0.6656
Non-classical mono.	3.8	0.7471	0.7620	0.7548	0.7422	0.7271	0.7605

Appendix

Table 15: Cross-donor marker-program consistency in held-out normal blood. Share is the annotated population percentage; ambiguous T-cell labels are excluded from subtype shares. Scores are Spearman correlations between cell program scores and their 30-neighbor cross-donor means. Macro mean weights the nine programs equally. Bold and italics mark the highest and second-highest values.

Cohort	Metric	scVI	Geneformer	scGPT	CellPLM	scPRINT	scTrilemma
Colon	Purity	0.6432	0.6371	0.6312	0.6310	0.6249	0.6583
	Accuracy	0.7276	0.7267	0.7274	0.7093	0.7299	0.7507
	Balanced acc.	0.6499	0.6511	0.6496	0.6220	0.6482	0.6733
ILD	Purity	0.6083	0.6082	0.5994	0.5952	0.5962	0.6385
	Accuracy	0.6920	0.6941	0.6868	0.6831	0.6830	0.7359
	Balanced acc.	0.5802	0.5828	0.5750	0.5659	0.5667	0.6294

Appendix

Table 16: Cohort-level within-cell-type disease-state retrieval in the supplementary biological-state analysis. Each cohort contains 12,000 sampled cells, and same-donor neighbors are excluded. The final block reports equal-cohort means, with sample SD across cohorts shown in parentheses. ILD, BPD, and UC denote interstitial lung disease, bronchopulmonary dysplasia, and ulcerative colitis. Bold and italics mark the highest and second-highest values within each cohort and metric, respectively.

Metric	scVI	Geneformer	scGPT	CellPLM	scPRINT	scTrilemma
Modality $R^{2}$ $\downarrow$	$0.4192$	$\mathit{0.3645}$	$0.5357$	$0.7190$	$0.3913$	$\mathbf{0.2677}$
Modality ASW $\uparrow$	$0.4265$	$\mathit{0.4765}$	$0.3169$	$0.1696$	$0.4523$	$\mathbf{0.5932}$
Conditional iLISI $\uparrow$	$\mathbf{0.0755}$	$0.0324$	$0.0259$	$\mathit{0.0356}$	$0.0270$	$0.0329$

Appendix

Table 17: Matched scRNA-seq/snRNA-seq context invariance. Values are means across 20 balanced cell resamples. Arrows indicate the favorable direction. Bold and italics mark the best and second-best values, respectively.

Figure 6: Stratum-level matched-modality results in RQ3. Values are means over 20 balanced cell resamples. (A) Modality

R^{2}

and (B) modality ASW for each donor–cell-type stratum; donor labels omit the common BCM_22_ prefix. Black points are scTrilemma , and colored points are the strongest external result per stratum among scVI, Geneformer, scGPT, CellPLM, and scPRINT, connected to scTrilemma by a line. EdV, endothelial cell of venule; FEC, fenestrated endothelial cell; Fib., fibroblast; Mac., macrophage; Mel., melanocyte; Per., pericyte; Mono., monocyte; RPE, retinal pigment epithelial cell.

Cohort	Method	logFC Spearman $\uparrow$	top-100 DEG Jaccard $\uparrow$	top-100 DEG Sign $\uparrow$	Pathway Jaccard $\uparrow$	Pathway score Spearman $\uparrow$
OAG	scVI	0.1573	0.1379	0.6367	0.2406	-0.1662
	scPRINT	0.5559	0.2163	0.7550	0.2535	-0.1124
	CellPLM	0.2194	0.1655	0.7000	0.2594	-0.1456
	scTrilemma	0.4302	0.2498	0.8683	0.3261	0.1589
IFALD	scVI	0.1932	0.1401	0.7233	0.2247	-0.1954
	scPRINT	0.4062	0.2528	0.7600	0.3823	0.2449

Appendix

Table 18: Cohort-level and equal-cohort differential-expression and pathway concordance in the supplementary expression-fidelity analysis. Each cohort value averages its within-cell-type disease-versus-normal contrasts; the final block reports equal-cohort means with sample SD across cohorts in parentheses. AD, Alzheimer disease; CRC, colorectal cancer; IFALD, intestinal failure-associated liver disease; mCRC, metastatic colorectal carcinoma; OAG, open-angle glaucoma; ON, obstructive nephropathy; T2D, type 2 diabetes mellitus. Arrows indicate the favorable direction, and bold and italics mark the highest and second-highest values within each block and metric.

All shared genes
Method	Pearson $\uparrow$	Spearman $\uparrow$	MAE $\downarrow$	MSE $\downarrow$
scVI	0.3722 $\pm$ 0.0846	0.3098 $\pm$ 0.0602	1.0041 $\pm$ 0.1157	1.5788 $\pm$ 0.3013
scPRINT	0.4751 $\pm$ 0.0577	0.3962 $\pm$ 0.0617	1.0475 $\pm$ 0.2220	1.6569 $\pm$ 0.5433
CellPLM	0.5785 $\pm$ 0.1045	0.4404 $\pm$ 0.0759	0.8135 $\pm$ 0.2247	1.3003 $\pm$ 0.5483
scTrilemma	0.6380 $\pm$ 0.0863	0.4524 $\pm$ 0.0584	1.4162 $\pm$ 0.2310	2.2828 $\pm$ 0.6831
Genes detected in $\geq$ 20% of cells

Appendix

Table 19: Direct reconstruction agreement across the release-held-out datasets. Values are equal-dataset means

\pm

sample standard deviations. Arrows indicate the favorable direction, and bold and italics mark the highest and second-highest values within each block. The second block restricts each dataset to genes detected in at least 20% of its cells.

Push identity (centroid shrinkage)
	NMI $\uparrow$	Label ASW $\uparrow$	BRAS $\uparrow$	iLISI $\uparrow$	Recon. Pearson $\uparrow$	MSE $\downarrow$
$\alpha=0.25$	$-0.002$ (40/89)	$+0.014$ (86/89)	$-0.024$ (0/75)	$+0.003$ (65/75)	$-0.007$ (0/89)	$+0.018$ (88/89)
$\alpha=0.5$	$-0.002$ (46/89)	$+0.031$ (85/89)	$-0.071$ (0/75)	$+0.009$ (69/75)	$-0.023$ (0/89)	$+0.048$ (89/89)
$\alpha=0.75$	$-0.001$ (46/89)	$+0.054$ (77/89)	$-0.175$ (0/75)	$+0.011$ (69/75)	$-0.053$ (0/89)	$+0.087$ (89/89)
Push identity (label centroids)
$\alpha=0.25$	$+0.105$ (89/89)	$+0.030$ (89/89)	$-0.003$ (1/75)	$+0.003$ (61/75)	$-0.004$ (7/89)	$+0.021$ (89/89)

Appendix

Table 20: Demand-targeted interventions. Each block pushes one demand at strength

\alpha

; the label-centroid block shrinks cells toward annotated cell-type centroids instead of k-means centroids, and each entry is the mean of the per-dataset change from the unmodified model, with the number of datasets in which the metric increases in parentheses.

	$\Delta$ NMI $\uparrow$	$\Delta$ ASW $\uparrow$	$\Delta$ BRAS $\uparrow$	$\Delta$ iLISI $\uparrow$	$\Delta$ HVG Pearson $\uparrow$	$\Delta$ MSE $\downarrow$
Backbone	$+0.0008$ $(0.0015)$	$+0.0000$ $(0.0005)$	$\mathbf{-0.0035}^{*}$ $(0.0009)$	$-0.0005$ $(0.0004)$	$\mathbf{+0.0571}^{*}$ $(0.0061)$	$+0.0026$ $(0.0020)$
E-Gate only	$-0.0014$ $(0.0014)$	$\mathbf{+0.0008}^{*}$ $(0.0002)$	$\mathbf{-0.0090}^{*}$ $(0.0005)$	$\mathbf{-0.0024}^{*}$ $(0.0005)$	$\mathbf{-0.0341}^{*}$ $(0.0037)$	$\mathbf{-0.0447}^{*}$ $(0.0044)$
E-Gate + C-Route (full model)	$\mathbf{+0.0122}^{*}$ $(0.0012)$	$\mathbf{+0.0018}^{*}$ $(0.0002)$	$\mathbf{+0.0020}^{*}$ $(0.0005)$	$-0.0007$ $(0.0005)$	$\mathbf{+0.0046}^{*}$ $(0.0011)$	$\mathbf{-0.0234}^{*}$ $(0.0031)$

Appendix

Table 21: PB-Cond switched off and on within three routing configurations (backbone, E-Gate only, E-Gate + C-Route). Values are dataset-level means of the on-minus-off change, with SEM in parentheses. Asterisks mark paired two-sided Wilcoxon tests after BH-FDR correction (

{}^{*}q<0.05

). MSE is the cell-wise mean squared error over the genes detected in at least 2% of the sampled cells, before restriction to the four-model common list of Appendix A.2.5 , so its absolute value differs slightly from Table 19 .

	NMI $\uparrow$	ASW $\uparrow$	BRAS $\uparrow$	iLISI $\uparrow$	Recon. Pearson $\uparrow$	MSE $\downarrow$
scVI, no batch covariate	$0.6019$	$0.5597$	$0.7579$	$0.1765$	$0.5127$	$0.3488$
+ decoder batch covariate	$0.5973$	$0.5634^{**}$	$0.7452^{***}$	$0.1775$	$0.5243^{***}$	$0.3497$
+ PB-conditioned prior	$0.5765^{***}$	$0.5509^{***}$	$0.7425^{***}$	$0.1832^{***}$	$0.5185^{*}$	$0.3524^{**}$

Appendix

Table 22: Matched scVI configurations that receive the same dataset–donor pseudo-bulk code. The three configurations share the seed, training data, training budget, and VAE recipe and differ only in where group information enters the model. Values are means across the held-out datasets; reconstruction Pearson and MSE are the cell-wise Pearson and mean squared error over all genes on the 76 datasets with complete reconstructions. Asterisks mark paired two-sided Wilcoxon tests against the configuration without a batch covariate (

{}^{*}p<0.05

;

{}^{**}p<0.01

;

{}^{***}p<0.001

	Matched	N = 50	Rare-80%	Dominant-80%	Global shuffle	Zero code
$\Delta$ KL (%)	$0.000$	$+0.027$	$+0.436$	$+0.086$	$+3.931$	$+3.821$
Mean $\|\Delta\mathbf{c}_{b}\|$	$0.0000$	$0.0012$	$0.0070$	$0.0024$	$0.0254$	$0.0313$
Max $\|\Delta\|$ outputs	–	$0$	$0$	$0$	$0$	$0$

Appendix

Table 23: Target-side code perturbations applied to the frozen scTrilemma model. Only the pseudo-bulk code is recomputed; the model, cells, and centroids are fixed.

\Delta

KL is the change in the KL term relative to the matched code, and the last row reports the maximum absolute change in the cell representation, posterior mean, and reconstruction (Max

|\Delta|

outputs).

Metric	Full	$\Delta$ Shuffled codes	$\Delta$ w/o PB-Cond
NMI $\uparrow$	$0.6511$	$+0.0021^{**}$	$-0.0129^{**}$
Label ASW $\uparrow$	$0.5404$	$-0.0006^{**}$	$-0.0018^{**}$
BRAS $\uparrow$	$0.8396$	$+0.0025^{**}$	$-0.0020^{**}$
iLISI $\uparrow$	$0.1814$	$-0.0006^{*}$	$+0.0007$
HVG Pearson $\uparrow$	$0.8852$	$-0.0043^{**}$	$-0.0046^{**}$
MSE $\downarrow$	$2.2754$	$+0.0199^{**}$	$+0.0234^{**}$

Appendix

Table 24: Retraining with a globally shuffled donor–code correspondence, scored with the protocol of Figure 3 .

\Delta

is the mean of the per-dataset difference from the full model for the model trained with shuffled codes and for the model without PB-Cond. Asterisks mark paired two-sided Wilcoxon tests after BH-FDR correction (

{}^{*}q<0.05

;

{}^{**}q<0.01

Figure 7: Hyperparameter sensitivity. (A) KL coefficient: zero-shot scores (mean

\pm

SEM across datasets) for the reference

\lambda_{\mathrm{KL}}=5\times 10^{-4}

with a 2,500-step warmup, a ten-fold larger and smaller coefficient, and no warmup, with the same seed, batch size, and 25,000 training steps; asterisks mark paired two-sided Wilcoxon tests against the reference after BH-FDR correction (

{}^{*}q<0.05

;

{}^{**}q<0.01

). (B) PB-Cond codebook size: zero-shot scores for

M=16

M=32

, and

M=64

centroids under the same protocol as (A), with the same tests against

M=32

Figure 8: Gate profile and same-checkpoint counterfactuals of E-Gate in the full model on the held-out blood/protocol subset. (A) Mean gate value per bin of observed expression: zero counts, and input values of at most 1, 1 to 2, 2 to 4, and above 4. (B) Reconstruction and NMI retained relative to the native gate after replacing the gate of the trained full model with a norm-matched scalar (RMS scalar), a norm-matched shuffle (RMS shuffle), its feature-wise mean (Mean gate), a shuffle across cell–gene positions (Shuffle), or one (Gene only). (C) Mean token norm (bars) with reconstruction and NMI (lines) for the native gate and the two norm-matched controls.

Figure 9: Expression-specificity quadrant analysis for E-Gate in the full model on the held-out blood/protocol subset. Genes are stratified within each cell type by observed expression magnitude and one-vs-rest cell-type specificity. Bars show the mean gate attenuation across analyzed cell types, and points show per-cell-type values.

Figure 10: C-Route leave-one-out comparison. Each point is one retrained model in which a single C-Route component of the full model (Full) is replaced: attention pooling (attn-pool), an additive decoder query (additive query), or no pooled contribution to the decoder queries (no query). Axes give the change from Full in dataset means, with Full at the origin, for each pairwise projection of the three demands: (A)

\Delta

NMI against

\Delta

BRAS, (B)

\Delta

NMI against

\Delta

HVG Pearson, and (C)

\Delta

BRAS against

\Delta

HVG Pearson. Error bars are standard errors of the paired per-dataset differences.

Figure 11: Evaluation-time counterfactuals of latent-token readout and decoder-query routing on the held-out blood/protocol subset. Colors denote the trained model (Full, attn-pool, additive query) and hatched bars denote an intervention applied to the same checkpoint at evaluation time. (A) Cell/dataset silhouette gap of Full, of attn-pool under its native attention readout, and of attn-pool with a forced mean readout of the same latent tokens. (B) Cell-type residual recovery of Full and additive query with the pooled representation kept in the decoder queries (on) or removed (off);

\Delta

gives off minus on. (C) Correlation between query-modulation magnitude and the dataset residual for additive query and Full; points are per-dataset values, and CV is the across-gene coefficient of variation of the modulation magnitude.

Figure 12: Encoder cross-attention enrichment over gene co-expression modules after controlling for mean expression. (A) Observed-over-expected attention mass for every latent query (rows) and co-expression module (columns) in the last encoder layer; queries are sorted by their most enriched module. (B) Distribution across latent queries of the maximum corrected enrichment attained by each query.

Module	Top genes	Source-backed enriched terms	Descriptive label
54	PTPRD , NPAS3 , NRXN1 , NLGN1 , PCDH9 , LSAMP	Brain and oligodendrocyte-precursor cell-marker terms; immature-neuron terms; neuron projection and adhesion genes	Neuronal projection / adhesion
22	MT-CO1 , MT-CO2 , MT-ATP6 , MT-ND2 , CALM2 , B2M	rRNA processing; metabolism of RNA; mitochondrial RNA processing and degradation	Translation / mitochondrial RNA axis
16	EDIL3 , NCAM2 , SLC1A2 , PTN , PLP1 , PTPRZ1	Oligodendrocyte and oligodendrocyte-progenitor marker terms; astrocyte and radial-glia terms	Oligodendrocyte / glial
1	HLA-B , PTPRC , CD83 , CCL4 , MT-ATP8 , MT-ND4L	Dendritic-cell, NK-cell, and T-cell marker terms; IL-2/STAT5 and TNF-alpha signaling	Immune / inflammatory

Appendix

Table 25: Post-hoc enrichment-backed annotation of recurrent latent attention modules. Gene sets are constructed from the top protein-coding genes receiving encoder cross-attention in each module. Enriched terms are used only to describe the modules qualitatively.

Figure 13: PB-Cond ablation and prior-geometry diagnostics. (A) Paired per-dataset NMI, BRAS, and HVG Pearson deltas for PB-Cond relative to the E-Gate+C-Route configuration, with dot size proportional to cell count. (B) NMI, BRAS, and HVG-reconstruction deltas relative to the same configuration for PB-Cond and for memory-based and learned-identifier context routes. (C) Posterior–prior geometry of PB-Cond and learned-identifier prior controls: bars give the matched-random prior margin and points the prior-centered posterior MSE.

Figure 14: Post-hoc metadata annotation of expression-derived PB-Cond clusters. Each row is a pseudo-bulk codebook cluster, and each column is a top tissue-general metadata term. Dot size denotes the fraction of cells in the cluster assigned to that tissue term, while color denotes enrichment relative to the global background. Metadata are used only for annotation after clustering.

Explore similar work

CellMSA: Context Modeling for Single-Cell Representation Learning
Sep 30, 2026Suyuan Zhao, Minghao Liu, Yizhen Luo +1Transcriptomics Representation Learning
Institute for AI Industry Research (AIR), Tsinghua University · Department of Computer Science and Technology, Tsinghua University · Tsinghua Institute of Multidisciplinary Biomedical Research (TIMBR), Tsinghua University +2
Prototype Guided Post-pretraining for Single-Cell Representation Learning
May 8, 2026Sachini Weerasekara, Natasha Darras, Sagar Kamarthi +2Transcriptomics Single-Cell Foundation Models
Sachini Weerasekara∗ · Natasha Darras · Sagar Kamarthi +2
Beyond Gene Reconstruction: Learning Cell Representations through Complementary Transcriptomic Views
Aug 2, 2026Jiaqi Xiong, Yuntao hu, Yu Zheng +3Single-Cell Rna Transcriptomics
University of Oxford, Oxford, United Kingdom · Tongji University, Shanghai, China · Hunan University, Changsha, China +2

Sep 30, 2026, q-bio.GNJ/K move · Enter open · S save

CellMSA: Context Modeling for Single-Cell Representation Learning

Suyuan Zhao, Minghao Liu, Yizhen Luo, Zaiqing Nie

Institute for AI Industry Research (AIR), Tsinghua University · Department of Computer Science and Technology, Tsinghua University · Tsinghua Institute of Multidisciplinary Biomedical Research (TIMBR), Tsinghua University+2

Transcriptomics Representation Learning

Open paper AskPDF arXiv

Single-cell transcriptomics enables profiling of cellular states at unprecedented resolution, but its high dimensionality, sparsity, and technical batch effects pose significant challenges for representation learning. Existing single-cell foundation models typically encode each cell independently or only model cells from the same batch for denoising, thereby underutilizing the rich relational information across batches and cell types to model gene expression patterns. We argue that single-cell models can benefit from more informative cell-context modeling. By comparing consistency and variation across cells, models can capture fine-grained gene-gene dependencies associated with cell states, which are essential for learning high-quality representations. Inspired by the use of multiple sequence alignment (MSA) context in protein modeling, we propose CellMSA, a single-cell representation learning framework that introduces an MSA-inspired inductive bias into transcriptomic modeling. For each target cell, CellMSA retrieves relevant cells from different batches and biologically related cell types as context, and summarizes cross-cell patterns into a context-dependent gene-pair representation. This representation is then injected into a pair-aware target-cell encoder for fine-grained representation learning. We pretrain CellMSA on a large-scale human single-cell corpus of approximately 109 million cell observations, including 65.6 million primary observations. Experiments show that our framework consistently outperforms existing methods across multiple benchmarks. Code is available at the following repository: https://github.com/PharMolix/CellMSA.

Abstract

Explore similar work

CellMSA: Context Modeling for Single-Cell Representation Learning

Prototype Guided Post-pretraining for Single-Cell Representation Learning

Beyond Gene Reconstruction: Learning Cell Representations through Complementary Transcriptomic Views