Adversarial Batch Representation Augmentation for Batch Correction in High-Content Cellular Screening
Authors: Lei Tong, Xujing Yao, Adam Corrigan, Long Chen, Navin Rathna Kumar, Kerry Hallbrook, Jonathan Orme, Yinhai Wang, Huiyu Zhou
Organizations: Data Sciences and Quantitative Biology, Discovery Sciences, AstraZeneca R&D, Cambridge, UK · School of Computing and Mathematical Sciences, University of Leicester, Leicester, UK · College of Computer and Information Engineering, Nanjing Tech University, Nanjing, China · Institute of Clinical Sciences, Faculty of Medicine, Imperial College London, London, UK
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.
Figures & tables
Figure 1: Motivational comparison of AdaBN, AdvStyle and ABRA. (a) AdaBN standardizes features by mapping them to historical bio-batch statistics (i.e., mean and standard deviation). (b) AdvStyle introduces instance-wise perturbations into the representation space, utilizing adversarial signals driven solely by classification likelihood. (c) ABRA explicitly perturbs the bio-batch representation space while jointly enforcing the classification objective and an angular geometry constraint to ensure robust and discriminative feature learning.
Figure 2: Framework of Adversarial Batch Representation Augmentation (ABRA). Given a bio-batch of cell painting images, a CNN encoder produces the corresponding representation X . The ABRA module computes batch channel-wise statistics μc(X) and σc(X) , and models uncertainty in the statistics space using learnable parameters {Kμ,Kσ} . Gaussian reparameterization instantiates the perturbations to transform the clean representation X into the perturbed representation Xt . A hybrid objective optimizes the framework: the cross-entropy loss LCE promotes inter-class separability, and the ArcFace loss Larc enforces intra-class compactness and inter-class separation through an additive angular margin. The stability terms RJS further align the discriminative distributions of clean and perturbed representations to reduce representation collapse. Optimization alternates between two phases: (1) maximize Ladv with respect to K while freezing θ ; (2) minimize Lrob with respect to θ to obtain a robust model.
Dataset
Cell Lines
Split (Train/Val/Test)
RxRx1
HEPG2, HUVEC, RPE, U2OS
7/ ∼ /4; 16/ ∼ /8; 7/ ∼ /4; 3/ ∼ /2
RxRx1-wilds
HEPG2, HUVEC, RPE, U2OS
7/1/3; 16/1/7; 7/1/3; 3/1/1
Table 1: Summary of datasets used in this study. Both datasets consist of a total of 51 biological batches. "Split Cases" refers to the number of training/validation/test batches per cell line. A semicolon separates different cell lines, and " ∼ " indicates the absence of a validation set.
Figure 3: Representative images from the RxRx1 dataset, depicting phenotypic variations in HEPG2 and HUVEC cell lines resulting from three siRNA perturbations across 5 experimental batches. Each perturbation elicits distinct changes in cell morphology, count, and spatial distribution.
Figure 4: Sample images from the RxRx1-wilds dataset, showcasing the phenotypic diversity elicited by siRNA perturbations in HEPG2 and HUVEC cell lines across five experimental batches.
Methods
Use of TTA
HEPG2
HUVEC
RPE
U2OS
Total
P-Values
ERM
✗
71.5±0.2
78.9±0.2
72.0±0.1
29.9±0.7
70.3±0.2
–
Simclr [ 51 ]
✗
18.3±2.2
35.6±3.4
28.2 ±1.4
5.8±1.5
26.8±2.4
4.3×10−7
BYOL [ 34 ]
✗
25.9±1.8
40.2±2.1
34.7±0.8
7.5±0.7
32.2±1.6
1.2×10−7
DinoV2 [ 35 ]
✗
43.3±1.1
55.0±1.2
44.6±0.5
12.1±0.9
45.3±1.0
6.7×10−8
DSU [ 23 ]
✗
76.7±0.4
79.9±0.2
71.2±0.5
37.8±0.5
72.6±0.3
2.0×10−7
AdvStyle [ 24 ]
✗
72.6±0.6
79.4±0.8
72.1±0.6
35.0±0.8
71.3±0.5
2.8×10−4
Table 2: Performance comparison on RxRx1. Results are reported as mean ± standard deviation over 5 random seeds. P-Values are computed on the Total metric using a two-sided paired t -test across the 5 seeds’ results. For methods without TTA, comparisons are made against ERM; for methods with TTA, comparisons are made against AdaBN.
Methods
Use of TTA
HEPG2
HUVEC
RPE
U2OS
Total
P-Values
ERM
✗
17.4±0.8
40.7±0.5
18.4±1.2
8.7±3.4
28.7±1.3
–
Simclr [ 51 ]
✗
7.2±0.7
13.9±0.3
6.4 ±1.2
3.0±0.9
10.1±0.6
1.1×10−7
BYOL [ 34 ]
✗
6.5±0.5
11.1±0.8
6.2±1.1
3.4±0.7
8.5±0.8
2.2×10−8
DinoV2 [ 35 ]
✗
8.4±0.5
17.3±0.4
10.4±0.9
5.5±0.5
13.0±0.5
3.7×10−7
DSU [ 23 ]
✗
20.9±0.4
46.1±0.2
21.4±0.4
10.2±0.1
32.9±0.3
2.9×10−4
AdvStyle [ 24 ]
✗
10.1±0.3
32.8±1.1
10.7±1.3
4.5±2.1
21.2±1.2
2.1×10−5
Table 3: Performance comparison on RxRx1-wilds out-of-distribution test set. Results are reported as mean ± standard deviation over 5 random seeds. P-Values are computed on the Total metric using a two-sided paired t -test across the 5 seeds’ results.
Methods
Use of TTA
HEPG2
HUVEC
RPE
U2OS
Total
P-Values
ERM
✗
22.0±0.7
47.7±0.6
23.8±0.9
15.2±2.8
34.1±0.9
–
Simclr [ 51 ]
✗
7.6±0.9
15.9±0.5
7.0±1.1
3.7±0.9
11.4±0.7
3.6×10−10
BYOL [ 34 ]
✗
6.8±0.7
12.3±1.1
6.3±1.1
4.3±0.5
9.3±1.0
1.6×10−11
DinoV2 [ 35 ]
✗
9.7±0.4
20.0±0.7
10.8±1.0
6.0±0.5
14.8±0.7
6.8×10−10
DSU [ 23 ]
✗
20.6±0.6
50.6±0.2
20.9±0.5
12.1±0.4
34.3±0.4
8.6×10−1
AdvStyle [ 24 ]
✗
12.5±0.4
40.5±1.3
12.1±1.1
6.7±2.5
25.4±1.2
9.1×10−8
Table 4: Performance comparison on RxRx1-wilds in-distribution test set. Results are reported as mean ± standard deviation over 5 random seeds. P-Values are computed on the Total metric using a two-sided paired t -test across the 5 seeds’ results.
Figure 5: Investigating the optimal location within the network architecture to apply ABRA. Designations such as ’res34’ indicate that ABRA is inserted after both the 3rd and 4th residual blocks of the ResNet-50 backbone.
Figure 6: Ablation of the hybrid loss components. ’CE’ refers to standard Cross-Entropy loss, ’Arc’ denotes the ArcFace margin loss, and ’JS’ represents the Jensen-Shannon divergence.
Methods
Model
Validation Accuracy
Test ID Accuracy
Test OOD Accuracy
ABRA (Ours)
ResNet-50
22.6±0.2
51.5±0.4
39.6±0.2
IID Representation Learning [ 29 ]
ResNet-50
23.9±0.3
49.9±0.5
39.2±0.2
ERM (CutMix) [ 29 ]
ResNet-50
23.6±0.3
47.4±1.0
38.4±0.2
LISA [ 56 ]
ResNet-50
20.1±0.4
41.1±1.3
31.9±1.0
ARM-BN [ 22 ]
ResNet-50
20.9±0.2
34.9±0.2
31.2±0.1
ERM (Grid Search) [ 27 ]
ResNet-50
19.4±0.2
35.9±0.4
29.9±0.4
Table 5: Comparison of ABRA against the official RxRx1-wilds Leaderboard. All methods utilize a ResNet-50 backbone and are evaluated without Test-Time Adaptation (TTA) to the strict leaderboard protocol.
Methods
TTA
Model
HEPG2
HUVEC
RPE
U2OS
Total
Gradient Reversal [ 11 ]
✗
DenseNet-161
74.0±0.6
83.8±0.2
78.1±0.5
24.3±0.7
71.2±0.4
ABRA (Ours)
✗
DenseNet-161
76.6±0.3
80.3±0.2
79.5±0.3
40.3±0.4
74.9±0.3
AdaBN [ 11 ]
✓
DenseNet-161
86.2±0.2
92.1±0.2
87.2±0.0
68.2±0.1
87.1±0.2
Gradient Reversal [ 11 ]
✓
DenseNet-161
85.6±0.1
82.0±0.0
87.5±0.0
66.9±0.3
86.2±0.3
ABRA (Ours)
✓
DenseNet-161
86.4±0.1
93.2±0.0
87.8±0.1
64.9±0.1
87.4±0.1
Table 6: Comparison of ABRA against the standard RxRx1 Leaderboard baselines. All methods utilize a DenseNet-161 backbone. Results highlight performance under both standard inductive (w/o TTA) and adaptive (w/ TTA) settings.
Figure 7: Comparison of classification accuracy across varying inference mini-batch sizes on (a) RxRx1-wilds and (b) RxRx1. The violin plots (generated across 10 independent runs per batch size) illustrate the performance distribution for TTA-based methods, which suffer from high variance and degradation at small batch sizes due to noisy statistical estimation. Horizontal dashed lines represent the performance of non-TTA baselines, which remain strictly invariant to the test batch size.
Figure 8: Layer-wise Batch Normalization (BN) statistics of ABRA (w/ TTA), measured by the KL divergence (top) and MMD distance (bottom) between the source and inference batches. The red curve represents the standard RxRx1 dataset, while the blue curve represents the RxRx1-wilds dataset. Note the logarithmic scale on the y-axis, illustrating the difference in statistical domain shifts.
Figure 9: UMAP visualization of the learned embedding spaces on the RxRx1 dataset for ABRA, AdvStyle, and ERM, alongside their respective Test-Time Adaptation (TTA) variants. Points represent instance-level image embeddings from U2OS experimental batches across both training and test sets. Points are colored by their specific experimental batch. (Similar alignment behaviors are observed across other cell types).
Figure 10: UMAP visualization of the learned embedding spaces on the RxRx1-wilds dataset. Points represent instance-level image embeddings from U2OS experimental batches, colored by experimental batch to highlight domain alignment efficacy.
ELLIS Unit Linz and LIT AI Lab, Institute for Machine Learning, Johannes Kepler University, Linz, Austria · Clinical Research Institute for Medical AI, Johannes Kepler University, Linz, Austria
ELLIS Unit Linz, LIT AI Lab and Institute for Machine Learning, JKU Linz, Austria · Institute of Artificial Intelligence, Center for Medical Data Science, Medical University Vienna · Clinical Research Center for Medical AI, JKU Linz, Austria