[Paper Review] A Consistent Extension of Discrete Optimal Transport Maps for Machine Learning Applications

Lucas de Lara, Alberto González-Sanz+1 more|arXiv (Cornell University)|Feb 17, 2021

Markov Chains and Monte Carlo Methods25 references4 citations

TL;DR

This paper proposes a statistically consistent method to extend discrete optimal transport maps to new, unseen data by learning a continuous approximation of the transport map from empirical samples. The approach ensures asymptotic convergence to the true continuous transport map, enabling consistent counterfactual explanations in fairness-aware machine learning with theoretical guarantees.

ABSTRACT

Optimal transport maps define a one-to-one correspondence between probability distributions, and as such have grown popular for machine learning applications. However, these maps are generally defined on empirical observations and cannot be generalized to new samples while preserving asymptotic properties. We extend a novel method to learn a consistent estimator of a continuous optimal transport map from two empirical distributions. The consequences of this work are two-fold: first, it enables to extend the transport plan to new observations without computing again the discrete optimal transport map; second, it provides statistical guarantees to machine learning applications of optimal transport. We illustrate the strength of this approach by deriving a consistent framework for transport-based counterfactual explanations in fairness.

Motivation & Objective

• To address the limitation of discrete optimal transport maps, which are only defined on training samples and cannot generalize to new data.
• To develop a statistically consistent estimator of the continuous optimal transport map from empirical distributions.
• To enable the use of optimal transport in machine learning applications requiring generalization, such as counterfactual explanations.
• To provide theoretical guarantees for OT-based fairness auditing and model interpretability.

Proposed method

• The method extends the interpolation framework from del Barrio et al. (2020a) to construct a continuous approximation of the discrete optimal transport map.
• It introduces a T-admissible estimator that converges almost surely to the true Brenier map under mild regularity conditions.
• The approach leverages the cyclic monotonicity of optimal transport maps and convex analysis to ensure consistency.
• It uses the subdifferential of a convex function to represent the transport map and ensures convergence via the dominated convergence theorem.
• The framework incorporates assumptions on classifier separability and negligible discontinuity sets to ensure almost sure convergence of flip probabilities.
• It applies the continuous mapping theorem and dominated convergence to prove consistency of key fairness metrics like flip sets and sign vectors.

Experimental results

Research questions

• RQ1 Can a discrete optimal transport map be consistently extended to new, unseen data points while preserving asymptotic convergence to the true continuous map?
• RQ2 Does the proposed continuous approximation maintain statistical consistency as sample size increases?
• RQ3 Can this consistent framework be applied to derive reliable counterfactual explanations for fairness auditing in binary classifiers?
• RQ4 What conditions ensure the almost sure convergence of flip set probabilities and sign vectors in the context of OT-based explanations?
• RQ5 How can theoretical guarantees be established for OT-based fairness metrics when using empirical estimates?

Key findings

• The proposed method achieves almost sure convergence of the estimated transport map to the true continuous Brenier map as sample size increases.
• The flip set probabilities, including both negative and positive flip sets, converge almost surely to their true values under the continuous map.
• The mean sign vector and reference-based difference metrics also converge almost surely, provided the discontinuity set has measure zero under the source measure.
• The framework ensures statistical consistency for fairness metrics such as Δ⁻_diff, Δ⁺_diff, Δ_ref_diff, Δ⁻_sign, and Δ_ref_sign.
• Theoretical guarantees hold under practical assumptions: classifier separability and negligible discontinuity sets, which are satisfied in most real-world ML applications.
• The method enables consistent, generalizable counterfactual explanations for fairness auditing without re-computing the transport map for each new sample.