Stabilizing In-Context Multi-Source Domain Adaptation for Biomedical Images Through Controls
Abstract
Biomedical imaging data presents enormous potential for deep learning models to predict invaluable properties, such as diseases and drug effects. However, unavoidable alterations of the technical conditions cause batch effects: variations between groups of samples that are not due to any biological signal of interest. Batch effects greatly hinder the generalization abilities of deep learning models, preventing their practical use in the real world. Unsupervised Domain Adaptation (UDA) methods have been proposed to mitigate batch effects, but they usually assume that the data is comprised of only one source domain and one target domain, whereas biological datasets are comprised of multiple domains, both at training and at inference time. While Batch Normalization–based test-time and meta-learning adaptation methods offer a promising mechanism for domain alignment, we show that existing approaches exhibit degraded performance under the usual inference scenarios of small target batch sizes and label shift. We address these limitations by leveraging negative control samples, which are consistently present in every experimental batch in biological datasets, as stable context for adaptation. We propose CS-ARM-BN, a meta-learning BN adaptation method that uses controls both during training and inference to stabilize domain statistics. We perform a suite of experiments of Mechanism-Of-Action (MoA) classification, a crucial task for drug discovery, on the large JUMP-CP imaging dataset. Our experiments show that CS-ARM-BN substantially improves robustness to batch size and class distribution shifts, enabling practical use of deep learning models for biomedical images.
Lay Summary
Biomedical images can reveal important information about diseases and how drugs affect cells. Deep learning models could help scientists extract this information automatically, which would be especially useful for drug discovery. However, these models often struggle when images are collected under different experimental conditions. These unwanted technical differences, known as batch effects, can make samples representing the same biological effect look different across experiments. Existing methods can adjust a model when a new batch arrives, but we found that they become unreliable when the new batch contains only a few samples or a different mix of drug effects than the training data. To address this, we use negative control samples, which are images of untreated cells, and are included in every experimental batch, to provide a stable reference point. Our method, CS-ARM-BN, uses these controls to better adjust deep learning models across batches. On a large biomedical imaging dataset, it improves robustness and makes these models more practical for real-world drug discovery.