Skip to content

Rethinking Class Imbalance for Single-Cell Foundation Models: A Systematic Benchmark Across Architectures and Long-Tail Loss Functions

Sep 2026 · 0 citations · 48 references
Computer Science

TL;DR

This work presents a systematic benchmark of six long-tail loss functions (cross-entropy, weighted CE, class-balanced loss, focal loss, LDAM, logit-adjusted softmax), giving both a reusable benchmark and mechanism-grounded practical guidelines for combining foundation models with imbalanced biological data.

Abstract

Single-cell foundation models (scGPT, scBERT, Geneformer) achieve cell-type classification accuracy up to 97.5% in our experiments, yet this aggregate accuracy can mask systematic failure on rare, often disease-relevant cell populations that long-tail loss functions are widely assumed to address. We present a systematic benchmark of six long-tail loss functions (cross-entropy, weighted CE, class-balanced loss, focal loss, LDAM, logit-adjusted softmax) across three architectures and three datasets (Multiple Sclerosis, Zheng68K, human Pancreas), totaling 162 controlled training runs (3 backbones x 3 datasets x 6 losses x 3 seeds). The gap between overall accuracy, Macro-F1, and rare-class recall under plain cross-entropy is consistent across all nine (architecture, dataset) settings, driven by dataset structure rather than pretraining. Rare-class failure itself splits into two regimes with distinct embedding-geometry signatures, visible before any loss is chosen: some classes are recoverable by the right loss, while others retain linear separability yet are absorbed into unrelated classes'neighborhoods under every evaluated loss and architecture. Among the recoverable classes, the efficacy of reweighting is predicted by a class's absolute training-set size, rather than its share of the dataset or the dataset's overall imbalance ratio. Class-balanced loss and LDAM are the most consistent choices across all nine settings, while logit adjustment trades rare-class precision for recall rather than improving both. Our results give both a reusable benchmark and mechanism-grounded practical guidelines for combining foundation models with imbalanced biological data.

View source

Similar papers

Open access Sep 2026

scBalFlow: a staged flow matching framework for imbalanced single-cell drug perturbation prediction

Abstract Motivation Conventional drug perturbation prediction models typically employ end-to-end encoder-decoder architectures, directly mapping control samples and perturbation conditions to post-perturbation gene expression profiles. However, these approaches widely overlook the severe class imbalance inherent in per...

H. Lyu, Jia-Wei Luo · 0 citations
Open access Oct 2026

Benchmarking single-cell foundation models reveals an advantage in the low-data regime for cell-type and disease-state prediction

Single-cell foundation models perform well on cell-type classification, but their value for disease-state prediction and their advantage over classical representations remains unclear. We benchmarked cell-representations from four foundation models using logistic regression, on 868 disease-prediction tasks (82 datasets...

Haya Benhayon, Michael M. Danziger, Danny Ben-Zvi et al. · 0 citations
Conference Sep 2026

sc2Flow: Mitigating Mean Prediction Bias in Single-Cell Perturbation with Dual-Stage Flow Matching

Predicting single-cell gene responses to chemical perturbations is vital for personalized therapy, yet existing deep learning models face significant hurdles. Standard regression-based approaches suffer from "mean prediction bias," failing to capture cellular heterogeneity, while current Flow Matching (FM) methods stru...

H. Lyu, Jia-Wei Luo · 0 citations
Open access Oct 2026

Beyond the Final Layer: Evaluating Single-Cell Foundation Model Representations for Molecular Prediction

Single-cell foundation models (scFMs) are commonly benchmarked using final-layer embeddings, although useful depth may depend on the task. We evaluate Geneformer, scBERT, scGPT, and scFoundation as frozen feature extractors for cell-type and perturbation-response prediction using layer-wise probes, leave-one-dataset-ou...

Ming Wen, Yiwei Lyu, Andrew Robbins · 0 citations
Open access Sep 2026

Cell-level random splits leak group-owned answers in single-cell benchmarks

This work shows how to estimate leakage before training any model, from two properties of the data: exposure, the fraction of test cells with relatives in training, and retrievability, how often a nearest-neighbor search returns such a relative rather than an unrelated cell.

Cang Hu, Sha Sun · 0 citations
Open access Aug 2026

Toward Imbalanced Molecular Property Regression: A Benchmark Study and Interval-Aware Mixture of Experts

Molecular property prediction is a key task in AI-driven drug discovery, yet the prevalence and impact of label imbalances in molecular property regression remain poorly understood. Through a systematic benchmark of widely used molecular property data sets, we show that target values are often highly imbalanced and tha...

Y. Sun, Yu Shi, Alana Deng et al. · 0 citations

Related blog posts

MIT News · Artificial Intelligence Oct 23, 2026

Q&A: How AI came to model the “messy” human mind

Assistant Professor Benjamin Lindquist studies how engineers deliberately injected randomness into early AI, giving rise to deep-learning models that sometimes hallucinate.

We use cookies to run the site and, with your consent, for analytics and to show ads. See our Cookie Policy.