Geometric Sparse Coding in Wasserstein Space

TL;DR

This paper introduces a geometrically sparse regularizer for Wasserstein dictionary learning that produces sparser, more interpretable representations and improves downstream task performance by addressing non-uniqueness and interpretability issues.

Contribution

It proposes a novel regularizer promoting sparsity in Wasserstein space, enhancing interpretability and recovery of generating distributions in Wasserstein dictionary learning.

Findings

Sparser and more interpretable dictionaries are achieved.

Regularizer improves recovery of generating distributions.

Enhanced performance in downstream applications.

Abstract

Wasserstein dictionary learning is an unsupervised approach to learning a collection of probability distributions that generate observed distributions as Wasserstein barycentric combinations. Existing methods for Wasserstein dictionary learning optimize an objective that seeks a dictionary with sufficient representation capacity via barycentric interpolation to approximate the observed training data, but without imposing additional structural properties on the coefficients associated to the dictionary. This leads to dictionaries that densely represent the observed data, which makes interpretation of the coefficients challenging and may also lead to poor empirical performance when using the learned coefficients in downstream tasks. In contrast and motivated by sparse dictionary learning in Euclidean spaces, we propose a geometrically sparse regularizer for Wasserstein space that promotes representations of a data point using only nearby dictionary elements. We show this approach leads to sparse representations in Wasserstein space and addresses the problem of non-uniqueness of barycentric representation. Moreover, when data is generated as Wasserstein barycenters of fixed distributions, this regularizer facilitates the recovery of the generating distributions in cases that are ill-posed for unregularized Wasserstein dictionary learning. Through experimentation on synthetic and real data, we show that our geometrically regularized approach yields sparser and more interpretable dictionaries in Wasserstein space, which perform better in downstream applications.