[Paper Review] Correcting Nuisance Variation using Wasserstein Distance

Gil Tabak, Minjie Fan+3 more|arXiv (Cornell University)|Nov 2, 2017

Cell Image Analysis Techniques27 references3 citations

TL;DR

This paper proposes a novel framework to correct nuisance variation in image embeddings by minimizing Wasserstein distance between distributional embeddings across experimental domains, preserving biological signal while reducing batch effects. The method improves downstream phenotypic profiling by enhancing k-NN MOA assignment and Silhouette scores without assuming Gaussian distributions.

ABSTRACT

Profiling cellular phenotypes from microscopic imaging can provide meaningful biological information resulting from various factors affecting the cells. One motivating application is drug development: morphological cell features can be captured from images, from which similarities between different drug compounds applied at different doses can be quantified. The general approach is to find a function mapping the images to an embedding space of manageable dimensionality whose geometry captures relevant features of the input images. An important known issue for such methods is separating relevant biological signal from nuisance variation. For example, the embedding vectors tend to be more correlated for cells that were cultured and imaged during the same week than for those from different weeks, despite having identical drug compounds applied in both cases. In this case, the particular batch in which a set of experiments were conducted constitutes the domain of the data; an ideal set of image embeddings should contain only the relevant biological information (e.g. drug effects). We develop a general framework for adjusting the image embeddings in order to `forget' domain-specific information while preserving relevant biological information. To achieve this, we minimize a loss function based on distances between marginal distributions (such as the Wasserstein distance) of embeddings across domains for each replicated treatment. For the dataset we present results with, the only replicated treatment happens to be the negative control treatment, for which we do not expect any treatment-induced cell morphology changes. We find that for our transformed embeddings (i) the underlying geometric structure is not only preserved but the embeddings also carry improved biological signal; and (ii) less domain-specific information is present.

Motivation & Objective

• To address nuisance variation in image embeddings caused by experimental batch effects, such as differences in culture conditions or imaging sessions.
• To develop a domain-invariant embedding transformation that preserves relevant biological signal while removing domain-specific artifacts.
• To enable more accurate phenotypic profiling of drug compounds by ensuring embeddings reflect biological treatment effects rather than experimental variability.
• To provide a flexible, non-parametric alternative to existing batch correction methods that does not assume Gaussian-distributed data.
• To extend batch correction beyond negative controls by leveraging multiple treatment replicates across domains, when available.

Proposed method

• The method uses a minimax optimization framework to learn domain-specific transformations $ A_d $ that align the marginal distributions of embeddings across different experimental domains.
• The loss function is based on the Wasserstein distance between the distributions of embeddings for the same treatment across different domains, promoting domain invariance.
• The Wasserstein distance is approximated using a neural network, enabling the method to handle general, non-Gaussian distributions of embeddings.
• The framework is trained to minimize domain-specific distributional differences while preserving the geometric structure of the embedding space.
• The method implicitly assumes that transformations learned from negative controls can be generalized to other treatments, though it is extensible to multiple treatment replicates.
• Regularization or early stopping is used to prevent embedding collapse during training, ensuring meaningful separation of biologically relevant signals.

Experimental results

Research questions

• RQ1 Can Wasserstein distance be effectively used to align image embeddings across experimental domains while preserving biological signal?
• RQ2 Does the proposed method improve the accuracy of phenotypic profiling tasks such as drug mechanism of action (MOA) prediction?
• RQ3 How does the method compare to existing batch correction techniques in terms of preserving embedding geometry and reducing domain-specific variation?
• RQ4 Can the framework be generalized beyond negative controls to utilize replicates of multiple treatments?
• RQ5 What are the effects of using more complex transformation functions or alternative loss formulations on performance?

Key findings

• The transformed embeddings preserve the underlying geometric structure of the original embedding space, as confirmed by stable k-NN MOA assignment metrics.
• The Silhouette score improved after correction, indicating better clustering of biologically similar treatments.
• The method successfully reduced domain-specific variation, with embeddings for the same treatment becoming more similar across different batches.
• The approach does not require assumptions about Gaussian distributional forms, making it suitable for complex, real-world biological data.
• The method demonstrated improved biological signal retention compared to baseline normalization and existing batch correction methods.
• While the current implementation relies on negative controls for alignment, the framework is extensible to multiple treatment replicates, which could enhance performance in future applications.