[Paper Review] InfoOT: Information Maximizing Optimal Transport

Ching-Yao Chuang, Stefanie Jegelka+1 more|arXiv (Cornell University)|Oct 6, 2022

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

InfoOT proposes an information-theoretic extension to optimal transport that maximizes mutual information between source and target domains while minimizing geometric distances, enabling robust, out-of-sample generalization via a novel conditional projection method. It improves alignment quality across domain adaptation, cross-domain retrieval, and single-cell omics benchmarks by preserving cluster structure and handling class imbalance.

ABSTRACT

Optimal transport aligns samples across distributions by minimizing the transportation cost between them, e.g., the geometric distances. Yet, it ignores coherence structure in the data such as clusters, does not handle outliers well, and cannot integrate new data points. To address these drawbacks, we propose InfoOT, an information-theoretic extension of optimal transport that maximizes the mutual information between domains while minimizing geometric distances. The resulting objective can still be formulated as a (generalized) optimal transport problem, and can be efficiently solved by projected gradient descent. This formulation yields a new projection method that is robust to outliers and generalizes to unseen samples. Empirically, InfoOT improves the quality of alignments across benchmarks in domain adaptation, cross-domain retrieval, and single-cell alignment.

Motivation & Objective

• To address the limitations of standard optimal transport in preserving data coherence, handling outliers, and generalizing to unseen samples.
• To integrate global structural information—such as cluster structure—into optimal transport via mutual information maximization.
• To develop a new projection method that enables out-of-sample generalization and robustness to class imbalance and outliers.
• To extend optimal transport to align distributions in different metric spaces, such as those with varying modalities or dimensions.

Proposed method

• InfoOT formulates a generalized optimal transport problem that combines geometric cost with a mutual information term estimated via kernel density estimation.
• The transportation cost is defined as the log ratio of the estimated joint density to the product of marginal densities, capturing dependence between source and target samples.
• A projected gradient descent algorithm efficiently solves the resulting optimization problem, maintaining the structure of optimal transport.
• Conditional projection is introduced as a new mapping method that computes target predictions via importance sampling over source-conditional densities, enabling out-of-sample generalization.
• The method leverages intra-domain distances only, making it applicable to distributions in different metric spaces, similar to Gromov-Wasserstein but with information-theoretic regularization.
• The approach uses kernel density estimation to estimate joint and marginal distributions from paired samples in the transportation plan.

Experimental results

Research questions

• RQ1 Can mutual information be effectively integrated into optimal transport to improve alignment quality and preserve data structure?
• RQ2 How can optimal transport be made robust to outliers and class imbalance in real-world datasets?
• RQ3 Can a new projection method be developed that generalizes to unseen samples, enabling streaming or large-scale applications?
• RQ4 To what extent can InfoOT align distributions across different modalities or dimensionality spaces?

Key findings

• InfoOT significantly outperforms barycentric projection and standard OT baselines in domain adaptation on Office-Caltech, especially under class imbalance, with a 68.9% label transfer accuracy on scGEM.
• On the ImageClef and Office-Caltech datasets, F-InfoOT achieves higher top-k precision than baselines, demonstrating superior performance in unsupervised cross-domain retrieval.
• In single-cell alignment, InfoOT achieves a 0.178 FOSCTTM score on scGEM and 98.8% label transfer accuracy, outperforming UnionCom, MMD-MA, and SCOT (GW).
• The conditional projection method in InfoOT shows better robustness to outliers and class imbalance than barycentric projection, as evidenced by t-SNE visualizations preserving cluster structure.
• InfoOT enables effective out-of-sample generalization, allowing similarity scoring for new source samples without re-solving the full OT problem.
• The method generalizes well across diverse settings, including cross-domain retrieval, domain adaptation, and single-cell multi-omics data integration.