Energy-based Modelling for Single-cell Data Annotation

Single-cell sequencing has provided profound insights into understanding heterogeneous cellular activities by measuring sequence information at the individual cell resolution. Accurately annotating a single-cell RNA sequencing (scRNA-seq) dataset is a crucial step for the single-cell analysis pipeline. In particular, previously unobserved cell types and cellular states frequently appear in scRNA-seq experiments and carry valuable information. These highlights the need for reliable annotation tools with out-of-distribution (OOD) detection capability. In this work, we introduce energy-based models (EBMs), a family of probabilistic models, for scRNA-seq annotation and OOD detection, which results in more accurate, calibrated, and robust cell type predictions. Recent advancements in energy-based modelling have made it possible to design and deploy EBMs for joint discriminative and generative tasks. Particularly, we present CLAMS, an EBM instance based on the joint energy-based model (JEM), for single-cell data hybrid modelling. Our experiments reveal that hybrid modelling with EBMs maintains the strong discriminative power of baseline classifiers and outperforms the state-of-the-art by integrating generative capabilities in data annotation and OOD detection tasks. In addition, we provide a diagnosis of training JEM and propose effective regularization methods to boost JEM's performance. To the best of our knowledge, we are the first work that applies EBMs to single-cell data modelling.