Wasserstein Distances Made Explainable: Insights into Dataset Shifts and Transport Phenomena

TL;DR

The paper addresses explaining what factor drives a Wasserstein distance between distributions, introducing $ abla$Wax, a framework that neuralizes the distance and uses LRP-style backpropagation to attribute the distance to data points and input features. It generalizes to $\mathcal{W}_p$ and Sinkhorn variants, supports subspace-based explanations via $\boldsymbol{U}$-$\nabla$Wax, and demonstrates high fidelity and transport-phenomena insight across synthetic and real datasets. Empirical results show superior attribution quality (SRG) and meaningful transport insights in aging- and dataset-difference use cases, with subspace decompositions revealing interpretable, domain-relevant factors. The work has practical implications for diagnosing dataset shifts, validating transport models, and guiding the design of more informative Wasserstein-based analyses.

Abstract

Wasserstein distances provide a powerful framework for comparing data distributions. They can be used to analyze processes over time or to detect inhomogeneities within data. However, simply calculating the Wasserstein distance or analyzing the corresponding transport map (or coupling) may not be sufficient for understanding what factors contribute to a high or low Wasserstein distance. In this work, we propose a novel solution based on Explainable AI that allows us to efficiently and accurately attribute Wasserstein distances to various data components, including data subgroups, input features, or interpretable subspaces. Our method achieves high accuracy across diverse datasets and Wasserstein distance specifications, and its practical utility is demonstrated in two use cases.

Figures (4)

• Figure 1: High-level illustration of $\mathcal{W}$aX, which goes beyond a classical inspection of the Wasserstein-derived transport map $\gamma^\star$, by attributing the Wasserstein distance on data points and input features. In this example, a fluid moves rightward and is observed at times $t=1$ and $t=2$. This phenomenon is represented by a Wasserstein model ($p=3$, $q=2$). $\mathcal{W}$aX identifies the bottleneck at the center as the primary contributor to the high Wasserstein distance---something that is not easily detectable and quantifiable from analyzing the transport map alone. In Supplementary Note A, we provide further examples highlighting how using $\mathcal{W}$aX and simply analyzing the transport map strongly differ in their sensitivity to different Wasserstein model specifications.
• Figure 2: Diagram of $\mathcal{W}$aX's computational workflow. The Wasserstein distance is viewed as a two-layer network with the coupling fixed to its optimum $\gamma^\star$. Explanation proceeds through a backward pass, leading first to an attribution on instance pairs, and then on input features.
• Figure 3: Insight into an abalone aging process through the $\mathcal{W}$aX and $\boldsymbol{U}$-$\mathcal{W}$aX methods. The Wasserstein distance models the shift in a simulated abalone cohort measured twice at approximately a one-year interval. (a) Instance-wise ($R_{kl}$) and feature-wise attributions ($R_i$) calculated by $\mathcal{W}$aX. (b) Attribution on subspaces (i.e. $R^c_{kl}$ and $R^c_i$) extracted by $\boldsymbol{U}$-$\mathcal{W}$aX ($r=4$). (c) Similar analysis where $\boldsymbol{U}$-$\mathcal{W}$aX is substituted by a simpler clustering baseline kulinskiExplainingDistributionShifts2023 to generate 'subspaces'. Instance-wise attributions are summarized in three data subgroups based on the cluster assignments of (c) and overlaid on the t-SNE plot of the data. Transport arrows in each t-SNE show the coupled sample pairs that are strongest in the subspace, or assigned to the cluster. The bar sizes indicate the contribution of each data component to the Wasserstein distance.
• Figure 4: Shift between the CelebA and LFW datasets as modeled by the Wasserstein distance ($p,q=2$) and characterized by $\boldsymbol{U}$-$\mathcal{W}$aX. We show $\boldsymbol{U}$-$\mathcal{W}$aX with $r=2$ in (a) and with $r=3$ in (b). While $r=2$ provides a high-level/global view of the shift by modeling the main transport trend, $r=3$ reveals diverse local sub-shifts. For ease of visualization, the subspace relevances $R_c$ from \ref{['main:eq:subspace-attribution-concepts']} are provided as an average over the intermediate range when labeled with vertical dots, and $S_\perp$ indicates the unexplained residual relevance. The bars between the coupled samples indicate the sample relevance $R^c_{kl}$ from \ref{['main:eq:subspace-attribution-sample']} within each subspace. For each subspace, we also show the t-SNE plots visualizing the localization and spread of the shifts, and the words with the highest absolute cosine similarity.

Theorems & Definitions (2)

• Proposition 1
• Proposition 2