Single-cell models reveal challenges with rare cell type recognition

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

Machine Learning

Summary

Classifying cell types from biological data is easier for common cells but tougher for rare ones, which may be important in disease. The authors tested different methods designed to help detect these rare cells across various models and datasets. They found that some rare cell types can be identified better with certain loss functions, but others remain hard to separate correctly. Also, the usefulness of techniques that give more weight to rare classes depends on the actual number of rare training examples, not just their proportion. This work helps guide better use of AI models with imbalanced biological data.

What this means in practice

  • For biomedical data engineers: Choose effective loss functions like class-balanced loss or LDAM to improve rare cell type detection in single-cell classification models.
  • For bioinformatics platform developers: Benchmark different long-tail loss functions with single-cell foundation models to design more robust pipelines under imbalanced data conditions.

Authors

Zeyu Dong, Jiahui Zhong

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.