Optimal Transport Feature Alignment for Cross-Domain Cell Image Classification

TMLR Paper9358 Authors

Abstract: Classifying single-cell images across datasets collected under varying imaging conditions remains a fundamental challenge in computational cytology. Differences in microscope hardware, illumination, staining, and media introduce substantial domain shifts, leading to poor generalization of models trained on a single dataset—even when underlying cell types are shared. To address this, we propose a robust pipeline that integrates (i) multi-encoder feature extraction using foundation models, (ii) unsupervised identification of shared cell types between datasets, (iii) alignment of embedding spaces using graph-regularized optimal transport (OT), and (iv) final classification via ensemble-based voting across encoders. By explicitly filtering unmatched classes and transporting only comparable subsets, our method mitigates the impact of non-overlapping cell types and structural misalignment. Experimental results across multiple real-world cell image datasets demonstrate improved accuracy (+20% in average over all datasets considered in the paper when adding class filtering and OT), robustness to domain shift, and reliable performance in both fully unsupervised and low-label regimes—without requiring fine-tuning or retraining on the target domain.

Submission Type: Long submission (more than 12 pages of main content)

Previous TMLR Submission Url: https://openreview.net/forum?id=YUYlEFYahx&noteId=YUYlEFYahx

Changes Since Last Submission:

We remove the funding information at the end of the paper.
We also went through the TMLR guidelines to align the paper style with what is required by the journal.

Assigned Action Editor: ~Ilan_Shomorony1

Submission Number: 9358