Global Ground Metric Learning with Applications to scRNA data

2025-06-18

AI Summary

In conventional optimal transport, predefined ground metrics (e.g., Euclidean distance) ignore intrinsic data geometry and feature relevance; existing supervised metric learning methods suffer from poor generalizability, reliance on shared support sets, or costly pointwise annotations. To address these limitations, we propose the first distribution-level supervised framework for learning global ground metrics—requiring only coarse class labels, without pointwise supervision or support-set alignment, and enabling cross-class generalization. Integrating optimal transport theory with metric learning, our approach formulates an end-to-end differentiable model that directly optimizes the Wasserstein distance objective. The learned metric is structure-aware and task-robust. Evaluated on multi-disease single-cell RNA-seq data, it significantly improves embedding quality, clustering purity, and classification accuracy, while yielding biologically interpretable distance structures.

Technology Category

• Machine Learning: Learning with Manifolds
• Search and Optimization: Distributed Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Search and Retrieval-Augmented AI: Web evaluation methodologies and metrics
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Economics, Online Markets and Human Computation: Data quality aspects of human-annotated datasets

Abstract

Optimal transport provides a robust framework for comparing probability distributions. Its effectiveness is significantly influenced by the choice of the underlying ground metric. Traditionally, the ground metric has either been (i) predefined, e.g., as the Euclidean distance, or (ii) learned in a supervised way, by utilizing labeled data to learn a suitable ground metric for enhanced task-specific performance. Yet, predefined metrics typically cannot account for the inherent structure and varying importance of different features in the data, and existing supervised approaches to ground metric learning often do not generalize across multiple classes or are restricted to distributions with shared supports. To address these limitations, we propose a novel approach for learning metrics for arbitrary distributions over a shared metric space. Our method provides a distance between individual points like a global metric, but requires only class labels on a distribution-level for training. The learned global ground metric enables more accurate optimal transport distances, leading to improved performance in embedding, clustering and classification tasks. We demonstrate the effectiveness and interpretability of our approach using patient-level scRNA-seq data spanning multiple diseases.

Problem

• Learning global ground metrics for arbitrary distributions
• Improving optimal transport distance accuracy
• Enhancing performance in embedding, clustering, classification

Innovation

• Learns global ground metric for distributions
• Uses class labels on distribution-level
• Improves optimal transport distance accuracy