## Scaling anndata training to the terabyte scale with annbatch

The demand for AI in omics has grown at an unprecedented rate, with
state-of-the-art models now routinely trained on datasets exceeding
the terabyte scale. To make that process more efficient, we developed
"annbatch",[1] a high-performance data loader built on "anndata" that
enables loading speeds of 60k samples/second and more, at least a
factor of 3 higher than the fastest recent alternatives.

Since "anndata"[2] released a first disk-backed data loader
("AnnCollection") in 2019, better implementations have been developed,
for instance, as used in SCimilarity[3], Cellarium[4] or BioNEMO.[5]
Based on these improvements, some of us helped develop
"MappedCollection"[6] to address the need for true weighted random
sampling in 2023. This came, however, at a significant performance
cost compared to approaches that load contiguous chunks, such as
NVIDIA Merlin[7] or the "tiledbsoma" loader of CELLxGENE.[8] In 2025,
"scDataset"[9] and SLAF[10] have been introduced with significant
performance improvements.

With "annbatch",[1] we developed an "anndata"-based loader that
optimizes loading contiguous chunks, assumes pre-shuffling, and uses
the popular ".zarr" array format.[11] It reaches 60k samples/second
and more[1] on the Tahoe-100M dataset,[12] which stores
transcriptional profiles of 100M cells (**Figure 1**). For
reproducibility and to showcase the conversion of the original
collection of ".h5ad" files to a collection of ".zarr" stores,
benchmarks were tracked with data lineage (**Figure 2**).

**Figure 1 (source)**: Dataloader throughput on the Tahoe-100M dataset
across three loaders, with "scDataset"[9] shown both with a matched
block/chunk size and with its recommended settings. By clicking on
"source", you can navigate to the runs that produced the results. For
example, the run producing the "annbatch" results is here with
information about parameters, environment, and hardware
("ml.m5.24xlarge" on AWS).

**Figure 2 (explore)**: Processing pipeline from the originally
published Tahoe-100M, over pre-shuffled datasets, to running the data
loader, to plotting Figure 1.

While the benchmarks here cache arrays from their cloud storage
location to local disk, streaming data directly from S3 is a relevant
use case, albeit resulting in lower loading speeds. A recent benchmark
by Ryan Conrad[13] found "annbatch" to perform well in this case, too.

# Data & code availability

* Paper: arXiv:2604.01949

* Repo: github.com/scverse/annbatch

* Benchmarking database: lamin.ai/laminlabs/arrayloader-benchmarks

* Benchmarking repo: github.com/laminlabs/arrayloader-benchmarks

# Acknowledgements

We are grateful to Raaghav Pillai for re-running benchmarks with the
latest versions after a long development process. We thank Sergei
Rybakov for early discussions, following the development of
"MappedCollection". We thank Pavan Ramkumar for feedback and for
validating early benchmarks. We thank Davide D'Ascenzo and Sebastiano
Cultrera di Montesano for discussions related to "scDataset".

# References

---

[1] Gold I, Fischer F, Arnoldt L, Wolf FA & Theis FJ (2026). MCML -
 Annbatch Unlocks Terabyte-Scale Training of Biological Data in
 Anndata. arXiv.

[2] Virshup I, Rybakov S, Theis FJ, Angerer P & Wolf FA (2024).
 anndata: Access and store annotated data matrices. *Journal of
 Open Source Software*, 9(101), 4371. doi:10.21105/joss.04371.

[3] Heimberg G, Kuo T, DePianto DJ, Salem O, Heigl T, Diamant N,
 Scalia G, Biancalani T, Turley SJ, Rock JR, Corrada Bravo H,
 Kaminker J, Vander Heiden JA & Regev A (2025). A cell atlas
 foundation model for scalable search of similar human cells.
 Nature.

[4] Cellarium AI (2022). Cellarium-ML: Distributed single-cell data
 analysis. GitHub.

[5] St. John P et al. (2024). BioNeMo Framework: a modular, high-
 performance library for AI model development in drug discovery.
 arXiv.

[6] Rybakov S, Fischer F, Wiatrak M, Gold I, Rosen Y, Sun S,
 Sriworarat C, Theis F, Kalfon J & Wolf A (2024). MappedCollection:
 Weighted random sampling from large collections of scRNA-seq
 datasets. Lamin Blog. blog.lamin.ai/mapped-collection.

[7] Oldridge E, Perez J, Frederickson B, Koumchatzky N, Lee M, Wang Z,
 Wu L, Yu F, Zamora R, Yilmaz O, Gunny A & Nguyen V (2020). Merlin:
 A GPU Accelerated Recommendation Framework. ACM.

[8] CELLxGENE Census experimental PyTorch data pipeline docs.
 chanzuckerberg.github.io/cellxgene-
 census/notebooks/experimental/pytorch.html.

[9] D'Ascenzo D & Cultrera di Montesano S (2025). scDataset: Scalable
 Data Loading for Deep Learning on Large-Scale Single-Cell Omics.
 arXiv.

[10] Pavan Ramkumar (2025). SLAF: Sparse Lazy Array Format. slaf-
 project.github.io.

[11] Zarr developers (2024). Zarr storage format specification v2.
 zarr-specs.readthedocs.io.

[12] Zhang JQ et al. (2025). Tahoe-100M: A Giga-Scale Single-Cell
 Perturbation Atlas for Context-Dependent Gene Function and
 Cellular Modeling. bioRxiv.

[13] Ryan Conrad (2026). Re-benchmark "AnnData" for remote stores.
 GitHub.