Scalable Batch Correction for Cell Painting via Batch-Dependent Kernels and Adaptive Sampling
Aditya Narayan Ravi, Snehal Vadvalkar, Abhishek Pandey, Ilan Shomorony
TL;DR
This work introduces BALANS, a scalable batch-correction method for Cell Painting data that integrates batch-aware local kernel scaling with adaptive landmark sampling. By constructing a sparse affinity operator from a small, coverage-proportionate subset of points and applying a low-rank smoothing, BALANS achieves near-linear runtime while preserving biological signal across highly heterogeneous batches. Theoretical results guarantee coverage across biological clusters and provide a spectral error bound for the recovered affinity, underpinning empirical gains. Extensive experiments on real JUMP and BBBC datasets, plus synthetic benchmarks, demonstrate BALANS's superior or competitive performance relative to existing methods with substantial scalability benefits, enabling robust batch correction at multi-million-point scales.
Abstract
Cell Painting is a microscopy-based, high-content imaging assay that produces rich morphological profiles of cells and can support drug discovery by quantifying cellular responses to chemical perturbations. At scale, however, Cell Painting data is strongly affected by batch effects arising from differences in laboratories, instruments, and protocols, which can obscure biological signal. We present BALANS (Batch Alignment via Local Affinities and Subsampling), a scalable batch-correction method that aligns samples across batches by constructing a smoothed affinity matrix from pairwise distances. Given $n$ data points, BALANS builds a sparse affinity matrix $A \in \mathbb{R}^{n \times n}$ using two ideas. (i) For points $i$ and $j$, it sets a local scale using the distance from $i$ to its $k$-th nearest neighbor within the batch of $j$, then computes $A_{ij}$ via a Gaussian kernel calibrated by these batch-aware local scales. (ii) Rather than forming all $n^2$ entries, BALANS uses an adaptive sampling procedure that prioritizes rows with low cumulative neighbor coverage and retains only the strongest affinities per row, yielding a sparse but informative approximation of $A$. We prove that this sampling strategy is order-optimal in sample complexity and provides an approximation guarantee, and we show that BALANS runs in nearly linear time in $n$. Experiments on diverse real-world Cell Painting datasets and controlled large-scale synthetic benchmarks demonstrate that BALANS scales to large collections while improving runtime over native implementations of widely used batch-correction methods, without sacrificing correction quality.