Partial Distribution Matching via Partial Wasserstein Adversarial Networks

TL;DR

This work tackles robust distribution matching when data contain heavy outliers by formulating partial distribution matching (PDM) and introducing Partial Wasserstein Adversarial Network (PWAN). PWAN leverages Kantorovich–Rubinstein duality for partial Wasserstein divergences $\\mathcal{L}_{M,m}$ and $\\mathcal{L}_{D,h}$, approximating the optimal potentials with neural networks and performing end-to-end gradient-based optimization. It unifies a mass- and a distance-type perspective, provides a mini-batch training scheme that avoids common OT batch errors, and reveals a direct link to Wasserstein GANs in the full-mass limit while enabling outlier omission in practice. The method is demonstrated on 3D point-set registration and partial domain adaptation, showing strong robustness to outliers, scalability to large datasets, and competitive or superior performance against state-of-the-art approaches. Overall, PWAN offers a principled, scalable framework for robust partial alignment across low- and high-dimensional domains, with practical impact in computer vision and machine learning tasks involving noisy or incomplete data.

Abstract

This paper studies the problem of distribution matching (DM), which is a fundamental machine learning problem seeking to robustly align two probability distributions. Our approach is established on a relaxed formulation, called partial distribution matching (PDM), which seeks to match a fraction of the distributions instead of matching them completely. We theoretically derive the Kantorovich-Rubinstein duality for the partial Wasserstain-1 (PW) discrepancy, and develop a partial Wasserstein adversarial network (PWAN) that efficiently approximates the PW discrepancy based on this dual form. Partial matching can then be achieved by optimizing the network using gradient descent. Two practical tasks, point set registration and partial domain adaptation are investigated, where the goals are to partially match distributions in 3D space and high-dimensional feature space respectively. The experiment results confirm that the proposed PWAN effectively produces highly robust matching results, performing better or on par with the state-of-the-art methods.

Figures (32)

• Figure 1: Comparison between DM and PDM. DM aims to completely match two distributions, but PDM only aims to match a certain fraction of them.
• Figure 2: The computed correspondence $\pi$ between $\alpha$ (blue) and $\beta_\theta$ (red).
• Figure 3: Comparison of the primal form and our approximated KR form on discrete distributions $\alpha$ (blue) and $\beta_\theta$ (red). The solutions to these two forms are presented in the 1-st and 2-nd row respectively, and the gradients of the potentials are presented in the 3-rd row.
• Figure 4: Comparison between PWAN and mini-POT under various batch sizes. $\alpha$ and $\beta$ are visualized on the left panel.
• Figure 5: Comparison of different discrepancies on a toy point set.

Theorems & Definitions (37)

• Theorem 1
• Theorem 2
• Corollary 1
• Corollary 2
• Proposition 1: Dual form of $\mathcal{L}_{M,m}$
• proof
• Proposition 2: Dual form of $\mathcal{L}_{D,h}$
• Proposition 3: KR form of $\mathcal{L}_{M,m}$
• proof
• Definition 1: $\mathbf{L}_{M,m}$ and $\overline{\mathbf{L}_{M,m}}$