TabVI: Leveraging Lightweight Transformer Architectures to Learn Biologically Meaningful Cellular Representations

Abstract

Transformer-based foundation models are changing the landscape of natural language processing (NLP), computer vision, and audio, achieving human-level performance across a variety of tasks. Extending these models to single-cell genomics holds significant potential for revealing the cellular and molecular perturbations associated with disease. However, unlike the sequential structure of language, the functional organization of genes is hierarchical and modular. This fundamental difference necessitates the development of meaningful feature selection strategies to adapt NLP transformer architectures effectively. In contrast to many large-scale foundation models, probabilistic models have shown success in learning complex cellular representations from single-cell datasets. In this work, we present TabVI, a probabilistic deep generative model that leverages tabular transformer architectures to improve latent embedding learning. We validate TabVIs performance in cell type annotation and integration benchmarks. We demonstrate that TabVI improves performance across down-stream tasks and is robust to scaling dataset sizes, producing interpretable, sample-specific feature attention masks. TabVI is a lightweight, scientifically-meaningful, transformer architecture for single-cell analysis that excels where large scale foundation models are less effective.