scDataset: Scalable Data Loading for Deep Learning on Large-Scale Single-Cell Omics

Training deep learning models on single-cell datasets with hundreds of millions of cells requires loading data from disk, as these datasets exceed available memory. While random sampling provides the data diversity needed for effective training, it is prohibitively slow due to the random access pattern overhead, whereas sequential streaming achieves high throughput but introduces biases that degrade model performance. We present scDataset, a PyTorch data loader that enables efficient training from on-disk data with seamless integration across diverse storage formats. Our approach combines block sampling and batched fetching to achieve quasi-random sampling that balances I/O efficiency with minibatch diversity. On Tahoe-100M, a dataset of 100 million cells, scDataset achieves more than two orders of magnitude speedup compared to true random sampling while working directly with AnnData files. We provide theoretical bounds on minibatch diversity and empirically show that scDataset matches the performance of true random sampling across multiple classification tasks and model architectures.

Modern single-cell biology generates datasets profiling the gene expression of millions to hundreds of millions of individual cells. Training machine learning models on these datasets holds enormous promise for understanding disease, drug response, and cellular biology, but it faces a fundamental data engineering challenge: the datasets are too large to fit in computer memory, and loading them efficiently from disk is the primary bottleneck. Training effective models requires presenting data in a random order, but random access to disk is extremely slow. At the scale of a 100-million-cell dataset, iterating through the data once in this way can take weeks to months. The practical alternative, reading data sequentially, is fast but introduces systematic biases: cells from the same experimental batch arrive together, causing the model to repeatedly forget what it learned about earlier conditions as new ones appear, a phenomenon known as catastrophic forgetting. scDataset resolves this tension through two complementary techniques. Block sampling reads groups of nearby cells together, dramatically reducing the number of expensive disk seeks. Batched fetching retrieves many such groups at once, shuffles them in memory, and distributes them across many training steps. Together, these mechanisms recover the statistical benefits of random shuffling at a small fraction of the I/O cost, with formal guarantees on how much diversity is preserved. On Tahoe-100M, a publicly available dataset of 100 million single-cell profiles across 50 cancer cell lines and 380 drugs, scDataset achieves more than 200-fold faster data loading than true random sampling, while producing models that perform identically on downstream biological prediction tasks. The software integrates directly with the existing single-cell analysis ecosystem, requiring no data format conversion or specialized computing infrastructure, and is freely available as an open-source Python library.