Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
Ana Sanchez-Fernandez 1 Thomas Pinetz 2 Werner Zellinger 1 Günter Klambauer 1 3
in-domain
Abstract 1.00
MoA Classification Accuracy
The central problem in biomedical imaging are batch effects: systematic technical variations unrelated to the biological signal of interest. These batch effects critically undermine experimental reproducibility and are the primary cause of failure of deep learning systems on new experimental batches, preventing their practical use in the real world. Despite years of research, no method has succeeded in closing this performance gap for deep learning models. We propose ControlStabilized Adaptive Risk Minimization via Batch Normalization (CS-ARM-BN), a meta-learning adaptation method that exploits negative control samples. Such unperturbed reference images are present in every experimental batch by design and serve as stable context for adaptation. We validate our novel method on Mechanismof-Action (MoA) classification, a crucial task for drug discovery, on the large-scale JUMP-CP dataset. The accuracy of standard ResNets drops from 0.939 ± 0.005, on the training domain, to 0.862 ± 0.060 on data from new experimental batches. Foundation models, even after Typical Variation Normalization, fail to close this gap. We are the first to show that meta-learning approaches close the domain gap by achieving 0.935 ± 0.018. If the new experimental batches exhibit strong domain shifts, such as being generated in a different lab, meta-learning approaches can be stabilized with control samples, which are always available in biomedical experiments. Our work shows that batch effects in bioimaging data can be effectively neutralized through principled in-context adaptation, which also makes them practically usable and efficient.
arXiv:2604.20824v1 [cs.LG] 22 Apr 2026
new domain
0.95
0.939
0.935
domain gap 0.893
0.90
0.862
0.85
= 0.077
0.80 0.747
0.75 0.70
ResNet
No adapt
FM+ TVN
CORAL
ARM-BN
Figure 1. Performance of MoA classifier on JUMP-CP data. Error bars represent variance across five cross-validation folds. Green bar: within the training domain, the performance of the classifier is high. Orange bars: The performance of the classifier on images from new experimental batches ("new domain"). Even foundation models with normalization (FM+TVN) suffer performance declines, and domain adaptation methods like CORAL only slightly counter the domain shift. Only the meta-learning method ARM-BN closes the domain gap. Details in Tab. 1
1. Introduction Batch effects severely degrade model generalization for biomedical datasets. The growing scale of publicly available biomedical imaging datasets offer an unprecedented opportunity for deep learning models to predict clinically and biologically meaningful properties, such as disease states, treatment responses, and drug effects (Chandrasekaran et al., 2023; Zhang et al., 2025; Peidli et al., 2023). However, a key challenge arises because biological data has to be acquired in experimental batches, which are groups of samples that were obtained under different technical conditions. Even in carefully controlled environments, images from different batches present variations due to non-biological factors that might be more prominent than the biological signal in question (Leek et al., 2010; Sypetkowski et al., 2023; Luecken et al., 2021). These variations are called batch effects and are a well-known artifact, such that unperturbed samples, known as negative controls, are included in every experimental batch as a reference (Bray et al., 2017). Negative controls have been long used in the field of cellu-
1
ELLIS Unit Linz, LIT AI Lab and Institute for Machine Learning, Johannes Kepler University Linz, Austria 2 Medical University Vienna 3 Clinical Research Center for Medical AI, Johannes Kepler University Linz, Austria. Correspondence to: Günter Klambauer <[email protected]>. Preprint. April 23, 2026.
1
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
MoA 1
MoA 2
MoA 3
Control
Figure 2. Representation of batch effects in microscopy imaging data considered as a multi-source domain adaptation (MSDA) problem. In this setting, each source consists of different experimental conditions, e.g. different plates. In each domain, an image of the same class is depicted, which represents a particular mechanism-of-action (MoA). Control samples are unperturbed samples, that are present in every domain, and which we represent with the c symbol.
lar imaging profiling to try to disentangle biological signal from confounding factors, often by standardizing features using controls from the corresponding experimental batch (Arevalo et al., 2024). However, such post-hoc corrections do not allow exploiting the capabilities of end-to-end deep learning models, and traditional deep learning models perform poorly on new batches (Chen et al., 2020; Kim et al., 2025).
stead as a multi-source domain adaptation (MSDA) problem (see Figure 2), where each experimental batch is a source domain (Farahani et al., 2020; Blanchard et al., 2011). Multi-source domain adaptation and in-context test-time adaptation. Recent works have been developed to extend classical UDA methods to a multi-source setting (Zhao et al., 2018; Peng et al., 2019; Wen et al., 2020; Yang et al., 2020; Venkat et al., 2020; Guo et al., 2018). Nevertheless, just like UDA, MSDA methods require re-training the model every time a new target domain needs to be evaluated. In microscopy imaging data, new experimental batches come in sequentially, making these methods highly inefficient. An ideal setting for microscopy imaging data would be one where the multi-source nature of the data can be exploited and also the model can be dynamically adapted at test time without need for retraining. Adaptive Risk Minimization (ARM) (Zhang et al., 2021) is a meta-learning framework (Hochreiter et al., 2001; Finn et al., 2017) that enables such rapid adaptation.
Domain adaptation provides a natural framework for batch effect correction. We could consider batch effect correction in machine learning as an unsupervised domain adaptation (UDA) problem, where the source domain is the training data available and the target domain is a new experimental batch. Classical UDA aligns source and target distributions using approaches such as a) importance weighting (Shimodaira, 2000; Ben-David et al., 2006; Kimura and Hino, 2024), b) distance-based matching (Sun et al., 2015; Zellinger et al., 2017), c) adversarial training (Ganin and Lempitsky, 2014), or d) generative translation (Liu and Tuzel, 2016; Taigman et al., 2016; Bousmalis et al., 2017; Chung et al., 2024). In the field of microscopy imaging data, test-time training (TTT) strategies have been explored (Haslum et al., 2023). However, traditional UDA and TTT methods assume a small number of source domains, each of them with a large amount of samples, and require retraining the model every time a new domain needs to be evaluated. In contrast, microscopy imaging data is comprised of many domains with few samples (Chandrasekaran et al., 2023; Bray et al., 2017), which makes UDA methods unsuitable in this setting. Batch effect correction can be formalized in-
Batch Normalization as a domain adaptation mechanism. Batch Normalization (BatchNorm) is a widely established normalization technique, that allowed training deep neural networks with skip connections (Ioffe and Szegedy, 2015), and that can additionally be used as a module for domain adaptation (Li et al., 2016). Conventionally, when the assumption that both training and test data are sampled from the same distribution is made, the stored training statistics are used at inference time. In contrast, when the test samples are known to come from a different distribution, adapting 2
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
2. Problem setting and Background
these layers by using the statistics of the new domain at inference time provide a straight-forward test-time domain adaptation method (Li et al., 2016; Wang et al., 2021). However, it has also been shown that these methods fail when not enough samples are available to accurately estimate the new domain statistics or when the class distribution is shifted at test time (Zhao et al., 2023; Park et al., 2023).
We consider the problem of adapting from B source batches of size N to a single target batch PKof size M under the crossentropy loss LCE (y, ŷ) = − k=1 yk log(ŷk ). This problem can be extended to any loss function L, many target batches and different batch sizes, in a straight forward way. We a source dataset Dtrain = (1) are given N (B) D ,...,D ⊂ (X × Y) consisting of B ex(b) (b) perimental batches D(b) = (xn , yn )N n=1 drawn from distributions pbxy , b ∈ {1, . . . , B}, respectively. Following Baxter (1998); Zhang et al. (2021), we assume that the batches are independent and drawn from a meta-distribution µ.
Control samples provide an opportunity to stabilize the context in meta-learning methods. In addition to the methods presented above, there are also meta-learning methods which include adapting the batch normalization layers in a meta-learning setting (Zhang et al., 2021). The main difference between BatchNorm adaptation (Li et al., 2016) and meta-learning BatchNorm adaptation (Zhang et al., 2021) is that the adaptation does not only occur at test-time but also during training, such that the model learns to adapt to new batches. This has been shown to provide a significant improvement with respect to its pure test-time counterparts. Nevertheless, in this work, we show that these methods also underperform in the presence of small batch sizes and under label shifts. Both scenarios, a) small sample sizes and b) labels shifts in a new experimental batch are very plausible, because of the typical design of such bioimaging studies (Hughes et al., 2011; Knowles and Gromo, 2003). However, almost all biomedical experiments guarantee that a specific type of samples will be present in every new experimental batch (i.e., domain): negative controls (Chandrasekaran et al., 2023; Bray et al., 2017). Therefore, we propose to exploit the presence of negative controls and use them as context samples both during training and for inference of a meta-learning model, showing that this approach makes the performance of our method robust to small numbers of new experimental batches and to shifts in the class distribution.
The goal is to find a prediction model g(x, w) = ŷ together e that with an adaptation method Adapt w, g, Dβ = w allow us, for any new unlabeled experimental target batch e Dβ = (xβn )M n=1 , to adapt the weights w to new weights w with a small target batch risk M 1 X e , ℓ ym , g (xm , w) M m=1
(1)
whenever the target batch is drawn from the marginal pβx of a new target distribution from µ. ℓ is an appropriate loss function. Feature extractor and classifier networks. We denote the following components that are normally used by domain adaptation methods: a) an encoder or feature extractor network f : X 7→ RD or fw with weights w, which maps the input images x ∈ X to an embedding a ∈ RD , and b) a classifier c : RD 7→ Y, which maps the embeddings to a class label ŷ ∈ Y in the output space Y. The model g : X 7→ Y is usually a composition of the feature extractor with the classifier g = c ◦ f : X 7→ Y.
Contributions. Our main contributions are the following: • We propose CS-ARM-BN, a control-stabilized metalearned BatchNorm adaptation method that leverages negative control samples available in every experimental batch, and show that it robustly alleviates the failure modes of previous methods and nearly closes the generalization gap between in-domain and unseen-batch performance.
2.1. Domain Adaptation Methods for Biomedical Images We first introduce a set of domain adaptation methods for biomedical images that will be compared to our work and mentioned throughout the manuscript. Supervised baselines. ResNet (He et al., 2015) and Foundation Model CA-MAE (Kraus et al., 2024). We train a standard classifier g on Dtrain by minimizing
• We demonstrate that adaptive Batch Normalization– based methods, particularly when combined with metalearning and in-context adaptation, provide an effective and computationally efficient approach to mitigate severe batch effects in biomedical microscopy images.
min w
• We systematically analyze and expose key limitations of existing BatchNorm-based adaptation methods, including AdaBN and ARM-BN, showing that their performance degrades in realistic settings such as label shift.
B X
X
LCE (y, g(x, w))
(2)
b=1 (x,y)∈D (b)
with cross-entropy loss LCE and using a ResNet50 architecture with either BatchNorm (BN) or InstanceNorm (IN) layers, and a linear layer as a classifier. This classifier is not adapted to new experimental batches. 3
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
ARM-BN
Training
batch 1
Inference
a. Balanced data
batch 1
Training
Image encoder
Image encoder
b.b.Small Smallbatch batchsizesize ... batch B
c.c.Label Labelshift shift
Image encoder
CS-ARM-BN (ours)
Inference
a. Balanced data
Image encoder
Image encoder
b. Small batch size Image encoder
Image encoder
... batch B Image encoder
Image encoder
c. Label shift Image encoder
Figure 3. Graphic representation of our method, CS-ARM-BN, and comparison to ARM-BN (Zhang et al., 2021). Both are meta-learning methods that are be modified at test-time by using the BN statistics from the target domain (lilac). CS-ARM-BN uses control samples both at training and at inference time, which provides stability when (b) the number of perturbed samples is small or (c) the label distribution is shifted.
2017). As in the foundation model baseline, we use the embeddings from a frozen CA-MAE encoder, and perform a post-hoc correction with TVN. For TVN, given a set of batch-β control embeddings from the target batch, {zjβ }, µctrl,β and Σctrl,β are estimated and used to apply a perbatch whitening: − 1 (5) ã = Σctrl,β 2 a − µctrl,β
For the foundation model baseline, we use a frozen CAMAE encoder fMAE to obtain embeddings a = fMAE (x) and perform linear probing with a classifier cθ trained with CE-loss. Domain-Adversarial Neural Networks (DANN) (Ganin and Lempitsky, 2014). DANN uses a domain discriminator dv : RD → S B−1 over features a = fw (x) and use gradient reversal to promote domain invariance: min max w
v
B X
is performed. Then, the covariance is aligned to the covariance of the training data Σtrain . The classifier cθ is then trained with the resulting re-aligned representations ã. This matches first and second moments to the typical, i.e. control, distribution of each batch.
h LCE (y, g(x, w)) +
X
b=1 (x,y)∈D (b)
i λ LCE b, dv (a) ,
(3)
Adaptive Batch Normalization (AdaBN) (Li et al., 2016). In Adaptive Batch Normalization, the adaptation of the model is simply carried out by using the BN statistics of the target domain instead of those of the training data. For microscopy images Lin and Lu (2022) proposed adapting the batchnorm statistics based on control samples.
where the second loss term is cross-entropy on domain labels and λ > 0 trades classification performance against domain objectives. Correlation Alignment (CORAL) (Sun et al., 2015). Let AS = {ai } and AT = {a′j } be source and target features in a training batch; let CS and CT denote their covariance matrices. CORAL adds a moment-matching penalty to the supervised loss: min w
B X
X
Fully Test-Time Adaptation by Entropy Minimization (TENT) (Wang et al., 2021). In TENT, the adaptation is also carried out at test-time, and, as in AdaBN, the BN statistics are computed using samples from the new domain. However, in this method an additional adaptation is included, consisting in minimizing the entropy H of the model predictions ŷ by updating only the BatchNorm affine parameters γ and β :
LCE (y, g(x, w)) +
b=1 (x,y)∈D (b)
γ ∥CS − CT ∥2F ,
(4)
aligning second moments of source and target features during training. We trained CORAL on considering the samples in the training set to stem from the source distribution and the samples in the test set, from the target distribution.
ŷ = g(x),
H(ŷ) = −
K X
ŷk log(ŷk )
(6)
k=1
Channel-Agnostic MAE (CA-MAE) + Typical Variation Normalization (TVN) (Kraus et al., 2024; Ando et al.,
min H(ŷ) γ,β
4
(7)
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
Generative methods: StyleID (Chung et al., 2024). We used the diffusion-based style transfer method proposed by Chung et al. (2024) to convert all images to the style of one specific experimental batch (Liu and Tuzel, 2016). Then, as in the supervised baselines, a ResNet50 and a linear classifier are trained to minimize the cross-entropy loss LCE in the generated training samples.
of available samples is small or when the class distribution at test time differs from that observed during training, the resulting batch normalization (BN) statistics become noisy and confounded by label shift (Park et al., 2023; Zhao et al., 2023), leading to unstable or degraded performance. We confirm this occurrence in our experiments shown on Sections 5.3 and 5.4 and Appendix Sections D.1 and D.2.
Adaptive Risk Minimization (ARM) / In-context adaptation (Zhang et al., 2021).
In high-content screening experiments, and generally in biomedical experiments, negative control samples are available (Chandrasekaran et al., 2023; Arevalo et al., 2024). ARM uses a meta-learning objective (Hochreiter et al., 2001; For every experimental batch β, a set of control images Finn et al., 2017): {zcβ }C c=1 is acquired under the same technical conditions as the perturbed samples, but without inducing any biological B X 1 X 1 b effect. These controls provide a direct estimate of batchmin LCE ym , g(xm ; hϕ (w, x1:M ) , w,ϕ B M specific technical variation that is independent of the class b b=1 (xm ,ym )∈DM composition of the perturbed samples. (8) Control-stabilized adaptation. In CS-ARM-BN, we extend ARM-BN by incorporating negative control samples into the context used for adaptation. Instead of estimating BatchNorm (BN) statistics solely from the unlabeled perturbed samples xβ1:L , we compute the adaptation statistics using the union of control and perturbed samples from the target batch,
b where DM is a random subset of M < N pairs of batch b, considered as context set. M is usually in the size of the expected future dataset from the new domain, in our case the size of an experimental batch.
The central part is the test-time adaptation function hϕ : W × X M → W,
wβ = hϕ w, xβ1:M ,
(9)
β L C β = {zcβ }C c=1 ∪ {xl }l=1 .
which maps the current model parameters w and context (the target-batch samples xβ1:M ) to the adapted parameters wβ . The test-time adaptation function hϕ has itself parameters ϕ, which are learned with the usual gradient descent techniques. At test time on a new experimental batch β, we first compute wβ = hϕ (w, xβ1:M ) and then predict with ŷ = g(x; wβ ) for any x from batch β.
Concretely, for each BN layer, the adapted mean and variance are computed over this combined context set as µβBN =
1 X u, |C β | β
(10)
2 1 X u − µβBN , β |C | β
(11)
u∈C
β σBN =
ARM-BN. In this variant of ARM the adaption function hϕ only replaces the batch normalization statistics µBN , σBN of all batchnorm-layers of the main network g with the statistics from the batch β from the new domain.
u∈C
where u ∈ C β denotes the BN-layer activations corresponding to either a control image zcβ or a perturbed image xβm . These adapted statistics replace the running training statistics in all BN layers of the prediction model g when evaluating samples from batch β.
ARM-CML. In this version, the adaptation function hϕ uses the unlabeled samples from the new batch to compute a compact representation, a context vector, that captures the specific characteristics of the target domain. Unlike ARM-BN, which only updates normalization statistics, ARM-CML performs a lightweight parameter adaptation conditioned on the new data, enabling the model to better align with the target domain without requiring labels or retraining.
Another option to mitigate the effects of label shift and small batch sizes could be to use only control samples to adapt the Batch Normalization statistics, as proposed by Batch Effect Normalization (BEN) (Lin and Lu, 2022). However, in the following paragraph, we argue that including both perturbed and control samples provide the best trade-off between bias and variance. Empirically, we found that adapting BN layers using control samples only at evaluation time performed poorly, but when only controls were already included during training, the performance improved (see "BEN w/o meta-learning" and "ARM-BEN" in Table 4).
3. CS-ARM-BN: Control-Stabilized Adaptive Risk Minimization with Batch Normalization A key limitation of BN-based test-time and meta-learned adaptation methods is their reliance on statistics estimated from a set of unlabeled target samples. When the number
Let M = C + L be the total number of samples in a new target batch, where C is the number of control samples and 5
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
L are the non-control or perturbed samples. We model the BN activation mean for sample x as:
Algorithm 1 Pseudo code for CS-ARM-BN vs traditional image classifier inference pipelines
µobs,x = µdomain + µclass(x)
### Standard Image Classifier Pipeline ### # no calibration of classifier to data x = load_data() # your biomedical images model = load_model() # public classifier pred = model.predict(x) # domain gap!
where µdomain is the batch-specific technical offset we wish to estimate, and µclass(x) = 0 for controls and µclass(x) ̸= 0 for perturbed samples.
### CS-ARM-BN Classifier Pipeline ### # adapt/calibrate classifier to your data x = load_data() # your biomedical images z = load_controls() # available controls model = load_model() # public classifier model.adapt(concat([x, z])) # calibrate pred = model.predict(x) # neutralized gap
The three estimators of µdomain and their mean squared errors (MSEs) are: 1. ARM-BN or AdaBN estimator (perturbed samples only): σ2 MSE(µ̂ARM-BN ) = µ̄2class + | {z } L
batch b, a control samples and perturbed samples are used as context to adapt BN statistics, and performance is optimized on perturbed query samples from the same batch. This trains the base network to expect and exploit controlconditioned normalization at test time, rather than treating it as a heuristic post-hoc correction.
(12)
large bias2
The bias µ̄class grows with the degree of label shift, explaining the performance collapse observed in Tables 3 and 4 .
Compared to standard ARM-BN and other domain adaptation methods, CS-ARM-BN offers three key advantages:
2. Control-based estimator (controls only): MSE(µ̂ctrl ) = |{z} 0 + no bias
σ2 C
(13)
• Stability under small context sizes: since all samples, perturbed plus control samples are used in each batch, the BN statistics can be estimated reliably even when only few perturbed samples are available.
The estimator is unbiased, but suffers from high variance when C, the number of controls, is small.
• Robustness to label shift: control samples are independent of the class distribution in the new domain, and are used to decrease the bias in the estimator at test time.
3. Our estimator CS-ARM-BN (controls + perturbed samples): CS-ARM-BN
MSE(µ̂
)= |
L M
2
σ2 µ̄2class + M {z }
(14) • No additional supervision or re-training: as in ARM-BN, adaptation remains fully unlabeled and lightweight, requiring only a single forward pass over the samples from the new domain (see Algorithm 1 and 2), i.e. experimental batch, and no expensive test-time training.
small bias2
exploits all available samples M = C + L and allows for a small bias, but decreases variance by a factor of 1/M . In typical settings, CS-ARM-BN achieves the lowest MSE by trading a small residual bias, shrunk by the factor (L/M )2 relative to ARM-BN, against a reduced variance, since M > C. This bias-variance tradeoff is most favorable precisely in the operating regime of real drug screening: availability of few controls (C) and moderate-to-severe label shift, which is confirmed empirically in Table 4. Details in Appendix F.
4. Related work Traditional batch effect correction methods. Batch effect correction methods that are based based on either linear transformations (Johnson et al., 2007; Ando et al., 2017), nearest neighbors (Haghverdi et al., 2018; Stuart et al., 2019; Hie et al., 2019) or a mixture of PCA and clustering (Korsunsky et al., 2019) have been previously proposed. One of the closest related works is Arevalo et al. (2024), who investigate batch correction methods for handcrafted features for microscopy images. These methods, however, have been originally designed to correct batch effects in feature
Relation to ARM. Importantly, our CS-ARM-BN method remains fully within the Adaptive Risk Minimization (ARM) framework. As shown in Figure 3, during training, the model is meta-learned episodically: for each source 6
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples Table 1. MoA prediction accuracy on in-domain and unseen experimental batches from a new domain. Column "TTA" indicates that the method can adapt itself given samples from a new domain at test time, i.e. test time adaptation. “In-context” means that the model adapts its parameters based on the test samples through in-context learning. The accuracy values denote mean and standard deviation over 5-fold cross-validation, thus five training re-runs of each method. "FM" indicates a foundation model.
Method category
Method
In-domain accuracy
New experimental batch accuracy
TTA
In-context
Baselines
no adaptation (BatchNorm) no adaptation (InstanceNorm) CA-MAE (FM) (Kraus et al., 2024) CellProfiler (Stirling et al., 2021)
0.939 ± 0.005 0.949 ± 0.007 0.747 ± 0.006 0.923 ± 0.031
0.862 ± 0.060 0.891 ± 0.043 0.709 ± 0.041 0.860 ± 0.032
✗ ✗ ✗ ✗
✗ ✗ ✗ ✗
Adversarial
DANN (Ganin and Lempitsky, 2014)
-
0.911 ± 0.020
✗
✗
CORAL (Sun et al., 2015) CA-MAE (FM) + TVN (Ando et al., 2017) TENT (Wang et al., 2021) AdaBN (Li et al., 2016)
– – – –
0.893 ± 0.054 0.747 ± 0.013 0.926 ± 0.014 0.928 ± 0.013
✗ ✓ ✓ ✓
✗ ✗ ✗ ✗
Style-ID (Chung et al., 2024)
–
0.781 ± 0.015
✓
✗
ARM-CML (Zhang et al., 2021) ARM-BN (Zhang et al., 2021) CS-ARM-BN (ours)
– – –
0.843 ± 0.074 0.935 ± 0.018 0.930 ± 0.019
✓ ✓ ✓
✓ ✓ ✓
Distance
Generative Meta-learning
5. Experiments and Results
vectors and cannot be used together with a model trained end-to-end directly from images.
We performed a set of experiments to evaluate the effectiveness of a wide selection of methods for addressing batch effects under different scenarios: a) Mild domain shift: new experimental batches. We evaluate adaptation performance when the target data originate from previously unseen batches generated under similar experimental conditions, reflecting routine batch-to-batch variability. b) Strong domain shift: transfer across sources. We consider a more challenging setting in which the target data come from different institutions or data-generating sources, inducing substantial domain shifts due to variations in protocols, equipment, and experimental practices. c) Realistic scenario 1: new experimental batch with label shifts. We study the impact of changes in class proportions, evaluating how well methods maintain performance when label distributions differ at inference time. This is the most realistic scenario since, novel sets of compounds are usually screened, which do not maintain the label distribution of the training data. One obvious example are drug discovery efforts, where the compounds that are tested in the lab are suspected to be active towards one specific protein target. d) Realistic scenario 2: new imaging source with label shifts. Finally, we evaluate the performance of our method and the baselines in a setting where test samples exhibit both class distribution shift and domain shift, but are additionally acquired from a new imaging source, yielding a more challenging scenario than the previous one.
Deep learning and batch effects in microscopy imaging data. While Lin and Lu (2022) also proposed using negative controls to estimate BN statistics, their work is not framed as meta-learning, and they do not show the scenarios where control-based statistics are beneficial in comparison to using only perturbed samples. Palma et al. (2025); Lopez et al. (2018); Chung et al. (2024) presented generative models that could be used to address batch effects. Due to computational constrains, we use Chung et al. (2024) as representative for this class of methods. In-context learning for biological data. Recently, in the field of single cell transcriptomic data, Dong et al. (2026) introduces a foundation model that can perform in-context predictions from unlabeled cells. However, their work focuses on tabular data and the transformer architecture. Test-time adaptation and robustness to label shift. Aside from TENT (Wang et al., 2021), additional test-time adaptation methods have been developed since. Wang et al. (2022); Gong et al. (2022); Boudiaf et al. (2022) propose methods to overcome catastrophic forgetting in online TTA, while our proposed method does not consist on continually adapting the model weights. Park et al. (2023) introduced a hypernetwork to address label shift, but it needs additional training of the model and it is not specifically tailored for biological data where control samples are available.
Dataset. We conduct experiments on the large JUMP-CP dataset (Chandrasekaran et al., 2023). We use source 3, and
7
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples Table 2. Comparison of MoA classification accuracy of different methods across target sources. The element before the arrow indicates the training domain, the element after the arrow indicates the evaluation domain. The accuracy values denote the mean and standard deviation over 3 re-runs with different seeds.
Source → Target
Method S3 → S3
S3 → S8
S8 → S8
S8 → S3
0.939 ± 0.005 0.949 ± 0.007 0.747 ± 0.006 0.923 ± 0.031
0.670±0.016 0.796±0.018 0.389±0.015 0.430±0.074
0.974 ± 0.002 0.984 ± 0.001 0.840 ± 0.004 0.991 ± 0.004
0.498±0.011 0.552±0.014 0.356±0.011 0.550±0.029
DANN
–
0.639±0.175
–
0.324±0.079
CORAL CA-MAE (FM) + TVN TENT AdaBN
– – – –
0.804±0.018 0.688±0.005 0.813±0.041 0.838±0.037
– – – –
0.697±0.010 0.488±0.010 0.768±0.001 0.787±0.002
StyleID
–
0.774±0.004
–
0.713±0.016
ARM-CML ARM-BN CS-ARM-BN (ours)
– – –
0.645±0.034 0.874±0.008 0.884±0.005
– – –
0.569±0.030 0.795±0.002 0.776±0.016
no adaptation (BatchNorm) no adaptation (InstanceNorm) CA-MAE (FM) CellProfiler
in a strictly unlabeled test-time adaptation setting unless explicitly stated.
source 8 in Exp. 5.2, and perform a classification task to predict eight well-defined MoAs. Therefore, we work with samples perturbed with eight different compounds, each one of them known to produce one of these MoAs (see Appendix Table A6), and which are present in every plate of the dataset. Each microscopy image has five channels; we use the preprocessing described in Appendix Section C.
5.1. Mild domain shift: new experimental batches within a source We first study the most common practical scenario: training and test data originate from the same imaging source, but the test batches correspond to new experimental batches. This setting reflects mild domain shift: acquisition conditions drift across experiments, but the imaging platform remains unchanged. Results are summarized in Table 1.
Data splits. To evaluate robustness to batch effects, we consider two data splitting strategies reflecting different levels of domain shift (see Appendix Section D.4). In most experiments, we adopt an experimental-batch split, where training and evaluation are performed on disjoint experimental batches, representing a challenging but realistic batch generalization setting. For the strong domain shift scenario, we instead evaluate on images acquired from a completely unseen imaging source, corresponding to a cross-platform or cross-institution transfer setting. Under the experimentalbatch split, we use 5-fold cross-validation and report mean accuracy ± standard deviation across folds, corresponding to five independent training runs per method. Unless otherwise specified, accuracy is reported at the image level.
Batch effects degrade generalization performance. A standard ResNet50 achieves high in-domain accuracy (e.g., 0.939 ± 0.005 with BN), but drops substantially when evaluated on unseen experimental batches (0.862 ± 0.060). This gap quantifies the severity of batch effects even under mild shifts. Classical UDA baselines are insufficient or impractical. Domain-invariant learning (DANN) and discrepancy/moment matching (CORAL) do not reliably close the gap in this few-shot multi-source regime. Moreover, these approaches conceptually assume access to the target distribution during training and are thus inefficient in sequential settings where new batches continually arrive.
Methods and implementation. We compare supervised baselines (ResNet50 with BN/IN; CA-MAE linear probe; CellProfiler features), classical UDA/MSDA baselines (DANN, CORAL), generative style transfer (Style-ID), and test-time adaptation / in-context methods (AdaBN, TENT, ARM-CML, ARM-BN). Training details and hyperparameter search spaces are provided in Appendix Table A5. For methods that require test-time samples to adapt (AdaBN, TENT, ARM-BN, ARM-CML), we evaluate them
Test-time and in-context adaptation close the gap. Methods that adapt using the unlabeled target batch at inference time provide the strongest gains. In particular, adapting normalization statistics (AdaBN) and entropy-minimizing BN adaptation (TENT) markedly improve robustness. Meta8
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples Table 3. Comparison of MoA classification accuracy of different methods in a new experimental batch under different degrees of label shift. A smaller α value corresponds to a larger degree of label shift.
learned in-context approaches further improve performance: ARM-BN achieves the best accuracy (Table 1), approaching the in-domain baseline and thereby largely neutralizing mild batch effects. This supports our hypothesis that, in microscopy, batch effects can often be corrected by fast, context-conditioned adaptation rather than retraining.
α=1
Label shift degree α = 0.1
α = 0.01
ERM (no adaptation)
0.862 ± 0.060
0.862 ± 0.061
0.862 ± 0.061
TENT AdaBN ARM-BN CS-ARM-BN (ours)
0.832 ± 0.014 0.851 ± 0.016 0.831 ± 0.017 0.924 ± 0.018
0.778 ± 0.016 0.452 ± 0.010 0.416 ± 0.020 0.906 ± 0.021
0.156 ± 0.010 0.229 ± 0.009 0.228 ± 0.009 0.894 ± 0.022
Method
5.2. Strong domain shift: transfer across sources Next, we evaluate whether the above conclusions hold under strong domain shift across laboratories, which might have different data acquisition protocols and microscopes. We consider the setting where the main model is trained using labeled images obtained from only one source, i.e. a pharmaceutical company, and the evaluation is performed in samples from another source. Specifically, we use source 3 (S3) and source 8 (S8), and we perform experiments where the source (labeled) domain is S3 and the target domain is S8, and vice versa, and we also compare to the in-domain performance for each of them. Table 2 reports MoA accuracy for each method under source transfer.
only a subset of MoAs may appear in a new batch. If the adaptation statistics are computed over the perturbed samples, then changing the class mixture changes the feature statistics and can mislead adaptation. Protocol. We simulate increasing levels of label shift by subsampling the target batch to enforce different class-mixture skew strengths. We report performance at three shift degrees (Table 3): from mild to severe label shift. In all of the experiments, we use a constant number of labeled samples so the decrease in accuracy is not due to a smaller batch size. In every TTA method, the adaptation is unlabeled; only the class mixture is manipulated. While for TENT, AdaBN and ARM-BN, the accuracy decreases rapidly when label shift increases, the performance for CS-ARM-BN remains robust. Moreover, we also compare to AdaBN and ARM-BN including negative controls so that every method uses the same amount of samples to compute the BN statistics, and it can be seen that CS-ARM-BN still presents the highest performance.
Cross-source shifts are substantially harder. All methods degrade under cross-source transfer compared to withinsource evaluation, confirming that laboratory-to-laboratory changes introduce stronger shifts than batch-to-batch drift within a single source. Normalization-based adaptation only at test time remains competitive but is less stable. AdaBN and TENT improve over non-adaptive baselines in some transfer directions, but the improvements are not consistent across all source→target pairs. This variability is expected: under strong shift, the target feature statistics may differ not only due to style or illumination differences, but also due to more complex, structured differences (e.g., staining intensity distributions, microscope optics, and site-specific pipelines), and may require more robust context construction.
BN-adaptation degrades under label shift. As label shift increases, AdaBN, TENT and ARM-BN lose accuracy because their estimated BN statistics entangle (i) batch-specific nuisance variation with (ii) class-mixture-specific feature shifts. In other words, the adaptation signal becomes confounded by the changing MoA composition.
Meta-learned in-context adaptation improves crosssource robustness. ARM-BN tends to provide stronger cross-source performance than purely test-time methods because the model is trained to adapt to new domains via episodic context sampling during training. Overall, Experiment 5.2 highlights that cross-source robustness remains challenging, but learning to adapt is a promising direction compared to static training or target-dependent retraining.
Controls stabilize adaptation under label shift. CS-ARMBN addresses this confounding by computing the adaptation statistics from negative controls, which are independent of the MoA mixture by design. Consequently, the adaptation remains anchored to a consistent reference distribution per batch, and performance becomes substantially more stable across label-shift degrees. In Appendix Section D.1, we show that CS-ARM-BN stabilizes the performance under small batch sizes.
5.3. Realistic scenario 1: new experimental batch with label shifts
5.4. Realistic scenario 2: new imaging source with label shifts
A known failure mode of BN-based adaptation is label shift (i.e., changes in the class proportions between training and test) (Park et al., 2023; Zhao et al., 2023). In microscopy screening, this occurs naturally: many experiments are designed to test only a narrow set of hypotheses, and thus
As in the previous experiment, the evaluation samples in this scenario exhibit a class distribution that differs from the training data. However, in this case, the domain shift 9
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples Table 4. Comparison of MoA classification accuracy of different methods across target sources (S8, S3, and S6) under different degrees of label shift. A lower α denotes a higher degree of label shift. The accuracy values denote mean and standard deviation over re-runs with three different seeds. S3 → S8
Method
S8 → S3
S3 → S6
S8 → S6
α=1
α = 0.01
α=1
α = 0.01
α=1
α = 0.01
α=1
α = 0.01
no adaptation (BN)
0.669±0.016
0.670±0.016
0.500±0.012
0.498±0.012
0.537±0.069
0.536±0.069
0.568±0.005
0.568±0.006
TENT AdaBN BEN (w/o meta-learning) ARM-BEN ARM-BN CS-ARM-BN (ours)
0.706±0.055 0.736±0.046 0.125±0.000 0.664±0.031 0.747±0.010 0.878±0.006
0.139±0.006 0.171±0.009 0.125±0.000 0.663±0.031 0.173±0.004 0.825±0.007
0.659±0.005 0.688±0.004 0.125±0.000 0.544±0.031 0.680±0.007 0.788±0.014
0.136±0.000 0.164±0.002 0.125±0.000 0.543±0.031 0.159±0.001 0.770±0.045
0.552±0.027 0.580±0.023 0.124±0.001 0.547±0.040 0.550±0.011 0.657±0.003
0.135±0.004 0.159±0.006 0.122±0.004 0.546±0.040 0.158±0.003 0.606±0.006
0.571±0.006 0.586±0.005 0.125±0.000 0.566±0.024 0.579±0.010 0.648±0.011
0.122±0.001 0.149±0.002 0.125±0.000 0.565±0.024 0.154±0.003 0.611±0.015
is more pronounced, as the test samples originate from a new imaging source. In Table 4, we show that, similarly to Experiment 5.3, the performance of methods that do not leverage negative controls (TENT, AdaBN, and ARM-BN) degrades as the degree of label shift increases.
with limited target data and minimal computation, reflecting realistic deployment conditions. Why was it possible at all to close the domain gap? ARM-BN and CS-ARM-BN close or nearly close the performance gap between in-domain and out-of-domain evaluation. On unseen experimental batches, ARM-BN and CS-ARM-BN reach 0.935 ± 0.018 and 0.930 ± 0.019 accuracy, respectively, matching the 0.939 ± 0.005 accuracy observed within the training distribution. This demonstrates that batch effects, a long-standing problem for AI systems for bioimaging data, can be effectively neutralized. We hypothesize that three components were necessary to allow for closing the gap: a) Many sources/domains available for meta-learning: the JUMP-CP (Chandrasekaran et al., 2023) dataset contains thousands of different sources within the training set: each plate can be considered a source domain. Meta-learning methods (Hochreiter et al., 2001; Finn et al., 2017) can effectively learn to adapt in such scenarios. b) Stable context for adaptation methods: Each plate contains sufficient samples and controls which can stabilize estimation of the adaption parameters. This is akin to a calibration phase that many instruments in medicine and molecular biology require and for which controls must be available. c) Powerful machine learning methods: Lastly, in-context meta-learning methods are powerful learners that can profit from the structure of the JUMP-CP dataset and finally allow for closing the domain gap. We believe that this will be not possible for many other cases, but might help the community to collect datasets that exhibit characteristics like JUMP-CP.
We evaluate the performance of BEN, which relies exclusively on control samples for adaptation, for which we also extend BEN to meta-learning. While the version without meta-learning performs poorly, the meta-learned variant, despite achieving lower accuracy than CS-ARM-BN, outperforms the other methods under strong label shift. Consistent with the previous experiment, CS-ARM-BN stabilizes the batch normalization statistics and achieves the highest accuracy across both mild and severe domain shifts.
6. Discussion and outlook We exemplified a real use-case scenario of batch effect correction for mechanism of action (MoA) classification of microscopy imaging data. Our study shows that in-context adaptation, where models adapt some of their parameters on a small set of unlabeled samples from a new batch, clearly outperforms established approaches. Across four experiments, we find that (i) batch effects cause a large generalization gap even under mild within-source shifts, (ii) cross-source transfer is substantially harder but the domain gap is notably narrowed by meta-learned incontext adaptation, (iii) standard BN-based test-time adaptation methods are weak when target statistics are estimated from few or label-skewed perturbed samples. By leveraging the experimental design of high-content screening, i.e. the presence of negative controls in every batch, CS-ARM-BN, yields strong accuracy while remaining robust to small context sizes and label shift. While it could be argued that this increase in performance is solely due to a larger batch size, we also show that CS-ARM-BN outperforms its counterparts without controls even when using the same amount of samples to compute BN statistics. We intentionally focus on lightweight, fast adaptation methods that can operate
Our findings highlight that in-context test-time domain adaptation methods can make deep learning models bioimaging data not only accurate but also practical for deployment across experimental batches.
Software and data Method and data are available at https://github. com/ml-jku/cs-arm-bn. 10
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
Acknowledgments
J. D., Brino, L., Byrne, P. J., Ceulemans, H., Ch’ng, C., Cimini, B. A., Clevert, D.-A., Deflaux, N., Doench, J. G., Dorval, T., Doyonnas, R., Dragone, V., Engkvist, O., Faloon, P. W., Fritchman, B., Fuchs, F., Garg, S., Gilbert, T. J., Glazer, D., Gnutt, D., Goodale, A., Grignard, J., Guenther, J., Han, Y., Hanifehlou, Z., Hariharan, S., Hernandez, D., Horman, S. R., Hormel, G., Huntley, M., Icke, I., Iida, M., Jacob, C. B., Jaensch, S., Khetan, J., KostAlimova, M., Krawiec, T., Kuhn, D., Lardeau, C.-H., Lembke, A., Lin, F., Little, K. D., Lofstrom, K. R., Lotfi, S., Logan, D. J., Luo, Y., Madoux, F., Zapata, P. A. M., Marion, B. A., Martin, G., McCarthy, N. J., Mervin, L., Miller, L., Mohamed, H., Monteverde, T., Mouchet, E., Nicke, B., Ogier, A., Ong, A.-L., Osterland, M., Otrocka, M., Peeters, P. J., Pilling, J., Prechtl, S., Qian, C., Rataj, K., Root, D. E., Sakata, S. K., Scrace, S., Shimizu, H., Simon, D., Sommer, P., Spruiell, C., Sumia, I., Swalley, S. E., Terauchi, H., Thibaudeau, A., Unruh, A., de Waeter, J. V., Dyck, M. V., van Staden, C., Warchoł, M., Weisbart, E., Weiss, A., Wiest-Daessle, N., Williams, G., Yu, S., Zapiec, B., Żyła, M., Singh, S., and Carpenter, A. E. (2023). Jump cell painting dataset: morphological impact of 136,000 chemical and genetic perturbations. bioRxiv, page 2023.03.23.534023.
The ELLIS Unit Linz, the LIT AI Lab, the Institute for Machine Learning, are supported by the Federal State Upper Austria. We thank the projects FWF AIRI FG 9-N (10.55776/FG9), AI4GreenHeatingGrids (FFG- 899943), Stars4Waters (HORIZON-CL6-2021-CLIMATE-01-01), FWF Bilateral Artificial Intelligence (10.55776/COE12). We thank Merck Healthcare KGaA, NXAI GmbH, Audi AG, Silicon Austria Labs (SAL), GLS (Univ. Waterloo), TÜV Holding GmbH, Software Competence Center Hagenberg GmbH, dSPACE GmbH, TRUMPF SE + Co. KG.
References Ando, D. M., McLean, C. Y., and Berndl, M. (2017). Improving phenotypic measurements in high-content imaging screens. bioRxiv, page 161422. Arevalo, J., Su, E., Ewald, J. D., van Dijk, R., Carpenter, A. E., and Singh, S. (2024). Evaluating batch correction methods for image-based cell profiling. Nature Communications 2024 15:1, 15:1–12. Baxter, J. (1998). Theoretical models of learning to learn. In Learning to learn, pages 71–94. Springer.
Chen, W., Zhao, Y., Chen, X., Yang, Z., Xu, X., Bi, Y., Chen, V., Li, J., Choi, H., Ernest, B., Tran, B., Mehta, M., Kumar, P., Farmer, A., Mir, A., Mehra, U. A., Li, J. L., Moos, M., Xiao, W., and Wang, C. (2020). A multicenter study benchmarking single-cell rna sequencing technologies using reference samples. Nature Biotechnology 2020 39:9, 39:1103–1114.
Ben-David, S., Blitzer, J., Crammer, K., and Pereira, F. (2006). Analysis of representations for domain adaptation. Advances in neural information processing systems, 19. Blanchard, G., Lee, G., and Scott, C. (2011). Generalizing from several related classification tasks to a new unlabeled sample. In Shawe-Taylor, J., Zemel, R., Bartlett, P., Pereira, F., and Weinberger, K., editors, Advances in Neural Information Processing Systems, volume 24. Curran Associates, Inc.
Chung, J., Hyun, S., and Heo, J. P. (2024). Style injection in diffusion: A training-free approach for adapting largescale diffusion models for style transfer. Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition, pages 8795–8805.
Boudiaf, M., Mueller, R., Ayed, I. B., and Bertinetto, L. (2022). Parameter-free online test-time adaptation. Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2022June:8334–8343.
Dong, M., Adduri, A., Gautam, D., Carpenter, C., Shah, R., Ricci-Tam, C., Kluger, Y., Burke, D. P., and Roohani, Y. H. (2026). Stack: In-context learning of single-cell biology. bioRxiv.
Bousmalis, K., Silberman, N., Dohan, D., Erhan, D., and Krishnan, D. (2017). Unsupervised pixel-level domain adaptation with generative adversarial networks. Proceedings - 30th IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2017, 2017-January:95–104.
Farahani, A., Voghoei, S., Rasheed, K., and Arabnia, H. R. (2020). A brief review of domain adaptation. ArXiv, pages 877–894. Finn, C., Abbeel, P., and Levine, S. (2017). Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pages 1126– 1135. PMLR.
Bray, M.-A., Carpenter, A., of MIT, B. I., and Platform, H. I. (2017). Advanced assay development guidelines for image-based high content screening and analysis. Assay Guidance Manual.
Ganin, Y. and Lempitsky, V. (2014). Unsupervised domain adaptation by backpropagation. 32nd International Conference on Machine Learning, ICML 2015, 2:1180–1189.
Chandrasekaran, S. N., Ackerman, J., Alix, E., Ando, D. M., Arevalo, J., Bennion, M., Boisseau, N., Borowa, A., Boyd, 11
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
Gong, T., Jeong, J., Kim, T., Kim, Y., Shin, J., and Lee, S. J. (2022). Note: Robust continual test-time adaptation against temporal correlation. Advances in Neural Information Processing Systems, 35.
(2025). Self-supervision advances morphological profiling by unlocking powerful image representations. Scientific Reports 2025 15:1, 15:1–15. Kimura, M. and Hino, H. (2024). A short survey on importance weighting for machine learning. arXiv preprint arXiv:2403.10175.
Graham, B., El-Nouby, A., Touvron, H., Stock, P., Joulin, A., Jégou, H., and Douze, M. (2021). Levit: a vision transformer in convnet’s clothing for faster inference. Proceedings of the IEEE International Conference on Computer Vision, pages 12239–12249.
Knowles, J. and Gromo, G. (2003). A guide to drug discovery: Target selection in drug discovery. Nature reviews. Drug discovery, 2:63–69.
Guo, J., Shah, D. J., and Barzilay, R. (2018). Multi-source domain adaptation with mixture of experts. arXiv preprint arXiv:1809.02256.
Korsunsky, I., Millard, N., Fan, J., Slowikowski, K., Zhang, F., Wei, K., Baglaenko, Y., Brenner, M., ru Loh, P., and Raychaudhuri, S. (2019). Fast, sensitive and accurate integration of single-cell data with harmony. Nature Methods 2019 16:12, 16:1289–1296.
Haghverdi, L., Lun, A. T., Morgan, M. D., and Marioni, J. C. (2018). Batch effects in single-cell rna-sequencing data are corrected by matching mutual nearest neighbors. Nature Biotechnology 2018 36:5, 36:421–427.
Kraus, O., Kenyon-Dean, K., Saberian, S., Fallah, M., McLean, P., Leung, J., Sharma, V., Khan, A., Balakrishnan, J., Celik, S., Beaini, D., Sypetkowski, M., Cheng, C. V., Morse, K., Makes, M., Mabey, B., and Earnshaw, B. (2024). Masked autoencoders for microscopy are scalable learners of cellular biology. Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition, pages 11757–11768.
Haslum, J. F., Matsoukas, C., Leuchowius, K. J., and Smith, K. (2023). Bridging generalization gaps in high content imaging through online self-supervised domain adaptation. IEEE Workshop/Winter Conference on Applications of Computer Vision, pages 7723–7732. He, K., Zhang, X., Ren, S., and Sun, J. (2015). Deep residual learning for image recognition. Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2016-December:770–778.
Lee, J., Jung, D., Lee, S., Park, J., Shin, J., Hwang, U., and Yoon, S. (2024). Entropy is not enough for test-time adaptation: From the perspective of disentangled factors. 12th International Conference on Learning Representations, ICLR 2024.
Hie, B., Bryson, B., and Berger, B. (2019). Efficient integration of heterogeneous single-cell transcriptomes using scanorama. Nature biotechnology, 37:685–691. Hochreiter, S., Younger, A. S., and Conwell, P. R. (2001). Learning to learn using gradient descent. In International conference on artificial neural networks, pages 87–94. Springer.
Leek, J. T., Scharpf, R. B., Bravo, H. C., Simcha, D., Langmead, B., Johnson, W. E., Geman, D., Baggerly, K., and Irizarry, R. A. (2010). Tackling the widespread and critical impact of batch effects in high-throughput data. Nature reviews. Genetics, 11:10.1038/nrg2825.
Hughes, J. P., Rees, S. S., Kalindjian, S. B., and Philpott, K. L. (2011). Principles of early drug discovery. British Journal of Pharmacology, 162:1239.
Li, Y., Wang, N., Shi, J., Liu, J., and Hou, X. (2016). Revisiting batch normalization for practical domain adaptation. International Conference on Learning Representations.
Ioffe, S. and Szegedy, C. (2015). Batch normalization: accelerating deep network training by reducing internal covariate shift. In Proceedings of the 32nd International Conference on International Conference on Machine Learning Volume 37, ICML’15, page 448–456. JMLR.org.
Lin, A. and Lu, A. (2022). Incorporating knowledge of plates in batch normalization improves generalization of deep learning for microscopy images. In Knowles, D. A., Mostafavi, S., and Lee, S.-I., editors, Proceedings of the 17th Machine Learning in Computational Biology meeting, volume 200 of Proceedings of Machine Learning Research, pages 74–93. PMLR.
Johnson, W. E., Li, C., and Rabinovic, A. (2007). Adjusting batch effects in microarray expression data using empirical bayes methods. Biostatistics (Oxford, England), 8:118–127.
Liu, M.-Y. and Tuzel, O. (2016). Coupled generative adversarial networks. In Lee, D., Sugiyama, M., Luxburg, U., Guyon, I., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc.
Kim, V., Adaloglou, N., Osterland, M., Morelli, F. M., Halawa, M., König, T., Gnutt, D., and Zapata, P. A. M. 12
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
Lopez, R., Regier, J., Cole, M. B., Jordan, M. I., and Yosef, N. (2018). Deep generative modeling for single-cell transcriptomics. Nature methods, 15:1053–1058.
J., Sereshkeh, A. R., Haque, I., and Earnshaw, B. (2023). RxRx1: A dataset for evaluating experimental batch correction methods. IEEE Computer Society Conference on Computer Vision and Pattern Recognition Workshops, 2023-June:4285–4294.
Luecken, M. D., Büttner, M., Chaichoompu, K., Danese, A., Interlandi, M., Mueller, M. F., Strobl, D. C., Zappia, L., Dugas, M., Colomé-Tatché, M., and Theis, F. J. (2021). Benchmarking atlas-level data integration in single-cell genomics. Nature Methods 2021 19:1, 19:41–50.
Taigman, Y., Polyak, A., and Wolf, L. (2016). Unsupervised cross-domain image generation. ArXiv, abs/1611.02200. Venkat, N., Kundu, J. N., Singh, D., Revanur, A., et al. (2020). Your classifier can secretly suffice multi-source domain adaptation. Advances in Neural Information Processing Systems, 33:4647–4659.
Marsden, R. A., Döbler, M., and Yang, B. (2023). Universal test-time adaptation through weight ensembling, diversity weighting, and prior correction. Proceedings - 2024 IEEE Winter Conference on Applications of Computer Vision, WACV 2024, pages 2543–2553.
Wang, D., Shelhamer, E., Liu, S., Olshausen, B., and Darrell, T. (2021). TENT: Fully test-time adaptation by entropy minimization. In International Conference on Learning Representations.
Palma, A., Theis, F. J., and Lotfollahi, M. (2025). Predicting cell morphological responses to perturbations using generative modeling. Nature Communications, 16:1–19.
Wang, Q., Fink, O., Gool, L. V., and Dai, D. (2022). Continual test-time domain adaptation. Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2022-June:7191–7201.
Park, S., Yang, S., Choo, J., and Yun, S. (2023). Label shift adapter for test-time adaptation under covariate and label shifts. Proceedings of the IEEE International Conference on Computer Vision, pages 16375–16385.
Wen, J., Greiner, R., and Schuurmans, D. (2020). Domain aggregation networks for multi-source domain adaptation. In International conference on machine learning, pages 10214–10224. PMLR.
Peidli, S., Green, T. D., Shen, C., Gross, T., Min, J., Garda, S., Yuan, B., Schumacher, L. J., Taylor-King, J. P., Marks, D. S., Luna, A., Blüthgen, N., and Sander, C. (2023). scPerturb: Harmonized single-cell perturbation data. bioRxiv, page 2022.08.20.504663.
Yang, L., Balaji, Y., Lim, S.-N., and Shrivastava, A. (2020). Curriculum manager for source selection in multi-source domain adaptation. In European conference on computer vision, pages 608–624. Springer.
Peng, X., Bai, Q., Xia, X., Huang, Z., Saenko, K., and Wang, B. (2019). Moment matching for multi-source domain adaptation. In Proceedings of the IEEE/CVF international conference on computer vision, pages 1406– 1415.
Zellinger, W., Grubinger, T., Lughofer, E., Natschläger, T., and Saminger-Platz, S. (2017). Central moment discrepancy (cmd) for domain-invariant representation learning. In International Conference on Learning Representations (ICLR).
Shimodaira, H. (2000). Improving predictive inference under covariate shift by weighting the log-likelihood function. Journal of statistical planning and inference, 90(2):227–244.
Zhang, J., Ubas, A. A., de Borja, R., Svensson, V., Thomas, N., Thakar, N., Lai, I., Winters, A., Khan, U., Jones, M. G., Tran, V., Pangallo, J., Papalexi, E., Sapre, A., Nguyen, H., Sanderson, O., Nigos, M., Kaplan, O., Schroeder, S., Hariadi, B., Marrujo, S., Salvino, C. C. A., Gallareta Olivares, G., Koehler, R., Geiss, G., Rosenberg, A., Roco, C., Merico, D., Alidoust, N., Goodarzi, H., and Yu, J. (2025). Tahoe-100M: A giga-scale single-cell perturbation atlas for context-dependent gene function and cellular modeling. bioRxiv.
Stirling, D. R., Swain-Bowden, M. J., Lucas, A. M., Carpenter, A. E., Cimini, B. A., and Goodman, A. (2021). Cellprofiler 4: improvements in speed, utility and usability. BMC Bioinformatics, 22:433. Stuart, T., Butler, A., Hoffman, P., Hafemeister, C., Papalexi, E., Mauck, W. M., Hao, Y., Stoeckius, M., Smibert, P., and Satija, R. (2019). Comprehensive integration of single-cell data. Cell, 177:1888–1902.e21.
Zhang, M., Marklund, H., Dhawan, N., Gupta, A., Levine, S., and Finn, C. (2021). Adaptive Risk Minimization: learning to adapt to domain shift. In Proceedings of the 35th International Conference on Neural Information Processing Systems, NeurIPS ’21, Red Hook, NY, USA. Curran Associates Inc.
Sun, B., Feng, J., and Saenko, K. (2015). Return of frustratingly easy domain adaptation. 30th AAAI Conference on Artificial Intelligence, AAAI 2016, pages 2058–2065. Sypetkowski, M., Rezanejad, M., Saberian, S., Kraus, O., Urbanik, J., Taylor, J., Mabey, B., Victors, M., Yosinski, 13
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
Zhao, H., Liu, Y., Alahi, A., and Lin, T. (2023). On pitfalls of test-time adaptation. Proceedings of Machine Learning Research, 202:42058–42080. Zhao, H., Zhang, S., Wu, G., Moura, J. M., Costeira, J. P., and Gordon, G. J. (2018). Adversarial Multiple Source Domain Adaptation. Advances in neural information processing systems, 31.
14
Contents A Notation table
16
B Training details
17
C Dataset details
17
D Additional experiments D.1 Failure mode of BN adaptation — dependence on adaptation batch size . . . . . . . . . . . . . . . . . . D.2 Label shift experiment. Supplementary information. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.3 Evaluation of CS-ARM-BN with a Vision Transformer-based architecture . . . . . . . . . . . . . . . . . . D.4 Microscopy imaging data structure and domain alignment . . . . . . . . . . . . . . . . . . . . . . . . . . .
19 19 20 21 21
E The in-context learning view on CS-ARM-BN
22
F Details on estimators F.1 Setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . F.2 Estimator 1: ARM-BN and AdaBN (Perturbed Samples Only) . . . . . . . . . . . . . . . . . . . . . . . F.3 Estimator 2: Controls Only . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . F.4 Estimator 3: CS-ARM-BN (Controls + Perturbed Samples) . . . . . . . . . . . . . . . . . . . . . . . . . F.5 Comparison . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
22 22 23 23 23 24
15
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
A. Notation table Definition Scalars number of source domains (batches) number of samples per source domain (simplified) number of samples from one of the target domains number of control samples from target domain number of non-control samples from target domain feature dimension number of classes DANN domain-loss trade-off CORAL penalty weight
Symbol
Type
B N M C L D K λ γ
N N N N N N N R≥0 R≥0
Sets, spaces, and distributions input/image space label space / class set training data (all sources) data from source domain b unseen target-batch data domain-specific data distribution meta-distribution over domains
Dtrain D(b) Dβ pbxy µ
Representations and tensors image / sample / input control image / control sample label feature embedding TVN-normalized embedding control embeddings (batch β) source/target feature covariances control mean/covariance (TVN)
x z y a = f (x; w′ ) ã {zjβ } CS , C T µctrl,β , Σctrl,β
Models and parameters feature extractor classifier prediction model model parameters adaptation function (test-time) adaptation parameters adapted parameters for batch β domain discriminator (DANN)
f , fw ,f (·; w′ ) c(·; θ) g(x; w) = (c ◦ f )(x; w) w hϕ (w, xβ1:M ) ϕ wβ dv
Losses and objectives supervised loss (cross-entropy) domain loss (DANN) CORAL discrepancy
LCE (y, g(x; w)) LCE (b, dv (h)) ∥CS − CT ∥2F
X Y
BatchNorm (BN) statistics test-batch mean/var (per channel, batch β)
µBN , σBN 16
SB
(b) b=1 D (b) (b) {(xn , yn )}N n=1 {xβ1 , . . . , xβM }
distribution on X × Y distribution on domains X X Y RD RD RD RD×D RD , RD×D X → RD RD → Y X →Y W W × XM → W parameters of h W RD → S B−1 R R R≥0 RD , R D >0
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
B. Training details In our experiments, we used either Nvidia A40 or A100 GPUs and all experiments use a single GPU. When reporting results, for the new experimental batch evaluation (experiment 5.1) and for the in-domain evaluation of source 3 (S3→S3 in Table 2), we perform 5-fold cross-validation and report the mean and standard deviation across folds. For the in-domain evaluation of source 8 (S8→S8), due to computational constrains, we perform re-runs in two different random splits. For all methods, in Experiments 5.1 and 5.3 we set a maximum number of training epochs to 400 and implement an early stopping strategy, setting the patience to 80 epochs, so that if the validation loss does not decrease with respect to the current minimum loss for 80 epochs, the training is stopped. In Experiments 5.2 and 5.4, due to computational constraints, we set we set a maximum number of training epochs to 200 and a patience of 20. In all experiments, the checkpoint used for evaluating the test set is the one that achieved the lowest validation loss. By default, we use a ResNet50 as a feature encoder. As an optimizer, unless otherwise stated, we use Adam and the selected learning rate is decreased using a cosine annealing scheduler. The hyperparameters were manually tuned, and the exploration space for each method can be found in Supplementary Table A5 For the methods based in the CA-MAE foundation model, linear probing is performed, so that the weights of the feature encoder are frozen and only the classifier, which consists in one linear layer, is trained. The pretrained weights of this model are the weights of a ViT-S/16, provided in HuggingFace by Kraus et al. (2024), trained with the RxRx3 dataset and subsets of the JUMP-CP with samples treated for gene knockout and overexpression. In the style transfer baseline, as the images were generated using a pre-trained diffusion model, the input for the feature extractor consisted in 3 channels instead of 5. Moreover, for computational efficiency, the images were generated with a resolution of 256x256. For performing the TVN correction, control images from the corresponding matching plate of each perturbed sample are used for centering and whitening.
C. Dataset details As explained in Section 5, the experiments were carried out on the JUMP-CP dataset, using source 3 and source 8. The classification task consists in classifying the MoA of the given image, similar to Haslum et al. (2023). As we need images that are correctly annotated, we use images that were treated with compounds that have a clear and well-defined MoA. Specifically, we use the images treated with the compounds used as positive controls, which are all present in every plate. The name of these compounds, along with their mechanism of action can be found in Supplementary Table A6. Source 3 contains 9 experimental batches and consists of 49,491 microscopy images (with 5 channels each) and 172 different plates. Source 8 contains 4 experimental batches and 56,160 microscopy images (with 5 channels each) and 195 different plates. For experiments with the CellProfiler features, we download the provided images in https://registry.opendata.aws/cellpaintinggallery/, and perform standardization by using the training features. We use these features as provided, and they were calculated on a well-level basis, so that this is the only baseline where accuracy is evaluated on a well level. Preprocessing. As a preprocessing step, the images were downsized to a resolution of 512x512 and we performed illumination correction by clipping 0.01% of the brightest pixels. During training, the images were further resized to 256x256 and random vertical and horizontal flipping was performed. Moreover, in all methods, except for the baseline with instance norm and the MAE-based methods, the images were normalized using the mean and standard deviation calculated for the training split. In the remaining methods, the images are preprocessed by normalizing the channel of each image individually, in order to match the pretraining procedure in the case of the CA-MAE baselines (Kraus et al., 2024). For the TVN correction, additional negative controls, i.e. samples in which only DMSO was applied, were used. In total, the negative controls present in the 172 plates from the dataset described above, consist of 49,536 images.
17
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
Table A5. Considered hyperparameter space for all compared methods. The selected configurations for mechanism of action prediction based on manual search on validation set are shown in bold.
Hyperparameter
Explored space
Baseline (ResNet)
Learning rate Scheduler Batch size
{0.001, 0.005} {Cosine annealing, None} {64, 128}
Baseline (CA-MAE)
Learning rate Batch size
{0.001, 0.002, 0.003, 0.004, 0.005, 0.006, 0.007, 0.009, 0.01} {128}
Learning rate Batch size λ Domain classifier learning rate Unlabeled batch size
DANN
{0.001} {128} {0.1, 0.5, 0.7, 2, 3, 4} {0.0001} {128}
CORAL
Learning rate Batch size γ Unlabeled batch size
CA-MAE (FM) + TVN
Learning rate Batch size
{0.001, 0.002, 0.003, 0.004, 0.005, 0.006, 0.007, 0.009, 0.01} {128}
Style-ID
Learning rate Batch size
{0.001} {128}
ARM-BN
Learning rate Batch size
{0.001} {8, 16, 32, 64, 128}
CS-ARM-BN
ARM-CML
{0.001} {128} {0.05, 0.2. 0.3, 1} {128}
Learning rate Batch size Negative controls batch size
{0.001} {8, 16, 32, 64} {64, 128}
Learning rate Batch size Context channels Kernel size
{0.001} {128} {12} {5}
Table A6. Compounds used as mechanism of action representatives for the classification task.
Compound name Aloxistatin FK-866 AMG900 Dexamethasone LY2109761 NVS-PAK1-1 Quinidine TC-S-7004
Mechanism of action Cysteine protease inhibitor NAMPT inhibitor Aurora kinase inhibitor Glucocorticoid agonist TGF-β inhibitor PAK1 inhibitor Sodium channel inhibitor DYRK1A/B inhibitor
18
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
D. Additional experiments D.1. Failure mode of BN adaptation — dependence on adaptation batch size BN-based test-time adaptation implicitly relies on accurate estimation of target statistics. When only a few samples are available from a new batch (or when adaptation must be performed in small micro-batches due to memory constraints), these estimates become noisy. To stress-test this failure mode, we vary the number of target samples used to compute adaptation statistics (adaptation batch size / context size M ) and evaluate methods that adapt BN parameters. Table (BatchNorm size) reports the results. Previous BN-adaptation methods collapse for small M . Both AdaBN and TENT perform poorly when only very few target samples are available to estimate BN statistics (e.g., M ∈ {1, 2, 4}). Performance improves monotonically with M , approaching a plateau only once statistics can be reliably estimated (larger M ). This explains why BN-based test-time adaptation can appear strong in benchmarks that provide a full target batch, yet fail in operational settings with limited or streaming access. Control-stabilized context mitigates the small-M regime. Our control-stabilized variant (CS-ARM-BN in the table) remains strong even for extremely small numbers of perturbed samples. The key is that, while perturbed-class composition may be sparse, negative controls are plentiful and consistently available in every batch. Using controls to estimate batch statistics yields a stable estimate of batch-specific nuisance variation, making adaptation substantially less sensitive to the number of available perturbed samples. Table A7. Comparison of MoA classification accuracy of different methods in a new experimental batch with different batch sizes. The accuracy values denote mean and standard deviation over 5-fold cross-validation, thus five training re-runs of each method. Method ERM (no adaptation) TENT AdaBN ARM-BN CS-ARM-BN
Labeled batch size 8 16
1
2
4
0.862 ± 0.060 0.126 ± 0.001 0.130 ± 0.005 0.126 ± 0.007 0.923 ± 0.018
0.862 ± 0.060 0.175 ± 0.013 0.447 ± 0.095 0.514 ± 0.092 0.923 ± 0.018
0.862 ± 0.060 0.230 ± 0.011 0.709 ± 0.052 0.716 ± 0.052 0.924 ± 0.019
0.862 ± 0.060 0.323 ± 0.012 0.839 ± 0.024 0.832 ± 0.028 0.925 ± 0.018
0.862 ± 0.060 0.550 ± 0.026 0.898 ± 0.014 0.896 ± 0.019 0.927 ± 0.019
32
64
128
Full domain
0.862 ± 0.060 0.794 ± 0.021 0.919 ± 0.014 0.923 ± 0.018 0.929 ± 0.019
0.862 ± 0.060 0.881 ± 0.025 0.924 ± 0.014 0.930 ± 0.018 0.930 ± 0.019
0.862 ± 0.060 0.889 ± 0.031 0.927 ± 0.013 0.933 ± 0.017 0.929 ± 0.020
0.862 ± 0.060 0.908 ± 0.032 0.928 ± 0.013 0.935 ± 0.018 0.923 ± 0.023
D.1.1. A BLATION STUDY. E FFECT OF THE NEGATIVE CONTROLS SET SAMPLE SIZE In the previous experiment, we assume that other methods do not consider the presence of negative controls in biological data, so that the available samples are only perturbed ones, while in CS-ARM-BN we make use of negative controls, which results in a larger effective batch size. However, in this experiment we also study the effectiveness of standard ARM-BN in the presence of negative controls, so that we compare both methods using the same total batch size. In Table A8, we show that CS-ARM-BN presents a higher accuracy even when the same amount of unlabeled data is available, probably due to the fact that during the meta-learning phase, the model has already learned to normalize the intermediate features of the model using both negative controls and perturbed samples. Table A8. Comparison between ARM-BN and CS-ARM-BN classification accuracy using the same number of samples in every batch, with the same proportion of negative controls per perturbed sample. The accuracy values denote mean and standard deviation over 5-fold cross-validation, thus five training re-runs of each method.
Method ARM-BN CS-ARM-BN (ours)
5 [1 : 4]
Total batch size [perturbed : controls] 10 [2 : 8] 20 [4 : 16] 40 [8 : 32]
0.799 ± 0.020 0.848 ± 0.035
0.827 ± 0.024 0.875 ± 0.029
19
0.846 ± 0.025 0.891 ± 0.028
0.860 ± 0.026 0.899 ± 0.027
80 [16 : 64] 0.870 ± 0.028 0.905 ± 0.026
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
D.2. Label shift experiment. Supplementary information. D.2.1. E XPERIMENT DETAILS In experiment 5.3 we show a comparison between methods under label shifts. To achieve batches that have a certain level of label shift, we sample a class probability vector from a Dirichlet distribution with a certain α parameter. Then we use this class probability vector to define a multinomial distribution and draw N samples from it, that we will use as input to our model. In total we use a labeled batch size of 36 in every method, and we include 288 negative samples in CS-ARM-BN. One could argue that performance differences might be only due to a larger sample size, therefore, we show in Table 3 other baselines including the same amount of negative controls in each batch, as explained in the following section. D.2.2. A BLATION STUDY. E FFECT OF THE NEGATIVE CONTROLS SET SAMPLE SIZE In our method, CS-ARM-BN, negative controls are used both at training and at inference time, so the total batch size used to compute the BN parameters is larger than for the other methods. To demonstrate that the gains in performance in our method are not only due to a larger batch size but also due to the fact that the models were trained in a meta-learning fashion, we run the AdaBN and ARM-BN with the same amount of negative controls and same total batch size than for CS-ARM-BN. We don’t perform this experiment for TENT because we argue that there is no reason why minimizing the entropy of the model outputs for control samples should be beneficial. In Table A9 we show that, even if including controls increases the performance of both methods, CS-ARM-BN still shows the best performance. Table A9. Comparison of MoA classification accuracy of different methods in a new experimental batch under different degrees of label shift. A smaller α value corresponds to a larger degree of label shift.
Method
Label shift degree α=1 α = 0.1
α=4
α = 0.01
ERM (no adaptation)
0.862 ± 0.060
0.862 ± 0.060
0.862 ± 0.061
0.862 ± 0.061
TENT Ada-BN ARM-BN Ada-BN (perturbed + controls) ARM-BN (perturbed + controls) CS-ARM-BN (ours)
0.887 ± 0.014 0.899 ± 0.012 0.894 ± 0.018 0.879 ± 0.018 0.861 ± 0.032 0.927 ± 0.017
0.832 ± 0.014 0.851 ± 0.016 0.831 ± 0.017 0.871 ± 0.019 0.850 ± 0.033 0.924 ± 0.018
0.778 ± 0.016 0.452 ± 0.010 0.416 ± 0.020 0.815 ± 0.022 0.785 ± 0.058 0.906 ± 0.021
0.156 ± 0.010 0.229 ± 0.009 0.228 ± 0.009 0.785 ± 0.033 0.752 ± 0.068 0.894 ± 0.022
20
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
D.3. Evaluation of CS-ARM-BN with a Vision Transformer-based architecture Although, throughout the main manuscript, we have evaluated our method and the baselines using a ResNet50, in this section we also benchmark the best performing methods using LeViT (Graham et al., 2021), an architecture based on the Vision Transformer that employs Batch Normalization layers instead of the usual Layer Normalization. Particularly, we evaluate the methods on the most challenging task, under both label shift and strong domain shift. In Table A10, it can be seen that adapting the Batch Normalization layers and including negative controls in a meta-learning fashion a is an effective strategy to mitigate domain and label shifts also when using a Vision Transformer architecture. We evaluated ERM, TENT, AdaBN and ARM-BN as in previous experiments, and additionally included two more recent TTA baselines, DEYO (Lee et al., 2024) and ROID (Marsden et al., 2023). Table A10. Comparison of MoA classification accuracy of different methods across target sources under different degrees of label shift, using a LeViT, a vision transformer-based architecture. A lower α denotes a higher degree of label shift. The accuracy values denote mean and standard deviation over re-runs with three different seeds.
S3 → S8
Method
S8 → S3
α=1
α = 0.01
α=1
α = 0.01
(LeViT) no adaptation (BatchNorm)
0.607±0.023
0.608±0.019
0.536±0.062
0.536±0.051
(LeViT) TENT (LeViT) DEYO (LeViT) ROID (LeViT) AdaBN (LeViT) ARM-BN (LeViT) CS-ARM-BN (ours)
0.675±0.017 0.674±0.014 0.694±0.018 0.699±0.015 0.735±0.024 0.791±0.025
0.137±0.002 0.137±0.002 0.138±0.003 0.165±0.002 0.166±0.003 0.707±0.021
0.640±0.013 0.639±0.015 0.657±0.015 0.665±0.014 0.702±0.119 0.722±0.015
0.145±0.005 0.145±0.005 0.147±0.005 0.168±0.004 0.165±0.001 0.610±0.017
D.4. Microscopy imaging data structure and domain alignment Microscopy imaging data presents a hierarchical structure. In each imaging platform, also called source throughout this manuscript, reagents and cell cultures change over time and have to be replaced, which leads to different experimental batches. Moreover, within each one of this experimental batches, cells have to be laid out in what are known as plates, which are then placed under the microscope and images are obtained. As we have explained in Section 1, there are visible differences between plates, batches and sources. Moreover, each one of these plates has different positions, known as wells. Given this hierarchy, a natural question is at which level to perform the domain alignment. While each source or batch has a larger number of samples, they are a set of other domains which might lead to a non-optimal alignment. In the following experiment, we show the adaptation results for AdaBN at the possible different levels: source, plate and batch, and we show that aligning at the plate level yields better results, even if fewer samples are available per domain. Table A11. MoA classification accuracy of the AdaBN method adapting the BatchNorm statistics per source, per batch and per plate. Even if each plate consists of, naturally, fewer samples than each batch or source, the resulting accuracy is higher when the domain adaptation is performed per plate. Accuracy is an average over the S3→S8 and S8→S3 transfer experiments, shown in Table 5.2
Alignment level Source Batch Plate
Accuracy
Number of domains
Samples per domain (avg.)
0.773 ± 0.033 0.790 ± 0.014 0.798 ± 0.014
1 4 195
56,160 14,040 288
21
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
E. The in-context learning view on CS-ARM-BN Our method CS-ARM-BN can be understood as a form of in-context learning: rather than retraining the model, the classifier conditions its predictions on a small context of unlabeled samples from the new experimental batch, without any gradient update. The negative controls play the role of a stable, class-agnostic reference that directly reflects the batch-specific technical variation. As shown in Algorithm 2, this can be written in two equivalent ways: either as an explicit two-step procedure of calibration followed by inference (model.adapt), or as a single context-conditioned forward pass (model.cpredict), making the analogy to in-context learning explicit.
Algorithm 2 Pseudo code for CS-ARM-BN vs traditional image classifier pipelines | In-context view on CS-ARM-BN ### Standard Image Classifier Pipeline ### # no calibration of classifier to data x = load_data() # your biomedical images model = load_model() # public classifier pred = model.predict(x) # domain gap! ### CS-ARM-BN Classifier Pipeline ### # adapt/calibrate classifier to your data x = load_data() # your biomedical images z = load_controls() # available controls model = load_model() # public classifier model.adapt(concat([x, z])) # calibrate pred = model.predict(x) # neutralized gap ### In-context learning view on CS-ARM-BN # adapt/calibrate classifier to your data x = load_data() # your biomedical images z = load_controls() # available controls model = load_model() # public classifier context = concat([x, z]) # assemble context pred = model.cpredict(x, context=context)
F. Details on estimators F.1. Setup We model the BN activation mean for a sample x as: µobs,x = µdomain + µclass(x) + εx ,
(15)
where µdomain is the batch-specific technical offset we wish to estimate, µclass(x) = 0 for control samples, µclass(x) ̸= 0 for perturbed samples, and εx ∼ (0, σ 2 ) is additive noise. We have C control samples and L perturbed samples, with M = C + L total. The mean class effect across all perturbed samples is L
µ̄class =
1X µclass(l) . L
(16)
l=1
For any estimator µ̂, the MSE decomposes as: MSE(µ̂) = Bias2 (µ̂) + Var(µ̂). 22
(17)
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
F.2. Estimator 1: ARM-BN and AdaBN (Perturbed Samples Only) The estimator is the sample mean over the L perturbed samples: µ̂ARM−BN =
L
L
l=1
l=1
1X 1X µobs,l = (µdomain + µclass(l) + εl ). L L
(18)
Bias. E[µ̂ARM−BN ] = µdomain + µ̄class , ARM−BN
Bias = E[µ̂
(19)
] − µdomain = µ̄class .
(20)
The bias equals the mean class effect, which grows large under label shift. Variance. L
ARM−BN
Var(µ̂
) = Var
1X εl L
! =
l=1
σ2 . L
(21)
MSE. MSE(µ̂ARM−BN ) = µ̄2class +
σ2 . L
(22)
F.3. Estimator 2: Controls Only The estimator is the sample mean over the C control samples. Since controls satisfy µclass = 0: µ̂ctrl =
C C 1 X 1 X µobs,c = (µdomain + εc ). C c=1 C c=1
(23)
Bias. E[µ̂ctrl ] = µdomain ,
(24)
Bias = 0.
(25)
The estimator is unbiased because controls carry no class-specific signal by definition. Variance. Var(µ̂ctrl ) =
σ2 . C
(26)
This is large when C is small, which is the typical regime in drug-screening experiments. MSE. MSE(µ̂ctrl ) = 0 +
σ2 . C
(27)
F.4. Estimator 3: CS-ARM-BN (Controls + Perturbed Samples) The estimator is the sample mean over all M = C + L samples: 1 X 1 µ̂CS−ARM−BN = µobs,u = M M β u∈C
23
C X c=1
µobs,c +
L X l=1
! µobs,l
.
(28)
Closing the Domain Gap in Biomedical Imaging by In-Context Control Samples
Expanding each term: C X
1 µ̂CS−ARM−BN = M
(µdomain + εc ) +
c=1
L X
! (µdomain + µclass(l) + εl )
l=1 L
1 X 1 X µclass(l) + εu M M u l=1 L 1 X = µdomain + µ̄class + εu . M M u
= µdomain +
(29)
Bias. E[µ̂CS−ARM−BN ] = µdomain + Bias =
L µ̄class , M
(30)
L µ̄class . M
(31)
The bias is attenuated by the factor L/M < 1 relative to Estimator 1. The more controls are included (larger C, smaller L/M ), the smaller the bias. Variance. CS−ARM−BN
Var(µ̂
) = Var
1 X εu M u
! =
σ2 . M
(32)
Since M > L and M > C, this variance is strictly smaller than those of Estimators 1 and 2. MSE. CS−ARM−BN
MSE(µ̂
)=
L M
2
µ̄2class +
σ2 . M
(33)
F.5. Comparison Table A12. MSE decomposition for the three estimators of µdomain .
Estimator ARM-BN and AdaBN Controls-only CS-ARM-BN
Bias2 µ̄2class 0 (L/M )2 µ̄2class
Variance σ 2 /L σ 2 /C σ 2 /M
Notes Large bias under label shift Large variance when C is small Trade-off between the two above
CS-ARM-BN achieves the lowest MSE by trading a residual bias shrunk by (L/M )2 relative to ARM-BN against the lowest variance of all three estimators (since M ≥ max(L, C)). This bias-variance trade-off is most favourable in the typical drug-screening regime: few controls (C small) and moderate-to-severe label shift.
24