BA-scVI

Pooling single-cell RNA studies into one atlas runs into batch effects: the same cell type measured on 10x, Microwell-seq, or Seq-Well lands in different regions of expression space, and methods that scrub those differences often scrub biological resolution along with them. Judging which method works has been its own problem, because integration is typically scored on a handful of datasets with a battery of loosely related metrics.

BA-scVI (Batch Adversarial single-cell Variational Inference) comes from Phenomic AI and was presented at the LMRL workshop at ICLR 2025 alongside two things built to answer that scoring problem: the K-Neighbors Intersection (KNI) score, which folds cross-study cell-type label accuracy and a batch-mixing penalty into a single number, and two curated human benchmarks — scMARK (11 studies) and scREF (46 studies) — with author labels standardized across them. BA-scVI is the model that came out of optimizing against that score.

Architecturally it is a variant of scVI with two changes. A discriminator predicts batch identity from the layers adjacent to the central embedding rather than from the embedding itself, which the authors found stabilizes adversarial training, and the batch-ID concatenation that standard scVI feeds its encoder is removed entirely. That second change is the consequential one: without a batch ID at inference, a study the model has never seen can be projected into the existing space directly, which is what reference-based mapping requires. The quickest path to a result is the phenomic-ai PyPI package, whose pai embed command and PaiEmbeddings API take an .h5ad file and one of 32 tissue/organ choices and return cell embeddings and cell-type predictions from Phenomic's hosted scREF service. Little is held back behind that service: the bascvi repository carries the model definition, a PyTorch Lightning training loop, the h5ad_to_zarr.py and shuffle_zarrs.py preprocessing chain, prediction scripts, and the KNI scoring code, and the trained human checkpoint, the scREF corpus, and scMARK are all published separately.

#Key Features

• Inference without a batch ID: Dropping batch-ID conditioning from the encoder lets a new dataset be embedded with no retraining and no batch covariate, unlike standard scVI.
• Adversarial batch penalty on both halves: Discriminators act on encoder outputs and decoder inputs, suppressing batch signal on the way in and on the way out rather than only at the bottleneck.
• Turnkey annotation workflow: pip install phenomic-ai gives a one-command CLI taking an .h5ad file and a tissue label and returning embeddings plus predicted cell types.
• Preserved cell-type granularity: Original author subtype labels remain recoverable from the aligned space, addressing the atlas-building community's standing worry that integration flattens fine cell-type structure.

#Technical Details

The inference network keeps scVI's zero-inflated negative binomial likelihood and its KL terms for the latent and library encoders, adding a term that maximizes discriminator loss; the two networks are updated in alternating steps. The reported configuration uses an adversarial weight of 1000, Adam with learning rates of 5e-5 for the inference network and 1e-2 for the discriminator, batch size 64, and up to 100 epochs, embedding into 10 dimensions. scREF spans 46 human studies, 2,359 samples, and 36 tissues across droplet and plate-based chemistries, standardized to 60 cell-type labels and study-stratified down to 1.21 million cells.

On scREF, BA-scVI reaches a KNI of 0.632 against 0.578 for a tuned four-layer scVI, 0.488 for Harmony, and 0.483 for PCA; on scMARK it reaches 0.711 against 0.646 for Harmony and 0.643 for four-layer scVI. Transformer-based single-cell models fine-tuned per their own protocols scored lower on scREF — 0.399 for Geneformer and 0.468 for scGPT — and substituting BA-scVI for the recommended fine-tuning head improved those alignments without matching the untransformed result. Two-fold cross-validated KNN accuracy on held-out cells reached 83% on a large breast study, 96% once numbered subtypes are merged, 99% on a brain study, and 83% on a blood study.

#Applications

The practical use is annotating and integrating a freshly generated dataset against an organism-scale human reference without curating a training run: a lab points pai embed at its .h5ad, names the tissue, and gets back a 10-dimensional embedding plus predicted cell types consistent with 46 published studies. The same embedding places new cells alongside existing ones, supporting cross-study comparison, and the released checkpoint and training code allow the model to be run locally or retrained on other corpora.

#Impact

BA-scVI's contribution is less a new architecture than a demonstration that a well-tuned adversarial VAE still beats transformer-based single-cell foundation models at organism-scale alignment, measured by a metric that refuses to reward batch mixing and label accuracy separately. The accompanying scREF atlas and KNI score are arguably the more durable artifacts, giving the field a 46-study yardstick where previous comparisons used a handful of datasets. The caveats are real: this is a workshop paper rather than a peer-reviewed article, the benchmarks and released checkpoint are human-only, the KNI score is proposed only for unsupervised alignment because a supervised model could game it, and the CLI path depends on a company-operated backend. Licensing is the loose end: scREF is MIT on HuggingFace and scMARK is CC BY 4.0 on Zenodo, but the bascvi repository ships no LICENSE file, so its MIT grant survives only as a line of README prose — the bundled pai client is the one component with the license text attached.