Point Cloud Classification via Deep Set Linearized Optimal Transport

TL;DR

This paper presents a novel deep learning method that embeds point clouds into an L^2 space using linearized optimal transport, enabling efficient classification while preserving geometric structures.

Contribution

It introduces Deep Set Linearized Optimal Transport, combining ICNNs and Wasserstein distances for improved point cloud embedding and classification.

Findings

Outperforms standard deep set methods on flow cytometry data

Efficient approximation of Wasserstein-2 distances using ICNNs

Effective classification of point clouds with limited labels

Abstract

We introduce Deep Set Linearized Optimal Transport, an algorithm designed for the efficient simultaneous embedding of point clouds into an $L^2-$space. This embedding preserves specific low-dimensional structures within the Wasserstein space while constructing a classifier to distinguish between various classes of point clouds. Our approach is motivated by the observation that $L^2-$distances between optimal transport maps for distinct point clouds, originating from a shared fixed reference distribution, provide an approximation of the Wasserstein-2 distance between these point clouds, under certain assumptions. To learn approximations of these transport maps, we employ input convex neural networks (ICNNs) and establish that, under specific conditions, Euclidean distances between samples from these ICNNs closely mirror Wasserstein-2 distances between the true distributions. Additionally, we train a discriminator network that attaches weights these samples and creates a permutation invariant classifier to differentiate between different classes of point clouds. We showcase the advantages of our algorithm over the standard deep set approach through experiments on a flow cytometry dataset with a limited number of labeled point clouds.