Federated-inspired Single-cell Batch Integration in Latent Space

TL;DR

This work addresses batch effects in large-scale single-cell RNA-seq by proposing scBatchProx, a post-hoc, federated-inspired framework that refines precomputed latent embeddings through batch-specific FiLM adapters without accessing raw data or retraining upstream models. The method treats each batch as a federated client, using proximal regularization (FedProx) to stabilize local updates and aggregating via FedAvg to align embeddings in latent space while preserving shared biology. Key contributions include a lightweight, deployable latent-space correction layer that supports cumulative retraining and continual training under dataset evolution, demonstrated to yield 3–8% gains in overall embedding quality with minimal runtime overhead. This approach enables practical refinement and incremental integration of evolving single-cell datasets, aligning distributed studies without centralized data sharing.

Abstract

Advances in single-cell RNA sequencing enable the rapid generation of massive, high-dimensional datasets, yet the accumulation of data across experiments introduces batch effects that obscure true biological signals. Existing batch correction approaches either insufficiently correct batch effects or require centralized retraining on the complete dataset, limiting their applicability in distributed and continually evolving single-cell data settings. We introduce scBatchProx, a post-hoc optimization method inspired by federated learning principles for refining cell-level embeddings produced by arbitrary upstream methods. Treating each batch as a client, scBatchProx learns batch-conditioned adapters under proximal regularization, correcting batch structure directly in latent space without requiring raw expression data or centralized optimization. The method is lightweight and deployable, optimizing batch-specific adapter parameters only. Extensive experiments show that scBatchProx consistently yields relative gains of approximately 3-8% in overall embedding quality, with batch correction and biological conservation improving in 90% and 85% of data-method pairs, respectively. We envision this work as a step toward the practical refinement of learned representations in dynamic single-cell data systems.

Figures (6)

• Figure 1: Framework of scBatchProx.
• Figure 2: Two dataset evolution scenarios in single-cell analysis. Cumulative retraining recomputes embeddings by retraining the upstream model on all available data, whereas continual training fits newly arriving datasets into a fixed reference embedding without retraining on previously processed data.
• Figure 3: Aggregate scIB scores for cumulative retraining under protocol arrival: (a) technological progression order (inDrop $\to$ CEL-Seq $\to$ CEL-Seq2 $\to$ Fluidigm C1 $\to$ Smart-seq2) and (b) reverse order (Smart-seq2 $\to$ Fluidigm C1 $\to$ CEL-Seq2 $\to$ CEL-Seq $\to$ inDrop). While cumulative retraining is supported by most existing methods, it is computationally expensive. scBatchProx consistently improves overall embedding quality across both arrival orders.
• Figure 4: Aggregate scIB scores for continual training under protocol arrival: (a) technological progression order (inDrop $\to$ CEL-Seq $\to$ CEL-Seq2 $\to$ Fluidigm C1 $\to$ Smart-seq2) and (b) reverse order (Smart-seq2 $\to$ Fluidigm C1 $\to$ CEL-Seq2 $\to$ CEL-Seq $\to$ inDrop). Continual training is commonly encountered in practice, yet most end-to-end batch-aware models do not support it. Despite this constraint, scBatchProx improves embedding quality across both arrival orders.
• Figure 5: Ablations on 2-batch PBMC (top) and HPMS (bottom) with the proximal coefficient set to $\mu = 0.0$ (no proximal regularization). Complete ablation results are provided in Appendix \ref{['app:complete_res']}. Removing proximal regularization degrades overall performance, primarily due to reduced biological conservation.