Optimal transport errors tied to class size not just data layout
Disentangling Marginal Capacity from Geometry in Optimal-Transport Label Transfer
Machine Learning
Summary
Optimal transport is a math tool that matches data from one group to another, but it can struggle when some groups are rare and cause mistakes. The authors found that a key problem is how much 'capacity' each group has to send its labels, which depends on its original size rather than how well it matches the new data. They show adjusting this capacity to better fit the new data's composition can greatly reduce errors in label transferring. Their work separates the effects of 'how much data there is' from 'how well the data matches' in causing errors.
What this means in practice
- •For bioinformatics teams: Improve annotation accuracy in single-cell data analysis by adjusting transport capacities to match target cell-type frequencies.
- •For machine learning engineers: Enhance label transfer between datasets with differing class distributions by tuning transport marginals to better reflect target composition.
Tested on one dataset.
Authors
Ali Khodabandeh Yalabadi, Azim Dehghani Amirabad, Artem Moskalev, Yun-Ching Chen, Mehdi Pirooznia, Scott Oloff, Rui Liao
Abstract
Optimal transport (OT) is widely used to align distributions and transfer labels between unpaired source and target data, but source-target composition shifts can cause rare classes to be systematically under-predicted even with high transport confidence. We identify a source of this failure in the transport marginal: under the standard uniform weighting of source points, each class has transport capacity proportional to its source frequency, which can be insufficient to meet its target demand even when separated well in the embedding space. We show that displaced mass is preferentially absorbed by high-capacity classes rather than geometrically similar classes, producing confident errors that are difficult to detect from the transport plan. Downstream read-outs, entropic smoothing, and marginal relaxation provide limited recovery, whereas matching the transport marginal to the target composition substantially improves label transfer. On single-cell data, the oracle composition-matched marginal raises macro-F1 from 0.322 to 0.623, and the best deployable estimator closes 69% of the oracle gap. Correcting the marginal also separates a coverage-dependent component of error from residual error associated with class separability, showing that transport capacity and representation geometry are distinct sources of transfer error.