[Paper Review] Feature Robust Optimal Transport for High-dimensional Data

Mathis Petrovich, Chao Liang+7 more|arXiv (Cornell University)|May 25, 2020

Advanced Neural Network Applications49 references4 citations

TL;DR

This paper proposes Feature-robust Optimal Transport (FROT), a novel method for high-dimensional data that selects discriminative features via min–max optimization to mitigate the curse of dimensionality. FROT uses a convex formulation solved by a Frank–Wolfe algorithm with Sinkhorn subproblems, achieving state-of-the-art performance in semantic correspondence tasks by identifying key layers in deep neural networks.

ABSTRACT

Optimal transport is a machine learning problem with applications including distribution comparison, feature selection, and generative adversarial networks. In this paper, we propose feature-robust optimal transport (FROT) for high-dimensional data, which solves high-dimensional OT problems using feature selection to avoid the curse of dimensionality. Specifically, we find a transport plan with discriminative features. To this end, we formulate the FROT problem as a min--max optimization problem. We then propose a convex formulation of the FROT problem and solve it using a Frank--Wolfe-based optimization algorithm, whereby the subproblem can be efficiently solved using the Sinkhorn algorithm. Since FROT finds the transport plan from selected features, it is robust to noise features. To show the effectiveness of FROT, we propose using the FROT algorithm for the layer selection problem in deep neural networks for semantic correspondence. By conducting synthetic and benchmark experiments, we demonstrate that the proposed method can find a strong correspondence by determining important layers. We show that the FROT algorithm achieves state-of-the-art performance in real-world semantic correspondence datasets.

Motivation & Objective

• To address the curse of dimensionality in optimal transport (OT) for high-dimensional data, especially in deep neural network feature spaces.
• To develop a method that selects informative feature groups (e.g., network layers) while being robust to noisy or irrelevant features.
• To formulate a convex, tractable optimization problem that enables efficient computation of transport plans in high-dimensional settings.
• To demonstrate the effectiveness of FROT in real-world semantic correspondence tasks using benchmark datasets.
• To show that feature selection via group sparsity improves OT robustness and performance compared to standard OT or prior robust OT methods.

Proposed method

• FROT formulates the optimal transport problem as a min–max optimization over transport plans and feature group selection, where the max step selects the most discriminative features.
• The method introduces a convex relaxation of the min–max problem using group sparsity regularization, enabling efficient optimization.
• A Frank–Wolfe algorithm is employed to solve the convexified FROT problem, with each subproblem solved via the Sinkhorn algorithm for entropic regularized OT.
• Feature groups (e.g., activations from specific DNN layers) are treated as atomic units, and the method selects a subset of these groups to construct the transport plan.
• The regularization parameter η controls the sparsity of selected feature groups, balancing robustness and performance.
• The framework naturally identifies important features (e.g., key network layers) without additional cost, enabling interpretability.

Experimental results

Research questions

• RQ1 Can feature selection via group sparsity improve the robustness of optimal transport in high-dimensional data with noisy features?
• RQ2 Does FROT outperform standard optimal transport and prior robust OT methods in semantic correspondence tasks on real-world datasets?
• RQ3 Can FROT effectively identify important layers in deep neural networks for feature matching, and is this selection interpretable?
• RQ4 How does the choice of regularization parameter η affect the performance and feature selection in FROT?
• RQ5 Is FROT computationally efficient and scalable to high-dimensional data, such as deep network activations?

Key findings

• FROT achieves state-of-the-art performance on the SPair-71k benchmark, with a 35.0% accuracy at η=0.2, outperforming OT (30.1%) and SRW (25.3%).
• At η=0.5, FROT achieves 40.0% accuracy on SPair-71k, significantly outperforming the next best method (SRW with 25.3%).
• FROT with η=0.5 achieves 61.9% accuracy on the PASCAL-Part dataset, surpassing OT (38.0%) and SRW (38.4%).
• The method demonstrates robustness to noise: in synthetic experiments, FROT recovers the true transport plan even when 80% of features are noisy, unlike standard OT.
• FROT successfully identifies key DNN layers for semantic correspondence, with optimal performance achieved at moderate η values (e.g., η=0.5), indicating effective feature group selection.
• The ablation study shows that performance degrades with very high η (e.g., η=5.0), confirming the importance of controlled sparsity in feature selection.