Reliable Modeling of Distribution Shifts via Displacement-Reshaped Optimal Transport

TL;DR

ReshapeOT enhances optimal transport models for distribution shifts by reshaping the ground metric with observed displacements, leading to more reliable transport solutions in practical scenarios.

Contribution

It introduces a novel Mahalanobis distance-based reshaping of the ground metric in OT, improving alignment with observed data displacements.

Findings

ReshapeOT achieves significant improvements in transport reliability on synthetic and real data.

The method is computationally lightweight and easily integrable into existing OT frameworks.

ReshapeOT enhances practical applications of OT in modeling distribution shifts.

Abstract

Optimal transport (OT) is a central framework for modeling distribution shifts. Because OT compares distributions directly in input space, a well-designed ground metric between observations is essential to ensure that the optimizer does not violate the true geometry of change. We propose Displacement-Reshaped Optimal Transport (ReshapeOT), a method that reshapes the ground metric by integrating observed sample displacements as an additional source of knowledge. Technically, ReshapeOT replaces the Euclidean metric with a Mahalanobis distance estimated from displacement second moments. This effectively carves expressways through the input space, inviting transport solutions that better align with observed displacements. Our method is computationally lightweight, integrates seamlessly into any OT solver that operates on a cost matrix, and can be kernelized for further flexibility. Experiments on synthetic and real-world data show that ReshapeOT achieves substantial gains in transport reliability. We further demonstrate our method's usefulness in two practical use cases.