Adversarial Batch Representation Augmentation for Batch Correction in High-Content Cellular Screening

TL;DR

This paper introduces ABRA, an adversarial augmentation method for batch correction in high-content cellular screening, improving model generalization across diverse experimental batches.

Contribution

ABRA models batch effects as structured uncertainties and synthesizes worst-case perturbations to enhance domain generalization in cellular image analysis.

Findings

ABRA achieves state-of-the-art results on RxRx1 benchmarks.

It effectively mitigates bio-batch effects without prior knowledge.

The method preserves class discriminability during adversarial augmentation.

Abstract

High-Content Screening routinely generates massive volumes of cell painting images for phenotypic profiling. However, technical variations across experimental executions inevitably induce biological batch (bio-batch) effects. These cause covariate shifts and degrade the generalization of deep learning models on unseen data. Existing batch correction methods typically rely on additional prior knowledge (e.g., treatment or cell culture information) or struggle to generalize to unseen bio-batches. In this work, we frame bio-batch mitigation as a Domain Generalization (DG) problem and propose Adversarial Batch Representation Augmentation (ABRA). ABRA explicitly models batch-wise statistical fluctuations by parameterizing feature statistics as structured uncertainties. Through a min-max optimization framework, it actively synthesizes worst-case bio-batch perturbations in the representation space, guided by a strict angular geometric margin to preserve fine-grained class discriminability. To prevent representation collapse during this adversarial exploration, we introduce a synergistic distribution alignment objective. Extensive evaluations on the large-scale RxRx1 and RxRx1-WILDS benchmarks demonstrate that ABRA establishes a new state-of-the-art for siRNA perturbation classification.